chore: optimize getHostInfoByStrategy in PluginManager (#332)

aws · Dec 10, 2024 · 02bf0c9 · 02bf0c9
1 parent 95f1b6f
commit 02bf0c9
Showing 1 changed file with 51 additions and 13 deletions.
diff --git a/common/lib/plugin_manager.ts b/common/lib/plugin_manager.ts
@@ -63,6 +63,7 @@ class PluginChain<T> {
 
 export class PluginManager {
   private static readonly PLUGIN_CHAIN_CACHE = new Map<[string, HostInfo], PluginChain<any>>();
+  private static readonly STRATEGY_PLUGIN_CHAIN_CACHE = new Map<ConnectionPlugin[], Set<ConnectionPlugin>>();
   private static readonly ALL_METHODS: string = "*";
   private static readonly CONNECT_METHOD = "connect";
   private static readonly FORCE_CONNECT_METHOD = "forceConnect";
@@ -272,28 +273,65 @@ export class PluginManager {
   }
 
   acceptsStrategy(role: HostRole, strategy: string) {
-    for (const plugin of this._plugins) {
-      const pluginSubscribedMethods = plugin.getSubscribedMethods();
-      const isSubscribed =
-        pluginSubscribedMethods.has(PluginManager.ALL_METHODS) || pluginSubscribedMethods.has(PluginManager.ACCEPTS_STRATEGY_METHOD);
+    let chain: Set<ConnectionPlugin> = PluginManager.STRATEGY_PLUGIN_CHAIN_CACHE.get(this._plugins);
+    if (!chain) {
+      chain = new Set();
+      let acceptsStrategy: boolean = false;
+
+      for (const plugin of this._plugins) {
+        if (
+          plugin.getSubscribedMethods().has(PluginManager.ALL_METHODS) ||
+          plugin.getSubscribedMethods().has(PluginManager.ACCEPTS_STRATEGY_METHOD)
+        ) {
+          chain.add(plugin);
+          if (!acceptsStrategy && plugin.acceptsStrategy(role, strategy)) {
+            acceptsStrategy = true;
+          }
+        }
+      }
 
-      if (isSubscribed && plugin.acceptsStrategy(role, strategy)) {
-        return true;
+      PluginManager.STRATEGY_PLUGIN_CHAIN_CACHE.set(this._plugins, chain);
+      return acceptsStrategy;
+    } else {
+      for (const plugin of chain) {
+        if (plugin.acceptsStrategy(role, strategy)) {
+          return true;
+        }
       }
     }
 
     return false;
   }
 
   getHostInfoByStrategy(role: HostRole, strategy: string, hosts?: HostInfo[]): HostInfo {
-    for (const plugin of this._plugins) {
-      const pluginSubscribedMethods = plugin.getSubscribedMethods();
-      const isSubscribed =
-        pluginSubscribedMethods.has(PluginManager.ALL_METHODS) || pluginSubscribedMethods.has(PluginManager.GET_HOST_INFO_BY_STRATEGY_METHOD);
-
-      if (isSubscribed) {
+    let chain: Set<ConnectionPlugin> = PluginManager.STRATEGY_PLUGIN_CHAIN_CACHE.get(this._plugins);
+    if (!chain) {
+      chain = new Set();
+      let host: HostInfo;
+
+      for (const plugin of this._plugins) {
+        if (
+          plugin.getSubscribedMethods().has(PluginManager.ALL_METHODS) ||
+          plugin.getSubscribedMethods().has(PluginManager.GET_HOST_INFO_BY_STRATEGY_METHOD)
+        ) {
+          chain.add(plugin);
+          if (!host) {
+            try {
+              host = plugin.getHostInfoByStrategy(role, strategy, hosts);
+            } catch (error) {
+              // This plugin does not support the provided strategy, ignore the exception and move on
+            }
+          }
+        }
+      }
+      PluginManager.STRATEGY_PLUGIN_CHAIN_CACHE.set(this._plugins, chain);
+      if (host) {
+        return host;
+      }
+    } else {
+      for (const plugin of chain) {
         try {
-          const host = plugin.getHostInfoByStrategy(role, strategy, hosts);
+          const host: HostInfo = plugin.getHostInfoByStrategy(role, strategy, hosts);
           if (host) {
             return host;
           }