From b1fdd6e34949706457cdddea983df3665323cac7 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 16 Sep 2026 09:41:31 +0800 Subject: [PATCH] Datanode: only serve a query's result to the session that submitted it fetchResultsV2, fetchResults and closeOperation resolve the queryId sent by the client in the coordinator wide map of running queries, without checking which session submitted that query. Add a per session validation of the queryId and return NO_PERMISSION when it was not issued to the calling session, so that the running query is neither read nor released by another session. - add IClientSession.containsQueryId, implemented on ClientSession and InternalClientSession over the statementId -> queryId bookkeeping that already exists, and on MqttClientSession/RestClientSession which cannot submit queries - check the queryId in fetchResultsV2, fetchResults and closeOperation; a queryId that is no longer running keeps the previous behaviour - a rejected fetch does not record latency or clean up the query - add the en/zh message and a unit test covering the two fetch APIs, the close path and the session level bookkeeping --- .../iotdb/db/i18n/DataNodeMiscMessages.java | 2 + .../iotdb/db/i18n/DataNodeMiscMessages.java | 2 + .../db/protocol/session/ClientSession.java | 19 ++ .../db/protocol/session/IClientSession.java | 3 + .../session/InternalClientSession.java | 5 + .../protocol/session/MqttClientSession.java | 5 + .../protocol/session/RestClientSession.java | 5 + .../thrift/impl/ClientRPCServiceImpl.java | 92 +++++-- .../protocol/session/QueryOwnershipTest.java | 231 ++++++++++++++++++ 9 files changed, 338 insertions(+), 26 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index e76211672ff71..4cf4a787d6ed0 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/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 INVALID_PIPE_NAME = "Invalid pipeName"; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index c9a3f49710f80..68c556c59c7e0 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/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 INVALID_PIPE_NAME = "无效的 pipeName"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java index bad4ddd6c7dca..c061099ab81cc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java @@ -90,6 +90,25 @@ 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> statementIdToQueryId, Long statementId, long queryId) { + if (statementId == null) { + for (Set queryIds : statementIdToQueryId.values()) { + if (queryIds != null && queryIds.contains(queryId)) { + return true; + } + } + return false; + } + Set queryIds = statementIdToQueryId.get(statementId); + return queryIds != null && queryIds.contains(queryId); + } + @Override public void removeQueryId(Long statementId, Long queryId) { removeQueryId(statementIdToQueryId, statementId, queryId); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java index bac4b15e342dd..5114ecb0d78a9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java +++ b/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); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java index 460ec9319f2dc..8bbb248e9897f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java +++ b/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); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java index 65e2c9a5b5d49..0fb830bcca3bb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java +++ b/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(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java index 58bae05fffe32..7fe2ceece08a6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java +++ b/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(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index 790e022ded36d..f0b760efd17a7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -1561,6 +1561,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 { @@ -1570,13 +1571,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); @@ -1605,19 +1615,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(); @@ -1711,8 +1723,22 @@ public TSStatus cancelOperation(TSCancelOperationReq req) { @Override public TSStatus closeOperation(TSCloseOperationReq req) { + IClientSession clientSession = SESSION_MANAGER.getCurrSession(); + if (req.isSetQueryId() + && clientSession != null + && clientSession.isLogin() + && COORDINATOR.getQueryExecution(req.queryId) != null + && !clientSession.containsQueryId( + req.isSetStatementId() ? req.getStatementId() : null, req.queryId)) { + // The queryId indexes the process-wide map of running queries, so only the session that + // submitted the query may release it. Queries that are no longer running keep the previous + // behaviour: releasing an unknown queryId stays a no-op. + return RpcUtils.getStatus( + TSStatusCode.NO_PERMISSION, + DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237); + } return SESSION_MANAGER.closeOperation( - SESSION_MANAGER.getCurrSession(), + clientSession, req.queryId, req.statementId, req.isSetStatementId(), @@ -2319,6 +2345,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 { @@ -2333,6 +2360,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(); @@ -2360,19 +2398,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(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java new file mode 100644 index 0000000000000..593c03f244188 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java @@ -0,0 +1,231 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.protocol.session; + +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.protocol.thrift.impl.ClientRPCServiceImpl; +import org.apache.iotdb.db.queryengine.plan.Coordinator; +import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.TSCloseOperationReq; +import org.apache.iotdb.service.rpc.thrift.TSFetchResultsReq; + +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; +import org.mockito.Mockito; + +import java.lang.reflect.Field; +import java.net.Socket; +import java.util.Map; + +public class QueryOwnershipTest { + + private static final long STATEMENT_ID = 1L; + private static final long QUERY_ID = 2L; + + private static int previousDataNodeId; + + @BeforeClass + public static void setUp() { + // the coordinator builds its query id generator from the data node id of this node + previousDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + } + + @AfterClass + public static void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId); + } + + @Test + public void testQueryIdsAreBoundToTheSessionThatSubmittedTheQuery() { + ClientSession owner = createSession("user"); + owner.addStatementId(STATEMENT_ID); + owner.addQueryId(STATEMENT_ID, QUERY_ID); + + ClientSession anotherSession = createSession("user"); + anotherSession.addStatementId(STATEMENT_ID); + + Assert.assertTrue(owner.containsQueryId(STATEMENT_ID, QUERY_ID)); + // clients that do not send a statement id together with the query id are still served + Assert.assertTrue(owner.containsQueryId(null, QUERY_ID)); + Assert.assertFalse(anotherSession.containsQueryId(STATEMENT_ID, QUERY_ID)); + Assert.assertFalse(anotherSession.containsQueryId(null, QUERY_ID)); + Assert.assertFalse(owner.containsQueryId(STATEMENT_ID + 1, QUERY_ID)); + } + + @Test + public void testFetchResultsRejectsQueryOfAnotherSession() throws Exception { + ClientSession anotherSession = createSession("user"); + anotherSession.addStatementId(STATEMENT_ID); + anotherSession.setLogin(true); + + Map queryExecutionMap = getQueryExecutionMap(); + queryExecutionMap.put(QUERY_ID, mockQueryExecution()); + try { + withCurrentSession( + anotherSession, + () -> { + ClientRPCServiceImpl service = new ClientRPCServiceImpl(); + TSFetchResultsReq request = createFetchResultsReq(anotherSession); + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + service.fetchResults(request).getStatus().getCode()); + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + service.fetchResultsV2(request).getStatus().getCode()); + }); + // the rejected requests must not release the query of the session that submitted it + Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + + @Test + public void testFetchResultsOfOwnQueryIsStillServed() throws Exception { + ClientSession owner = createSession("user"); + owner.addStatementId(STATEMENT_ID); + owner.addQueryId(STATEMENT_ID, QUERY_ID); + owner.setLogin(true); + + IQueryExecution queryExecution = mockQueryExecution(); + Map queryExecutionMap = getQueryExecutionMap(); + queryExecutionMap.put(QUERY_ID, queryExecution); + try { + withCurrentSession( + owner, + () -> + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + new ClientRPCServiceImpl() + .fetchResultsV2(createFetchResultsReq(owner)) + .getStatus() + .getCode())); + // a fully consumed query is released, and the client closes it afterwards + Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID)); + Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + + @Test + public void testCloseOperationRejectsQueryOfAnotherSession() throws Exception { + ClientSession anotherSession = createSession("user"); + anotherSession.addStatementId(STATEMENT_ID); + anotherSession.setLogin(true); + + Map queryExecutionMap = getQueryExecutionMap(); + queryExecutionMap.put(QUERY_ID, mockQueryExecution()); + try { + withCurrentSession( + anotherSession, + () -> + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + new ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode())); + // the rejected request must not release the query of the session that submitted it + Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + + @Test + public void testCloseOperationReleasesQueryOfOwnSession() throws Exception { + ClientSession owner = createSession("user"); + owner.addStatementId(STATEMENT_ID); + owner.addQueryId(STATEMENT_ID, QUERY_ID); + owner.setLogin(true); + + Map queryExecutionMap = getQueryExecutionMap(); + queryExecutionMap.put(QUERY_ID, mockQueryExecution()); + try { + withCurrentSession( + owner, + () -> + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + new ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode())); + Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID)); + Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + + @Test + public void testCloseOperationOfQueryThatIsNoLongerRunningStaysANoOp() { + // a client that consumed a result set completely sends closeOperation after the query has + // already been released, and that request has to succeed as it always did + ClientSession session = createSession("user"); + session.addStatementId(STATEMENT_ID); + session.setLogin(true); + + withCurrentSession( + session, + () -> + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + new ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode())); + } + + private IQueryExecution mockQueryExecution() { + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Mockito.when(queryExecution.getQueryId()).thenReturn("query"); + return queryExecution; + } + + private TSFetchResultsReq createFetchResultsReq(ClientSession session) { + return new TSFetchResultsReq(session.getId(), "select 1", 1024, QUERY_ID, true) + .setStatementId(STATEMENT_ID); + } + + private TSCloseOperationReq createCloseOperationReq() { + return new TSCloseOperationReq().setStatementId(STATEMENT_ID).setQueryId(QUERY_ID); + } + + private void withCurrentSession(ClientSession session, Runnable body) { + SessionManager sessionManager = SessionManager.getInstance(); + IClientSession previousSession = sessionManager.getCurrSession(); + sessionManager.setCurrSession(session); + try { + body.run(); + } finally { + sessionManager.restoreSession(previousSession, session); + } + } + + private ClientSession createSession(String username) { + ClientSession session = new ClientSession(Mockito.mock(Socket.class)); + session.setUsername(username); + return session; + } + + @SuppressWarnings("unchecked") + private Map getQueryExecutionMap() throws Exception { + Field field = Coordinator.class.getDeclaredField("queryExecutionMap"); + field.setAccessible(true); + return (Map) field.get(Coordinator.getInstance()); + } +}