Apollo Federated GraphQL Cache Headers

March 15, 2021

Preface

At Pocket we’ve started moving to a Federated GraphQL API. We decided to use Apollo to manage this system because the documentation makes it look easy. Now we aren’t experts yet at GraphQL, but we have enough knowledge to read the docs and get services up and running.

Cache Control Directives

One of the most exciting features we saw in the documentation was that Apollo would be smart enough to read cache directives on the schema and automatically send the necessary cache headers in the response.

From the docs: “Apollo Server enables you to define cache control settings (maxAge and scope) for each field in your schema”

type Post {
  id: ID!
  title: String
  author: Author
  votes: Int @cacheControl(maxAge: 30)  comments: [Comment]
  readByCurrentUser: Boolean! @cacheControl(maxAge: 10, scope: PRIVATE)}
}

Perfect we thought! We can use this and let the server decide our cache strategy based on the models that were sent in the response. This solves a big problem for us at Pocket. Our current strategy relies on us manually setting the cache, if we need it, per request in our current restful design. So you can imagine when we saw the ability to let the server figure out the cache strategy based on the domain model we were wildly enthused.

So we started implementing these directives in all of our domain models in each of our implementing services that will build our new Pocket Graph. As we went through each service it was working great! Every service we added the directives to started responding with appropriate cache control headers. Everything was working as expected, that is until we wired it up to our Federated Graph.

Simple Caching You Say?

We learned quickly, things are never as good as they seem. Once we were wired into the Federated Graph we noticed that the final merged result was only using the default cache control header that was set in our Apollo config. So why wasn’t it working? All our implementing services were fully working when we hit them individually, but not when it was merged together.

Enter the Github issue search https://github.com/apollographql/federation/issues/356. Turns out we were not alone! The current version of Apollo Federation and Gateway do not support the cacheControl directives 😱.

One En-genius Solution

Thankfully one user in the comments came up with a pretty en-genius solution. They proposed code that would do the following:

Store the cache control max age header from all downstream services that were called
Sort the max-age values received grabbing the lowest value received
If there was a value received, set the lowest max-age as our response cache header

This was great, but it didn’t really support our use case. The code did not take into account private vs public cache headers, which we wanted to take advantage of, and it also manually set the header instead of using Apollo’s cache control request context.

A Simple Rewrite

Utilizing the code from the GitHub comment we began a simple rewrite that would:

Convert us to Typescript (because let’s be honest types rule)
Support Public & Private Cache Control headers
Utilize the built in overallCachePolicy context in Apollo

RemoteGraphQLDataSource

The first step was to save off our cacheControl headers that our down stream services responded with. To do this we write a new RemoteGraphQLDataSource that we would tell Apollo to use when calling our downstream services.

An empty RemoteGraphQLDataSource looks like this:

import { RemoteGraphQLDataSource } from '@apollo/gateway';
import {
  GraphQLResponse,
  GraphQLRequest,
  BaseContext,
} from 'apollo-server-types';

export class CachedDataSource extends RemoteGraphQLDataSource {
  async didReceiveResponse({
    response,
    request,
    context,
  }: {
    response: GraphQLResponse;
    request: GraphQLRequest;
    context: BaseContext;
  }): Promise<GraphQLResponse> {
    if (super.didReceiveResponse) {
      return super.didReceiveResponse({ response, request, context });
    }
    return response;
  }
}

First we add in some code to get the response’s header:

const cacheControl = response.http.headers.get('Cache-Control');

Then we check to see if we already have an array of cache headers in our Apollo Context and if not we create it.

if (!context.cacheControl || !Array.isArray(context.cacheControl)) {
  context.cacheControl = [];
}

And finally we add the header we received from our response into the context of the overall request.

context.cacheControl.push(cacheControl);

The final result then looks like:

import { RemoteGraphQLDataSource } from '@apollo/gateway';
import {
  GraphQLResponse,
  GraphQLRequest,
  BaseContext,
} from 'apollo-server-types';
export class CachedDataSource extends RemoteGraphQLDataSource {
  async didReceiveResponse({
    response,
    request,
    context,
  }: {
    response: GraphQLResponse;
    request: GraphQLRequest;
    context: BaseContext;
  }): Promise<GraphQLResponse> {
    const cacheControl = response.http.headers.get('Cache-Control');

    if (!context.cacheControl || !Array.isArray(context.cacheControl)) {
      context.cacheControl = [];
    }

    context.cacheControl.push(cacheControl);
    if (super.didReceiveResponse) {
      return super.didReceiveResponse({ response, request, context });
    }

    return response;
  }
}

And using it is as simple as adding it to the Gateway configuration when it builds the service:

const gateway = new ApolloGateway({
  buildService({ url }) {
    return new CachedDataSource({url});
  },
});

This will tell the gateway to use the CachedDataSource when it builds the downstream services instead of RemoteGraphQLDataSource.

CacheControlHeaderPlugin

So at this point we have Apollo collecting the downstream cache headers but we haven’t done anything with them yet. Thankfully Apollo has a great plugin system.

All we need to do is write a plugin that takes our saved values in the context and calculates what the overall cache strategy should be from that.

