Update CronjobService to become CronjobController

packages/controllers/src/services/CronjobService.ts → packages/controllers/src/services/CronjobController.ts

@@ -1,4 +1,5 @@
 import {
+  BaseControllerV2 as BaseController,
   RestrictedControllerMessenger,
   HasPermission,
   GetPermissions,
@@ -29,10 +30,6 @@ export type CronjobServiceArgs = {
   messenger: CronjobServiceMessenger;
 };
 
-export type CronjobData = {
-  jobs: Cronjob[];
-};
-
 export type CronjobDefinition = {
   expression: string;
   request: {
@@ -46,21 +43,42 @@ export type Cronjob = {
   id: string;
 } & CronjobDefinition;
 
+export type CronjobControllerState = {
+  snaps: Record<SnapId, string[]>;
+  jobs: Record<string, CronjobDefinition>;
+};
+
+const controllerName = 'cronjobController';
+
 /**
  * Use this service to register and schedule periodically executed jobs
  * using RPC method hooks.
  */
-export class CronjobService {
+export class CronjobController extends BaseController<
+  typeof controllerName,
+  CronjobControllerState,
+  CronjobServiceMessenger
+> {
   private _messenger: CronjobServiceMessenger;
 
-  private _snaps: Map<SnapId, string[]>;
-
   private _jobs: Map<string, Cronjob>;
 
   constructor({ messenger }: CronjobServiceArgs) {
-    this._snaps = new Map();
+    super({
+      messenger,
+      metadata: {
+        snaps: { persist: true, anonymous: false },
+        jobs: { persist: true, anonymous: false },
+      },
+      name: controllerName,
+      state: {
+        snaps: {},
+        jobs: {},
+      },
+    });
     this._jobs = new Map();
     this._messenger = messenger;
+    this.initializeJobSchedule();
   }
 
   /**
@@ -97,10 +115,9 @@ export class CronjobService {
       return { ...definition, id };
     });
 
-    this._snaps.set(
-      snapId,
-      jobs?.map((job) => job.id),
-    );
+    this.update((state) => {
+      state.snaps[snapId] = jobs?.map((job) => job.id);
+    });
 
     jobs.forEach((job) => this.schedule(snapId, job));
   }
@@ -130,6 +147,13 @@ export class CronjobService {
       });
       this.schedule(snapId, job);
     });
+
+    this.update((state) => {
+      state.jobs[job.id] = {
+        expression: job.expression,
+        request: job.request,
+      };
+    });
     this._jobs.set(job.id, { ...job, timer });
   }
 
@@ -139,13 +163,16 @@ export class CronjobService {
    * @param snapId - ID of a snap.
    */
   unregister(snapId: SnapId) {
-    const jobs = this._snaps.get(snapId);
+    const jobs = this.state.snaps[snapId];
     jobs?.forEach((id) => {
       const job = this._jobs.get(id);
       job?.timer?.cancel();
       this._jobs.delete(job?.id as string);
     });
-    this._snaps.delete(snapId);
+
+    this.update((state) => {
+      delete state.snaps[snapId];
+    });
   }
 
   /**
@@ -158,13 +185,35 @@ export class CronjobService {
     return this._jobs.get(jobId);
   }
 
+  /**
+   * Retrieve job specification from state by given job ID.
+   *
+   * @param jobId - ID of a job.
+   * @returns Plain job specification.
+   */
+  getJobFromState(jobId: string): CronjobDefinition {
+    return this.state.jobs[jobId];
+  }
+
   /**
    * Retrieve an array of job ID registered for a given snap identified by ID.
    *
    * @param snapId - ID of a snap.
    * @returns All registered job ID for single snap.
    */
-  getRegisteredJobs(snapId: string): string[] | undefined {
-    return this._snaps.get(snapId);
+  getRegisteredJobs(snapId: string): string[] {
+    return this.state.snaps[snapId] ? this.state.snaps[snapId] : [];
+  }
+
+  /**
+   * Initialize (schedule) jobs when controller starts.
+   */
+  initializeJobSchedule() {
+    Object.keys(this.state.snaps).forEach((snapId) => {
+      this.state.snaps[snapId].forEach((jobId) => {
+        const jobDefinition = this.getJobFromState(jobId);
+        this.schedule(snapId, { id: jobId, ...jobDefinition });
+      });
+    });
   }
 }