chore: add last_scheduled_task_id column to schedules (NangoHQ#2382)

DueSchedule query is too slow because for each schedules we are going
through the tasks table to find out if there is any running tasks or if
any task has been running recently.
In order to make the query faster we are adding a last_scheduled_task_id
column to schedules which is populated when scheduling a task. This
column will allow us to join schedules with tasks based on this column
and look at the state and created_at of the tasks to find due schedules

## Describe your changes

## Issue ticket number and link

## Checklist before requesting a review (skip if just adding/editing
APIs & templates)
- [ ] I added tests, otherwise the reason is: 
- [ ] I added observability, otherwise the reason is:
- [ ] I added analytics, otherwise the reason is:

Author: TBonnin

packages/orchestrator/lib/clients/validate.ts

@@ -152,7 +152,8 @@ export function validateSchedule(schedule: Schedule): Result<OrchestratorSchedul
             heartbeatTimeoutSecs: z.number().int(),
             createdAt: z.coerce.date(),
             updatedAt: z.coerce.date(),
-            deletedAt: z.coerce.date().nullable()
+            deletedAt: z.coerce.date().nullable(),
+            lastScheduledTaskId: z.string().uuid().nullable()
         })
         .strict();
     const getNextDueDate = (startsAt: Date, frequencyMs: number) => {

packages/orchestrator/lib/routes/v1/postRecurring.ts

@@ -69,7 +69,8 @@ const handler = (scheduler: Scheduler) => {
             retryMax: req.body.retry.max,
             createdToStartedTimeoutSecs: req.body.timeoutSettingsInSecs.createdToStarted,
             startedToCompletedTimeoutSecs: req.body.timeoutSettingsInSecs.startedToCompleted,
-            heartbeatTimeoutSecs: req.body.timeoutSettingsInSecs.heartbeat
+            heartbeatTimeoutSecs: req.body.timeoutSettingsInSecs.heartbeat,
+            lastScheduledTaskId: null
         });
         if (schedule.isErr()) {
             return res.status(500).json({ error: { code: 'recurring_failed', message: schedule.error.message } });

packages/scheduler/lib/db/migrations/20240620163500_schedules_last_task.ts

@@ -0,0 +1,17 @@
+import type { Knex } from 'knex';
+import { SCHEDULES_TABLE } from '../../models/schedules.js';
+
+export async function up(knex: Knex): Promise<void> {
+    await knex.raw(`
+        ALTER TABLE ${SCHEDULES_TABLE}
+        ADD COLUMN IF NOT EXISTS last_scheduled_task_id uuid NULL;
+    `);
+}
+
+export async function down(knex: Knex): Promise<void> {
+    await knex.raw(`
+        ALTER TABLE ${SCHEDULES_TABLE}
+        DROP COLUMN IF EXISTS last_scheduled_task_id;
+
+    `);
+}

packages/scheduler/lib/models/schedules.integration.test.ts

@@ -3,6 +3,7 @@ import * as schedules from './schedules.js';
 import { getTestDbClient } from '../db/helpers.test.js';
 import type { Schedule } from '../types.js';
 import type knex from 'knex';
+import { uuidv7 } from 'uuidv7';
 
 describe('Schedules', () => {
     const dbClient = getTestDbClient();
@@ -25,7 +26,8 @@ describe('Schedules', () => {
             frequencyMs: 300_000,
             createdAt: expect.toBeIsoDateTimezone(),
             updatedAt: expect.toBeIsoDateTimezone(),
-            deletedAt: null
+            deletedAt: null,
+            lastScheduledTaskId: null
         });
     });
     it('should be successfully retrieved', async () => {
@@ -60,7 +62,7 @@ describe('Schedules', () => {
     });
     it('should be successfully updated', async () => {
         const schedule = await createSchedule(db);
-        const updated = (await schedules.update(db, { id: schedule.id, frequencyMs: 600_000, payload: { i: 2 } })).unwrap();
+        const updated = (await schedules.update(db, { id: schedule.id, frequencyMs: 600_000, payload: { i: 2 }, lastScheduledTaskId: uuidv7() })).unwrap();
         expect(updated.frequencyMs).toBe(600_000);
         expect(updated.payload).toMatchObject({ i: 2 });
         expect(updated.updatedAt.getTime()).toBeGreaterThan(schedule.updatedAt.getTime());
@@ -90,7 +92,8 @@ async function createSchedule(db: knex.Knex): Promise<Schedule> {
             retryMax: 1,
             createdToStartedTimeoutSecs: 1,
             startedToCompletedTimeoutSecs: 1,
-            heartbeatTimeoutSecs: 1
+            heartbeatTimeoutSecs: 1,
+            lastScheduledTaskId: null
         })
     ).unwrap();
 }

packages/scheduler/lib/models/schedules.ts

@@ -46,6 +46,7 @@ export interface DbSchedule {
     readonly created_at: Date;
     updated_at: Date;
     deleted_at: Date | null;
+    last_scheduled_task_id: string | null;
 }
 
 // knex uses https://github.com/bendrucker/postgres-interval
@@ -84,7 +85,8 @@ export const DbSchedule = {
         heartbeat_timeout_secs: schedule.heartbeatTimeoutSecs,
         created_at: schedule.createdAt,
         updated_at: schedule.updatedAt,
-        deleted_at: schedule.deletedAt
+        deleted_at: schedule.deletedAt,
+        last_scheduled_task_id: schedule.lastScheduledTaskId
     }),
     from: (dbSchedule: DbSchedule): Schedule => ({
         id: dbSchedule.id,
@@ -100,7 +102,8 @@ export const DbSchedule = {
         heartbeatTimeoutSecs: dbSchedule.heartbeat_timeout_secs,
         createdAt: dbSchedule.created_at,
         updatedAt: dbSchedule.updated_at,
-        deletedAt: dbSchedule.deleted_at
+        deletedAt: dbSchedule.deleted_at,
+        lastScheduledTaskId: dbSchedule.last_scheduled_task_id
     })
 };
 
@@ -115,7 +118,8 @@ export async function create(db: knex.Knex, props: ScheduleProps): Promise<Resul
         frequencyMs: props.frequencyMs,
         createdAt: now,
         updatedAt: now,
-        deletedAt: null
+        deletedAt: null,
+        lastScheduledTaskId: null
     };
     try {
         const inserted = await db.from<DbSchedule>(SCHEDULES_TABLE).insert(DbSchedule.to(newSchedule)).returning('*');
@@ -166,11 +170,15 @@ export async function transitionState(db: knex.Knex, scheduleId: string, to: Sch
     }
 }
 
