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..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,6 +548,9 @@ public void onSuccess(PublishResponse result) { @Override public void onFailure(Throwable t) { + // 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()) { messagesBatchLock.lock(); @@ -556,6 +561,7 @@ public void onFailure(Throwable t) { outstanding.publishResult.setException( SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION); } + cancelledMessagesCount = messagesBatch.getMessagesCount(); messagesBatches.remove(outstandingBatch.orderingKey); } } finally { @@ -564,7 +570,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..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,6 +643,51 @@ public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws E } } + @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, 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"); + 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(