diff --git a/handwritten/firestore/dev/src/pool.ts b/handwritten/firestore/dev/src/pool.ts index 5bb4550171b1..2e64f3c8dec6 100644 --- a/handwritten/firestore/dev/src/pool.ts +++ b/handwritten/firestore/dev/src/pool.ts @@ -107,6 +107,34 @@ export class ClientPool { }); } + /** + * Transitions the pool to gRPC and eagerly garbage-collects idle REST clients. + * REST clients with in-flight requests will be garbage collected upon request completion. + */ + private transitionToGrpc(requestTag: string): void { + this.grpcEnabled = true; + logger( + `ClientPool[${this.instanceId}].acquire`, + requestTag, + 'Transitioning pool to gRPC (requiresGrpc: true)', + ); + + for (const [client, metadata] of this.activeClients) { + if (!metadata.grpcEnabled && metadata.activeRequestCount === 0) { + this.activeClients.delete(client); + this.failedClients.delete(client); + void Promise.resolve(this.clientDestructor(client)).catch(err => { + logger( + `ClientPool[${this.instanceId}].transitionToGrpc`, + requestTag, + 'Failed to destroy client: %s', + err, + ); + }); + } + } + } + /** * Returns an already existing client if it has less than the maximum number * of concurrent operations or initializes and returns a new client. @@ -118,21 +146,25 @@ export class ClientPool { let selectedClient: T | null = null; let selectedClientRequestCount = -1; - // Transition to grpc when we see the first operation that requires grpc. - this.grpcEnabled = this.grpcEnabled || requiresGrpc; + if (!this.grpcEnabled && requiresGrpc) { + this.transitionToGrpc(requestTag); + } // Require a grpc client for this operation if we have transitioned to grpc. requiresGrpc = requiresGrpc || this.grpcEnabled; for (const [client, metadata] of this.activeClients) { + const isEligible = metadata.grpcEnabled || !requiresGrpc; + if (!isEligible || this.failedClients.has(client)) { + continue; + } + // Use the "most-full" client that can still accommodate the request // in order to maximize the number of idle clients as operations start to // complete. if ( - !this.failedClients.has(client) && metadata.activeRequestCount > selectedClientRequestCount && - metadata.activeRequestCount < this.concurrentOperationLimit && - (metadata.grpcEnabled || !requiresGrpc) + metadata.activeRequestCount < this.concurrentOperationLimit ) { selectedClient = client; selectedClientRequestCount = metadata.activeRequestCount; @@ -270,8 +302,10 @@ export class ClientPool { // more than 100 idle capacity with default settings). let idleCapacityCount = 0; for (const [, metadata] of this.activeClients) { - idleCapacityCount += - this.concurrentOperationLimit - metadata.activeRequestCount; + if (metadata.grpcEnabled === this.grpcEnabled) { + idleCapacityCount += + this.concurrentOperationLimit - metadata.activeRequestCount; + } } const maxIdleCapacityCount = diff --git a/handwritten/firestore/dev/test/pool.ts b/handwritten/firestore/dev/test/pool.ts index 3a7b188a5605..58a0786f917a 100644 --- a/handwritten/firestore/dev/test/pool.ts +++ b/handwritten/firestore/dev/test/pool.ts @@ -19,6 +19,7 @@ import * as chaiAsPromised from 'chai-as-promised'; import {ClientPool, CLIENT_TERMINATED_ERROR_MSG} from '../src/pool'; import {Deferred} from '../src/util'; +import {setLogFunction, setLibVersion} from '../src/logger'; use(chaiAsPromised); @@ -575,6 +576,106 @@ describe('Client pool', () => { expect(clientPool.size).to.equal(1); }); + it('garbage collects idle REST clients upon gRPC transition and reuses gRPC client', async () => { + let clientCount = 0; + const clientPool = new ClientPool<{}>(100, 1, () => { + ++clientCount; + return {}; + }); + + // Run REST operation + await clientPool.run(REQUEST_TAG, USE_REST, () => Promise.resolve()); + expect(clientCount).to.equal(1); + + // Run gRPC operation + await clientPool.run(REQUEST_TAG, USE_GRPC, () => Promise.resolve()); + expect(clientCount).to.equal(2); + + // Run subsequent operations. + await clientPool.run(REQUEST_TAG, USE_REST, () => Promise.resolve()); + await clientPool.run(REQUEST_TAG, USE_REST, () => Promise.resolve()); + + // Assert that the active client in the pool is gRPC enabled, and that no extra clients were created. + expect( + clientPool.size, + 'Pool size should equal 1 after eager eviction of REST client', + ).to.equal(1); + const activeClientMetadata = Array.from(clientPool._activeClients.values()); + const gRPCClientRemaining = activeClientMetadata.some(m => m.grpcEnabled); + + expect( + gRPCClientRemaining, + 'The active client in pool must be gRPC-enabled', + ).to.be.true; + expect( + clientCount, + 'Subsequent operations should reuse the gRPC client without creating new clients', + ).to.equal(2); + }); + + it('garbage collects active REST clients upon release after gRPC transition', async () => { + let createdCount = 0; + const destroyedClientIds: string[] = []; + + const clientPool = new ClientPool<{id: string}>( + 100, // concurrentOperationLimit + 1, // maxIdleClients + requiresGrpc => { + createdCount++; + return {id: `client-${createdCount}`}; + }, + async client => { + destroyedClientIds.push(client.id); + }, + ); + + const restDeferred = new Deferred(); + + // Run REST operation that stays active (client-1 active) + const restOpPromise = clientPool.run( + REQUEST_TAG, + USE_REST, + () => restDeferred.promise, + ); + expect(createdCount).to.equal(1); + expect(clientPool.size).to.equal(1); + + // Run gRPC operation (transitions pool to gRPC, creates client-2) + // client-1 is active so it is NOT eagerly evicted during transition. + await clientPool.run(REQUEST_TAG, USE_GRPC, () => Promise.resolve()); + expect(createdCount).to.equal(2); + expect(clientPool.size).to.equal(2); + + // Resolve the active REST operation (triggers release for client-1) + restDeferred.resolve(); + await restOpPromise; + + // Assert that client-1 (REST client) is destroyed immediately upon completion + expect(destroyedClientIds).to.include('client-1'); + expect(clientPool.size).to.equal(1); + + const activeClients = Array.from(clientPool._activeClients.keys()); + expect(activeClients.map(c => c.id)).to.not.include('client-1'); + expect(activeClients.map(c => c.id)).to.include('client-2'); + }); + + it('logs transition to gRPC', async () => { + const logs: string[] = []; + setLogFunction(msg => logs.push(msg)); + + try { + const clientPool = new ClientPool<{}>(100, 1, () => ({})); + + await clientPool.run('op-rest', USE_REST, () => Promise.resolve()); + await clientPool.run('op-grpc', USE_GRPC, () => Promise.resolve()); + + expect(logs.some(l => l.includes('Transitioning pool to gRPC'))).to.be + .true; + } finally { + setLogFunction(null); + } + }); + it('keeps pool of idle clients', async () => { const clientPool = new ClientPool<{}>( /* concurrentOperationLimit= */ 1,