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()); + } +}