diff --git a/router-tests/modules/cache-tags/module.go b/router-tests/modules/cache-tags/module.go new file mode 100644 index 0000000000..aad6b7d206 --- /dev/null +++ b/router-tests/modules/cache-tags/module.go @@ -0,0 +1,230 @@ +// Package cache_tags provides a custom router module that collects cache tags +// and Cache-Control policies from subgraph responses. It targets synchronous +// GraphQL responses; HTTP headers cannot be changed after a deferred or +// streaming response has started. +package cache_tags + +import ( + "fmt" + "net/http" + "sort" + "strconv" + "strings" + "sync" + + "github.com/wundergraph/cosmo/router/core" + "golang.org/x/net/http/httpguts" +) + +const ( + moduleID = "cacheTagsModule" + requestStateContextKey = "cache_tags_module_request_state" + cacheControlHeaderName = "Cache-Control" + cacheTagValueDelimiter = "," + cacheTagOutputDelimiter = "," + cacheControlJoinDelimiter = ", " +) + +// CacheTagsModule collects the configured header from every subgraph response +// and exposes the merged values on the federated response. +type CacheTagsModule struct { + HeaderName string `mapstructure:"header_name"` +} + +func (m *CacheTagsModule) Provision(_ *core.ModuleContext) error { + headerName := strings.TrimSpace(m.HeaderName) + if headerName == "" { + return fmt.Errorf("header_name must not be empty") + } + if !httpguts.ValidHeaderFieldName(headerName) { + return fmt.Errorf("header_name %q is not a valid HTTP header name", headerName) + } + + m.HeaderName = http.CanonicalHeaderKey(headerName) + switch m.HeaderName { + case cacheControlHeaderName, "Connection", "Content-Length", "Trailer", "Transfer-Encoding": + return fmt.Errorf("header_name must not be a reserved response header: %s", m.HeaderName) + } + + return nil +} + +// Middleware initializes the aggregation state before any subgraph requests +// are made. Module instances are shared by all requests, so the mutable state +// must live on the request context instead of the module. +func (m *CacheTagsModule) Middleware(ctx core.RequestContext, next http.Handler) { + state := &requestState{ + responseWriter: ctx.ResponseWriter(), + tags: make(map[string]struct{}), + } + ctx.Set(requestStateContextKey, state) + + next.ServeHTTP(ctx.ResponseWriter(), ctx.Request()) +} + +// OnOriginRequest intentionally leaves the request unchanged. Registering a +// pre-origin handler makes the router disable request deduplication unless it +// is force-enabled. That guarantees every client request receives its own +// OnOriginResponse callbacks and therefore its own collected headers. +func (m *CacheTagsModule) OnOriginRequest(request *http.Request, _ core.RequestContext) (*http.Request, *http.Response) { + return request, nil +} + +// OnOriginResponse is called concurrently for subgraph responses. It adds the +// response's cache tags to the request-local set and retains the most +// restrictive Cache-Control policy. +func (m *CacheTagsModule) OnOriginResponse(response *http.Response, ctx core.RequestContext) *http.Response { + if response == nil { + return nil + } + + value, ok := ctx.Get(requestStateContextKey) + if !ok { + return nil + } + state, ok := value.(*requestState) + if !ok { + return nil + } + + state.merge(m.HeaderName, response.Header) + + // Returning a non-nil response would short-circuit the remaining + // post-origin handlers. This module only observes the response. + return nil +} + +func (m *CacheTagsModule) Module() core.ModuleInfo { + return core.ModuleInfo{ + ID: moduleID, + Priority: 1, + New: func() core.Module { + return &CacheTagsModule{} + }, + } +} + +type requestState struct { + mu sync.Mutex + responseWriter http.ResponseWriter + tags map[string]struct{} + cacheControl cacheControlPolicy +} + +func (s *requestState) merge(tagHeaderName string, subgraphHeaders http.Header) { + s.mu.Lock() + defer s.mu.Unlock() + + for _, value := range subgraphHeaders.Values(tagHeaderName) { + for tag := range strings.SplitSeq(value, cacheTagValueDelimiter) { + tag = strings.TrimSpace(tag) + if tag != "" { + s.tags[tag] = struct{}{} + } + } + } + + if policy, ok := parseCacheControl(subgraphHeaders.Values(cacheControlHeaderName)); ok { + s.cacheControl.merge(policy) + } + + s.applyHeaders(tagHeaderName) +} + +func (s *requestState) applyHeaders(tagHeaderName string) { + if len(s.tags) > 0 { + tags := make([]string, 0, len(s.tags)) + for tag := range s.tags { + tags = append(tags, tag) + } + sort.Strings(tags) + s.responseWriter.Header().Set(tagHeaderName, strings.Join(tags, cacheTagOutputDelimiter)) + } + + if s.cacheControl.present { + s.responseWriter.Header().Set(cacheControlHeaderName, s.cacheControl.headerValue()) + } +} + +type cacheControlPolicy struct { + present bool + noStore bool + noCache bool + private bool + public bool + hasMaxAge bool + maxAge uint64 +} + +// parseCacheControl extracts the directives needed to build the combined +// policy. Multiple max-age directives are reduced to their lowest value. +func parseCacheControl(values []string) (cacheControlPolicy, bool) { + policy := cacheControlPolicy{} + for directive := range strings.SplitSeq(strings.Join(values, cacheControlJoinDelimiter), ",") { + name, value, hasValue := strings.Cut(strings.TrimSpace(directive), "=") + switch { + case strings.EqualFold(name, "no-store"): + policy.present = true + policy.noStore = true + case strings.EqualFold(name, "no-cache"): + policy.present = true + policy.noCache = true + case strings.EqualFold(name, "private"): + policy.present = true + policy.private = true + case strings.EqualFold(name, "public"): + policy.present = true + policy.public = true + case strings.EqualFold(name, "max-age") && hasValue: + maxAge, err := strconv.ParseUint(strings.Trim(strings.TrimSpace(value), `"`), 10, 64) + if err == nil && (!policy.hasMaxAge || maxAge < policy.maxAge) { + policy.present = true + policy.hasMaxAge = true + policy.maxAge = maxAge + } + } + } + + return policy, policy.present +} + +func (p *cacheControlPolicy) merge(other cacheControlPolicy) { + p.present = p.present || other.present + p.noStore = p.noStore || other.noStore + p.noCache = p.noCache || other.noCache + p.private = p.private || other.private + p.public = p.public || other.public + if other.hasMaxAge && (!p.hasMaxAge || other.maxAge < p.maxAge) { + p.hasMaxAge = true + p.maxAge = other.maxAge + } +} + +func (p cacheControlPolicy) headerValue() string { + if p.noStore { + return "no-store" + } + + parts := make([]string, 0, 2) + if p.noCache { + parts = append(parts, "no-cache") + } else if p.hasMaxAge { + parts = append(parts, fmt.Sprintf("max-age=%d", p.maxAge)) + } + + if p.private { + parts = append(parts, "private") + } else if p.public { + parts = append(parts, "public") + } + + return strings.Join(parts, cacheControlJoinDelimiter) +} + +var ( + _ core.Module = (*CacheTagsModule)(nil) + _ core.Provisioner = (*CacheTagsModule)(nil) + _ core.RouterMiddlewareHandler = (*CacheTagsModule)(nil) + _ core.EnginePreOriginHandler = (*CacheTagsModule)(nil) + _ core.EnginePostOriginHandler = (*CacheTagsModule)(nil) +) diff --git a/router-tests/modules/cache_tags_test.go b/router-tests/modules/cache_tags_test.go new file mode 100644 index 0000000000..3638b994ac --- /dev/null +++ b/router-tests/modules/cache_tags_test.go @@ -0,0 +1,341 @@ +package module_test + +import ( + "fmt" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + cachetags "github.com/wundergraph/cosmo/router-tests/modules/cache-tags" + "github.com/wundergraph/cosmo/router-tests/testenv" + "github.com/wundergraph/cosmo/router/core" + "github.com/wundergraph/cosmo/router/pkg/config" +) + +const ( + cacheTagsModuleID = "cacheTagsModule" + + queryEmployeeWithHobbies = `{ + employee(id: 1) { + id + hobbies { + ... on Gaming { + name + } + } + } + }` +) + +func TestCacheTagsModuleConfiguration(t *testing.T) { + t.Parallel() + + for _, headerName := range []string{"", "invalid header", "Cache-Control", "Content-Length"} { + t.Run(headerName, func(t *testing.T) { + t.Parallel() + + module := &cachetags.CacheTagsModule{HeaderName: headerName} + require.Error(t, module.Provision(&core.ModuleContext{})) + }) + } + + module := &cachetags.CacheTagsModule{HeaderName: " x-cache-tags "} + require.NoError(t, module.Provision(&core.ModuleContext{})) + require.Equal(t, "X-Cache-Tags", module.HeaderName) +} + +func TestCacheTagsModule(t *testing.T) { + t.Parallel() + + t.Run("merges configured cache tags and uses the lowest max age", func(t *testing.T) { + t.Parallel() + + const tagHeader = "X-Cache-Tags" + + testenv.Run(t, &testenv.Config{ + RouterOptions: cacheTagsModuleOptions(tagHeader), + Subgraphs: testenv.SubgraphsConfig{ + Employees: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"employee:1, shared", "employee:2"}, + "Cache-Control": {"max-age=120, public"}, + }), + }, + Hobbies: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"hobby:gaming, shared"}, + "Cache-Control": {"max-age=90", "max-age=60, public"}, + }), + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: queryEmployeeWithHobbies, + }) + + require.Equal(t, http.StatusOK, res.Response.StatusCode) + require.Equal(t, "employee:1,employee:2,hobby:gaming,shared", res.Response.Header.Get(tagHeader)) + require.Equal(t, "max-age=60, public", res.Response.Header.Get("Cache-Control")) + require.JSONEq(t, `{"data":{"employee":{"id":1,"hobbies":[{},{"name":"Counter Strike"},{},{},{}]}}}`, res.Body) + }) + }) + + t.Run("combines restrictive cache control directives", func(t *testing.T) { + t.Parallel() + + const tagHeader = "X-Cache-Tags" + + testenv.Run(t, &testenv.Config{ + RouterOptions: cacheTagsModuleOptions(tagHeader), + Subgraphs: testenv.SubgraphsConfig{ + Employees: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"employees"}, + "Cache-Control": {"max-age=600, private"}, + }), + }, + Hobbies: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"hobbies"}, + "Cache-Control": {"max-age=300, public"}, + }), + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: queryEmployeeWithHobbies, + }) + + require.Equal(t, "employees,hobbies", res.Response.Header.Get(tagHeader)) + require.Equal(t, "max-age=300, private", res.Response.Header.Get("Cache-Control")) + }) + }) + + t.Run("only collects the configured tag header", func(t *testing.T) { + t.Parallel() + + const ( + tagHeader = "X-Entity-Tags" + otherHeader = "X-Cache-Tags" + ) + + testenv.Run(t, &testenv.Config{ + RouterOptions: cacheTagsModuleOptions(tagHeader), + Subgraphs: testenv.SubgraphsConfig{ + Employees: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"employee:1"}, + otherHeader: {"not-collected:employees"}, + }), + }, + Hobbies: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"hobby:gaming"}, + otherHeader: {"not-collected:hobbies"}, + }), + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: queryEmployeeWithHobbies, + }) + + require.Equal(t, "employee:1,hobby:gaming", res.Response.Header.Get(tagHeader)) + require.Empty(t, res.Response.Header.Get(otherHeader)) + require.Empty(t, res.Response.Header.Get("Cache-Control")) + }) + }) + + t.Run("is race safe for parallel subgraph responses", func(t *testing.T) { + t.Parallel() + + const tagHeader = "X-Cache-Tags" + + testenv.Run(t, &testenv.Config{ + RouterOptions: cacheTagsModuleOptions(tagHeader), + Subgraphs: testenv.SubgraphsConfig{ + Employees: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"employees"}, + "Cache-Control": {"max-age=300"}, + }), + }, + Hobbies: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"hobbies"}, + "Cache-Control": {"max-age=0"}, + }), + }, + Availability: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"availability"}, + "Cache-Control": {"max-age=120"}, + }), + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: `{ + employees { + id + isAvailable + hobbies { + ... on Gaming { + name + } + } + } + }`, + }) + + require.Equal(t, http.StatusOK, res.Response.StatusCode) + require.Equal(t, "availability,employees,hobbies", res.Response.Header.Get(tagHeader)) + require.Equal(t, "max-age=0", res.Response.Header.Get("Cache-Control")) + }) + }) + + t.Run("returns collected headers to concurrent clients", func(t *testing.T) { + t.Parallel() + + const tagHeader = "X-Cache-Tags" + + routerOptions := append( + cacheTagsModuleOptions(tagHeader), + core.WithEngineExecutionConfig(config.EngineExecutionConfiguration{ + EnableSingleFlight: true, + ForceEnableSingleFlight: false, + MaxConcurrentResolvers: 0, + }), + ) + + testenv.Run(t, &testenv.Config{ + RouterOptions: routerOptions, + Subgraphs: testenv.SubgraphsConfig{ + GlobalDelay: 100 * time.Millisecond, + Employees: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"employees"}, + "Cache-Control": {"max-age=120"}, + }), + }, + Hobbies: testenv.SubgraphConfig{ + Middleware: subgraphHeadersMiddleware(http.Header{ + tagHeader: {"hobbies"}, + "Cache-Control": {"max-age=60"}, + }), + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + responses := makeConcurrentGraphQLRequests(t, xEnv, queryEmployeeWithHobbies, 5) + for _, res := range responses { + require.Equal(t, "employees,hobbies", res.Response.Header.Get(tagHeader)) + require.Equal(t, "max-age=60", res.Response.Header.Get("Cache-Control")) + } + }) + }) + + t.Run("keeps aggregation isolated between requests", func(t *testing.T) { + t.Parallel() + + const tagHeader = "X-Cache-Tags" + var requestCount atomic.Uint32 + + testenv.Run(t, &testenv.Config{ + RouterOptions: cacheTagsModuleOptions(tagHeader), + Subgraphs: testenv.SubgraphsConfig{ + Employees: testenv.SubgraphConfig{ + Middleware: func(handler http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestNumber := requestCount.Add(1) + w.Header().Set(tagHeader, fmt.Sprintf("request:%d", requestNumber)) + handler.ServeHTTP(w, r) + }) + }, + }, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + first := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: `{ employee(id: 1) { id } }`, + }) + second := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: `{ employee(id: 2) { id } }`, + }) + + require.Equal(t, "request:1", first.Response.Header.Get(tagHeader)) + require.Equal(t, "request:2", second.Response.Header.Get(tagHeader)) + }) + }) +} + +func makeConcurrentGraphQLRequests( + t *testing.T, + xEnv *testenv.Environment, + query string, + requestCount int, +) []*testenv.TestResponse { + t.Helper() + + var ready, done sync.WaitGroup + ready.Add(requestCount) + done.Add(requestCount) + + trigger := make(chan struct{}) + errs := make(chan error, requestCount) + responses := make([]*testenv.TestResponse, requestCount) + for i := range requestCount { + go func() { + defer done.Done() + ready.Done() + <-trigger + + response, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{Query: query}) + if err != nil { + errs <- err + return + } + responses[i] = response + }() + } + + ready.Wait() + close(trigger) + done.Wait() + close(errs) + + for err := range errs { + require.NoError(t, err) + } + for _, response := range responses { + require.NotNil(t, response) + require.Equal(t, http.StatusOK, response.Response.StatusCode) + } + + return responses +} + +func cacheTagsModuleOptions(headerName string) []core.Option { + return []core.Option{ + core.WithModulesConfig(map[string]any{ + cacheTagsModuleID: map[string]any{ + "header_name": headerName, + }, + }), + core.WithCustomModules(&cachetags.CacheTagsModule{}), + } +} + +func subgraphHeadersMiddleware(headers http.Header) func(http.Handler) http.Handler { + return func(handler http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + for name, values := range headers { + for _, value := range values { + w.Header().Add(name, value) + } + } + handler.ServeHTTP(w, r) + }) + } +}