Start with a simple Apollo Plugin definition that will let us access the requestContext before Apollo decides to write the response.

import { PluginDefinition } from 'apollo-server-core';
export const CacheControlHeaderPlugin: PluginDefinition = {
  requestDidStart() {
    return {
      willSendResponse(requestContext) {
        
      },
    };
  },
};

From there we eventually want to look like this:

export const CacheControlHeaderPlugin: PluginDefinition = {
  requestDidStart() {
    return {
      willSendResponse(requestContext) {
        requestContext.overallCachePolicy = calculateOverallCachePolicy(
          requestContext.context.cacheControl,
          config.apollo.defaultMaxAge
        );
      },
    };
  },
};

Here we grab the context from the requestContext (I know dual context??, weird right 🤷🏼) and look cacheControl! That’s the cacheControl array that we created and added to our DataSource. From there we take that cacheControl array and our default value we like to set at Pocket and pass it to a function to help us calculate the overallCachePolicy.

By setting the overallCachePolicy on the outer request context, Apollo will automatically set the headers when it builds the response. For the curious overallCachePolicy has the following interface:

export interface CacheHint {
  maxAge?: number;
  scope?: CacheScope;
}

Scope here is an enum of Private or Public.

Armed with all this knowledge lets build our calculation function, but let’s do it by breaking it up.

/**
 * Calculates the minimum max-age and privacy between the headers, returns a new overall cache policy
 *
 * @param {string[]} cacheControl - array of cache-control headers
 */
export const calculateOverallCachePolicy = (
  cacheControl: string[] = [],
  defaultMaxAge?: number
): Required<CacheHint> | undefined => {
  const maxAge = calculateMaxAge(cacheControl, defaultMaxAge);
  return maxAge
    ? { maxAge: maxAge, scope: calculatePrivateOrPublic(cacheControl) }
    : undefined;
};

To calculate max age:

Run all the headers through our regex
Filter out the empty matches
Ensure there is more than 1 match group, if so return the match, otherwise set it to 0
Sort the results and return the lowest max-age

We are particularly using the lowest max-age here because we do not want to cache longer then the lowest allowed cache value, since the entire response will be cached.

/**
 * The regext for finding the max-age and public/private value
 */
const CacheHeaderRegex = /^max-age=([0-9]+), (public|private)$/;

/**
 * Calculates the minimum max-age from all the implementing service headers
 * @param cacheControl
 * @param defaultMaxAge
 */
export const calculateMaxAge = (
  cacheControl = [],
  defaultMaxAge?: number
): number | null => {
  let maxAge = cacheControl
    .map((h) => CacheHeaderRegex.exec(h))
    .map((matches) => matches || [])
    .map((matches) => (matches.length > 1 ? matches[1] : 0)) // eslint-disable-line no-magic-numbers
    .reduce((acc, val) => Math.min(acc, val), +Infinity);

  //If there is no cache control we need to use our default max age
  if ((!maxAge || maxAge == Infinity) && defaultMaxAge && defaultMaxAge > 0) {
    maxAge = defaultMaxAge;
  } else if (maxAge == Infinity || maxAge == 0) {
    return null;
  }

  return maxAge;
};

To calculate private or public, we follow almost the exact same steps as max-age.

Run all the headers through our regex
Filter out the empty matches
Ensure there are more than 2 match groups, if so return the match, otherwise set it to null
If the results contain private then set the scope to Private, otherwise set to Public.

Much like max-age we need to set the overall cache header to the most limiting value we have since the entire response will be cached.

const CacheHeaderRegex = /^max-age=([0-9]+), (public|private)$/;

/**
 * Calculates the cache privacy header from all the implementing service headers
 *
 * @param cacheControl
 * @param defaultMaxAge
 */
export const calculatePrivateOrPublic = (cacheControl = []): CacheScope => {
  const headerValues = cacheControl
    .map((h) => CacheHeaderRegex.exec(h))
    .map((matches) => matches || [])
    .map((matches) => (matches.length > 2 ? matches[2] : null));

  // Just having one private value makes the whole request private at the gateway level
  return headerValues.includes('private')
    ? CacheScope.Private
    : CacheScope.Public;
};

And finally we use our plugin in combination with the default apollo cache plugin:

import responseCachePlugin from 'apollo-server-plugin-response-cache';
import { CacheControlHeaderPlugin } from './';

// Pass the ApolloGateway to the ApolloServer constructor
const server = new ApolloServer({
  gateway,
  // Set a default cache control of 5 seconds so it will send cache headers.
  // Individual schemas can define headers on directives that the Gateway will then merge
  cacheControl: { defaultMaxAge: config.apollo.defaultMaxAge },
  plugins: [
    CacheControlHeaderPlugin,
    responseCachePlugin,
  ],
});

A Hopeful Future

Phew, that was a lot! Hopefully our stumbling helps you out on your federated graph adventure. Thankfully this hack shouldn’t be needed for too long. Cache Control support in Apollo federation is stated in the Apollo Server 3 roadmap, which at least means Apollo is aware of the need and is working towards giving us a more supported future. Until next time!

~ daniel (bass_rock), backend team

Tagged with: #graphql, #apollo, #federation, #caching