refactor: [OM-387] governance entitlement performance improvement 2

e2e/Makefile

@@ -50,6 +50,23 @@ bench-governance-trace: ## Capture traces for ONE small cell (default metered 10
 	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_KIND=$(TRACE_KIND) \
 		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$/kind=$(TRACE_KIND)/$(BENCH_CELL)$$' -benchmem -benchtime=$(TRACE_BENCHTIME) -count=1 ./...
 
+# --- Over-compute demonstration ---
+# Selective query: seed 50 entitlements/customer but request only 10. Reveals that GetAccess
+# computes all 50 balances (50 ClickHouse queries) while the result uses 10 — cost scales with a
+# customer's TOTAL entitlements, not what the query asks for. The regular cells hide this (they
+# request every seeded feature). Requires GOV_BENCH_OVERCOMPUTE=1 (set by these targets).
+.PHONY: bench-governance-overcompute
+bench-governance-overcompute: ## Latency: selective query (50 seeded / 10 queried), 1 and 10 customers
+	$(call print-target)
+	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_KIND=metered GOV_BENCH_OVERCOMPUTE=1 \
+		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$/kind=metered/' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./...
+
+.PHONY: bench-governance-overcompute-trace
+bench-governance-overcompute-trace: ## Capture ONE trace of the over-compute cell (1 customer, 50 seeded, 10 queried)
+	$(call print-target)
+	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_KIND=metered GOV_BENCH_OVERCOMPUTE=1 \
+		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$/kind=metered/customers=1$$/features=50$$/query=10$$' -benchmem -benchtime=$(TRACE_BENCHTIME) -count=1 ./...
+
 .PHONY: test-local
 test-local: ## Run tests against local openmeter
 	$(call print-target)

e2e/governance_bench_test.go

