[shuwenwei]

Problem

TableSessionPool.getSession() borrows a raw Session from the underlying SessionPool, but the table-session path did not register that Session in occupied. A later pool.close() therefore only saw idle sessions and sessions tracked for old tree-model result sets. Active table sessions could remain open, and returning one after the pool had closed could put it back into the closed pool.

A concrete failure sequence is:

ITableSessionPool pool = new TableSessionPoolBuilder()
    .nodeUrls(nodeUrls)
    .user(user)
    .password(password)
    .maxSize(1)
    .build();

ITableSession session = pool.getSession();
pool.close();
session.close();

Before this change:

• pool.close() did not see the borrowed table session because it was absent from both queue and occupied.
• session.close() then returned the connection to the already-closed pool, leaving it unmanaged and open.
• Closing the wrapper twice, or closing it after a connection failure, could throw an NPE.
• A failure while restoring SQL dialect/database could close the physical session without restoring the configured pool capacity.

After this change, the borrowed session is tracked until release; pool.close() closes it, and releasing a session into a closed pool closes it directly instead of publishing it back into queue.

The same cleanup logic also fixes an expected exception leak in tree-model sortTablet():

maxSize=1
sortTablet(tablet) throws IllegalArgumentException

Before: the borrowed session was not returned, so subsequent operations waited until timeout even though size == 1. After: try/finally returns the session, allowing the next operation to borrow it normally.

Changes

• Track pooled table sessions in occupied from acquisition until wrapper release.
• Make table-session close idempotent and safe after connection failures.
• Release sessions through one path that removes occupied, reuses the connection only when safe, and replaces failed sessions while the pool remains open.
• Close sessions directly when they are returned after the pool has closed.
• Preserve pool capacity when SQL dialect/database restoration fails.
• Ensure sortTablet() returns its session when sorting succeeds or throws.
• Add lifecycle tests for pool close, repeated table-session close with maxSize=1, and sortTablet() failure recovery.
• Add deterministic latch-based tests that force pool.close() between queue removal and occupied registration, and immediately before session release.

Validation

• git diff --check passed.
• Spotless checks passed for iotdb-client/session and integration-test.
• Tests were added but were not executed locally. CI will exercise them.

[codecov]

Codecov Report

❌ Patch coverage is 65.00000% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 49.04%. Comparing base (7b407ec) to head (665c926).
⚠️ Report is 13 commits behind head on master.

Files with missing lines	Patch %	Lines
...apache/iotdb/session/pool/TableSessionWrapper.java	41.02%	23 Missing ⚠️
...ava/org/apache/iotdb/session/pool/SessionPool.java	80.32%	12 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19247      +/-   ##
============================================
+ Coverage     48.89%   49.04%   +0.14%     
  Complexity      689      689              
============================================
  Files          5588     5590       +2     
  Lines        413444   414011     +567     
  Branches      54537    54645     +108     
============================================
+ Hits         202154   203052     +898     
+ Misses       211290   210959     -331     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[sonarqubecloud]

Quality Gate failed

Failed conditions
C Reliability Rating on New Code (required ≥ A)

See analysis details on SonarQube Cloud

Catch issues before they fail your Quality Gate with our IDE extension SonarQube for IDE

[JackieTien97]

The lifecycle problems described in this PR are real. I checked the fixes with a fake Thrift server probe (a mock IClientRPCService.Iface behind a TThreadPoolServer, plus mocked Sessions), running the same probes against the merge base and against this head:

Scenario	merge base	this PR
pool.close() while a table session is borrowed	session left open	session closed
sql_dialect restore fails, maxSize=1	size=1 and empty queue, so the pool is exhausted for good	capacity kept
tree-model result set closed after pool.close()	session pushed into the closed pool's queue	session closed
ITableSession.close() after a connection failure	NPE	returns normally

There is one new problem, though. pool.close() now waits for in-flight RPCs on borrowed table sessions, and it holds the pool monitor while it waits (see the comment on close()). On this head, with one table session running a statement that the server holds for about 4.5 s, close() took 4528 ms, and a concurrent getSession() was blocked on the monitor for 1.5 s before it failed with Session pool is closed. On the merge base, close() returned in 0 ms. Because connectionTimeoutInMs defaults to 0 (no socket read timeout), one stuck RPC can hang close(), and every other borrow and return, indefinitely. I think this should be fixed before merging.

The other comments are smaller:

• The tree-model query path still has the close race that this PR fixes for table sessions.
• The wrapper still throws an NPE if it is reused after a connection failure.
• Two of the capacity fixes have no test.
• A few nits, marked as such.

CI: the only failing IT job is IoTDBLoadTsFileClusterIT.testRestartedDataNodeIsSyncedWithTheLoadedFile (tree model, 1C3D), which looks unrelated. The new unit tests and table-model ITs ran and passed.

[JackieTien97]

Borrowed table sessions are now tracked in occupied, so this loop closes sessions that other threads may be using at that moment, and it holds the pool monitor while doing so. Session.close() calls SessionConnection.close(), which sends closeSession through the synchronized client proxy (SynchronizedHandler locks on the client). The call therefore waits until the other thread's in-flight RPC returns.

Probe: a fake server, a table pool with maxSize=2, and one borrowed session running a statement that the server holds for about 4.5 s.

