[yamass]

Description

TL;DR

SessionPool only tracks sessions held by open result sets (the occupied map), but treats that map as "all sessions in use".
As a result, currentOccupiedSize(), six setters and close() ignore sessions executing a request.
closeResultSet() can also release a session that a later query already has taken, leaks the session when closing fails,
and throws an NPE for a wrapper without a session.

All these bugs stem from the same root cause. Each bug has one or more tests in SessionPoolTest that fails on master. See table below.

Longer Description

Problem

SessionPool keeps idle sessions in queue. Apart from that, it has an additional collection: the occupied map.
It contains only sessions whose query result set is still open, until the result set is closed.
A session executing a request – an insert, a non-query statement, or a query whose result set has not been returned yet – is in neither queue nor occupied.

Note that in the following, I will just refer sessions executing requests instead of spelling out that these requests are inserts, non-query statements, or queries whose result set has not been returned yet.

However, several parts of the class treat occupied as if it held every session taken from the pool,
which it doesn't:

1. currentOccupiedSize() ignores sessions executing a request. On a write-heavy pool whose sessions are all
   writing, it reports 0 occupied and 0 available. We saw exactly this in production: "occupied" stayed at 0–1 and "available" was 0,
   all while writes queued up.
2. setFetchSize, setTimeZone, setEnableRedirection, setEnableQueryRedirection,
   setVersion and setQueryTimeout apply the new value to queue and occupied only. A session executing a request
   misses the change and keeps the old value for its whole life after being put back, so the pool ends up with mixed settings.
3. close() closes queue and occupied only. A session executing a request at that moment is later put
   back into the already cleared queue, and its connection is never closed.

closeResultSet() has problems of its own:

4. A late second release can steal another thread's session. A result set is usually released
   twice: by SessionDataSetWrapper.hasNext() once it is exhausted, and again by the caller's
   closeResultSet()/close() in a finally block. The second call only checks whether the
   session is still in occupied. If another thread took that session in between and opened a
   new result set with it, the late call puts the session back while the other thread is still
   reading, and two threads end up sharing one connection.
5. A failed close leaks the session. When closing the operation handle fails, a replacement
   session is opened, but the failed session is neither put back nor closed, so its connection
   leaks. The failed close also leaves the data set open, so the second release (see 4) sends the
   close request again over the broken session and opens yet another replacement for it.
   tryConstructNewSession() does not count replacements in size, so the pool can end up with
   more sessions than maxSize.
6. A wrapper without a session throws a NullPointerException (ConcurrentHashMap.remove(null)).

Fix

Only one implementation file was changed: SessionPool.

• occupied is now a concurrent set of all sessions handed out by getSession(). A session leaves
  it when it is put back (putBack), closed after a connection error (closeSession, used by both
  cleanSessionAndMayThrowConnectionException overloads), or released by closeResultSet().
• currentOccupiedSize() and the timeout warning report occupied.size() (1).
• The six setters and close() iterate occupied instead of the result-set sessions (2, 3).
  The setters only change client-side fields of the session, as they already did for sessions
  held by result sets. A session closed by close() makes its running request fail with a
  connection error, the same as reading a result set of a closed pool already did. The retry then
  fails with "session pool is closed". Session.close() returns early when already closed, so a
  second close by the request's error handling is harmless.
• The former occupied map is renamed to occupiedByResultSet and maps each session to the
  SessionDataSetWrapper holding it; occupy(session, wrapper) stores it.
• closeResultSet() first claims the session with occupiedByResultSet.remove(session, wrapper)
  and returns if that fail. So only the first release of this result set closes its operation
  handle and hands the session back (4), even if that close failed (5). A failed close goes through
  cleanSessionAndMayThrowConnectionException(session), which closes the session and opens a
  replacement (5). A wrapper without a session is ignored (6).

No public API changes.

Tests

New tests in SessionPoolTest. Eight of them fail on master, at least one per bug above; the other
four cover every way a session goes back, so the count cannot drift.

