Clean up after extension give up

googleapis · Sep 1, 2018 · 430938f · 430938f
1 parent b9b4aff
commit 430938f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/...ients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java b/...ients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java
@@ -455,6 +455,10 @@ void extendDeadlines() {
         int sec = Math.max(1, (int) now.until(totalExpiration, ChronoUnit.SECONDS));
         modacks.add(new PendingModifyAckDeadline(sec, ackId));
         count++;
+      } else {
+        // TODO: How to get the byte count?
+        flowController.release(1, 0);
+        messagesWaiter.incrementPendingMessages(-1);
       }
     }
     modacks.add(modack);

diff --git a/...s/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/MessageDispatcherTest.java b/...s/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/MessageDispatcherTest.java
@@ -56,6 +56,7 @@ public void run() {
   private List<String> sentAcks;
   private List<ModAckItem> sentModAcks;
   private FakeClock clock;
+  private FlowController flowController;
 
   @AutoValue
   abstract static class ModAckItem {
@@ -101,6 +102,12 @@ public void sendAckOperations(
     systemExecutor.shutdownNow();
 
     clock = new FakeClock();
+    flowController =
+        new FlowController(
+            FlowControlSettings.newBuilder()
+              .setMaxOutstandingElementCount(1L)
+              .setLimitExceededBehavior(FlowController.LimitExceededBehavior.ThrowException)
+              .build());
 
     dispatcher =
         new MessageDispatcher(
@@ -109,7 +116,7 @@ public void sendAckOperations(
             Duration.ofSeconds(5),
             Duration.ofMinutes(60),
             new Distribution(Subscriber.MAX_ACK_DEADLINE_SECONDS + 1),
-            new FlowController(FlowControlSettings.newBuilder().build()),
+            flowController,
             new LinkedList<MessageDispatcher.OutstandingMessageBatch>(),
             MoreExecutors.directExecutor(),
             systemExecutor,
@@ -182,6 +189,10 @@ public void testExtension_GiveUp() throws Exception {
     clock.advance(1, TimeUnit.DAYS);
     dispatcher.extendDeadlines();
     assertThat(sentModAcks).isEmpty();
+
+    // We should be able to reserve another item in the flow controller and not block shutdown
+    flowController.reserve(1, 0);
+    dispatcher.stop();
   }
 
   @Test