Add cache option to config to expose ApolloGateway's queryPlanStore (#…

…2385)

Introduce new `queryPlannerConfig.cache` option to gateway constructor
to allow query plan cache configuration. Export new `QueryPlanCache` type
from `@apollo/query-planner` for implementors.

.changeset/forty-gifts-burn.md

@@ -0,0 +1,38 @@
+---
+"@apollo/query-planner": minor
+"@apollo/gateway": minor
+---
+
+This change introduces a configurable query plan cache. This option allows
+developers to provide their own query plan cache like so:
+
+```
+new ApolloGateway({
+  queryPlannerConfig: {
+    cache: new MyCustomQueryPlanCache(),
+  },
+});
+```
+
+The current default implementation is effectively as follows:
+```
+import { InMemoryLRUCache } from "@apollo/utils.keyvaluecache";
+
+const cache = new InMemoryLRUCache<string>({
+  maxSize: Math.pow(2, 20) * 30,
+  sizeCalculation<T>(obj: T): number {
+    return Buffer.byteLength(JSON.stringify(obj), "utf8");
+  },
+});
+```
+
+TypeScript users should implement the `QueryPlanCache` type which is now
+exported by `@apollo/query-planner`:
+```
+import { QueryPlanCache } from '@apollo/query-planner';
+
+class MyCustomQueryPlanCache implements QueryPlanCache {
+  // ...
+}
+```
+

gateway-js/package.json

@@ -39,7 +39,6 @@
     "@types/node-fetch": "^2.6.2",
     "async-retry": "^1.3.3",
     "loglevel": "^1.6.1",
-    "lru-cache": "^7.13.1",
     "make-fetch-happen": "^11.0.0",
     "node-abort-controller": "^3.0.1",
     "node-fetch": "^2.6.7"

gateway-js/src/__tests__/gateway/endToEnd.test.ts

@@ -8,6 +8,14 @@ import { GraphQLSchemaModule } from '@apollo/subgraph/src/schema-helper';
 import { buildSchema, ObjectType, ServiceDefinition } from '@apollo/federation-internals';
 import gql from 'graphql-tag';
 import { printSchema } from 'graphql';
