[PabloPauling]

Fixes #40720.

ActiveBundle.close() waited for the ProcessBundleResponse and only then
consumed the bundle's output. The inbound queue holds 100 messages, and since
#38407 the Python SDK's send queue is bounded too (100 MB). The SDK does not
complete a bundle until its output has been sent. So a bundle whose output
exceeds both buffers deadlocks: the SDK blocks sending, and the runner blocks
waiting for a response that never comes. On Flink batch, one bundle is a whole
partition, so this triggers at about 1 GB of output per subtask.

Changes in close():

• It now calls awaitCompletion() on the inbound observer before
  MoreFutures.get(response), so output is consumed while the SDK is still
  producing it.
• Elements embedded in the response are delivered from an asynchronous
  whenCompleteAsync callback. It has to be async: if the response is already
  complete, a synchronous callback would run on the draining thread and could
  block putting into the queue that thread must drain.
• A failed response closes the inbound observer, so the drain does not wait for
  an is_last that will never arrive. The response's error stays the primary
  exception, and the cancelled-queue exception is added as suppressed. The
  existing handleCleanup* tests still pass unchanged.

Tests:

• testOutputLargerThanTheInboundQueueDoesNotDeadlock: the SDK sends 150 data
  messages and only then the response. Times out without the fix.
• testOutputEmbeddedInAnAlreadyCompletedResponseIsDelivered: covers the
  embedded-elements path.

Also adds a CHANGES.md entry under 2.78.0 Bugfixes.

cc @scwhittle (author of #38407)

🤖 Generated with Claude Code

[github-actions]

📋 Review Summary

This PR resolves a deadlock in SdkHarnessClient.ActiveBundle.close() where bundles producing more output than the combined SDK outbound buffer and runner BeamFnDataInboundObserver queue capacity (~100 messages) would block indefinitely because the runner waited for ProcessBundleResponse before draining the inbound data queue. Draining the inbound observer prior to awaiting the response, delivering embedded response elements via whenCompleteAsync, and closing the inbound observer on response failure are well-reasoned and thoroughly tested.

🔍 General Feedback

• The ordering of inboundObserver.awaitCompletion(), MoreFutures.get(response), and exception aggregation preserves primary SDK failure causes while preventing deadlocks on both the data stream and embedded response paths.
• One minor edge case worth considering is avoiding completion/checkpoint/finalization callbacks when drainException != null (i.e., when draining the output failed even though the SDK completed ProcessBundleResponse normally).

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

🟢 Low: If inboundObserver.awaitCompletion() fails (e.g., a downstream output receiver throws an exception or the data stream closes prematurely) while the SDK's ProcessBundleResponse completes normally, drainException is non-null but exception remains null during the MoreFutures.get(response) block above (lines 556–567). As a result, progressHandler.onCompleted, checkpointHandler.onCheckpoint, and finalizationHandler.requestsFinalization will still be invoked before stateRegistration.abort() is called and drainException is thrown.

Consider guarding the completion/checkpoint/finalization callbacks at lines 561–567 with if (drainException == null) so that a bundle whose output drain failed does not report completion or register finalization/residuals.

[PabloPauling]

Good point, fixed in 5b2c223. The completion, checkpoint and finalization callbacks now only run when the output drain succeeded. A bundle whose drain failed reports neither residuals nor finalization (its state registration was already aborted). Covered by testFailedOutputDrainDoesNotReportCompletion. The original code also ran these callbacks before draining, so this tightens pre-existing behaviour too.

[PabloPauling]

/gemini review

[github-actions]

Assigning reviewers:

R: @shunping added as fallback since no labels match configuration
Backup reviewer: Abacn (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[github-actions]

📋 Review Summary

This PR resolves a deadlock in SdkHarnessClient.ActiveBundle.close() where bundles producing more output than the bounded inbound queue and SDK send buffer could hold would hang indefinitely. By draining BeamFnDataInboundObserver before awaiting the ProcessBundleResponse and handling embedded response elements and failure states asynchronously, the fix cleanly breaks the circular wait.

🔍 General Feedback

• The asynchronous handling of ProcessBundleResponse via whenCompleteAsync properly prevents deadlocking on the draining thread when the response is already completed with embedded elements.
• Error propagation and suppression order are well-preserved, and the new unit tests thoroughly cover large outputs, embedded response elements, and receiver failures during draining.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.