From b923ba2e1de410f51159917c966bfa970de3404b Mon Sep 17 00:00:00 2001 From: Yann Massard Date: Mon, 2 Nov 2026 17:58:27 +0000 Subject: [PATCH] Fix SessionPool's bookkeeping of sessions in use SessionPool only tracked sessions held by an open query result set (the "occupied" map). A session executing a request (an insert, a non-query statement, or a query whose result set has not been returned yet) was in neither that map nor the idle queue. Several places treated the map as if it held every session taken from the pool: - currentOccupiedSize() ignored sessions executing a request, so a saturated write pool reported 0 occupied and 0 available. - setFetchSize, setTimeZone, setEnableRedirection, setEnableQueryRedirection, setVersion and setQueryTimeout did not reach sessions executing a request, which kept the old value after being put back. - close() did not close sessions executing a request; they were later put back into the cleared queue and their connections leaked. "occupied" is now a set of all sessions handed out by getSession() until they are put back, cleaned after a connection error or released by closeResultSet(), and the places above use it. The result-set map is renamed to occupiedByResultSet and maps each session to the result set holding it. closeResultSet() releases a session only if it still belongs to that result set. Before, the second release of a result set (hasNext() auto-releases an exhausted one, the caller's finally releases it again) could put back a session that another thread had meanwhile taken for a new query. It also ignores a wrapper without a session instead of throwing an NPE, and closes the session when closing the operation handle fails instead of leaking its connection. --- .../iotdb/session/pool/SessionPool.java | 81 ++++--- .../iotdb/session/pool/SessionPoolTest.java | 219 ++++++++++++++++++ 2 files changed, 269 insertions(+), 31 deletions(-) diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 16b79c6415c4b..7465ea859c562 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -56,6 +56,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.ConcurrentMap; @@ -100,8 +101,14 @@ public class SessionPool implements ISessionPool { private static final int FINAL_RETRY = RETRY - 1; private final ConcurrentLinkedDeque queue = new ConcurrentLinkedDeque<>(); - // for session whose resultSet is not released. - private final ConcurrentMap occupied = new ConcurrentHashMap<>(); + // Every session taken by getSession() and not yet returned (put back or cleaned), whatever it + // is used for. + private final Set occupied = ConcurrentHashMap.newKeySet(); + // The occupied sessions whose result set is not released yet, each with that result set. A + // session can be released, taken again and hold a new result set while the old one is still + // referenced, so a release must match the result set, not only the session. + private final ConcurrentMap occupiedByResultSet = + new ConcurrentHashMap<>(); private int size = 0; private int maxSize = 0; private final long waitToGetSessionTimeoutInMs; @@ -702,6 +709,7 @@ private ISession getSession() throws IoTDBConnectionException { throw new IoTDBConnectionException(SESSION_POOL_IS_CLOSED); } if (session != null) { + occupied.add(session); return session; } @@ -790,6 +798,7 @@ private ISession getSession() throws IoTDBConnectionException { } } + occupied.add(session); return session; } @@ -802,6 +811,10 @@ public int currentAvailableSize() { return queue.size(); } + /** + * The number of sessions currently taken from the pool: those executing a request (insert, + * non-query, query, ...) and those held by a query result set that has not been closed yet. + */ @Override public int currentOccupiedSize() { return occupied.size(); @@ -809,6 +822,7 @@ public int currentOccupiedSize() { @SuppressWarnings({"squid:S2446"}) protected void putBack(ISession session) { + occupied.remove(session); queue.push(session); synchronized (this) { // we do not need to notifyAll as any waited thread can continue to work after waked up. @@ -821,8 +835,8 @@ protected void putBack(ISession session) { } } - private void occupy(ISession session) { - occupied.put(session, session); + private void occupy(ISession session, SessionDataSetWrapper wrapper) { + occupiedByResultSet.put(session, wrapper); } /** close all connections in the pool */ @@ -836,7 +850,8 @@ public synchronized void close() { LOGGER.warn(CLOSE_THE_SESSION_FAILED, e); } } - for (ISession session : occupied.keySet()) { + // sessions executing a request are closed as well; that request then fails + for (ISession session : occupied) { try { session.close(); } catch (IoTDBConnectionException e) { @@ -855,23 +870,26 @@ public synchronized void close() { LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); this.closed = true; queue.clear(); + occupiedByResultSet.clear(); occupied.clear(); } @Override public void closeResultSet(SessionDataSetWrapper wrapper) { - boolean putback = true; + ISession session = wrapper.getSession(); + // SessionDataSetWrapper.hasNext() already releases an exhausted result set, so this is often + // called twice for the same one. Only the first call may hand its session back; a later one + // must not release the session again, which may already hold another result set by then. + if (session == null || !occupiedByResultSet.remove(session, wrapper)) { + return; + } try { wrapper.getSessionDataSet().closeOperationHandle(); } catch (IoTDBConnectionException | StatementExecutionException e) { - tryConstructNewSession(); - putback = false; - } finally { - ISession session = occupied.remove(wrapper.getSession()); - if (putback && session != null) { - putBack(wrapper.getSession()); - } + cleanSessionAndMayThrowConnectionException(session); + return; } + putBack(session); } @SuppressWarnings({"squid:S2446"}) @@ -905,6 +923,7 @@ private void tryConstructNewSession() { private void closeSession(ISession session) { if (session != null) { + occupied.remove(session); try { session.close(); } catch (Exception e2) { @@ -3077,7 +3096,7 @@ public SessionDataSetWrapper executeQueryStatement(String sql) try { SessionDataSet resp = session.executeQueryStatement(sql); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3115,7 +3134,7 @@ public SessionDataSetWrapper executeQueryStatement(String sql, long timeoutInMs) try { SessionDataSet resp = session.executeQueryStatement(sql, timeoutInMs); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3179,7 +3198,7 @@ public SessionDataSetWrapper executeRawDataQuery( try { SessionDataSet resp = session.executeRawDataQuery(paths, startTime, endTime, timeOut); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3206,7 +3225,7 @@ public SessionDataSetWrapper executeLastDataQuery(List paths, long lastT try { SessionDataSet resp = session.executeLastDataQuery(paths, lastTime); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3233,7 +3252,7 @@ public SessionDataSetWrapper executeLastDataQuery(List paths, long lastT try { SessionDataSet resp = session.executeLastDataQuery(paths, lastTime, timeOut); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3260,7 +3279,7 @@ public SessionDataSetWrapper executeLastDataQuery(List paths) try { SessionDataSet resp = session.executeLastDataQuery(paths); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3287,7 +3306,7 @@ public SessionDataSetWrapper executeFastLastDataQueryForOnePrefixPath(final List try { SessionDataSet resp = session.executeFastLastDataQueryForOnePrefixPath(prefixes); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3316,7 +3335,7 @@ public SessionDataSetWrapper executeLastDataQueryForOneDevice( SessionDataSet resp = session.executeLastDataQueryForOneDevice(db, device, sensors, isLegalPathNodes); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3344,7 +3363,7 @@ public SessionDataSetWrapper executeAggregationQuery( try { SessionDataSet resp = session.executeAggregationQuery(paths, aggregations); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3373,7 +3392,7 @@ public SessionDataSetWrapper executeAggregationQuery( SessionDataSet resp = session.executeAggregationQuery(paths, aggregations, startTime, endTime); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3406,7 +3425,7 @@ public SessionDataSetWrapper executeAggregationQuery( SessionDataSet resp = session.executeAggregationQuery(paths, aggregations, startTime, endTime, interval); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3441,7 +3460,7 @@ public SessionDataSetWrapper executeAggregationQuery( session.executeAggregationQuery( paths, aggregations, startTime, endTime, interval, slidingStep); SessionDataSetWrapper wrapper = new SessionDataSetWrapper(resp, session, this); - occupy(session); + occupy(session, wrapper); return wrapper; } catch (IoTDBConnectionException e) { // TException means the connection is broken, remove it and get a new one. @@ -3491,7 +3510,7 @@ public void setFetchSize(int fetchSize) { for (ISession session : queue) { session.setFetchSize(fetchSize); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setFetchSize(fetchSize); } } @@ -3522,7 +3541,7 @@ public void setTimeZone(String zoneId) for (ISession session : queue) { session.setTimeZoneOfSession(zoneId); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setTimeZoneOfSession(zoneId); } } @@ -3552,7 +3571,7 @@ public void setEnableRedirection(boolean enableRedirection) { for (ISession session : queue) { session.setEnableRedirection(enableRedirection); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setEnableRedirection(enableRedirection); } } @@ -3568,7 +3587,7 @@ public void setEnableQueryRedirection(boolean enableQueryRedirection) { for (ISession session : queue) { session.setEnableQueryRedirection(enableQueryRedirection); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setEnableQueryRedirection(enableQueryRedirection); } } @@ -3630,7 +3649,7 @@ public void setVersion(Version version) { for (ISession session : queue) { session.setVersion(version); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setVersion(version); } } @@ -3646,7 +3665,7 @@ public void setQueryTimeout(long timeoutInMs) { for (ISession session : queue) { session.setQueryTimeout(timeoutInMs); } - for (ISession session : occupied.keySet()) { + for (ISession session : occupied) { session.setQueryTimeout(timeoutInMs); } } diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java index d08625cb978c4..2e8c83c40f064 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java @@ -1634,4 +1634,223 @@ private List FakedFirstFetchTsBlockResult() { return Collections.singletonList(tsBlock); } + + @Test + public void test_CurrentOccupiedSize_WhenInsertIsRunning_CountsItsSession() + throws IoTDBConnectionException, StatementExecutionException { + int[] occupiedDuringInsert = {-1}; + Mockito.doAnswer( + invocation -> { + occupiedDuringInsert[0] = sessionPool.currentOccupiedSize(); + return null; + }) + .when(session) + .insertTablet(any(Tablet.class), Mockito.anyBoolean()); + + sessionPool.insertTablet(createBooleanTablet()); + + assertEquals(1, occupiedDuringInsert[0]); + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_CurrentOccupiedSize_WhenNonQueryIsRunning_CountsItsSession() + throws IoTDBConnectionException, StatementExecutionException { + int[] occupiedDuringNonQuery = {-1}; + Mockito.doAnswer( + invocation -> { + occupiedDuringNonQuery[0] = sessionPool.currentOccupiedSize(); + return null; + }) + .when(session) + .executeNonQueryStatement(any(String.class)); + + sessionPool.executeNonQueryStatement("flush"); + + assertEquals(1, occupiedDuringNonQuery[0]); + assertEquals(0, sessionPool.currentOccupiedSize()); + } + + @Test + public void test_CurrentOccupiedSize_WhenInsertWasRejected_IsZero() + throws IoTDBConnectionException, StatementExecutionException { + doThrow(new StatementExecutionException("rejected")) + .when(session) + .insertTablet(any(Tablet.class), Mockito.anyBoolean()); + + try { + sessionPool.insertTablet(createBooleanTablet()); + Assert.fail("expected StatementExecutionException"); + } catch (StatementExecutionException e) { + // expected + } + + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_CurrentOccupiedSize_WhenBrokenSessionWasCleanedUp_IsZero() + throws IoTDBConnectionException, StatementExecutionException { + doThrow(new IoTDBConnectionException("broken")) + .when(session) + .insertTablet(any(Tablet.class), Mockito.anyBoolean()); + + try { + sessionPool.insertTablet(createBooleanTablet()); + Assert.fail("expected IoTDBConnectionException"); + } catch (IoTDBConnectionException e) { + // expected + } + + assertEquals(0, sessionPool.currentOccupiedSize()); + } + + @Test + public void test_CurrentOccupiedSize_WhenResultSetIsOpen_CountsItsSessionUntilClosed() + throws IoTDBConnectionException, StatementExecutionException { + SessionDataSet dataSet = Mockito.mock(SessionDataSet.class); + Mockito.when(session.executeQueryStatement(any(String.class))).thenReturn(dataSet); + + SessionDataSetWrapper wrapper = sessionPool.executeQueryStatement("select * from root.**"); + assertEquals(1, sessionPool.currentOccupiedSize()); + + sessionPool.closeResultSet(wrapper); + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_CloseResultSet_WhenCalledTwice_PutsSessionBackIntoPoolOnlyOnce() + throws IoTDBConnectionException, StatementExecutionException { + SessionDataSet dataSet = Mockito.mock(SessionDataSet.class); + Mockito.when(session.executeQueryStatement(any(String.class))).thenReturn(dataSet); + + SessionDataSetWrapper wrapper = sessionPool.executeQueryStatement("select * from root.**"); + sessionPool.closeResultSet(wrapper); + sessionPool.closeResultSet(wrapper); + + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_CloseResultSet_WhenCalledAgainAfterNextQueryTookSession_LeavesItToNextQuery() + throws IoTDBConnectionException, StatementExecutionException { + SessionDataSet firstDataSet = Mockito.mock(SessionDataSet.class); + SessionDataSet secondDataSet = Mockito.mock(SessionDataSet.class); + Mockito.when(session.executeQueryStatement(any(String.class))) + .thenReturn(firstDataSet, secondDataSet); + + SessionDataSetWrapper first = sessionPool.executeQueryStatement("select * from root.a"); + sessionPool.closeResultSet(first); // as SessionDataSetWrapper.hasNext() does when exhausted + SessionDataSetWrapper second = sessionPool.executeQueryStatement("select * from root.b"); + Assert.assertSame(first.getSession(), second.getSession()); + + sessionPool.closeResultSet(first); // e.g. in the caller's finally block + + assertEquals(0, sessionPool.currentAvailableSize()); + assertEquals(1, sessionPool.currentOccupiedSize()); + Mockito.verify(secondDataSet, Mockito.never()).closeOperationHandle(); + + sessionPool.closeResultSet(second); + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_CloseResultSet_WhenClosingOperationHandleFails_ClosesSessionAndStopsCountingIt() + throws IoTDBConnectionException, StatementExecutionException { + SessionDataSet dataSet = Mockito.mock(SessionDataSet.class); + doThrow(new IoTDBConnectionException("broken")).when(dataSet).closeOperationHandle(); + Mockito.when(session.executeQueryStatement(any(String.class))).thenReturn(dataSet); + + SessionDataSetWrapper wrapper = sessionPool.executeQueryStatement("select * from root.**"); + sessionPool.closeResultSet(wrapper); + + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(0, sessionPool.currentAvailableSize()); + Mockito.verify(session).close(); + } + + @Test + public void test_CloseResultSet_WhenCalledAgainAfterFailedClose_DoesNotRetryClose() + throws IoTDBConnectionException, StatementExecutionException { + SessionDataSet dataSet = Mockito.mock(SessionDataSet.class); + doThrow(new IoTDBConnectionException("broken")).when(dataSet).closeOperationHandle(); + Mockito.when(session.executeQueryStatement(any(String.class))).thenReturn(dataSet); + + SessionDataSetWrapper wrapper = sessionPool.executeQueryStatement("select * from root.**"); + sessionPool.closeResultSet(wrapper); + sessionPool.closeResultSet(wrapper); + + // a failed close leaves the data set open, so a second attempt would send the close request + // again over the broken session and open a second replacement session for it + Mockito.verify(dataSet, Mockito.times(1)).closeOperationHandle(); + } + + @Test + public void test_CloseResultSet_WhenWrapperHasNoSession_DoesNothing() { + SessionDataSet dataSet = Mockito.mock(SessionDataSet.class); + + sessionPool.closeResultSet(new SessionDataSetWrapper(dataSet, null, sessionPool)); + + assertEquals(0, sessionPool.currentOccupiedSize()); + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_Setters_WhenCalledWhileInsertIsRunning_ReachSessionExecutingInsert() + throws IoTDBConnectionException, StatementExecutionException { + // setTimeZone first sends the new zone to the server, through an idle session + Session idleSession = Mockito.mock(Session.class); + ((ConcurrentLinkedDeque) Whitebox.getInternalState(sessionPool, "queue")) + .add(idleSession); + Mockito.doAnswer( + invocation -> { + sessionPool.setTimeZone("+02:00"); + sessionPool.setFetchSize(4321); + sessionPool.setEnableRedirection(false); + sessionPool.setEnableQueryRedirection(true); + sessionPool.setVersion(Version.V_0_13); + sessionPool.setQueryTimeout(1234); + return null; + }) + .when(session) + .insertTablet(any(Tablet.class), Mockito.anyBoolean()); + + sessionPool.insertTablet(createBooleanTablet()); + + Mockito.verify(session).setFetchSize(4321); + Mockito.verify(session).setEnableRedirection(false); + Mockito.verify(session).setEnableQueryRedirection(true); + Mockito.verify(session).setVersion(Version.V_0_13); + Mockito.verify(session).setQueryTimeout(1234); + Mockito.verify(session).setTimeZoneOfSession("+02:00"); + } + + @Test + public void test_Close_WhenInsertIsRunning_ClosesSessionExecutingInsert() + throws IoTDBConnectionException, StatementExecutionException { + Mockito.doAnswer( + invocation -> { + sessionPool.close(); + return null; + }) + .when(session) + .insertTablet(any(Tablet.class), Mockito.anyBoolean()); + + sessionPool.insertTablet(createBooleanTablet()); + + Mockito.verify(session).close(); + } + + private static Tablet createBooleanTablet() { + List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema("pressure", TSDataType.BOOLEAN)); + long[] timestamp = new long[] {1L, 2L}; + boolean[][] values = new boolean[][] {{true, false}}; + return new Tablet("device1", schemas, timestamp, values, new BitMap[1], 2); + } }