@@ -98,44 +98,70 @@ func BenchmarkGovernanceQuery(b *testing.B) {
 	type size struct {
 		name      string
 		customers int
-		features  int
+		// features is how many features/entitlements are SEEDED per customer (the real cost
+		// driver: GetAccess computes a balance for every one of a customer's entitlements).
+		features int
+		// queryFeatures is how many of those features the query actually REQUESTS; 0 means all.
+		// When queryFeatures < features the query is "selective": the caller asks about a few
+		// features of a customer that has many entitlements. This decouples what's requested (M)
+		// from what's computed (N) to expose the over-compute — GetAccess currently resolves all
+		// N regardless of M. The default cells keep queryFeatures==features (N==M), which HIDES
+		// the over-compute; only GOV_BENCH_OVERCOMPUTE cells reveal it.
+		queryFeatures int
 	}
 
 	// Diagonal + fixed-overhead baseline (default).
 	sizes := []size{
-		{"customers=1/features=1", 1, 1},
-		{"customers=10/features=10", 10, 10},
-		{"customers=50/features=50", 50, 50},
-		{"customers=100/features=100", 100, 100},
+		{"customers=1/features=1", 1, 1, 0},
+		{"customers=10/features=10", 10, 10, 0},
+		{"customers=50/features=50", 50, 50, 0},
+		{"customers=100/features=100", 100, 100, 0},
 	}
 
 	// Full 3x3 matrix isolates the customer axis from the feature axis.
 	if os.Getenv("GOV_BENCH_FULL_MATRIX") != "" {
 		sizes = nil
-		sizes = append(sizes, size{"customers=1/features=1", 1, 1})
+		sizes = append(sizes, size{"customers=1/features=1", 1, 1, 0})
 		for _, c := range []int{10, 50, 100} {
 			for _, f := range []int{10, 50, 100} {
-				sizes = append(sizes, size{fmt.Sprintf("customers=%d/features=%d", c, f), c, f})
+				sizes = append(sizes, size{fmt.Sprintf("customers=%d/features=%d", c, f), c, f, 0})
 			}
 		}
 	}
 
+	// Over-compute cells: seed N entitlements/customer but query only M<N features. Reveals that
+	// GetAccess computes all N balances (N ClickHouse queries) while the result uses only M — a
+	// selective query against a customer with many entitlements. customers=1 gives one clean trace
+	// (N GetEntitlementBalance spans, M-feature result); customers=10 gives a latency signal.
+	if os.Getenv("GOV_BENCH_OVERCOMPUTE") != "" {
+		sizes = []size{
+			{"customers=1/features=50/query=10", 1, 50, 10},
+			{"customers=10/features=50/query=10", 10, 50, 10},
+		}
+	}
+
 	for _, kind := range selectedKinds() {
 		for _, s := range sizes {
 			b.Run(fmt.Sprintf("kind=%s/%s", kind, s.name), func(b *testing.B) {
 				custKeys, featKeys := seedGovernanceFixture(b, client, s.customers, s.features, kind)
 
+				// Query all seeded features unless the cell requests a selective subset.
+				queryKeys := featKeys
+				if s.queryFeatures > 0 && s.queryFeatures < len(featKeys) {
+					queryKeys = featKeys[:s.queryFeatures]
+				}
+
 				reqBody := apiv3.GovernanceQueryRequest{
 					Customer: apiv3.GovernanceQueryRequestCustomers{Keys: custKeys},
-					Feature:  &apiv3.GovernanceQueryRequestFeatures{Keys: featKeys},
+					Feature:  &apiv3.GovernanceQueryRequestFeatures{Keys: queryKeys},
 				}
 
 				// Warm-up + correctness gate: a wrong result (e.g. missing customers)
 				// would make the latency number meaningless.
 				status, resp, problem := v3.QueryGovernance(reqBody)
 				require.Equalf(b, http.StatusOK, status, "governance query failed: %+v", problem)
 				require.Lenf(b, resp.Data, s.customers, "expected %d resolved customers", s.customers)
-				require.Lenf(b, resp.Data[0].Features, s.features, "expected %d features per customer", s.features)
+				require.Lenf(b, resp.Data[0].Features, len(queryKeys), "expected %d features per customer", len(queryKeys))
 
 				b.ReportAllocs()
 				b.ResetTimer()

openmeter/entitlement/connector.go

@@ -84,9 +84,14 @@ type Service interface {
 	// For consistency, it is forbidden for entitlements to be created for featueres the keys of which could be mistaken for entitlement IDs.
 	GetEntitlementOfCustomerAt(ctx context.Context, namespace string, customerID string, idOrFeatureKey string, at time.Time) (*Entitlement, error)
 
-	// GetAccess returns the access of a customer.
-	// It returns a map of featureKey to entitlement value + ID.
-	GetAccess(ctx context.Context, namespace string, customerID string) (Access, error)
+	// GetAccess returns the access of a customer as a map of featureKey to entitlement value + ID.
+	//
+	// If featureKeys is empty, all of the customer's entitlements are resolved. If featureKeys is
+	// non-empty, only entitlements for those feature keys are resolved — a caller that only needs a
+	// subset (e.g. a filtered governance query) then avoids computing balances (and their ClickHouse
+	// usage queries) for entitlements it would discard. The returned map only ever contains feature
+	// keys the customer actually has an entitlement for.
+	GetAccess(ctx context.Context, namespace string, customerID string, featureKeys ...string) (Access, error)
 }
 
 type ListEntitlementsWithCustomerResult struct {

openmeter/entitlement/metered/grant_owner_adapter.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
-	"sync"
 	"time"
 
 	"github.com/samber/lo"
@@ -20,6 +19,7 @@ import (
 	"github.com/openmeterio/openmeter/pkg/framework/entutils"
 	"github.com/openmeterio/openmeter/pkg/framework/transaction"
 	"github.com/openmeterio/openmeter/pkg/models"
+	"github.com/openmeterio/openmeter/pkg/syncx"
 	"github.com/openmeterio/openmeter/pkg/timeutil"
 )
 
@@ -53,61 +53,53 @@ func NewEntitlementGrantOwnerAdapter(
 	}
 }
 
-// entitlementCacheCtxKey scopes a per-operation memo of owner-entitlement lookups. Within a
-// single balance calculation the same owner entitlement is read by several methods here
-// (DescribeOwner, GetStartOfMeasurement, GetUsagePeriodStartAt, GetResetTimelineInclusive),
-// each issuing its own GetEntitlement query for the same immutable row. The entitlement cannot
-// change for the duration of that read, so the cache lets those methods share one fetch.
-type entitlementCacheCtxKey struct{}
-
-// entitlementCacheEntry memoizes one owner's GetEntitlement result (value or error) so
-// concurrent readers of the same owner resolve to a single DB fetch.
-type entitlementCacheEntry struct {
-	once sync.Once
-	ent  *entitlement.Entitlement
-	err  error
-}
+// ownerEntitlementMemo dedups the owner-entitlement lookups made by this adapter within a single
+// balance calculation: DescribeOwner, GetStartOfMeasurement, GetUsagePeriodStartAt, and
+// GetResetTimelineInclusive all read the same immutable entitlement row. Installed by
+// withEntitlementCache.
+var ownerEntitlementMemo = syncx.NewContextMemo[string, *entitlement.Entitlement]()
 
-// withEntitlementCache installs a scoped cache for the owner-entitlement lookups made by this
-// adapter. It is opt-in: callers that do not install it (e.g. one-off value reads) hit the
-// database directly and are unaffected. Nesting is a no-op so the outermost scope wins.
-//
-// INVARIANT — install ONLY around read-only operations. Cached entries are never invalidated
-// within their scope, so if this cache is installed around a call stack that MUTATES the owner
-// entitlement, any getEntitlement after that write returns the stale, pre-write row. This is
-// safe today because the sole installer is GetEntitlementBalance, whose call tree performs no
-// entitlement writes (it only writes balance_snapshot); the only entitlement mutation on this
-// adapter, EndCurrentUsagePeriod, lives in the reset flow, which the cache never wraps.
+// ownerCustomerMemo dedups the owner-customer lookup made by DescribeOwner across an
+// entitlement.Service.GetAccess fan-out, which resolves many entitlements of the SAME customer and
+// would otherwise fetch that customer once per entitlement. Installed by WithCustomerCache.
+var ownerCustomerMemo = syncx.NewContextMemo[string, *customer.Customer]()
+
+// withEntitlementCache installs the owner-entitlement memo for a single balance calculation.
 //
-// It is unexported deliberately: keeping installation inside this package prevents callers from
-// enabling the cache around a write-containing stack. If a future read-only path wants the same
-// memo, install it here at that entry point — do not expose it. If you ever need it around a
-// stack that can write an entitlement, add an eviction (delete the owner key) on the write.
+// Unexported deliberately: the sole installer is GetEntitlementBalance (same package), a pure
+// read whose call tree performs no entitlement writes (it only writes balance_snapshot; the one
+// entitlement mutation on this adapter, EndCurrentUsagePeriod, lives in the reset flow the cache
+// never wraps). Keeping it in-package prevents callers from enabling it around a write-containing
+// stack — see the ContextMemo read-only invariant.
 func withEntitlementCache(ctx context.Context) context.Context {
-	if _, ok := ctx.Value(entitlementCacheCtxKey{}).(*sync.Map); ok {
-		return ctx
-	}
+	return ownerEntitlementMemo.Install(ctx)
+}
 
-	return context.WithValue(ctx, entitlementCacheCtxKey{}, &sync.Map{})
+// WithCustomerCache installs the owner-customer memo. Exported because it is installed one layer
+// up, in entitlement.Service.GetAccess, which fans out over a customer's entitlements. Opt-in:
+// without it, DescribeOwner fetches the customer directly.
+//
+// INVARIANT (see syncx.ContextMemo): install ONLY around read-only operations. GetAccess is a pure
+// read, so the memoized customer cannot go stale relative to a write within the call. Do NOT
+// install it around a stack that mutates the customer, or reads after that write serve a stale row.
+func WithCustomerCache(ctx context.Context) context.Context {
+	return ownerCustomerMemo.Install(ctx)
 }
 
-// getEntitlement fetches the owner entitlement, serving it from the scoped cache when one is
-// installed (see withEntitlementCache). Without a cache it falls back to a direct fetch,
-// preserving the previous behavior for callers that do not opt in.
+// getEntitlement fetches the owner entitlement, serving it from the scoped memo when installed
+// (see withEntitlementCache) and falling back to a direct fetch otherwise.
 func (e *entitlementGrantOwner) getEntitlement(ctx context.Context, owner models.NamespacedID) (*entitlement.Entitlement, error) {
-	cache, ok := ctx.Value(entitlementCacheCtxKey{}).(*sync.Map)
-	if !ok {
+	return ownerEntitlementMemo.GetOrLoad(ctx, owner.Namespace+"/"+owner.ID, func(ctx context.Context) (*entitlement.Entitlement, error) {
 		return e.entitlementRepo.GetEntitlement(ctx, owner)
-	}
-
-	v, _ := cache.LoadOrStore(owner.Namespace+"/"+owner.ID, &entitlementCacheEntry{})
-	entry := v.(*entitlementCacheEntry)
-
-	entry.once.Do(func() {
-		entry.ent, entry.err = e.entitlementRepo.GetEntitlement(ctx, owner)
 	})
+}
 
-	return entry.ent, entry.err
+// getCustomer fetches the owner customer, serving it from the scoped memo when installed (see
+// WithCustomerCache) and falling back to a direct fetch otherwise.
+func (e *entitlementGrantOwner) getCustomer(ctx context.Context, id customer.CustomerID) (*customer.Customer, error) {
+	return ownerCustomerMemo.GetOrLoad(ctx, id.Namespace+"/"+id.ID, func(ctx context.Context) (*customer.Customer, error) {
+		return e.customerService.GetCustomer(ctx, customer.GetCustomerInput{CustomerID: &id})
+	})
 }
 
 func (e *entitlementGrantOwner) DescribeOwner(ctx context.Context, id models.NamespacedID) (grant.Owner, error) {
@@ -153,11 +145,9 @@ func (e *entitlementGrantOwner) DescribeOwner(ctx context.Context, id models.Nam
 		FilterGroupBy: feature.MeterGroupByFilters,
 	}
 
-	cust, err := e.customerService.GetCustomer(ctx, customer.GetCustomerInput{
-		CustomerID: &customer.CustomerID{
-			Namespace: id.Namespace,
-			ID:        ent.CustomerID,
-		},
+	cust, err := e.getCustomer(ctx, customer.CustomerID{
+		Namespace: id.Namespace,
+		ID:        ent.CustomerID,
 	})
 	if err != nil {
 		return def, fmt.Errorf("failed to get customer: %w", err)

openmeter/entitlement/service/service.go

@@ -334,18 +334,36 @@ func (c *service) expandCustomers(ctx context.Context, entitlements []entitlemen
 	return custs, nil
 }
 
-func (c *service) GetAccess(ctx context.Context, namespace string, customerId string) (entitlement.Access, error) {
+func (c *service) GetAccess(ctx context.Context, namespace string, customerId string, featureKeys ...string) (entitlement.Access, error) {
 	now := clock.Now()
 
 	entitlements, err := c.GetEntitlementsOfCustomer(ctx, namespace, customerId, now)
 	if err != nil {
 		return entitlement.Access{}, err
 	}
 
+	// When the caller only needs a subset of features, drop the rest before the fan-out so we
+	// don't compute (and issue ClickHouse usage queries for) balances that would be discarded.
+	// Empty featureKeys means "all" — the fetch is a single query regardless, so filtering its
+	// result in memory is enough; the per-entitlement value calc is the cost this avoids.
+	if len(featureKeys) > 0 {
+		wanted := lo.SliceToMap(featureKeys, func(k string) (string, struct{}) { return k, struct{}{} })
+		entitlements = lo.Filter(entitlements, func(ent entitlement.Entitlement, _ int) bool {
+			_, ok := wanted[ent.FeatureKey]
+			return ok
+		})
+	}
+
 	if len(entitlements) == 0 {
 		return entitlement.Access{}, nil
 	}
 
+	// This fans out over all of the customer's entitlements; each metered entitlement's
+	// DescribeOwner would otherwise re-fetch the same customer once per entitlement. Memoize the
+	// customer for the duration of this (read-only) call so the fan-out shares one lookup.
+	// Installed before errgroup.WithContext so the derived ctx carries the memo into every goroutine.
+	ctx = meteredentitlement.WithCustomerCache(ctx)
+
 	var result sync.Map
 
 	g, ctx := errgroup.WithContext(ctx)

openmeter/governance/service/service.go

@@ -419,7 +419,10 @@ func (s *service) resolveAccess(ctx context.Context, input governance.QueryAcces
 
 		for i, rc := range customers {
 			g.Go(func() error {
-				access, err := s.entitlementService.GetAccess(ctx, input.Namespace, rc.customer.ID)
+				// Pass the query's feature filter so GetAccess only resolves the requested
+				// entitlements. When FeatureKeys is empty (all-org path) this is a no-op and all
+				// entitlements are resolved, matching the all-org buildFeatureAccess branch.
+				access, err := s.entitlementService.GetAccess(ctx, input.Namespace, rc.customer.ID, input.FeatureKeys...)
 				if err != nil {
 					return fmt.Errorf("failed to get access for customer %s: %w", rc.customer.ID, err)
 				}

openmeter/server/server_test.go

@@ -1119,7 +1119,7 @@ func (n NoopEntitlementConnector) GetEntitlementOfCustomerAt(ctx context.Context
 	return &entitlement.Entitlement{}, nil
 }
 
-func (n NoopEntitlementConnector) GetAccess(ctx context.Context, namespace string, subjectKey string) (entitlement.Access, error) {
+func (n NoopEntitlementConnector) GetAccess(ctx context.Context, namespace string, subjectKey string, featureKeys ...string) (entitlement.Access, error) {
 	return entitlement.Access{}, nil
 }
 

pkg/syncx/contextmemo.go

@@ -0,0 +1,71 @@
+package syncx
+
+import (
+	"context"
+	"sync"
+)
+
+// ContextMemo is a request-scoped, keyed memoizer.
+//
+// Install it into a context (Install) to enable caching for that context's lifetime; without
+// installation GetOrLoad falls back to a direct call, so callers that do not opt in are
+// unaffected. Each ContextMemo value carries its own private context key (pointer identity), so
+// several memos — even of the same K/V types — coexist in one context without colliding. Entries
+// are computed at most once per key (sync.Once) and shared across concurrent goroutines.
+//
+// Typical use: dedup repeated lookups of the same immutable entity within one request/operation
+// (e.g. the same customer fetched once per entitlement across a fan-out).
+//
+// INVARIANT — install ONLY around read-only operations. Cached values are never invalidated
+// within a scope, so installing a memo around a call stack that MUTATES the cached entity would
+// make reads after the write return the stale, pre-write value.
+type ContextMemo[K comparable, V any] struct {
+	key *contextMemoKey
+}
+
+// contextMemoKey is an unexported context-key type. It carries a field so it is NOT zero-sized:
+// Go does not guarantee distinct addresses for zero-size allocations, so an empty struct would let
+// two memos share a key and collide. A non-zero size makes each &contextMemoKey{} address unique,
+// giving every ContextMemo its own context key even across identical K/V type parameters.
+type contextMemoKey struct{ _ byte }
+
+type contextMemoEntry[V any] struct {
+	once sync.Once
+	val  V
+	err  error
+}
+
+// NewContextMemo returns a ContextMemo with its own private context key. Create one per logical
+// cache (e.g. package-level var) and reuse it; the value is safe for concurrent use.
+func NewContextMemo[K comparable, V any]() *ContextMemo[K, V] {
+	return &ContextMemo[K, V]{key: &contextMemoKey{}}
+}
+
+// Install returns a context with this memo's store enabled. Nested installs are a no-op: the
+// outermost scope wins and its store is shared by everything below it.
+func (m *ContextMemo[K, V]) Install(ctx context.Context) context.Context {
+	if _, ok := ctx.Value(m.key).(*sync.Map); ok {
+		return ctx
+	}
+
+	return context.WithValue(ctx, m.key, &sync.Map{})
+}
+
+// GetOrLoad returns the memoized value for k when a store is installed in ctx; otherwise it calls
+// load directly (no-cache fallback). load runs at most once per key within an installed scope,
+// even under concurrency; both the value and the error are cached for the scope's lifetime.
+func (m *ContextMemo[K, V]) GetOrLoad(ctx context.Context, k K, load func(context.Context) (V, error)) (V, error) {
+	store, ok := ctx.Value(m.key).(*sync.Map)
+	if !ok {
+		return load(ctx)
+	}
+
+	v, _ := store.LoadOrStore(k, &contextMemoEntry[V]{})
+	entry := v.(*contextMemoEntry[V])
+
+	entry.once.Do(func() {
+		entry.val, entry.err = load(ctx)
+	})
+
+	return entry.val, entry.err
+}