From ddf7c70faeeefbc83f6e50661e1cc195ed814eb8 Mon Sep 17 00:00:00 2001 From: Shane Jonas Date: Wed, 27 Sep 2023 10:03:17 -0400 Subject: [PATCH 1/7] Fixed PollingController to be a mixin so it can be used with both V1 and V2 controllers --- .../src/PollingController.test.ts | 21 +- .../src/PollingController.ts | 226 ++++++++---------- 2 files changed, 112 insertions(+), 135 deletions(-) diff --git a/packages/polling-controller/src/PollingController.test.ts b/packages/polling-controller/src/PollingController.test.ts index 26c583d2c68..bf3f4417589 100644 --- a/packages/polling-controller/src/PollingController.test.ts +++ b/packages/polling-controller/src/PollingController.test.ts @@ -1,7 +1,7 @@ import { ControllerMessenger } from '@metamask/base-controller'; import type { PollingCompleteType } from './PollingController'; -import PollingController from './PollingController'; +import { PollingController } from './PollingController'; const TICK_TIME = 1000; @@ -27,7 +27,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); jest.advanceTimersByTime(TICK_TIME); @@ -48,7 +47,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); const pollingToken = controller.start('mainnet'); jest.advanceTimersByTime(TICK_TIME); @@ -69,7 +67,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); const pollingToken1 = controller.start('mainnet'); controller.start('mainnet'); @@ -93,7 +90,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); expect(() => { @@ -113,7 +109,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); expect(() => { @@ -136,7 +131,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); jest.advanceTimersByTime(TICK_TIME); @@ -158,7 +152,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); controller.start('mainnet'); @@ -182,20 +175,19 @@ describe('PollingController', () => { PollingCompleteType >(); - mockMessenger.subscribe(`${name}:pollingComplete`, pollingComplete); const controller = new MyGasFeeController({ messenger: mockMessenger, metadata: {}, name, state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); + controller.onPollingComplete(pollingComplete); const pollingToken = controller.start('mainnet'); controller.stop(pollingToken); expect(pollingComplete).toHaveBeenCalledTimes(1); }); - it('should poll at the interval length passed via the constructor', async () => { + it('should poll at the interval length when set via setIntervalLength', async () => { jest.useFakeTimers(); class MyGasFeeController extends PollingController { @@ -208,8 +200,8 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME * 3, }); + controller.setIntervalLength(TICK_TIME * 3); controller.start('mainnet'); jest.advanceTimersByTime(TICK_TIME); await Promise.resolve(); @@ -238,7 +230,6 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME, }); controller.start('mainnet'); controller.start('rinkeby'); @@ -259,7 +250,7 @@ describe('PollingController', () => { controller.stopAll(); }); - it('should poll multiple networkClientIds at the interval length passed via the constructor', async () => { + it('should poll multiple networkClientIds when setting interval length', async () => { jest.useFakeTimers(); class MyGasFeeController extends PollingController { @@ -272,8 +263,8 @@ describe('PollingController', () => { metadata: {}, name: 'PollingController', state: { foo: 'bar' }, - pollingIntervalLength: TICK_TIME * 2, }); + controller.setIntervalLength(TICK_TIME * 2); controller.start('mainnet'); jest.advanceTimersByTime(TICK_TIME); await Promise.resolve(); diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index a9595cc4e84..4a36609f2b1 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -1,10 +1,5 @@ -import { BaseControllerV2 } from '@metamask/base-controller'; -import type { - RestrictedControllerMessenger, - StateMetadata, -} from '@metamask/base-controller'; +import { BaseController, BaseControllerV2 } from '@metamask/base-controller'; import type { NetworkClientId } from '@metamask/network-controller'; -import type { Json } from '@metamask/utils'; import { v4 as random } from 'uuid'; export type PollingCompleteType = { @@ -12,139 +7,130 @@ export type PollingCompleteType = { payload: [string]; }; +type Constructor = new (...args: any[]) => {}; + /** - * PollingController is an abstract class that implements the polling - * functionality for a controller. It is meant to be extended by a controller - * that needs to poll for data by networkClientId. + * PollingControllerMixin * + * @param Base - The base class to mix onto. + * @returns The mixin. */ -export default abstract class PollingController< - Name extends string, - State extends Record, - messenger extends RestrictedControllerMessenger< - Name, - any, - PollingCompleteType | any, - string, - string - >, -> extends BaseControllerV2 { - readonly #intervalLength: number; +function PollingControllerMixin(Base: TBase) { + /** + * PollingController is an abstract class that implements the polling + * functionality for a controller. It is meant to be extended by a controller + * that needs to poll for data by networkClientId. + * + */ + abstract class PollingControllerBase extends Base { + readonly #networkClientIdTokensMap: Map> = + new Map(); - private readonly networkClientIdTokensMap: Map> = - new Map(); + readonly #intervalIds: Record = {}; - private readonly intervalIds: Record = {}; + #callbacks: Set<(networkClientId: NetworkClientId) => void> = new Set(); - constructor({ - name, - state, - messenger, - metadata, - pollingIntervalLength, - }: { - name: Name; - state: State; - metadata: StateMetadata; - messenger: messenger; - pollingIntervalLength: number; - }) { - super({ - name, - state, - messenger, - metadata, - }); + #intervalLength = 1000; - if (!pollingIntervalLength) { - throw new Error('pollingIntervalLength required for PollingController'); + getIntervalLength() { + return this.#intervalLength; } - this.#intervalLength = pollingIntervalLength; - } + setIntervalLength(length: number) { + this.#intervalLength = length; + } - /** - * Starts polling for a networkClientId - * - * @param networkClientId - The networkClientId to start polling for - * @returns void - */ - start(networkClientId: NetworkClientId) { - const innerPollToken = random(); - if (this.networkClientIdTokensMap.has(networkClientId)) { - const set = this.networkClientIdTokensMap.get(networkClientId); - set?.add(innerPollToken); - } else { - const set = new Set(); - set.add(innerPollToken); - this.networkClientIdTokensMap.set(networkClientId, set); + /** + * Starts polling for a networkClientId + * + * @param networkClientId - The networkClientId to start polling for + * @returns void + */ + start(networkClientId: NetworkClientId) { + const innerPollToken = random(); + if (this.#networkClientIdTokensMap.has(networkClientId)) { + const set = this.#networkClientIdTokensMap.get(networkClientId); + set?.add(innerPollToken); + } else { + const set = new Set(); + set.add(innerPollToken); + this.#networkClientIdTokensMap.set(networkClientId, set); + } + this.#poll(networkClientId); + return innerPollToken; } - this.#poll(networkClientId); - return innerPollToken; - } - /** - * Stops polling for all networkClientIds - */ - stopAll() { - this.networkClientIdTokensMap.forEach((tokens, _networkClientId) => { - tokens.forEach((token) => { - this.stop(token); + /** + * Stops polling for all networkClientIds + */ + stopAll() { + this.#networkClientIdTokensMap.forEach((tokens, _networkClientId) => { + tokens.forEach((token) => { + this.stop(token); + }); }); - }); - } - - /** - * Stops polling for a networkClientId - * - * @param pollingToken - The polling token to stop polling for - */ - stop(pollingToken: string) { - if (!pollingToken) { - throw new Error('pollingToken required'); } - let found = false; - this.networkClientIdTokensMap.forEach((tokens, networkClientId) => { - if (tokens.has(pollingToken)) { - found = true; - this.networkClientIdTokensMap - .get(networkClientId) - ?.delete(pollingToken); - if (this.networkClientIdTokensMap.get(networkClientId)?.size === 0) { - clearTimeout(this.intervalIds[networkClientId]); - delete this.intervalIds[networkClientId]; - this.networkClientIdTokensMap.delete(networkClientId); - this.messagingSystem.publish( - `${this.name}:pollingComplete`, - networkClientId, - ); + + /** + * Stops polling for a networkClientId + * + * @param pollingToken - The polling token to stop polling for + */ + stop(pollingToken: string) { + if (!pollingToken) { + throw new Error('pollingToken required'); + } + let found = false; + this.#networkClientIdTokensMap.forEach((tokens, networkClientId) => { + if (tokens.has(pollingToken)) { + found = true; + this.#networkClientIdTokensMap + .get(networkClientId) + ?.delete(pollingToken); + if (this.#networkClientIdTokensMap.get(networkClientId)?.size === 0) { + clearTimeout(this.#intervalIds[networkClientId]); + delete this.#intervalIds[networkClientId]; + this.#networkClientIdTokensMap.delete(networkClientId); + this.#callbacks.forEach((callback) => { + callback(networkClientId); + }); + this.#callbacks.clear(); + } } + }); + if (!found) { + throw new Error('pollingToken not found'); } - }); - if (!found) { - throw new Error('pollingToken not found'); } - } - /** - * Executes the poll for a networkClientId - * - * @param networkClientId - The networkClientId to execute the poll for - */ - abstract executePoll(networkClientId: NetworkClientId): Promise; + /** + * Executes the poll for a networkClientId + * + * @param networkClientId - The networkClientId to execute the poll for + */ + abstract executePoll(networkClientId: NetworkClientId): Promise; - #poll(networkClientId: NetworkClientId) { - if (this.intervalIds[networkClientId]) { - clearTimeout(this.intervalIds[networkClientId]); - delete this.intervalIds[networkClientId]; - } - this.intervalIds[networkClientId] = setTimeout(async () => { - try { - await this.executePoll(networkClientId); - } catch (error) { - console.error(error); + #poll(networkClientId: NetworkClientId) { + if (this.#intervalIds[networkClientId]) { + clearTimeout(this.#intervalIds[networkClientId]); + delete this.#intervalIds[networkClientId]; } - this.#poll(networkClientId); - }, this.#intervalLength); + this.#intervalIds[networkClientId] = setTimeout(async () => { + try { + await this.executePoll(networkClientId); + } catch (error) { + console.error(error); + } + this.#poll(networkClientId); + }, this.#intervalLength); + } + + onPollingComplete(callback: (networkClientId: NetworkClientId) => void) { + this.#callbacks.add(callback); + } } + return PollingControllerBase; } + +export const PollingController = PollingControllerMixin(BaseControllerV2); +export const PollingControllerV1 = PollingControllerMixin(BaseController); From 82255b837595e9e52e655a9dd22e1993f31e3973 Mon Sep 17 00:00:00 2001 From: Shane Jonas Date: Wed, 27 Sep 2023 12:12:27 -0400 Subject: [PATCH 2/7] Changed onPollingComplete to be by networkClientId --- .../src/PollingController.test.ts | 3 +-- .../src/PollingController.ts | 22 ++++++++++++++----- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/packages/polling-controller/src/PollingController.test.ts b/packages/polling-controller/src/PollingController.test.ts index bf3f4417589..6c699a2c5d3 100644 --- a/packages/polling-controller/src/PollingController.test.ts +++ b/packages/polling-controller/src/PollingController.test.ts @@ -175,14 +175,13 @@ describe('PollingController', () => { PollingCompleteType >(); - const controller = new MyGasFeeController({ messenger: mockMessenger, metadata: {}, name, state: { foo: 'bar' }, }); - controller.onPollingComplete(pollingComplete); + controller.onPollingComplete('mainnet', pollingComplete); const pollingToken = controller.start('mainnet'); controller.stop(pollingToken); expect(pollingComplete).toHaveBeenCalledTimes(1); diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index 4a36609f2b1..29eacad18ed 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -28,7 +28,10 @@ function PollingControllerMixin(Base: TBase) { readonly #intervalIds: Record = {}; - #callbacks: Set<(networkClientId: NetworkClientId) => void> = new Set(); + #callbacks: Map< + NetworkClientId, + Set<(networkClientId: NetworkClientId) => void> + > = new Map(); #intervalLength = 1000; @@ -91,10 +94,10 @@ function PollingControllerMixin(Base: TBase) { clearTimeout(this.#intervalIds[networkClientId]); delete this.#intervalIds[networkClientId]; this.#networkClientIdTokensMap.delete(networkClientId); - this.#callbacks.forEach((callback) => { + this.#callbacks.get(networkClientId)?.forEach((callback) => { callback(networkClientId); }); - this.#callbacks.clear(); + this.#callbacks.get(networkClientId)?.clear(); } } }); @@ -125,8 +128,17 @@ function PollingControllerMixin(Base: TBase) { }, this.#intervalLength); } - onPollingComplete(callback: (networkClientId: NetworkClientId) => void) { - this.#callbacks.add(callback); + onPollingComplete( + networkClientId: NetworkClientId, + callback: (networkClientId: NetworkClientId) => void, + ) { + if (this.#callbacks.has(networkClientId)) { + this.#callbacks.get(networkClientId)?.add(callback); + } else { + const set = new Set(); + set.add(callback); + this.#callbacks.set(networkClientId, set); + } } } return PollingControllerBase; From 0b7980721ad78c3a3a23571dfe9f202bdf562b59 Mon Sep 17 00:00:00 2001 From: Shane Jonas Date: Wed, 27 Sep 2023 12:28:37 -0400 Subject: [PATCH 3/7] Fixed linting --- packages/polling-controller/src/PollingController.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index 29eacad18ed..39b6e793601 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -2,11 +2,7 @@ import { BaseController, BaseControllerV2 } from '@metamask/base-controller'; import type { NetworkClientId } from '@metamask/network-controller'; import { v4 as random } from 'uuid'; -export type PollingCompleteType = { - type: `${N}:pollingComplete`; - payload: [string]; -}; - +// eslint-disable-next-line @typescript-eslint/ban-types type Constructor = new (...args: any[]) => {}; /** From 2a2e00a9934114ede290e9bf4c1bc9e7f5215298 Mon Sep 17 00:00:00 2001 From: Shane Jonas Date: Wed, 27 Sep 2023 12:45:51 -0400 Subject: [PATCH 4/7] Removed usused type in PollingController test --- packages/polling-controller/src/PollingController.test.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/polling-controller/src/PollingController.test.ts b/packages/polling-controller/src/PollingController.test.ts index 6c699a2c5d3..394d48af2a0 100644 --- a/packages/polling-controller/src/PollingController.test.ts +++ b/packages/polling-controller/src/PollingController.test.ts @@ -1,6 +1,5 @@ import { ControllerMessenger } from '@metamask/base-controller'; -import type { PollingCompleteType } from './PollingController'; import { PollingController } from './PollingController'; const TICK_TIME = 1000; @@ -170,10 +169,7 @@ describe('PollingController', () => { } const name = 'PollingController'; - const mockMessenger = new ControllerMessenger< - any, - PollingCompleteType - >(); + const mockMessenger = new ControllerMessenger(); const controller = new MyGasFeeController({ messenger: mockMessenger, From ae7372d3a4aeed8dcdd7c641562ac41cf6c66a4b Mon Sep 17 00:00:00 2001 From: Shane Date: Thu, 28 Sep 2023 06:23:20 -0700 Subject: [PATCH 5/7] Update packages/polling-controller/src/PollingController.ts Co-authored-by: Mark Stacey --- packages/polling-controller/src/PollingController.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index 39b6e793601..ceb551f68ec 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -2,8 +2,9 @@ import { BaseController, BaseControllerV2 } from '@metamask/base-controller'; import type { NetworkClientId } from '@metamask/network-controller'; import { v4 as random } from 'uuid'; -// eslint-disable-next-line @typescript-eslint/ban-types -type Constructor = new (...args: any[]) => {}; +// Mixin classes require a constructor with an `...any[]` parameter +// See TS2545 +type Constructor = new (...args: any[]) => object; /** * PollingControllerMixin From 3f882654eccea3b605210a6194d6c9d2821e4003 Mon Sep 17 00:00:00 2001 From: Shane Date: Thu, 28 Sep 2023 06:23:40 -0700 Subject: [PATCH 6/7] Update packages/polling-controller/src/PollingController.ts Co-authored-by: Mark Stacey --- packages/polling-controller/src/PollingController.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index ceb551f68ec..38a8825a308 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -115,6 +115,9 @@ function PollingControllerMixin(Base: TBase) { clearTimeout(this.#intervalIds[networkClientId]); delete this.#intervalIds[networkClientId]; } + // setTimeout is not `await`ing this async function, which is expected + // We're just using async here for improved stack traces + // eslint-disable-next-line @typescript-eslint/no-misused-promises this.#intervalIds[networkClientId] = setTimeout(async () => { try { await this.executePoll(networkClientId); From b28ebd3abf2f4fbf16b577a4c04db2a03f504f34 Mon Sep 17 00:00:00 2001 From: Shane Jonas Date: Thu, 28 Sep 2023 09:31:23 -0400 Subject: [PATCH 7/7] Changed JSDoc for mixin return --- .../polling-controller/src/PollingController.ts | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts index 38a8825a308..03c61f94baa 100644 --- a/packages/polling-controller/src/PollingController.ts +++ b/packages/polling-controller/src/PollingController.ts @@ -10,7 +10,7 @@ type Constructor = new (...args: any[]) => object; * PollingControllerMixin * * @param Base - The base class to mix onto. - * @returns The mixin. + * @returns The composed class. */ function PollingControllerMixin(Base: TBase) { /** @@ -36,6 +36,11 @@ function PollingControllerMixin(Base: TBase) { return this.#intervalLength; } + /** + * Sets the length of the polling interval + * + * @param length - The length of the polling interval in milliseconds + */ setIntervalLength(length: number) { this.#intervalLength = length; } @@ -128,6 +133,12 @@ function PollingControllerMixin(Base: TBase) { }, this.#intervalLength); } + /** + * Adds a callback to execute when polling is complete + * + * @param networkClientId - The networkClientId to listen for polling complete events + * @param callback - The callback to execute when polling is complete + */ onPollingComplete( networkClientId: NetworkClientId, callback: (networkClientId: NetworkClientId) => void,