Skip to content

Commit e38bbc0

Browse files
authored
Simplify exchange release logic (#128084)
1 parent f4b6086 commit e38bbc0

File tree

1 file changed

+5
-7
lines changed

1 file changed

+5
-7
lines changed

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)