feat(controlplane): run federated composition in worker pool to unblock main event loop

pepol · pepol · commit cb7ec4893acc · 2025-11-14T10:38:59.000+01:00
diff --git a/controlplane/package.json b/controlplane/package.json
@@ -90,6 +90,7 @@
     "stream-json": "^1.8.0",
     "stripe": "^14.19.0",
     "tiny-lru": "^11.2.11",
+    "tinypool": "^2.0.0",
     "uid": "^2.0.2",
     "uuid": "^10.0.0",
     "zod": "^3.22.4",
diff --git a/controlplane/src/core/bufservices/contract/createContract.ts b/controlplane/src/core/bufservices/contract/createContract.ts
@@ -206,6 +206,7 @@ export function createContract(
         chClient: opts.chClient!,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs: [{ ...contractGraph, contract }],
+        composeWorkerPool: opts.composeWorkerPool,
       });
 
       compositionErrors.push(...composition.compositionErrors);
diff --git a/controlplane/src/core/bufservices/contract/updateContract.ts b/controlplane/src/core/bufservices/contract/updateContract.ts
@@ -140,6 +140,7 @@ export function updateContract(
       },
       labelMatchers: [],
       chClient: opts.chClient!,
+      composeWorkerPool: opts.composeWorkerPool,
       compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
     });
 
@@ -155,6 +156,7 @@ export function updateContract(
       },
       blobStorage: opts.blobStorage,
       chClient: opts.chClient!,
+      composeWorkerPool: opts.composeWorkerPool,
       compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
       federatedGraphs: [
         {
diff --git a/controlplane/src/core/bufservices/feature-flag/createFeatureFlag.ts b/controlplane/src/core/bufservices/feature-flag/createFeatureFlag.ts
@@ -205,6 +205,7 @@ export function createFeatureFlag(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs,
       });
diff --git a/controlplane/src/core/bufservices/feature-flag/deleteFeatureFlag.ts b/controlplane/src/core/bufservices/feature-flag/deleteFeatureFlag.ts
@@ -135,6 +135,7 @@ export function deleteFeatureFlag(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs,
       });
diff --git a/controlplane/src/core/bufservices/feature-flag/enableFeatureFlag.ts b/controlplane/src/core/bufservices/feature-flag/enableFeatureFlag.ts
@@ -119,6 +119,7 @@ export function enableFeatureFlag(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs,
       });
diff --git a/controlplane/src/core/bufservices/feature-flag/updateFeatureFlag.ts b/controlplane/src/core/bufservices/feature-flag/updateFeatureFlag.ts
@@ -174,6 +174,7 @@ export function updateFeatureFlag(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs: allFederatedGraphsToCompose,
       });
diff --git a/controlplane/src/core/bufservices/federated-graph/createFederatedGraph.ts b/controlplane/src/core/bufservices/federated-graph/createFederatedGraph.ts
@@ -232,6 +232,7 @@ export function createFederatedGraph(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs: [federatedGraph],
       });
diff --git a/controlplane/src/core/bufservices/federated-graph/migrateFromApollo.ts b/controlplane/src/core/bufservices/federated-graph/migrateFromApollo.ts
@@ -156,6 +156,7 @@ export function migrateFromApollo(
           webhookJWTSecret: opts.admissionWebhookJWTSecret,
         },
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: {
           disableResolvabilityValidation: true,
         },
diff --git a/controlplane/src/core/bufservices/federated-graph/moveFederatedGraph.ts b/controlplane/src/core/bufservices/federated-graph/moveFederatedGraph.ts
@@ -125,6 +125,7 @@ export function moveFederatedGraph(
           jwtSecret: opts.admissionWebhookJWTSecret,
         },
         opts.chClient!,
+        opts.composeWorkerPool,
       );
 
       const allDeploymentErrors: PlainMessage<DeploymentError>[] = [];
@@ -163,6 +164,7 @@ export function moveFederatedGraph(
             jwtSecret: opts.admissionWebhookJWTSecret,
           },
           opts.chClient!,
+          opts.composeWorkerPool,
         );
 
         allCompositionErrors.push(...contractErrors);
