jasperjiaguo commented on code in PR #9171:
URL: https://github.com/apache/pinot/pull/9171#discussion_r942950227
##########
pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java:
##########
@@ -103,7 +115,68 @@ public QueryScheduler(PinotConfiguration config,
QueryExecutor queryExecutor, Re
* @return Listenable future for query result representing serialized
response. It is possible that the
* future may return immediately or be scheduled for execution at a later
time.
*/
- public abstract ListenableFuture<byte[]> submit(ServerQueryRequest
queryRequest);
+ public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ ListenableFuture<byte[]> future = submitInternal(queryRequest);
+ if (_enableQueryCancellation) {
+ String queryId = queryRequest.getQueryId();
+ // Track the running query for cancellation.
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Keep track of running query: {}", queryId);
+ }
+ _queryFuturesById.put(queryId, future);
+ // And remove the track when the query ends.
+ Futures.addCallback(future, new FutureCallback<byte[]>() {
+ @Override
+ public void onSuccess(@Nullable byte[] ignored) {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Remove track of running query: {} on success",
queryId);
+ }
+ _queryFuturesById.remove(queryId);
+ }
+
+ @Override
+ public void onFailure(Throwable ignored) {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Remove track of running query: {} on failure",
queryId);
+ }
+ _queryFuturesById.remove(queryId);
+ }
+ }, MoreExecutors.directExecutor());
+ }
+ return future;
+ }
+
+ protected abstract ListenableFuture<byte[]>
submitInternal(ServerQueryRequest queryRequest);
+
+ /**
+ * Cancel a query as identified by the queryId. This method is non-blocking
and the query may still run for a while
+ * after calling this method. This method can be called multiple times.
+ *
+ * @param queryId a unique Id to find the query
+ * @return true if a running query exists for the given queryId.
+ */
+ public boolean cancelQuery(String queryId) {
+ // Keep the future as it'll be cleaned up by the thread executing the
query.
+ Future<byte[]> future = _queryFuturesById.get(queryId);
+ if (future == null) {
+ return false;
+ }
+ boolean done = future.isDone();
+ if (!done) {
+ future.cancel(true);
Review Comment:
Hey, we plan to use the same cancellation mechanism for cpu time/memory
usage based query pre-emption. Do you think the current frequency of checking
on `if (Thread.interrupted())` at each `nextBlock()` is insufficient to
preempt a query in timely manner @npawar ?
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
