Max ack extension

google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/spi/v1/FakeScheduledExecutorService.java

@@ -20,6 +20,8 @@
 import com.google.common.primitives.Ints;
 import com.google.common.util.concurrent.SettableFuture;
 import java.util.ArrayList;
+import java.util.Deque;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.PriorityQueue;
 import java.util.concurrent.AbstractExecutorService;
@@ -79,6 +81,26 @@ public ScheduledFuture<?> scheduleWithFixedDelay(
             Duration.millis(unit.toMillis(initialDelay)), command, PendingCallableType.FIXED_DELAY));
   }
 
+  private Deque<Duration> expectedWorkQueue = new LinkedList<>();
+
+  public void setupScheduleExpectation(Duration delay) {
+    synchronized (expectedWorkQueue) {
+      expectedWorkQueue.add(delay);
+    }
+  }
+
+  public void waitForExpectedWork() {
+    synchronized (expectedWorkQueue) {
+      while (!expectedWorkQueue.isEmpty()) {
+        try {
+          expectedWorkQueue.wait();
+        } catch (InterruptedException e) {
+          // Wait uninterruptibly
+        }
+      }
+    }
+  }
+
   /**
    * This will advance the reference time of the executor and execute (in the same thread) any
    * outstanding callable which execution time has passed.
@@ -100,7 +122,7 @@ private void work() {
         callable = pendingCallables.poll();
       }
       if (callable != null) {
-        try{
+        try {
           callable.call();
         } catch (Exception e) {
           // We ignore any callable exception, which should be set to the future but not relevant to
@@ -182,6 +204,13 @@ <V> ScheduledFuture<V> schedulePendingCallable(PendingCallable<V> callable) {
       pendingCallables.add(callable);
     }
     work();
+    synchronized (expectedWorkQueue) {
+      if (!expectedWorkQueue.isEmpty() && expectedWorkQueue.peek().equals(callable.delay)) {
+        expectedWorkQueue.poll();
+      }
+      expectedWorkQueue.notifyAll();
+    }
+
     return callable.getScheduledFuture();
   }
 

google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/spi/v1/SubscriberImplTest.java

@@ -284,9 +284,12 @@ public void testModifyAckDeadline() throws Exception {
     // Send messages to be acked
     List<String> testAckIdsBatch = ImmutableList.of("A", "B", "C");
     testReceiver.setExplicitAck(true);
+    // A modify ack deadline should be schedule for the next 9s
+    fakeExecutor.setupScheduleExpectation(Duration.standardSeconds(9)); 
     sendMessages(testAckIdsBatch);
+    // To ensure first modify ack deadline got scheduled
+    fakeExecutor.waitForExpectedWork(); 
 
-    // Trigger modify ack deadline sending - 10s initial stream ack deadline - 1 padding
     fakeExecutor.advanceTime(Duration.standardSeconds(9));
 
     assertEquivalentWithTransformation(
@@ -299,17 +302,16 @@ public ModifyAckDeadline apply(String ack) {
           }
         });
 
-    // Trigger modify ack deadline sending - 2s of the renewed deadlines
-    fakeExecutor.advanceTime(Duration.standardSeconds(2));
+    fakeExecutor.advanceTime(Duration.standardSeconds(1));
 
     assertEquivalentWithTransformation(
         testAckIdsBatch,
         fakeSubscriberServiceImpl.waitAndConsumeModifyAckDeadlines(3),
         new Function<String, ModifyAckDeadline>() {
           @Override
           public ModifyAckDeadline apply(String ack) {
-            return new ModifyAckDeadline(ack, 2); // It is expected that the deadline is renewed
-                                                  // only two more seconds to not pass the max
+            return new ModifyAckDeadline(ack, 3); // It is expected that the deadline is renewed
+                                                  // only three more seconds to not pass the max
                                                   // ack deadline ext.
           }
         });
@@ -332,9 +334,13 @@ public void testModifyAckDeadline_defaultMaxExtensionPeriod() throws Exception {
     // Send messages to be acked
     List<String> testAckIdsBatch = ImmutableList.of("A", "B", "C");
     testReceiver.setExplicitAck(true);
+    // A modify ack deadline should be schedule for the next 9s
+    fakeExecutor.setupScheduleExpectation(Duration.standardSeconds(9)); 
     sendMessages(testAckIdsBatch);
+    // To ensure the first modify ack deadlines got scheduled
+    fakeExecutor.waitForExpectedWork(); 
 
-    // Trigger modify ack deadline sending - 10s initial stream ack deadline - 1 padding
+    // Next modify ack deadline should be schedule in the next 1s
     fakeExecutor.advanceTime(Duration.standardSeconds(9));
 
     assertEquivalentWithTransformation(
@@ -347,12 +353,12 @@ public ModifyAckDeadline apply(String ack) {
           }
         });
 
-    int timeIncrementSecs = INITIAL_ACK_DEADLINE_EXTENSION_SECS * 2; // Second time increment
+    fakeExecutor.advanceTime(Duration.standardSeconds(1));
+    int timeIncrementSecs = INITIAL_ACK_DEADLINE_EXTENSION_SECS; // Second time increment
 
     // Check ack deadline extensions while the current time has not reached 60 minutes
-    while (fakeExecutor.getClock().millisTime() + (timeIncrementSecs * 1000) < 1000 * 60 * 60) {
-      fakeExecutor.advanceTime(Duration.standardSeconds(timeIncrementSecs));
-
+    while (fakeExecutor.getClock().millisTime() + timeIncrementSecs - 1 < 1000 * 60 * 60) {
+      timeIncrementSecs *= 2;
       final int expectedIncrementSecs = Math.min(600, timeIncrementSecs);
       assertEquivalentWithTransformation(
           testAckIdsBatch,
@@ -363,7 +369,7 @@ public ModifyAckDeadline apply(String ack) {
               return new ModifyAckDeadline(ack, expectedIncrementSecs);
             }
           });
-      timeIncrementSecs *= 2; 
+      fakeExecutor.advanceTime(Duration.standardSeconds(timeIncrementSecs - 1));
     }
 
     // No more modify ack deadline extension should be triggered at this point