diff --git a/controlplane/src/core/bufservices/federated-graph/updateFederatedGraph.ts b/controlplane/src/core/bufservices/federated-graph/updateFederatedGraph.ts
@@ -119,6 +119,7 @@ export function updateFederatedGraph(
       admissionWebhookURL: req.admissionWebhookURL,
       blobStorage: opts.blobStorage,
       chClient: opts.chClient!,
+      composeWorkerPool: opts.composeWorkerPool,
       compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
       labelMatchers: req.labelMatchers,
       namespaceId: federatedGraph.namespaceId,
diff --git a/controlplane/src/core/bufservices/graph/setGraphRouterCompatibilityVersion.ts b/controlplane/src/core/bufservices/graph/setGraphRouterCompatibilityVersion.ts
@@ -140,6 +140,7 @@ export function setGraphRouterCompatibilityVersion(
         },
         blobStorage: opts.blobStorage,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
         compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
         federatedGraphs: [federatedGraph],
       });
diff --git a/controlplane/src/core/bufservices/monograph/publishMonograph.ts b/controlplane/src/core/bufservices/monograph/publishMonograph.ts
@@ -152,6 +152,7 @@ export function publishMonograph(
           webhookJWTSecret: opts.admissionWebhookJWTSecret,
         },
         opts.chClient!,
+        opts.composeWorkerPool,
       );
 
     for (const graph of updatedFederatedGraphs) {
diff --git a/controlplane/src/core/bufservices/monograph/updateMonograph.ts b/controlplane/src/core/bufservices/monograph/updateMonograph.ts
@@ -150,6 +150,7 @@ export function updateMonograph(
         admissionWebhookURL: req.admissionWebhookURL,
         admissionWebhookSecret: req.admissionWebhookSecret,
         chClient: opts.chClient!,
+        composeWorkerPool: opts.composeWorkerPool,
       });
 
       await subgraphRepo.update(
@@ -173,6 +174,7 @@ export function updateMonograph(
           webhookJWTSecret: opts.admissionWebhookJWTSecret,
         },
         opts.chClient!,
+        opts.composeWorkerPool,
       );
 
       await auditLogRepo.addAuditLog({
diff --git a/controlplane/src/core/bufservices/subgraph/deleteFederatedSubgraph.ts b/controlplane/src/core/bufservices/subgraph/deleteFederatedSubgraph.ts
@@ -186,6 +186,7 @@ export function deleteFederatedSubgraph(
           chClient: opts.chClient!,
           compositionOptions: newCompositionOptions(req.disableResolvabilityValidation),
           federatedGraphs: affectedFederatedGraphs,
+          composeWorkerPool: opts.composeWorkerPool,
         });
 
         return { affectedFederatedGraphs, compositionErrors, deploymentErrors, compositionWarnings };
diff --git a/controlplane/src/core/bufservices/subgraph/moveSubgraph.ts b/controlplane/src/core/bufservices/subgraph/moveSubgraph.ts
@@ -117,6 +117,7 @@ export function moveSubgraph(
               jwtSecret: opts.admissionWebhookJWTSecret,
             },
             opts.chClient!,
+            opts.composeWorkerPool,
             newCompositionOptions(req.disableResolvabilityValidation),
           );
 
diff --git a/controlplane/src/core/bufservices/subgraph/publishFederatedSubgraph.ts b/controlplane/src/core/bufservices/subgraph/publishFederatedSubgraph.ts
@@ -574,6 +574,7 @@ export function publishFederatedSubgraph(
           webhookJWTSecret: opts.admissionWebhookJWTSecret,
         },
         opts.chClient!,
+        opts.composeWorkerPool,
         newCompositionOptions(req.disableResolvabilityValidation),
       );
 
diff --git a/controlplane/src/core/bufservices/subgraph/updateSubgraph.ts b/controlplane/src/core/bufservices/subgraph/updateSubgraph.ts
@@ -184,6 +184,7 @@ export function updateSubgraph(
           webhookJWTSecret: opts.admissionWebhookJWTSecret,
         },
         opts.chClient!,
+        opts.composeWorkerPool,
         newCompositionOptions(req.disableResolvabilityValidation),
       );
 
