From 6916c4b291b0fd0cec26bb4b3f1154089c00240b Mon Sep 17 00:00:00 2001
From: blam <ben@blam.sh>
Date: Tue, 26 Sep 2023 14:54:40 +0200
Subject: [PATCH] feat: move the implementation to the scheduler instead

Signed-off-by: blam <ben@blam.sh>
---
 .../src/tasks/PluginTaskSchedulerImpl.ts      | 44 +++++++++++++++++--
 .../backend-tasks/src/tasks/TaskWorker.ts     | 23 +---------
 2 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/packages/backend-tasks/src/tasks/PluginTaskSchedulerImpl.ts b/packages/backend-tasks/src/tasks/PluginTaskSchedulerImpl.ts
index b9ac80cb47..83acc9e972 100644
--- a/packages/backend-tasks/src/tasks/PluginTaskSchedulerImpl.ts
+++ b/packages/backend-tasks/src/tasks/PluginTaskSchedulerImpl.ts
@@ -28,7 +28,8 @@ import {
   TaskSettingsV2,
 } from './types';
 import { validateId } from './util';
-
+import { TaskFunction } from './types';
+import { metrics, Counter, Histogram } from '@opentelemetry/api';
 /**
  * Implements the actual task management.
  */
@@ -36,10 +37,22 @@ export class PluginTaskSchedulerImpl implements PluginTaskScheduler {
   private readonly localTasksById = new Map<string, LocalTaskWorker>();
   private readonly allScheduledTasks: TaskDescriptor[] = [];
 
+  private readonly counter: Counter;
+  private readonly duration: Histogram;
+
   constructor(
     private readonly databaseFactory: () => Promise<Knex>,
     private readonly logger: Logger,
-  ) {}
+  ) {
+    const meter = metrics.getMeter('default');
+    this.counter = meter.createCounter('backend_task_runs_total', {
+      description: 'Total number of times a task has been run',
+    });
+    this.duration = meter.createHistogram('backend_task_run_duration', {
+      description: 'Histogram of task run durations',
+      unit: 'seconds',
+    });
+  }
 
   async triggerTask(id: string): Promise<void> {
     const localTask = this.localTasksById.get(id);
@@ -70,7 +83,7 @@ export class PluginTaskSchedulerImpl implements PluginTaskScheduler {
       const knex = await this.databaseFactory();
       const worker = new TaskWorker(
         task.id,
-        task.fn,
+        this.wrapInMetrics(task.fn, { labels: { taskId: task.id, scope } }),
         knex,
         this.logger.child({ task: task.id }),
       );
@@ -78,7 +91,7 @@ export class PluginTaskSchedulerImpl implements PluginTaskScheduler {
     } else {
       const worker = new LocalTaskWorker(
         task.id,
-        task.fn,
+        this.wrapInMetrics(task.fn, { labels: { taskId: task.id, scope } }),
         this.logger.child({ task: task.id }),
       );
       worker.start(settings, { signal: task.signal });
@@ -103,6 +116,29 @@ export class PluginTaskSchedulerImpl implements PluginTaskScheduler {
   async getScheduledTasks(): Promise<TaskDescriptor[]> {
     return this.allScheduledTasks;
   }
+
+  private wrapInMetrics(
+    fn: TaskFunction,
+    opts: { labels: Record<string, string> },
+  ): TaskFunction {
+    return async abort => {
+      this.counter.add(1, { ...opts.labels, result: 'started' });
+
+      const startTime = process.hrtime();
+
+      try {
+        await fn(abort);
+        this.counter.add(1, { ...opts.labels, result: 'completed' });
+      } catch (ex) {
+        this.counter.add(1, { ...opts.labels, result: 'failed' });
+        throw ex;
+      } finally {
+        const delta = process.hrtime(startTime);
+        const endTime = delta[0] + delta[1] / 1e9;
+        this.duration.record(endTime, { ...opts.labels });
+      }
+    };
+  }
 }
 
 export function parseDuration(
diff --git a/packages/backend-tasks/src/tasks/TaskWorker.ts b/packages/backend-tasks/src/tasks/TaskWorker.ts
index ecf71c1db4..4e10f0b5b2 100644
--- a/packages/backend-tasks/src/tasks/TaskWorker.ts
+++ b/packages/backend-tasks/src/tasks/TaskWorker.ts
@@ -23,7 +23,6 @@ import { Logger } from 'winston';
 import { DbTasksRow, DB_TASKS_TABLE } from '../database/tables';
 import { TaskFunction, TaskSettingsV2, taskSettingsV2Schema } from './types';
 import { delegateAbortController, nowPlus, sleep } from './util';
-import { metrics, Counter, Histogram } from '@opentelemetry/api';
 
 const DEFAULT_WORK_CHECK_FREQUENCY = Duration.fromObject({ seconds: 5 });
 
@@ -33,25 +32,13 @@ const DEFAULT_WORK_CHECK_FREQUENCY = Duration.fromObject({ seconds: 5 });
  * @private
  */
 export class TaskWorker {
-  private readonly counter: Counter;
-  private readonly duration: Histogram;
-
   constructor(
     private readonly taskId: string,
     private readonly fn: TaskFunction,
     private readonly knex: Knex,
     private readonly logger: Logger,
     private readonly workCheckFrequency: Duration = DEFAULT_WORK_CHECK_FREQUENCY,
-  ) {
-    const meter = metrics.getMeter('default');
-    this.counter = meter.createCounter('backend_task_runs_total', {
-      description: 'Total number of times a task has been run',
-    });
-    this.duration = meter.createHistogram('backend_task_run_duration', {
-      description: 'Histogram of task run durations',
-      unit: 'seconds',
-    });
-  }
+  ) {}
 
   async start(settings: TaskSettingsV2, options?: { signal?: AbortSignal }) {
     try {
@@ -63,7 +50,6 @@ export class TaskWorker {
     this.logger.info(
       `Task worker starting: ${this.taskId}, ${JSON.stringify(settings)}`,
     );
-    this.counter.add(1, { task_id: this.taskId, result: 'started' });
 
     let attemptNum = 1;
     (async () => {
@@ -157,17 +143,10 @@ export class TaskWorker {
     }, Duration.fromISO(taskSettings.timeoutAfterDuration).as('milliseconds'));
 
     try {
-      const startTime = process.hrtime();
       await this.fn(taskAbortController.signal);
-      const delta = process.hrtime(startTime);
-      const endTime = delta[0] + delta[1] / 1e9;
-
-      this.duration.record(endTime, { task_id: this.taskId });
-      this.counter.add(1, { task_id: this.taskId, result: 'completed' });
       taskAbortController.abort(); // releases resources
     } catch (e) {
       this.logger.error(e);
-      this.counter.add(1, { task_id: this.taskId, result: 'failed' });
       await this.tryReleaseTask(ticket, taskSettings);
       return { result: 'failed' };
     } finally {