Skip to content

Fix SessionPool's bookkeeping of sessions in use #19092

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yamass wants to merge 1 commit into
base: master
Choose a base branch
from
+269 −31
Open

Fix SessionPool's bookkeeping of sessions in use #19092

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 50 additions & 31 deletions iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -100,8 +101,14 @@ public class SessionPool implements ISessionPool {
private static final int FINAL_RETRY = RETRY - 1;

private final ConcurrentLinkedDeque<ISession> queue = new ConcurrentLinkedDeque<>();
// for session whose resultSet is not released.
private final ConcurrentMap<ISession, ISession> occupied = new ConcurrentHashMap<>();
// Every session taken by getSession() and not yet returned (put back or cleaned), whatever it
// is used for.
private final Set<ISession> occupied = ConcurrentHashMap.newKeySet();

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

// 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<ISession, SessionDataSetWrapper> occupiedByResultSet =
new ConcurrentHashMap<>();
private int size = 0;
private int maxSize = 0;
private final long waitToGetSessionTimeoutInMs;
Expand Down Expand Up @@ -702,6 +709,7 @@ private ISession getSession() throws IoTDBConnectionException {
throw new IoTDBConnectionException(SESSION_POOL_IS_CLOSED);
}
if (session != null) {
occupied.add(session);
return session;
}

Expand Down Expand Up @@ -790,6 +798,7 @@ private ISession getSession() throws IoTDBConnectionException {
}
}

occupied.add(session);
return session;
}

Expand All @@ -802,13 +811,18 @@ 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();
}

@SuppressWarnings({"squid:S2446"})
protected void putBack(ISession session) {
occupied.remove(session);

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

queue.push(session);
synchronized (this) {
// we do not need to notifyAll as any waited thread can continue to work after waked up.
Expand All @@ -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 */
Expand All @@ -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) {

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

try {
session.close();
} catch (IoTDBConnectionException e) {
Expand All @@ -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();

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

} 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"})
Expand Down Expand Up @@ -905,6 +923,7 @@ private void tryConstructNewSession() {

private void closeSession(ISession session) {
if (session != null) {
occupied.remove(session);
try {
session.close();
} catch (Exception e2) {
Expand Down Expand Up @@ -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);

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

return wrapper;
} catch (IoTDBConnectionException e) {
// TException means the connection is broken, remove it and get a new one.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -3206,7 +3225,7 @@ public SessionDataSetWrapper executeLastDataQuery(List<String> 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.
Expand All @@ -3233,7 +3252,7 @@ public SessionDataSetWrapper executeLastDataQuery(List<String> 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.
Expand All @@ -3260,7 +3279,7 @@ public SessionDataSetWrapper executeLastDataQuery(List<String> 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.
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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);
}
}
Expand Down Expand Up @@ -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);

@JackieTien97 JackieTien97 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

}
}
Expand Down Expand Up @@ -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);
}
}
Expand All @@ -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);
}
}
Expand Down Expand Up @@ -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);
}
}
Expand All @@ -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);
}
}
Expand Down
Loading