Skip to content

Commit

Permalink
implement cancellation logic for StreamProducer's BatchedJob
Browse files Browse the repository at this point in the history
  • Loading branch information
rishabhmaurya committed Nov 20, 2024
1 parent 9d95b02 commit 90ed36a
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import org.opensearch.arrow.StreamProducer;
import org.opensearch.arrow.StreamTicket;

import java.io.IOException;

/**
* ProxyStreamProvider acts as forward proxy for FlightStream.
* It creates a BatchedJob to handle the streaming of data from the remote FlightStream.
Expand All @@ -38,6 +40,15 @@ public BatchedJob createJob(BufferAllocator allocator) {
return new ProxyBatchedJob(remoteStream);
}

@Override
public void close() {
try {
remoteStream.close();
} catch (Exception e) {
throw new RuntimeException(e);
}
}

public static class ProxyBatchedJob implements BatchedJob {

private final FlightStream remoteStream;
Expand Down Expand Up @@ -66,5 +77,12 @@ public void onCancel() {
throw new RuntimeException(e);
}
}

@Override
public boolean isCancelled() {
// Proxy stream don't have any business logic to set this flag,
// they piggyback on remote stream getting cancelled.
return false;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
import org.opensearch.tasks.TaskAwareRequest;
import org.opensearch.tasks.TaskManager;

import java.io.IOException;

public class StreamManagerWrapper implements StreamManager {

private final StreamManager streamManager;
Expand Down Expand Up @@ -96,6 +98,11 @@ public int estimatedRowCount() {
return streamProducer.estimatedRowCount();
}

@Override
public void close() throws IOException {
streamProducer.close();
}

static class BatchedJobTaskWrapper implements BatchedJob, TaskAwareRequest {
private final BatchedJob batchedJob;
private final TaskManager taskManager;
Expand Down Expand Up @@ -140,6 +147,11 @@ public void onCancel() {
batchedJob.onCancel();
}

@Override
public boolean isCancelled() {
return batchedJob.isCancelled();
}

@Override
public void setParentTask(TaskId taskId) {
this.parentTaskId = taskId;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import org.apache.logging.log4j.Logger;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.Query;
import org.opensearch.OpenSearchException;
import org.opensearch.arrow.StreamManager;
import org.opensearch.arrow.StreamProducer;
import org.opensearch.arrow.StreamTicket;
Expand Down Expand Up @@ -112,7 +113,14 @@ private boolean searchWithCollector(
if (streamManager == null) {
throw new RuntimeException("StreamManager not setup");
}
final boolean[] isCancelled = {false};
StreamTicket ticket = streamManager.registerStream(new StreamProducer() {

@Override
public void close() {
isCancelled[0] = true;
}

@Override
public BatchedJob createJob(BufferAllocator allocator) {
return new BatchedJob() {
Expand All @@ -123,6 +131,9 @@ public void run(VectorSchemaRoot root, StreamProducer.FlushSignal flushSignal) {
Collector collector = QueryCollectorContext.createQueryCollector(collectors);
final ArrowDocIdCollector arrowDocIdCollector = new ArrowDocIdCollector(collector, root, flushSignal, 1000);
try {
searcher.addQueryCancellation(() -> {if (isCancelled[0] == true) {
throw new OpenSearchException("Stream for query results cancelled.");
}});
searcher.search(query, arrowDocIdCollector);
} catch (EarlyTerminatingCollector.EarlyTerminationException e) {
// EarlyTerminationException is not caught in ContextIndexSearcher to allow force termination of
Expand Down Expand Up @@ -153,7 +164,12 @@ public void run(VectorSchemaRoot root, StreamProducer.FlushSignal flushSignal) {

@Override
public void onCancel() {
isCancelled[0] = true;
}

@Override
public boolean isCancelled() {
return searchContext.isCancelled() || isCancelled();
}
};
}
Expand Down

0 comments on commit 90ed36a

Please sign in to comment.