Test	Covers	On master
test_CurrentOccupiedSize_WhenInsertIsRunning_CountsItsSession	(1)	fails
test_CurrentOccupiedSize_WhenNonQueryIsRunning_CountsItsSession	(1)	fails
test_Setters_WhenCalledWhileInsertIsRunning_ReachSessionExecutingInsert	(2)	fails
test_Close_WhenInsertIsRunning_ClosesSessionExecutingInsert	(3)	fails
test_CloseResultSet_WhenCalledAgainAfterNextQueryTookSession_LeavesItToNextQuery	(4)	fails
test_CloseResultSet_WhenClosingOperationHandleFails_ClosesSessionAndStopsCountingIt	(5)	fails
test_CloseResultSet_WhenCalledAgainAfterFailedClose_DoesNotRetryClose	(5)	fails
test_CloseResultSet_WhenWrapperHasNoSession_DoesNothing	(6)	fails
test_CurrentOccupiedSize_WhenInsertWasRejected_IsZero	return via putBack	passes
test_CurrentOccupiedSize_WhenBrokenSessionWasCleanedUp_IsZero	return via cleanup	passes
test_CurrentOccupiedSize_WhenResultSetIsOpen_CountsItsSessionUntilClosed	return via closeResultSet	passes
test_CloseResultSet_WhenCalledTwice_PutsSessionBackIntoPoolOnlyOnce	double release	passes

---

This PR has:

• been self-reviewed.
• added comments explaining the "why" and the intent of the code wherever would not be obvious
  for an unfamiliar reader.
• added unit tests or modified existing tests to cover new code paths, ensuring the threshold
  for code coverage.

---

Key changed/added classes (or packages if there are too many classes) in this PR

• SessionPool.java + tests

[JackieTien97]

Thanks for the thorough analysis and the regression tests. The diagnosis is right: occupied only tracked result-set sessions, but was treated as "all sessions in use". Claiming the session with occupiedByResultSet.remove(session, wrapper) is a neat fix for the double release. I checked the tests locally: running this PR's SessionPoolTest against master's SessionPool, 8 of the new tests fail; on this branch all 105 pass, and Spotless is clean.

My main concern is the new close(), which now closes sessions that other threads are still using. In a reproduction against a fake server, close() blocked until the in-flight RPC finished, holding the pool monitor the whole time. With the default connectionTimeoutInMs = 0, that wait has no upper bound. The inline comment on close() has the details and a suggested alternative: close sessions when they are returned to a closed pool. The other comments are smaller:

• a narrow window between queue and occupied,
• TableSessionWrapper still bypassing the new invariant,
• setTimeZone robustness,
• test nits.

[JackieTien97]

Closing sessions that are still in use from another thread has side effects beyond making the running request fail:

1. close() now waits for every in-flight RPC while holding the pool monitor. Session.close() → SessionConnection.close() sends closeSession through the client from RpcUtils.newSynchronizedClient. SynchronizedHandler holds synchronized (client) for the whole duration of a call, so closeSession has to queue behind the in-flight call (an insert, flush, a long query, …). Meanwhile, every putBack() blocks on the pool monitor, and so does any getSession() that has to wait or create a session. With the default connectionTimeoutInMs = 0 there is no read timeout, so a hung DataNode or a half-open connection makes close() hang indefinitely.
   I reproduced this with a fake Thrift server (a mocked IClientRPCService.Iface behind TThreadPoolServer) and a real SessionPool, with one insertTablet held server-side for ~3 s:
   • on this branch, close() returned after ~4 s, and another thread that tried to synchronize on the pool was blocked the whole time;
   • on master, close() returned immediately, leaking that session (bug 3).
2. Sessions borrowed from TableSessionPool are closed under the user. A later ITableSession.close() (e.g. at the end of a try-with-resources block) throws IoTDBConnectionException: Session is not open, please invoke Session.open() first if the user changed the database. The reason is that TableSessionWrapper.close() runs use <pool database> on the closed session and only catches StatementExecutionException. Reproduced with the same fake server: pool built with database("db1"), then use db2, then pool.close(), then tableSession.close() throws. On master it returns normally.
3. A retrying request can bring the closed session back to life. If the request is between attempts in SessionConnection.callWithRetry (by default up to 60 × 500 ms), the next reconnect() re-inits the transport. It then calls session.setDefaultSessionConnection(this) on a Session that is already marked closed. From then on Session.close() returns early, so the new connection is never closed: the same symptom as bug 3.

Suggestion: instead of closing sessions while they are in use, close them when they come back to a closed pool:

