feat: fix bug where blockingpullsubscriber doesn't refill flowcontrol…

… quota. (#449)
googleapis · Jan 12, 2021 · 109bd83 · 109bd83
1 parent 2ab6e3e
commit 109bd83
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/...ublite/src/main/java/com/google/cloud/pubsublite/internal/BlockingPullSubscriberImpl.java b/...ublite/src/main/java/com/google/cloud/pubsublite/internal/BlockingPullSubscriberImpl.java
@@ -33,7 +33,6 @@
 import java.util.ArrayDeque;
 import java.util.Collection;
 import java.util.Deque;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 
@@ -113,7 +112,13 @@ public synchronized Optional<SequencedMessage> messageIfAvailable() throws Check
     if (messages.isEmpty()) {
       return Optional.empty();
     }
-    return Optional.of(Objects.requireNonNull(messages.pollFirst()));
+    SequencedMessage msg = messages.remove();
+    underlying.allowFlow(
+        FlowControlRequest.newBuilder()
+            .setAllowedMessages(1)
+            .setAllowedBytes(msg.byteSize())
+            .build());
+    return Optional.of(msg);
   }
 
   @Override

diff --git a/...te/src/test/java/com/google/cloud/pubsublite/internal/BlockingPullSubscriberImplTest.java b/...te/src/test/java/com/google/cloud/pubsublite/internal/BlockingPullSubscriberImplTest.java
@@ -148,10 +148,17 @@ public void onDataSuccess() throws Exception {
 
   @Test
   public void pullMessage() throws Exception {
+    int byteSize = 30;
     SequencedMessage message =
-        SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(12), 30);
+        SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(12), byteSize);
     messageConsumer.accept(ImmutableList.of(message));
     assertThat(Optional.of(message)).isEqualTo(subscriber.messageIfAvailable());
+    verify(underlying)
+        .allowFlow(
+            FlowControlRequest.newBuilder()
+                .setAllowedBytes(byteSize)
+                .setAllowedMessages(1)
+                .build());
   }
 
   @Test