speed up the refresh loop

Signed-off-by: Fredrik Adelöw <freben@gmail.com>
2024-11-15 16:17:09 +01:00
parent 50e5d95e49
commit 5efde17329
8 changed files with 139 additions and 35 deletions
@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend': patch
+---
+
+Internal refactor to slightly speed up the processing loop
@@ -198,40 +198,46 @@ export class DefaultProcessingDatabase implements ProcessingDatabase {
  }

  async getProcessableEntities(
-    txOpaque: Transaction,
+    maybeTx: Transaction | Knex,
    request: { processBatchSize: number },
  ): Promise<GetProcessableEntitiesResult> {
-    const tx = txOpaque as Knex.Transaction;
+    const knex = maybeTx as Knex.Transaction | Knex;

-    let itemsQuery = tx<DbRefreshStateRow>('refresh_state').select();
+    let itemsQuery = knex<DbRefreshStateRow>('refresh_state').select([
+      'entity_id',
+      'entity_ref',
+      'unprocessed_entity',
+      'result_hash',
+      'cache',
+      'errors',
+      'location_key',
+      'next_update_at',
+    ]);

    // This avoids duplication of work because of race conditions and is
    // also fast because locked rows are ignored rather than blocking.
    // It's only available in MySQL and PostgreSQL
-    if (['mysql', 'mysql2', 'pg'].includes(tx.client.config.client)) {
+    if (['mysql', 'mysql2', 'pg'].includes(knex.client.config.client)) {
      itemsQuery = itemsQuery.forUpdate().skipLocked();
    }

    const items = await itemsQuery
-      .where('next_update_at', '<=', tx.fn.now())
+      .where('next_update_at', '<=', knex.fn.now())
      .limit(request.processBatchSize)
      .orderBy('next_update_at', 'asc');

    const interval = this.options.refreshInterval();

    const nextUpdateAt = (refreshInterval: number) => {
-      if (tx.client.config.client.includes('sqlite3')) {
-        return tx.raw(`datetime('now', ?)`, [`${refreshInterval} seconds`]);
+      if (knex.client.config.client.includes('sqlite3')) {
+        return knex.raw(`datetime('now', ?)`, [`${refreshInterval} seconds`]);
+      } else if (knex.client.config.client.includes('mysql')) {
+        return knex.raw(`now() + interval ${refreshInterval} second`);
      }
-
-      if (tx.client.config.client.includes('mysql')) {
-        return tx.raw(`now() + interval ${refreshInterval} second`);
-      }
-
-      return tx.raw(`now() + interval '${refreshInterval} seconds'`);
+      return knex.raw(`now() + interval '${refreshInterval} seconds'`);
    };

-    await tx<DbRefreshStateRow>('refresh_state')
+    await knex<DbRefreshStateRow>('refresh_state')
      .whereIn(
        'entity_ref',
        items.map(i => i.entity_ref),
@@ -247,16 +253,12 @@ export class DefaultProcessingDatabase implements ProcessingDatabase {
            id: i.entity_id,
            entityRef: i.entity_ref,
            unprocessedEntity: JSON.parse(i.unprocessed_entity) as Entity,
-            processedEntity: i.processed_entity
-              ? (JSON.parse(i.processed_entity) as Entity)
-              : undefined,
            resultHash: i.result_hash || '',
            nextUpdateAt: timestampToDateTime(i.next_update_at),
-            lastDiscoveryAt: timestampToDateTime(i.last_discovery_at),
            state: i.cache ? JSON.parse(i.cache) : undefined,
            errors: i.errors,
            locationKey: i.location_key,
-          } as RefreshStateItem),
+          } satisfies RefreshStateItem),
      ),
    };
  }
@@ -23,6 +23,7 @@ import {
 } from '@backstage/plugin-catalog-node';
 import { DbRelationsRow } from './tables';
 import { RefreshKeyData } from '../processing/types';
