Merge pull request #13817 from howlowck/howlowck/non-blocking

Scaffolder: `concurrentTasksLimit` for TaskWorker

.changeset/little-chairs-dance.md

@@ -0,0 +1,21 @@
+---
+'@backstage/plugin-scaffolder-backend': minor
+---
+
+Deprecated the `taskWorkers` option in RouterOptions in favor of `concurrentTasksLimit` which sets the limit of concurrent tasks in a single TaskWorker
+
+TaskWorker can now run multiple (defaults to 10) tasks concurrently using the `concurrentTasksLimit` option available in both `RouterOptions` and `CreateWorkerOptions`.
+
+To use the option to create a TaskWorker:
+
+```diff
+const worker = await TaskWorker.create({
+    taskBroker,
+    actionRegistry,
+    integrations,
+    logger,
+    workingDirectory,
+    additionalTemplateFilters,
++   concurrentTasksLimit: 10 // (1 to Infinity)
+});
+```

plugins/scaffolder-backend/api-report.md

@@ -479,6 +479,7 @@ export type CreateWorkerOptions = {
   workingDirectory: string;
   logger: Logger;
   additionalTemplateFilters?: Record<string, TemplateFilter>;
+  concurrentTasksLimit?: number;
   additionalTemplateGlobals?: Record<string, TemplateGlobal>;
 };
 
@@ -573,6 +574,7 @@ export interface RouterOptions {
   additionalTemplateGlobals?: Record<string, TemplateGlobal>;
   // (undocumented)
   catalogClient: CatalogApi;
+  concurrentTasksLimit?: number;
   // (undocumented)
   config: Config;
   // (undocumented)
@@ -587,7 +589,7 @@ export interface RouterOptions {
   scheduler?: PluginTaskScheduler;
   // (undocumented)
   taskBroker?: TaskBroker;
-  // (undocumented)
+  // @deprecated (undocumented)
   taskWorkers?: number;
 }
 
@@ -819,6 +821,8 @@ export class TaskWorker {
   // (undocumented)
   static create(options: CreateWorkerOptions): Promise<TaskWorker>;
   // (undocumented)
+  protected onReadyToClaimTask(): Promise<void>;
+  // (undocumented)
   runOneTask(task: TaskContext): Promise<void>;
   // (undocumented)
   start(): void;

plugins/scaffolder-backend/package.json

@@ -72,6 +72,7 @@
     "octokit": "^2.0.0",
     "octokit-plugin-create-pull-request": "^3.10.0",
     "p-limit": "^3.1.0",
+    "p-queue": "^6.6.2",
     "prom-client": "^14.0.1",
     "uuid": "^8.2.0",
     "vm2": "^3.9.11",

plugins/scaffolder-backend/src/scaffolder/tasks/TaskWorker.test.ts

@@ -14,14 +14,16 @@
  * limitations under the License.
  */
 
+import os from 'os';
 import { getVoidLogger, DatabaseManager } from '@backstage/backend-common';
 import { ConfigReader } from '@backstage/config';
 import { DatabaseTaskStore } from './DatabaseTaskStore';
 import { StorageTaskBroker } from './StorageTaskBroker';
-import { TaskWorker } from './TaskWorker';
+import { TaskWorker, TaskWorkerOptions } from './TaskWorker';
 import { ScmIntegrations } from '@backstage/integration';
 import { TemplateActionRegistry } from '../actions';
 import { NunjucksWorkflowRunner } from './NunjucksWorkflowRunner';
+import { TaskBroker, TaskContext, WorkflowRunner } from './types';
 
 jest.mock('./NunjucksWorkflowRunner');
 const MockedNunjucksWorkflowRunner =
@@ -127,3 +129,130 @@ describe('TaskWorker', () => {
     expect(event?.body.output).toEqual({ testOutput: 'testmockoutput' });
   });
 });