diff --git a/controlplane/src/core/build-server.ts b/controlplane/src/core/build-server.ts
@@ -8,6 +8,7 @@ import { compressionBrotli, compressionGzip } from '@connectrpc/connect-node';
 import fastifyGracefulShutdown from 'fastify-graceful-shutdown';
 import { App } from 'octokit';
 import { Worker } from 'bullmq';
+import { Tinypool } from 'tinypool';
 import routes from './routes.js';
 import fastifyHealth from './plugins/health.js';
 import fastifyMetrics, { MetricsPluginOptions } from './plugins/metrics.js';
@@ -53,6 +54,7 @@ import {
   createReactivateOrganizationWorker,
   ReactivateOrganizationQueue,
 } from './workers/ReactivateOrganizationWorker.js';
+import { getWorkerPool } from '../workers/compose.js';
 
 export interface BuildConfig {
   logger: LoggerOptions;
@@ -136,6 +138,9 @@ export interface BuildConfig {
       key?: string; // e.g. string or '/path/to/my/client-key.pem'
     };
   };
+  composeWorkers?: {
+    maxCount: number;
+  };
 }
 
 export interface MetricsOptions {
@@ -469,6 +474,8 @@ export default async function build(opts: BuildConfig) {
     keycloakRealm: opts.keycloak.realm,
   });
 
+  const composeWorkerPool: Tinypool = getWorkerPool(opts.composeWorkers?.maxCount)!;
+
   // Must be registered after custom fastify routes
   // Because it registers an all-catch route for connect handlers
 
