diff --git a/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.test.ts b/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.test.ts index 8ef685e6b0..3d1f887e1f 100644 --- a/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.test.ts +++ b/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.test.ts @@ -14,8 +14,9 @@ * limitations under the License. */ +import { getVoidLogger } from '@backstage/backend-common'; import type { Entity } from '@backstage/catalog-model'; -import type { Database } from '../database'; +import { Database, DatabaseManager } from '../database'; import { DatabaseEntitiesCatalog } from './DatabaseEntitiesCatalog'; describe('DatabaseEntitiesCatalog', () => { @@ -58,7 +59,7 @@ describe('DatabaseEntitiesCatalog', () => { db.entities.mockResolvedValue([]); db.addEntities.mockResolvedValue([{ entity }]); - const catalog = new DatabaseEntitiesCatalog(db); + const catalog = new DatabaseEntitiesCatalog(db, getVoidLogger()); const result = await catalog.addOrUpdateEntity(entity); expect(db.entityByName).toHaveBeenCalledTimes(1); @@ -97,7 +98,7 @@ describe('DatabaseEntitiesCatalog', () => { }); db.updateEntity.mockResolvedValue({ entity }); - const catalog = new DatabaseEntitiesCatalog(db); + const catalog = new DatabaseEntitiesCatalog(db, getVoidLogger()); const result = await catalog.addOrUpdateEntity(entity); expect(db.entities).toHaveBeenCalledTimes(0); @@ -149,7 +150,7 @@ describe('DatabaseEntitiesCatalog', () => { db.entityByName.mockResolvedValue({ entity: existing }); db.updateEntity.mockResolvedValue({ entity: existing }); - const catalog = new DatabaseEntitiesCatalog(db); + const catalog = new DatabaseEntitiesCatalog(db, getVoidLogger()); const result = await catalog.addOrUpdateEntity(added); expect(db.entityByName).toHaveBeenCalledTimes(1); @@ -180,4 +181,38 @@ describe('DatabaseEntitiesCatalog', () => { expect(result).toEqual(existing); }); }); + + describe('batchAddOrUpdateEntities', () => { + it('both adds and updates', async () => { + const catalog = new DatabaseEntitiesCatalog( + await DatabaseManager.createTestDatabase(), + getVoidLogger(), + ); + const entities: Entity[] = []; + for (let i = 0; i < 500; ++i) { + entities.push({ + apiVersion: 'a', + kind: 'k', + metadata: { name: `n${i}` }, + }); + } + + await catalog.batchAddOrUpdateEntities(entities); + const afterFirst = await catalog.entities(); + expect(afterFirst.length).toBe(500); + + entities[40].metadata.op = 'changed'; + entities.push({ + apiVersion: 'a', + kind: 'k', + metadata: { name: `n500`, op: 'added' }, + }); + + await catalog.batchAddOrUpdateEntities(entities); + const afterSecond = await catalog.entities(); + expect(afterSecond.length).toBe(501); + expect(afterSecond.find(e => e.metadata.op === 'changed')).toBeDefined(); + expect(afterSecond.find(e => e.metadata.op === 'added')).toBeDefined(); + }); + }); }); diff --git a/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.ts b/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.ts index 8619355c91..0a213dd186 100644 --- a/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.ts +++ b/plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.ts @@ -14,18 +14,46 @@ * limitations under the License. */ -import { NotFoundError } from '@backstage/backend-common'; +import { ConflictError, NotFoundError } from '@backstage/backend-common'; import type { Entity } from '@backstage/catalog-model'; import { + entityHasChanges, generateUpdatedEntity, getEntityName, LOCATION_ANNOTATION, + serializeEntityRef, } from '@backstage/catalog-model'; +import { chunk, groupBy } from 'lodash'; +import limiterFactory from 'p-limit'; +import { Logger } from 'winston'; import type { Database, DbEntityResponse, EntityFilters } from '../database'; +import { durationText } from '../util/timing'; import type { EntitiesCatalog } from './types'; +type BatchContext = { + kind: string; + namespace: string; + locationId?: string; +}; + +// Some locations return tens or hundreds of thousands of entities. To make +// those payloads more manageable, we break work apart in batches of this +// many entities and write them to storage per batch. +const BATCH_SIZE = 100; + +// When writing large batches, there's an increasing chance of contention in +// the form of conflicts where we compete with other writes. Each batch gets +// this many attempts at being written before giving up. +const BATCH_ATTEMPTS = 3; + +// The number of batches that may be ongoing at the same time. +const BATCH_CONCURRENCY = 3; + export class DatabaseEntitiesCatalog implements EntitiesCatalog { - constructor(private readonly database: Database) {} + constructor( + private readonly database: Database, + private readonly logger: Logger, + ) {} async entities(filters?: EntityFilters): Promise { const items = await this.database.transaction(tx => @@ -102,4 +130,139 @@ export class DatabaseEntitiesCatalog implements EntitiesCatalog { return undefined; }); } + + /** + * Writes a number of entities efficiently to storage. + * + * @param entities Some entities + * @param locationId The location that they all belong to + */ + async batchAddOrUpdateEntities(entities: Entity[], locationId?: string) { + // Group the entities by unique kind+namespace combinations + const entitiesByKindAndNamespace = groupBy(entities, entity => { + const name = getEntityName(entity); + return `${name.kind}:${name.namespace}`.toLowerCase(); + }); + + const limiter = limiterFactory(BATCH_CONCURRENCY); + const tasks: Promise[] = []; + + for (const groupEntities of Object.values(entitiesByKindAndNamespace)) { + const { kind, namespace } = getEntityName(groupEntities[0]); + + // Go through the new entities in reasonable chunk sizes (sometimes, + // sources produce tens of thousands of entities, and those are too large + // batch sizes to reasonably send to the database) + for (const batch of chunk(groupEntities, BATCH_SIZE)) { + tasks.push( + limiter(async () => { + const first = serializeEntityRef(batch[0]); + const last = serializeEntityRef(batch[batch.length - 1]); + this.logger.debug( + `Considering batch ${first}-${last} (${batch.length} entries)`, + ); + + // Retry the batch write a few times to deal with contention + const context = { kind, namespace, locationId }; + for (let attempt = 1; attempt <= BATCH_ATTEMPTS; ++attempt) { + try { + const { toAdd, toUpdate } = await this.analyzeBatch( + batch, + context, + ); + if (toAdd.length) await this.batchAdd(toAdd, context); + if (toUpdate.length) await this.batchUpdate(toUpdate, context); + break; + } catch (e) { + if (e instanceof ConflictError && attempt < BATCH_ATTEMPTS) { + this.logger.warn( + `Failed to write batch at attempt ${attempt}/${BATCH_ATTEMPTS}, ${e}`, + ); + } else { + throw e; + } + } + } + }), + ); + } + } + + await Promise.all(tasks); + } + + // Given a batch of entities that were just read from a location, take them + // into consideration by comparing against the existing catalog entities and + // produce the list of entities to be added, and the list of entities to be + // updated + private async analyzeBatch( + newEntities: Entity[], + { kind, namespace }: BatchContext, + ): Promise<{ + toAdd: Entity[]; + toUpdate: Entity[]; + }> { + const markTimestamp = process.hrtime(); + + const names = newEntities.map(e => e.metadata.name); + const oldEntities = await this.entities({ + kind: kind, + 'metadata.namespace': namespace, + 'metadata.name': names, + }); + + const oldEntitiesByName = new Map( + oldEntities.map(e => [e.metadata.name, e]), + ); + + const toAdd: Entity[] = []; + const toUpdate: Entity[] = []; + + for (const newEntity of newEntities) { + const oldEntity = oldEntitiesByName.get(newEntity.metadata.name); + if (!oldEntity) { + toAdd.push(newEntity); + } else if (entityHasChanges(oldEntity, newEntity)) { + // TODO(freben): This currently uses addOrUpdateEntity under the hood, + // but should probably calculate the end result entity right here + // instead and call a dedicated batch update database method instead + toUpdate.push(newEntity); + } + } + + this.logger.debug( + `Found ${toAdd.length} entities to add, ${ + toUpdate.length + } entities to update in ${durationText(markTimestamp)}`, + ); + + return { toAdd, toUpdate }; + } + + // Efficiently adds the given entities to storage, under the assumption that + // they do not conflict with any existing entities + private async batchAdd(entities: Entity[], { locationId }: BatchContext) { + const markTimestamp = process.hrtime(); + + await this.addEntities(entities, locationId); + + this.logger.debug( + `Added ${entities.length} entities in ${durationText(markTimestamp)}`, + ); + } + + // Efficiently updates the given entities into storage, under the assumption + // that there already exist entities with the same names + private async batchUpdate(entities: Entity[], { locationId }: BatchContext) { + const markTimestamp = process.hrtime(); + + // TODO(freben): Still not batched + for (const entity of entities) { + await this.addOrUpdateEntity(entity, locationId); + } + + this.logger.debug( + `Updated ${entities.length} entities in ${durationText(markTimestamp)}`, + ); + } } diff --git a/plugins/catalog-backend/src/catalog/types.ts b/plugins/catalog-backend/src/catalog/types.ts index 810d51a6c1..ca73eddce9 100644 --- a/plugins/catalog-backend/src/catalog/types.ts +++ b/plugins/catalog-backend/src/catalog/types.ts @@ -26,6 +26,17 @@ export type EntitiesCatalog = { addOrUpdateEntity(entity: Entity, locationId?: string): Promise; addEntities(entities: Entity[], locationId?: string): Promise; removeEntityByUid(uid: string): Promise; + + /** + * Writes a number of entities efficiently to storage. + * + * @param entities Some entities + * @param locationId The location that they all belong to + */ + batchAddOrUpdateEntities( + entities: Entity[], + locationId?: string, + ): Promise; }; // diff --git a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.test.ts b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.test.ts index 909139d678..7fe164d05f 100644 --- a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.test.ts +++ b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.test.ts @@ -15,12 +15,7 @@ */ import { getVoidLogger } from '@backstage/backend-common'; -import { - Entity, - ENTITY_DEFAULT_NAMESPACE, - Location, - LocationSpec, -} from '@backstage/catalog-model'; +import { Entity, Location, LocationSpec } from '@backstage/catalog-model'; import { EntitiesCatalog, LocationsCatalog } from '../catalog'; import { LocationUpdateStatus } from '../catalog/types'; import { DatabaseLocationUpdateLogStatus } from '../database/types'; @@ -39,6 +34,7 @@ describe('HigherOrderOperations', () => { addOrUpdateEntity: jest.fn(), addEntities: jest.fn(), removeEntityByUid: jest.fn(), + batchAddOrUpdateEntities: jest.fn(), }; locationsCatalog = { addLocation: jest.fn(), @@ -190,8 +186,7 @@ describe('HigherOrderOperations', () => { entities: [{ entity: desc, location }], errors: [], }); - entitiesCatalog.entities.mockResolvedValue([]); - entitiesCatalog.addEntities.mockResolvedValue(undefined); + entitiesCatalog.batchAddOrUpdateEntities.mockResolvedValue(undefined); await expect( higherOrderOperation.refreshAllLocations(), @@ -203,14 +198,8 @@ describe('HigherOrderOperations', () => { type: 'some', target: 'thing', }); - expect(entitiesCatalog.entities).toHaveBeenCalledTimes(1); - expect(entitiesCatalog.entities).toHaveBeenNthCalledWith(1, { - kind: 'Component', - 'metadata.namespace': ENTITY_DEFAULT_NAMESPACE, - 'metadata.name': ['c1'], - }); - expect(entitiesCatalog.addEntities).toHaveBeenCalledTimes(1); - expect(entitiesCatalog.addEntities).toHaveBeenNthCalledWith( + expect(entitiesCatalog.batchAddOrUpdateEntities).toHaveBeenCalledTimes(1); + expect(entitiesCatalog.batchAddOrUpdateEntities).toHaveBeenNthCalledWith( 1, [expect.objectContaining({ metadata: { name: 'c1' } })], '123', diff --git a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts index fa73c58348..6df6d9b4b7 100644 --- a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts +++ b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts @@ -14,17 +14,8 @@ * limitations under the License. */ -import { ConflictError, InputError } from '@backstage/backend-common'; -import { - Entity, - entityHasChanges, - getEntityName, - Location, - LocationSpec, - serializeEntityRef, -} from '@backstage/catalog-model'; -import { chunk, groupBy } from 'lodash'; -import limiterFactory from 'p-limit'; +import { InputError } from '@backstage/backend-common'; +import { Entity, Location, LocationSpec } from '@backstage/catalog-model'; import { v4 as uuidv4 } from 'uuid'; import { Logger } from 'winston'; import { EntitiesCatalog, LocationsCatalog } from '../catalog'; @@ -35,25 +26,6 @@ import { LocationReader, } from './types'; -type BatchContext = { - kind: string; - namespace: string; - location: Location; -}; - -// Some locations return tens or hundreds of thousands of entities. To make -// those payloads more manageable, we break work apart in batches of this -// many entities and write them to storage per batch. -const BATCH_SIZE = 100; - -// When writing large batches, there's an increasing chance of contention in -// the form of conflicts where we compete with other writes. Each batch gets -// this many attempts at being written before giving up. -const BATCH_ATTEMPTS = 3; - -// The number of batches that may be ongoing at the same time. -const BATCH_CONCURRENCY = 3; - /** * Placeholder for operations that span several catalogs and/or stretches out * in time. @@ -192,10 +164,30 @@ export class HigherOrderOperations implements HigherOrderOperation { startTimestamp = process.hrtime(); - await this.batchAddOrUpdateEntities( - readerOutput.entities.map(e => e.entity), - location, - ); + try { + await this.entitiesCatalog.batchAddOrUpdateEntities( + readerOutput.entities.map(e => e.entity), + location.id, + ); + } catch (e) { + for (const entity of readerOutput.entities) { + await this.locationsCatalog.logUpdateFailure( + location.id, + e, + entity.entity.metadata.name, + ); + } + throw e; + } + + this.logger.info(`Posting update success markers`); + + for (const entity of readerOutput.entities) { + await this.locationsCatalog.logUpdateSuccess( + location.id, + entity.entity.metadata.name, + ); + } this.logger.info( `Wrote ${readerOutput.entities.length} entities from location ${ @@ -203,148 +195,4 @@ export class HigherOrderOperations implements HigherOrderOperation { }:${location.target} in ${durationText(startTimestamp)}`, ); } - - /** - * Writes a number of entities efficiently to storage. - * - * @param entities Some entities - * @param location The location that they all belong to - */ - async batchAddOrUpdateEntities(entities: Entity[], location: Location) { - // Group the entities by unique kind+namespace combinations - const entitiesByKindAndNamespace = groupBy(entities, entity => { - const name = getEntityName(entity); - return `${name.kind}:${name.namespace}`.toLowerCase(); - }); - - const limiter = limiterFactory(BATCH_CONCURRENCY); - const tasks: Promise[] = []; - - for (const groupEntities of Object.values(entitiesByKindAndNamespace)) { - const { kind, namespace } = getEntityName(groupEntities[0]); - - // Go through the new entities in reasonable chunk sizes (sometimes, - // sources produce tens of thousands of entities, and those are too large - // batch sizes to reasonably send to the database) - for (const batch of chunk(groupEntities, BATCH_SIZE)) { - tasks.push( - limiter(async () => { - const first = serializeEntityRef(batch[0]); - const last = serializeEntityRef(batch[batch.length - 1]); - this.logger.debug( - `Considering batch ${first}-${last} (${batch.length} entries)`, - ); - - // Retry the batch write a few times to deal with contention - const context = { kind, namespace, location }; - for (let attempt = 1; attempt <= BATCH_ATTEMPTS; ++attempt) { - try { - const { toAdd, toUpdate } = await this.analyzeBatch( - batch, - context, - ); - if (toAdd.length) await this.batchAdd(toAdd, context); - if (toUpdate.length) await this.batchUpdate(toUpdate, context); - break; - } catch (e) { - if (e instanceof ConflictError && attempt < BATCH_ATTEMPTS) { - this.logger.warn( - `Failed to write batch at attempt ${attempt}/${BATCH_ATTEMPTS}, ${e}`, - ); - } else { - throw e; - } - } - } - }), - ); - } - } - - await Promise.all(tasks); - } - - // Given a batch of entities that were just read from a location, take them - // into consideration by comparing against the existing catalog entities and - // produce the list of entities to be added, and the list of entities to be - // updated - private async analyzeBatch( - newEntities: Entity[], - { kind, namespace }: BatchContext, - ): Promise<{ - toAdd: Entity[]; - toUpdate: Entity[]; - }> { - const markTimestamp = process.hrtime(); - - const oldEntities = await this.entitiesCatalog.entities({ - kind: kind, - 'metadata.namespace': namespace, - 'metadata.name': newEntities.map(e => e.metadata.name), - }); - - const oldEntitiesByName = new Map( - oldEntities.map(e => [e.metadata.name, e]), - ); - - const toAdd: Entity[] = []; - const toUpdate: Entity[] = []; - - for (const newEntity of newEntities) { - const oldEntity = oldEntitiesByName.get(newEntity.metadata.name); - if (!oldEntity) { - toAdd.push(newEntity); - } else if (entityHasChanges(oldEntity, newEntity)) { - toUpdate.push(newEntity); - } - } - - this.logger.debug( - `Found ${toAdd.length} entities to add, ${ - toUpdate.length - } entities to update in ${durationText(markTimestamp)}`, - ); - - return { toAdd, toUpdate }; - } - - // Efficiently adds the given entities to storage, under the assumption that - // they do not conflict with any existing entities - private async batchAdd(entities: Entity[], { location }: BatchContext) { - const markTimestamp = process.hrtime(); - - await this.entitiesCatalog.addEntities(entities, location.id); - - // TODO(freben): Still not batched - for (const entity of entities) { - await this.locationsCatalog.logUpdateSuccess( - location.id, - entity.metadata.name, - ); - } - - this.logger.debug( - `Added ${entities.length} entities in ${durationText(markTimestamp)}`, - ); - } - - // Efficiently updates the given entities into storage, under the assumption - // that there already exist entities with the same names - private async batchUpdate(entities: Entity[], { location }: BatchContext) { - const markTimestamp = process.hrtime(); - - // TODO(freben): Still not batched - for (const entity of entities) { - await this.entitiesCatalog.addOrUpdateEntity(entity); - - await this.locationsCatalog.logUpdateSuccess( - location.id, - entity.metadata.name, - ); - } - - this.logger.debug( - `Updated ${entities.length} entities in ${durationText(markTimestamp)}`, - ); - } } diff --git a/plugins/catalog-backend/src/service/CatalogBuilder.ts b/plugins/catalog-backend/src/service/CatalogBuilder.ts index e02005b817..e1d76bc3cd 100644 --- a/plugins/catalog-backend/src/service/CatalogBuilder.ts +++ b/plugins/catalog-backend/src/service/CatalogBuilder.ts @@ -370,7 +370,7 @@ export class CatalogBuilder { { logger }, ); - const entitiesCatalog = new DatabaseEntitiesCatalog(db); + const entitiesCatalog = new DatabaseEntitiesCatalog(db, this.env.logger); const locationsCatalog = new DatabaseLocationsCatalog(db); const higherOrderOperation = new HigherOrderOperations( entitiesCatalog, diff --git a/plugins/catalog-backend/src/service/router.test.ts b/plugins/catalog-backend/src/service/router.test.ts index f7ab4878f8..414f8a93de 100644 --- a/plugins/catalog-backend/src/service/router.test.ts +++ b/plugins/catalog-backend/src/service/router.test.ts @@ -35,6 +35,7 @@ describe('createRouter', () => { addOrUpdateEntity: jest.fn(), addEntities: jest.fn(), removeEntityByUid: jest.fn(), + batchAddOrUpdateEntities: jest.fn(), }; locationsCatalog = { addLocation: jest.fn(),