diff --git a/src/client/BaseSessionPool.ts b/src/client/BaseSessionPool.ts index 8580daf..cfd0ff0 100644 --- a/src/client/BaseSessionPool.ts +++ b/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 = new Denque(); protected activeSessions: Set = 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); diff --git a/tests/unit/BaseSessionPoolLifecycle.test.ts b/tests/unit/BaseSessionPoolLifecycle.test.ts new file mode 100644 index 0000000..64bf2e6 --- /dev/null +++ b/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 | null = null; + private counter = 0; + + protected async createPoolSession(): Promise { + 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 | 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 { + 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).map( + (ps: any) => ps.session, + ); + } +} + +function newPool(overrides: Record): 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((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((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(); + }); +});