+import { Knex } from 'knex';

 /**
 * An abstraction for transactions of the underlying database technology.
@@ -59,10 +60,8 @@ export type RefreshStateItem = {
  id: string;
  entityRef: string;
  unprocessedEntity: Entity;
-  processedEntity?: Entity;
  resultHash: string;
  nextUpdateAt: DateTime;
-  lastDiscoveryAt: DateTime; // remove?
  state?: JsonObject;
  errors?: string;
  locationKey?: string;
@@ -116,7 +115,7 @@ export interface ProcessingDatabase {
  transaction<T>(fn: (tx: Transaction) => Promise<T>): Promise<T>;

  getProcessableEntities(
-    txOpaque: Transaction,
+    txOpaque: Transaction | Knex,
    request: { processBatchSize: number },
  ): Promise<GetProcessableEntitiesResult>;

@@ -92,7 +92,6 @@ describe('DefaultCatalogProcessingEngine', () => {
            resultHash: '',
            state: [] as any,
            nextUpdateAt: DateTime.now(),
-            lastDiscoveryAt: DateTime.now(),
          },
        ],
      });
@@ -161,7 +160,6 @@ describe('DefaultCatalogProcessingEngine', () => {
            resultHash: '',
            state: { cache: { myProcessor: { myKey: 'myValue' } } },
            nextUpdateAt: DateTime.now(),
-            lastDiscoveryAt: DateTime.now(),
          },
        ],
      });
@@ -135,13 +135,10 @@ export class DefaultCatalogProcessingEngine {
      pollingIntervalMs: this.pollingIntervalMs,
      loadTasks: async count => {
        try {
-          const { items } = await this.processingDatabase.transaction(
-            async tx => {
-              return this.processingDatabase.getProcessableEntities(tx, {
-                processBatchSize: count,
-              });
-            },
-          );
+          const { items } =
+            await this.processingDatabase.getProcessableEntities(this.knex, {
+              processBatchSize: count,
+            });
          return items;
        } catch (error) {
          this.logger.warn('Failed to load processing items', error);
@@ -28,7 +28,6 @@ import {
 import { catalogProcessingExtensionPoint } from '@backstage/plugin-catalog-node/alpha';
 import { Knex } from 'knex';
 import { default as catalogPlugin } from '../..';
-import { applyDatabaseMigrations } from '../../database/migrations';
 import {
  SyntheticLoadEntitiesProcessor,
  SyntheticLoadEntitiesProvider,
@@ -57,8 +56,6 @@ async function createBackend(
  numberOfEntities: number;
  stop: () => Promise<void>;
 }> {
-  await applyDatabaseMigrations(knex);
-
  const numberOfEntities =
    load.baseEntitiesCount + load.baseEntitiesCount * load.childrenCount;

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2024 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 { TestDatabases, mockServices } from '@backstage/backend-test-utils';
+import { Knex } from 'knex';
+import { DefaultProcessingDatabase } from '../../database/DefaultProcessingDatabase';
+import { applyDatabaseMigrations } from '../../database/migrations';
+import { describePerformanceTest, performanceTraceEnabled } from './lib/env';
+
+// #region Helpers
+
+jest.setTimeout(600_000);
+
+const databases = TestDatabases.create({
+  ids: [/* 'MYSQL_8', */ 'POSTGRES_16', /* 'POSTGRES_12',*/ 'SQLITE_3'],
+  disableDocker: false,
+});
+
+const traceLog: typeof console.log = performanceTraceEnabled
+  ? console.log
+  : () => {};
+
+async function setupDatabase(knex: Knex): Promise<void> {
+  await applyDatabaseMigrations(knex);
+
+  traceLog(`Creating test dataset`);
+
+  const now = knex.fn.now();
+  const largeEntity = `"${'x'.repeat(3 * 1024)}"`;
+
+  for (let i = 0; i < 100; ++i) {
+    await knex.batchInsert(
+      'refresh_state',
+      new Array(100).fill(null).map((_, j) => ({
+        entity_id: `id-${i}-${j}`,
+        entity_ref: `entity-${i}-${j}`,
+        unprocessed_entity: largeEntity,
+        processed_entity: largeEntity,
+        errors: '{}',
+        next_update_at: now,
+        last_discovery_at: now,
+      })),
+    );
+  }
+}
+
+// #endregion
+// #region Tests
+
+describePerformanceTest('getProcessableEntities', () => {
+  let knex: Knex;
+
+  describe.each(databases.eachSupportedId())('%p', databaseId => {
+    beforeAll(async () => {
+      knex = await databases.init(databaseId);
+      await setupDatabase(knex);
+    });
+
+    afterAll(async () => {
+      await knex.destroy();
+    });
+
+    it.each([2, 5, 10])(
+      'reads as fast as possible, batch size %p',
+      async processBatchSize => {
+        const sut = new DefaultProcessingDatabase({
+          database: knex,
+          logger: mockServices.logger.mock(),
+          refreshInterval: () => 0,
+        });
+
+        const start = Date.now();
+        let total = 0;
+        while (total < 10000) {
+          const result = await sut.getProcessableEntities(knex, {
+            processBatchSize,
+          });
+          total += result.items.length;
+        }
+
+        const perSecond = Math.round(total / ((Date.now() - start) / 1000));
+        traceLog(
+          `${databaseId} processed ${perSecond} entities per second at a batch size of ${processBatchSize}`,
+        );
+
+        expect(true).toBe(true);
+      },
+    );
+  });
+});
+
+// #endregion
@@ -22,6 +22,7 @@ import {
 } from '@backstage/backend-test-utils';
 import { catalogProcessingExtensionPoint } from '@backstage/plugin-catalog-node/alpha';
 import { Knex } from 'knex';
+import { default as catalogPlugin } from '../..';
 import { applyDatabaseMigrations } from '../../database/migrations';
 import {
  SyntheticLoadEntitiesProcessor,
@@ -176,7 +177,7 @@ describePerformanceTest('stitchingPerformance', () => {

      const backend = await startTestBackend({
        features: [
-          import('@backstage/plugin-catalog-backend/alpha'),
+          catalogPlugin,
          mockServices.rootConfig.factory({ data: config }),
          mockServices.database.factory({ knex }),
          createBackendModule({