Pub/Sub: Shutting down Subscriber executor (#5468)

sduskis · web-flow · commit a560bb578f1f · 2019-06-14T10:37:03.000-04:00
`MessageDispatcher` has a scheduled future which needed to be shut down so that the `Subscriber` can shutdown properly.  Added a test to make sure that this doesn't happen again.
diff --git a/google-cloud-clients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java b/google-cloud-clients/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java
@@ -86,6 +86,7 @@ class MessageDispatcher {
   private final AtomicBoolean extendDeadline = new AtomicBoolean(true);
   private final Lock jobLock;
   private ScheduledFuture<?> backgroundJob;
+  private ScheduledFuture<?> setExtendedDeadlineFuture;
 
   // To keep track of number of seconds the receiver takes to process messages.
   private final Distribution ackLatencyDistribution;
@@ -239,11 +240,15 @@ public void run() {
                       int newDeadlineSec = computeDeadlineSeconds();
                       messageDeadlineSeconds.set(newDeadlineSec);
                       extendDeadlines();
-                      // Don't bother cancelling this when we stop. It'd just set an atomic boolean.
-                      systemExecutor.schedule(
-                          setExtendDeadline,
-                          newDeadlineSec - ackExpirationPadding.getSeconds(),
-                          TimeUnit.SECONDS);
+                      if (setExtendedDeadlineFuture != null && !backgroundJob.isDone()) {
+                        setExtendedDeadlineFuture.cancel(true);
+                      }
+
+                      setExtendedDeadlineFuture =
+                          systemExecutor.schedule(
+                              setExtendDeadline,
+                              newDeadlineSec - ackExpirationPadding.getSeconds(),
+                              TimeUnit.SECONDS);
                     }
                     processOutstandingAckOperations();
                   } catch (Throwable t) {
@@ -266,8 +271,12 @@ void stop() {
     try {
       if (backgroundJob != null) {
         backgroundJob.cancel(false);
-        backgroundJob = null;
       }
+      if (setExtendedDeadlineFuture != null) {
+        setExtendedDeadlineFuture.cancel(true);
+      }
+      backgroundJob = null;
+      setExtendedDeadlineFuture = null;
     } finally {
       jobLock.unlock();
     }
diff --git a/google-cloud-clients/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/SubscriberTest.java b/google-cloud-clients/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/SubscriberTest.java
@@ -19,6 +19,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
+import com.google.api.gax.core.ExecutorProvider;
 import com.google.api.gax.core.FixedExecutorProvider;
 import com.google.api.gax.core.InstantiatingExecutorProvider;
 import com.google.api.gax.core.NoCredentialsProvider;
@@ -36,6 +37,9 @@
 import io.grpc.StatusException;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -112,6 +116,38 @@ public void testFailedChannel_recoverableError_channelReopened() throws Exceptio
     subscriber.stopAsync().awaitTerminated();
   }
 
+  @Test
+  public void testFailedChannel_testTerminated() throws Exception {
+    final ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(5);
+    ExecutorProvider provider =
+        new ExecutorProvider() {
+          @Override
+          public boolean shouldAutoClose() {
+            return true;
+          }
+
+          @Override
+          public ScheduledExecutorService getExecutor() {
+            return scheduledExecutorService;
+          }
+        };
+
+    try {
+      Subscriber subscriber =
+          startSubscriber(
+              getTestSubscriberBuilder(testReceiver).setSystemExecutorProvider(provider));
+
+      // wait long enough for the MessageDispatcher to set up, which at one point
+      // caused shutdown problems.
+      Thread.sleep(100);
+      subscriber.stopAsync().awaitTerminated();
+
+      assertTrue(scheduledExecutorService.awaitTermination(10, TimeUnit.SECONDS));
+    } finally {
+      scheduledExecutorService.shutdownNow();
+    }
+  }
+
   @Test(expected = IllegalStateException.class)
   public void testFailedChannel_fatalError_subscriberFails() throws Exception {
     Subscriber subscriber =
