feat: improve cache warmer (wundergraph#1530)

Co-authored-by: Ludwig <ludwig.bedacht@gmail.com>

controlplane/clickhouse/migrations/20250101213018_operation_planning_metrics_5_30_mv.sql

@@ -24,7 +24,7 @@ SELECT
     maxSimpleState(Max) AS MaxDuration
 FROM otel_metrics_histogram
 -- Only works with the same bounds for all buckets. If bounds are different, we can't add them together
-WHERE ScopeName = 'cosmo.router' AND ScopeVersion = '0.0.1' AND MetricName = 'router.graphql.operation.planning_time' AND OrganizationID != '' AND FederatedGraphID != ''
+WHERE ScopeName = 'cosmo.router' AND ScopeVersion = '0.0.1' AND MetricName = 'router.graphql.operation.planning_time' AND Attributes['wg.engine.plan_cache_hit'] == 'false' AND OrganizationID != '' AND FederatedGraphID != ''
 GROUP BY
     OperationName,
     OperationHash,

controlplane/clickhouse/migrations/20250122211233_planning_time_exclude_cache_hits.sql

@@ -0,0 +1,79 @@
+-- migrate:up
+
+ALTER TABLE cosmo.operation_planning_metrics_5_30_mv MODIFY QUERY
+    SELECT
+        toStartOfFiveMinute(TimeUnix) as Timestamp,
+        toLowCardinality(Attributes [ 'wg.operation.name' ]) as OperationName,
+        Attributes [ 'wg.operation.hash' ] as OperationHash,
+        toLowCardinality(Attributes [ 'wg.operation.type' ]) as OperationType,
+        Attributes [ 'wg.operation.persisted_id' ] as OperationPersistedID,
+        toLowCardinality(Attributes [ 'wg.router.config.version']) as RouterConfigVersion,
+        toLowCardinality(Attributes [ 'wg.federated_graph.id']) as FederatedGraphID,
+        toLowCardinality(Attributes [ 'wg.organization.id' ]) as OrganizationID,
+        toLowCardinality(Attributes [ 'wg.client.name' ]) as ClientName,
+        toLowCardinality(Attributes [ 'wg.client.version' ]) as ClientVersion,
+        -- Sum up the bucket counts on the same index which produces the overall count of samples of the histogram
+        sumForEachState(BucketCounts) as BucketCounts,
+        -- Populate the bounds so we have a base value for quantile calculations
+        ExplicitBounds,
+        sumSimpleState(Sum) AS Sum,
+        sumSimpleState(Count) AS Count,
+        minSimpleState(Min) AS MinDuration,
+        maxSimpleState(Max) AS MaxDuration
+    FROM otel_metrics_histogram
+    -- Only works with the same bounds for all buckets. If bounds are different, we can't add them together
+    WHERE ScopeName = 'cosmo.router' AND ScopeVersion = '0.0.1' AND MetricName = 'router.graphql.operation.planning_time' AND Attributes['wg.engine.plan_cache_hit'] == 'false' AND OrganizationID != '' AND FederatedGraphID != ''
+    GROUP BY
+        OperationName,
+        OperationHash,
+        OperationPersistedID,
+        FederatedGraphID,
+        RouterConfigVersion,
+        OrganizationID,
+        OperationType,
+        Timestamp,
+        ClientName,
+        ClientVersion,
+        ExplicitBounds
+    ORDER BY
+        Timestamp;
+
+-- migrate:down
+
+ALTER TABLE cosmo.operation_planning_metrics_5_30_mv MODIFY QUERY
+    SELECT
+        toStartOfFiveMinute(TimeUnix) as Timestamp,
+        toLowCardinality(Attributes [ 'wg.operation.name' ]) as OperationName,
+        Attributes [ 'wg.operation.hash' ] as OperationHash,
+        toLowCardinality(Attributes [ 'wg.operation.type' ]) as OperationType,
+        Attributes [ 'wg.operation.persisted_id' ] as OperationPersistedID,
+        toLowCardinality(Attributes [ 'wg.router.config.version']) as RouterConfigVersion,
+        toLowCardinality(Attributes [ 'wg.federated_graph.id']) as FederatedGraphID,
+        toLowCardinality(Attributes [ 'wg.organization.id' ]) as OrganizationID,
+        toLowCardinality(Attributes [ 'wg.client.name' ]) as ClientName,
+        toLowCardinality(Attributes [ 'wg.client.version' ]) as ClientVersion,
+        -- Sum up the bucket counts on the same index which produces the overall count of samples of the histogram
+        sumForEachState(BucketCounts) as BucketCounts,
+        -- Populate the bounds so we have a base value for quantile calculations
+        ExplicitBounds,
+        sumSimpleState(Sum) AS Sum,
+        sumSimpleState(Count) AS Count,
+        minSimpleState(Min) AS MinDuration,
+        maxSimpleState(Max) AS MaxDuration
+    FROM otel_metrics_histogram
+    -- Only works with the same bounds for all buckets. If bounds are different, we can't add them together
+    WHERE ScopeName = 'cosmo.router' AND ScopeVersion = '0.0.1' AND MetricName = 'router.graphql.operation.planning_time' AND OrganizationID != '' AND FederatedGraphID != ''
+    GROUP BY
+        OperationName,
+        OperationHash,
+        OperationPersistedID,
+        FederatedGraphID,
+        RouterConfigVersion,
+        OrganizationID,
+        OperationType,
+        Timestamp,
+        ClientName,
+        ClientVersion,
+        ExplicitBounds
+    ORDER BY
+        Timestamp;

controlplane/src/core/bufservices/cache-warmer/pushCacheWarmerOperation.ts

@@ -166,15 +166,15 @@ export function pushCacheWarmerOperation(
     }
 
     await cacheWarmerRepo.addCacheWarmerOperations({
-      federatedGraphId: federatedGraph.id,
-      organizationId: authContext.organizationId,
-      createdById: authContext.userId,
-      isManuallyAdded: true,
       operations: [
         {
-          name: req.operationName,
-          persistedId: req.operationPersistedId,
-          content: req.operationContent,
+          operationName: req.operationName,
+          operationPersistedID: req.operationPersistedId,
+          operationContent: req.operationContent,
+          federatedGraphId: federatedGraph.id,
+          organizationId: authContext.organizationId,
+          createdById: authContext.userId,
+          isManuallyAdded: true,
           clientName,
         },
       ],

controlplane/src/core/repositories/CacheWarmerRepository.ts

@@ -15,6 +15,7 @@ import { DateRange } from '../../types/index.js';
 import { BlobStorage } from '../blobstorage/index.js';
 import { ClickHouseClient } from '../clickhouse/index.js';
 import { S3RouterConfigMetadata } from '../composition/composer.js';
+import { CacheWarmupOperation } from '../../db/models.js';
 import { getDateRange, isoDateRangeToTimestamps } from './analytics/util.js';
 
 interface ComputeCacheWarmerOperationsProps {
@@ -24,16 +25,6 @@ interface ComputeCacheWarmerOperationsProps {
   federatedGraphId: string;
 }
 
-interface DBCacheWarmerOperation {
-  content?: string;
-  hash?: string;
-  name?: string;
-  persistedId?: string;
-  clientName?: string;
-  clientVersion?: string;
-  planningTime?: number;
-}
-
 export class CacheWarmerRepository {
   constructor(
     private client: ClickHouseClient,
@@ -49,13 +40,14 @@ export class CacheWarmerRepository {
     const parsedDateRange = isoDateRangeToTimestamps(dateRange, rangeInHours);
     const [start, end] = getDateRange(parsedDateRange);
     const quantile = 0.9;
-    const minPlanningTimeInMs = 0;
+    const minPlanningTimeInMs = 1;
 
     const query = `
       WITH
         toDateTime('${start}') AS startDate,
         toDateTime('${end}') AS endDate
       SELECT
+        max(MaxDuration) as maxDuration,
         OperationHash as operationHash,
         OperationName as operationName,
         OperationPersistedID as operationPersistedID,
@@ -76,7 +68,7 @@ export class CacheWarmerRepository {
       AND OrganizationID = '${organizationId}'
       AND OperationName != 'IntrospectionQuery'
       GROUP BY OperationHash, OperationName, OperationPersistedID, ClientName, ClientVersion
-      HAVING planningTime > ${minPlanningTimeInMs}
+      HAVING maxDuration >= ${minPlanningTimeInMs}
       ORDER BY planningTime DESC LIMIT 100
     `;
 
@@ -144,7 +136,7 @@ export class CacheWarmerRepository {
     const topOperationsByPlanningTime = await this.getTopOperationsByPlanningTime(props);
 
     const computedOperations: Operation[] = [];
-    const dbCacheWarmerOperations: DBCacheWarmerOperation[] = [];
+    const dbCacheWarmerOperations: CacheWarmupOperation[] = [];
 
     const manuallyAddedOperations = await this.getCacheWarmerOperations({
       organizationId: props.organizationId,
@@ -235,12 +227,15 @@ export class CacheWarmerRepository {
         );
 
         dbCacheWarmerOperations.push({
-          name: operation.operationName,
-          hash: operation.operationHash,
-          persistedId: operation.operationPersistedID,
+          operationName: operation.operationName,
+          operationHash: operation.operationHash,
+          operationPersistedID: operation.operationPersistedID,
           clientName: operation.clientName,
           clientVersion: operation.clientVersion,
           planningTime: operation.planningTime,
+          federatedGraphId: props.federatedGraphId,
+          organizationId: props.organizationId,
+          isManuallyAdded: false,
         });
         continue;
       }
@@ -252,13 +247,16 @@ export class CacheWarmerRepository {
       }
 
       dbCacheWarmerOperations.push({
-        content: operationContent,
-        name: operation.operationName,
-        hash: operation.operationHash,
-        persistedId: operation.operationPersistedID,
+        operationName: operation.operationName,
+        operationHash: operation.operationHash,
+        operationPersistedID: operation.operationPersistedID,
         clientName: operation.clientName,
         clientVersion: operation.clientVersion,
         planningTime: operation.planningTime,
+        federatedGraphId: props.federatedGraphId,
+        organizationId: props.organizationId,
+        operationContent,
+        isManuallyAdded: false,
       });
 
       computedOperations.push(
@@ -285,9 +283,6 @@ export class CacheWarmerRepository {
       });
 
       await cacheWarmerRepo.addCacheWarmerOperations({
-        organizationId: props.organizationId,
-        federatedGraphId: props.federatedGraphId,
-        isManuallyAdded: false,
         operations: dbCacheWarmerOperations,
       });
     });
@@ -328,45 +323,12 @@ export class CacheWarmerRepository {
       .then((res) => res.length > 0);
   }
 
-  public async addCacheWarmerOperations({
-    organizationId,
-    federatedGraphId,
-    isManuallyAdded,
-    operations,
-    createdById,
-  }: {
-    organizationId: string;
-    federatedGraphId: string;
-    isManuallyAdded: boolean;
-    operations: {
-      content?: string;
-      hash?: string;
-      name?: string;
-      persistedId?: string;
-      clientName?: string;
-      clientVersion?: string;
-      planningTime?: number;
-    }[];
-    createdById?: string;
-  }) {
+  public async addCacheWarmerOperations({ operations }: { operations: CacheWarmupOperation[] }) {
     if (!operations || operations.length === 0) {
       return;
     }
-    const data = operations.map((operation) => ({
-      federatedGraphId,
-      organizationId,
-      isManuallyAdded,
-      operationContent: operation.content || null,
-      operationHash: operation.hash || null,
-      operationName: operation.name || null,
-      operationPersistedID: operation.persistedId || null,
-      clientName: operation.clientName || null,
-      clientVersion: operation.clientVersion || null,
-      planningTime: operation.planningTime,
-      createdById,
-    }));
-
-    await this.db.insert(cacheWarmerOperations).values(data);
+
+    await this.db.insert(cacheWarmerOperations).values(operations);
   }
 
   public getCacheWarmerOperations({

controlplane/src/db/models.ts

@@ -12,6 +12,7 @@ import {
   websocketSubprotocolEnum,
   webhookDeliveries,
   graphPruningRulesEnum,
+  cacheWarmerOperations,
 } from './schema.js';
 
 export type FederatedGraph = typeof federatedGraphs.$inferSelect;
@@ -26,6 +27,7 @@ export type MemberRole = (typeof memberRoleEnum.enumValues)[number];
 export type LintRuleEnum = (typeof lintRulesEnum.enumValues)[number];
 export type GraphPruningRuleEnum = (typeof graphPruningRulesEnum.enumValues)[number];
 export type WebsocketSubprotocol = (typeof websocketSubprotocolEnum.enumValues)[number];
+export type CacheWarmupOperation = typeof cacheWarmerOperations.$inferInsert;
 
 export type WebhookDeliveryInfo = typeof webhookDeliveries.$inferInsert;