From 1ef5c2720dff91e898212ee95d967cc39ed94c71 Mon Sep 17 00:00:00 2001
From: blam <ben@blam.sh>
Date: Fri, 22 Sep 2023 15:02:11 +0200
Subject: [PATCH] feat: turn on some simple monitoring with prometheus for now

Signed-off-by: blam <ben@blam.sh>
---
 packages/backend-tasks/package.json           |  1 +
 .../backend-tasks/src/tasks/TaskWorker.ts     | 26 ++++++++++-
 packages/backend-tasks/src/tasks/metrics.ts   | 46 +++++++++++++++++++
 yarn.lock                                     |  1 +
 4 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 packages/backend-tasks/src/tasks/metrics.ts

diff --git a/packages/backend-tasks/package.json b/packages/backend-tasks/package.json
index f4118b8098..221e1f3a69 100644
--- a/packages/backend-tasks/package.json
+++ b/packages/backend-tasks/package.json
@@ -41,6 +41,7 @@
     "knex": "^2.0.0",
     "lodash": "^4.17.21",
     "luxon": "^3.0.0",
+    "prom-client": "^14.0.1",
     "uuid": "^8.0.0",
     "winston": "^3.2.1",
     "zod": "^3.21.4"
diff --git a/packages/backend-tasks/src/tasks/TaskWorker.ts b/packages/backend-tasks/src/tasks/TaskWorker.ts
index 6e3a88410b..b5c4162f68 100644
--- a/packages/backend-tasks/src/tasks/TaskWorker.ts
+++ b/packages/backend-tasks/src/tasks/TaskWorker.ts
@@ -18,9 +18,11 @@ import { ConflictError, NotFoundError } from '@backstage/errors';
 import { CronTime } from 'cron';
 import { Knex } from 'knex';
 import { DateTime, Duration } from 'luxon';
+import { Counter, Summary } from 'prom-client';
 import { v4 as uuid } from 'uuid';
 import { Logger } from 'winston';
 import { DbTasksRow, DB_TASKS_TABLE } from '../database/tables';
+import { createCounterMetric, createSummaryMetric } from './metrics';
 import { TaskFunction, TaskSettingsV2, taskSettingsV2Schema } from './types';
 import { delegateAbortController, nowPlus, sleep } from './util';
 
@@ -32,13 +34,27 @@ const DEFAULT_WORK_CHECK_FREQUENCY = Duration.fromObject({ seconds: 5 });
  * @private
  */
 export class TaskWorker {
+  private readonly counter: Counter;
+  private readonly duration: Summary;
+
   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,
-  ) {}
+  ) {
+    this.counter = createCounterMetric({
+      name: 'backstage_task_runs_total',
+      help: 'Total number of times a task has been run',
+      labelNames: ['task_id', 'result'],
+    });
+    this.duration = createSummaryMetric({
+      name: 'backstage_task_run_duration_seconds',
+      help: 'Duration of task runs in seconds',
+      labelNames: ['task_id', 'result'],
+    });
+  }
 
   async start(settings: TaskSettingsV2, options?: { signal?: AbortSignal }) {
     try {
@@ -50,9 +66,13 @@ export class TaskWorker {
     this.logger.info(
       `Task worker starting: ${this.taskId}, ${JSON.stringify(settings)}`,
     );
+    this.counter.inc({ task_id: this.taskId, result: 'started' }, 1);
+
     let attemptNum = 1;
     (async () => {
       for (;;) {
+        const endDuration = this.duration.startTimer();
+
         try {
           if (settings.initialDelayDuration) {
             await sleep(
@@ -71,6 +91,8 @@ export class TaskWorker {
           }
 
           this.logger.info(`Task worker finished: ${this.taskId}`);
+          this.counter.inc({ task_id: this.taskId, result: 'completed' }, 1);
+          endDuration({ task_id: this.taskId, result: 'completed' });
           attemptNum = 0;
           break;
         } catch (e) {
@@ -78,6 +100,8 @@ export class TaskWorker {
           this.logger.warn(
             `Task worker failed unexpectedly, attempt number ${attemptNum}, ${e}`,
           );
+          this.counter.inc({ task_id: this.taskId, result: 'failed' }, 1);
+          endDuration({ task_id: this.taskId, result: 'failed' });
           await sleep(Duration.fromObject({ seconds: 1 }));
         }
       }
diff --git a/packages/backend-tasks/src/tasks/metrics.ts b/packages/backend-tasks/src/tasks/metrics.ts
new file mode 100644
index 0000000000..574e64afcb
--- /dev/null
+++ b/packages/backend-tasks/src/tasks/metrics.ts
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2023 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  Counter,
+  CounterConfiguration,
+  register,
+  Summary,
+  SummaryConfiguration,
+} from 'prom-client';
+
+export function createCounterMetric<T extends string>(
+  config: CounterConfiguration<T>,
+): Counter<T> {
+  let metric = register.getSingleMetric(config.name);
+  if (!metric) {
+    metric = new Counter<T>(config);
+    register.registerMetric(metric);
+  }
+  return metric as Counter<T>;
+}
+
+export function createSummaryMetric<T extends string>(
+  config: SummaryConfiguration<T>,
+): Summary<T> {
+  let metric = register.getSingleMetric(config.name);
+  if (!metric) {
+    metric = new Summary<T>(config);
+    register.registerMetric(metric);
+  }
+
+  return metric as Summary<T>;
+}
diff --git a/yarn.lock b/yarn.lock
index de2a8d5ea7..fee4c29a3a 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -3628,6 +3628,7 @@ __metadata:
     knex: ^2.0.0
     lodash: ^4.17.21
     luxon: ^3.0.0
+    prom-client: ^14.0.1
     uuid: ^8.0.0
     wait-for-expect: ^3.0.2
     winston: ^3.2.1