From 788111156adfe3b4eb5b8882dcf77ef041eaff5a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 6 Aug 2026 12:02:58 +0900 Subject: [PATCH 1/2] 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 --- .../com/google/cloud/pubsub/v1/Publisher.java | 8 ++- .../cloud/pubsub/v1/PublisherImplTest.java | 53 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index 56c920bcfdc1..cc9420a07fff 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java +++ b/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)); } } }; diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 8e6efaf372c9..1d415b3fe20a 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java +++ b/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 publishFuture1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA"); + ApiFuture publishFuture2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA"); + ApiFuture 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 sendTestMessageWithOrderingKey( Publisher publisher, String data, String orderingKey) { return publisher.publish( From 38b4d423d6018a672c5a7a2c2aee4b4e54fc7eed Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 6 Aug 2026 12:15:45 +0900 Subject: [PATCH 2/2] 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. --- .../java/com/google/cloud/pubsub/v1/Publisher.java | 11 ++++++----- .../google/cloud/pubsub/v1/PublisherImplTest.java | 12 ++---------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index cc9420a07fff..edd1f8fd863f 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java @@ -322,6 +322,10 @@ public ApiFuture 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 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()) { diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 1d415b3fe20a..512181a507b9 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java +++ b/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 publishFuture1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA"); ApiFuture publishFuture2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA"); ApiFuture publishFuture3 = sendTestMessageWithOrderingKey(publisher, "m3", "orderA");