Copilot commented on code in PR #3844:
URL: https://github.com/apache/celeborn/pull/3844#discussion_r4018347012
##########
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;
Review Comment:
When `waitOnTermination()` is interrupted while `activeProducers > 0`, this
path interrupts `pushThread` but leaves the lifecycle in `CLOSING`. The push
thread exits on interruption without setting `exceptionRef`, so an admitted
producer blocked in `dataPushQueue.addPushTask()` can loop forever (or enqueue
successfully after the pusher has stopped) because its retry loop only stops on
an exception. Keep the pusher alive for admitted producers or explicitly
fail/wake them and complete the lifecycle on this cancellation path.
##########
client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java:
##########
@@ -178,6 +182,110 @@ 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);
+
+ AtomicReference<Throwable> terminationFailure = new AtomicReference<>();
+ Thread terminationThread =
+ new Thread(
+ () -> {
+ try {
+ dataPusher.waitOnTermination();
+ } catch (Throwable t) {
+ terminationFailure.set(t);
+ }
+ });
+ terminationThread.start();
+ Assert.assertTrue(pushStarted.await(10, TimeUnit.SECONDS));
+ // Give the termination thread time to close the lifecycle admission
barrier.
+ Thread.sleep(100);
Review Comment:
`Thread.sleep(100)` does not synchronize with the transition to `CLOSING`;
if the termination thread is delayed before acquiring `lifecycleLock`, the
second `addTask()` can be admitted and this test will fail intermittently. Use
a deterministic signal or await tied to the admission-barrier transition
instead of a fixed sleep.
--
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]
