@@ -72,7 +72,7 @@ class MessageDispatcher {
   private final AckProcessor ackProcessor;
 
   private final FlowController flowController;
-  private final MessageWaiter messagesWaiter;
+  private final Waiter messagesWaiter;
 
   // Maps ID to "total expiration time". If it takes longer than this, stop extending.
   private final ConcurrentMap<String, AckHandler> pendingMessages = new ConcurrentHashMap<>();
@@ -145,7 +145,7 @@ private void forget() {
         return;
       }
       flowController.release(1, outstandingBytes);
-      messagesWaiter.incrementPendingMessages(-1);
+      messagesWaiter.incrementPendingCount(-1);
     }
 
     @Override
@@ -205,7 +205,7 @@ void sendAckOperations(
     // 601 buckets of 1s resolution from 0s to MAX_ACK_DEADLINE_SECONDS
     this.ackLatencyDistribution = ackLatencyDistribution;
     jobLock = new ReentrantLock();
-    messagesWaiter = new MessageWaiter();
+    messagesWaiter = new Waiter();
     this.clock = clock;
     this.sequentialExecutor = new SequentialExecutorService.AutoExecutor(executor);
   }
@@ -268,7 +268,7 @@ public void run() {
   }
 
   void stop() {
-    messagesWaiter.waitNoMessages();
+    messagesWaiter.waitComplete();
     jobLock.lock();
     try {
       if (backgroundJob != null) {
@@ -331,9 +331,9 @@ void processReceivedMessages(List<ReceivedMessage> messages) {
   }
 
   private void processBatch(List<OutstandingMessage> batch) {
-    messagesWaiter.incrementPendingMessages(batch.size());
+    messagesWaiter.incrementPendingCount(batch.size());
     for (OutstandingMessage message : batch) {
-      // This is a blocking flow controller.  We have already incremented MessageWaiter, so
+      // This is a blocking flow controller.  We have already incremented messagesWaiter, so
       // shutdown will block on processing of all these messages anyway.
       try {
         flowController.reserve(1, message.receivedMessage.getMessage().getSerializedSize());
 
@@ -104,7 +104,7 @@ public class Publisher {
 
   private final AtomicBoolean shutdown;
   private final BackgroundResource backgroundResources;
-  private final MessageWaiter messagesWaiter;
+  private final Waiter messagesWaiter;
   private ScheduledFuture<?> currentAlarmFuture;
   private final ApiFunction<PubsubMessage, PubsubMessage> messageTransform;
 
@@ -173,7 +173,7 @@ private Publisher(Builder builder) throws IOException {
     backgroundResourceList.add(publisherStub);
     backgroundResources = new BackgroundResourceAggregation(backgroundResourceList);
     shutdown = new AtomicBoolean(false);
-    messagesWaiter = new MessageWaiter();
+    messagesWaiter = new Waiter();
   }
 
   /** Topic which the publisher publishes to. */
@@ -249,7 +249,7 @@ public ApiFuture<String> publish(PubsubMessage message) {
       messagesBatchLock.unlock();
     }
 
-    messagesWaiter.incrementPendingMessages(1);
+    messagesWaiter.incrementPendingCount(1);
 
     // For messages without ordering keys, it is okay to send batches without holding
     // messagesBatchLock.
@@ -423,7 +423,7 @@ public void onSuccess(PublishResponse result) {
                 }
               }
             } finally {
-              messagesWaiter.incrementPendingMessages(-outstandingBatch.size());
+              messagesWaiter.incrementPendingCount(-outstandingBatch.size());
             }
           }
 
@@ -432,7 +432,7 @@ public void onFailure(Throwable t) {
             try {
               outstandingBatch.onFailure(t);
             } finally {
-              messagesWaiter.incrementPendingMessages(-outstandingBatch.size());
+              messagesWaiter.incrementPendingCount(-outstandingBatch.size());
             }
           }
         };
 
@@ -72,6 +72,7 @@ final class StreamingSubscriberConnection extends AbstractApiService implements
 
   private final AtomicLong channelReconnectBackoffMillis =
       new AtomicLong(INITIAL_CHANNEL_RECONNECT_BACKOFF.toMillis());
+  private final Waiter ackOperationsWaiter = new Waiter();
 
   private final Lock lock = new ReentrantLock();
   private ClientStream<StreamingPullRequest> clientStream;
@@ -116,6 +117,7 @@ protected void doStart() {
   @Override
   protected void doStop() {
     messageDispatcher.stop();
+    ackOperationsWaiter.waitComplete();
 
     lock.lock();
     try {
@@ -273,16 +275,18 @@ public void sendAckOperations(
         new ApiFutureCallback<Empty>() {
           @Override
           public void onSuccess(Empty empty) {
-            // noop
+            ackOperationsWaiter.incrementPendingCount(-1);
           }
 
           @Override
           public void onFailure(Throwable t) {
+            ackOperationsWaiter.incrementPendingCount(-1);
             Level level = isAlive() ? Level.WARNING : Level.FINER;
             logger.log(level, "failed to send operations", t);
           }
         };
 
+    int pendingOperations = 0;
     for (PendingModifyAckDeadline modack : ackDeadlineExtensions) {
       for (List<String> idChunk : Lists.partition(modack.ackIds, MAX_PER_REQUEST_CHANGES)) {
         ApiFuture<Empty> future =
@@ -294,6 +298,7 @@ public void onFailure(Throwable t) {
                         .setAckDeadlineSeconds(modack.deadlineExtensionSeconds)
                         .build());
         ApiFutures.addCallback(future, loggingCallback, directExecutor());
+        pendingOperations++;
       }
     }
 
@@ -306,6 +311,9 @@ public void onFailure(Throwable t) {
                       .addAllAckIds(idChunk)
                       .build());
       ApiFutures.addCallback(future, loggingCallback, directExecutor());
+      pendingOperations++;
     }
+
+    ackOperationsWaiter.incrementPendingCount(pendingOperations);
   }
 }
@@ -305,6 +305,7 @@ public void run() {
                   // stop connection is no-op if connections haven't been started.
                   stopAllStreamingConnections();
                   shutdownBackgroundResources();
+                  subStub.shutdownNow();
                   notifyStopped();
                 } catch (Exception e) {
                   notifyFailed(e);
 
@@ -18,25 +18,28 @@
 
 import com.google.api.core.InternalApi;
 
-/** A barrier kind of object that helps to keep track and synchronously wait on pending messages. */
-class MessageWaiter {
-  private int pendingMessages;
+/**
+ * A barrier kind of object that helps keep track of pending actions and synchronously wait until
+ * all have completed.
+ */
+class Waiter {
+  private int pendingCount;
 
-  MessageWaiter() {
-    pendingMessages = 0;
+  Waiter() {
+    pendingCount = 0;
   }
 
-  public synchronized void incrementPendingMessages(int messages) {
-    this.pendingMessages += messages;
-    if (pendingMessages == 0) {
+  public synchronized void incrementPendingCount(int delta) {
+    this.pendingCount += delta;
+    if (pendingCount == 0) {
       notifyAll();
     }
   }
 
-  public synchronized void waitNoMessages() {
+  public synchronized void waitComplete() {
     boolean interrupted = false;
     try {
-      while (pendingMessages > 0) {
+      while (pendingCount > 0) {
         try {
           wait();
         } catch (InterruptedException e) {
@@ -52,7 +55,7 @@ public synchronized void waitNoMessages() {
   }
 
   @InternalApi
-  public int pendingMessages() {
-    return pendingMessages;
+  public int pendingCount() {
+    return pendingCount;
   }
 }
@@ -22,14 +22,14 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Tests for {@link MessageWaiter}. */
+/** Tests for {@link Waiter}. */
 @RunWith(JUnit4.class)
-public class MessageWaiterTest {
+public class WaiterTest {
 
   @Test
   public void test() throws Exception {
-    final MessageWaiter waiter = new MessageWaiter();
-    waiter.incrementPendingMessages(1);
+    final Waiter waiter = new Waiter();
+    waiter.incrementPendingCount(1);
 
     final Thread mainThread = Thread.currentThread();
     Thread t =
@@ -40,14 +40,14 @@ public void run() {
                 while (mainThread.getState() != Thread.State.WAITING) {
                   Thread.yield();
                 }
-                waiter.incrementPendingMessages(-1);
+                waiter.incrementPendingCount(-1);
               }
             });
     t.start();
 
-    waiter.waitNoMessages();
+    waiter.waitComplete();
     t.join();
 
-    assertEquals(0, waiter.pendingMessages());
+    assertEquals(0, waiter.pendingCount());
   }
 }
Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ public class Publisher {`
`104`	`104`
`105`	`105`	`private final AtomicBoolean shutdown;`
`106`	`106`	`private final BackgroundResource backgroundResources;`
`107`		`- private final MessageWaiter messagesWaiter;`
	`107`	`+ private final Waiter messagesWaiter;`
`108`	`108`	`private ScheduledFuture<?> currentAlarmFuture;`
`109`	`109`	`private final ApiFunction<PubsubMessage, PubsubMessage> messageTransform;`
`110`	`110`
`@@ -173,7 +173,7 @@ private Publisher(Builder builder) throws IOException {`
`173`	`173`	`backgroundResourceList.add(publisherStub);`
`174`	`174`	`backgroundResources = new BackgroundResourceAggregation(backgroundResourceList);`
`175`	`175`	`shutdown = new AtomicBoolean(false);`
`176`		`- messagesWaiter = new MessageWaiter();`
	`176`	`+ messagesWaiter = new Waiter();`
`177`	`177`	`}`
`178`	`178`
`179`	`179`	`/** Topic which the publisher publishes to. */`
`@@ -249,7 +249,7 @@ public ApiFuture<String> publish(PubsubMessage message) {`
`249`	`249`	`messagesBatchLock.unlock();`
`250`	`250`	`}`
`251`	`251`
`252`		`- messagesWaiter.incrementPendingMessages(1);`
	`252`	`+ messagesWaiter.incrementPendingCount(1);`
`253`	`253`
`254`	`254`	`// For messages without ordering keys, it is okay to send batches without holding`
`255`	`255`	`// messagesBatchLock.`
`@@ -423,7 +423,7 @@ public void onSuccess(PublishResponse result) {`
`423`	`423`	`}`
`424`	`424`	`}`
`425`	`425`	`} finally {`
`426`		`- messagesWaiter.incrementPendingMessages(-outstandingBatch.size());`
	`426`	`+ messagesWaiter.incrementPendingCount(-outstandingBatch.size());`
`427`	`427`	`}`
`428`	`428`	`}`
`429`	`429`
`@@ -432,7 +432,7 @@ public void onFailure(Throwable t) {`
`432`	`432`	`try {`
`433`	`433`	`outstandingBatch.onFailure(t);`
`434`	`434`	`} finally {`
`435`		`- messagesWaiter.incrementPendingMessages(-outstandingBatch.size());`
	`435`	`+ messagesWaiter.incrementPendingCount(-outstandingBatch.size());`
`436`	`436`	`}`
`437`	`437`	`}`
`438`	`438`	`};`