Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
48 changes: 41 additions & 7 deletions handwritten/firestore/dev/src/pool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,34 @@ export class ClientPool<T extends object> {
});
}

/**
* 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,
);
});
Comment on lines +126 to +133

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.

high

The clientDestructor parameter is optional in the ClientPool constructor and can be undefined (as demonstrated in some of the unit tests where only 3 arguments are passed). Calling this.clientDestructor(client) directly without checking if it is defined will throw a TypeError at runtime and crash the pool transition.

Additionally, wrapping the call in Promise.resolve(this.clientDestructor(client)) does not catch synchronous exceptions thrown by the destructor function itself, which would also propagate and crash the caller.

To fix both issues, we should check if this.clientDestructor is defined, and safely invoke it within a .then() block of a resolved promise chain to ensure any synchronous errors are safely caught and logged.

        if (this.clientDestructor) {
          const destructor = this.clientDestructor;
          void Promise.resolve()
            .then(() => destructor(client))
            .catch(err => {
              logger(
                'ClientPool[' + this.instanceId + '].transitionToGrpc',
                requestTag,
                'Failed to destroy client: %s',
                err,
              );
            });
        }

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.

This seems to be a false positive. The clientDestructor is not optional in the ClientPool constructor.

}
}
}

/**
* Returns an already existing client if it has less than the maximum number
* of concurrent operations or initializes and returns a new client.
Expand All @@ -118,21 +146,25 @@ export class ClientPool<T extends object> {
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;
Expand Down Expand Up @@ -270,8 +302,10 @@ export class ClientPool<T extends object> {
// 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) {

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.

high

Active REST Clients Leak After Transition

There is a potential connection/resource leak for REST clients that are active during the transition to gRPC.

The PR description notes:

REST clients with active requests during transition will complete and are garbage collected on release(). They are not eagerly evicted.

However, under the new logic in shouldGarbageCollectClient, when an active REST client completes and release() is called, shouldGarbageCollectClient is invoked. Inside this method, idleCapacityCount only sums the idle capacity of clients matching this.grpcEnabled (which is now true for gRPC).

Since the REST client's metadata.grpcEnabled is false, its idle capacity is not counted. If the total idle capacity of the gRPC clients does not exceed maxIdleCapacityCount, shouldGarbageCollectClient will return false for the REST client.

As a result, the REST client will not be garbage collected and will remain in activeClients indefinitely as an idle zombie, because it can never be acquired again.

Suggested Fix

To prevent this leak, we should immediately return true (garbage collect) if the client being evaluated does not match the pool's current grpcEnabled status.

You can add this check at the very beginning of shouldGarbageCollectClient (which is outside the current diff hunk):

private shouldGarbageCollectClient(client: T): boolean {
  const clientMetadata = this.activeClients.get(client);
  if (clientMetadata && clientMetadata.grpcEnabled !== this.grpcEnabled) {
    return true;
  }
  // ... rest of the existing logic
}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

The idleCapacityCount check only happens after we check

if (this.grpcEnabled !== clientMetadata.grpcEnabled)

which is true since this.grpcEnabled is true, and clientMetadata.grpcEnabled would be false- in which case we immediately return PoolIsTransitioningFromGrpc(), which has shouldGarbageCollectClient: true.

I added a test to verify that the rest connection is garbage collected.

idleCapacityCount +=
this.concurrentOperationLimit - metadata.activeRequestCount;
}
}

const maxIdleCapacityCount =
Expand Down
101 changes: 101 additions & 0 deletions handwritten/firestore/dev/test/pool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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<void>();

// 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,
Expand Down
Loading