聊聊flink的AsyncWaitOperator

@Internal
public class AsyncWaitOperator<IN, OUT>
        extends AbstractUdfStreamOperator<OUT, AsyncFunction<IN, OUT>>
        implements OneInputStreamOperator<IN, OUT>, OperatorActions {
    private static final long serialVersionUID = 1L;

    private static final String STATE_NAME = "_async_wait_operator_state_";

    /** Capacity of the stream element queue. */
    private final int capacity;

    /** Output mode for this operator. */
    private final AsyncDataStream.OutputMode outputMode;

    /** Timeout for the async collectors. */
    private final long timeout;

    protected transient Object checkpointingLock;

    /** {@link TypeSerializer} for inputs while making snapshots. */
    private transient StreamElementSerializer<IN> inStreamElementSerializer;

    /** Recovered input stream elements. */
    private transient ListState<StreamElement> recoveredStreamElements;

    /** Queue to store the currently in-flight stream elements into. */
    private transient StreamElementQueue queue;

    /** Pending stream element which could not yet added to the queue. */
    private transient StreamElementQueueEntry<?> pendingStreamElementQueueEntry;

    private transient ExecutorService executor;

    /** Emitter for the completed stream element queue entries. */
    private transient Emitter<OUT> emitter;

    /** Thread running the emitter. */
    private transient Thread emitterThread;

    public AsyncWaitOperator(
            AsyncFunction<IN, OUT> asyncFunction,
            long timeout,
            int capacity,
            AsyncDataStream.OutputMode outputMode) {
        super(asyncFunction);
        chainingStrategy = ChainingStrategy.ALWAYS;

        Preconditions.checkArgument(capacity > 0, "The number of concurrent async operation should be greater than 0.");
        this.capacity = capacity;

        this.outputMode = Preconditions.checkNotNull(outputMode, "outputMode");

        this.timeout = timeout;
    }

    @Override
    public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) {
        super.setup(containingTask, config, output);

        this.checkpointingLock = getContainingTask().getCheckpointLock();

        this.inStreamElementSerializer = new StreamElementSerializer<>(
            getOperatorConfig().<IN>getTypeSerializerIn1(getUserCodeClassloader()));

        // create the operators executor for the complete operations of the queue entries
        this.executor = Executors.newSingleThreadExecutor();

        switch (outputMode) {
            case ORDERED:
                queue = new OrderedStreamElementQueue(
                    capacity,
                    executor,
                    this);
                break;
            case UNORDERED:
                queue = new UnorderedStreamElementQueue(
                    capacity,
                    executor,
                    this);
                break;
            default:
                throw new IllegalStateException("Unknown async mode: " + outputMode + '.');
        }
    }

    @Override
    public void open() throws Exception {
        super.open();

        // create the emitter
        this.emitter = new Emitter<>(checkpointingLock, output, queue, this);

        // start the emitter thread
        this.emitterThread = new Thread(emitter, "AsyncIO-Emitter-Thread (" + getOperatorName() + ')');
        emitterThread.setDaemon(true);
        emitterThread.start();

        // process stream elements from state, since the Emit thread will start as soon as all
        // elements from previous state are in the StreamElementQueue, we have to make sure that the
        // order to open all operators in the operator chain proceeds from the tail operator to the
        // head operator.
        if (recoveredStreamElements != null) {
            for (StreamElement element : recoveredStreamElements.get()) {
                if (element.isRecord()) {
                    processElement(element.<IN>asRecord());
                }
                else if (element.isWatermark()) {
                    processWatermark(element.asWatermark());
                }
                else if (element.isLatencyMarker()) {
                    processLatencyMarker(element.asLatencyMarker());
                }
                else {
                    throw new IllegalStateException("Unknown record type " + element.getClass() +
                        " encountered while opening the operator.");
                }
            }
            recoveredStreamElements = null;
        }

    }

    @Override
    public void processElement(StreamRecord<IN> element) throws Exception {
        final StreamRecordQueueEntry<OUT> streamRecordBufferEntry = new StreamRecordQueueEntry<>(element);

        if (timeout > 0L) {
            // register a timeout for this AsyncStreamRecordBufferEntry
            long timeoutTimestamp = timeout + getProcessingTimeService().getCurrentProcessingTime();

            final ScheduledFuture<?> timerFuture = getProcessingTimeService().registerTimer(
                timeoutTimestamp,
                new ProcessingTimeCallback() {
                    @Override
                    public void onProcessingTime(long timestamp) throws Exception {
                        userFunction.timeout(element.getValue(), streamRecordBufferEntry);
                    }
                });

            // Cancel the timer once we've completed the stream record buffer entry. This will remove
            // the register trigger task
            streamRecordBufferEntry.onComplete(
                (StreamElementQueueEntry<Collection<OUT>> value) -> {
                    timerFuture.cancel(true);
                },
                executor);
        }

        addAsyncBufferEntry(streamRecordBufferEntry);

        userFunction.asyncInvoke(element.getValue(), streamRecordBufferEntry);
    }

    @Override
    public void processWatermark(Watermark mark) throws Exception {
        WatermarkQueueEntry watermarkBufferEntry = new WatermarkQueueEntry(mark);

        addAsyncBufferEntry(watermarkBufferEntry);
    }

    @Override
    public void snapshotState(StateSnapshotContext context) throws Exception {
        super.snapshotState(context);

        ListState<StreamElement> partitionableState =
            getOperatorStateBackend().getListState(new ListStateDescriptor<>(STATE_NAME, inStreamElementSerializer));
        partitionableState.clear();

        Collection<StreamElementQueueEntry<?>> values = queue.values();

        try {
            for (StreamElementQueueEntry<?> value : values) {
                partitionableState.add(value.getStreamElement());
            }

            // add the pending stream element queue entry if the stream element queue is currently full
            if (pendingStreamElementQueueEntry != null) {
                partitionableState.add(pendingStreamElementQueueEntry.getStreamElement());
            }
        } catch (Exception e) {
            partitionableState.clear();

            throw new Exception("Could not add stream element queue entries to operator state " +
                "backend of operator " + getOperatorName() + '.', e);
        }
    }

    @Override
    public void initializeState(StateInitializationContext context) throws Exception {
        super.initializeState(context);
        recoveredStreamElements = context
            .getOperatorStateStore()
            .getListState(new ListStateDescriptor<>(STATE_NAME, inStreamElementSerializer));

    }

    @Override
    public void close() throws Exception {
        try {
            assert(Thread.holdsLock(checkpointingLock));

            while (!queue.isEmpty()) {
                // wait for the emitter thread to output the remaining elements
                // for that he needs the checkpointing lock and thus we have to free it
                checkpointingLock.wait();
            }
        }
        finally {
            Exception exception = null;

            try {
                super.close();
            } catch (InterruptedException interrupted) {
                exception = interrupted;

                Thread.currentThread().interrupt();
            } catch (Exception e) {
                exception = e;
            }

            try {
                // terminate the emitter, the emitter thread and the executor
                stopResources(true);
            } catch (InterruptedException interrupted) {
                exception = ExceptionUtils.firstOrSuppressed(interrupted, exception);

                Thread.currentThread().interrupt();
            } catch (Exception e) {
                exception = ExceptionUtils.firstOrSuppressed(e, exception);
            }

            if (exception != null) {
                LOG.warn("Errors occurred while closing the AsyncWaitOperator.", exception);
            }
        }
    }

    @Override
    public void dispose() throws Exception {
        Exception exception = null;

        try {
            super.dispose();
        } catch (InterruptedException interrupted) {
            exception = interrupted;

            Thread.currentThread().interrupt();
        } catch (Exception e) {
            exception = e;
        }

        try {
            stopResources(false);
        } catch (InterruptedException interrupted) {
            exception = ExceptionUtils.firstOrSuppressed(interrupted, exception);

            Thread.currentThread().interrupt();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }

        if (exception != null) {
            throw exception;
        }
    }

    private void stopResources(boolean waitForShutdown) throws InterruptedException {
        emitter.stop();
        emitterThread.interrupt();

        executor.shutdown();

        if (waitForShutdown) {
            try {
                if (!executor.awaitTermination(365L, TimeUnit.DAYS)) {
                    executor.shutdownNow();
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();

                Thread.currentThread().interrupt();
            }

            /*
             * FLINK-5638: If we have the checkpoint lock we might have to free it for a while so
             * that the emitter thread can complete/react to the interrupt signal.
             */
            if (Thread.holdsLock(checkpointingLock)) {
                while (emitterThread.isAlive()) {
                    checkpointingLock.wait(100L);
                }
            }

            emitterThread.join();
        } else {
            executor.shutdownNow();
        }
    }

    private <T> void addAsyncBufferEntry(StreamElementQueueEntry<T> streamElementQueueEntry) throws InterruptedException {
        assert(Thread.holdsLock(checkpointingLock));

        pendingStreamElementQueueEntry = streamElementQueueEntry;

        while (!queue.tryPut(streamElementQueueEntry)) {
            // we wait for the emitter to notify us if the queue has space left again
            checkpointingLock.wait();
        }

        pendingStreamElementQueueEntry = null;
    }

    @Override
    public void failOperator(Throwable throwable) {
        getContainingTask().getEnvironment().failExternally(throwable);
    }
}

