chore: storage service refactoring (#611)

Author: karenc-bq

.github/workflows/integration_tests.yml

@@ -14,6 +14,9 @@ on:
       - "docs/**"
       - "ISSUE_TEMPLATE/**"
       - "**/remove-old-artifacts.yml"
+  pull_request:
+    branches:
+      - dev/v3
 
 permissions:
   id-token: write   # This is required for requesting the JWT

.github/workflows/main.yml

@@ -4,10 +4,10 @@ on:
   push:
     branches:
       - main
-      - dev/v3
   pull_request:
     branches:
       - "*"
+      - dev/v3
 
 permissions:
   contents: read

common/lib/aws_client.ts

@@ -35,11 +35,14 @@ import { Messages } from "./utils/messages";
 import { HostListProviderService } from "./host_list_provider_service";
 import { SessionStateClient } from "./session_state_client";
 import { DriverConnectionProvider } from "./driver_connection_provider";
+import { StorageService } from "./utils/storage/storage_service";
+import { CoreServicesContainer } from "./utils/core_services_container";
 
 const { EventEmitter } = pkgStream;
 
 export abstract class AwsClient extends EventEmitter implements SessionStateClient {
   private _defaultPort: number = -1;
+  private readonly storageService: StorageService;
   protected telemetryFactory: TelemetryFactory;
   protected pluginManager: PluginManager;
   protected pluginService: PluginService;
@@ -64,6 +67,8 @@ export abstract class AwsClient extends EventEmitter implements SessionStateClie
 
     this.properties = new Map<string, any>(Object.entries(config));
 
+    this.storageService = CoreServicesContainer.getInstance().getStorageService();
+
     const profileName = WrapperProperties.PROFILE_NAME.get(this.properties);
     if (profileName && profileName.length > 0) {
       this._configurationProfile = DriverConfigurationProfiles.getProfileConfiguration(profileName);

common/lib/host_list_provider/monitoring/cluster_topology_monitor.ts

@@ -25,6 +25,9 @@ import { ClientWrapper } from "../../client_wrapper";
 import { AwsWrapperError } from "../../utils/errors";
 import { MonitoringRdsHostListProvider } from "./monitoring_host_list_provider";
 import { Messages } from "../../utils/messages";
+import { CoreServicesContainer } from "../../utils/core_services_container";
+import { Topology } from "../topology";
+import { StorageService, StorageServiceImpl } from "../../utils/storage/storage_service";
 
 export interface ClusterTopologyMonitor {
   forceRefresh(client: ClientWrapper, timeoutMs: number): Promise<HostInfo[]>;
@@ -45,8 +48,8 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
   private readonly _hostListProvider: MonitoringRdsHostListProvider;
   private readonly refreshRateMs: number;
   private readonly highRefreshRateMs: number;
+  private readonly storageService: StorageService;
 
-  private topologyMap: CacheMap<string, HostInfo[]>;
   private writerHostInfo: HostInfo = null;
   private isVerifiedWriterConnection: boolean = false;
   private monitoringClient: ClientWrapper = null;
@@ -72,7 +75,6 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
 
   constructor(
     clusterId: string,
-    topologyMap: CacheMap<string, HostInfo[]>,
     initialHostInfo: HostInfo,
     props: Map<string, any>,
     pluginService: PluginService,
@@ -81,7 +83,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
     highRefreshRateMs: number
   ) {
     this.clusterId = clusterId;
-    this.topologyMap = topologyMap;
+    this.storageService = CoreServicesContainer.getInstance().getStorageService(); // TODO: store serviceContainer instead
     this.initialHostInfo = initialHostInfo;
     this._pluginService = pluginService;
     this._hostListProvider = hostListProvider;
@@ -125,7 +127,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
   async forceMonitoringRefresh(shouldVerifyWriter: boolean, timeoutMs: number): Promise<HostInfo[] | null> {
     if (Date.now() < this.ignoreNewTopologyRequestsEndTimeMs) {
       // Previous failover has just completed, use results without triggering new update.
-      const currentHosts = this.topologyMap.get(this.clusterId);
+      const currentHosts = this.getStoredHosts();
       if (currentHosts !== null) {
         logger.info(Messages.get("ClusterTopologyMonitoring.ignoringNewTopologyRequest"));
         return currentHosts;
@@ -158,7 +160,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
     // Signal to any monitor that might be in delay, that topology should be updated.
     this.requestToUpdateTopology = true;
 
-    const currentHosts: HostInfo[] = this.topologyMap.get(this.clusterId);
+    const currentHosts: HostInfo[] = this.getStoredHosts();
 
     if (timeoutMs === 0) {
       logger.info(logTopology(currentHosts, Messages.get("ClusterTopologyMonitoring.timeoutSetToZero")));
@@ -168,7 +170,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
     const endTime = Date.now() + timeoutMs;
     let latestHosts: HostInfo[];
 
-    while ((latestHosts = this.topologyMap.get(this.clusterId)) === currentHosts && Date.now() < endTime) {
+    while ((latestHosts = this.getStoredHosts()) === currentHosts && Date.now() < endTime) {
       await sleep(1000);
     }
 
@@ -244,7 +246,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
   }
 
   updateTopologyCache(hosts: HostInfo[]): void {
-    this.topologyMap.put(this.clusterId, hosts, ClusterTopologyMonitorImpl.TOPOLOGY_CACHE_EXPIRATION_NANOS);
+    this.storageService.set(this.clusterId, new Topology(hosts));
     this.requestToUpdateTopology = false;
   }
 
@@ -296,7 +298,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
             this.hostMonitorsLatestTopology = [];
 
             // Use any client to gather topology information.
-            let hosts: HostInfo[] = this.topologyMap.get(this.clusterId);
+            let hosts: HostInfo[] = this.getStoredHosts();
             if (!hosts) {
               hosts = await this.openAnyClientAndUpdateTopology();
             }
@@ -392,6 +394,11 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
     }
   }
 
+  private getStoredHosts(): HostInfo[] | null {
+    const topology = this.storageService.get(Topology, this.clusterId);
+    return topology == null ? null : topology.hosts;
+  }
+
   private async delay(useHighRefreshRate: boolean) {
     if (Date.now() < this.highRefreshRateEndTimeMs) {
       useHighRefreshRate = true;
@@ -404,7 +411,7 @@ export class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
   }
 
   logTopology(msgPrefix: string) {
-    const hosts: HostInfo[] = this.topologyMap.get(this.clusterId);
+    const hosts: HostInfo[] = this.getStoredHosts();
     if (hosts && hosts.length !== 0) {
       logger.debug(logTopology(hosts, msgPrefix));
     }

common/lib/host_list_provider/monitoring/monitoring_host_list_provider.ts

@@ -89,7 +89,6 @@ export class MonitoringRdsHostListProvider extends RdsHostListProvider implement
       () =>
         new ClusterTopologyMonitorImpl(
           this.clusterId,
-          MonitoringRdsHostListProvider.topologyCache,
           this.initialHost,
           this.properties,
           this.pluginService,

common/lib/host_list_provider/rds_host_list_provider.ts

@@ -30,10 +30,15 @@ import { CacheMap } from "../utils/cache_map";
 import { isDialectTopologyAware, logTopology } from "../utils/utils";
 import { DatabaseDialect } from "../database_dialect/database_dialect";
 import { ClientWrapper } from "../client_wrapper";
+import { CoreServicesContainer } from "../utils/core_services_container";
+import { StorageService } from "../utils/storage/storage_service";
+import { Topology } from "./topology";
+import { ExpirationCache } from "../utils/storage/expiration_cache";
 
 export class RdsHostListProvider implements DynamicHostListProvider {
   private readonly originalUrl: string;
   private readonly rdsHelper: RdsUtils;
+  private readonly storageService: StorageService;
   protected readonly properties: Map<string, any>;
   private rdsUrlType: RdsUrlType;
   private initialHostList: HostInfo[];
@@ -45,9 +50,7 @@ export class RdsHostListProvider implements DynamicHostListProvider {
   protected readonly hostListProviderService: HostListProviderService;
 
   public static readonly suggestedPrimaryClusterIdCache: CacheMap<string, string> = new CacheMap<string, string>();
-
   public static readonly primaryClusterIdCache: CacheMap<string, boolean> = new CacheMap<string, boolean>();
-  public static readonly topologyCache: CacheMap<string, HostInfo[]> = new CacheMap<string, HostInfo[]>();
   public clusterId: string = Date.now().toString();
   public isInitialized: boolean = false;
   public isPrimaryClusterId?: boolean;
@@ -59,6 +62,7 @@ export class RdsHostListProvider implements DynamicHostListProvider {
     this.connectionUrlParser = hostListProviderService.getConnectionUrlParser();
     this.originalUrl = originalUrl;
     this.properties = properties;
+    this.storageService = CoreServicesContainer.getInstance().getStorageService(); // TODO: store the service container instead.
 
     let port = WrapperProperties.PORT.get(properties);
     if (port == null) {
@@ -197,7 +201,7 @@ export class RdsHostListProvider implements DynamicHostListProvider {
       this.isPrimaryClusterId = true;
     }
 
-    const cachedHosts: HostInfo[] | null = RdsHostListProvider.topologyCache.get(this.clusterId);
+    const cachedHosts: HostInfo[] | null = this.getStoredTopology();
 
     // This clusterId is a primary one and is about to create a new entry in the cache.
     // When a primary entry is created it needs to be suggested for other (non-primary) entries.
@@ -211,7 +215,7 @@ export class RdsHostListProvider implements DynamicHostListProvider {
 
       const hosts = await this.queryForTopology(targetClient, this.hostListProviderService.getDialect());
       if (hosts && hosts.length > 0) {
-        RdsHostListProvider.topologyCache.put(this.clusterId, hosts, this.refreshRateNano);
+        this.storageService.set(this.clusterId, new Topology(hosts));
         if (needToSuggest) {
           this.suggestPrimaryCluster(hosts);
         }
@@ -227,14 +231,18 @@ export class RdsHostListProvider implements DynamicHostListProvider {
   }
 
   private getSuggestedClusterId(hostAndPort: string): ClusterSuggestedResult | null {
-    for (const [key, hosts] of RdsHostListProvider.topologyCache.getEntries()) {
+    const cache: ExpirationCache<string, Topology> = this.storageService.getAll(Topology) as ExpirationCache<string, Topology>;
+    if (!cache) {
+      return null;
+    }
+    for (const [key, hosts] of cache.getEntries()) {
       const isPrimaryCluster: boolean = RdsHostListProvider.primaryClusterIdCache.get(key, false, this.suggestedClusterIdRefreshRateNano) ?? false;
       if (key === hostAndPort) {
         return new ClusterSuggestedResult(hostAndPort, isPrimaryCluster);
       }
 
       if (hosts) {
-        for (const hostInfo of hosts) {
+        for (const hostInfo of hosts.hosts) {
           if (hostInfo.hostAndPort === hostAndPort) {
             logger.debug(Messages.get("RdsHostListProvider.suggestedClusterId", key, hostAndPort));
             return new ClusterSuggestedResult(key, isPrimaryCluster);
@@ -255,7 +263,11 @@ export class RdsHostListProvider implements DynamicHostListProvider {
       primaryClusterHostUrls.add(hostInfo.url);
     });
 
-    for (const [clusterId, clusterHosts] of RdsHostListProvider.topologyCache.getEntries()) {
+    const cache: ExpirationCache<string, Topology> = this.storageService.getAll(Topology) as ExpirationCache<string, Topology>;
+    if (!cache) {
+      return;
+    }
+    for (const [clusterId, clusterHosts] of cache.getEntries()) {
       const isPrimaryCluster: boolean | null = RdsHostListProvider.primaryClusterIdCache.get(
         clusterId,
         false,
@@ -266,7 +278,7 @@ export class RdsHostListProvider implements DynamicHostListProvider {
         continue;
       }
 
-      for (const clusterHost of clusterHosts) {
+      for (const clusterHost of clusterHosts.hosts) {
         if (primaryClusterHostUrls.has(clusterHost.url)) {
           RdsHostListProvider.suggestedPrimaryClusterIdCache.put(clusterId, this.clusterId, this.suggestedClusterIdRefreshRateNano);
           break;
@@ -343,22 +355,24 @@ export class RdsHostListProvider implements DynamicHostListProvider {
     return host.replace("?", hostName);
   }
 
-  getCachedTopology(): HostInfo[] | null {
+  getStoredTopology(): HostInfo[] | null {
     if (!this.clusterId) {
       return null;
     }
-    return RdsHostListProvider.topologyCache.get(this.clusterId) ?? null;
+
+    const topology: Topology = this.storageService.get(Topology, this.clusterId);
+
+    return topology == null ? null : topology.hosts;
   }
 
   static clearAll(): void {
-    RdsHostListProvider.topologyCache.clear();
     RdsHostListProvider.primaryClusterIdCache.clear();
     RdsHostListProvider.suggestedPrimaryClusterIdCache.clear();
   }
 
   clear(): void {
     if (this.clusterId) {
-      RdsHostListProvider.topologyCache.delete(this.clusterId);
+      CoreServicesContainer.getInstance().getStorageService().remove(Topology, this.clusterId);
     }
   }
 

common/lib/host_list_provider/topology.ts

@@ -0,0 +1,29 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed 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 { HostInfo } from "../host_info";
+
+export class Topology {
+  private readonly _hosts: HostInfo[];
+
+  constructor(hosts: HostInfo[]) {
+    this._hosts = hosts;
+  }
+
+  get hosts(): HostInfo[] {
+    return this._hosts;
+  }
+}

common/lib/types.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed 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.
+ */
+
+/**
+ * Type representing a constructor for any class.
+ */
+export type Constructor<T = unknown> = new (...args: unknown[]) => T;
+
+/**
+ * Function type that determines whether an item should be disposed when expired.
+ * @param item The item to check
+ * @returns true if the item should be disposed, false otherwise
+ */
+export type ShouldDisposeFunc<V> = (item: V) => boolean;
+
+/**
+ * Function type that defines how to dispose of an item when it is removed.
+ * @param item The item to dispose
+ */
+export type ItemDisposalFunc<V> = (item: V) => void;

common/lib/utils/core_services_container.ts

@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed 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 { StorageService, StorageServiceImpl } from "./storage/storage_service";
+
+/**
+ * A singleton container object used to instantiate and access core universal services. This class should be used
+ * instead of directly instantiating core services so that only one instance of each service is instantiated.
+ *
+ * @see FullServicesContainer for a container that holds both connection-specific services and core universal
+ *     services.
+ */
+export class CoreServicesContainer {
+  private static readonly INSTANCE = new CoreServicesContainer();
+
+  // private readonly monitorService: MonitorService; // TODO: implement monitor service
+  private readonly storageService: StorageService;
+
+  private constructor() {
+    this.storageService = new StorageServiceImpl();
+    // this.monitorService = new MonitorServiceImpl();
+  }
+
+  static getInstance(): CoreServicesContainer {
+    return CoreServicesContainer.INSTANCE;
+  }
+
+  getStorageService(): StorageService {
+    return this.storageService;
+  }
+
+  // getMonitorService(): MonitorService {
+  //   return this.monitorService;
+  // }
+}