fix(pubsub): return cancelled messages to the publisher's waiter

When a batch for an ordering key fails, the failure callback also cancels
the messages still accumulating in that key's un-flushed MessagesBatch and
drops the batch, but decrements messagesWaiter only by the in-flight batch's
size. Those cancelled messages each incremented the waiter when they were
published and never become part of any OutstandingBatch, so nothing ever
decrements for them and pendingCount can no longer reach zero.

Publisher.shutdown() waits on that counter uninterruptibly and without a
timeout, so it never returns. awaitTermination(timeout, unit) is documented
to be called after shutdown(), so its bound is never reached either.

Return the cancelled count to the waiter alongside the batch's own.

Fixes #14001

Author: laughingman7743

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java

@@ -546,6 +546,10 @@ public void onSuccess(PublishResponse result) {
 
           @Override
           public void onFailure(Throwable t) {
+            // Messages cancelled below are dropped without ever becoming part of an
+            // OutstandingBatch, so they are owed back to messagesWaiter here; nothing else will
+            // ever decrement for them.
+            int cancelledMessagesCount = 0;
             try {
               if (outstandingBatch.orderingKey != null && !outstandingBatch.orderingKey.isEmpty()) {
                 messagesBatchLock.lock();
@@ -556,6 +560,7 @@ public void onFailure(Throwable t) {
                       outstanding.publishResult.setException(
                           SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION);
                     }
+                    cancelledMessagesCount = messagesBatch.getMessagesCount();
                     messagesBatches.remove(outstandingBatch.orderingKey);
                   }
                 } finally {
@@ -564,7 +569,8 @@ public void onFailure(Throwable t) {
               }
               outstandingBatch.onFailure(t);
             } finally {
-              messagesWaiter.incrementPendingCount(-outstandingBatch.size());
+              messagesWaiter.incrementPendingCount(
+                  -(outstandingBatch.size() + cancelledMessagesCount));
             }
           }
         };

java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java

@@ -643,6 +643,59 @@ public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws E
     }
   }
 
+  /**
+   * When a batch for an ordering key fails, its failure callback also cancels the messages still
+   * accumulating in that key's un-flushed batch. Those messages incremented {@code messagesWaiter}
+   * when they were published and never become part of any {@code OutstandingBatch}, so they have to
+   * be returned to the waiter there — otherwise {@code pendingCount} can never reach zero again and
+   * {@code shutdown()}, which waits on it uninterruptibly and without a timeout, never returns.
+   */
+  @Test(timeout = 60_000)
+  public void testShutdownAfterOrderingKeyFailureWithMoreOfThatKeyStillBatched() throws Exception {
+    Publisher publisher =
+        getTestPublisherBuilder()
+            .setBatchingSettings(
+                Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder()
+                    .setElementCountThreshold(2L)
+                    .setDelayThresholdDuration(Duration.ofSeconds(100))
+                    .build())
+            .setEnableMessageOrdering(true)
+            .build();
+
+    // Queued before publishing, so the fake never blocks in publishResponses.take() (see #13394).
+    testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT));
+
+    // m1 and m2 meet the threshold and are popped into an outstanding batch, but the request only
+    // leaves once the fake executor runs — so m3 is published into the un-flushed batch for the
+    // same key first, and is still there when the failure lands.
+    ApiFuture<String> publishFuture1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA");
+    ApiFuture<String> publishFuture2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA");
+    ApiFuture<String> publishFuture3 = sendTestMessageWithOrderingKey(publisher, "m3", "orderA");
+    assertFalse(publishFuture3.isDone());
+
+    fakeExecutor.advanceTime(Duration.ZERO);
+
+    try {
+      publishFuture1.get();
+      fail("This should fail.");
+    } catch (ExecutionException e) {
+    }
+    try {
+      publishFuture2.get();
+      fail("This should fail.");
+    } catch (ExecutionException e) {
+    }
+    try {
+      publishFuture3.get();
+      fail("This should fail.");
+    } catch (ExecutionException e) {
+      assertEquals(SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION, e.getCause());
+    }
+
+    // Hangs here without the accounting fix: m3's increment was never returned.
+    shutdownTestPublisher(publisher);
+  }
+
   private ApiFuture<String> sendTestMessageWithOrderingKey(
       Publisher publisher, String data, String orderingKey) {
     return publisher.publish(