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..6f338f46cf46 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,12 @@ public ApiFuture publish(PubsubMessage message) { } batchesToSend = messagesBatch.add(outstandingPublish); + // Counted while messagesBatchLock is still held, so that "in a MessagesBatch" and "counted" + // are one state. The failure callback decrements for the messages it cancels out of a + // MessagesBatch, so one that is visible there but not yet counted would take pendingCount + // below zero. Lock ordering is messagesBatchLock -> Waiter monitor here and nowhere the + // reverse, and incrementPendingCount never blocks. + messagesWaiter.incrementPendingCount(1); if (!batchesToSend.isEmpty() && messagesBatch.isEmpty()) { messagesBatches.remove(orderingKey); } @@ -340,8 +346,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,6 +550,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 +564,7 @@ public void onFailure(Throwable t) { outstanding.publishResult.setException( SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION); } + cancelledMessagesCount = messagesBatch.getMessagesCount(); messagesBatches.remove(outstandingBatch.orderingKey); } } finally { @@ -564,7 +573,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(