[CELEBORN-2459] Fix race causing DataPusher to accept data during termination and drop it

[rjvkr2021]

What changes were proposed in this pull request?

This PR fixes a race causing DataPusher to accept data during termination and drop it.

Why are the changes needed?

Background

When a task wants to push shuffle data to a Celeborn worker, it does not push the
data directly. Instead, the data is copied into a reusable buffer and enqueued
for asynchronous processing.

The producer returns successfully after the buffer is accepted by the queue. A
background thread later dequeues the buffer, pushes its data to the worker, and
returns the buffer to the reusable buffer pool.

Race

The race occurs when the background thread is being terminated. In the original
implementation, termination waits for the reusable buffer pool (idleQueue)
to become full and then sets terminated=true.

The check and the state change are not an atomic admission barrier. After the
buffer-pool check succeeds but before terminated is set, a producer can
borrow a reusable buffer, copy data into it, enqueue it successfully, and
return from addTask().

The background thread can then observe terminated, stop without processing
the newly enqueued buffer, and termination can clear the queue.

Impact

The producer has reported success even though the data was never pushed to the
worker. This causes silent data loss.

Fix

Task admission and background thread termination must be ordered through an explicit lifecycle:

RUNNING -> CLOSING -> TERMINATED

When termination begins, the pusher transitions to CLOSING before waiting for
existing work to finish. New calls to addTask() are rejected in CLOSING,
while producers admitted in RUNNING are allowed to complete. After all
admitted producers and queued tasks have been drained, the pusher transitions
to TERMINATED.

This guarantees that every successful addTask() corresponds to work admitted
before termination and therefore handled before termination completes.

Does this PR resolve a correctness bug?

• Yes

Does this PR introduce any user-facing change?

No

How was this patch tested?

Added following unit tests:

1. testDataPusherRejectsTasksWhileClosing
2. testDataPusherDrainsTaskAdmittedBeforeTermination

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Pull request overview

Fixes a race in DataPusher’s shutdown path where tasks could be admitted during termination and then dropped, by introducing an explicit lifecycle state machine and admission barrier.

Changes:

• Add RUNNING -> CLOSING -> TERMINATED lifecycle gating for addTask() vs waitOnTermination().
• Track and drain in-flight producers before allowing termination to proceed.
• Add unit tests covering admission rejection while closing and draining of pre-admitted work.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
client/src/main/java/org/apache/celeborn/client/write/DataPusher.java	Introduces lifecycle state + producer-drain coordination to prevent task admission during termination.
client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java	Adds unit tests that exercise the new closing/termination behavior under concurrency.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[rjvkr2021]

updated the exception message accordingly

[rjvkr2021]

@cxzl25 can you please review the pr? i have addressed all the comments left by copilot.

[rjvkr2021]

@cxzl25 just a gentle reminder.

[cxzl25]

Is this a real production environment case?

[rjvkr2021]

I don’t know. Actually, I was investigating something else when I discovered this race.