Re: [PR] Enable query cancellation for MSQE + cancel using client-provided id [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14823:
URL: https://github.com/apache/pinot/pull/14823#discussion_r1946664556


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java:
##########
@@ -223,4 +245,65 @@ protected static void augmentStatistics(RequestContext statistics, BrokerRespons
     statistics.setExplainPlanNumMatchAllFilterSegments(response.getExplainPlanNumMatchAllFilterSegments());
     statistics.setTraceInfo(response.getTraceInfo());
   }
+
+  @Override
+  public Map<Long, String> getRunningQueries() {
+    Preconditions.checkState(isQueryCancellationEnabled(), "Query cancellation is not enabled on broker");
+    return new HashMap<>(_queriesById);
+  }
+
+  @Override
+  public boolean cancelQuery(long queryId, int timeoutMs, Executor executor, HttpClientConnectionManager connMgr,
+      Map<String, Integer> serverResponses)
+      throws Exception {
+    Preconditions.checkState(isQueryCancellationEnabled(), "Query cancellation is not enabled on broker");
+    try {
+      return handleCancel(queryId, timeoutMs, executor, connMgr, serverResponses);
+    } finally {
+      maybeRemoveQuery(queryId);
+    }
+  }
+
+  @Override
+  public boolean cancelQueryByClientId(String clientQueryId, int timeoutMs, Executor executor,
+      HttpClientConnectionManager connMgr, Map<String, Integer> serverResponses)
+      throws Exception {
+    Preconditions.checkState(isQueryCancellationEnabled(), "Query cancellation is not enabled on broker");
+    Optional<Long> requestId = _clientQueryIds.entrySet().stream()
+        .filter(e -> clientQueryId.equals(e.getValue())).map(Map.Entry::getKey).findFirst();
+    if (requestId.isPresent()) {
+      return cancelQuery(requestId.get(), timeoutMs, executor, connMgr, serverResponses);
+    } else {
+      LOGGER.warn("Query cancellation cannot be performed due to unknown client query id: {}", clientQueryId);
+      return false;

Review Comment:
   PinotClientRequest is pretty inconsistent in the way it handles the errors, for sure



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org