Datanode: validate that a query id belongs to the calling session

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java

@@ -21,6 +21,8 @@
 
 /** Compile-time i18n constants for DataNode misc subsystems (English). */
 public final class DataNodeMiscMessages {
+  public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
+      "The requested query does not belong to the current session.";
 
   public static final String MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
       "Missing Load TsFile slice metadata: %s";

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java

@@ -21,6 +21,8 @@
 
 /** 编译时国际化常量 - DataNode 杂项子系统（中文）。 */
 public final class DataNodeMiscMessages {
+  public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
+      "请求的查询不属于当前会话。";
 
   public static final String MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
       "缺少 Load TsFile 分片元数据：%s";

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java

@@ -90,6 +90,29 @@ public void addQueryId(Long statementId, long queryId) {
     queryIds.add(queryId);
   }
 
+  @Override
+  public boolean containsQueryId(Long statementId, long queryId) {
+    return containsQueryId(statementIdToQueryId, statementId, queryId);
+  }
+
+  public static boolean containsQueryId(
+      Map<Long, Set<Long>> statementIdToQueryId, Long statementId, long queryId) {
+    // Set#contains takes an Object, so box the primitive queryId once: a client that does not send
+    // a statement id makes this method visit every statement set of the session, and a box per
+    // visited set would allocate once per statement on every fetched page.
+    Long boxedQueryId = queryId;
+    if (statementId == null) {
+      for (Set<Long> queryIds : statementIdToQueryId.values()) {
+        if (queryIds != null && queryIds.contains(boxedQueryId)) {
+          return true;
+        }
+      }
+      return false;
+    }
+    Set<Long> queryIds = statementIdToQueryId.get(statementId);
+    return queryIds != null && queryIds.contains(boxedQueryId);
+  }
+
   @Override
   public void removeQueryId(Long statementId, Long queryId) {
     removeQueryId(statementIdToQueryId, statementId, queryId);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java

@@ -164,6 +164,9 @@ public ConnectionInfo convertToConnectionInfo() {
 
   public abstract void addQueryId(Long statementId, long queryId);
 
+  // statementId could be null
+  public abstract boolean containsQueryId(Long statementId, long queryId);
+
   // statementId could be null
   public abstract void removeQueryId(Long statementId, Long queryId);
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java

@@ -90,6 +90,11 @@ public void addQueryId(Long statementId, long queryId) {
     queryIds.add(queryId);
   }
 
+  @Override
+  public boolean containsQueryId(Long statementId, long queryId) {
+    return ClientSession.containsQueryId(statementIdToQueryId, statementId, queryId);
+  }
+
   @Override
   public void removeQueryId(Long statementId, Long queryId) {
     ClientSession.removeQueryId(statementIdToQueryId, statementId, queryId);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java

@@ -77,6 +77,11 @@ public void addQueryId(Long statementId, long queryId) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public boolean containsQueryId(Long statementId, long queryId) {
+    return false;
+  }
+
   @Override
   public void removeQueryId(Long statementId, Long queryId) {
     throw new UnsupportedOperationException();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java

@@ -79,6 +79,11 @@ public void addQueryId(Long statementId, long queryId) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public boolean containsQueryId(Long statementId, long queryId) {
+    return false;
+  }
+
   @Override
   public void removeQueryId(Long statementId, Long queryId) {
     throw new UnsupportedOperationException();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java

@@ -1533,6 +1533,7 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
     String statementType = null;
     Throwable t = null;
     IQueryExecution queryExecution = null;
+    boolean queryOwnedBySession = false;
     IClientSession clientSession = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
     Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
     try {
@@ -1542,13 +1543,22 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
       }
 
       queryExecution = COORDINATOR.getQueryExecution(req.queryId);
-
       if (queryExecution == null) {
         TSStatus noQueryExecutionStatus = new TSStatus(QUERY_WAS_KILLED.getStatusCode());
         noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
         return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
       }
 
+      if (!clientSession.containsQueryId(statementId, req.queryId)) {
+        // The query is still running, but it was submitted by another session: do not stream its
+        // result and do not release it, so that the query which owns it is left untouched.
+        return RpcUtils.getTSFetchResultsResp(
+            RpcUtils.getStatus(
+                TSStatusCode.NO_PERMISSION,
+                DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
+      }
+      queryOwnedBySession = true;
+
       TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
 
       queryExecution.updateCurrentRpcStartTime(startTime);
@@ -1577,19 +1587,21 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
       throw error;
     } finally {
 
-      long currentOperationCost = System.nanoTime() - startTime;
-      COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
-
-      // record each operation time cost
-      CommonUtils.addStatementExecutionLatency(
-          OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+      if (queryOwnedBySession) {
+        long currentOperationCost = System.nanoTime() - startTime;
+        COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
 
-      if (finished) {
-        // record total time cost for one query
-        long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
-        CommonUtils.addQueryLatency(
-            StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
-        clearUp(clientSession, statementId, req.queryId, req, t);
+        // record each operation time cost
+        CommonUtils.addStatementExecutionLatency(
+            OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+
+        if (finished) {
+          // record total time cost for one query
+          long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
+          CommonUtils.addQueryLatency(
+              StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
+          clearUp(clientSession, statementId, req.queryId, req, t);
+        }
       }
 
       SESSION_MANAGER.updateIdleTime();
@@ -1683,8 +1695,27 @@ public TSStatus cancelOperation(TSCancelOperationReq req) {
 
   @Override
   public TSStatus closeOperation(TSCloseOperationReq req) {
+    IClientSession clientSession = SESSION_MANAGER.getCurrSession();
+    if (req.isSetQueryId()
+        && req.isSetStatementId()
+        && clientSession != null
+        && clientSession.isLogin()
+        && !clientSession.containsQueryId(req.getStatementId(), req.queryId)) {
+      // The queryId indexes the process-wide map of running queries, so only the session that
+      // submitted the query may release it.
+      if (COORDINATOR.getQueryExecution(req.queryId) != null) {
+        return RpcUtils.getStatus(
+            TSStatusCode.NO_PERMISSION,
+            DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237);
+      }
+      // A queryId that is no longer running keeps the previous behaviour: releasing it stays a
+      // no-op. It must not fall through to the global cleanup below: query ids are allocated
+      // before their execution is published, so the session that owns this queryId can register
+      // it between the lookup above and the cleanup.
+      return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
+    }
     return SESSION_MANAGER.closeOperation(
-        SESSION_MANAGER.getCurrSession(),
+        clientSession,
         req.queryId,
         req.statementId,
         req.isSetStatementId(),
@@ -2291,6 +2322,7 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
     String statementType = null;
     Throwable t = null;
     IQueryExecution queryExecution = null;
+    boolean queryOwnedBySession = false;
     IClientSession clientSession = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
     Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
     try {
@@ -2305,6 +2337,17 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
         noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
         return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
       }
+
+      if (!clientSession.containsQueryId(statementId, req.queryId)) {
+        // The query is still running, but it was submitted by another session: do not stream its
+        // result and do not release it, so that the query which owns it is left untouched.
+        return RpcUtils.getTSFetchResultsResp(
+            RpcUtils.getStatus(
+                TSStatusCode.NO_PERMISSION,
+                DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
+      }
+      queryOwnedBySession = true;
+
       queryExecution.updateCurrentRpcStartTime(startTime);
       statementType = queryExecution.getStatementType();
 
@@ -2332,19 +2375,21 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
       throw error;
     } finally {
 
-      long currentOperationCost = System.nanoTime() - startTime;
-      COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
+      if (queryOwnedBySession) {
+        long currentOperationCost = System.nanoTime() - startTime;
+        COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
 
-      // record each operation time cost
-      CommonUtils.addStatementExecutionLatency(
-          OperationType.FETCH_RESULTS, statementType, currentOperationCost);
-
-      if (finished) {
-        // record total time cost for one query
-        long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
-        CommonUtils.addQueryLatency(
-            StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
-        clearUp(clientSession, statementId, req.queryId, req, t);
+        // record each operation time cost
+        CommonUtils.addStatementExecutionLatency(
+            OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+
+        if (finished) {
+          // record total time cost for one query
+          long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
+          CommonUtils.addQueryLatency(
+              StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
+          clearUp(clientSession, statementId, req.queryId, req, t);
+        }
       }
 
       SESSION_MANAGER.updateIdleTime();