Generalize polling abstraction

packages/assets-controllers/src/CurrencyRateController.ts

@@ -11,7 +11,8 @@ import type {
   NetworkClientId,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import { PollingController } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import { Mutex } from 'async-mutex';
 
 import { fetchExchangeRate as defaultFetchExchangeRate } from './crypto-compare';
@@ -82,11 +83,14 @@ const defaultState = {
  * Controller that passively polls on a set interval for an exchange rate from the current network
  * asset to the user's preferred currency.
  */
-export class CurrencyRateController extends PollingController<
-  typeof name,
-  CurrencyRateState,
-  CurrencyRateMessenger
-> {
+export class CurrencyRateController
+  extends StaticIntervalPollingController<
+    typeof name,
+    CurrencyRateState,
+    CurrencyRateMessenger
+  >
+  implements IPollingController
+{
   private readonly mutex = new Mutex();
 
   private readonly fetchExchangeRate;

packages/assets-controllers/src/NftDetectionController.ts

@@ -11,7 +11,8 @@ import type {
   NetworkState,
   NetworkClient,
 } from '@metamask/network-controller';
-import { PollingControllerV1 } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
 
@@ -147,10 +148,10 @@ export interface NftDetectionConfig extends BaseConfig {
 /**
  * Controller that passively polls on a set interval for NFT auto detection
  */
-export class NftDetectionController extends PollingControllerV1<
-  NftDetectionConfig,
-  BaseState
-> {
+export class NftDetectionController
+  extends StaticIntervalPollingControllerV1<NftDetectionConfig, BaseState>
+  implements IPollingController
+{
   private intervalId?: ReturnType<typeof setTimeout>;
 
   private getOwnerNftApi({

packages/assets-controllers/src/TokenDetectionController.ts

@@ -8,7 +8,8 @@ import type {
   NetworkController,
   NetworkState,
 } from '@metamask/network-controller';
-import { PollingControllerV1 } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
 
@@ -44,10 +45,10 @@ export interface TokenDetectionConfig extends BaseConfig {
 /**
  * Controller that passively polls on a set interval for Tokens auto detection
  */
-export class TokenDetectionController extends PollingControllerV1<
-  TokenDetectionConfig,
-  BaseState
-> {
+export class TokenDetectionController
+  extends StaticIntervalPollingControllerV1<TokenDetectionConfig, BaseState>
+  implements IPollingController
+{
   private intervalId?: ReturnType<typeof setTimeout>;
 
   /**

packages/assets-controllers/src/TokenListController.ts

@@ -10,7 +10,8 @@ import type {
   NetworkState,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import { PollingController } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { Mutex } from 'async-mutex';
 
@@ -91,11 +92,14 @@ const defaultState: TokenListState = {
 /**
  * Controller that passively polls on a set interval for the list of tokens from metaswaps api
  */
-export class TokenListController extends PollingController<
-  typeof name,
-  TokenListState,
-  TokenListMessenger
-> {
+export class TokenListController
+  extends StaticIntervalPollingController<
+    typeof name,
+    TokenListState,
+    TokenListMessenger
+  >
+  implements IPollingController
+{
   private readonly mutex = new Mutex();
 
   private intervalId?: ReturnType<typeof setTimeout>;

packages/assets-controllers/src/TokenRatesController.ts

@@ -10,7 +10,8 @@ import type {
   NetworkController,
   NetworkState,
 } from '@metamask/network-controller';
-import { PollingControllerV1 } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
 import { isDeepStrictEqual } from 'util';
@@ -136,10 +137,10 @@ async function getCurrencyConversionRate({
  * Controller that passively polls on a set interval for token-to-fiat exchange rates
  * for tokens stored in the TokensController
  */
-export class TokenRatesController extends PollingControllerV1<
-  TokenRatesConfig,
-  TokenRatesState
-> {
+export class TokenRatesController
+  extends StaticIntervalPollingControllerV1<TokenRatesConfig, TokenRatesState>
+  implements IPollingController
+{
   private handle?: ReturnType<typeof setTimeout>;
 
   #pollState = PollState.Inactive;

packages/gas-fee-controller/src/GasFeeController.ts

@@ -18,7 +18,8 @@ import type {
   NetworkState,
   ProviderProxy,
 } from '@metamask/network-controller';
-import { PollingController } from '@metamask/polling-controller';
+import type { IPollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { v1 as random } from 'uuid';
 
@@ -253,11 +254,14 @@ const defaultState: GasFeeState = {
 /**
  * Controller that retrieves gas fee estimate data and polls for updated data on a set interval
  */
-export class GasFeeController extends PollingController<
-  typeof name,
-  GasFeeState,
-  GasFeeMessenger
-> {
+export class GasFeeController
+  extends StaticIntervalPollingController<
+    typeof name,
+    GasFeeState,
+    GasFeeMessenger
+  >
+  implements IPollingController
+{
   private intervalId?: ReturnType<typeof setTimeout>;
 
   private readonly intervalDelay;

packages/polling-controller/src/AbstractPollingController.ts

@@ -0,0 +1,138 @@
+import type { NetworkClientId } from '@metamask/network-controller';
+import type { Json } from '@metamask/utils';
+import stringify from 'fast-json-stable-stringify';
+import { v4 as random } from 'uuid';
+
+export type IPollingController = {
+  startPollingByNetworkClientId(
+    networkClientId: NetworkClientId,
+    options: Json,
+  ): string;
+
+  stopAllPolling(): void;
+
+  stopPollingByPollingToken(pollingToken: string): void;
+
+  onPollingCompleteByNetworkClientId(
+    networkClientId: NetworkClientId,
+    callback: (networkClientId: NetworkClientId) => void,
+    options: Json,
+  ): void;
+
+  _executePoll(networkClientId: NetworkClientId, options: Json): Promise<void>;
+  _startPollingByNetworkClientId(
+    networkClientId: NetworkClientId,
+    options: Json,
+  ): void;
+  _stopPollingByPollingTokenSetId(key: PollingTokenSetId): void;
+};
+
+export const getKey = (
+  networkClientId: NetworkClientId,
+  options: Json,
+): PollingTokenSetId => `${networkClientId}:${stringify(options)}`;
+
+export type PollingTokenSetId = `${NetworkClientId}:${string}`;
+
+type Constructor = new (...args: any[]) => object;
+
+/**
+ * AbstractPollingControllerBaseMixin
+ *
+ * @param Base - The base class to mix onto.
+ * @returns The composed class.
+ */
+export function AbstractPollingControllerBaseMixin<TBase extends Constructor>(
+  Base: TBase,
+) {
+  abstract class AbstractPollingControllerBase
+    extends Base
+    implements IPollingController
+  {
+    readonly #pollingTokenSets: Map<PollingTokenSetId, Set<string>> = new Map();
+
+    #callbacks: Map<
+      PollingTokenSetId,
+      Set<(PollingTokenSetId: PollingTokenSetId) => void>
+    > = new Map();
+
+    abstract _executePoll(
+      networkClientId: NetworkClientId,
+      options: Json,
+    ): Promise<void>;
+
+    abstract _startPollingByNetworkClientId(
+      networkClientId: NetworkClientId,
+      options: Json,
+    ): void;
+
+    abstract _stopPollingByPollingTokenSetId(key: PollingTokenSetId): void;
+
+    startPollingByNetworkClientId(
+      networkClientId: NetworkClientId,
+      options: Json = {},
+    ): string {
+      const pollToken = random();
+      const key = getKey(networkClientId, options);
+      const pollingTokenSet =
+        this.#pollingTokenSets.get(key) || new Set<string>();
+      pollingTokenSet.add(pollToken);
+      this.#pollingTokenSets.set(key, pollingTokenSet);
+
+      if (pollingTokenSet.size === 1) {
+        // Start new polling only if it's the first token for this key
+        this._startPollingByNetworkClientId(networkClientId, options);
+      }
+
+      return pollToken;
+    }
+
+    stopAllPolling() {
+      this.#pollingTokenSets.forEach((tokenSet, _key) => {
+        tokenSet.forEach((token) => {
+          this.stopPollingByPollingToken(token);
+        });
+      });
+    }
+
+    stopPollingByPollingToken(pollingToken: string) {
+      if (!pollingToken) {
+        throw new Error('pollingToken required');
+      }
+
+      let keyToDelete: PollingTokenSetId | null = null;
+      for (const [key, tokenSet] of this.#pollingTokenSets) {
+        if (tokenSet.delete(pollingToken)) {
+          if (tokenSet.size === 0) {
+            keyToDelete = key;
+          }
+          break;
+        }
+      }
+
+      if (keyToDelete) {
+        // TODO figure out why this is necessary
+        const nonNullKey = keyToDelete;
+        this._stopPollingByPollingTokenSetId(nonNullKey);
+        this.#pollingTokenSets.delete(nonNullKey);
+        this.#callbacks.get(nonNullKey)?.forEach((callback) => {
+          // for some reason this typescript can't tell that keyToDelete is not null here
+          callback(nonNullKey);
+        });
+        this.#callbacks.get(nonNullKey)?.clear();
+      }
+    }
+
+    onPollingCompleteByNetworkClientId(
+      networkClientId: NetworkClientId,
+      callback: (networkClientId: NetworkClientId) => void,
+      options: Json = {},
+    ) {
+      const key = getKey(networkClientId, options);
+      const callbacks = this.#callbacks.get(key) || new Set<typeof callback>();
+      callbacks.add(callback);
+      this.#callbacks.set(key, callbacks);
+    }
+  }
+  return AbstractPollingControllerBase;
+}