@Internal
public class Emitter<OUT> implements Runnable {

    private static final Logger LOG = LoggerFactory.getLogger(Emitter.class);

    /** Lock to hold before outputting. */
    private final Object checkpointLock;

    /** Output for the watermark elements. */
    private final Output<StreamRecord<OUT>> output;

    /** Queue to consume the async results from. */
    private final StreamElementQueue streamElementQueue;

    private final OperatorActions operatorActions;

    /** Output for stream records. */
    private final TimestampedCollector<OUT> timestampedCollector;

    private volatile boolean running;

    public Emitter(
            final Object checkpointLock,
            final Output<StreamRecord<OUT>> output,
            final StreamElementQueue streamElementQueue,
            final OperatorActions operatorActions) {

        this.checkpointLock = Preconditions.checkNotNull(checkpointLock, "checkpointLock");
        this.output = Preconditions.checkNotNull(output, "output");
        this.streamElementQueue = Preconditions.checkNotNull(streamElementQueue, "streamElementQueue");
        this.operatorActions = Preconditions.checkNotNull(operatorActions, "operatorActions");

        this.timestampedCollector = new TimestampedCollector<>(this.output);
        this.running = true;
    }

    @Override
    public void run() {
        try {
            while (running) {
                LOG.debug("Wait for next completed async stream element result.");
                AsyncResult streamElementEntry = streamElementQueue.peekBlockingly();

                output(streamElementEntry);
            }
        } catch (InterruptedException e) {
            if (running) {
                operatorActions.failOperator(e);
            } else {
                // Thread got interrupted which means that it should shut down
                LOG.debug("Emitter thread got interrupted, shutting down.");
            }
        } catch (Throwable t) {
            operatorActions.failOperator(new Exception("AsyncWaitOperator's emitter caught an " +
                "unexpected throwable.", t));
        }
    }

