Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.airlift.concurrent;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Ticker;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.ThreadSafe;
import jakarta.annotation.Nullable;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.ToLongFunction;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Verify.verify;
import static java.util.Objects.requireNonNull;
import static java.util.Objects.requireNonNullElseGet;
import static java.util.concurrent.TimeUnit.NANOSECONDS;
/**
* Size constrained queue that utilizes a dynamic element size function. To prevent
* starvation for adding large elements, the queue will only block new elements if
* the total size has already been reached or exceeded. This means in the normal
* case, the queue should be no larger than the max size plus the size of one element.
* Callers also have the additional option to force insert further elements without
* regard for size constraints. In the current implementation, elements are required
* to have positive sizes (they cannot have zero size). This implementation is designed
* to closely mirror the method signatures of {@link java.util.concurrent.BlockingQueue}.
*/
@ThreadSafe
public class DynamicSizeBoundQueue
{
private final AtomicLong size = new AtomicLong();
private final Queue> queue = new ConcurrentLinkedQueue<>();
private final AtomicReference> enqueueFuture = new AtomicReference<>();
private final AtomicReference> dequeueFuture = new AtomicReference<>();
private final long maxSize;
private final ToLongFunction elementSizeFunction;
private final Ticker ticker;
public DynamicSizeBoundQueue(long maxSize, ToLongFunction elementSizeFunction)
{
this(maxSize, elementSizeFunction, Ticker.systemTicker());
}
public DynamicSizeBoundQueue(long maxSize, ToLongFunction elementSizeFunction, Ticker ticker)
{
checkArgument(maxSize > 0, "maxSize must be positive");
this.maxSize = maxSize;
this.elementSizeFunction = requireNonNull(elementSizeFunction, "elementSizeFunction is null");
this.ticker = requireNonNull(ticker, "ticker is null");
}
public long getMaxSize()
{
return maxSize;
}
/**
* Gets the current size of the queue. The size is guaranteed to be no larger than max size plus
* the size of one element if {@link DynamicSizeBoundQueue#forcePut(Object)} is not used.
*/
public long getSize()
{
return size.get();
}
public boolean offer(T element)
{
long elementSize = elementSizeFunction.applyAsLong(element);
return offer(element, elementSize);
}
private boolean offer(T element, long elementSize)
{
requireNonNull(element, "element is null");
checkArgument(elementSize > 0, "element size must be positive");
if (!tryAcquireSizeReservation(elementSize)) {
return false;
}
queue.add(new ElementAndSize<>(element, elementSize));
notifyIfNecessary(enqueueFuture);
return true;
}
private boolean tryAcquireSizeReservation(long elementSize)
{
// Add the element as long as there is any space available
if (size.get() >= maxSize) {
return false;
}
long newSize;
try {
newSize = getAndAddOverflowChecked(size, elementSize);
}
catch (ArithmeticException e) { // Numeric overflow
// While numeric overflow is extremely unlikely given typical numerical sizes,
// even the largest possible element of size Long.MAX_VALUE can eventually fit
// without numeric overflow as long as the queue can be emptied.
return false;
}
if (newSize >= maxSize) {
verify(size.addAndGet(-elementSize) >= 0);
return false;
}
return true;
}
/**
* Version of {@link java.util.concurrent.atomic.AtomicLong#getAndAdd} that throws {@link ArithmeticException}
* on numeric overflow. This is slightly less efficient than the normal getAndAdd (which often has intrinsic
* support). If this ever becomes a performance bottleneck, it is possible to use the original getAndAdd if
* the caller can guarantee no risk of numeric overflow.
*/
private static long getAndAddOverflowChecked(AtomicLong atomicLong, long delta)
{
return atomicLong.getAndAccumulate(delta, Math::addExact);
}
public boolean offer(T element, long timeout, TimeUnit unit)
throws InterruptedException
{
long elementSize = elementSizeFunction.applyAsLong(element);
long remainingTimeoutNs = unit.toNanos(timeout);
while (!offer(element, elementSize)) {
ListenableFuture future = getOrCreateFuture(dequeueFuture);
// Check again in case we already missed the relevant dequeue event
if (offer(element, elementSize)) {
break;
}
long startTimeNs = ticker.read();
if (remainingTimeoutNs <= 0 || !awaitDequeueFuture(future, remainingTimeoutNs, NANOSECONDS)) {
// Timed out
return false;
}
remainingTimeoutNs -= ticker.read() - startTimeNs;
}
return true;
}
public void put(T element)
throws InterruptedException
{
long elementSize = elementSizeFunction.applyAsLong(element);
while (!offer(element, elementSize)) {
ListenableFuture future = getOrCreateFuture(dequeueFuture);
// Check again in case we already missed the relevant dequeue event
if (offer(element, elementSize)) {
break;
}
awaitDequeueFuture(future);
}
}
/**
* Enqueue the element if there is space, otherwise returns a ListenableFuture that will complete
* when space becomes available for the element. If a future is returned, the element was not inserted.
*/
public Optional> offerWithBackoff(T element)
{
long elementSize = elementSizeFunction.applyAsLong(element);
if (offer(element, elementSize)) {
return Optional.empty();
}
ListenableFuture future = getOrCreateFuture(dequeueFuture);
// Check again in case we already missed the relevant dequeue event
if (offer(element, elementSize)) {
return Optional.empty();
}
return Optional.of(Futures.nonCancellationPropagating(future));
}
/**
* Insert without regard to the max size (potentially exceeding the max limit). This can throw an
* {@link IllegalStateException} if the forced element triggers a numeric overflow, in which case
* the element is not inserted.
*/
public void forcePut(T element)
{
long elementSize = elementSizeFunction.applyAsLong(element);
checkArgument(elementSize > 0, "element size must be positive");
try {
getAndAddOverflowChecked(size, elementSize);
}
catch (ArithmeticException e) { // Numeric overflow
throw new IllegalStateException("Forced element triggered queue size numeric overflow");
}
queue.add(new ElementAndSize<>(element, elementSize));
notifyIfNecessary(enqueueFuture);
}
@Nullable
public T poll()
{
ElementAndSize elementAndSize = queue.poll();
if (elementAndSize == null) {
return null;
}
verify(size.addAndGet(-elementAndSize.size()) >= 0);
notifyIfNecessary(dequeueFuture);
return elementAndSize.element();
}
public T poll(long timeout, TimeUnit unit)
throws InterruptedException
{
long remainingTimeoutNs = unit.toNanos(timeout);
while (true) {
T element = poll();
if (element != null) {
return element;
}
ListenableFuture future = getOrCreateFuture(enqueueFuture);
// Check again in case we already missed the relevant enqueue event
element = poll();
if (element != null) {
return element;
}
long startTimeNs = ticker.read();
if (remainingTimeoutNs <= 0 || !awaitEnqueueFuture(future, remainingTimeoutNs, NANOSECONDS)) {
// Timed out
return null;
}
remainingTimeoutNs -= ticker.read() - startTimeNs;
}
}
public T take()
throws InterruptedException
{
while (true) {
T element = poll();
if (element != null) {
return element;
}
ListenableFuture future = getOrCreateFuture(enqueueFuture);
// Check again in case we already missed the relevant enqueue event
element = poll();
if (element != null) {
return element;
}
awaitEnqueueFuture(future);
}
}
private static ListenableFuture getOrCreateFuture(AtomicReference> reference)
{
return reference.updateAndGet(current -> requireNonNullElseGet(current, SettableFuture::create));
}
private static void notifyIfNecessary(AtomicReference> reference)
{
SettableFuture future = reference.getAndSet(null);
if (future != null) {
future.set(null);
}
}
@VisibleForTesting
void preEnqueueAwaitHook() {}
@VisibleForTesting
void preDequeueAwaitHook() {}
private void awaitDequeueFuture(Future future)
throws InterruptedException
{
preDequeueAwaitHook();
awaitFutureUnchecked(future);
}
private boolean awaitDequeueFuture(Future future, long timeout, TimeUnit timeUnit)
throws InterruptedException
{
preDequeueAwaitHook();
return awaitFutureUnchecked(future, timeout, timeUnit);
}
private void awaitEnqueueFuture(Future future)
throws InterruptedException
{
preEnqueueAwaitHook();
awaitFutureUnchecked(future);
}
private boolean awaitEnqueueFuture(Future future, long timeout, TimeUnit timeUnit)
throws InterruptedException
{
preEnqueueAwaitHook();
return awaitFutureUnchecked(future, timeout, timeUnit);
}
private static void awaitFutureUnchecked(Future future)
throws InterruptedException
{
try {
future.get();
}
catch (ExecutionException e) {
throw new RuntimeException(e);
}
}
private static boolean awaitFutureUnchecked(Future future, long timeout, TimeUnit timeUnit)
throws InterruptedException
{
try {
future.get(timeout, timeUnit);
return true;
}
catch (ExecutionException e) {
throw new RuntimeException(e);
}
catch (TimeoutException e) {
return false;
}
}
private record ElementAndSize(T element, long size)
{
private ElementAndSize
{
requireNonNull(element, "element is null");
checkArgument(size > 0, "size must be positive");
}
}
}
