Copilot commented on code in PR #3844:
URL: https://github.com/apache/celeborn/pull/3844#discussion_r3967051077
##########
client/src/main/java/org/apache/celeborn/client/write/DataPusher.java:
##########
@@ -170,20 +191,46 @@ public void addTask(int partitionId, byte[] buffer, int
size)
logger.error("DataPusher thread interrupted while adding push task.");
pushThread.interrupt();
throw e;
+ } finally {
+ lifecycleLock.lock();
+ try {
+ activeProducers--;
+ if (activeProducers == 0) {
+ producersDrained.signalAll();
+ }
+ } finally {
+ lifecycleLock.unlock();
+ }
}
}
public void waitOnTermination() throws IOException, InterruptedException {
+ lifecycleLock.lockInterruptibly();
try {
- idleLock.lockInterruptibly();
- waitIdleQueueFullWithLock();
+ if (lifecycleState == LifecycleState.RUNNING) {
+ lifecycleState = LifecycleState.CLOSING;
+ }
+ while (activeProducers > 0) {
+ producersDrained.await();
+ }
} catch (InterruptedException e) {
logger.error("DataPusher thread interrupted while waitOnTermination.");
pushThread.interrupt();
throw e;
+ } finally {
+ lifecycleLock.unlock();
}
- terminated = true;
+ idleLock.lockInterruptibly();
+ waitIdleQueueFullWithLock();
+
+ lifecycleLock.lock();
+ try {
+ lifecycleState = LifecycleState.TERMINATED;
+ terminated = true;
+ } finally {
+ lifecycleLock.unlock();
+ }
Review Comment:
`idleLock` is acquired but never released in this method (at least in the
shown hunk). If `waitIdleQueueFullWithLock()` does not reliably unlock (its
name suggests it expects the lock to be held), this will leave `idleLock`
permanently locked, potentially deadlocking producers / the push thread. Wrap
the `idleLock.lockInterruptibly()` + `waitIdleQueueFullWithLock()` section in a
`try/finally` that always unlocks `idleLock`, and avoid holding `idleLock`
while acquiring `lifecycleLock` (unlock `idleLock` before taking
`lifecycleLock`) to prevent lock-order deadlocks.
##########
client/src/main/java/org/apache/celeborn/client/write/DataPusher.java:
##########
@@ -170,20 +191,46 @@ public void addTask(int partitionId, byte[] buffer, int
size)
logger.error("DataPusher thread interrupted while adding push task.");
pushThread.interrupt();
throw e;
+ } finally {
+ lifecycleLock.lock();
+ try {
+ activeProducers--;
+ if (activeProducers == 0) {
+ producersDrained.signalAll();
+ }
+ } finally {
+ lifecycleLock.unlock();
+ }
}
}
public void waitOnTermination() throws IOException, InterruptedException {
+ lifecycleLock.lockInterruptibly();
try {
- idleLock.lockInterruptibly();
- waitIdleQueueFullWithLock();
+ if (lifecycleState == LifecycleState.RUNNING) {
+ lifecycleState = LifecycleState.CLOSING;
+ }
Review Comment:
PR description says this does **not** resolve a correctness bug, but it
explicitly fixes a silent data-loss race (successful `addTask()` followed by
dropped work). Please update the PR description/checkbox to reflect that this
is a correctness fix so downstream release notes and triage remain accurate.
##########
client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java:
##########
@@ -178,6 +181,103 @@ protected void pushData(PushTask task) throws IOException
{
client.shutdown();
}
+ @Test
+ public void testDataPusherRejectsTasksWhileClosing() throws Exception {
+ CelebornConf conf = new CelebornConf();
+ DummyShuffleClient client =
+ new DummyShuffleClient(conf, new File(tempDir,
UUID.randomUUID().toString()));
+ client.initReducePartitionMap(0, 1, 1);
+ LongAdder[] mapStatusLengths = {new LongAdder()};
+ CountDownLatch pushStarted = new CountDownLatch(1);
+ CountDownLatch allowPush = new CountDownLatch(1);
+ DataPusher dataPusher =
+ new DataPusher(0, 0, 0, 0, 1, 1, conf, client, null, integer -> {},
mapStatusLengths) {
+ @Override
+ protected void pushData(PushTask task) throws IOException {
+ pushStarted.countDown();
+ try {
+ allowPush.await(10, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ }
+ };
+
+ dataPusher.addTask(0, new byte[1], 1);
+
+ Thread terminationThread =
+ new Thread(
+ () -> {
+ try {
+ dataPusher.waitOnTermination();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
Review Comment:
Exceptions thrown inside `terminationThread` won’t automatically fail the
JUnit test, so these tests can pass even if `waitOnTermination()` throws.
Capture any thrown exception in a shared holder (e.g.,
`AtomicReference<Throwable>` or a `Future` via `ExecutorService`) and assert it
is null after `join()`. This will make the tests reliably detect failures
instead of silently ignoring them.
##########
client/src/main/java/org/apache/celeborn/client/write/DataPusher.java:
##########
@@ -153,8 +164,18 @@ public void run() {
public void addTask(int partitionId, byte[] buffer, int size)
throws IOException, InterruptedException {
- client.computeBatchCRC(shuffleId, mapId, attemptId, partitionId, buffer,
0, size);
+ lifecycleLock.lockInterruptibly();
try {
+ if (lifecycleState != LifecycleState.RUNNING) {
+ throw new IOException("DataPusher is closing or terminated");
Review Comment:
The new error message is ambiguous and loses useful debugging context.
Consider including the actual state (e.g., `lifecycleState`) and the operation
being rejected (e.g., `addTask rejected because state=CLOSING`) to make
production logs actionable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