@@ -502,6 +509,7 @@ export default async function build(opts: BuildConfig) {
       stripeSecretKey: opts.stripe?.secret,
       admissionWebhookJWTSecret: opts.admissionWebhook.secret,
       cdnBaseUrl: opts.cdnBaseUrl,
+      composeWorkerPool: composeWorkerPool,
     }),
     contextValues(req) {
       return createContextValues().set<FastifyBaseLogger>({ id: fastifyLoggerId, defaultValue: req.log }, req.log);
@@ -522,6 +530,12 @@ export default async function build(opts: BuildConfig) {
   });
 
   fastify.gracefulShutdown(async () => {
+    fastify.log.debug('Shutting down compose worker pool');
+
+    await composeWorkerPool.destroy();
+
+    fastify.log.debug('Compose worker pool shut down');
+
     fastify.log.debug('Shutting down bull workers');
 
     await Promise.all(bullWorkers.map((worker) => worker.close()));
diff --git a/controlplane/src/core/repositories/FederatedGraphRepository.ts b/controlplane/src/core/repositories/FederatedGraphRepository.ts
@@ -29,6 +29,7 @@ import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import { FastifyBaseLogger } from 'fastify';
 import { parse } from 'graphql';
 import { generateKeyPair, importPKCS8, SignJWT } from 'jose';
+import { Tinypool } from 'tinypool';
 import { uid } from 'uid/secure';
 import {
   ContractTagOptions,
@@ -205,6 +206,7 @@ export class FederatedGraphRepository {
     targetId: string;
     updatedBy: string;
     chClient: ClickHouseClient;
+    composeWorkerPool: Tinypool;
     admissionWebhookSecret?: string;
     admissionWebhookURL?: string;
     compositionOptions?: CompositionOptions;
@@ -337,6 +339,7 @@ export class FederatedGraphRepository {
           actorId: data.updatedBy,
           chClient: data.chClient,
           compositionOptions: data.compositionOptions,
+          composeWorkerPool: data.composeWorkerPool,
         });
 
         return {
@@ -369,6 +372,7 @@ export class FederatedGraphRepository {
       cdnBaseUrl: string;
     },
     chClient: ClickHouseClient,
+    composeWorkerPool: Tinypool,
     compositionOptions?: CompositionOptions,
   ): Promise<{
     compositionErrors: PlainMessage<CompositionError>[];
@@ -430,6 +434,7 @@ export class FederatedGraphRepository {
           chClient,
           compositionOptions,
           federatedGraphs: [movedContractGraph],
+          composeWorkerPool,
         });
 
         return {
@@ -449,6 +454,7 @@ export class FederatedGraphRepository {
         },
         chClient,
         compositionOptions,
+        composeWorkerPool,
       });
 
       return {
@@ -1505,6 +1511,7 @@ export class FederatedGraphRepository {
     chClient,
     blobStorage,
     federatedGraphs,
+    composeWorkerPool,
   }: {
     actorId: string;
     admissionConfig: {
@@ -1515,6 +1522,7 @@ export class FederatedGraphRepository {
     chClient: ClickHouseClient;
     federatedGraphs: FederatedGraphDTO[];
     compositionOptions?: CompositionOptions;
+    composeWorkerPool: Tinypool;
   }) => {
     return this.db.transaction(async (tx) => {
       const subgraphRepo = new SubgraphRepository(this.logger, tx, this.organizationId);
@@ -1524,7 +1532,7 @@ export class FederatedGraphRepository {
       const graphCompositionRepo = new GraphCompositionRepository(this.logger, tx);
       const composer = new Composer(
         this.logger,
-        this.db,
+        tx,
         fedGraphRepo,
         subgraphRepo,
         contractRepo,
@@ -1577,7 +1585,8 @@ export class FederatedGraphRepository {
         const contractBaseCompositionDataByContractId = new Map<string, ContractBaseCompositionData>();
 
         for (const subgraphsToCompose of allSubgraphsToCompose) {
-          const result: FederationResult | FederationResultWithContracts = getFederationResultWithPotentialContracts(
+          const result: FederationResult | FederationResultWithContracts = await getFederationResultWithPotentialContracts(
+            composeWorkerPool,
             federatedGraph,
             subgraphsToCompose,
             tagOptionsByContractName,
diff --git a/controlplane/src/core/repositories/SubgraphRepository.ts b/controlplane/src/core/repositories/SubgraphRepository.ts
@@ -13,6 +13,7 @@ import { and, asc, count, desc, eq, getTableName, gt, inArray, like, lt, notInAr
 import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import { FastifyBaseLogger } from 'fastify';
 import { GraphQLSchema } from 'graphql';
+import { Tinypool } from 'tinypool';
 import { DBSubgraphType, WebsocketSubprotocol } from '../../db/models.js';
 import * as schema from '../../db/schema.js';
 import {
@@ -242,6 +243,7 @@ export class SubgraphRepository {
       cdnBaseUrl: string;
     },
     chClient: ClickHouseClient,
+    composeWorkerPool: Tinypool,
     compositionOptions?: CompositionOptions,
   ): Promise<{
     compositionErrors: PlainMessage<CompositionError>[];
@@ -470,6 +472,7 @@ export class SubgraphRepository {
         chClient,
         compositionOptions,
         federatedGraphs: updatedFederatedGraphs.filter((g) => !g.contract),
+        composeWorkerPool,
       });
 
       compositionErrors.push(...cErrors);
@@ -501,6 +504,7 @@ export class SubgraphRepository {
       cdnBaseUrl: string;
     },
     chClient: ClickHouseClient,
+    composeWorkerPool: Tinypool,
     compositionOptions?: CompositionOptions,
   ): Promise<{
     compositionErrors: PlainMessage<CompositionError>[];
@@ -554,6 +558,7 @@ export class SubgraphRepository {
         actorId: data.updatedBy,
         chClient,
         compositionOptions,
+        composeWorkerPool,
       });
 
       return { compositionErrors, updatedFederatedGraphs, deploymentErrors, compositionWarnings };
diff --git a/controlplane/src/core/routes.ts b/controlplane/src/core/routes.ts
@@ -5,6 +5,7 @@ import { PlatformService } from '@wundergraph/cosmo-connect/dist/platform/v1/pla
 import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import pino from 'pino';
 import { App } from 'octokit';
+import { Tinypool } from 'tinypool';
 import * as schema from '../db/schema.js';
 import NodeServiceImpl from './bufservices/NodeService.js';
 import PlatformServiceImpl from './bufservices/PlatformService.js';
@@ -51,6 +52,7 @@ export interface RouterOptions {
   };
   stripeSecretKey?: string;
   cdnBaseUrl: string;
+  composeWorkerPool: Tinypool;
 }
 const handlerOptions: Partial<ConnectRouterOptions> = {
   maxTimeoutMs: 80_000,
diff --git a/controlplane/src/core/util.ts b/controlplane/src/core/util.ts
diff --git a/controlplane/src/workers/compose.ts b/controlplane/src/workers/compose.ts
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
