feat: improve performance of subgraph batch-publish

controlplane/src/core/bufservices/subgraph/publishFederatedSubgraphs.ts

@@ -436,7 +436,7 @@ async function runBatchPublish({
   );
 
   const { compositionErrors, compositionWarnings, deploymentErrors } =
-    await compositionService.recomposeAndDeployAffected({
+    await compositionService.recomposeAndDeployAffectedBatch({
       actorId: authContext.userId,
       affectedFederatedGraphs,
       affectedFeatureFlags,

controlplane/src/core/repositories/FeatureFlagRepository.ts

@@ -65,6 +65,29 @@ export type CheckConstituentFeatureSubgraphsResult = {
   featureSubgraphIds: Array<string>;
 };
 
+export type FeatureFlagCollectCaches = {
+  featureFlagsByBaseSubgraphId: Map<string, ReturnType<FeatureFlagRepository['getFeatureFlagsByBaseSubgraphId']>>;
+  matchedFeatureFlagsByLabelKey: Map<string, ReturnType<FeatureFlagRepository['getMatchedFeatureFlags']>>;
+  featureSubgraphsByFlagId: Map<string, ReturnType<FeatureFlagRepository['getFeatureSubgraphsByFeatureFlagId']>>;
+};
+
+/** Get-or-compute a cached promise (single-flight). With no cache, just computes. */
+function memoizePromise<T>(
+  cache: Map<string, Promise<T>> | undefined,
+  key: string,
+  compute: () => Promise<T>,
+): Promise<T> {
+  if (!cache) {
+    return compute();
+  }
+  let promise = cache.get(key);
+  if (!promise) {
+    promise = compute();
+    cache.set(key, promise);
+  }
+  return promise;
+}
+
 @traced
 export class FeatureFlagRepository {
   constructor(
@@ -1107,37 +1130,39 @@ export class FeatureFlagRepository {
     baseSubgraphNames,
     fedGraphLabelMatchers,
     excludeDisabled,
+    caches,
   }: {
     baseSubgraphId: string;
     namespaceId: string;
     baseSubgraphNames: string[];
     fedGraphLabelMatchers: string[];
     excludeDisabled: boolean;
+    caches?: FeatureFlagCollectCaches;
   }): Promise<FeatureFlagWithFeatureSubgraphs[]> {
     const featureFlagWithEnabledFeatureGraphs: FeatureFlagWithFeatureSubgraphs[] = [];
-    const featureFlagsBySubgraphId = await this.getFeatureFlagsByBaseSubgraphId({
-      baseSubgraphId,
-      namespaceId,
-      excludeDisabled,
-    });
 
-    // gets all the feature flags that match the label matchers
-    const matchedFeatureFlags = await this.getMatchedFeatureFlags({
-      namespaceId,
-      fedGraphLabelMatchers,
-      excludeDisabled,
-    });
+    const featureFlagsBySubgraphId = await memoizePromise(
+      caches?.featureFlagsByBaseSubgraphId,
+      `${namespaceId}:${excludeDisabled}:${baseSubgraphId}`,
+      () => this.getFeatureFlagsByBaseSubgraphId({ baseSubgraphId, namespaceId, excludeDisabled }),
+    );
+
+    const matchedFeatureFlags = await memoizePromise(
+      caches?.matchedFeatureFlagsByLabelKey,
+      `${namespaceId}:${excludeDisabled}:${[...fedGraphLabelMatchers].sort().join(';')}`,
+      () => this.getMatchedFeatureFlags({ namespaceId, fedGraphLabelMatchers, excludeDisabled }),
+    );
 
     for (const featureFlag of featureFlagsBySubgraphId) {
       const matched = matchedFeatureFlags.some((m) => m.id === featureFlag.id);
       if (!matched) {
         continue;
       }
 
-      const featureSubgraphsByFlag = await this.getFeatureSubgraphsByFeatureFlagId({
-        featureFlagId: featureFlag.id,
-        namespaceId,
-      });
+      // Feature subgraphs of the flag — memoized by flag id (the same flags recur across the batch).
+      const featureSubgraphsByFlag = await memoizePromise(caches?.featureSubgraphsByFlagId, featureFlag.id, () =>
+        this.getFeatureSubgraphsByFeatureFlagId({ featureFlagId: featureFlag.id, namespaceId }),
+      );
 
       // if there are no feature subgraphs in the flag, then skip the flag
       if (featureSubgraphsByFlag.length === 0) {

controlplane/src/core/repositories/SubgraphRepository.ts

@@ -73,7 +73,7 @@ import { OrganizationWebhookService } from '../webhooks/OrganizationWebhookServi
 import { traced } from '../tracing.js';
 import type { CompositionService } from '../services/CompositionService.js';
 import { ContractRepository } from './ContractRepository.js';
-import { FeatureFlagRepository } from './FeatureFlagRepository.js';
+import { FeatureFlagCollectCaches, FeatureFlagRepository } from './FeatureFlagRepository.js';
 import { FederatedGraphRepository } from './FederatedGraphRepository.js';
 import { GraphCompositionRepository } from './GraphCompositionRepository.js';
 import { OperationsRepository } from './OperationsRepository.js';
@@ -370,6 +370,12 @@ export class SubgraphRepository {
     tx: PostgresJsDatabase<typeof schema>,
     data: UpdateSubgraphSchemaData,
     splitConfigFeature?: Feature,
+    // When provided (batch path), `listByFederatedGraph` reads are memoized per federated graph across the whole batch.
+    // Without it, every changed feature subgraph re-loads ALL subgraphs (with their SDL) of the same federated graph,
+    // making the collect step scale with (changed subgraphs × subgraphs in the graph).
+    listByFederatedGraphCache?: Map<string, Promise<SubgraphDTO[]>>,
+    // When provided (batch path), the feature-flag sub-queries are memoized across the whole batch.
+    featureFlagCaches?: FeatureFlagCollectCaches,
   ): Promise<{
     subgraph: SubgraphDTO | undefined;
     affectedFederatedGraphById: Map<string, FederatedGraphDTO>;
@@ -568,18 +574,25 @@ export class SubgraphRepository {
         });
 
         for (const federatedGraphDTO of federatedGraphDTOs) {
-          // Retrieve all the subgraphs that compose the federated graph to retrieve the feature flags
-          const subgraphs = await subgraphRepo.listByFederatedGraph({
-            federatedGraphTargetId: federatedGraphDTO.targetId,
-            published: true,
-          });
+          // Retrieve all the subgraphs that compose the federated graph.
+          // Memoized across the batch (see `listByFederatedGraphCache`), so the same federated graph is loaded once, rather than once per changed subgraph.
+          let subgraphsPromise = listByFederatedGraphCache?.get(federatedGraphDTO.targetId);
+          if (!subgraphsPromise) {
+            subgraphsPromise = subgraphRepo.listByFederatedGraph({
+              federatedGraphTargetId: federatedGraphDTO.targetId,
+              published: true,
+            });
+            listByFederatedGraphCache?.set(federatedGraphDTO.targetId, subgraphsPromise);
+          }
+          const subgraphs = await subgraphsPromise;
 
           const enabledFeatureFlags = await featureFlagRepo.getFeatureFlagsByBaseSubgraphIdAndLabelMatchers({
             baseSubgraphId: baseSubgraph[0].id,
             namespaceId: data.namespaceId,
             fedGraphLabelMatchers: federatedGraphDTO.labelMatchers || [],
             baseSubgraphNames: subgraphs.map((subgraph) => subgraph.name),
             excludeDisabled: true,
+            caches: featureFlagCaches,
           });
 
           // If an enabled feature flag includes the feature graph that has just been published, push it to the array
@@ -674,8 +687,25 @@ export class SubgraphRepository {
 
     await this.db.transaction(async (tx) => {
       // Write every schema version and collect the affected graphs/flags. NO composition happens here.
+      // Memoize `listByFederatedGraph` across the batch so a federated graph's subgraphs are loaded once, not once per
+      // changed feature subgraph (the dominant cost when many feature subgraphs of the same graph change at once).
+      const listByFederatedGraphCache = new Map<string, Promise<SubgraphDTO[]>>();
+      // Memoize the feature-flag sub-queries across the batch (matched-flags per graph, feature-subgraphs per flag).
+      const featureFlagCaches: FeatureFlagCollectCaches = {
+        featureFlagsByBaseSubgraphId: new Map(),
+        matchedFeatureFlagsByLabelKey: new Map(),
+        featureSubgraphsByFlagId: new Map(),
+      };
       const results = await Promise.all(
-        items.map((item) => this.writeSchemaAndCollectAffected(tx, item, splitConfigFeature)),
+        items.map((item) =>
+          this.writeSchemaAndCollectAffected(
+            tx,
+            item,
+            splitConfigFeature,
+            listByFederatedGraphCache,
+            featureFlagCaches,
+          ),
+        ),
       );
 
       for (const [index, result] of results.entries()) {