Merge pull request #27915 from backstage/freben/stream-entities

Implement `/entities` in terms of `queryEntities`

.changeset/fifty-countries-decide.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend': minor
+---
+
+Implement `/entities` in terms of `queryEntities` to not run into memory and performance problems on large catalogs

plugins/catalog-backend/src/catalog/types.ts

@@ -203,6 +203,7 @@ export interface QueryEntitiesInitialRequest {
     term: string;
     fields?: string[];
   };
+  skipTotalItems?: boolean;
 }
 
 /**

plugins/catalog-backend/src/service/DefaultEntitiesCatalog.test.ts

@@ -1478,6 +1478,52 @@ describe('DefaultEntitiesCatalog', () => {
       },
     );
 
+    it.each(databases.eachSupportedId())(
+      'can skip totalItems, %p',
+      async databaseId => {
+        await createDatabase(databaseId);
+
+        await Promise.all(
+          Array(15)
+            .fill(0)
+            .map(() =>
+              addEntityToSearch({
+                apiVersion: 'a',
+                kind: 'k',
+                metadata: { name: v4() },
+              }),
+            ),
+        );
+
+        const catalog = new DefaultEntitiesCatalog({
+          database: knex,
+          logger: mockServices.logger.mock(),
+          stitcher,
+        });
+
+        const request: QueryEntitiesInitialRequest = {
+          limit: 10,
+          credentials: mockCredentials.none(),
+          skipTotalItems: true,
+        };
+        let response = await catalog.queryEntities(request);
+        expect(response).toEqual({
+          totalItems: 0,
+          items: expect.objectContaining({ length: 10 }),
+          pageInfo: { nextCursor: expect.anything() },
+        });
+        response = await catalog.queryEntities({
+          ...request,
+          cursor: response.pageInfo.nextCursor!,
+        });
+        expect(response).toEqual({
+          totalItems: 0,
+          items: expect.objectContaining({ length: 5 }),
+          pageInfo: { prevCursor: expect.anything() },
+        });
+      },
+    );
+
     it.each(databases.eachSupportedId())(
       'should paginate results accordingly in case of clashing items, %p',
       async databaseId => {

plugins/catalog-backend/src/service/DefaultEntitiesCatalog.ts

@@ -14,11 +14,7 @@
  * limitations under the License.
  */
 
-import {
-  Entity,
-  parseEntityRef,
-  stringifyEntityRef,
-} from '@backstage/catalog-model';
+import { Entity, stringifyEntityRef } from '@backstage/catalog-model';
 import { InputError, NotFoundError } from '@backstage/errors';
 import { Knex } from 'knex';
 import { chunk as lodashChunk, isEqual } from 'lodash';
