From de7513e7fa1eb7f1eebb038988982fd8f9c86792 Mon Sep 17 00:00:00 2001
From: Bogdan Nechyporenko <bnechyporenko@bol.com>
Date: Sat, 14 Oct 2023 23:13:04 +0200
Subject: [PATCH] Made shut down stale tasks configurable.

Signed-off-by: Bogdan Nechyporenko <bnechyporenko@bol.com>
---
 .../src/ScaffolderPlugin.ts                   | 24 ++++++++++++++++++-
 .../src/scaffolder/tasks/DatabaseTaskStore.ts | 17 +++++++++----
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/plugins/scaffolder-backend/src/ScaffolderPlugin.ts b/plugins/scaffolder-backend/src/ScaffolderPlugin.ts
index beabaa50aa..b9d5d69791 100644
--- a/plugins/scaffolder-backend/src/ScaffolderPlugin.ts
+++ b/plugins/scaffolder-backend/src/ScaffolderPlugin.ts
@@ -32,7 +32,7 @@ import {
   scaffolderTaskBrokerExtensionPoint,
   scaffolderTemplatingExtensionPoint,
 } from '@backstage/plugin-scaffolder-node/alpha';
-import { createBuiltinActions } from './scaffolder';
+import { createBuiltinActions, DatabaseTaskStore } from './scaffolder';
 import { createRouter } from './service/router';
 
 /**
@@ -75,6 +75,7 @@ export const scaffolderPlugin = createBackendPlugin({
       deps: {
         logger: coreServices.logger,
         config: coreServices.rootConfig,
+        lifecycle: coreServices.lifecycle,
         reader: coreServices.urlReader,
         permissions: coreServices.permissions,
         database: coreServices.database,
@@ -84,6 +85,7 @@ export const scaffolderPlugin = createBackendPlugin({
       async init({
         logger,
         config,
+        lifecycle,
         reader,
         database,
         httpRouter,
@@ -104,6 +106,26 @@ export const scaffolderPlugin = createBackendPlugin({
           }),
         ];
 
+        lifecycle.addShutdownHook(async () => {
+          const databaseTaskStore = await DatabaseTaskStore.create({
+            database,
+          });
+          const { tasks: pendingTasks } = await databaseTaskStore.list({
+            status: 'processing',
+          });
+
+          if (pendingTasks.length > 0) {
+            await Promise.all(
+              pendingTasks.map(task =>
+                databaseTaskStore.shutdownTask({ taskId: task.id }),
+              ),
+            );
+            logger.info(
+              `Successfully shut ${pendingTasks.length} processing tasks down.`,
+            );
+          }
+        });
+
         const actionIds = actions.map(action => action.id).join(', ');
         log.info(
           `Starting scaffolder with the following actions enabled ${actionIds}`,
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
index 40ffdb17d4..8cb45240cc 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
@@ -23,10 +23,6 @@ import { ConflictError, NotFoundError } from '@backstage/errors';
 import { Knex } from 'knex';
 import { v4 as uuid } from 'uuid';
 import {
-  SerializedTaskEvent,
-  SerializedTask,
-  TaskStatus,
-  TaskEventType,
   TaskStore,
   TaskStoreEmitOptions,
   TaskStoreListEventsOptions,
@@ -34,6 +30,12 @@ import {
   TaskStoreCreateTaskResult,
   TaskStoreShutDownTaskOptions,
 } from './types';
+import {
+  SerializedTaskEvent,
+  SerializedTask,
+  TaskStatus,
+  TaskEventType,
+} from '@backstage/plugin-scaffolder-node';
 import { DateTime } from 'luxon';
 
 const migrationsDir = resolvePackagePath(
@@ -147,6 +149,7 @@ export class DatabaseTaskStore implements TaskStore {
 
   async list(options: {
     createdBy?: string;
+    status?: TaskStatus;
   }): Promise<{ tasks: SerializedTask[] }> {
     const queryBuilder = this.db<RawDbTaskRow>('tasks');
 
@@ -156,6 +159,12 @@ export class DatabaseTaskStore implements TaskStore {
       });
     }
 
+    if (options.status) {
+      queryBuilder.where({
+        status: options.status,
+      });
+    }
+
     const results = await queryBuilder.orderBy('created_at', 'desc').select();
 
     const tasks = results.map(result => ({