Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
laughingman7743 wants to merge 2 commits into
base: main
Choose a base branch
from
+66 −3
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
330af86
fix(pubsub): return cancelled messages to the publisher's waiter
laughingman7743 Aug 6, 2026
4f3fcd8
fix(pubsub): count a published message while its batch lock is held
laughingman7743 Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,12 @@ public ApiFuture<String> 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);
}
Expand All @@ -340,8 +346,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()) {
Expand Down Expand Up @@ -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;

@gemini-code-assist gemini-code-assist Bot Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

While this fix correctly accounts for cancelled messages, there is a potential race condition because publish() increments the waiter after releasing messagesBatchLock. If a message is added to messagesBatch but not yet counted when a failure occurs, cancelledMessagesCount will include it, causing pendingCount to temporarily go negative (or hit zero prematurely). This can cause waitComplete() to return early before the publishing thread actually finishes.

To make "batched" and "counted" atomic and completely eliminate this race, consider moving the messagesWaiter.incrementPendingCount(1) call inside the messagesBatchLock block in the publish() method.

@laughingman7743 laughingman7743 Aug 6, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, and done in 4f3fcd8 — kept as a separate commit so it can be dropped independently of the accounting fix.

incrementPendingCount(1) now happens immediately after messagesBatch.add(...), still under messagesBatchLock, so "in a MessagesBatch" and "counted" are one state. Two things I checked rather than assumed:

  • Lock ordering. It is messagesBatchLock → Waiter monitor here and nowhere the reverse — shutdown() calls publishAllOutstanding() and waitComplete() sequentially, not nested — and incrementPendingCount is a non-blocking synchronized method, so there is no cycle.
  • The paused-key path still does not increment, since it returns at the keyHasError check before reaching the new position.

Also measured that the two commits are genuinely independent: with the increment moved but the accounting fix reverted, the new test still times out in shutdownTestPublisher → shutdown → Object.wait. So this commit is hardening, not a second fix for the same bug.

Full PublisherImplTest is green with both applied (33 run, 2 pre-existing @Ignores), and fmt-maven-plugin:check reports 0 non-complying files.

try {
if (outstandingBatch.orderingKey != null && !outstandingBatch.orderingKey.isEmpty()) {
messagesBatchLock.lock();
Expand All @@ -556,6 +564,7 @@ public void onFailure(Throwable t) {
outstanding.publishResult.setException(
SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION);
}
cancelledMessagesCount = messagesBatch.getMessagesCount();
messagesBatches.remove(outstandingBatch.orderingKey);
}
} finally {
Expand All @@ -564,7 +573,8 @@ public void onFailure(Throwable t) {
}
outstandingBatch.onFailure(t);
} finally {
messagesWaiter.incrementPendingCount(-outstandingBatch.size());
messagesWaiter.incrementPendingCount(
-(outstandingBatch.size() + cancelledMessagesCount));
}
}
};
Expand Down
53 changes: 53 additions & 0 deletions ...ubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Comment on lines +664 to +696

@gemini-code-assist gemini-code-assist Bot Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent resource leaks (such as active background threads or executors) if any assertions fail or unexpected exceptions are thrown during the test, wrap the test execution in a try-finally block to ensure shutdownTestPublisher(publisher) is always executed.

    try {
      // 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());
      }
    } finally {
      // Hangs here without the accounting fix: m3's increment was never returned.
      shutdownTestPublisher(publisher);
    }

@laughingman7743 laughingman7743 Aug 6, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would rather not, for a reason specific to this test — happy to be overruled by a maintainer.

shutdownTestPublisher(publisher) is the assertion here: without the production change it blocks forever, which is what makes the test fail. Moving it into a finally means that when one of the earlier assertions fails, the test hangs in the finally and JUnit reports TestTimedOut instead of the assertion — the diagnostic that actually says what went wrong is replaced by the one that does not.

On the leak itself: tearDown() already shuts the in-process server down and closes the channel unconditionally, and the publisher's executor here is the FakeScheduledExecutorService, not a real pool — so a failed assertion leaves no live thread behind. The other 15 tests in this class call shutdownTestPublisher(publisher) as the last statement without a try/finally, so this also keeps the file consistent.

}

private ApiFuture<String> sendTestMessageWithOrderingKey(
Publisher publisher, String data, String orderingKey) {
return publisher.publish(
Expand Down
Loading