@@ -49,6 +45,7 @@ import {
 import { Stitcher } from '../stitching/types';
 
 import {
+  expandLegacyCompoundRelationRefsInResponse,
   isQueryEntitiesCursorRequest,
   isQueryEntitiesInitialRequest,
 } from './util';
@@ -280,27 +277,7 @@ export class DefaultEntitiesCatalog implements EntitiesCatalog {
       entities = entities.map(e => request.fields!(e));
     }
 
-    // TODO(freben): This is added as a compatibility guarantee, until we can be
-    // sure that all adopters have re-stitched their entities so that the new
-    // targetRef field is present on them, and that they have stopped consuming
-    // the now-removed old field
-    // TODO(jhaals): Remove this in April 2022
-    for (const entity of entities) {
-      if (entity.relations) {
-        for (const relation of entity.relations as any) {
-          if (!relation.targetRef && relation.target) {
-            // This is the case where an old-form entity, not yet stitched with
-            // the updated code, was in the database
-            relation.targetRef = stringifyEntityRef(relation.target);
-          } else if (!relation.target && relation.targetRef) {
-            // This is the case where a new-form entity, stitched with the
-            // updated code, was in the database but we still want to produce
-            // the old data shape as well for compatibility reasons
-            relation.target = parseEntityRef(relation.targetRef);
-          }
-        }
-      }
-    }
+    expandLegacyCompoundRelationRefsInResponse(entities);
 
     return {
       entities,
@@ -354,6 +331,7 @@ export class DefaultEntitiesCatalog implements EntitiesCatalog {
 
     const cursor: Omit<Cursor, 'orderFieldValues'> & {
       orderFieldValues?: (string | null)[];
+      skipTotalItems: boolean;
     } = {
       orderFields: [],
       isPrevious: false,
@@ -364,7 +342,8 @@ export class DefaultEntitiesCatalog implements EntitiesCatalog {
     // request. The result is then embedded into the cursor for subsequent
     // requests. Threfore this can be undefined here, but will then get
     // populated further down.
-    const shouldComputeTotalItems = cursor.totalItems === undefined;
+    const shouldComputeTotalItems =
+      cursor.totalItems === undefined && !cursor.skipTotalItems;
     const isFetchingBackwards = cursor.isPrevious;
 
     if (cursor.orderFields.length > 1) {
@@ -555,8 +534,16 @@ export class DefaultEntitiesCatalog implements EntitiesCatalog {
 
     const rows = shouldComputeTotalItems || limit > 0 ? await dbQuery : [];
 
-    const totalItems =
-      cursor.totalItems ?? (rows.length ? Number(rows[0].count) : 0);
+    let totalItems: number;
+    if (cursor.totalItems !== undefined) {
+      totalItems = cursor.totalItems;
+    } else if (cursor.skipTotalItems) {
+      totalItems = 0;
+    } else if (rows.length) {
+      totalItems = Number(rows[0].count);
+    } else {
+      totalItems = 0;
+    }
 
     if (isFetchingBackwards) {
       rows.reverse();
@@ -832,15 +819,26 @@ export const cursorParser: z.ZodSchema<Cursor> = z.object({
 
 function parseCursorFromRequest(
   request?: QueryEntitiesRequest,
-): Partial<Cursor> {
+): Partial<Cursor> & { skipTotalItems: boolean } {
   if (isQueryEntitiesInitialRequest(request)) {
-    const { filter, orderFields: sortFields = [], fullTextFilter } = request;
-    return { filter, orderFields: sortFields, fullTextFilter };
+    const {
+      filter,
+      orderFields: sortFields = [],
+      fullTextFilter,
+      skipTotalItems = false,
+    } = request;
+    return { filter, orderFields: sortFields, fullTextFilter, skipTotalItems };
   }
   if (isQueryEntitiesCursorRequest(request)) {
-    return request.cursor;
+    return {
+      ...request.cursor,
+      // Doesn't matter here
+      skipTotalItems: false,
+    };
   }
-  return {};
+  return {
+    skipTotalItems: false,
+  };
 }
 
 function invertOrder(order: EntityOrder['order']) {

plugins/catalog-backend/src/service/createRouter.test.ts

@@ -135,9 +135,10 @@ describe('createRouter readonly disabled', () => {
         { apiVersion: 'a', kind: 'b', metadata: { name: 'n' } },
       ];
 
-      entitiesCatalog.entities.mockResolvedValueOnce({
-        entities: [entities[0]],
-        pageInfo: { hasNextPage: false },
+      entitiesCatalog.queryEntities.mockResolvedValueOnce({
+        items: [entities[0]],
+        pageInfo: {},
+        totalItems: 1,
       });
 
       const response = await request(app).get('/entities');
@@ -147,17 +148,18 @@ describe('createRouter readonly disabled', () => {
     });
 
     it('parses single and multiple request parameters and passes them down', async () => {
-      entitiesCatalog.entities.mockResolvedValueOnce({
-        entities: [],
-        pageInfo: { hasNextPage: false },
+      entitiesCatalog.queryEntities.mockResolvedValueOnce({
+        items: [],
+        pageInfo: {},
+        totalItems: 0,
       });
       const response = await request(app).get(
         '/entities?filter=a=1,a=2,b=3&filter=c=4',
       );
 
       expect(response.status).toEqual(200);
-      expect(entitiesCatalog.entities).toHaveBeenCalledTimes(1);
-      expect(entitiesCatalog.entities).toHaveBeenCalledWith({
+      expect(entitiesCatalog.queryEntities).toHaveBeenCalledTimes(1);
+      expect(entitiesCatalog.queryEntities).toHaveBeenCalledWith({
         filter: {
           anyOf: [
             {
@@ -169,7 +171,9 @@ describe('createRouter readonly disabled', () => {
             { key: 'c', values: ['4'] },
           ],
         },
+        limit: 10000,
         credentials: mockCredentials.user(),
+        skipTotalItems: true,
       });
     });
   });
@@ -913,9 +917,10 @@ describe('createRouter readonly enabled', () => {
         { apiVersion: 'a', kind: 'b', metadata: { name: 'n' } },
       ];
 
-      entitiesCatalog.entities.mockResolvedValueOnce({
-        entities: [entities[0]],
-        pageInfo: { hasNextPage: false },
+      entitiesCatalog.queryEntities.mockResolvedValueOnce({
+        items: [entities[0]],
+        pageInfo: {},
+        totalItems: 1,
       });
 
       const response = await request(app).get('/entities');

plugins/catalog-backend/src/service/createRouter.ts

@@ -27,7 +27,7 @@ import { InputError, NotFoundError, serializeError } from '@backstage/errors';
 import express from 'express';
 import yn from 'yn';
 import { z } from 'zod';
-import { EntitiesCatalog } from '../catalog/types';
+import { Cursor, EntitiesCatalog } from '../catalog/types';
 import { CatalogProcessingOrchestrator } from '../processing/types';
 import { validateEntityEnvelope } from '../processing/util';
 import {
@@ -41,8 +41,10 @@ import { parseEntityFacetParams } from './request/parseEntityFacetParams';
 import { parseEntityOrderParams } from './request/parseEntityOrderParams';
 import { LocationService, RefreshService } from './types';
 import {
+  createEntityArrayJsonStream,
   disallowReadonlyMode,
   encodeCursor,
+  expandLegacyCompoundRelationRefsInResponse,
   locationInput,
   validateRequestBody,
 } from './util';
@@ -57,6 +59,7 @@ import {
 } from '@backstage/backend-plugin-api';
 import { LocationAnalyzer } from '@backstage/plugin-catalog-node';
 import { AuthorizedValidationService } from './AuthorizedValidationService';
+import { DeferredPromise, createDeferred } from '@backstage/types';
 
 /**
  * Options used by {@link createRouter}.
@@ -135,24 +138,99 @@ export async function createRouter(
   if (entitiesCatalog) {
     router
       .get('/entities', async (req, res) => {
-        const { entities, pageInfo } = await entitiesCatalog.entities({
-          filter: parseEntityFilterParams(req.query),
-          fields: parseEntityTransformParams(req.query),
-          order: parseEntityOrderParams(req.query),
-          pagination: parseEntityPaginationParams(req.query),
-          credentials: await httpAuth.credentials(req),
-        });
+        const filter = parseEntityFilterParams(req.query);
+        const fields = parseEntityTransformParams(req.query);
+        const order = parseEntityOrderParams(req.query);
+        const pagination = parseEntityPaginationParams(req.query);
+        const credentials = await httpAuth.credentials(req);
 
-        // Add a Link header to the next page
-        if (pageInfo.hasNextPage) {
-          const url = new URL(`http://ignored${req.url}`);
-          url.searchParams.delete('offset');
-          url.searchParams.set('after', pageInfo.endCursor);
-          res.setHeader('link', `<${url.pathname}${url.search}>; rel="next"`);
+        // When pagination parameters are passed in, use the legacy slow path
+        // that loads all entities into memory
+
+        if (pagination) {
+          const { entities, pageInfo } = await entitiesCatalog.entities({
+            filter,
+            fields,
+            order,
+            pagination,
+            credentials,
+          });
+
+          // Add a Link header to the next page
+          if (pageInfo.hasNextPage) {
+            const url = new URL(`http://ignored${req.url}`);
+            url.searchParams.delete('offset');
+            url.searchParams.set('after', pageInfo.endCursor);
+            res.setHeader('link', `<${url.pathname}${url.search}>; rel="next"`);
+          }
+
+          res.json(entities);
+          return;
         }
 
-        // TODO(freben): encode the pageInfo in the response
-        res.json(entities);
+        // For other read-the-entire-world cases, use queryEntities and stream
+        // out results.
+
+        // The write lock is used for back pressure, preventing slow readers
+        // from forcing our read loop to pile up response data in userspace
+        // buffers faster than the kernel buffer is emptied.
+        // https://nodejs.org/api/http.html#http_response_write_chunk_encoding_callback
+        const locks: { writeLock?: DeferredPromise } = {};
+        const controller = new AbortController();
+        const signal = controller.signal;
+        req.on('end', () => {
+          controller.abort(new Error('Client closed connection'));
+          locks.writeLock?.resolve();
+          delete locks.writeLock;
+        });
+
+        const responseStream = createEntityArrayJsonStream(res);
+        const limit = 10000;
+        let cursor: Cursor | undefined;
+
+        try {
+          do {
+            const result = await entitiesCatalog.queryEntities(
+              !cursor
+                ? {
+                    credentials,
+                    fields,
+                    limit,
+                    filter,
+                    orderFields: order,
+                    skipTotalItems: true,
+                  }
+                : { credentials, fields, limit, cursor },
+            );
+
+            if (result.items.length) {
+              await locks?.writeLock;
+
+              signal.throwIfAborted();
+
+              expandLegacyCompoundRelationRefsInResponse(result.items);
+              if (!responseStream.send(result.items)) {
+                // The kernel buffer is full. Create the lock but do not await it
+                // yet - we can better spend our time going to the next round of
+                // the loop and read from the database while we wait for it to
+                // drain.
+                locks.writeLock = createDeferred();
+                res.once('drain', () => {
+                  locks.writeLock?.resolve();
+                  delete locks.writeLock;
+                });
+              }
+            }
+
+            signal.throwIfAborted();
+
+            cursor = result.pageInfo?.nextCursor;
+          } while (cursor);
+
+          responseStream.complete();
+        } finally {
+          responseStream.close();
+        }
       })
       .get('/entities/by-query', async (req, res) => {
         const { items, pageInfo, totalItems } =

plugins/catalog-backend/src/service/util.ts

@@ -15,7 +15,7 @@
  */
 
 import { InputError, NotAllowedError } from '@backstage/errors';
-import { Request } from 'express';
+import { Request, Response } from 'express';
 import lodash from 'lodash';
 import { z } from 'zod';
 import {
@@ -25,6 +25,11 @@ import {
   QueryEntitiesRequest,
 } from '../catalog/types';
 import { EntityFilter } from '@backstage/plugin-catalog-node';
+import {
+  Entity,
+  parseEntityRef,
+  stringifyEntityRef,
+} from '@backstage/catalog-model';
 
 export async function requireRequestBody(req: Request): Promise<unknown> {
   const contentType = req.header('content-type');
@@ -138,3 +143,80 @@ export function decodeCursor(encodedCursor: string) {
     throw new InputError(`Malformed cursor: ${e}`);
   }
 }
+
+// TODO(freben): This is added as a compatibility guarantee, until we can be
+// sure that all adopters have re-stitched their entities so that the new
+// targetRef field is present on them, and that they have stopped consuming
+// the now-removed old field
+// TODO(jhaals): Remove this in April 2022
+export function expandLegacyCompoundRelationRefsInResponse(
+  entities: Entity[],
+): void {
+  for (const entity of entities) {
+    if (entity.relations) {
+      for (const relation of entity.relations as any) {
+        if (!relation.targetRef && relation.target) {
+          // This is the case where an old-form entity, not yet stitched with
+          // the updated code, was in the database
+          relation.targetRef = stringifyEntityRef(relation.target);
+        } else if (!relation.target && relation.targetRef) {
+          // This is the case where a new-form entity, stitched with the
+          // updated code, was in the database but we still want to produce
+          // the old data shape as well for compatibility reasons
+          relation.target = parseEntityRef(relation.targetRef);
+        }
+      }
+    }
+  }
+}
+
+export interface EntityArrayJsonStream {
+  send(entities: Entity[]): boolean;
+  complete(): void;
+  close(): void;
+}
+
+// Helps stream Entity[] as a JSON response stream to avoid performance issues
+export function createEntityArrayJsonStream(
+  res: Response,
+): EntityArrayJsonStream {
+  // Imitate the httpRouter behavior of pretty-printing in development
+  const prettyPrint = process.env.NODE_ENV === 'development';
+  let firstSend = true;
+  let completed = false;
+
+  return {
+    send(entities) {
+      if (firstSend) {
+        res.setHeader('Content-Type', 'application/json; charset=utf-8');
+        res.status(200);
+        res.flushHeaders();
+      }
+
+      let data: string;
+      if (prettyPrint) {
+        data = JSON.stringify(entities, null, 2);
+        data = firstSend ? data.slice(0, -2) : `,\n${data.slice(2, -2)}`;
+      } else {
+        data = JSON.stringify(entities);
+        data = firstSend ? data.slice(0, -1) : `,${data.slice(1, -1)}`;
+      }
+
+      firstSend = false;
+      return res.write(data);
+    },
+    complete() {
+      if (firstSend) {
+        res.json([]);
+      } else {
+        res.end(prettyPrint ? '\n]' : ']');
+      }
+      completed = true;
+    },
+    close() {
+      if (!completed) {
+        res.end();
+      }
+    },
+  };
+}