Invalidate hosted pending caches after mutations

Author: julianknutsen

cmd/wl/cmd_serve.go

@@ -423,6 +423,11 @@ func runServeHosted(cmd *cobra.Command, stdout, _ io.Writer) error {
 	// Build the API server with hosted workspace resolution.
 	apiServer := api.NewHostedWorkspace(hosted.NewClientFunc(), hosted.NewWorkspaceFunc())
 	apiServer.SetEnvironment(environment)
+	apiServer.SetMutationInvalidator(func(ctx context.Context) {
+		if connectionID, ok := hosted.ConnectionIDFromContext(ctx); ok {
+			resolver.InvalidateConnection(connectionID)
+		}
+	})
 
 	// Public read-only RemoteDB against the canonical hosted upstream (no token needed).
 	publicDB := newHostedPublicDB()

internal/api/handlers.go

@@ -342,6 +342,9 @@ func canonicalBrowseKey(r *http.Request) string {
 // invalidateReadCaches busts browse caches for the active upstream and detail
 // caches for the affected item on that upstream.
 func (s *Server) invalidateReadCaches(r *http.Request, client *sdk.Client, wantedID string) {
+	if s.mutationInvalidator != nil {
+		s.mutationInvalidator(r.Context())
+	}
 	if s.hosted && !s.hasCanonicalHostedReadIdentity(r) {
 		s.invalidateAllCaches()
 		return
@@ -352,6 +355,9 @@ func (s *Server) invalidateReadCaches(r *http.Request, client *sdk.Client, wante
 }
 
 func (s *Server) invalidateBrowseReadCaches(r *http.Request, client *sdk.Client) {
+	if s.mutationInvalidator != nil {
+		s.mutationInvalidator(r.Context())
+	}
 	if s.hosted && !s.hasCanonicalHostedReadIdentity(r) {
 		s.invalidateAllCaches()
 		return
@@ -375,6 +381,9 @@ func (s *Server) invalidateDetailCaches(upstream, wantedID string) {
 }
 
 func (s *Server) invalidateUpstreamReadCaches(r *http.Request, client *sdk.Client) {
+	if s.mutationInvalidator != nil {
+		s.mutationInvalidator(r.Context())
+	}
 	if s.hosted && !s.hasCanonicalHostedReadIdentity(r) {
 		s.invalidateAllCaches()
 		return

internal/api/infrastructure_test.go

@@ -675,6 +675,38 @@ func TestInvalidateReadCaches_TargetsOnlyMatchingUpstreamAndItem(t *testing.T) {
 	}
 }
 
+func TestInvalidateReadCaches_CallsMutationInvalidator(t *testing.T) {
+	srv := NewHostedWorkspace(func(*http.Request) (*sdk.Client, error) {
+		return sdk.New(sdk.ClientConfig{
+			RigHandle: "alice",
+			Upstream:  "hop/wl-commons",
+			Mode:      "wild-west",
+		}), nil
+	}, nil)
+
+	var calls atomic.Int32
+	srv.SetMutationInvalidator(func(context.Context) {
+		calls.Add(1)
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/api/wanted/w-1/reject-upstream", nil)
+	req = req.WithContext(WithResolvedReadIdentity(req.Context(), ResolvedReadIdentity{
+		Upstream: "hop/wl-commons",
+		Viewer:   "alice",
+	}))
+	client := sdk.New(sdk.ClientConfig{
+		RigHandle: "alice",
+		Upstream:  "hop/wl-commons",
+		Mode:      "wild-west",
+	})
+
+	srv.invalidateReadCaches(req, client, "w-1")
+
+	if got := calls.Load(); got != 1 {
+		t.Fatalf("mutation invalidator calls = %d, want 1", got)
+	}
+}
+
 func TestInvalidateUpstreamReadCaches_TargetsOnlyMatchingUpstream(t *testing.T) {
 	srv := NewHostedWorkspace(func(*http.Request) (*sdk.Client, error) {
 		return sdk.New(sdk.ClientConfig{

internal/api/server.go

@@ -5,6 +5,7 @@
 package api
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -21,18 +22,19 @@ type WorkspaceFunc func(r *http.Request) (*sdk.Workspace, error)
 
 // Server is the HTTP API server wrapping an SDK client.
 type Server struct {
-	clientFunc       ClientFunc
-	workspaceFunc    WorkspaceFunc
-	pile             pile.RowQuerier
-	scoreboard       *CachedEndpoint
-	scoreboardDetail *CachedEndpoint
-	scoreboardDump   *CachedEndpoint
-	publicClient     *sdk.Client // anonymous fallback for public reads (hosted mode)
-	browseCache      *ReadCache  // keyed by canonicalized query string
-	detailCache      *ReadCache  // keyed by item ID
-	environment      string
-	mux              *http.ServeMux
-	hosted           bool // true when running in multi-tenant hosted mode
+	clientFunc          ClientFunc
+	workspaceFunc       WorkspaceFunc
+	mutationInvalidator func(context.Context)
+	pile                pile.RowQuerier
+	scoreboard          *CachedEndpoint
+	scoreboardDetail    *CachedEndpoint
+	scoreboardDump      *CachedEndpoint
+	publicClient        *sdk.Client // anonymous fallback for public reads (hosted mode)
+	browseCache         *ReadCache  // keyed by canonicalized query string
+	detailCache         *ReadCache  // keyed by item ID
+	environment         string
+	mux                 *http.ServeMux
+	hosted              bool // true when running in multi-tenant hosted mode
 }
 
 // New creates a Server backed by the given SDK client.
@@ -115,6 +117,13 @@ func (s *Server) SetEnvironment(environment string) {
 	s.environment = environment
 }
 
+// SetMutationInvalidator registers a callback that runs after successful
+// mutations invalidate API read caches. Hosted mode uses this to evict
+// resolver-owned caches that live beneath the HTTP layer.
+func (s *Server) SetMutationInvalidator(fn func(context.Context)) {
+	s.mutationInvalidator = fn
+}
+
 // ScoreboardHandler returns an http.HandlerFunc for the scoreboard endpoint.
 func (s *Server) ScoreboardHandler() http.HandlerFunc {
 	return s.handleScoreboard

internal/hosted/auth.go

@@ -13,8 +13,9 @@ import (
 type contextKey string
 
 const (
-	clientContextKey    contextKey = "hosted-client"
-	workspaceContextKey contextKey = "hosted-workspace"
+	clientContextKey     contextKey = "hosted-client"
+	workspaceContextKey  contextKey = "hosted-workspace"
+	connectionContextKey contextKey = "hosted-connection-id"
 )
 
 // ClientFromContext extracts the sdk.Client injected by auth middleware.
@@ -29,6 +30,13 @@ func WorkspaceFromContext(ctx context.Context) (*sdk.Workspace, bool) {
 	return ws, ok
 }
 
+// ConnectionIDFromContext extracts the active hosted connection ID injected by
+// auth middleware.
+func ConnectionIDFromContext(ctx context.Context) (string, bool) {
+	connectionID, ok := ctx.Value(connectionContextKey).(string)
+	return connectionID, ok && connectionID != ""
+}
+
 // AuthMiddleware protects /api/* routes (excluding /api/auth/*).
 // It resolves the session cookie, looks up the Nango connection, and injects
 // the per-user sdk.Workspace and active sdk.Client into the request context.
@@ -152,6 +160,7 @@ func (s *Server) AuthMiddleware(next http.Handler) http.Handler {
 
 		// Inject both workspace and client into context.
 		ctx := r.Context()
+		ctx = context.WithValue(ctx, connectionContextKey, session.ConnectionID)
 		ctx = context.WithValue(ctx, workspaceContextKey, workspace)
 		ctx = context.WithValue(ctx, clientContextKey, client)
 		ctx = api.WithResolvedReadIdentity(ctx, api.ResolvedReadIdentity{

internal/hosted/authservice_resolver.go

@@ -184,8 +184,22 @@ func (wr *AuthServiceWorkspaceResolver) resolveFromConnection(_ context.Context,
 // InvalidateConnection evicts any cached workspace for the given connection.
 func (wr *AuthServiceWorkspaceResolver) InvalidateConnection(connectionID string) {
 	wr.mu.Lock()
-	defer wr.mu.Unlock()
 	delete(wr.cache, connectionID)
+	wr.mu.Unlock()
+
+	prefix := connectionID + ":"
+	wr.pendingMu.Lock()
+	caches := make([]*pendingUpstreamCache, 0, len(wr.pendingCache))
+	for key, cache := range wr.pendingCache {
+		if strings.HasPrefix(key, prefix) {
+			caches = append(caches, cache)
+			delete(wr.pendingCache, key)
+		}
+	}
+	wr.pendingMu.Unlock()
+	for _, cache := range caches {
+		cache.Stop()
+	}
 }
 
 func (wr *AuthServiceWorkspaceResolver) cachedWorkspace(connectionID string) (*sdk.Workspace, bool) {

internal/hosted/authservice_resolver_test.go

@@ -0,0 +1,26 @@
+package hosted
+
+import (
+	"testing"
+	"time"
+)
+
+func TestAuthServiceWorkspaceResolver_InvalidateConnectionClearsPendingCaches(t *testing.T) {
+	resolver := NewAuthServiceWorkspaceResolver(nil, NewSessionStore())
+	cache1 := newPendingUpstreamCache(nil, "hop", "wl-commons", time.Hour)
+	cache2 := newPendingUpstreamCache(nil, "gastownhall", "gascity", time.Hour)
+	defer cache1.Stop()
+	defer cache2.Stop()
+
+	resolver.pendingCache["conn-1:hop/wl-commons"] = cache1
+	resolver.pendingCache["conn-2:gastownhall/gascity"] = cache2
+
+	resolver.InvalidateConnection("conn-1")
+
+	if _, ok := resolver.pendingCache["conn-1:hop/wl-commons"]; ok {
+		t.Fatal("expected conn-1 pending cache to be evicted")
+	}
+	if _, ok := resolver.pendingCache["conn-2:gastownhall/gascity"]; !ok {
+		t.Fatal("expected unrelated pending cache to remain")
+	}
+}

internal/hosted/authservice_server.go

@@ -554,6 +554,7 @@ func (s *AuthServiceServer) AuthMiddleware(next http.Handler) http.Handler {
 			Upstream: upstream,
 			Viewer:   workspace.RigHandle(),
 		})
+		ctx = withConnectionID(ctx, session.ConnectionID)
 		ctx = withWorkspaceAndClient(ctx, workspace, client)
 		next.ServeHTTP(w, r.WithContext(ctx))
 	})
@@ -578,3 +579,7 @@ func withWorkspaceAndClient(ctx context.Context, workspace *sdk.Workspace, clien
 	ctx = context.WithValue(ctx, clientContextKey, client)
 	return ctx
 }
+
+func withConnectionID(ctx context.Context, connectionID string) context.Context {
+	return context.WithValue(ctx, connectionContextKey, connectionID)
+}

internal/hosted/resolver.go

@@ -388,8 +388,35 @@ func (wr *WorkspaceResolver) resolveFromMetadata(_ context.Context, connectionID
 // InvalidateConnection removes the cached workspace for a connection.
 func (wr *WorkspaceResolver) InvalidateConnection(connectionID string) {
 	wr.mu.Lock()
-	defer wr.mu.Unlock()
+	var upstreams []string
+	if cached, ok := wr.cache[connectionID]; ok && cached != nil && cached.workspace != nil {
+		infos := cached.workspace.Upstreams()
+		upstreams = make([]string, 0, len(infos))
+		for _, info := range infos {
+			if info.Upstream != "" {
+				upstreams = append(upstreams, info.Upstream)
+			}
+		}
+	}
 	delete(wr.cache, connectionID)
+	wr.mu.Unlock()
+
+	if len(upstreams) == 0 {
+		return
+	}
+
+	wr.pendingMu.Lock()
+	caches := make([]*pendingUpstreamCache, 0, len(wr.pendingCache))
+	for _, upstream := range upstreams {
+		if cache, ok := wr.pendingCache[upstream]; ok {
+			caches = append(caches, cache)
+			delete(wr.pendingCache, upstream)
+		}
+	}
+	wr.pendingMu.Unlock()
+	for _, cache := range caches {
+		cache.Stop()
+	}
 }
 
 func (wr *WorkspaceResolver) cachedWorkspace(connectionID string) (*sdk.Workspace, bool) {

internal/hosted/resolver_test.go

@@ -439,8 +439,18 @@ func TestWorkspaceResolver_InvalidateConnection(t *testing.T) {
 		t.Fatalf("first resolve: %v", err)
 	}
 
+	cache := resolver.getOrCreatePendingCache(nil, "wasteland", "wl-commons")
+	defer cache.Stop()
+	if _, ok := resolver.pendingCache["wasteland/wl-commons"]; !ok {
+		t.Fatal("expected pending cache to be stored before invalidation")
+	}
+
 	resolver.InvalidateConnection("conn-1")
 
+	if _, ok := resolver.pendingCache["wasteland/wl-commons"]; ok {
+		t.Fatal("expected pending cache to be evicted for invalidated connection")
+	}
+
 	ws2, err := resolver.Resolve(session)
 	if err != nil {
 		t.Fatalf("second resolve: %v", err)