-export async function update(db: knex.Knex, props: Partial<Pick<ScheduleProps, 'frequencyMs' | 'payload'>> & { id: string }): Promise<Result<Schedule>> {
+export async function update(
+    db: knex.Knex,
+    props: Partial<Pick<ScheduleProps, 'frequencyMs' | 'payload' | 'lastScheduledTaskId'>> & { id: string }
+): Promise<Result<Schedule>> {
     try {
         const newValues = {
             ...(props.frequencyMs ? { frequency: `${props.frequencyMs} milliseconds` } : {}),
             ...(props.payload ? { payload: props.payload } : {}),
+            ...(props.lastScheduledTaskId ? { last_scheduled_task_id: props.lastScheduledTaskId } : {}),
             updated_at: new Date()
         };
         const updated = await db.from<DbSchedule>(SCHEDULES_TABLE).where('id', props.id).update(newValues).returning('*');

packages/scheduler/lib/scheduler.integration.test.ts

@@ -168,7 +168,8 @@ async function recurring({ scheduler, state = 'PAUSED' }: { scheduler: Scheduler
         retryCount: 0,
         createdToStartedTimeoutSecs: 3600,
         startedToCompletedTimeoutSecs: 3600,
-        heartbeatTimeoutSecs: 600
+        heartbeatTimeoutSecs: 600,
+        lastScheduledTaskId: null
     };
     return (await scheduler.recurring(recurringProps)).unwrap();
 }

packages/scheduler/lib/types.ts

@@ -48,4 +48,5 @@ export interface Schedule {
     readonly createdAt: Date;
     readonly updatedAt: Date;
     readonly deletedAt: Date | null;
+    readonly lastScheduledTaskId: string | null;
 }

packages/scheduler/lib/workers/scheduling/scheduling.integration.test.ts

@@ -108,7 +108,8 @@ async function addSchedule(db: knex.Knex, params?: { state?: ScheduleState; star
         heartbeat_timeout_secs: 1,
         created_at: new Date(),
         updated_at: new Date(),
-        deleted_at: params?.state === 'DELETED' ? new Date() : null
+        deleted_at: params?.state === 'DELETED' ? new Date() : null,
+        last_scheduled_task_id: null
     };
     const res = await db.from<DbSchedule>(SCHEDULES_TABLE).insert(schedule).returning('*');
     const inserted = res[0];
@@ -151,5 +152,8 @@ async function addTask(
     if (!inserted) {
         throw new Error('Failed to insert task');
     }
+    if (params?.scheduleId) {
+        await db.from<DbSchedule>(SCHEDULES_TABLE).where('id', params.scheduleId).update({ last_scheduled_task_id: inserted.id });
+    }
     return DbTask.from(inserted);
 }

packages/scheduler/lib/workers/scheduling/scheduling.worker.ts

@@ -8,6 +8,7 @@ import type knex from 'knex';
 import { logger } from '../../utils/logger.js';
 import { dueSchedules } from './scheduling.js';
 import * as tasks from '../../models/tasks.js';
+import * as schedules from '../../models/schedules.js';
 import tracer from 'dd-trace';
 
 interface CreatedTasksMessage {
@@ -109,17 +110,17 @@ export class SchedulingChild {
 
                 if (lockGranted) {
                     const dueSchedulesSpan = tracer.startSpan('scheduler.scheduling.due_schedules', { childOf: span });
-                    const schedules = await tracer.scope().activate(dueSchedulesSpan, async () => {
+                    const getDueSchedules = await tracer.scope().activate(dueSchedulesSpan, async () => {
                         return dueSchedules(trx);
                     });
                     dueSchedulesSpan.finish();
 
-                    if (schedules.isErr()) {
-                        return Err(`Failed to get due schedules: ${stringifyError(schedules.error)}`);
+                    if (getDueSchedules.isErr()) {
+                        return Err(`Failed to get due schedules: ${stringifyError(getDueSchedules.error)}`);
                     } else {
                         const tasksCreationSpan = tracer.startSpan('scheduler.scheduling.tasks_creation', { childOf: span });
                         await tracer.scope().activate(tasksCreationSpan, async () => {
-                            const createTasks = schedules.value.map((schedule) =>
+                            const createTasks = getDueSchedules.value.map((schedule) =>
                                 tasks.create(trx, {
                                     scheduleId: schedule.id,
                                     startsAfter: new Date(),
@@ -140,7 +141,11 @@ export class SchedulingChild {
                                 } else if (taskRes.value.isErr()) {
                                     logger.error(`Failed to schedule task: ${stringifyError(taskRes.value.error)}`);
                                 } else {
-                                    taskIds.push(taskRes.value.value.id);
+                                    const task = taskRes.value.value;
+                                    if (task.scheduleId) {
+                                        schedules.update(trx, { id: task.scheduleId, lastScheduledTaskId: task.id });
+                                    }
+                                    taskIds.push(task.id);
                                 }
                             }
                         });