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