Generalize polling abstraction

packages/assets-controllers/src/CurrencyRateController.ts

@@ -11,7 +11,7 @@ import type {
   NetworkClientId,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import { PollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import { Mutex } from 'async-mutex';
 
 import { fetchExchangeRate as defaultFetchExchangeRate } from './crypto-compare';
@@ -82,7 +82,7 @@ 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<
+export class CurrencyRateController extends StaticIntervalPollingController<
   typeof name,
   CurrencyRateState,
   CurrencyRateMessenger

packages/assets-controllers/src/NftDetectionController.ts

@@ -11,7 +11,7 @@ import type {
   NetworkState,
   NetworkClient,
 } from '@metamask/network-controller';
-import { PollingControllerV1 } from '@metamask/polling-controller';
+import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
 
@@ -147,7 +147,7 @@ export interface NftDetectionConfig extends BaseConfig {
 /**
  * Controller that passively polls on a set interval for NFT auto detection
  */
-export class NftDetectionController extends PollingControllerV1<
+export class NftDetectionController extends StaticIntervalPollingControllerV1<
   NftDetectionConfig,
   BaseState
 > {

packages/assets-controllers/src/TokenDetectionController.ts

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

packages/assets-controllers/src/TokenListController.ts

@@ -10,7 +10,7 @@ import type {
   NetworkState,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import { PollingController } from '@metamask/polling-controller';
+import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { Mutex } from 'async-mutex';
 
@@ -91,7 +91,7 @@ const defaultState: TokenListState = {
 /**
  * Controller that passively polls on a set interval for the list of tokens from metaswaps api
  */
-export class TokenListController extends PollingController<
+export class TokenListController extends StaticIntervalPollingController<
   typeof name,
   TokenListState,
   TokenListMessenger

packages/assets-controllers/src/TokenRatesController.ts

@@ -10,7 +10,7 @@ import type {
   NetworkController,
   NetworkState,
 } from '@metamask/network-controller';
-import { PollingControllerV1 } from '@metamask/polling-controller';
+import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
 
@@ -135,7 +135,7 @@ 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<
+export class TokenRatesController extends StaticIntervalPollingControllerV1<
   TokenRatesConfig,
   TokenRatesState
 > {

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

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

packages/polling-controller/src/AbstractPollingController.ts

@@ -0,0 +1,111 @@
+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 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 {
+    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;
+}