    private void output(AsyncResult asyncResult) throws InterruptedException {
        if (asyncResult.isWatermark()) {
            synchronized (checkpointLock) {
                AsyncWatermarkResult asyncWatermarkResult = asyncResult.asWatermark();

                LOG.debug("Output async watermark.");
                output.emitWatermark(asyncWatermarkResult.getWatermark());

                // remove the peeked element from the async collector buffer so that it is no longer
                // checkpointed
                streamElementQueue.poll();

                // notify the main thread that there is again space left in the async collector
                // buffer
                checkpointLock.notifyAll();
            }
        } else {
            AsyncCollectionResult<OUT> streamRecordResult = asyncResult.asResultCollection();

            if (streamRecordResult.hasTimestamp()) {
                timestampedCollector.setAbsoluteTimestamp(streamRecordResult.getTimestamp());
            } else {
                timestampedCollector.eraseTimestamp();
            }

            synchronized (checkpointLock) {
                LOG.debug("Output async stream element collection result.");

                try {
                    Collection<OUT> resultCollection = streamRecordResult.get();

                    if (resultCollection != null) {
                        for (OUT result : resultCollection) {
                            timestampedCollector.collect(result);
                        }
                    }
                } catch (Exception e) {
                    operatorActions.failOperator(
                        new Exception("An async function call terminated with an exception. " +
                            "Failing the AsyncWaitOperator.", e));
                }

                // remove the peeked element from the async collector buffer so that it is no longer
                // checkpointed
                streamElementQueue.poll();

                // notify the main thread that there is again space left in the async collector
                // buffer
                checkpointLock.notifyAll();
            }
        }
    }

    public void stop() {
        running = false;
    }
}

@Internal
public interface StreamElementQueue {

    <T> void put(StreamElementQueueEntry<T> streamElementQueueEntry) throws InterruptedException;

    <T> boolean tryPut(StreamElementQueueEntry<T> streamElementQueueEntry) throws InterruptedException;

    AsyncResult peekBlockingly() throws InterruptedException;

    AsyncResult poll() throws InterruptedException;

    Collection<StreamElementQueueEntry<?>> values() throws InterruptedException;

    boolean isEmpty();

    int size();
}

@Internal
public interface AsyncResult {

    boolean isWatermark();

    boolean isResultCollection();

    AsyncWatermarkResult asWatermark();

    <T> AsyncCollectionResult<T> asResultCollection();
}

@Internal
public abstract class StreamElementQueueEntry<T> implements AsyncResult {

    private final StreamElement streamElement;

    public StreamElementQueueEntry(StreamElement streamElement) {
        this.streamElement = Preconditions.checkNotNull(streamElement);
    }

    public StreamElement getStreamElement() {
        return streamElement;
    }

    public boolean isDone() {
        return getFuture().isDone();
    }

    public void onComplete(
            final Consumer<StreamElementQueueEntry<T>> completeFunction,
            Executor executor) {
        final StreamElementQueueEntry<T> thisReference = this;

        getFuture().whenCompleteAsync(
            // call the complete function for normal completion as well as exceptional completion
            // see FLINK-6435
            (value, throwable) -> completeFunction.accept(thisReference),
            executor);
    }

    protected abstract CompletableFuture<T> getFuture();

    @Override
    public final boolean isWatermark() {
        return AsyncWatermarkResult.class.isAssignableFrom(getClass());
    }

    @Override
    public final boolean isResultCollection() {
        return AsyncCollectionResult.class.isAssignableFrom(getClass());
    }

    @Override
    public final AsyncWatermarkResult asWatermark() {
        return (AsyncWatermarkResult) this;
    }

    @Override
    public final <T> AsyncCollectionResult<T> asResultCollection() {
        return (AsyncCollectionResult<T>) this;
    }
}