+import { InMemoryLRUCache } from '@apollo/utils.keyvaluecache';
+import { QueryPlan } from '@apollo/query-planner';
+import { createHash } from '@apollo/utils.createhash';
+import { QueryPlanCache } from '@apollo/query-planner';
+
+function approximateObjectSize<T>(obj: T): number {
+  return Buffer.byteLength(JSON.stringify(obj), 'utf8');
+}
 
 async function startFederatedServer(modules: GraphQLSchemaModule[]) {
   const schema = buildSubgraphSchema(modules);
@@ -25,7 +33,7 @@ let gateway: ApolloGateway;
 let gatewayServer: ApolloServer;
 let gatewayUrl: string;
 
-async function startServicesAndGateway(servicesDefs: ServiceDefinition[]) {
+async function startServicesAndGateway(servicesDefs: ServiceDefinition[], cache?: QueryPlanCache) {
   backendServers = [];
   const serviceList = [];
   for (const serviceDef of servicesDefs) {
@@ -34,7 +42,11 @@ async function startServicesAndGateway(servicesDefs: ServiceDefinition[]) {
     serviceList.push({ name: serviceDef.name, url });
   }
 
-  gateway = new ApolloGateway({ serviceList });
+  gateway = new ApolloGateway({
+    serviceList,
+    queryPlannerConfig: cache ? { cache } : undefined,
+  });
+
   gatewayServer = new ApolloServer({
     gateway,
   });
@@ -60,9 +72,31 @@ afterEach(async () => {
   }
 });
 
+
 describe('caching', () => {
+  const cache = new InMemoryLRUCache<QueryPlan>({maxSize: Math.pow(2, 20) * (30), sizeCalculation: approximateObjectSize});
   beforeEach(async () => {
-    await startServicesAndGateway(fixtures);
+    await startServicesAndGateway(fixtures, cache);
+  });
+
+  it(`cached query plan`, async () => {
+    const query = `
+      query {
+        me {
+          name {
+            first
+            last
+          }
+        }
+        topProducts {
+          name
+        }
+      }
+    `;
+
+    await queryGateway(query);
+    const queryHash:string = createHash('sha256').update(query).digest('hex');
+    expect(await cache.get(queryHash)).toBeTruthy();
   });
 
   it(`cache control`, async () => {

gateway-js/src/index.ts

@@ -1,7 +1,8 @@
 import { deprecate } from 'util';
 import { createHash } from '@apollo/utils.createhash';
 import type { Logger } from '@apollo/utils.logger';
-import LRUCache from 'lru-cache';
+import { QueryPlanCache } from '@apollo/query-planner'
+import { InMemoryLRUCache } from '@apollo/utils.keyvaluecache';
 import {
   GraphQLSchema,
   VariableDefinitionNode,
@@ -123,7 +124,7 @@ export class ApolloGateway implements GatewayInterface {
   private serviceMap: DataSourceMap = Object.create(null);
   private config: GatewayConfig;
   private logger: Logger;
-  private queryPlanStore: LRUCache<string, QueryPlan>;
+  private queryPlanStore: QueryPlanCache;
   private apolloConfig?: ApolloConfigFromAS3;
   private onSchemaChangeListeners = new Set<(schema: GraphQLSchema) => void>();
   private onSchemaLoadOrUpdateListeners = new Set<
@@ -189,14 +190,17 @@ export class ApolloGateway implements GatewayInterface {
   }
 
   private initQueryPlanStore(approximateQueryPlanStoreMiB?: number) {
+    if(this.config.queryPlannerConfig?.cache){
+      return this.config.queryPlannerConfig?.cache
+    }
     // Create ~about~ a 30MiB InMemoryLRUCache (or 50MiB if the full operation ASTs are
     // enabled in query plans as this requires plans to use more memory). This is
     // less than precise since the technique to calculate the size of a DocumentNode is
     // only using JSON.stringify on the DocumentNode (and thus doesn't account
     // for unicode characters, etc.), but it should do a reasonable job at
     // providing a caching document store for most operations.
     const defaultSize = this.config.queryPlannerConfig?.exposeDocumentNodeInFetchNode ? 50 : 30;
-    return new LRUCache<string, QueryPlan>({
+    return new InMemoryLRUCache<QueryPlan>({
       maxSize: Math.pow(2, 20) * (approximateQueryPlanStoreMiB || defaultSize),
       sizeCalculation: approximateObjectSize,
     });
@@ -768,7 +772,7 @@ export class ApolloGateway implements GatewayInterface {
             span.setStatus({ code: SpanStatusCode.ERROR });
             return { errors: validationErrors };
           }
-          let queryPlan = this.queryPlanStore.get(queryPlanStoreKey);
+          let queryPlan = await this.queryPlanStore.get(queryPlanStoreKey);
 
           if (!queryPlan) {
             queryPlan = tracer.startActiveSpan(
@@ -792,7 +796,7 @@ export class ApolloGateway implements GatewayInterface {
             );
 
             try {
-              this.queryPlanStore.set(queryPlanStoreKey, queryPlan);
+              await this.queryPlanStore.set(queryPlanStoreKey, queryPlan);
             } catch (err) {
               this.logger.warn(
                 'Could not store queryPlan' + ((err && err.message) || err),

query-planner-js/src/config.ts

@@ -1,4 +1,8 @@
 import { Concrete } from "@apollo/federation-internals";
+import { QueryPlan } from ".";
+import { InMemoryLRUCache, KeyValueCache } from '@apollo/utils.keyvaluecache';
+
+export type QueryPlanCache = KeyValueCache<QueryPlan> & { clear: () => void }
 
 export type QueryPlannerConfig = {
   /**
@@ -45,6 +49,7 @@ export type QueryPlannerConfig = {
      */
     enableDefer?: boolean,
   }
+  cache?: QueryPlanCache,
 }
 
 export function enforceQueryPlannerConfigDefaults(
@@ -56,6 +61,7 @@ export function enforceQueryPlannerConfigDefaults(
     incrementalDelivery: {
       enableDefer: false,
     },
+    cache: new InMemoryLRUCache<QueryPlan>({maxSize: Math.pow(2, 20) * 50 }),
     ...config,
   };
 }

query-planner-js/src/index.ts

@@ -3,5 +3,5 @@ export { prettyFormatQueryPlan } from './prettyFormatQueryPlan';
 
 export * from './QueryPlan';
 export { QueryPlanner } from './buildPlan';
-export { QueryPlannerConfig } from './config';
+export { QueryPlanCache, QueryPlannerConfig } from './config';