+
+describe('Concurrent TaskWorker', () => {
+  let storage: DatabaseTaskStore;
+
+  const integrations: ScmIntegrations = {} as ScmIntegrations;
+
+  const actionRegistry: TemplateActionRegistry = {} as TemplateActionRegistry;
+  const workingDirectory = os.tmpdir();
+  let asyncTasksCount = 0;
+
+  const workflowRunner: NunjucksWorkflowRunner = {
+    execute: () => {
+      asyncTasksCount++;
+      return new Promise(resolve => {
+        setTimeout(() => {
+          resolve({ output: { testOutput: 'testmockoutput' } });
+        }, 1000);
+      });
+    },
+  } as unknown as NunjucksWorkflowRunner;
+
+  beforeAll(async () => {
+    storage = await createStore();
+  });
+
+  beforeEach(() => {
+    asyncTasksCount = 0;
+    jest.resetAllMocks();
+    MockedNunjucksWorkflowRunner.mockImplementation(() => workflowRunner);
+  });
+
+  const logger = getVoidLogger();
+
+  it('should be able to run multiple tasks at once', async () => {
+    const broker = new StorageTaskBroker(storage, logger);
+
+    const dispatchANewTask = () =>
+      broker.dispatch({
+        spec: {
+          apiVersion: 'scaffolder.backstage.io/v1beta3',
+          steps: [{ id: 'test', name: 'test', action: 'not-found-action' }],
+          output: {
+            result: '{{ steps.test.output.testOutput }}',
+          },
+          parameters: {},
+        },
+      });
+
+    const expectedConcurrentTasks = 3;
+    const taskWorker = await TaskWorker.create({
+      logger,
+      workingDirectory,
+      integrations,
+      taskBroker: broker,
+      actionRegistry,
+      concurrentTasksLimit: expectedConcurrentTasks,
+    });
+
+    taskWorker.start();
+
+    await dispatchANewTask();
+    await dispatchANewTask();
+    await dispatchANewTask();
+    await dispatchANewTask();
+
+    expect(asyncTasksCount).toEqual(expectedConcurrentTasks);
+  });
+});
+
+describe('TaskWorker internals', () => {
+  const TaskWorkerConstructor = TaskWorker as unknown as {
+    new (options: TaskWorkerOptions): TaskWorker;
+  };
+
+  it('should not pick up tasks before it is ready to execute more work', async () => {
+    const inflightTasks = new Array<{
+      task: TaskContext;
+      resolve: () => void;
+    }>();
+    const workflowRunner: WorkflowRunner = {
+      async execute(task) {
+        await new Promise<void>(resolve => {
+          inflightTasks.push({ task, resolve });
+        });
+        return {
+          output: {},
+        };
+      },
+    };
+
+    let claimedTaskCount = 0;
+    const taskWorker = new TaskWorkerConstructor({
+      runners: { workflowRunner },
+      taskBroker: {
+        async claim() {
+          claimedTaskCount++;
+          return {
+            spec: {
+              apiVersion: 'scaffolder.backstage.io/v1beta3',
+            },
+            createdBy: `test-${claimedTaskCount}`,
+            async complete(_result, _metadata) {},
+          } as TaskContext;
+        },
+      } as TaskBroker,
+      concurrentTasksLimit: 2,
+    });
+
+    expect(claimedTaskCount).toBe(0);
+    taskWorker.start();
+
+    // This will wait for all higher priority promise ticks to complete
+    await new Promise(resolve => setTimeout(resolve));
+
+    // Once we start the worker it should pick up 2 tasks, since that's our limit
+    expect(claimedTaskCount).toBe(2);
+    expect(inflightTasks.length).toBe(2);
+
+    // This completes the first task, making space for one more
+    inflightTasks.shift()?.resolve();
+    await new Promise(resolve => setTimeout(resolve));
+
+    // We now expect one more task to have been claimed, and two tasks in the queue again
+    expect(claimedTaskCount).toBe(3);
+    expect(inflightTasks.length).toBe(2);
+  });
+});

plugins/scaffolder-backend/src/scaffolder/tasks/TaskWorker.ts

@@ -15,6 +15,7 @@
  */
 
 import { TaskContext, TaskBroker, WorkflowRunner } from './types';
+import PQueue from 'p-queue';
 import { NunjucksWorkflowRunner } from './NunjucksWorkflowRunner';
 import { Logger } from 'winston';
 import { TemplateActionRegistry } from '../actions';
