Pubsub ordering keys

google-cloud-clients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java

@@ -62,6 +62,7 @@ class MessageDispatcher {
   @InternalApi static final Duration PENDING_ACKS_SEND_DELAY = Duration.ofMillis(100);
 
   private final Executor executor;
+  private final SequentialExecutorService.AutoExecutor sequentialExecutor;
   private final ScheduledExecutorService systemExecutor;
   private final ApiClock clock;
 
@@ -205,6 +206,7 @@ void sendAckOperations(
     jobLock = new ReentrantLock();
     messagesWaiter = new MessageWaiter();
     this.clock = clock;
+    this.sequentialExecutor = new SequentialExecutorService.AutoExecutor(executor);
   }
 
   void start() {
@@ -349,7 +351,7 @@ public void nack() {
           }
         };
     ApiFutures.addCallback(response, ackHandler, MoreExecutors.directExecutor());
-    executor.execute(
+    Runnable deliverMessageTask =
         new Runnable() {
           @Override
           public void run() {
@@ -370,7 +372,12 @@ public void run() {
               response.setException(e);
             }
           }
-        });
+        };
+    if (message.getOrderingKey().isEmpty()) {
+      executor.execute(deliverMessageTask);
+    } else {
+      sequentialExecutor.submit(message.getOrderingKey(), deliverMessageTask);
+    }
   }
 
   /** Compute the ideal deadline, set subsequent modacks to this deadline, and return it. */

google-cloud-clients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java

