fix: Ensure messages are delivered in order (#1031)

The messages are ordered when they are processed by the wire subscriber. However, if the user message receiver is slow, different message batches may be concurrently processed by the receiver. The wire subscriber now serializes delivery of the message batches.

Author: tmdiep

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/SerialExecutor.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * 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
+ *
+ *     https://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 com.google.cloud.pubsublite.internal;
+
+import java.util.ArrayDeque;
+import java.util.Queue;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
+import javax.annotation.concurrent.GuardedBy;
+
+/** An executor that runs tasks sequentially. */
+public final class SerialExecutor implements AutoCloseable, Executor {
+  private final Executor executor;
+  private final AtomicBoolean isShutdown = new AtomicBoolean(false);
+
+  @GuardedBy("this")
+  private final Queue<Runnable> tasks;
+
+  @GuardedBy("this")
+  private boolean isTaskActive;
+
+  public SerialExecutor(Executor executor) {
+    this.executor = executor;
+    this.tasks = new ArrayDeque<>();
+    this.isTaskActive = false;
+  }
+
+  /**
+   * Shuts down the executor. No subsequent tasks will run, but any running task will be left to
+   * complete.
+   */
+  @Override
+  public void close() {
+    isShutdown.set(true);
+  }
+
+  @Override
+  public synchronized void execute(Runnable r) {
+    if (isShutdown.get()) {
+      return;
+    }
+    tasks.add(
+        () -> {
+          if (isShutdown.get()) {
+            return;
+          }
+          try {
+            r.run();
+          } finally {
+            scheduleNextTask();
+          }
+        });
+    if (!isTaskActive) {
+      scheduleNextTask();
+    }
+  }
+
+  private synchronized void scheduleNextTask() {
+    isTaskActive = !tasks.isEmpty();
+    if (isTaskActive) {
+      executor.execute(tasks.poll());
+    }
+  }
+}

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/SubscriberImpl.java

@@ -24,6 +24,7 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.CloseableMonitor;
 import com.google.cloud.pubsublite.internal.ProxyService;
+import com.google.cloud.pubsublite.internal.SerialExecutor;
 import com.google.cloud.pubsublite.internal.wire.StreamFactories.SubscribeStreamFactory;
 import com.google.cloud.pubsublite.proto.FlowControlRequest;
 import com.google.cloud.pubsublite.proto.InitialSubscribeRequest;
@@ -50,6 +51,9 @@ public class SubscriberImpl extends ProxyService
 
   private final InitialSubscribeRequest baseInitialRequest;
 
+  // Used to ensure messages are delivered to consumers in order.
+  private final SerialExecutor messageDeliveryExecutor;
+
   private final CloseableMonitor monitor = new CloseableMonitor();
 
   @GuardedBy("monitor.monitor")
@@ -84,6 +88,7 @@ public class SubscriberImpl extends ProxyService
     this.messageConsumer = messageConsumer;
     this.resetHandler = resetHandler;
     this.baseInitialRequest = baseInitialRequest;
+    this.messageDeliveryExecutor = new SerialExecutor(SystemExecutors.getFuturesExecutor());
     this.initialLocation = initialLocation;
     this.connection =
         new RetryingConnectionImpl<>(streamFactory, factory, this, getInitialRequest());
@@ -121,6 +126,7 @@ protected void stop() {
       shutdown = true;
       this.alarmFuture.ifPresent(future -> future.cancel(false));
       this.alarmFuture = Optional.empty();
+      messageDeliveryExecutor.close();
     }
   }
 
@@ -199,8 +205,8 @@ public void onClientResponse(List<SequencedMessage> messages) throws CheckedApiE
       if (shutdown) return;
       nextOffsetTracker.onMessages(messages);
       flowControlBatcher.onMessages(messages);
+      messageDeliveryExecutor.execute(() -> messageConsumer.accept(messages));
     }
