@@ -248,9 +248,6 @@ private class DataNodeRequestExecutor {
248248 }
249249
250250 void start () {
251- parentTask .addListener (
252- () -> exchangeService .finishSinkHandler (request .sessionId (), new TaskCancelledException (parentTask .getReasonCancelled ()))
253- );
254251 runBatch (0 );
255252 }
256253
@@ -419,7 +416,12 @@ private void runComputeOnDataNode(
419416 var parentListener = computeListener .acquireAvoid ();
420417 try {
421418 // run compute with target shards
419+ var externalSink = exchangeService .getSinkHandler (externalId );
422420 var internalSink = exchangeService .createSinkHandler (request .sessionId (), request .pragmas ().exchangeBufferSize ());
421+ task .addListener (() -> {
422+ exchangeService .finishSinkHandler (externalId , new TaskCancelledException (task .getReasonCancelled ()));
423+ exchangeService .finishSinkHandler (request .sessionId (), new TaskCancelledException (task .getReasonCancelled ()));
424+ });
423425 DataNodeRequestExecutor dataNodeRequestExecutor = new DataNodeRequestExecutor (
424426 request ,
425427 task ,
@@ -431,10 +433,6 @@ private void runComputeOnDataNode(
431433 );
432434 dataNodeRequestExecutor .start ();
433435 // run the node-level reduction
434- var externalSink = exchangeService .getSinkHandler (externalId );
435- task .addListener (
436- () -> exchangeService .finishSinkHandler (externalId , new TaskCancelledException (task .getReasonCancelled ()))
437- );
438436 var exchangeSource = new ExchangeSourceHandler (1 , esqlExecutor );
439437 exchangeSource .addRemoteSink (internalSink ::fetchPageAsync , true , () -> {}, 1 , ActionListener .noop ());
440438 var reductionListener = computeListener .acquireCompute ();
0 commit comments