software.amazon.awssdk.utils.async.SimplePublisher Maven / Gradle / Ivy
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/apache2.0
*
* or in the "license" file accompanying this file. This file 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 software.amazon.awssdk.utils.async;
import static software.amazon.awssdk.utils.async.SimplePublisher.QueueEntry.Type.CANCEL;
import static software.amazon.awssdk.utils.async.SimplePublisher.QueueEntry.Type.ON_COMPLETE;
import static software.amazon.awssdk.utils.async.SimplePublisher.QueueEntry.Type.ON_ERROR;
import static software.amazon.awssdk.utils.async.SimplePublisher.QueueEntry.Type.ON_NEXT;
import java.util.Queue;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;
import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import software.amazon.awssdk.annotations.SdkProtectedApi;
import software.amazon.awssdk.utils.Logger;
import software.amazon.awssdk.utils.Validate;
/**
* A {@link Publisher} to which callers can {@link #send(Object)} messages, simplifying the process of implementing a publisher.
*
* Operations
*
*
The {@code SimplePublisher} supports three simplified operations:
*
* - {@link #send(Object)} for sending messages
* - {@link #complete()} for indicating the successful end of messages
* - {@link #error(Throwable)} for indicating the unsuccessful end of messages
*
*
* Each of these operations returns a {@link CompletableFuture} for indicating when the message has been successfully sent.
*
* Callers are expected to invoke a series of {@link #send(Object)}s followed by a single {@link #complete()} or
* {@link #error(Throwable)}. See the documentation on each operation for more details.
*
*
This publisher will store an unbounded number of messages. It is recommended that callers limit the number of in-flight
* {@link #send(Object)} operations in order to bound the amount of memory used by this publisher.
*/
@SdkProtectedApi
public final class SimplePublisher implements Publisher {
private static final Logger log = Logger.loggerFor(SimplePublisher.class);
/**
* Track the amount of outstanding demand requested by the active subscriber.
*/
private final AtomicLong outstandingDemand = new AtomicLong();
/**
* The queue of events to be processed, in the order they should be processed. These events are lower priority than those
* in {@link #highPriorityQueue} and will be processed after that queue is empty.
*
* All logic within this publisher is represented using events in this queue. This ensures proper ordering of events
* processing and simplified reasoning about thread safety.
*/
private final Queue> standardPriorityQueue = new ConcurrentLinkedQueue<>();
/**
* The queue of events to be processed, in the order they should be processed. These events are higher priority than those
* in {@link #standardPriorityQueue} and will be processed first.
*
* Events are written to this queue to "skip the line" in processing, so it's typically reserved for terminal events,
* like subscription cancellation.
*
*
All logic within this publisher is represented using events in this queue. This ensures proper ordering of events
* processing and simplified reasoning about thread safety.
*/
private final Queue> highPriorityQueue = new ConcurrentLinkedQueue<>();
/**
* Whether the {@link #standardPriorityQueue} and {@link #highPriorityQueue}s are currently being processed. Only one thread
* may read events from the queues at a time.
*/
private final AtomicBoolean processingQueue = new AtomicBoolean(false);
/**
* The failure message that should be sent to future events.
*/
private final FailureMessage failureMessage = new FailureMessage();
/**
* The subscriber provided via {@link #subscribe(Subscriber)}. This publisher only supports a single subscriber.
*/
private Subscriber subscriber;
/**
* Send a message using this publisher.
*
* Messages sent using this publisher will eventually be sent to a downstream subscriber, in the order they were
* written. When the message is sent to the subscriber, the returned future will be completed successfully.
*
*
This method may be invoked concurrently when the order of messages is not important.
*
*
In the time between when this method is invoked and the returned future is not completed, this publisher stores the
* request message in memory. Callers are recommended to limit the number of sends in progress at a time to bound the
* amount of memory used by this publisher.
*
*
The returned future will be completed exceptionally if the downstream subscriber cancels the subscription, or
* if the {@code send} call was performed after a {@link #complete()} or {@link #error(Throwable)} call.
*
* @param value The message to send. Must not be null.
* @return A future that is completed when the message is sent to the subscriber.
*/
public CompletableFuture send(T value) {
log.trace(() -> "Received send() with " + value);
OnNextQueueEntry entry = new OnNextQueueEntry<>(value);
try {
Validate.notNull(value, "Null cannot be written.");
standardPriorityQueue.add(entry);
processEventQueue();
} catch (RuntimeException t) {
entry.resultFuture.completeExceptionally(t);
}
return entry.resultFuture;
}
/**
* Indicate that no more {@link #send(Object)} calls will be made, and that stream of messages is completed successfully.
*
* This can be called before any in-flight {@code send} calls are complete. Such messages will be processed before the
* stream is treated as complete. The returned future will be completed successfully when the {@code complete} is sent to
* the downstream subscriber.
*
*
After this method is invoked, any future {@link #send(Object)}, {@code complete()} or {@link #error(Throwable)}
* calls will be completed exceptionally and not be processed.
*
*
The returned future will be completed exceptionally if the downstream subscriber cancels the subscription, or
* if the {@code complete} call was performed after a {@code complete} or {@link #error(Throwable)} call.
*
* @return A future that is completed when the complete has been sent to the downstream subscriber.
*/
public CompletableFuture complete() {
log.trace(() -> "Received complete()");
OnCompleteQueueEntry entry = new OnCompleteQueueEntry<>();
try {
standardPriorityQueue.add(entry);
processEventQueue();
} catch (RuntimeException t) {
entry.resultFuture.completeExceptionally(t);
}
return entry.resultFuture;
}
/**
* Indicate that no more {@link #send(Object)} calls will be made, and that streaming of messages has failed.
*
* This can be called before any in-flight {@code send} calls are complete. Such messages will be processed before the
* stream is treated as being in-error. The returned future will be completed successfully when the {@code error} is
* sent to the downstream subscriber.
*
*
After this method is invoked, any future {@link #send(Object)}, {@link #complete()} or {@code #error(Throwable)}
* calls will be completed exceptionally and not be processed.
*
*
The returned future will be completed exceptionally if the downstream subscriber cancels the subscription, or
* if the {@code complete} call was performed after a {@link #complete()} or {@code error} call.
*
* @param error The error to send.
* @return A future that is completed when the exception has been sent to the downstream subscriber.
*/
public CompletableFuture error(Throwable error) {
log.trace(() -> "Received error() with " + error, error);
OnErrorQueueEntry entry = new OnErrorQueueEntry<>(error);
try {
standardPriorityQueue.add(entry);
processEventQueue();
} catch (RuntimeException t) {
entry.resultFuture.completeExceptionally(t);
}
return entry.resultFuture;
}
/**
* A method called by the downstream subscriber in order to subscribe to the publisher.
*/
@Override
public void subscribe(Subscriber s) {
if (subscriber != null) {
s.onSubscribe(new NoOpSubscription());
s.onError(new IllegalStateException("Only one subscription may be active at a time."));
}
this.subscriber = s;
s.onSubscribe(new SubscriptionImpl());
processEventQueue();
}
/**
* Process the messages in the event queue. This is invoked after every operation on the publisher that changes the state
* of the event queue.
*
* Internally, this method will only be executed by one thread at a time. Any calls to this method will another thread
* is processing the queue will return immediately. This ensures: (1) thread safety in queue processing, (2) mutual recursion
* between onSubscribe/onNext with {@link Subscription#request(long)} are impossible.
*/
private void processEventQueue() {
do {
if (!processingQueue.compareAndSet(false, true)) {
// Some other thread is processing the queue, so we don't need to.
return;
}
try {
doProcessQueue();
} catch (Throwable e) {
panicAndDie(e);
break;
} finally {
processingQueue.set(false);
}
// Once releasing the processing-queue flag, we need to double-check that the queue still doesn't need to be
// processed, because new messages might have come in since we decided to release the flag.
} while (shouldProcessQueueEntry(standardPriorityQueue.peek()) ||
shouldProcessQueueEntry(highPriorityQueue.peek()));
}
/**
* Pop events off of the queue and process them in the order they are given, returning when we can no longer process the
* event at the head of the queue.
*
*
Invoked only from within the {@link #processEventQueue()} method with the {@link #processingQueue} flag held.
*/
private void doProcessQueue() {
while (true) {
QueueEntry entry = highPriorityQueue.peek();
Queue sourceQueue = highPriorityQueue;
if (entry == null) {
entry = standardPriorityQueue.peek();
sourceQueue = standardPriorityQueue;
}
if (!shouldProcessQueueEntry(entry)) {
// We're done processing entries.
return;
}
if (failureMessage.isSet()) {
entry.resultFuture.completeExceptionally(failureMessage.get());
} else {
switch (entry.type()) {
case ON_NEXT:
OnNextQueueEntry onNextEntry = (OnNextQueueEntry) entry;
log.trace(() -> "Calling onNext() with " + onNextEntry.value);
subscriber.onNext(onNextEntry.value);
long newDemand = outstandingDemand.decrementAndGet();
log.trace(() -> "Decreased demand to " + newDemand);
break;
case ON_COMPLETE:
failureMessage.trySet(() -> new IllegalStateException("onComplete() was already invoked."));
log.trace(() -> "Calling onComplete()");
subscriber.onComplete();
break;
case ON_ERROR:
OnErrorQueueEntry onErrorEntry = (OnErrorQueueEntry) entry;
failureMessage.trySet(() -> new IllegalStateException("onError() was already invoked.",
onErrorEntry.failure));
log.trace(() -> "Calling onError() with " + onErrorEntry.failure, onErrorEntry.failure);
subscriber.onError(onErrorEntry.failure);
break;
case CANCEL:
failureMessage.trySet(() -> new CancellationException("subscription has been cancelled."));
subscriber = null; // Allow subscriber to be garbage collected after cancellation.
break;
default:
// Should never happen. Famous last words?
throw new IllegalStateException("Unknown entry type: " + entry.type());
}
entry.resultFuture.complete(null);
}
sourceQueue.remove();
}
}
/**
* Return true if we should process the provided queue entry.
*/
private boolean shouldProcessQueueEntry(QueueEntry entry) {
if (entry == null) {
// The queue is empty.
return false;
}
if (failureMessage.isSet()) {
return true;
}
if (subscriber == null) {
// We don't have a subscriber yet.
return false;
}
if (entry.type() != ON_NEXT) {
// This event isn't an on-next event, so we don't need subscriber demand in order to process it.
return true;
}
// This is an on-next event and we're not failing on-next events, so make sure we have demand available before
// processing it.
return outstandingDemand.get() > 0;
}
/**
* Invoked from within {@link #processEventQueue()} when we can't process the queue for some reason. This is likely
* caused by a downstream subscriber throwing an exception from {@code onNext}, which it should never do.
*
* Here we try our best to fail all of the entries in the queue, so that no callers have "stuck" futures.
*/
private void panicAndDie(Throwable cause) {
try {
// Create exception here instead of in supplier to preserve a more-useful stack trace.
RuntimeException failure = new IllegalStateException("Encountered fatal error in publisher", cause);
failureMessage.trySet(() -> failure);
subscriber.onError(cause instanceof Error ? cause : failure);
while (true) {
QueueEntry entry = standardPriorityQueue.poll();
if (entry == null) {
break;
}
entry.resultFuture.completeExceptionally(failure);
}
} catch (Throwable t) {
t.addSuppressed(cause);
log.error(() -> "Failed while processing a failure. This could result in stuck futures.", t);
}
}
/**
* The subscription passed to the first {@link #subscriber} that subscribes to this publisher. This allows the downstream
* subscriber to request for more {@code onNext} calls or to {@code cancel} the stream of messages.
*/
private class SubscriptionImpl implements Subscription {
@Override
public void request(long n) {
log.trace(() -> "Received request() with " + n);
if (n <= 0) {
// Create exception here instead of in supplier to preserve a more-useful stack trace.
IllegalArgumentException failure = new IllegalArgumentException("A downstream publisher requested an invalid "
+ "amount of data: " + n);
highPriorityQueue.add(new OnErrorQueueEntry<>(failure));
processEventQueue();
} else {
long newDemand = outstandingDemand.updateAndGet(current -> {
if (Long.MAX_VALUE - current < n) {
return Long.MAX_VALUE;
}
return current + n;
});
log.trace(() -> "Increased demand to " + newDemand);
processEventQueue();
}
}
@Override
public void cancel() {
log.trace(() -> "Received cancel()");
// Create exception here instead of in supplier to preserve a more-useful stack trace.
highPriorityQueue.add(new CancelQueueEntry<>());
processEventQueue();
}
}
/**
* A lazily-initialized failure message for future events sent to this publisher after a terminal event has
* occurred.
*/
private static final class FailureMessage {
private Supplier failureMessageSupplier;
private Throwable failureMessage;
private void trySet(Supplier supplier) {
if (failureMessageSupplier == null) {
failureMessageSupplier = supplier;
}
}
private boolean isSet() {
return failureMessageSupplier != null;
}
private Throwable get() {
if (failureMessage == null) {
failureMessage = failureMessageSupplier.get();
}
return failureMessage;
}
}
/**
* An entry in the {@link #standardPriorityQueue}.
*/
abstract static class QueueEntry {
/**
* The future that was returned to a {@link #send(Object)}, {@link #complete()} or {@link #error(Throwable)} message.
*/
protected final CompletableFuture resultFuture = new CompletableFuture<>();
/**
* Retrieve the type of this queue entry.
*/
protected abstract Type type();
protected enum Type {
ON_NEXT,
ON_COMPLETE,
ON_ERROR,
CANCEL
}
}
/**
* An entry added when we get a {@link #send(Object)} call.
*/
private static final class OnNextQueueEntry extends QueueEntry {
private final T value;
private OnNextQueueEntry(T value) {
this.value = value;
}
@Override
protected Type type() {
return ON_NEXT;
}
}
/**
* An entry added when we get a {@link #complete()} call.
*/
private static final class OnCompleteQueueEntry extends QueueEntry {
@Override
protected Type type() {
return ON_COMPLETE;
}
}
/**
* An entry added when we get an {@link #error(Throwable)} call.
*/
private static final class OnErrorQueueEntry extends QueueEntry {
private final Throwable failure;
private OnErrorQueueEntry(Throwable failure) {
this.failure = failure;
}
@Override
protected Type type() {
return ON_ERROR;
}
}
/**
* An entry added when we get a {@link SubscriptionImpl#cancel()} call.
*/
private static final class CancelQueueEntry extends QueueEntry {
@Override
protected Type type() {
return CANCEL;
}
}
/**
* A subscription that does nothing. This is used for signaling {@code onError} to subscribers that subscribe to this
* publisher for the second time. Only one subscriber is supported.
*/
private static final class NoOpSubscription implements Subscription {
@Override
public void request(long n) {
}
@Override
public void cancel() {
}
}
}