org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.flink.connector.base.source.reader.synchronization;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.util.FlinkRuntimeException;
import javax.annotation.concurrent.GuardedBy;
import java.lang.reflect.Field;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import static org.apache.flink.util.Preconditions.checkArgument;
/**
* A custom implementation of blocking queue in combination with a {@link CompletableFuture} that is
* used in the hand-over of data from a producing thread to a consuming thread. This
* FutureCompletingBlockingQueue has the following features:
*
* Consumer Notifications
*
* Rather than letting consumers block on the {@link #take()} method, or have them poll the
* {@link #poll()} method, this queue offers a {@link CompletableFuture}, obtained via the {@link
* #getAvailabilityFuture()} method) that gets completed whenever the queue is non-empty. A consumer
* can thus subscribe to asynchronous notifications for availability by adding a handler to the
* obtained {@code CompletableFuture}.
*
*
The future may also be completed by an explicit call to {@link #notifyAvailable()}. That way
* the consumer may be notified of a situation/condition without adding an element to the queue.
*
*
Availability is reset when a call to {@link #poll()} (or {@link #take()} finds an empty queue
* or results in an empty queue (takes the last element).
*
*
Note that this model generally assumes that false positives are okay, meaning that the
* availability future completes despite there being no data availabile in the queue. The consumer
* is responsible for polling data and obtaining another future to wait on. This is similar to the
* way that Java's Monitors and Conditions can have the spurious wakeup of the waiting
* threads and commonly need to be used in loop with the waiting condition.
*
*
Producer Wakeup
*
* The queue supports gracefully waking up producing threads that are blocked due to the queue
* capacity limits, without interrupting the thread. This is done via the {@link
* #wakeUpPuttingThread(int)} method.
*
* @param the type of the elements in the queue.
*/
public class FutureCompletingBlockingQueue {
/**
* A constant future that is complete, indicating availability. Using this constant in cases
* that are guaranteed available helps short-circuiting some checks and avoiding volatile memory
* operations.
*/
public static final CompletableFuture AVAILABLE = getAvailableFuture();
/** The default capacity for the queue. */
private static final int DEFAULT_CAPACITY = 2;
// ------------------------------------------------------------------------
/** The maximum capacity of the queue. */
private final int capacity;
/**
* The availability future. This doubles as a "non empty" condition. This value is never null.
*/
private CompletableFuture currentFuture;
/** The lock for synchronization. */
private final Lock lock;
/** The element queue. */
@GuardedBy("lock")
private final Queue queue;
/** The per-thread conditions that are waiting on putting elements. */
@GuardedBy("lock")
private final Queue notFull;
/** The per-thread conditions and wakeUp flags. */
@GuardedBy("lock")
private ConditionAndFlag[] putConditionAndFlags;
public FutureCompletingBlockingQueue() {
this(DEFAULT_CAPACITY);
}
public FutureCompletingBlockingQueue(int capacity) {
checkArgument(capacity > 0, "capacity must be > 0");
this.capacity = capacity;
this.queue = new ArrayDeque<>(capacity);
this.lock = new ReentrantLock();
this.putConditionAndFlags = new ConditionAndFlag[1];
this.notFull = new ArrayDeque<>();
// initially the queue is empty and thus unavailable
this.currentFuture = new CompletableFuture<>();
}
// ------------------------------------------------------------------------
// Future / Notification logic
// ------------------------------------------------------------------------
/**
* Returns the availability future. If the queue is non-empty, then this future will already be
* complete. Otherwise the obtained future is guaranteed to get completed the next time the
* queue becomes non-empty, or a notification happens via {@link #notifyAvailable()}.
*
* It is important that a completed future is no guarantee that the next call to {@link
* #poll()} will return a non-null element. If there are concurrent consumer, another consumer
* may have taken the available element. Or there was no element in the first place, because the
* future was completed through a call to {@link #notifyAvailable()}.
*
*
For that reason, it is important to call this method (to obtain a new future) every time
* again after {@link #poll()} returned null and you want to wait for data.
*/
public CompletableFuture getAvailabilityFuture() {
return currentFuture;
}
/**
* Makes sure the availability future is complete, if it is not complete already. All futures
* returned by previous calls to {@link #getAvailabilityFuture()} are guaranteed to be
* completed.
*
* All future calls to the method will return a completed future, until the point that the
* availability is reset via calls to {@link #poll()} that leave the queue empty.
*/
public void notifyAvailable() {
lock.lock();
try {
moveToAvailable();
} finally {
lock.unlock();
}
}
/** Internal utility to make sure that the current future futures are complete (until reset). */
@GuardedBy("lock")
private void moveToAvailable() {
final CompletableFuture current = currentFuture;
if (current != AVAILABLE) {
currentFuture = AVAILABLE;
current.complete(null);
}
}
/** Makes sure the availability future is incomplete, if it was complete before. */
@GuardedBy("lock")
private void moveToUnAvailable() {
if (currentFuture == AVAILABLE) {
currentFuture = new CompletableFuture<>();
}
}
// ------------------------------------------------------------------------
// Blocking Queue Logic
// ------------------------------------------------------------------------
/**
* Put an element into the queue. The thread blocks if the queue is full.
*
* @param threadIndex the index of the thread.
* @param element the element to put.
* @return true if the element has been successfully put into the queue, false otherwise.
* @throws InterruptedException when the thread is interrupted.
*/
public boolean put(int threadIndex, T element) throws InterruptedException {
if (element == null) {
throw new NullPointerException();
}
lock.lockInterruptibly();
try {
while (queue.size() >= capacity) {
if (getAndResetWakeUpFlag(threadIndex)) {
return false;
}
waitOnPut(threadIndex);
}
enqueue(element);
return true;
} finally {
lock.unlock();
}
}
/**
* Warning: This is a dangerous method and should only be used for testing convenience. A
* method that blocks until availability does not go together well with the concept of
* asynchronous notifications and non-blocking polling.
*
* Get and remove the first element from the queue. The call blocks if the queue is empty.
* The problem with this method is that it may loop internally until an element is available and
* that way eagerly reset the availability future. If a consumer thread is blocked in taking an
* element, it will receive availability notifications from {@link #notifyAvailable()} and
* immediately reset them by calling {@link #poll()} and finding the queue empty.
*
* @return the first element in the queue.
* @throws InterruptedException when the thread is interrupted.
*/
@VisibleForTesting
public T take() throws InterruptedException {
T next;
while ((next = poll()) == null) {
// use the future to wait for availability to avoid busy waiting
try {
getAvailabilityFuture().get();
} catch (ExecutionException | CompletionException e) {
// this should never happen, but we propagate just in case
throw new FlinkRuntimeException("exception in queue future completion", e);
}
}
return next;
}
/**
* Get and remove the first element from the queue. Null is returned if the queue is empty. If
* this makes the queue empty (takes the last element) or finds the queue already empty, then
* this resets the availability notifications. The next call to {@link #getAvailabilityFuture()}
* will then return a non-complete future that completes only the next time that the queue
* becomes non-empty or the {@link #notifyAvailable()} method is called.
*
* @return the first element from the queue, or Null if the queue is empty.
*/
public T poll() {
lock.lock();
try {
if (queue.size() == 0) {
moveToUnAvailable();
return null;
}
return dequeue();
} finally {
lock.unlock();
}
}
/**
* Get the first element from the queue without removing it.
*
* @return the first element in the queue, or Null if the queue is empty.
*/
public T peek() {
lock.lock();
try {
return queue.peek();
} finally {
lock.unlock();
}
}
/** Gets the size of the queue. */
public int size() {
lock.lock();
try {
return queue.size();
} finally {
lock.unlock();
}
}
/** Checks whether the queue is empty. */
public boolean isEmpty() {
lock.lock();
try {
return queue.isEmpty();
} finally {
lock.unlock();
}
}
/**
* Checks the remaining capacity in the queue. That is the difference between the maximum
* capacity and the current number of elements in the queue.
*/
public int remainingCapacity() {
lock.lock();
try {
return capacity - queue.size();
} finally {
lock.unlock();
}
}
/**
* Gracefully wakes up the thread with the given {@code threadIndex} if it is blocked in adding
* an element. to the queue. If the thread is blocked in {@link #put(int, Object)} it will
* immediately return from the method with a return value of false.
*
*
If this method is called, the next time the thread with the given index is about to be
* blocked in adding an element, it may immediately wake up and return.
*
* @param threadIndex The number identifying the thread.
*/
public void wakeUpPuttingThread(int threadIndex) {
lock.lock();
try {
maybeCreateCondition(threadIndex);
ConditionAndFlag caf = putConditionAndFlags[threadIndex];
if (caf != null) {
caf.setWakeUp(true);
caf.condition().signal();
}
} finally {
lock.unlock();
}
}
// --------------- private helpers -------------------------
@GuardedBy("lock")
private void enqueue(T element) {
final int sizeBefore = queue.size();
queue.add(element);
if (sizeBefore == 0) {
moveToAvailable();
}
if (sizeBefore < capacity - 1 && !notFull.isEmpty()) {
signalNextPutter();
}
}
@GuardedBy("lock")
private T dequeue() {
final int sizeBefore = queue.size();
final T element = queue.poll();
if (sizeBefore == capacity && !notFull.isEmpty()) {
signalNextPutter();
}
if (queue.isEmpty()) {
moveToUnAvailable();
}
return element;
}
@GuardedBy("lock")
private void waitOnPut(int fetcherIndex) throws InterruptedException {
maybeCreateCondition(fetcherIndex);
Condition cond = putConditionAndFlags[fetcherIndex].condition();
notFull.add(cond);
cond.await();
}
@GuardedBy("lock")
private void signalNextPutter() {
if (!notFull.isEmpty()) {
notFull.poll().signal();
}
}
@GuardedBy("lock")
private void maybeCreateCondition(int threadIndex) {
if (putConditionAndFlags.length < threadIndex + 1) {
putConditionAndFlags = Arrays.copyOf(putConditionAndFlags, threadIndex + 1);
}
if (putConditionAndFlags[threadIndex] == null) {
putConditionAndFlags[threadIndex] = new ConditionAndFlag(lock.newCondition());
}
}
@GuardedBy("lock")
private boolean getAndResetWakeUpFlag(int threadIndex) {
maybeCreateCondition(threadIndex);
if (putConditionAndFlags[threadIndex].getWakeUp()) {
putConditionAndFlags[threadIndex].setWakeUp(false);
return true;
}
return false;
}
// --------------- private per thread state ------------
private static class ConditionAndFlag {
private final Condition cond;
private boolean wakeUp;
private ConditionAndFlag(Condition cond) {
this.cond = cond;
this.wakeUp = false;
}
private Condition condition() {
return cond;
}
private boolean getWakeUp() {
return wakeUp;
}
private void setWakeUp(boolean value) {
wakeUp = value;
}
}
// ------------------------------------------------------------------------
// utilities
// ------------------------------------------------------------------------
@SuppressWarnings("unchecked")
private static CompletableFuture getAvailableFuture() {
// this is a way to obtain the AvailabilityProvider.AVAILABLE future until we decide to
// move the class from the runtime module to the core module
try {
final Class clazz =
Class.forName("org.apache.flink.runtime.io.AvailabilityProvider");
final Field field = clazz.getDeclaredField("AVAILABLE");
return (CompletableFuture) field.get(null);
} catch (Throwable t) {
return CompletableFuture.completedFuture(null);
}
}
}