From 5705424dfd79e6803873c60d3c66836cc890c4c7 Mon Sep 17 00:00:00 2001
From: Boris Bera <bbera@coveo.com>
Date: Tue, 23 Jul 2024 15:01:24 -0400
Subject: [PATCH] Wrap scheduler tasks in OTEL spans

Signed-off-by: Boris Bera <bbera@coveo.com>
---
 .changeset/rich-bees-tickle.md                |  5 +++
 .../scheduler/lib/PluginTaskSchedulerImpl.ts  | 41 ++++++++++++++-----
 .../src/entrypoints/scheduler/lib/util.ts     |  2 +
 3 files changed, 38 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/rich-bees-tickle.md

diff --git a/.changeset/rich-bees-tickle.md b/.changeset/rich-bees-tickle.md
new file mode 100644
index 0000000000..a2af89c867
--- /dev/null
+++ b/.changeset/rich-bees-tickle.md
@@ -0,0 +1,5 @@
+---
+'@backstage/backend-defaults': patch
+---
+
+Wrap scheduled tasks from the scheduler core service now in OpenTelemetry spans
diff --git a/packages/backend-defaults/src/entrypoints/scheduler/lib/PluginTaskSchedulerImpl.ts b/packages/backend-defaults/src/entrypoints/scheduler/lib/PluginTaskSchedulerImpl.ts
index 40aa9d521c..6bddfeca21 100644
--- a/packages/backend-defaults/src/entrypoints/scheduler/lib/PluginTaskSchedulerImpl.ts
+++ b/packages/backend-defaults/src/entrypoints/scheduler/lib/PluginTaskSchedulerImpl.ts
@@ -23,13 +23,21 @@ import {
   SchedulerServiceTaskRunner,
   SchedulerServiceTaskScheduleDefinition,
 } from '@backstage/backend-plugin-api';
-import { Counter, Histogram, metrics } from '@opentelemetry/api';
+import {
+  Counter,
+  Histogram,
+  metrics,
+  SpanStatusCode,
+  trace,
+} from '@opentelemetry/api';
 import { Knex } from 'knex';
 import { Duration } from 'luxon';
 import { LocalTaskWorker } from './LocalTaskWorker';
 import { TaskWorker } from './TaskWorker';
 import { TaskSettingsV2 } from './types';
-import { validateId } from './util';
+import { TRACER_ID, validateId } from './util';
+
+const tracer = trace.getTracer(TRACER_ID);
 
 /**
  * Implements the actual task management.
@@ -85,7 +93,7 @@ export class PluginTaskSchedulerImpl implements SchedulerService {
       const knex = await this.databaseFactory();
       const worker = new TaskWorker(
         task.id,
-        this.wrapInMetrics(task.fn, { labels: { taskId: task.id, scope } }),
+        this.instrumentedFunction(task, scope),
         knex,
         this.logger.child({ task: task.id }),
       );
@@ -93,7 +101,7 @@ export class PluginTaskSchedulerImpl implements SchedulerService {
     } else {
       const worker = new LocalTaskWorker(
         task.id,
-        this.wrapInMetrics(task.fn, { labels: { taskId: task.id, scope } }),
+        this.instrumentedFunction(task, scope),
         this.logger.child({ task: task.id }),
       );
       worker.start(settings, { signal: task.signal });
@@ -121,20 +129,33 @@ export class PluginTaskSchedulerImpl implements SchedulerService {
     return this.allScheduledTasks;
   }
 
-  private wrapInMetrics(
-    fn: SchedulerServiceTaskFunction,
-    opts: { labels: Record<string, string> },
+  private instrumentedFunction(
+    task: SchedulerServiceTaskInvocationDefinition,
+    scope: string,
   ): SchedulerServiceTaskFunction {
     return async abort => {
-      const labels = {
-        ...opts.labels,
+      const labels: Record<string, string> = {
+        taskId: task.id,
+        scope,
       };
       this.counter.add(1, { ...labels, result: 'started' });
 
       const startTime = process.hrtime();
 
       try {
-        await fn(abort);
+        await tracer.startActiveSpan(`task ${task.id}`, async span => {
+          try {
+            span.setAttributes(labels);
+            await task.fn(abort);
+          } catch (error) {
+            if (error instanceof Error) {
+              span.recordException(error);
+            }
+            throw error;
+          } finally {
+            span.end();
+          }
+        });
         labels.result = 'completed';
       } catch (ex) {
         labels.result = 'failed';
diff --git a/packages/backend-defaults/src/entrypoints/scheduler/lib/util.ts b/packages/backend-defaults/src/entrypoints/scheduler/lib/util.ts
index 70d67a9fbe..5984e6885e 100644
--- a/packages/backend-defaults/src/entrypoints/scheduler/lib/util.ts
+++ b/packages/backend-defaults/src/entrypoints/scheduler/lib/util.ts
@@ -18,6 +18,8 @@ import { InputError } from '@backstage/errors';
 import { Knex } from 'knex';
 import { DateTime, Duration } from 'luxon';
 
+export const TRACER_ID = 'backstage.scheduler';
+
 // Keep the IDs compatible with e.g. Prometheus labels
 export function validateId(id: string) {
   if (typeof id !== 'string' || !id.trim()) {