• This PR: close() returned after 4528 ms. A concurrent getSession() waited on the pool monitor for 1.5 s and then got Session pool is closed.
• Merge base: close() returned in 0 ms, because borrowed table sessions were not touched here.

connectionTimeoutInMs defaults to 0, which means no socket read timeout. One stuck RPC can therefore hang close() indefinitely. Everything that needs the monitor hangs with it: getSession(), putBack(), releaseSession(), and TableSessionWrapper.close() (through isPoolClosed()).

Suggestion: under the lock, set closed = true and take and clear snapshots of queue and occupied. Close the snapshots after leaving the monitor. For sessions that are still in use, consider closing the transport directly instead of queueing a closeSession RPC behind the in-flight call. SessionConnection already has forceCloseTransport().

[JackieTien97]

This fixes the borrow/close race for table sessions, but the tree-model query path still has it. The 12 query methods call occupy(session) after the RPC returns, and occupy() does not check closed. If close() runs while the query RPC is in flight, the session goes into occupied after close() has already cleared it.

Probe: executeQueryStatement blocked inside the RPC, then pool.close(), then the RPC returns. The wrapper is still returned, session.close() was called 0 times, and currentOccupiedSize() is 1 on the closed pool. The new releaseSession() does close the session once the caller closes the result set. If the caller relies on pool.close() for cleanup, for example by abandoning result sets at shutdown, the connection stays open.

Could occupy() do the same check under the lock as this code, closing the session and throwing if the pool is closed? Another option is to register every borrowed session in getSession() itself, under the lock that already checks closed. That would also remove the window between the two lock acquisitions here.

[JackieTien97]

releaseOnConnectionFailure() and close() both set session to null. If the same wrapper is called again, this line dereferences null and throws NullPointerException instead of an IoTDBConnectionException. Both executeQueryStatement overloads and executeNonQueryStatement behave the same way.

Probe: the first executeNonQueryStatement fails with IoTDBConnectionException. The second throws NullPointerException: Cannot invoke "org.apache.iotdb.session.Session.executeNonQueryStatement(String)" because "this.session" is null. The merge base behaves the same, so this is not a regression. Still, a caller that retries on the same ITableSession after a connection error will hit it, and this PR already makes close() safe after a connection failure. A guard at the top of these methods that throws IoTDBConnectionException, with the message as an i18n constant, would make the behavior consistent.

[JackieTien97]

Two of the fixes listed in the description have no test:

1. Restoring sql_dialect/database fails: the session is closed and pool capacity is kept (this line with reusable = false).
2. A connection failure in insert/execute*: releaseOnConnectionFailure() closes the session and replenishes the pool.

If either path went back to a bare session.close(), every test in this PR would still pass.

I checked (1) with a probe: maxSize=1, a getSqlDialect() that differs from the pool's, and a restore that throws StatementExecutionException. The merge base ends with size=1 and an empty queue, so the pool is exhausted for good. This PR ends with size=0: in the probe the replacement could not connect, so the slot was freed. In SessionPoolTest, mocking a failing executeNonQueryStatement and then asserting size, currentOccupiedSize() and currentAvailableSize() would protect both paths.

[JackieTien97]

SessionPool is public and not final, and this protected overload has existed since 2.0.1-beta. Removing it breaks any subclass that calls it. Recompiling against the new jar fails, and running an old build against the new jar throws NoSuchMethodError. If compatibility matters here, a @Deprecated overload that delegates to discardSession(session) would keep those subclasses working. Also note that releaseSession and isPoolClosed become new protected API.

[JackieTien97]

nit: only one session is pushed here, but notifyAll() wakes every waiting borrower. putBack() and tryConstructNewSession() use notify(). With many waiters, each release wakes all of them; all but one get null from poll(), and each of those calls notify() again. notify() is enough here and matches putBack().

[JackieTien97]

nit: every caller of this method is a tree-model operation, and its session is not yet in occupied when the exception happens. releaseSession(session, false) therefore always returns false here, and discardSession() does the actual work. discardSession() also duplicates the close-then-replenish tail of releaseSession(). A single private discard(session) that removes the session from occupied, closes it, and replenishes if the pool is open could serve both callers, so the two paths cannot drift apart.

[JackieTien97]

nit: closed = true is always set together with session = null. sessionToRelease can only be null here if closed was already true, so this branch is unreachable. Either use session == null as the closed state, or keep closed and drop this branch. The this.closed = false in the constructor is also redundant.

[JackieTien97]

nit (stale comments): the class Javadoc still says the wrapper calls cleanSessionAndMayThrowConnectionException on connection errors, but that overload is gone and this method releases through releaseSession. Likewise, the comment on occupied in SessionPool (// for session whose resultSet is not released.) is no longer accurate, because borrowed table sessions are tracked there now. That also changes what currentOccupiedSize() reports.

[JackieTien97]

nit: both race tests depend on implementation details, and if those change, the failures hide the cause.

• This pool.wait() works only because getSession() calls poll() while holding the pool monitor. If poll() moves out of the lock, pool.wait() throws IllegalMonitorStateException inside the borrow. The test then fails on assertTrue(e.getCause() instanceof IoTDBConnectionException) with no message, so the real exception is never shown.
• testTableSessionReleaseRacingWithPoolClose depends on getDatabase() being evaluated before sessionPool.database != null. If someone swaps that condition, which is a natural micro-optimization, the latch never counts down. The test then fails at assertTrue(aboutToRelease.await(...)) after 5 s, again with no message.