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..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 @@ -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> 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(boxedQueryId)) { + return true; + } + } + return false; + } + Set queryIds = statementIdToQueryId.get(statementId); + return queryIds != null && queryIds.contains(boxedQueryId); + } + @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..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 @@ -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(); 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..b8c53352c09a5 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java @@ -0,0 +1,343 @@ +/* + * 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(); + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + service.fetchResults(createFetchResultsReq(anotherSession)).getStatus().getCode()); + Assert.assertEquals( + TSStatusCode.NO_PERMISSION.getStatusCode(), + 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)); + } 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(createFetchResultsReqWithStatementId(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 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"); + 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(createFetchResultsReqWithStatementId(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 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"); + 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; + } + + /** 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); + } + + /** The V2 fetch request, which always carries the statement id. */ + private TSFetchResultsReq createFetchResultsReqWithStatementId(ClientSession session) { + return createFetchResultsReq(session).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()); + } +}