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

[laughingman7743]

Fixes #14001.

The defect

Publisher.publish counts every accepted message into messagesWaiter. Both batch callbacks
decrement by the size of the batch that was in flight. But before that finally, the failure
callback also cancels the messages still accumulating in the un-flushed MessagesBatch for the
failed ordering key, and drops the batch:

MessagesBatch messagesBatch = messagesBatches.get(outstandingBatch.orderingKey);
if (messagesBatch != null) {
  for (OutstandingPublish outstanding : messagesBatch.messages) {
    outstanding.publishResult.setException(
        SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION);
  }
  messagesBatches.remove(outstandingBatch.orderingKey);
}

Each of those incremented the waiter when it was published, and none is part of any
OutstandingBatch, so nothing ever decrements for them. pendingCount stays at least equal to the
number cancelled, for the life of the publisher.

Waiter.waitComplete() ignores interruption and wakes only on an exact zero, and
Publisher.shutdown() calls it directly — so the calling thread parks permanently.
awaitTermination(timeout, unit) is the API that takes a bound, but it is documented to be called
after shutdown(), so it is never reached.

The change

Count what the failure callback cancels and return it to the waiter alongside the batch's own size.
MessagesBatch.getMessagesCount() already exists, so the change is local to that callback.

Verification

PublisherImplTest.testShutdownAfterOrderingKeyFailureWithMoreOfThatKeyStillBatched is new and
discriminating — measured both ways, not assumed:

• without the production change: times out, with the stack ending in
  shutdownTestPublisher → Publisher.shutdown → Object.wait, which is messagesWaiter.waitComplete();
• with it: passes in 0.3 s.

The whole of PublisherImplTest passes (33 run, 2 pre-existing @Ignores), and
fmt-maven-plugin:check reports no non-complying files.

The test queues the error before publishing rather than after, so FakePublisherServiceImpl.publish
never blocks in publishResponses.take() — the hazard behind #13394. It needs no response delay:
with the fake executor, the request only leaves once advanceTime runs, so the third message is
deterministically in the un-flushed batch when the failure lands.

One thing worth a maintainer's opinion

publish() increments the waiter after releasing messagesBatchLock, so a message can be in a
MessagesBatch before it is counted. With this patch that can drive pendingCount transiently
negative — self-correcting, since the pending +1 lands and the counter reaches zero, but it means
waitComplete() could return marginally early for a message that was just cancelled anyway.

Moving the increment inside the lock would make "batched" and "counted" one atomic state. The lock
order is already messagesBatchLock → Waiter monitor everywhere, so there is no inversion. I have
deliberately not done that here — it widens the change beyond the accounting bug, and the call is
yours.

[gemini-code-assist]

Code Review

This pull request fixes a bug where the publisher's shutdown process hangs indefinitely after an ordering key failure when there are still un-flushed messages accumulating for that key. It resolves this by tracking and accounting for these cancelled messages in the pending messages waiter. A corresponding unit test was also added. The review feedback highlights a potential race condition where the waiter increment and batching are not atomic, which could cause the pending count to temporarily go negative, and suggests wrapping the new test's execution in a try-finally block to ensure the publisher is always cleaned up.

[gemini-code-assist]

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]

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.

[gemini-code-assist]

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]

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.