From 515385484bacfae6209573e715f1f2736750bf2d Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 16 Sep 2026 09:41:31 +0800 Subject: [PATCH 1/2] 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, the session level bookkeeping, and a fetch that leaves the result set unconsumed --- .../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 | 272 ++++++++++++++++++ 9 files changed, 379 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 c961031b38e02..405a5139e3466 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 MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA = "Missing Load TsFile slice metadata: %s"; 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 11bc2fc579f45..a87a49eaaf5f4 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 MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA = "缺少 Load TsFile 分片元数据：%s"; 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 ce0c5d6745004..660db73522f02 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 @@ -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,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(), @@ -2291,6 +2317,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 +2332,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 +2370,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..f811d55e51091 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java @@ -0,0 +1,272 @@ +/* + * 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.apache.iotdb.service.rpc.thrift.TSFetchResultsResp; + +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; +import java.util.Optional; + +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 testFetchResultsOfQueryWithMoreDataKeepsQueryAndSessionBinding() 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(true)); + try { + withCurrentSession( + owner, + () -> { + TSFetchResultsResp response = + new ClientRPCServiceImpl().fetchResultsV2(createFetchResultsReq(owner)); + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode()); + Assert.assertTrue(response.isMoreData()); + }); + // the result set is not consumed yet, so the query has to stay fetchable by its owner + Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID)); + Assert.assertTrue(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() throws Exception { + return mockQueryExecution(false); + } + + /** + * A mocked execution. {@code hasNextResult} says whether the mocked query still holds data after + * the batch that is about to be read, which decides if a single fetch consumes the whole result + * set and therefore releases the query. + */ + private IQueryExecution mockQueryExecution(boolean hasNextResult) throws Exception { + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Mockito.when(queryExecution.getQueryId()).thenReturn("query"); + // the mock carries no buffered batch, only the "is there anything left" flag + Mockito.when(queryExecution.getByteBufferBatchResult()).thenReturn(Optional.empty()); + Mockito.when(queryExecution.hasNextResult()).thenReturn(hasNextResult); + 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()); + } +} From b57e4e70e50dbfce336fc688311b7ffd5cb6b87e Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 23 Sep 2026 09:58:18 +0800 Subject: [PATCH 2/2] Datanode: never release a query id the session does not own closeOperation consulted the coordinator map before the session binding and fell through to the global cleanup when the lookup returned null. Query ids are allocated before their execution is published, so a query of another session can be registered between the lookup and the cleanup, and the foreign request released it. Check the session binding first and return the compatible no-op result directly when the query id is not running. Also box the primitive query id once in ClientSession#containsQueryId instead of once per visited statement set, and align the tests with the V1 and V2 fetch contracts: the V1 request carries no statement id, and a query registered while a foreign close is served must survive. --- .../db/protocol/session/ClientSession.java | 8 +- .../thrift/impl/ClientRPCServiceImpl.java | 21 +++-- .../protocol/session/QueryOwnershipTest.java | 85 +++++++++++++++++-- 3 files changed, 97 insertions(+), 17 deletions(-) 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 c061099ab81cc..2f08920b8d248 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 @@ -97,16 +97,20 @@ public boolean containsQueryId(Long statementId, long queryId) { public static boolean containsQueryId( Map> 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 queryIds : statementIdToQueryId.values()) { - if (queryIds != null && queryIds.contains(queryId)) { + if (queryIds != null && queryIds.contains(boxedQueryId)) { return true; } } return false; } Set queryIds = statementIdToQueryId.get(statementId); - return queryIds != null && queryIds.contains(queryId); + return queryIds != null && queryIds.contains(boxedQueryId); } @Override 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 660db73522f02..692b5bedd2d66 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 @@ -1697,17 +1697,22 @@ public TSStatus cancelOperation(TSCancelOperationReq req) { public TSStatus closeOperation(TSCloseOperationReq req) { IClientSession clientSession = SESSION_MANAGER.getCurrSession(); if (req.isSetQueryId() + && req.isSetStatementId() && clientSession != null && clientSession.isLogin() - && COORDINATOR.getQueryExecution(req.queryId) != null - && !clientSession.containsQueryId( - req.isSetStatementId() ? req.getStatementId() : null, req.queryId)) { + && !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. 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); + // 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( clientSession, 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 index f811d55e51091..b8c53352c09a5 100644 --- 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 @@ -88,13 +88,15 @@ public void testFetchResultsRejectsQueryOfAnotherSession() throws Exception { anotherSession, () -> { ClientRPCServiceImpl service = new ClientRPCServiceImpl(); - TSFetchResultsReq request = createFetchResultsReq(anotherSession); Assert.assertEquals( TSStatusCode.NO_PERMISSION.getStatusCode(), - service.fetchResults(request).getStatus().getCode()); + service.fetchResults(createFetchResultsReq(anotherSession)).getStatus().getCode()); Assert.assertEquals( TSStatusCode.NO_PERMISSION.getStatusCode(), - service.fetchResultsV2(request).getStatus().getCode()); + service + .fetchResultsV2(createFetchResultsReqWithStatementId(anotherSession)) + .getStatus() + .getCode()); }); // the rejected requests must not release the query of the session that submitted it Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID)); @@ -120,7 +122,7 @@ public void testFetchResultsOfOwnQueryIsStillServed() throws Exception { Assert.assertEquals( TSStatusCode.SUCCESS_STATUS.getStatusCode(), new ClientRPCServiceImpl() - .fetchResultsV2(createFetchResultsReq(owner)) + .fetchResultsV2(createFetchResultsReqWithStatementId(owner)) .getStatus() .getCode())); // a fully consumed query is released, and the client closes it afterwards @@ -131,6 +133,34 @@ public void testFetchResultsOfOwnQueryIsStillServed() throws Exception { } } + @Test + public void testFetchResultsOfOwnQueryWithoutStatementIdIsStillServed() 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 { + // the query id is bound to a statement of this session, the request does not mention it + withCurrentSession( + owner, + () -> + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + new ClientRPCServiceImpl() + .fetchResults(createFetchResultsReq(owner)) + .getStatus() + .getCode())); + // a fully consumed query is released for the session that submitted it + Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID)); + Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + @Test public void testFetchResultsOfQueryWithMoreDataKeepsQueryAndSessionBinding() throws Exception { ClientSession owner = createSession("user"); @@ -145,7 +175,8 @@ public void testFetchResultsOfQueryWithMoreDataKeepsQueryAndSessionBinding() thr owner, () -> { TSFetchResultsResp response = - new ClientRPCServiceImpl().fetchResultsV2(createFetchResultsReq(owner)); + new ClientRPCServiceImpl() + .fetchResultsV2(createFetchResultsReqWithStatementId(owner)); Assert.assertEquals( TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode()); Assert.assertTrue(response.isMoreData()); @@ -180,6 +211,41 @@ public void testCloseOperationRejectsQueryOfAnotherSession() throws Exception { } } + @Test + public void testCloseOperationDoesNotReleaseQueryRegisteredWhileTheRequestIsServed() + throws Exception { + Map queryExecutionMap = getQueryExecutionMap(); + IQueryExecution queryOfAnotherSession = mockQueryExecution(); + // Query ids are allocated before their execution is published, so the session that owns this + // queryId can register its execution between the ownership check of this request and the point + // where the request would release the query. Simulate that publication from inside the + // ownership check, which is where the request decides who may release the queryId. + ClientSession anotherSession = + new ClientSession(Mockito.mock(Socket.class)) { + @Override + public boolean containsQueryId(Long statementId, long queryId) { + queryExecutionMap.putIfAbsent(queryId, queryOfAnotherSession); + return super.containsQueryId(statementId, queryId); + } + }; + anotherSession.setUsername("user"); + anotherSession.addStatementId(STATEMENT_ID); + anotherSession.setLogin(true); + + try { + withCurrentSession( + anotherSession, + () -> + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + new ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode())); + // the request must never release the query that was just registered by its owner + Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID)); + } finally { + queryExecutionMap.remove(QUERY_ID); + } + } + @Test public void testCloseOperationReleasesQueryOfOwnSession() throws Exception { ClientSession owner = createSession("user"); @@ -237,9 +303,14 @@ private IQueryExecution mockQueryExecution(boolean hasNextResult) throws Excepti return queryExecution; } + /** The V1 fetch request of the legacy JDBC data set, which does not send a statement id. */ private TSFetchResultsReq createFetchResultsReq(ClientSession session) { - return new TSFetchResultsReq(session.getId(), "select 1", 1024, QUERY_ID, true) - .setStatementId(STATEMENT_ID); + return new TSFetchResultsReq(session.getId(), "select 1", 1024, QUERY_ID, true); + } + + /** The V2 fetch request, which always carries the statement id. */ + private TSFetchResultsReq createFetchResultsReqWithStatementId(ClientSession session) { + return createFetchResultsReq(session).setStatementId(STATEMENT_ID); } private TSCloseOperationReq createCloseOperationReq() {