Implements batching of flow control requests (#92)

* Implemented batching of flow control requests

* Minor renames of variables

* Tests and review comments

* Expose constant for flow request flush interval

* Review comments: revert name requestForRestart

Author: tmdiep

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

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2020 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.wire;
+
+import com.google.cloud.pubsublite.SequencedMessage;
+import com.google.cloud.pubsublite.proto.FlowControlRequest;
+import io.grpc.StatusException;
+import java.util.Collection;
+import java.util.Optional;
+
+// A FlowControlBatcher manages batching of FlowControlRequests.
+class FlowControlBatcher {
+  // If the pending flow request exceeds this ratio of the current outstanding tokens, the pending
+  // flow request should be flushed to the stream ASAP.
+  private static double EXPEDITE_BATCH_REQUEST_RATIO = 0.5;
+
+  // The current amount of outstanding byte and message flow control tokens.
+  private final TokenCounter clientTokens = new TokenCounter();
+
+  // The pending aggregate flow control request that needs to be sent to the stream.
+  private final TokenCounter pendingTokens = new TokenCounter();
+
+  void onClientFlowRequest(FlowControlRequest request) throws StatusException {
+    clientTokens.add(request.getAllowedBytes(), request.getAllowedMessages());
+    pendingTokens.add(request.getAllowedBytes(), request.getAllowedMessages());
+  }
+
+  void onMessages(Collection<SequencedMessage> received) throws StatusException {
+    long byteSize = received.stream().mapToLong(SequencedMessage::byteSize).sum();
+    clientTokens.sub(byteSize, received.size());
+  }
+
+  void onClientSeek() {
+    clientTokens.reset();
+    pendingTokens.reset();
+  }
+
+  // The caller must send the FlowControlRequest to the stream, as pending tokens are reset.
+  Optional<FlowControlRequest> requestForRestart() {
+    pendingTokens.reset();
+    return clientTokens.toFlowControlRequest();
+  }
+
+  // The caller must send the FlowControlRequest to the stream, as pending tokens are reset.
+  Optional<FlowControlRequest> releasePendingRequest() {
+    Optional<FlowControlRequest> request = pendingTokens.toFlowControlRequest();
+    pendingTokens.reset();
+    return request;
+  }
+
+  boolean shouldExpediteBatchRequest() {
+    if (exceedsExpediteRatio(pendingTokens.bytes(), clientTokens.bytes())) return true;
+    if (exceedsExpediteRatio(pendingTokens.messages(), clientTokens.messages())) return true;
+    return false;
+  }
+
+  private boolean exceedsExpediteRatio(long pending, long client) {
+    return client > 0 && ((double) pending / client) >= EXPEDITE_BATCH_REQUEST_RATIO;
+  }
+}

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

@@ -38,23 +38,32 @@
 import io.grpc.Status;
 import io.grpc.StatusException;
 import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import javax.annotation.concurrent.GuardedBy;
 
 public class SubscriberImpl extends ProxyService
     implements Subscriber, RetryingConnectionObserver<Response> {
+  @VisibleForTesting static final long FLOW_REQUESTS_FLUSH_INTERVAL_MS = 100;
+
   private final Consumer<ImmutableList<SequencedMessage>> messageConsumer;
 
   private final CloseableMonitor monitor = new CloseableMonitor();
 
+  private final ScheduledExecutorService executorService;
+  private Future<?> alarmFuture;
+
   @GuardedBy("monitor.monitor")
   private final RetryingConnection<ConnectedSubscriber> connection;
 
   @GuardedBy("monitor.monitor")
   private final NextOffsetTracker nextOffsetTracker = new NextOffsetTracker();
 
   @GuardedBy("monitor.monitor")
-  private final TokenCounter tokenCounter = new TokenCounter();
+  private final FlowControlBatcher flowControlBatcher = new FlowControlBatcher();
 
   @GuardedBy("monitor.monitor")
   private Optional<InFlightSeek> inFlightSeek = Optional.empty();
@@ -89,6 +98,7 @@ private static class InFlightSeek {
             factory,
             SubscribeRequest.newBuilder().setInitial(initialRequest).build(),
             this);
+    this.executorService = Executors.newSingleThreadScheduledExecutor();
     addServices(this.connection);
   }
 
@@ -112,10 +122,21 @@ protected void handlePermanentError(StatusException error) {
   }
 
   @Override
-  protected void start() {}
+  protected void start() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      alarmFuture =
+          executorService.scheduleWithFixedDelay(
+              this::processBatchFlowRequest,
+              FLOW_REQUESTS_FLUSH_INTERVAL_MS,
+              FLOW_REQUESTS_FLUSH_INTERVAL_MS,
+              TimeUnit.MILLISECONDS);
+    }
+  }
 
   @Override
   protected void stop() {
+    alarmFuture.cancel(false /* mayInterruptIfRunning */);
+    executorService.shutdown();
     try (CloseableMonitor.Hold h = monitor.enter()) {
       shutdown = true;
       inFlightSeek.ifPresent(
@@ -143,7 +164,7 @@ public boolean isSatisfied() {
       checkState(!inFlightSeek.isPresent(), "Seeked while seek is already in flight.");
       SettableApiFuture<Offset> future = SettableApiFuture.create();
       inFlightSeek = Optional.of(new InFlightSeek(request, future));
-      tokenCounter.onClientSeek();
+      flowControlBatcher.onClientSeek();
       connection.modifyConnection(
           connectedSubscriber ->
               connectedSubscriber.ifPresent(subscriber -> subscriber.seek(request)));
@@ -161,15 +182,16 @@ public boolean seekInFlight() {
     }
   }
 
-  // TODO: Consider batching these requests before sending to the stream.
   @Override
-  public void allowFlow(FlowControlRequest request) {
+  public void allowFlow(FlowControlRequest clientRequest) {
     try (CloseableMonitor.Hold h = monitor.enter()) {
       if (shutdown) return;
-      tokenCounter.onClientFlowRequest(request);
-      connection.modifyConnection(
-          connectedSubscriber ->
-              connectedSubscriber.ifPresent(subscriber -> subscriber.allowFlow(request)));
+      flowControlBatcher.onClientFlowRequest(clientRequest);
+      if (flowControlBatcher.shouldExpediteBatchRequest()) {
+        connection.modifyConnection(
+            connectedSubscriber ->
+                connectedSubscriber.ifPresent(subscriber -> flushBatchFlowRequest(subscriber)));
+      }
     } catch (StatusException e) {
       onPermanentError(e);
       throw e.getStatus().asRuntimeException();
@@ -197,7 +219,7 @@ public void triggerReinitialize() {
                         connectedSubscriber.get().seek(request);
                       });
             }
-            tokenCounter
+            flowControlBatcher
                 .requestForRestart()
                 .ifPresent(request -> connectedSubscriber.get().allowFlow(request));
           });
@@ -223,7 +245,7 @@ private Status onMessageResponse(ImmutableList<SequencedMessage> messages) {
         return Status.OK;
       }
       nextOffsetTracker.onMessages(messages);
-      tokenCounter.onMessages(messages);
+      flowControlBatcher.onMessages(messages);
     } catch (StatusException e) {
 
       onPermanentError(e);
@@ -252,4 +274,24 @@ private Status onSeekResponse(Offset seekOffset) {
       return e.getStatus();
     }
   }
+
+  @VisibleForTesting
+  void processBatchFlowRequest() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      if (shutdown) return;
+      connection.modifyConnection(
+          connectedSubscriber ->
+              connectedSubscriber.ifPresent(subscriber -> flushBatchFlowRequest(subscriber)));
+    } catch (StatusException e) {
+      onPermanentError(e);
+    }
+  }
+
+  private void flushBatchFlowRequest(ConnectedSubscriber subscriber) {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      flowControlBatcher
+          .releasePendingRequest()
+          .ifPresent(request -> subscriber.allowFlow(request));
+    }
+  }
 }

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