@@ -49,9 +49,13 @@
 import com.google.pubsub.v1.TopicNames;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -85,15 +89,17 @@ public class Publisher {
   private final String topicName;
 
   private final BatchingSettings batchingSettings;
+  private final boolean enableMessageOrdering;
 
   private final Lock messagesBatchLock;
-  private MessagesBatch messagesBatch;
+  final Map<String, MessagesBatch> messagesBatches;
 
   private final AtomicBoolean activeAlarm;
 
   private final PublisherStub publisherStub;
 
   private final ScheduledExecutorService executor;
+  final SequentialExecutorService.CallbackExecutor sequentialExecutor;
   private final AtomicBoolean shutdown;
   private final BackgroundResource backgroundResources;
   private final MessageWaiter messagesWaiter;
@@ -114,22 +120,33 @@ private Publisher(Builder builder) throws IOException {
     topicName = builder.topicName;
 
     this.batchingSettings = builder.batchingSettings;
+    this.enableMessageOrdering = builder.enableMessageOrdering;
     this.messageTransform = builder.messageTransform;
 
-    messagesBatch = new MessagesBatch(batchingSettings);
+    messagesBatches = new HashMap<>();
     messagesBatchLock = new ReentrantLock();
     activeAlarm = new AtomicBoolean(false);
     executor = builder.executorProvider.getExecutor();
+    sequentialExecutor = new SequentialExecutorService.CallbackExecutor(executor);
     List<BackgroundResource> backgroundResourceList = new ArrayList<>();
     if (builder.executorProvider.shouldAutoClose()) {
       backgroundResourceList.add(new ExecutorAsBackgroundResource(executor));
     }
 
     // Publisher used to take maxAttempt == 0 to mean infinity, but to GAX it means don't retry.
     // We post-process this here to keep backward-compatibility.
-    RetrySettings retrySettings = builder.retrySettings;
-    if (retrySettings.getMaxAttempts() == 0) {
-      retrySettings = retrySettings.toBuilder().setMaxAttempts(Integer.MAX_VALUE).build();
+    // Also, if "message ordering" is enabled, the publisher should retry sending the failed
+    // message infinitely rather than sending the next one.
+    RetrySettings.Builder retrySettingsBuilder = builder.retrySettings.toBuilder();
+    if (retrySettingsBuilder.getMaxAttempts() == 0) {
+      retrySettingsBuilder.setMaxAttempts(Integer.MAX_VALUE);
+    }
+    if (enableMessageOrdering) {
+      // TODO: is there a way to have the default retry settings for requests without an ordering
+      // key?
+      retrySettingsBuilder
+          .setMaxAttempts(Integer.MAX_VALUE)
+          .setTotalTimeout(Duration.ofNanos(Long.MAX_VALUE));
     }
 
     PublisherStubSettings.Builder stubSettings =
@@ -147,7 +164,7 @@ private Publisher(Builder builder) throws IOException {
             StatusCode.Code.RESOURCE_EXHAUSTED,
             StatusCode.Code.UNKNOWN,
             StatusCode.Code.UNAVAILABLE)
-        .setRetrySettings(retrySettings)
+        .setRetrySettings(retrySettingsBuilder.build())
         .setBatchingSettings(BatchingSettings.newBuilder().setIsEnabled(false).build());
     this.publisherStub = GrpcPublisherStub.create(stubSettings.build());
     backgroundResourceList.add(publisherStub);
@@ -194,13 +211,27 @@ public String getTopicNameString() {
   public ApiFuture<String> publish(PubsubMessage message) {
     Preconditions.checkState(!shutdown.get(), "Cannot publish on a shut-down publisher.");
 
+    final String orderingKey = message.getOrderingKey();
+    Preconditions.checkState(
+        orderingKey.isEmpty() || enableMessageOrdering,
+        "Cannot publish a message with an ordering key when message ordering is not enabled.");
+
     final OutstandingPublish outstandingPublish =
         new OutstandingPublish(messageTransform.apply(message));
     List<OutstandingBatch> batchesToSend;
     messagesBatchLock.lock();
     try {
+      // Check if the next message makes the current batch exceed the max batch byte size.
+      MessagesBatch messagesBatch = messagesBatches.get(orderingKey);
+      if (messagesBatch == null) {
+        messagesBatch = new MessagesBatch(batchingSettings, orderingKey);
+        messagesBatches.put(orderingKey, messagesBatch);
+      }
+
       batchesToSend = messagesBatch.add(outstandingPublish);
-      // Setup the next duration based delivery alarm if there are messages batched.
+      if (!batchesToSend.isEmpty() && messagesBatch.isEmpty()) {
+        messagesBatches.remove(orderingKey);
+      }
       setupAlarm();
     } finally {
       messagesBatchLock.unlock();
@@ -209,6 +240,10 @@ public ApiFuture<String> publish(PubsubMessage message) {
     messagesWaiter.incrementPendingMessages(1);
 
     if (!batchesToSend.isEmpty()) {
+      // TODO: if this is not an ordering keys scenario, will this do anything?
+      publishAllWithoutInflight();
+
+      // TODO: if this is an ordering keys scenario, is this safe without messagesBatchLock?
       for (final OutstandingBatch batch : batchesToSend) {
         logger.log(Level.FINER, "Scheduling a batch for immediate sending.");
         executor.execute(
@@ -224,8 +259,22 @@ public void run() {
     return outstandingPublish.publishResult;
   }
 
+  /**
+   * There may be non-recoverable problems with a request for an ordering key. In that case, all
+   * subsequent requests will fail until this method is called. If the key is not currently paused,
+   * calling this method will be a no-op.
+   *
+   * @param key The key for which to resume publishing.
+   */
+  // TODO: make this public when Ordering keys is live
+  @BetaApi
+  void resumePublish(String key) {
+    Preconditions.checkState(!shutdown.get(), "Cannot publish on a shut-down publisher.");
+    sequentialExecutor.resumePublish(key);
+  }
+
   private void setupAlarm() {
-    if (!messagesBatch.isEmpty()) {
+    if (!messagesBatches.isEmpty()) {
       if (!activeAlarm.getAndSet(true)) {
         long delayThresholdMs = getBatchingSettings().getDelayThreshold().toMillis();
         logger.log(Level.FINER, "Setting up alarm for the next {0} ms.", delayThresholdMs);
@@ -236,7 +285,7 @@ private void setupAlarm() {
                   public void run() {
                     logger.log(Level.FINER, "Sending messages based on schedule.");
                     activeAlarm.getAndSet(false);
-                    publishAllOutstanding();
+                    publishAllWithoutInflight();
                   }
                 },
                 delayThresholdMs,
@@ -257,16 +306,51 @@ public void run() {
    */
   public void publishAllOutstanding() {
     messagesBatchLock.lock();
-    OutstandingBatch batchToSend;
     try {
-      if (messagesBatch.isEmpty()) {
-        return;
+      for (MessagesBatch batch : messagesBatches.values()) {
+        if (!batch.isEmpty()) {
+          // TODO(kimkyung-goog): Do not release `messagesBatchLock` when publishing a batch. If
+          // it's released, the order of publishing cannot be guaranteed if `publish()` is called
+          // while this function is running. This locking mechanism needs to be improved if it
+          // causes any performance degradation.
+          publishOutstandingBatch(batch.popOutstandingBatch());
+        }
+      }
+      messagesBatches.clear();
+    } finally {
+      messagesBatchLock.unlock();
+    }
+  }
+
+  /**
+   * Publish any outstanding batches if non-empty and there are no other batches in flight. This
+   * method sends buffered messages, but does not wait for the send operations to complete. To wait
+   * for messages to send, call {@code get} on the futures returned from {@code publish}.
+   */
+  private void publishAllWithoutInflight() {
+    messagesBatchLock.lock();
+    try {
+      Iterator<Map.Entry<String, MessagesBatch>> it = messagesBatches.entrySet().iterator();
+      while (it.hasNext()) {
+        Map.Entry<String, MessagesBatch> entry = it.next();
+        MessagesBatch batch = entry.getValue();
+        String key = entry.getKey();
+        if (batch.isEmpty()) {
+          it.remove();
+        } else if (key.isEmpty() || !sequentialExecutor.hasTasksInflight(key)) {
+          // TODO(kimkyung-goog): Do not release `messagesBatchLock` when publishing a batch. If
+          // it's released, the order of publishing cannot be guaranteed if `publish()` is called
+          // while this function is running. This locking mechanism needs to be improved if it
+          // causes any performance degradation.
+
+          // TODO: Will this cause a performance problem for non-ordering keys scenarios?
+          publishOutstandingBatch(batch.popOutstandingBatch());
+          it.remove();
+        }
       }
-      batchToSend = messagesBatch.popOutstandingBatch();
     } finally {
       messagesBatchLock.unlock();
     }
-    publishOutstandingBatch(batchToSend);
   }
 
   private ApiFuture<PublishResponse> publishCall(OutstandingBatch outstandingBatch) {
@@ -280,12 +364,12 @@ private ApiFuture<PublishResponse> publishCall(OutstandingBatch outstandingBatch
   }
 
   private void publishOutstandingBatch(final OutstandingBatch outstandingBatch) {
-    ApiFutureCallback<PublishResponse> futureCallback =
+    final ApiFutureCallback<PublishResponse> futureCallback =
         new ApiFutureCallback<PublishResponse>() {
           @Override
           public void onSuccess(PublishResponse result) {
             try {
-              if (result.getMessageIdsCount() != outstandingBatch.size()) {
+              if (result == null || result.getMessageIdsCount() != outstandingBatch.size()) {
                 outstandingBatch.onFailure(
                     new IllegalStateException(
                         String.format(
@@ -311,20 +395,37 @@ public void onFailure(Throwable t) {
           }
         };
 
-    ApiFutures.addCallback(publishCall(outstandingBatch), futureCallback, directExecutor());
+    ApiFuture<PublishResponse> future;
+    if (outstandingBatch.orderingKey == null || outstandingBatch.orderingKey.isEmpty()) {
+      future = publishCall(outstandingBatch);
+    } else {
+      // If ordering key is specified, publish the batch using the sequential executor.
+      future =
+          sequentialExecutor.submit(
+              outstandingBatch.orderingKey,
+              new Callable<ApiFuture<PublishResponse>>() {
+                public ApiFuture<PublishResponse> call() {
+                  return publishCall(outstandingBatch);
+                }
+              });
+    }
+    ApiFutures.addCallback(future, futureCallback, directExecutor());
   }
 
   private static final class OutstandingBatch {
     final List<OutstandingPublish> outstandingPublishes;
     final long creationTime;
     int attempt;
     int batchSizeBytes;
+    final String orderingKey;
 
-    OutstandingBatch(List<OutstandingPublish> outstandingPublishes, int batchSizeBytes) {
+    OutstandingBatch(
+        List<OutstandingPublish> outstandingPublishes, int batchSizeBytes, String orderingKey) {
       this.outstandingPublishes = outstandingPublishes;
       attempt = 1;
       creationTime = System.currentTimeMillis();
       this.batchSizeBytes = batchSizeBytes;
+      this.orderingKey = orderingKey;
     }
 
     int size() {
@@ -468,7 +569,7 @@ public static final class Builder {
             .setRpcTimeoutMultiplier(2)
             .setMaxRpcTimeout(DEFAULT_RPC_TIMEOUT)
             .build();
-
+    static final boolean DEFAULT_ENABLE_MESSAGE_ORDERING = false;
     private static final int THREADS_PER_CPU = 5;
     static final ExecutorProvider DEFAULT_EXECUTOR_PROVIDER =
         InstantiatingExecutorProvider.newBuilder()
@@ -482,6 +583,8 @@ public static final class Builder {
 
     RetrySettings retrySettings = DEFAULT_RETRY_SETTINGS;
 
+    private boolean enableMessageOrdering = DEFAULT_ENABLE_MESSAGE_ORDERING;
+
     private TransportChannelProvider channelProvider =
         TopicAdminSettings.defaultGrpcTransportProviderBuilder().setChannelsPerCpu(1).build();
 
@@ -576,6 +679,14 @@ public Builder setRetrySettings(RetrySettings retrySettings) {
       return this;
     }
 
+    /** Sets the message ordering option. */
+    // TODO: make this public when Ordering keys is live
+    @BetaApi
+    Builder setEnableMessageOrdering(boolean enableMessageOrdering) {
+      this.enableMessageOrdering = enableMessageOrdering;
+      return this;
+    }
+
     /** Gives the ability to set a custom executor to be used by the library. */
     public Builder setExecutorProvider(ExecutorProvider executorProvider) {
       this.executorProvider = Preconditions.checkNotNull(executorProvider);
@@ -601,15 +712,17 @@ public Publisher build() throws IOException {
   private static class MessagesBatch {
     private List<OutstandingPublish> messages;
     private int batchedBytes;
+    private String orderingKey;
     private final BatchingSettings batchingSettings;
 
-    public MessagesBatch(BatchingSettings batchingSettings) {
+    private MessagesBatch(BatchingSettings batchingSettings, String orderingKey) {
       this.batchingSettings = batchingSettings;
+      this.orderingKey = orderingKey;
       reset();
     }
 
     private OutstandingBatch popOutstandingBatch() {
-      OutstandingBatch batch = new OutstandingBatch(messages, batchedBytes);
+      OutstandingBatch batch = new OutstandingBatch(messages, batchedBytes, orderingKey);
       reset();
       return batch;
     }