• In putBack(), do occupied.remove, the closed check, and then queue.push or session.close(), all inside the synchronized (this) block it already has.
• In getSession(), close the polled session when the pool turns out to be closed.
• Make closed volatile.

In-flight requests then finish normally, their connections are closed when they are returned, and close() never has to wait on another thread's RPC while holding the monitor. If you'd still like close() to wait for in-flight work, a bounded wait outside the monitor would be safer. Happy to share the probe test if it helps.

[JackieTien97]

occupied.remove(session) and queue.push(session) run outside the monitor, and so do queue.poll() → occupied.add(session) in getSession(). On the create path the gap is between the closed check (L782–L788) and L801. While a session is between the two calls, it is in neither collection, and a close() or setter that runs in that window misses it:

• close() doesn't close it, so it is later pushed into the cleared queue and its connection leaks.
• A setter doesn't update it, so the session keeps the old value.

These are the same symptoms as bugs (2) and (3), just in a narrower window. Moving the remove/push (and a closed check) into the existing synchronized (this) block would close the window on the putBack() side. For the getSession() side, see the suggestion on close(); for the setters, the pool could also re-apply its current settings when it hands out a session.

[JackieTien97]

TableSessionWrapper.close() (not part of this diff) still bypasses this invariant. When restoring sql_dialect or the database fails, it either calls session.close() directly (TableSessionWrapper.java L121, L137) or lets an IoTDBConnectionException escape (L114, L131). In both cases the session never leaves occupied, so currentOccupiedSize() can still drift. In a unit test it stays at 1 after such a failure; master reports 0.

Routing these branches through sessionPool.cleanSessionAndMayThrowConnectionException(session) would keep the count right. It would also stop the pool from losing capacity, because size is never decremented on those paths today.

[JackieTien97]

closeSession() removes a session from occupied and then closes it, but this loop uses a weakly consistent iterator, so it can still return a session that another thread has just removed and closed. That happens after a request fails with a connection error, or on the new closeResultSet() failure path. Session.setTimeZoneOfSession() then throws IoTDBConnectionException("Session is not open") from getDefaultSessionConnection(). By that point the server-side time zone and this.zoneId have already changed, but the rest of occupied is skipped, so the pool ends up with mixed time zones again. Catching and ignoring the failure per session would make the loop robust; that session is going away anyway. This is the only one of the six setters that can throw; the others only assign fields.

Unrelated to this PR, but in the same method: the retry loop above never breaks or returns after a successful session.setTimeZone(zoneId); putBack(session);, so it sends the RPC RETRY (3) times. If the 3rd attempt hits a connection error, cleanSessionAndMayThrowConnectionException(session, FINAL_RETRY, e) throws, even though the time zone was already set.

[JackieTien97]

Minor: the old finally guaranteed that the session was always put back or dropped. Now the session is claimed first (L883), and only IoTDBConnectionException/StatementExecutionException lead to cleanup. Any other RuntimeException/Error thrown here leaves the session neither put back nor closed, while it is still counted in occupied. One example is an NPE when getSessionDataSet() is null (setSessionDataSet is public). This is unlikely, because SynchronizedHandler wraps RPC-level failures into TException. Still, keeping a try/finally around the release, or also catching RuntimeException, would be cheap insurance.

[JackieTien97]

Minor: this test, test_CurrentOccupiedSize_WhenBrokenSessionWasCleanedUp_IsZero and test_CloseResultSet_WhenCalledAgainAfterFailedClose_DoesNotRetryClose all reach tryConstructNewSession(). Each of them therefore opens a real Session to the fake host:11, which took 0.7–1.3 s per test locally, spent waiting for DNS resolution to fail.

This assertion effectively checks that the replacement connection failed, which depends on the environment. With the default connectionTimeoutInMs = 0, a resolver that maps host to an address that silently drops packets would also make these tests hang until the OS TCP timeout. Consider one of these:

• point the pool at 127.0.0.1 and a port with nothing listening, so the replacement fails fast without a DNS lookup;
• assert on occupied instead of currentAvailableSize().

[JackieTien97]

Nit: wrapper construction followed by occupy(session, wrapper) is now repeated in 12 query methods. A small helper such as wrapResultSet(resp, session) that builds and registers the wrapper would replace occupy(). It would also keep a future query method from forgetting the registration; if that happened, closeResultSet() would silently return and the session would never go back to the pool.