Skip to content

Commit

Permalink
Add some refactoring and new features for the CronjobController
Browse files Browse the repository at this point in the history
  • Loading branch information
@david0xd
david0xd committed Oct 12, 2022
1 parent 0070f68 commit 6296ecb
Show file tree
Hide file tree
Showing 3 changed files with 217 additions and 56 deletions.
50 changes: 25 additions & 25 deletions ...llers/src/services/CronjobService.test.ts → ...rs/src/services/CronjobController.test.ts
View file
@@ -1,7 +1,7 @@
import { SnapEndowments } from '../snaps';
import { getSnapControllerMessenger } from '../test-utils';
import { getRestrictedCronjobControllerMessenger } from '../test-utils';
import { Timer } from '../snaps/Timer';
import { Cronjob, CronjobService } from './CronjobService';
import { Cronjob, CronjobController } from './CronjobController';

jest.mock('../snaps/Timer');
const MOCK_ORIGIN = 'metamask.io';
Expand Down Expand Up @@ -37,7 +37,7 @@ describe('CronjobService', () => {
const mockSnapId = 'snap-id';

it('registers a snap', async () => {
const controllerMessenger = getSnapControllerMessenger();
const controllerMessenger = getRestrictedCronjobControllerMessenger();

const callActionMock = jest
.spyOn(controllerMessenger, 'call')
Expand All @@ -50,11 +50,11 @@ describe('CronjobService', () => {
return false;
});

const cronjobService = new CronjobService({
const cronjobController = new CronjobController({
messenger: controllerMessenger,
});

await cronjobService.register(mockSnapId);
await cronjobController.register(mockSnapId);

expect(callActionMock).toHaveBeenCalledWith(
'PermissionController:hasPermission',
Expand All @@ -74,23 +74,23 @@ describe('CronjobService', () => {
id: 'cronjob-id',
request: { method: 'call_this_method', params: [] },
};
const controllerMessenger = getSnapControllerMessenger();
const controllerMessenger = getRestrictedCronjobControllerMessenger();

const cronjobService = new CronjobService({
const cronjobController = new CronjobController({
messenger: controllerMessenger,
});

await cronjobService.schedule(mockSnapId, mockJob);
await cronjobController.schedule(mockSnapId, mockJob);

expect(Timer.prototype.start).toHaveBeenCalled();
expect(cronjobService.getJob(mockJob.id)).toStrictEqual({
expect(cronjobController.getJob(mockJob.id)).toStrictEqual({
...mockJob,
timer: expect.anything(),
});
});

it('unregisters scheduled jobs for a snap', async () => {
const controllerMessenger = getSnapControllerMessenger();
const controllerMessenger = getRestrictedCronjobControllerMessenger();

jest.spyOn(controllerMessenger, 'call').mockImplementation((method) => {
if (method === 'PermissionController:hasPermission') {
Expand All @@ -101,36 +101,36 @@ describe('CronjobService', () => {
return false;
});

const cronjobService = new CronjobService({
const cronjobController = new CronjobController({
messenger: controllerMessenger,
});

await cronjobService.register(mockSnapId);
const registeredJobIds = cronjobService.getRegisteredJobs(mockSnapId);
await cronjobController.register(mockSnapId);
const registeredJobIds = cronjobController.getRegisteredJobs(mockSnapId);
const jobOneId = registeredJobIds?.[0] as string;
const jobTwoId = registeredJobIds?.[1] as string;
expect(registeredJobIds?.length).toStrictEqual(2);

cronjobService.unregister(mockSnapId);
expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined();
expect(cronjobService.getJob(jobOneId)).toBeUndefined();
expect(cronjobService.getJob(jobTwoId)).toBeUndefined();
cronjobController.unregister(mockSnapId);
expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]);
expect(cronjobController.getJob(jobOneId)).toBeUndefined();
expect(cronjobController.getJob(jobTwoId)).toBeUndefined();
});

it('should not register jobs if snap has no cronjob permission', async () => {
const controllerMessenger = getSnapControllerMessenger();
const controllerMessenger = getRestrictedCronjobControllerMessenger();

const callActionMock = jest
.spyOn(controllerMessenger, 'call')
.mockImplementation(() => {
return false;
});

const cronjobService = new CronjobService({
const cronjobController = new CronjobController({
messenger: controllerMessenger,
});

await cronjobService.register(mockSnapId);
await cronjobController.register(mockSnapId);

expect(callActionMock).toHaveBeenCalledWith(
'PermissionController:hasPermission',
Expand All @@ -143,11 +143,11 @@ describe('CronjobService', () => {
mockSnapId,
);

expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined();
expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]);
});

it('should not register any jobs if snap has no cronjob caveats', async () => {
const controllerMessenger = getSnapControllerMessenger();
const controllerMessenger = getRestrictedCronjobControllerMessenger();

const callActionMock = jest
.spyOn(controllerMessenger, 'call')
Expand All @@ -160,11 +160,11 @@ describe('CronjobService', () => {
return false;
});

const cronjobService = new CronjobService({
const cronjobController = new CronjobController({
messenger: controllerMessenger,
});

await cronjobService.register(mockSnapId);
await cronjobController.register(mockSnapId);

expect(callActionMock).toHaveBeenCalledWith(
'PermissionController:hasPermission',
Expand All @@ -177,6 +177,6 @@ describe('CronjobService', () => {
mockSnapId,
);

expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined();
expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]);
});
});
159 changes: 128 additions & 31 deletions packages/controllers/src/services/CronjobController.ts
View file
Expand Up @@ -4,30 +4,48 @@ import {
HasPermission,
GetPermissions,
} from '@metamask/controllers';
import { HandlerType, SnapId } from '@metamask/snap-utils';
import { HandlerType, SnapId, TruncatedSnap } from '@metamask/snap-utils';
import { parseExpression } from 'cron-parser';
import { nanoid } from 'nanoid';
import { GetSnap, HandleSnapRequest, SnapAdded, SnapEndowments } from '..';
import {
GetSnap,
HandleSnapRequest,
SnapAdded,
SnapBlocked,
SnapEndowments,
SnapInstalled,
SnapRemoved,
SnapTerminated,
SnapUnblocked,
SnapUpdated,
} from '..';
import { Timer } from '../snaps/Timer';

export type CronjobServiceActions =
export type CronjobControllerActions =
| GetSnap
| HandleSnapRequest
| HasPermission
| GetPermissions;

export type CronjobServiceEvents = SnapAdded;
export type CronjobControllerEvents =
| SnapAdded
| SnapBlocked
| SnapInstalled
| SnapRemoved
| SnapUnblocked
| SnapUpdated
| SnapTerminated;

export type CronjobServiceMessenger = RestrictedControllerMessenger<
'CronjobService',
CronjobServiceActions,
CronjobServiceEvents,
CronjobServiceActions['type'],
CronjobServiceEvents['type']
export type CronjobControllerMessenger = RestrictedControllerMessenger<
'CronjobController',
CronjobControllerActions,
CronjobControllerEvents,
CronjobControllerActions['type'],
CronjobControllerEvents['type']
>;

export type CronjobServiceArgs = {
messenger: CronjobServiceMessenger;
export type CronjobControllerArgs = {
messenger: CronjobControllerMessenger;
};

export type CronjobDefinition = {
Expand All @@ -43,27 +61,31 @@ export type Cronjob = {
id: string;
} & CronjobDefinition;

export type StoredJobInformation = {
lastRun: number;
};

export type CronjobControllerState = {
snaps: Record<SnapId, string[]>;
jobs: Record<string, CronjobDefinition>;
jobs: Record<string, StoredJobInformation>;
};

const controllerName = 'cronjobController';
const controllerName = 'CronjobController';

/**
* Use this service to register and schedule periodically executed jobs
* Use this controller to register and schedule periodically executed jobs
* using RPC method hooks.
*/
export class CronjobController extends BaseController<
typeof controllerName,
CronjobControllerState,
CronjobServiceMessenger
CronjobControllerMessenger
> {
private _messenger: CronjobServiceMessenger;
private _messenger: CronjobControllerMessenger;

private _jobs: Map<string, Cronjob>;

constructor({ messenger }: CronjobServiceArgs) {
constructor({ messenger }: CronjobControllerArgs) {
super({
messenger,
metadata: {
Expand All @@ -79,6 +101,26 @@ export class CronjobController extends BaseController<
this._jobs = new Map();
this._messenger = messenger;
this.initializeJobSchedule();

this._handleEventSnapInstalled = this._handleEventSnapInstalled.bind(this);
this._handleEventSnapRemoved = this._handleEventSnapRemoved.bind(this);
this._handleEventSnapUpdated = this._handleEventSnapUpdated.bind(this);

// Subscribe to Snap events
this.messagingSystem.subscribe(
'SnapController:snapInstalled',
this._handleEventSnapInstalled,
);

this.messagingSystem.subscribe(
'SnapController:snapRemoved',
this._handleEventSnapRemoved,
);

this.messagingSystem.subscribe(
'SnapController:snapInstalled',
this._handleEventSnapUpdated,
);
}

/**
Expand Down Expand Up @@ -138,6 +180,7 @@ export class CronjobController extends BaseController<
const ms = next.getTime() - now.getTime();
const timer = new Timer(ms);
timer.start(() => {
this.updateJobLastRunState(job.id, next.getTime());
this._messenger.call('SnapController:handleRequest', {
snapId,
// @todo Decide on origin for requests like this
Expand All @@ -148,12 +191,7 @@ export class CronjobController extends BaseController<
this.schedule(snapId, job);
});

this.update((state) => {
state.jobs[job.id] = {
expression: job.expression,
request: job.request,
};
});
this.updateJobLastRunState(job.id, 0); // 0 for init, never ran actually
this._jobs.set(job.id, { ...job, timer });
}

Expand All @@ -175,6 +213,20 @@ export class CronjobController extends BaseController<
});
}

/**
* Update time of a last run for the Cronjob specified by ID.
*
* @param jobId - ID of a cron job.
* @param lastRun - Unix timestamp when the job was last ran.
*/
updateJobLastRunState(jobId: string, lastRun: number) {
this.update((state) => {
state.jobs[jobId] = {
lastRun,
};
});
}

/**
* Retrieve job specification by given job ID.
*
Expand All @@ -191,7 +243,7 @@ export class CronjobController extends BaseController<
* @param jobId - ID of a job.
* @returns Plain job specification.
*/
getJobFromState(jobId: string): CronjobDefinition {
getJobInfoFromState(jobId: string): StoredJobInformation {
return this.state.jobs[jobId];
}

Expand All @@ -209,11 +261,56 @@ export class CronjobController extends BaseController<
* 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 });
});
});
// TODO: Do some init actions after restart?
}

/**
* Run controller teardown process and unsubscribe from everything.
*/
protected destroy() {
super.destroy();

this.messagingSystem.unsubscribe(
'SnapController:snapInstalled',
this._handleEventSnapInstalled,
);

this.messagingSystem.unsubscribe(
'SnapController:snapRemoved',
this._handleEventSnapRemoved,
);

this.messagingSystem.unsubscribe(
'SnapController:snapUpdated',
this._handleEventSnapUpdated,
);
}

/**
* Handle cron jobs on 'snapInstalled' event.
*
* @param snap - Basic Snap information.
*/
async _handleEventSnapInstalled(snap: TruncatedSnap) {
await this.register(snap.id);
}

/**
* Handle cron jobs on 'snapRemoved' event.
*
* @param snap - Basic Snap information.
*/
_handleEventSnapRemoved(snap: TruncatedSnap) {
this.unregister(snap.id);
}

/**
* Handle cron jobs on 'snapUpdated' event.
*
* @param snap - Basic Snap information.
*/
async _handleEventSnapUpdated(snap: TruncatedSnap) {
this.unregister(snap.id);
await this.register(snap.id);
}
}

0 comments on commit 6296ecb

Please sign in to comment.