@@ -35,6 +36,7 @@ export type TaskWorkerOptions = {
   runners: {
     workflowRunner: WorkflowRunner;
   };
+  concurrentTasksLimit: number;
 };
 
 /**
@@ -49,6 +51,19 @@ export type CreateWorkerOptions = {
   workingDirectory: string;
   logger: Logger;
   additionalTemplateFilters?: Record<string, TemplateFilter>;
+  /**
+   * The number of tasks that can be executed at the same time by the worker
+   * @defaultValue 10
+   * @example
+   * ```
+   * {
+   *   concurrentTasksLimit: 1,
+   *   // OR
+   *   concurrentTasksLimit: Infinity
+   * }
+   * ```
+   */
+  concurrentTasksLimit?: number;
   additionalTemplateGlobals?: Record<string, TemplateGlobal>;
 };
 
@@ -60,6 +75,10 @@ export type CreateWorkerOptions = {
 export class TaskWorker {
   private constructor(private readonly options: TaskWorkerOptions) {}
 
+  private taskQueue: PQueue = new PQueue({
+    concurrency: this.options.concurrentTasksLimit,
+  });
+
   static async create(options: CreateWorkerOptions): Promise<TaskWorker> {
     const {
       taskBroker,
@@ -68,6 +87,7 @@ export class TaskWorker {
       integrations,
       workingDirectory,
       additionalTemplateFilters,
+      concurrentTasksLimit = 10, // from 1 to Infinity
       additionalTemplateGlobals,
     } = options;
 
@@ -83,18 +103,33 @@ export class TaskWorker {
     return new TaskWorker({
       taskBroker: taskBroker,
       runners: { workflowRunner },
+      concurrentTasksLimit,
     });
   }
 
   start() {
     (async () => {
       for (;;) {
+        await this.onReadyToClaimTask();
         const task = await this.options.taskBroker.claim();
-        await this.runOneTask(task);
+        this.taskQueue.add(() => this.runOneTask(task));
       }
     })();
   }
 
+  protected onReadyToClaimTask(): Promise<void> {
+    if (this.taskQueue.pending < this.options.concurrentTasksLimit) {
+      return Promise.resolve();
+    }
+    return new Promise(resolve => {
+      // "next" event emits when a task completes
+      // https://github.com/sindresorhus/p-queue#next
+      this.taskQueue.once('next', () => {
+        resolve();
+      });
+    });
+  }
+
   async runOneTask(task: TaskContext) {
     try {
       if (task.spec.apiVersion !== 'scaffolder.backstage.io/v1beta3') {

plugins/scaffolder-backend/src/service/router.ts

@@ -68,7 +68,16 @@ export interface RouterOptions {
   scheduler?: PluginTaskScheduler;
 
   actions?: TemplateAction<any>[];
+  /**
+   * @deprecated taskWorkers is deprecated in favor of concurrentTasksLimit option with a single TaskWorker
+   * @defaultValue 1
+   */
   taskWorkers?: number;
+  /**
+   * Sets the number of concurrent tasks that can be run at any given time on the TaskWorker
+   * @defaultValue 10
+   */
+  concurrentTasksLimit?: number;
   taskBroker?: TaskBroker;
   additionalTemplateFilters?: Record<string, TemplateFilter>;
   additionalTemplateGlobals?: Record<string, TemplateGlobal>;
@@ -161,6 +170,7 @@ export async function createRouter(
     catalogClient,
     actions,
     taskWorkers,
+    concurrentTasksLimit,
     scheduler,
     additionalTemplateFilters,
     additionalTemplateGlobals,
@@ -203,7 +213,7 @@ export async function createRouter(
   const actionRegistry = new TemplateActionRegistry();
 
   const workers = [];
-  for (let i = 0; i < (taskWorkers || 3); i++) {
+  for (let i = 0; i < (taskWorkers || 1); i++) {
     const worker = await TaskWorker.create({
       taskBroker,
       actionRegistry,
@@ -212,6 +222,7 @@ export async function createRouter(
       workingDirectory,
       additionalTemplateFilters,
       additionalTemplateGlobals,
+      concurrentTasksLimit,
     });
     workers.push(worker);
   }

yarn.lock

@@ -7409,6 +7409,7 @@ __metadata:
     octokit: ^2.0.0
     octokit-plugin-create-pull-request: ^3.10.0
     p-limit: ^3.1.0
+    p-queue: ^6.6.2
     prom-client: ^14.0.1
     supertest: ^6.1.3
     uuid: ^8.2.0