feat: add scheduler (#30)

packages/server/package.json

@@ -31,6 +31,7 @@
     "@trpc/client": "^10.45.0",
     "@trpc/server": "^10.45.0",
     "commander": "^11.1.0",
+    "cron": "^3.1.6",
     "eventemitter3": "^5.0.1",
     "fastify": "^4.25.2",
     "jsonwebtoken": "^9.0.2",

packages/server/src/database/migrations/migration.schedule.ts

@@ -0,0 +1,22 @@
+import { Knex } from 'knex';
+
+const name = 'schedule-support';
+
+const up = async (knex: Knex) => {
+  await knex.schema.createTable('schedules', (table) => {
+    table.string('id').primary();
+    table.string('name').nullable();
+    table.string('description').nullable();
+    table.string('load').notNullable();
+    table.string('cron').notNullable();
+    table.string('input').nullable();
+    table.timestamp('createdAt').notNullable();
+    table.timestamp('updatedAt').notNullable();
+  });
+};
+
+const down = async (knex: Knex) => {
+  await knex.schema.dropTable('schedule');
+};
+
+export { name, up, down };

packages/server/src/database/migrations/migrations.source.ts

@@ -1,6 +1,7 @@
 import { Knex } from 'knex';
 
 import * as init from './migration.init.js';
+import * as scheduleSupport from './migration.schedule.js';
 
 type Migration = {
   name: string;
@@ -8,7 +9,7 @@ type Migration = {
   down: (knex: Knex) => Promise<void>;
 };
 
-const migrations = [init] satisfies Migration[];
+const migrations = [init, scheduleSupport] satisfies Migration[];
 
 const source: Knex.MigrationSource<Migration> = {
   getMigrations: async () => migrations,

packages/server/src/index.ts

@@ -6,6 +6,7 @@ const start = new Command('start');
 start.action(async () => {
   const port = 4500;
   const runtime = await Runtime.create();
+  await runtime.scheduler.start();
   const server = await createServer(runtime);
   await server.listen({
     port,

packages/server/src/knex.d.ts

@@ -43,5 +43,15 @@ declare module 'knex/types/tables.js' {
       createdAt: Date;
       updatedAt: Date;
     };
+    schedules: {
+      id: string;
+      name?: string;
+      description?: string;
+      load: string;
+      cron: string;
+      input?: string;
+      createdAt: Date;
+      updatedAt: Date;
+    };
   }
 }

packages/server/src/repos/repos.ts

@@ -4,6 +4,7 @@ import { ArtifactRepo } from './artifacts/artifacts.js';
 import { LoadRepo } from './loads/loads.js';
 import { LogRepo } from './logs/logs.js';
 import { RunRepo } from './runs/runs.js';
+import { ScheduleRepo } from './schedules/schedules.js';
 import { SecretRepo } from './secrets/secrets.js';
 
 type ReposOptions = {
@@ -17,6 +18,7 @@ class Repos {
   #logs: LogRepo;
   #artifacts: ArtifactRepo;
   #secrets: SecretRepo;
+  #schedule: ScheduleRepo;
 
   constructor({ database, config }: ReposOptions) {
     this.#loads = new LoadRepo({
@@ -36,6 +38,9 @@ class Repos {
     this.#secrets = new SecretRepo({
       database,
     });
+    this.#schedule = new ScheduleRepo({
+      database,
+    });
   }
 
   public get loads() {
@@ -57,8 +62,13 @@ class Repos {
   public get secrets() {
     return this.#secrets;
   }
+
+  public get schedules() {
+    return this.#schedule;
+  }
 }
 
+export { findSchedulesSchema, addScheduleSchema } from './schedules/schedules.js';
 export { findLogsSchema, addLogSchema } from './logs/logs.js';
 export { setLoadSchema, findLoadsSchema } from './loads/loads.js';
 export { createRunSchema, findRunsSchema } from './runs/runs.js';

packages/server/src/repos/schedules/schedules.schemas.ts

@@ -0,0 +1,22 @@
+import { z } from 'zod';
+
+const addScheduleSchema = z.object({
+  name: z.string().optional(),
+  description: z.string().optional(),
+  load: z.string(),
+  cron: z.string(),
+  input: z.string().optional(),
+});
+
+const findSchedulesSchema = z.object({
+  ids: z.array(z.string()).optional(),
+  loadIds: z.array(z.string()).optional(),
+  offset: z.number().optional(),
+  limit: z.number().optional(),
+});
+
+type AddScheduleOptions = z.infer<typeof addScheduleSchema>;
+type FindSchedulesOptions = z.infer<typeof findSchedulesSchema>;
+
+export type { AddScheduleOptions, FindSchedulesOptions };
+export { addScheduleSchema, findSchedulesSchema };

packages/server/src/repos/schedules/schedules.ts

@@ -0,0 +1,118 @@
+import { EventEmitter } from 'eventemitter3';
+import { Database } from '../../database/database.js';
+import { nanoid } from 'nanoid';
+import { AddScheduleOptions, FindSchedulesOptions } from './schedules.schemas.js';
+import { createHash } from 'crypto';
+
+type ScheduleRepoEvents = {
+  added: (id: string) => void;
+  removed: (id: string) => void;
+};
+
+type ScheduleRepoOptions = {
+  database: Database;
+};
+
+class ScheduleRepo extends EventEmitter<ScheduleRepoEvents> {
+  #options: ScheduleRepoOptions;
+
+  constructor(options: ScheduleRepoOptions) {
+    super();
+    this.#options = options;
+  }
+
+  public get = async (id: string) => {
+    const { database } = this.#options;
+    const db = await database.instance;
+    const result = await db('schedules').where('id', id).first();
+    return result;
+  };
+
+  public add = async (options: AddScheduleOptions) => {
+    const { database } = this.#options;
+    const db = await database.instance;
+    const id = nanoid();
+
+    await db('schedules').insert({
+      id,
+      name: options.name,
+      description: options.description,
+      cron: options.cron,
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    });
+
+    this.emit('added', id);
+
+    return id;
+  };
+
+  public prepareRemove = async (options: FindSchedulesOptions) => {
+    const { database } = this.#options;
+    const db = await database.instance;
+
+    const query = db('schedules').select('id');
+
+    if (options.ids) {
+      query.whereIn('id', options.ids);
+    }
+
+    if (options.loadIds) {
+      query.whereIn('loadId', options.loadIds);
+    }
+
+    const result = await query;
+    const ids = result.map((row) => row.id);
+    const token = ids.map((id) => Buffer.from(id).toString('base64')).join('|');
+    const hash = createHash('sha256').update(token).digest('hex');
+    return {
+      ids,
+      hash,
+    };
+  };
+
+  public remove = async (hash: string, ids: string[]) => {
+    const { database } = this.#options;
+    const db = await database.instance;
+    const token = ids.map((id) => Buffer.from(id).toString('base64')).join('|');
+    const actualHash = createHash('sha256').update(token).digest('hex');
+
+    if (hash !== actualHash) {
+      throw new Error('Invalid hash');
+    }
+
+    await db('schedules').whereIn('id', ids).delete();
+    ids.forEach((id) => {
+      this.emit('removed', id);
+    });
+  };
+
+  public find = async (options: FindSchedulesOptions) => {
+    const { database } = this.#options;
+    const db = await database.instance;
+
+    const query = db('schedules');
+
+    if (options.ids) {
+      query.whereIn('id', options.ids);
+    }
+
+    if (options.loadIds) {
+      query.whereIn('loadId', options.loadIds);
+    }
+
+    if (options.offset) {
+      query.offset(options.offset);
+    }
+
+    if (options.limit) {
+      query.limit(options.limit);
+    }
+
+    const results = await query;
+    return results;
+  };
+}
+
+export { addScheduleSchema, findSchedulesSchema } from './schedules.schemas.js';
+export { ScheduleRepo };

packages/server/src/router/router.schedules.ts

@@ -0,0 +1,53 @@
+import { z } from 'zod';
+import { addScheduleSchema, findSchedulesSchema } from '../repos/repos.js';
+import { publicProcedure, router } from './router.utils.js';
+
+const add = publicProcedure.input(addScheduleSchema).mutation(async ({ input, ctx }) => {
+  const { runtime } = ctx;
+  const { repos } = runtime;
+  const { schedules } = repos;
+
+  const result = await schedules.add(input);
+  return result;
+});
+
+const find = publicProcedure.input(findSchedulesSchema).query(async ({ input, ctx }) => {
+  const { runtime } = ctx;
+  const { repos } = runtime;
+  const { schedules } = repos;
+
+  const result = await schedules.find(input);
+  return result;
+});
+
+const prepareRemove = publicProcedure.input(findSchedulesSchema).query(async ({ input, ctx }) => {
+  const { runtime } = ctx;
+  const { repos } = runtime;
+  const { schedules } = repos;
+
+  return await schedules.prepareRemove(input);
+});
+
+const remove = publicProcedure
+  .input(
+    z.object({
+      hash: z.string(),
+      ids: z.array(z.string()),
+    }),
+  )
+  .mutation(async ({ input, ctx }) => {
+    const { runtime } = ctx;
+    const { repos } = runtime;
+    const { artifacts } = repos;
+
+    await artifacts.remove(input.hash, input.ids);
+  });
+
+const schedulesRouter = router({
+  add,
+  find,
+  remove,
+  prepareRemove,
+});
+
+export { schedulesRouter };

packages/server/src/router/router.ts

@@ -2,6 +2,7 @@ import { artifactsRouter } from './router.artifacts.js';
 import { loadsRouter } from './router.loads.js';
 import { logsRouter } from './router.logs.js';
 import { runsRouter } from './router.runs.js';
+import { schedulesRouter } from './router.schedules.js';
 import { secretsRouter } from './router.secrets.js';
 import { router } from './router.utils.js';
 
@@ -11,6 +12,7 @@ const rootRouter = router({
   logs: logsRouter,
   artifacts: artifactsRouter,
   secrets: secretsRouter,
+  schedules: schedulesRouter,
 });
 
 type RootRouter = typeof rootRouter;

packages/server/src/runtime/runtime.ts

@@ -4,17 +4,20 @@ import { Runner } from '../runner/runner.js';
 import { Config } from '../config/config.js';
 import { Auth } from '../auth/auth.js';
 import { resolve } from 'path';
+import { Scheduler } from '../scheduler/scheduler.js';
 
 class Runtime {
   #repos: Repos;
   #runner: Runner;
   #auth: Auth;
+  #scheduler: Scheduler;
 
   constructor(options: Config) {
     const database = new Database(options.database);
     this.#repos = new Repos({ database, config: options });
     this.#runner = new Runner({ repos: this.#repos, config: options });
     this.#auth = new Auth({ config: options });
+    this.#scheduler = new Scheduler({ runs: this.#repos.runs, schedules: this.#repos.schedules });
   }
 
   public get repos() {
@@ -29,6 +32,10 @@ class Runtime {
     return this.#auth;
   }
 
+  public get scheduler() {
+    return this.#scheduler;
+  }
+
   public static create = async () => {
     const runtime = new Runtime({
       database: {

packages/server/src/scheduler/scheduler.ts

@@ -0,0 +1,73 @@
+import { CronJob } from 'cron';
+import { ScheduleRepo } from '../repos/schedules/schedules.js';
+import { RunRepo } from '../repos/runs/runs.js';
+
+type SchedulerOptions = {
+  runs: RunRepo;
+  schedules: ScheduleRepo;
+};
+
+type RunningSchedule = {
+  id: string;
+  job: CronJob;
+  stop: () => Promise<void>;
+};
+
+class Scheduler {
+  #running: RunningSchedule[] = [];
+  #options: SchedulerOptions;
+
+  constructor(options: SchedulerOptions) {
+    this.#options = options;
+    const { schedules } = this.#options;
+    schedules.on('added', this.#add);
+    schedules.on('removed', this.#remove);
+  }
+
+  #remove = async (id: string) => {
+    const current = this.#running.filter((r) => r.id === id);
+    await Promise.all(current.map((r) => r.stop()));
+    this.#running = this.#running.filter((r) => r.id !== id);
+  };
+
+  #add = async (id: string) => {
+    const { schedules, runs } = this.#options;
+    const current = this.#running.filter((r) => r.id === id);
+    await Promise.all(current.map((r) => r.stop()));
+    const schedule = await schedules.get(id);
+    if (!schedule) {
+      return;
+    }
+    const job = new CronJob(schedule.cron, async () => {
+      await runs.create({
+        loadId: schedule.load,
+      });
+    });
+    const stop = async () => {
+      job.stop();
+    };
+    this.#running.push({
+      id: schedule.id,
+      job,
+      stop,
+    });
+  };
+
+  public stop = async () => {
+    for (const running of this.#running) {
+      await running.stop();
+      this.#running = this.#running.filter((r) => r !== running);
+    }
+  };
+
+  public start = async () => {
+    const { schedules } = this.#options;
+    await this.stop();
+    const all = await schedules.find({});
+    for (const schedule of all) {
+      await this.#add(schedule.id);
+    }
+  };
+}
+
+export { Scheduler };