feat: connection pool (#533)

Author: karenc-bq

common/lib/AllowedAndBlockedHosts.ts common/lib/allowed_and_blocked_hosts.tscommon/lib/AllowedAndBlockedHosts.ts renamed to common/lib/allowed_and_blocked_hosts.ts

@@ -22,7 +22,6 @@ import { HostListProvider } from "./host_list_provider/host_list_provider";
 import { PluginManager } from "./plugin_manager";
 
 import pkgStream from "stream";
-import { DriverConnectionProvider } from "./driver_connection_provider";
 import { ClientWrapper } from "./client_wrapper";
 import { ConnectionProviderManager } from "./connection_provider_manager";
 import { DefaultTelemetryFactory } from "./utils/telemetry/default_telemetry_factory";
@@ -35,10 +34,13 @@ import { AwsWrapperError } from "./utils/errors";
 import { Messages } from "./utils/messages";
 import { TransactionIsolationLevel } from "./utils/transaction_isolation_level";
 import { HostListProviderService } from "./host_list_provider_service";
+import { SessionStateClient } from "./session_state_client";
+import { ConnectionProvider } from "./connection_provider";
+import { DriverConnectionProvider } from "./driver_connection_provider";
 
 const { EventEmitter } = pkgStream;
 
-export abstract class AwsClient extends EventEmitter {
+export abstract class AwsClient extends EventEmitter implements SessionStateClient {
   private _defaultPort: number = -1;
   protected telemetryFactory: TelemetryFactory;
   protected pluginManager: PluginManager;
@@ -55,7 +57,8 @@ export abstract class AwsClient extends EventEmitter {
     dbType: DatabaseType,
     knownDialectsByCode: Map<string, DatabaseDialect>,
     parser: ConnectionUrlParser,
-    driverDialect: DriverDialect
+    driverDialect: DriverDialect,
+    connectionProvider?: ConnectionProvider
   ) {
     super();
     this.config = config;
@@ -110,7 +113,7 @@ export abstract class AwsClient extends EventEmitter {
     this.pluginManager = new PluginManager(
       container,
       this.properties,
-      new ConnectionProviderManager(new DriverConnectionProvider(), WrapperProperties.CONNECTION_PROVIDER.get(this.properties)),
+      new ConnectionProviderManager(connectionProvider ?? new DriverConnectionProvider(), WrapperProperties.CONNECTION_PROVIDER.get(this.properties)),
       this.telemetryFactory
     );
   }

common/lib/aws_client.ts

@@ -14,7 +14,7 @@
   limitations under the License.
 */
 
-export interface AwsPoolClient {
+export interface AwsInternalPoolClient {
   connect(): Promise<any>;
 
   end(): Promise<any>;

common/lib/aws_pool_client.ts

@@ -30,6 +30,7 @@ export class AwsPoolConfig {
    */
   readonly minConnections?: number | undefined;
   readonly allowExitOnIdle?: boolean | undefined;
+  readonly maxLifetimeSeconds?: number | undefined;
 
   constructor(props?: any) {
     this.maxConnections = props.maxConnections ?? 10;
@@ -39,5 +40,6 @@ export class AwsPoolConfig {
     this.queueLimit = props.queueLimit ?? 0;
     this.minConnections = props.minConnections ?? 0;
     this.allowExitOnIdle = props.allowExitOnIdle ?? false;
+    this.maxLifetimeSeconds = props.maxLifetimeSeconds ?? 0;
   }
 }

common/lib/aws_pool_config.ts

@@ -22,7 +22,9 @@ export interface ClientWrapper {
   readonly properties: Map<string, any>;
   readonly id: string;
 
+  // Internal method, executes wrapper-specific queries like the topology query.
   query(sql: string): Promise<any>;
+  query(config: any, values?: any, callback?: any): Promise<any>;
 
   end(): Promise<void>;
 

common/lib/client_wrapper.ts

@@ -17,14 +17,18 @@
 import { ConnectionProvider } from "./connection_provider";
 import { HostRole } from "./host_role";
 import { HostInfo } from "./host_info";
+import { CanReleaseResources } from "./can_release_resources";
+import { DriverConnectionProvider } from "./driver_connection_provider";
 
 export class ConnectionProviderManager {
+  private readonly forceConnectProvider: ConnectionProvider;
   private readonly defaultProvider: ConnectionProvider;
   private readonly effectiveProvider: ConnectionProvider | null;
 
   constructor(defaultProvider: ConnectionProvider, effectiveProvider: ConnectionProvider | null) {
     this.defaultProvider = defaultProvider;
     this.effectiveProvider = effectiveProvider;
+    this.forceConnectProvider = new DriverConnectionProvider();
   }
 
   getConnectionProvider(hostInfo: HostInfo | null, props: Map<string, any>): ConnectionProvider {
@@ -64,4 +68,22 @@ export class ConnectionProviderManager {
   getDefaultConnectionProvider(): ConnectionProvider {
     return this.defaultProvider;
   }
+
+  getForceConnectionProvider(): ConnectionProvider {
+    return this.forceConnectProvider;
+  }
+
+  private static implementsCanReleaseResources(provider: any): provider is CanReleaseResources {
+    return provider.releaseResources !== undefined;
+  }
+
+  async releaseResources(): Promise<void> {
+    if (this.effectiveProvider && ConnectionProviderManager.implementsCanReleaseResources(this.effectiveProvider)) {
+      await this.effectiveProvider.releaseResources();
+    }
+
+    if (ConnectionProviderManager.implementsCanReleaseResources(this.defaultProvider)) {
+      await this.defaultProvider.releaseResources();
+    }
+  }
 }

common/lib/connection_provider_manager.ts

@@ -16,7 +16,7 @@
 
 import { ClientWrapper } from "../client_wrapper";
 import { AwsPoolConfig } from "../aws_pool_config";
-import { AwsPoolClient } from "../aws_pool_client";
+import { AwsInternalPoolClient } from "../aws_pool_client";
 import { HostInfo } from "../host_info";
 
 export interface DriverDialect {
@@ -26,11 +26,13 @@ export interface DriverDialect {
 
   preparePoolClientProperties(props: Map<string, any>, poolConfig: AwsPoolConfig | undefined): any;
 
-  getAwsPoolClient(props: any): AwsPoolClient;
+  getAwsPoolClient(props: any): AwsInternalPoolClient;
 
   setConnectTimeout(props: Map<string, any>, wrapperConnectTimeout?: any): void;
 
   setQueryTimeout(props: Map<string, any>, sql?: any, wrapperConnectTimeout?: any): void;
 
   setKeepAliveProperties(props: Map<string, any>, keepAliveProps: any): void;
+
+  getQueryFromMethodArg(methodArg: any): string;
 }

common/lib/driver_dialect/driver_dialect.ts

@@ -47,4 +47,10 @@ export interface ErrorHandler {
    * @param clientWrapper a wrapper containing the target community client.
    */
   attachNoOpErrorListener(clientWrapper: ClientWrapper | undefined): void;
+
+  /**
+   * Remove an error event listener from a ClientWrapper.
+   * @param clientWrapper a wrapper containing the target community client.
+   */
+  removeErrorListener(clientWrapper: ClientWrapper | undefined): void;
 }

common/lib/error_handler.ts

@@ -21,7 +21,7 @@ import { DatabaseDialect } from "./database_dialect/database_dialect";
 import { HostInfoBuilder } from "./host_info_builder";
 import { ConnectionUrlParser } from "./utils/connection_url_parser";
 import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
-import { AllowedAndBlockedHosts } from "./AllowedAndBlockedHosts";
+import { AllowedAndBlockedHosts } from "./allowed_and_blocked_hosts";
 
 export interface HostListProviderService {
   getHostListProvider(): HostListProvider | null;

common/lib/host_list_provider_service.ts

@@ -33,7 +33,7 @@ import { HostSelector } from "./host_selector";
 import { RandomHostSelector } from "./random_host_selector";
 import { InternalPoolMapping } from "./utils/internal_pool_mapping";
 import { RoundRobinHostSelector } from "./round_robin_host_selector";
-import { AwsPoolClient } from "./aws_pool_client";
+import { AwsInternalPoolClient } from "./aws_pool_client";
 import { AwsPoolConfig } from "./aws_pool_config";
 import { LeastConnectionsHostSelector } from "./least_connections_host_selector";
 import { PoolClientWrapper } from "./pool_client_wrapper";
@@ -59,7 +59,7 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
   private readonly _poolMapping?: InternalPoolMapping;
   private readonly _poolConfig?: AwsPoolConfig;
   targetClient?: ClientWrapper;
-  internalPool: AwsPoolClient | undefined;
+  internalPool: AwsInternalPoolClient | undefined;
 
   private static poolExpirationCheckNanos: bigint = InternalPooledConnectionProvider.POOL_EXPIRATION_NANOS; // 30 minutes
 
@@ -80,6 +80,12 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
   }
 
   async connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ClientWrapper> {
+    const resultProps = new Map(props);
+    resultProps.set(WrapperProperties.HOST.name, hostInfo.host);
+    if (hostInfo.isPortSpecified()) {
+      resultProps.set(WrapperProperties.PORT.name, hostInfo.port);
+    }
+
     let connectionHostInfo: HostInfo = hostInfo;
     if (
       WrapperProperties.ENABLE_GREEN_HOST_REPLACEMENT.get(props) &&
@@ -96,6 +102,7 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
         // Green instance DNS doesn't exist
 
         const fixedHost: string = this.rdsUtil.removeGreenInstancePrefix(hostInfo.host);
+        resultProps.set(WrapperProperties.HOST.name, fixedHost);
         connectionHostInfo = new HostInfoBuilder({
           hostAvailabilityStrategy: hostInfo.hostAvailabilityStrategy
         })
@@ -106,9 +113,9 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
     }
 
     const dialect = pluginService.getDriverDialect();
-    const preparedConfig = dialect.preparePoolClientProperties(props, this._poolConfig);
+    const preparedConfig = dialect.preparePoolClientProperties(resultProps, this._poolConfig);
     this.internalPool = InternalPooledConnectionProvider.databasePools.computeIfAbsent(
-      new PoolKey(connectionHostInfo.url, this.getPoolKey(connectionHostInfo, props)).getPoolKeyString(),
+      new PoolKey(connectionHostInfo.url, this.getPoolKey(connectionHostInfo, resultProps)).getPoolKeyString(),
       () => dialect.getAwsPoolClient(preparedConfig),
       InternalPooledConnectionProvider.poolExpirationCheckNanos
     );
@@ -123,6 +130,14 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
   }
 
   async releaseResources() {
+    if (this.internalPool) {
+      try {
+        await this.internalPool.releaseResources();
+      } catch (error) {
+        // ignore
+      }
+    }
+
     await InternalPooledConnectionProvider.databasePools.clear();
   }