Skip to content

Commit 24d576a

Browse files
MrtenzMrtenz
committed
Use custom state manager for cronjob controller
1 parent 46c9eb9 commit 24d576a

File tree

2 files changed

+75
-6
lines changed

2 files changed

+75
-6
lines changed

‎packages/snaps-controllers/src/cronjob/CronjobController.test.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@ import { MOCK_ORIGIN, MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
55
import type { SemVerVersion } from '@metamask/utils';
66
import { Duration, inMilliseconds } from '@metamask/utils';
77

8+
import type {
9+
CronjobControllerState,
10+
CronjobControllerStateManager,
11+
} from './CronjobController';
812
import { CronjobController } from './CronjobController';
913
import { METAMASK_ORIGIN } from '../snaps/constants';
1014
import {
@@ -15,6 +19,22 @@ import { getCronjobPermission } from '../test-utils/cronjob';
1519

1620
const MOCK_VERSION = '1.0.0' as SemVerVersion;
1721

22+
/**
23+
* Get a mock state manager for the `CronjobController`.
24+
*
25+
* @returns A state manager object with `get` and `set` methods.
26+
*/
27+
function getMockStateManager(): CronjobControllerStateManager {
28+
let state: CronjobControllerState | undefined;
29+
30+
return {
31+
get: () => state,
32+
set: (newState) => {
33+
state = newState;
34+
},
35+
};
36+
}
37+
1838
describe('CronjobController', () => {
1939
const originalProcessNextTick = process.nextTick;
2040

@@ -44,6 +64,7 @@ describe('CronjobController', () => {
4464

4565
const cronjobController = new CronjobController({
4666
messenger: controllerMessenger,
67+
stateManager: getMockStateManager(),
4768
});
4869

4970
cronjobController.init();
@@ -92,6 +113,7 @@ describe('CronjobController', () => {
92113

93114
const cronjobController = new CronjobController({
94115
messenger: controllerMessenger,
116+
stateManager: getMockStateManager(),
95117
});
96118

97119
cronjobController.init();
@@ -129,6 +151,7 @@ describe('CronjobController', () => {
129151

130152
const cronjobController = new CronjobController({
131153
messenger: controllerMessenger,
154+
stateManager: getMockStateManager(),
132155
});
133156

134157
cronjobController.init();
@@ -161,6 +184,7 @@ describe('CronjobController', () => {
161184

162185
const cronjobController = new CronjobController({
163186
messenger: controllerMessenger,
187+
stateManager: getMockStateManager(),
164188
});
165189

166190
cronjobController.init();
@@ -216,6 +240,7 @@ describe('CronjobController', () => {
216240

217241
const cronjobController = new CronjobController({
218242
messenger: controllerMessenger,
243+
stateManager: getMockStateManager(),
219244
state: {
220245
events: {
221246
[`cronjob-${MOCK_SNAP_ID}-0`]: {
@@ -253,6 +278,7 @@ describe('CronjobController', () => {
253278
const secondCronjobController = new CronjobController({
254279
messenger: controllerMessenger,
255280
state: cronjobController.state,
281+
stateManager: getMockStateManager(),
256282
});
257283

258284
secondCronjobController.init();
@@ -282,6 +308,7 @@ describe('CronjobController', () => {
282308

283309
const cronjobController = new CronjobController({
284310
messenger: controllerMessenger,
311+
stateManager: getMockStateManager(),
285312
});
286313

287314
cronjobController.init();
@@ -329,6 +356,7 @@ describe('CronjobController', () => {
329356

330357
const cronjobController = new CronjobController({
331358
messenger: controllerMessenger,
359+
stateManager: getMockStateManager(),
332360
});
333361

334362
cronjobController.init();
@@ -386,6 +414,7 @@ describe('CronjobController', () => {
386414

387415
const cronjobController = new CronjobController({
388416
messenger: controllerMessenger,
417+
stateManager: getMockStateManager(),
389418
});
390419

391420
cronjobController.init();
@@ -403,6 +432,7 @@ describe('CronjobController', () => {
403432

404433
const cronjobController = new CronjobController({
405434
messenger: controllerMessenger,
435+
stateManager: getMockStateManager(),
406436
state: {
407437
events: {
408438
foo: {
@@ -450,6 +480,7 @@ describe('CronjobController', () => {
450480

451481
const cronjobController = new CronjobController({
452482
messenger: controllerMessenger,
483+
stateManager: getMockStateManager(),
453484
state: {
454485
events: {
455486
foo: {
@@ -498,6 +529,7 @@ describe('CronjobController', () => {
498529

499530
const cronjobController = new CronjobController({
500531
messenger: controllerMessenger,
532+
stateManager: getMockStateManager(),
501533
state: {
502534
events: {
503535
[`cronjob-${MOCK_SNAP_ID}-0`]: {
@@ -572,6 +604,7 @@ describe('CronjobController', () => {
572604

573605
const cronjobController = new CronjobController({
574606
messenger: controllerMessenger,
607+
stateManager: getMockStateManager(),
575608
state: {
576609
events: {
577610
foo: {
@@ -682,6 +715,7 @@ describe('CronjobController', () => {
682715

683716
const cronjobController = new CronjobController({
684717
messenger: controllerMessenger,
718+
stateManager: getMockStateManager(),
685719
});
686720

687721
cronjobController.init();
@@ -716,6 +750,7 @@ describe('CronjobController', () => {
716750

717751
const cronjobController = new CronjobController({
718752
messenger: controllerMessenger,
753+
stateManager: getMockStateManager(),
719754
});
720755

721756
cronjobController.init();
@@ -750,6 +785,7 @@ describe('CronjobController', () => {
750785

751786
const cronjobController = new CronjobController({
752787
messenger: controllerMessenger,
788+
stateManager: getMockStateManager(),
753789
state: {
754790
events: {
755791
[`cronjob-${MOCK_SNAP_ID}-0`]: {
@@ -823,6 +859,7 @@ describe('CronjobController', () => {
823859

824860
const cronjobController = new CronjobController({
825861
messenger: controllerMessenger,
862+
stateManager: getMockStateManager(),
826863
});
827864

828865
cronjobController.init();
@@ -860,6 +897,7 @@ describe('CronjobController', () => {
860897

861898
const cronjobController = new CronjobController({
862899
messenger: controllerMessenger,
900+
stateManager: getMockStateManager(),
863901
});
864902

865903
cronjobController.init();
@@ -893,6 +931,7 @@ describe('CronjobController', () => {
893931

894932
const cronjobController = new CronjobController({
895933
messenger: controllerMessenger,
934+
stateManager: getMockStateManager(),
896935
});
897936

898937
cronjobController.init();
@@ -944,6 +983,7 @@ describe('CronjobController', () => {
944983

945984
const cronjobController = new CronjobController({
946985
messenger: controllerMessenger,
986+
stateManager: getMockStateManager(),
947987
});
948988

949989
cronjobController.init();
@@ -975,6 +1015,7 @@ describe('CronjobController', () => {
9751015

9761016
const cronjobController = new CronjobController({
9771017
messenger: controllerMessenger,
1018+
stateManager: getMockStateManager(),
9781019
});
9791020

9801021
cronjobController.init();
@@ -1015,6 +1056,7 @@ describe('CronjobController', () => {
10151056

10161057
const cronjobController = new CronjobController({
10171058
messenger: controllerMessenger,
1059+
stateManager: getMockStateManager(),
10181060
});
10191061

10201062
cronjobController.init();
@@ -1055,6 +1097,7 @@ describe('CronjobController', () => {
10551097

10561098
const cronjobController = new CronjobController({
10571099
messenger: controllerMessenger,
1100+
stateManager: getMockStateManager(),
10581101
});
10591102

10601103
cronjobController.init();
@@ -1098,6 +1141,7 @@ describe('CronjobController', () => {
10981141

10991142
const cronjobController = new CronjobController({
11001143
messenger: controllerMessenger,
1144+
stateManager: getMockStateManager(),
11011145
state: {
11021146
events: {
11031147
[`cronjob-${MOCK_SNAP_ID}-0`]: {

‎packages/snaps-controllers/src/cronjob/CronjobController.ts‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -93,13 +93,26 @@ export type CronjobControllerMessenger = RestrictedMessenger<
9393

9494
export const DAILY_TIMEOUT = inMilliseconds(24, Duration.Hour);
9595

96+
export type CronjobControllerStateManager = {
97+
set(state: CronjobControllerState): void;
98+
get(): CronjobControllerState | undefined;
99+
};
100+
96101
export type CronjobControllerArgs = {
97102
messenger: CronjobControllerMessenger;
98103

99104
/**
100105
* Persisted state that will be used for rehydration.
101106
*/
102107
state?: CronjobControllerState;
108+
109+
/**
110+
* State manager for the controller.
111+
*
112+
* This is a temporary workaround to allow the controller to update the state
113+
* often without persisting all of the client state to disk.
114+
*/
115+
stateManager: CronjobControllerStateManager;
103116
};
104117

105118
/**
@@ -157,22 +170,26 @@ export class CronjobController extends BaseController<
157170
> {
158171
readonly #timers: Map<string, Timer>;
159172

173+
readonly #stateManager: CronjobControllerStateManager;
174+
160175
#dailyTimer: Timer = new Timer(DAILY_TIMEOUT);
161176

162-
constructor({ messenger, state }: CronjobControllerArgs) {
177+
constructor({ messenger, state, stateManager }: CronjobControllerArgs) {
163178
super({
164179
messenger,
165180
metadata: {
166-
events: { persist: true, anonymous: false },
181+
events: { persist: false, anonymous: false },
167182
},
168183
name: controllerName,
169184
state: {
170185
events: {},
171186
...state,
187+
...stateManager.get(),
172188
},
173189
});
174190

175191
this.#timers = new Map();
192+
this.#stateManager = stateManager;
176193

177194
this.messagingSystem.subscribe(
178195
'SnapController:snapInstalled',
@@ -375,10 +392,12 @@ export class CronjobController extends BaseController<
375392
scheduledAt: new Date().toISOString(),
376393
};
377394

378-
this.update((state) => {
395+
const { nextState } = this.update((state) => {
379396
state.events[internalEvent.id] = castDraft(internalEvent);
380397
});
381398

399+
this.#stateManager.set(nextState);
400+
382401
this.#schedule(internalEvent);
383402
return id;
384403
}
@@ -390,10 +409,12 @@ export class CronjobController extends BaseController<
390409
*/
391410
#schedule(event: InternalBackgroundEvent) {
392411
const date = getExecutionDate(event.schedule);
393-
this.update((state) => {
412+
const { nextState } = this.update((state) => {
394413
state.events[event.id].date = date;
395414
});
396415

416+
this.#stateManager.set(nextState);
417+
397418
this.#startTimer({
398419
...event,
399420
date,
@@ -459,10 +480,12 @@ export class CronjobController extends BaseController<
459480
// Non-recurring events are removed from the state after execution, and
460481
// recurring events are rescheduled.
461482
if (!event.recurring) {
462-
this.update((state) => {
483+
const { nextState } = this.update((state) => {
463484
delete state.events[event.id];
464485
});
465486

487+
this.#stateManager.set(nextState);
488+
466489
return;
467490
}
468491

@@ -480,9 +503,11 @@ export class CronjobController extends BaseController<
480503
timer?.cancel();
481504
this.#timers.delete(id);
482505

483-
this.update((state) => {
506+
const { nextState } = this.update((state) => {
484507
delete state.events[id];
485508
});
509+
510+
this.#stateManager.set(nextState);
486511
}
487512

488513
/**

0 commit comments

Comments
 (0)