-    messageConsumer.accept(messages);
   }
 
   private void processBatchFlowRequest() {

google-cloud-pubsublite/src/test/java/com/google/cloud/pubsublite/internal/SerialExecutorTest.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * 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 com.google.cloud.pubsublite.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class SerialExecutorTest {
+  private final SerialExecutor executor = new SerialExecutor(SystemExecutors.getFuturesExecutor());
+
+  @Test
+  public void serializesTasks() throws Exception {
+    final int numTasks = 100;
+    List<Integer> receivedSequences = new ArrayList<>();
+    CountDownLatch tasksDone = new CountDownLatch(numTasks);
+    for (int i = 0; i < numTasks; i++) {
+      int sequence = i;
+      executor.execute(
+          () -> {
+            synchronized (receivedSequences) {
+              receivedSequences.add(sequence);
+            }
+            tasksDone.countDown();
+          });
+    }
+    assertThat(tasksDone.await(30, SECONDS)).isTrue();
+
+    for (int i = 0; i < receivedSequences.size(); i++) {
+      assertThat(receivedSequences.get(i)).isEqualTo(i);
+    }
+  }
+
+  @Test
+  public void closeDiscardsTasks() throws Exception {
+    List<Integer> receivedSequences = new ArrayList<>();
+    CountDownLatch tasksDone = new CountDownLatch(1);
+    for (int i = 0; i < 10; i++) {
+      int sequence = i;
+      executor.execute(
+          () -> {
+            synchronized (receivedSequences) {
+              receivedSequences.add(sequence);
+            }
+            tasksDone.countDown();
+            executor.close();
+          });
+    }
+    assertThat(tasksDone.await(10, SECONDS)).isTrue();
+
+    assertThat(receivedSequences).containsExactly(0);
+  }
+}

google-cloud-pubsublite/src/test/java/com/google/cloud/pubsublite/internal/wire/SubscriberImplTest.java

@@ -19,6 +19,7 @@
 import static com.google.cloud.pubsublite.internal.ApiExceptionMatcher.assertThrowableMatches;
 import static com.google.cloud.pubsublite.internal.testing.RetryingConnectionHelpers.whenFailed;
 import static com.google.common.truth.Truth.assertThat;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -54,6 +55,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.util.Timestamps;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Future;
 import java.util.function.Consumer;
 import org.junit.Before;
@@ -97,6 +99,18 @@ private static SubscribeRequest initialRequest() {
   private ResponseObserver<List<SequencedMessage>> leakedResponseObserver;
   private Runnable leakedFlowControlAlarm;
 
+  private CountDownLatch countdownMessageBatches(int count) {
+    CountDownLatch received = new CountDownLatch(count);
+    doAnswer(
+            args -> {
+              received.countDown();
+              return null;
+            })
+        .when(mockMessageConsumer)
+        .accept(any());
+    return received;
+  }
+
   @Before
   public void setUp() throws CheckedApiException {
     initMocks(this);
@@ -216,7 +230,9 @@ public void messagesOrdered_Ok() throws Exception {
         ImmutableList.of(
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 10),
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 10));
+    CountDownLatch messagesReceived = countdownMessageBatches(1);
     leakedResponseObserver.onResponse(messages);
+    assertThat(messagesReceived.await(10, SECONDS)).isTrue();
 
     verify(mockMessageConsumer).accept(messages);
     assertThat(subscriber.isRunning()).isTrue();
@@ -237,7 +253,9 @@ public void messageResponseSubtracts() throws Exception {
     ImmutableList<SequencedMessage> messages2 =
         ImmutableList.of(
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(3), 2));
+    CountDownLatch messagesReceived = countdownMessageBatches(1);
     leakedResponseObserver.onResponse(messages1);
+    assertThat(messagesReceived.await(10, SECONDS)).isTrue();
     verify(mockMessageConsumer).accept(messages1);
     assertThat(subscriber.isRunning()).isTrue();
     leakedResponseObserver.onResponse(messages2);
@@ -254,7 +272,9 @@ public void reinitialize_reconnectsToNextOffset() throws Exception {
         ImmutableList.of(
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 10),
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 10));
+    CountDownLatch messagesReceived = countdownMessageBatches(1);
     leakedResponseObserver.onResponse(messages);
+    assertThat(messagesReceived.await(10, SECONDS)).isTrue();
     verify(mockMessageConsumer).accept(messages);
 
     final SubscribeRequest nextOffsetRequest =
@@ -288,7 +308,9 @@ public void reinitialize_handlesSuccessfulReset() throws Exception {
         ImmutableList.of(
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 10),
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 10));
+    CountDownLatch messagesReceived = countdownMessageBatches(1);
     leakedResponseObserver.onResponse(messages);
+    assertThat(messagesReceived.await(10, SECONDS)).isTrue();
     verify(mockMessageConsumer).accept(messages);
 
     doAnswer(
@@ -317,7 +339,9 @@ public void reinitialize_handlesIgnoredReset() throws Exception {
         ImmutableList.of(
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 10),
             SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 10));
+    CountDownLatch messagesReceived = countdownMessageBatches(1);
     leakedResponseObserver.onResponse(messages);
+    assertThat(messagesReceived.await(10, SECONDS)).isTrue();
     verify(mockMessageConsumer).accept(messages);
 
     final SubscribeRequest nextOffsetRequest =