fix(pubsub): count a published message while its batch lock is held

The failure callback decrements messagesWaiter for the messages it cancels
out of a MessagesBatch, but publish() incremented after releasing
messagesBatchLock. A message visible in the batch and not yet counted would
therefore be decremented for without ever having been counted, taking
pendingCount below zero and letting waitComplete() return early.

Incrementing while the lock is still held makes "in a MessagesBatch" and
"counted" one state. Lock ordering is messagesBatchLock -> Waiter monitor
here and nowhere the reverse, and incrementPendingCount never blocks. The
paused-key path still returns before the increment, as before.

Author: laughingman7743

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

@@ -322,6 +322,10 @@ public ApiFuture<String> publish(PubsubMessage message) {
       }
 
       batchesToSend = messagesBatch.add(outstandingPublish);
+      // Counted while messagesBatchLock is held, so that "in a MessagesBatch" and "counted" are
+      // one state: the failure callback decrements for what it cancels out of a MessagesBatch.
+      // Lock ordering is messagesBatchLock -> Waiter monitor here and nowhere the reverse.
+      messagesWaiter.incrementPendingCount(1);
       if (!batchesToSend.isEmpty() && messagesBatch.isEmpty()) {
         messagesBatches.remove(orderingKey);
       }
@@ -340,8 +344,6 @@ public ApiFuture<String> publish(PubsubMessage message) {
       messagesBatchLock.unlock();
     }
 
-    messagesWaiter.incrementPendingCount(1);
-
     // For messages without ordering keys, it is okay to send batches without holding
     // messagesBatchLock.
     if (!batchesToSend.isEmpty() && orderingKey.isEmpty()) {
@@ -546,9 +548,8 @@ 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.
+            // Cancelled below without ever becoming part of an OutstandingBatch, so nothing
+            // else will decrement for them.
             int cancelledMessagesCount = 0;
             try {
               if (outstandingBatch.orderingKey != null && !outstandingBatch.orderingKey.isEmpty()) {

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

@@ -643,13 +643,6 @@ 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 =
@@ -665,9 +658,8 @@ public void testShutdownAfterOrderingKeyFailureWithMoreOfThatKeyStillBatched() t
     // 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.
+    // m1 and m2 meet the threshold, but the request only leaves once the fake executor runs, so
+    // m3 lands in the un-flushed batch for the same key and is still there when the failure does.
     ApiFuture<String> publishFuture1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA");
     ApiFuture<String> publishFuture2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA");
     ApiFuture<String> publishFuture3 = sendTestMessageWithOrderingKey(publisher, "m3", "orderA");