Fix session pool lifecycle bugs in getSession/releaseSession/cleanupIdleSessions

src/client/BaseSessionPool.ts

@@ -50,7 +50,7 @@ export abstract class BaseSessionPool {
  protected config: PoolConfig;
  protected endPoints: EndPoint[];
  protected pool: PooledSession[] = [];
- protected waitQueue: Denque<(session: Session) => void> = new Denque();
+ protected waitQueue: Denque<(session: Session) => boolean> = new Denque();
  protected idleSessions: Denque<PooledSession> = new Denque();
  protected activeSessions: Set<PooledSession> = new Set();
  protected currentEndPointIndex = 0;
@@ -237,6 +237,11 @@ export abstract class BaseSessionPool {
  // Verify session is still open
  if (pooledSession.session.isOpen()) {
  this.activeSessions.add(pooledSession);
+ // Mark in use — the new-session and waiter branches both do this; the
+ // idle-reuse path omitting it left a reused, actively-in-use session
+ // with inUse===false, which syncDatabaseContextToPool (filters
+ // !inUse) would treat as idle and fire a concurrent USE on.
+ pooledSession.inUse = true;
  pooledSession.lastUsed = Date.now();
  const duration = Date.now() - startTime;
  logger.debug(
@@ -261,7 +266,15 @@ export abstract class BaseSessionPool {
  const session = await this.createSession();
  const pooledSession = this.pool.find((ps) => ps.session === session);
  if (pooledSession) {
- this.idleSessions.shift(); // Remove from idle since we just added it
+ // Remove *this* session from idle. createSession() pushed it to the
+ // back of idleSessions; a blind shift() removes the front, which under
+ // concurrent interleaving (a session released into idle while we were
+ // awaiting createSession) would evict a different session and leave
+ // this one double-tracked as both idle and active.
+ const idleIndex = this.idleSessions.toArray().indexOf(pooledSession);
+ if (idleIndex > -1) {
+ this.idleSessions.remove(idleIndex, 1);
+ }
  this.activeSessions.add(pooledSession);
  pooledSession.inUse = true;
  }
@@ -276,9 +289,31 @@ export abstract class BaseSessionPool {
  );
  const waitTimeout = this.config.waitTimeout || 60000;
  return new Promise((resolve, reject) => {
+ let settled = false;
+
+ // The queue stores this exact wrapper. On timeout we must remove *this*
+ // reference (not `resolve`, which is never in the queue), and the
+ // `settled` guard makes timeout and fulfillment mutually exclusive so a
+ // session is never handed to a waiter whose promise already rejected.
+ const waiter = (session: Session): boolean => {
+ if (settled) {
+ return false;
+ }
+ settled = true;
+ clearTimeout(timeoutId);
+ const duration = Date.now() - startTime;
+ logger.debug(`[PERF] getSession (waited): ${duration}ms`);
+ resolve(session);
+ return true;
+ };
+
  const timeoutId = setTimeout(() => {
+ if (settled) {
+ return;
+ }
+ settled = true;
  const waiters = this.waitQueue.toArray();
- const index = waiters.indexOf(resolve);
+ const index = waiters.indexOf(waiter);
  if (index > -1) {
  this.waitQueue.remove(index, 1);
  }
@@ -290,12 +325,7 @@ export abstract class BaseSessionPool {
  timeoutId.unref();
  }
 
- this.waitQueue.push((session: Session) => {
- clearTimeout(timeoutId);
- const duration = Date.now() - startTime;
- logger.debug(`[PERF] getSession (waited): ${duration}ms`);
- resolve(session);
- });
+ this.waitQueue.push(waiter);
  });
  }
 
@@ -313,22 +343,27 @@ export abstract class BaseSessionPool {
  pooledSession.inUse = false;
  pooledSession.lastUsed = Date.now();
 
- // Check if there are waiting requests
- if (this.waitQueue.length > 0) {
+ // Hand the session to the first waiter that is still pending. A waiter
+ // whose promise already settled (e.g. it timed out) returns false; skip
+ // it and try the next one, so a released session is never leaked to a
+ // dead waiter (which would leave it marked active but held by nobody).
+ while (this.waitQueue.length > 0) {
  const waiter = this.waitQueue.shift();
- if (waiter) {
- // Move to active for the waiter
- this.activeSessions.add(pooledSession);
- pooledSession.inUse = true;
- waiter(session);
- } else {
- // No waiter actually found, add back to idle
- this.idleSessions.push(pooledSession);
+ if (!waiter) {
+ continue;
  }
- } else {
- // No waiters, add back to idle
- this.idleSessions.push(pooledSession);
+ this.activeSessions.add(pooledSession);
+ pooledSession.inUse = true;
+ if (waiter(session)) {
+ return;
+ }
+ // Stale waiter; undo the active bookkeeping and try the next one.
+ this.activeSessions.delete(pooledSession);
+ pooledSession.inUse = false;
  }
+
+ // No live waiter; add back to idle.
+ this.idleSessions.push(pooledSession);
  }
  }
 
@@ -342,24 +377,33 @@ export abstract class BaseSessionPool {
  const idleArray = this.idleSessions.toArray();
 
  for (const ps of idleArray) {
- if (now - ps.lastUsed > maxIdleTime && this.pool.length > minSize) {
+ // Subtract already-queued removals so the pool never drops below
+ // minPoolSize: without this the guard sees the constant pre-cleanup size
+ // and can queue every idle session, collapsing the pool to 0.
+ if (
+ now - ps.lastUsed > maxIdleTime &&
+ this.pool.length - sessionsToRemove.length > minSize
+ ) {
  sessionsToRemove.push(ps);
  }
  }
 
  await Promise.all(
  sessionsToRemove.map(async (ps) => {
+ // Remove from pool + idle BEFORE closing. close() awaits the
+ // closeSession RPC and isOpen() stays true until it resolves, so a
+ // concurrent getSession() could otherwise shift() this session and
+ // hand out a connection that is about to be destroyed.
+ const poolIndex = this.pool.indexOf(ps);
+ if (poolIndex > -1) {
+ this.pool.splice(poolIndex, 1);
+ }
+ const idleIndex = this.idleSessions.toArray().indexOf(ps);
+ if (idleIndex > -1) {
+ this.idleSessions.remove(idleIndex, 1);
+ }
  try {
  await ps.session.close();
- const poolIndex = this.pool.indexOf(ps);
- if (poolIndex > -1) {
- this.pool.splice(poolIndex, 1);
- }
- // Remove from idle sessions deque
- const idleIndex = this.idleSessions.toArray().indexOf(ps);
- if (idleIndex > -1) {
- this.idleSessions.remove(idleIndex, 1);
- }
  logger.debug(`Removed idle session from ${this.getPoolName()}`);
  } catch (error) {
  logger.error("Error closing idle session:", error);

tests/unit/BaseSessionPoolLifecycle.test.ts

@@ -0,0 +1,205 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { SessionPool } from "../../src/client/SessionPool";
+
+/**
+ * Test pool that hands out lightweight fake sessions (no real IoTDB
+ * connection) so getSession()/releaseSession() lifecycle bookkeeping can be
+ * exercised deterministically. `createGate`, when set, lets a test hold a
+ * createSession() call open to interleave a concurrent release.
+ */
+class TestPool extends SessionPool {
+ public createGate: Promise<void> | null = null;
+ private counter = 0;
+
+ protected async createPoolSession(): Promise<any> {
+ if (this.createGate) {
+ await this.createGate;
+ }
+ const id = ++this.counter;
+ const session: any = {
+ id,
+ _closed: false,
+ // When set to a promise, close() parks on it while isOpen() keeps
+ // returning true — lets a test hold a close open to probe TOCTOU.
+ closeGate: null as Promise<void> | null,
+ isOpen: () => !session._closed,
+ close: async () => {
+ if (session.closeGate) {
+ await session.closeGate;
+ }
+ session._closed = true;
+ },
+ };
+ return session;
+ }
+
+ pooledFor(session: any): any {
+ return (this as any).pool.find((ps: any) => ps.session === session);
+ }
+ runCleanup(): Promise<void> {
+ return (this as any).cleanupIdleSessions();
+ }
+
+ idleSessionObjects(): any[] {
+ return (this as any).idleSessions.toArray().map((ps: any) => ps.session);
+ }
+ activeSessionObjects(): any[] {
+ return Array.from((this as any).activeSessions as Set<any>).map(
+ (ps: any) => ps.session,
+ );
+ }
+}
+
+function newPool(overrides: Record<string, unknown>): TestPool {
+ return new TestPool({
+ host: "localhost",
+ port: 6667,
+ minPoolSize: 0,
+ ...overrides,
+ } as any);
+}
+
+describe("BaseSessionPool lifecycle", () => {
+ it("does not leak a released session to a timed-out waiter (no starvation)", async () => {
+ const pool = newPool({ maxPoolSize: 1, waitTimeout: 50 });
+
+ const s1 = await pool.getSession(); // creates S1; pool is now full (1/1)
+
+ // Pool is full, so this acquisition waits and then times out.
+ await expect(pool.getSession()).rejects.toThrow(/Timeout/);
+
+ // Releasing S1 must return it to the pool, not hand it to the dead waiter
+ // (which would mark S1 active-but-held-by-nobody and starve the pool).
+ pool.releaseSession(s1);
+
+ // S1 must be acquirable again. On the buggy code this getSession() starves
+ // and rejects with a timeout.
+ const s2 = await pool.getSession();
+ expect(s2).toBe(s1);
+
+ await pool.close();
+ });
+
+ it("create-branch removes the new session, not a concurrently-released idle one", async () => {
+ const pool = newPool({ maxPoolSize: 3, waitTimeout: 1000 });
+
+ const s1 = await pool.getSession(); // create S1; active, idle=[]
+
+ // Hold the next createSession() open so we can release S1 mid-flight.
+ let openGate!: () => void;
+ pool.createGate = new Promise<void>((resolve) => {
+ openGate = resolve;
+ });
+
+ const acquireA = pool.getSession(); // enters create-branch, awaits the gate
+ // One yield is enough: getSession() runs synchronously up to the first
+ // await (the createSession call), so after this tick acquireA is parked on
+ // the gate and the release below interleaves before it resumes.
+ await new Promise((r) => setImmediate(r));
+
+ // Concurrent release pushes S1 to the FRONT of idle while A is awaiting.
+ pool.releaseSession(s1); // idle=[S1]
+
+ openGate(); // A's createSession resolves -> pushes S2 -> idle=[S1,S2]
+ const s2 = await acquireA;
+
+ expect(s2).not.toBe(s1);
+ // S1 must remain the idle session; S2 was handed to A and must not also
+ // linger in idle (the blind shift() bug evicted S1 and left S2 in idle).
+ expect(pool.idleSessionObjects()).toContain(s1);
+ expect(pool.idleSessionObjects()).not.toContain(s2);
+ expect(pool.activeSessionObjects()).toContain(s2);
+
+ await pool.close();
+ });
+
+ it("marks a reused idle session as inUse", async () => {
+ const pool = newPool({ maxPoolSize: 2 });
+
+ const s1 = await pool.getSession();
+ pool.releaseSession(s1); // back to idle, inUse=false
+ const s2 = await pool.getSession(); // idle-reuse branch
+
+ expect(s2).toBe(s1);
+ // The reused session is handed to a caller, so it must be inUse; the
+ // idle-reuse branch used to skip this, leaving it false while active.
+ expect(pool.pooledFor(s2).inUse).toBe(true);
+
+ await pool.close();
+ });
+
+ it("cleanupIdleSessions never shrinks the pool below minPoolSize", async () => {
+ const pool = newPool({ maxPoolSize: 5, minPoolSize: 1, maxIdleTime: 1 });
+
+ const a = await pool.getSession();
+ const b = await pool.getSession();
+ const c = await pool.getSession();
+ pool.releaseSession(a);
+ pool.releaseSession(b);
+ pool.releaseSession(c);
+ // Make every session look long-idle so all three qualify for cleanup.
+ for (const ps of (pool as any).pool) {
+ ps.lastUsed = 0;
+ }
+
+ await pool.runCleanup();
+
+ // Must retain minPoolSize; the buggy guard (constant pre-cleanup size)
+ // removed all three and collapsed the pool to 0.
+ expect(pool.getPoolSize()).toBe(1);
+
+ await pool.close();
+ });
+
+ it("cleanupIdleSessions does not hand out a session that is being closed", async () => {
+ // minPoolSize=1 (0 would coerce to 1 anyway), 2 idle sessions so cleanup
+ // removes exactly one (the first-queued, s1) and keeps one warm.
+ const pool = newPool({ maxPoolSize: 3, minPoolSize: 1, maxIdleTime: 1 });
+
+ const s1 = await pool.getSession();
+ const s2 = await pool.getSession();
+ pool.releaseSession(s1); // idle=[s1]
+ pool.releaseSession(s2); // idle=[s1, s2]
+ for (const ps of (pool as any).pool) {
+ ps.lastUsed = 0; // both qualify as long-idle
+ }
+
+ // Gate s1's close so it stays "closing" (isOpen()===true) across an await.
+ let openClose!: () => void;
+ (s1 as any).closeGate = new Promise<void>((r) => {
+ openClose = r;
+ });
+
+ const cleanup = pool.runCleanup(); // removes s1 (down to minSize=1), gated close
+ await new Promise((r) => setImmediate(r)); // let cleanup reach the close await
+
+ // A concurrent acquire must NOT receive the session being closed. On the
+ // buggy close-then-splice, s1 stays in idle during close() and shift()
+ // hands it out.
+ const acquired = await pool.getSession();
+ expect(acquired).not.toBe(s1);
+
+ openClose(); // let the close finish
+ await cleanup;
+
+ await pool.close();
+ });
+});