@@ -19,43 +19,50 @@
 import static com.google.cloud.pubsublite.internal.Preconditions.checkArgument;
 import static com.google.cloud.pubsublite.internal.Preconditions.checkState;
 
-import com.google.cloud.pubsublite.SequencedMessage;
 import com.google.cloud.pubsublite.proto.FlowControlRequest;
 import com.google.common.math.LongMath;
 import io.grpc.StatusException;
-import java.util.Collection;
 import java.util.Optional;
 
-// A TokenCounter counts the amount of outstanding byte and message flow control tokens that the
+// A TokenCounter stores the amount of outstanding byte and message flow control tokens that the
 // client believes exists for the stream.
 class TokenCounter {
   private long bytes = 0;
   private long messages = 0;
 
-  void onClientFlowRequest(FlowControlRequest request) throws StatusException {
-    checkArgument(request.getAllowedMessages() >= 0);
-    checkArgument(request.getAllowedBytes() >= 0);
+  void add(long deltaBytes, long deltaMessages) throws StatusException {
+    checkArgument(deltaBytes >= 0);
+    checkArgument(deltaMessages >= 0);
 
-    bytes = LongMath.saturatedAdd(bytes, request.getAllowedBytes());
-    messages = LongMath.saturatedAdd(messages, request.getAllowedMessages());
+    bytes = LongMath.saturatedAdd(bytes, deltaBytes);
+    messages = LongMath.saturatedAdd(messages, deltaMessages);
   }
 
-  void onMessages(Collection<SequencedMessage> received) throws StatusException {
-    long byteSize = received.stream().mapToLong(SequencedMessage::byteSize).sum();
+  void sub(long deltaBytes, long deltaMessages) throws StatusException {
+    checkArgument(deltaBytes >= 0);
+    checkArgument(deltaMessages >= 0);
     checkState(
-        byteSize <= bytes, "Received messages that account for more bytes than were requested.");
-    checkState(received.size() <= messages, "Received more messages than were requested.");
+        deltaBytes <= bytes, "Received messages that account for more bytes than were requested.");
+    checkState(deltaMessages <= messages, "Received more messages than were requested.");
 
-    bytes -= byteSize;
-    messages -= received.size();
+    bytes -= deltaBytes;
+    messages -= deltaMessages;
   }
 
-  void onClientSeek() {
+  void reset() {
     bytes = 0;
     messages = 0;
   }
 
-  Optional<FlowControlRequest> requestForRestart() {
+  long bytes() {
+    return bytes;
+  }
+
+  long messages() {
+    return messages;
+  }
+
+  Optional<FlowControlRequest> toFlowControlRequest() {
     if (bytes == 0 && messages == 0) return Optional.empty();
     return Optional.of(
         FlowControlRequest.newBuilder()

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

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2020 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.wire;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.pubsublite.Message;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.SequencedMessage;
+import com.google.cloud.pubsublite.proto.FlowControlRequest;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.util.Timestamps;
+import io.grpc.StatusException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class FlowControlBatcherTest {
+
+  private final FlowControlBatcher batcher = new FlowControlBatcher();
+
+  @Test
+  public void onClientFlowRequestIncrementsTokens() throws StatusException {
+    FlowControlRequest clientFlowRequest =
+        FlowControlRequest.newBuilder().setAllowedBytes(500).setAllowedMessages(10).build();
+    batcher.onClientFlowRequest(clientFlowRequest);
+
+    assertThat(batcher.releasePendingRequest().get()).isEqualTo(clientFlowRequest);
+    assertThat(batcher.releasePendingRequest().isPresent()).isFalse();
+    assertThat(batcher.requestForRestart().get()).isEqualTo(clientFlowRequest);
+    assertThat(batcher.requestForRestart().get()).isEqualTo(clientFlowRequest);
+  }
+
+  @Test
+  public void onMessagesDecrementsClientTokens() throws StatusException {
+    FlowControlRequest clientFlowRequest =
+        FlowControlRequest.newBuilder().setAllowedBytes(500).setAllowedMessages(10).build();
+    batcher.onClientFlowRequest(clientFlowRequest);
+    ImmutableList<SequencedMessage> messages =
+        ImmutableList.of(
+            SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 100),
+            SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 150));
+    batcher.onMessages(messages);
+
+    assertThat(batcher.releasePendingRequest().get()).isEqualTo(clientFlowRequest);
+    FlowControlRequest expectedRequestForRestart =
+        FlowControlRequest.newBuilder().setAllowedBytes(250).setAllowedMessages(8).build();
+    assertThat(batcher.requestForRestart().get()).isEqualTo(expectedRequestForRestart);
+  }
+
+  @Test
+  public void shouldExpediteBatchRequestChecksByteRatio() throws StatusException {
+    batcher.onClientFlowRequest(
+        FlowControlRequest.newBuilder().setAllowedBytes(100).setAllowedMessages(100).build());
+    batcher.releasePendingRequest();
+
+    batcher.onClientFlowRequest(FlowControlRequest.newBuilder().setAllowedBytes(10).build());
+    assertThat(batcher.shouldExpediteBatchRequest()).isFalse();
+
+    batcher.onClientFlowRequest(FlowControlRequest.newBuilder().setAllowedBytes(90).build());
+    assertThat(batcher.shouldExpediteBatchRequest()).isTrue();
+  }
+
+  @Test
+  public void shouldExpediteBatchRequestChecksMessageRatio() throws StatusException {
+    batcher.onClientFlowRequest(
+        FlowControlRequest.newBuilder().setAllowedBytes(100).setAllowedMessages(100).build());
+    batcher.releasePendingRequest();
+
+    batcher.onClientFlowRequest(FlowControlRequest.newBuilder().setAllowedMessages(80).build());
+    assertThat(batcher.shouldExpediteBatchRequest()).isFalse();
+
+    batcher.onClientFlowRequest(FlowControlRequest.newBuilder().setAllowedMessages(20).build());
+    assertThat(batcher.shouldExpediteBatchRequest()).isTrue();
+  }
+
+  @Test
+  public void shouldExpediteBatchRequestHandlesDivByZero() {
+    assertThat(batcher.shouldExpediteBatchRequest()).isFalse();
+  }
+}