refactor: [OM-387] governance entitlement performance improvement

e2e/Makefile

@@ -8,20 +8,48 @@ OPENMETER_ADDRESS ?= http://localhost:38888
 BENCHTIME ?= 20x
 # Number of times to repeat the whole benchmark (feed >1 into benchstat for variance).
 COUNT ?= 1
+# Entitlement kind the seeded features use: boolean (default, cheap algorithmic baseline),
+# metered (forces the ClickHouse balance query in GetAccess — production-representative),
+# mixed (~50/50 boolean+metered), or all (runs all three kinds back to back).
+GOV_BENCH_KIND ?= boolean
 OPENMETER_E2E_EXTRA_COMPOSE_FILES ?=
 
 .PHONY: bench-governance
 bench-governance: ## Benchmark governance query latency (1x1 baseline + 10/50/100 diagonal)
 	$(call print-target)
-	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) \
+	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_KIND=$(GOV_BENCH_KIND) \
 		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./...
 
 .PHONY: bench-governance-matrix
 bench-governance-matrix: ## Benchmark governance query latency (full 3x3 customers x features matrix)
 	$(call print-target)
-	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_FULL_MATRIX=1 \
+	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_FULL_MATRIX=1 GOV_BENCH_KIND=$(GOV_BENCH_KIND) \
 		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./...
 
+.PHONY: bench-governance-metered
+bench-governance-metered: ## Benchmark governance query latency, metered entitlements (diagonal)
+	$(MAKE) bench-governance GOV_BENCH_KIND=metered
+
+.PHONY: bench-governance-matrix-metered
+bench-governance-matrix-metered: ## Benchmark governance query latency, metered entitlements (full 3x3 matrix)
+	$(MAKE) bench-governance-matrix GOV_BENCH_KIND=metered
+
+# --- Trace capture ---
+# Capturing traces for big metered cells (50x50, 100x100) OOMs the local Tempo and the
+# traces are too large to open anyway (>16MiB query cap). This target captures ONE SMALL
+# cell, which shows the SAME per-entitlement query structure (query counts are per
+# entitlement). Low benchtime keeps the number of huge traces ingested down.
+# Override: TRACE_KIND=mixed, BENCH_CELL=customers=1/features=1, TRACE_BENCHTIME=1x.
+TRACE_KIND ?= metered
+BENCH_CELL ?= customers=10/features=10
+TRACE_BENCHTIME ?= 2x
+
+.PHONY: bench-governance-trace
+bench-governance-trace: ## Capture traces for ONE small cell (default metered 10x10, 2x) — safe for local Tempo
+	$(call print-target)
+	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 ./...
+
 .PHONY: test-local
 test-local: ## Run tests against local openmeter
 	$(call print-target)

e2e/governance_bench_test.go

@@ -6,14 +6,52 @@ import (
 	"net/http"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/require"
 
 	api "github.com/openmeterio/openmeter/api/client/go"
 	apiv3 "github.com/openmeterio/openmeter/api/v3"
+	"github.com/openmeterio/openmeter/pkg/convert"
 )
 
+// entKind selects which entitlement type the seeded features use. It is the third
+// benchmark axis (alongside customer and feature counts) and is what makes the
+// benchmark representative of the production cost: boolean entitlements skip the
+// metered balance path (ClickHouse), while metered entitlements force the
+// per-entitlement usage query inside GetAccess that dominates real latency.
+type entKind string
+
+const (
+	// kindBoolean grants every feature a boolean entitlement. GetAccess short-circuits
+	// the balance calculation, so this is the cheap, algorithmic-only baseline.
+	kindBoolean entKind = "boolean"
+	// kindMetered grants every feature a metered entitlement, so GetAccess runs a
+	// ClickHouse usage query per entitlement — the production-representative path.
+	kindMetered entKind = "metered"
+	// kindMixed grants alternating boolean/metered entitlements (~50/50), modeling a
+	// realistic tenant where only some features are usage-metered.
+	kindMixed entKind = "mixed"
+)
+
+// selectedKinds reads GOV_BENCH_KIND and returns the entitlement kinds to benchmark.
+// Default is boolean only, so the existing `make -C e2e bench-governance` behavior
+// (and its baseline numbers) is unchanged. Set GOV_BENCH_KIND=metered|mixed|all to
+// measure the metered balance path.
+func selectedKinds() []entKind {
+	switch os.Getenv("GOV_BENCH_KIND") {
+	case "metered":
+		return []entKind{kindMetered}
+	case "mixed":
+		return []entKind{kindMixed}
+	case "all":
+		return []entKind{kindBoolean, kindMetered, kindMixed}
+	default:
+		return []entKind{kindBoolean}
+	}
+}
+
 // BenchmarkGovernanceQuery measures end-to-end latency of
 // POST /api/v3/openmeter/governance/query against a running stack
 // (OPENMETER_ADDRESS). It is a benchmark, so it only runs under `go test -bench`;
@@ -23,11 +61,27 @@ import (
 // performance work — it exercises the real router, the OAS layer, real Postgres,
 // and the real entitlement GetAccess fan-out over HTTP.
 //
-// What it is NOT: a production-latency oracle. Entitlements seeded here are
-// boolean, so GetAccess skips the metered balance path (ClickHouse). That path is
-// the larger production cost and is instrumented separately (entitlement package).
-// A metered variant — seeding usage events and waiting for ClickHouse ingestion —
-// is the follow-up for measuring it.
+// Entitlement kind (GOV_BENCH_KIND) selects what GetAccess actually does:
+//   - boolean (default): the balance path is short-circuited — algorithmic baseline.
+//   - metered / mixed:   GetAccess runs a ClickHouse usage query per metered
+//     entitlement, the larger production cost. This is the path the deferred
+//     resolveAccess parallelization is meant to speed up, so measuring it here is
+//     what tells us whether that optimization moves the ceiling.
+//
+// Caveat: the metered path is exercised by the per-entitlement usage query inside
+// GetAccess; this fixture does NOT ingest usage events (that would require waiting
+// for ClickHouse ingestion and is flaky/slow). So row volume is minimal — the cost
+// measured is the per-entitlement query overhead × N customers, which is exactly the
+// serial structure parallelization targets, not absolute production row-scan time.
+//
+// The ClickHouse query is issued even with no grants and no events: with zero grants
+// the burndown engine still produces a single full-period phase (credit/engine
+// burnphase.go) and runBetweenResets calls QueryUsage for it (credit/engine run.go).
+// So there is no zero-grant/zero-usage short-circuit that would collapse the metered
+// path back to the boolean one — verified, which is why seeding grants or events is
+// unnecessary to make the metered benchmark meaningful. Ingesting events would only
+// add row-scan volume on top, material only at production data scale this harness is
+// not built to seed.
 //
 // Sizes scale customers x features. The customer-count axis drives the GetAccess
 // fan-out (the dominant cost); the feature-count axis drives the per-customer
@@ -66,55 +120,83 @@ func BenchmarkGovernanceQuery(b *testing.B) {
 		}
 	}
 
-	for _, s := range sizes {
-		b.Run(s.name, func(b *testing.B) {
-			custKeys, featKeys := seedGovernanceFixture(b, client, s.customers, s.features)
+	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)
 
-			reqBody := apiv3.GovernanceQueryRequest{
-				Customer: apiv3.GovernanceQueryRequestCustomers{Keys: custKeys},
-				Feature:  &apiv3.GovernanceQueryRequestFeatures{Keys: featKeys},
-			}
+				reqBody := apiv3.GovernanceQueryRequest{
+					Customer: apiv3.GovernanceQueryRequestCustomers{Keys: custKeys},
+					Feature:  &apiv3.GovernanceQueryRequestFeatures{Keys: featKeys},
+				}
+
+				// 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)
 
-			// 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)
-
-			b.ReportAllocs()
-			b.ResetTimer()
-			for i := 0; i < b.N; i++ {
-				status, _, _ := v3.QueryGovernance(reqBody)
-				if status != http.StatusOK {
-					b.Fatalf("governance query returned %d", status)
+				b.ReportAllocs()
+				b.ResetTimer()
+				for i := 0; i < b.N; i++ {
+					status, _, _ := v3.QueryGovernance(reqBody)
+					if status != http.StatusOK {
+						b.Fatalf("governance query returned %d", status)
+					}
 				}
-			}
-			b.StopTimer()
-		})
+				b.StopTimer()
+			})
+		}
 	}
 }
 
-// seedGovernanceFixture creates nFeatures boolean features and nCustomers
-// customers, granting each customer a boolean entitlement for every feature
-// (nCustomers x nFeatures entitlements). Keys carry a per-run unique prefix so
-// repeated runs against the same database do not collide. Returns the customer
+// seedGovernanceFixture creates nFeatures features and nCustomers customers, granting
+// each customer one entitlement per feature (nCustomers x nFeatures entitlements). The
+// entitlement type is chosen per feature by kind: boolean for kindBoolean, metered for
+// kindMetered, and alternating boolean/metered (~50/50) for kindMixed. Metered features
+// are backed by a single SUM meter created once per run. Keys carry a per-run unique
+// prefix so repeated runs against the same database do not collide. Returns the customer
 // keys and feature keys to query.
-func seedGovernanceFixture(b *testing.B, client *api.ClientWithResponses, nCustomers, nFeatures int) (custKeys, featKeys []string) {
+func seedGovernanceFixture(b *testing.B, client *api.ClientWithResponses, nCustomers, nFeatures int, kind entKind) (custKeys, featKeys []string) {
 	b.Helper()
 	ctx := b.Context()
 	run := uniqueKey("gov_bench")
 
+	// isMetered reports whether feature i is metered, given the kind.
+	isMetered := func(i int) bool {
+		switch kind {
+		case kindMetered:
+			return true
+		case kindMixed:
+			return i%2 == 1
+		default: // kindBoolean
+			return false
+		}
+	}
+
+	// Any metered feature needs a meter to point at; create one SUM meter for the run.
+	var meterSlug string
+	for i := 0; i < nFeatures; i++ {
+		if isMetered(i) {
+			meterSlug = createGovernanceMeter(b, client, ctx, run)
+			break
+		}
+	}
+
 	featKeys = make([]string, 0, nFeatures)
+	featMetered := make([]bool, 0, nFeatures)
 	for i := 0; i < nFeatures; i++ {
 		fkey := fmt.Sprintf("%s_feat_%d", run, i)
-		resp, err := client.CreateFeatureWithResponse(ctx, api.CreateFeatureJSONRequestBody{
-			Key:  fkey,
-			Name: fkey,
-		})
+		body := api.CreateFeatureJSONRequestBody{Key: fkey, Name: fkey}
+		if isMetered(i) {
+			body.MeterSlug = convert.ToPointer(meterSlug)
+		}
+		resp, err := client.CreateFeatureWithResponse(ctx, body)
 		require.NoError(b, err)
 		require.Equalf(b, http.StatusCreated, resp.StatusCode(), "create feature: %s", resp.Body)
 		featKeys = append(featKeys, fkey)
+		featMetered = append(featMetered, isMetered(i))
 	}
 
 	custKeys = make([]string, 0, nCustomers)
@@ -135,15 +217,40 @@ func seedGovernanceFixture(b *testing.B, client *api.ClientWithResponses, nCusto
 		require.Equalf(b, http.StatusCreated, custResp.StatusCode(), "create customer: %s", custResp.Body)
 		custID := custResp.JSON201.Id
 
-		for _, fkey := range featKeys {
-			grantBooleanEntitlement(b, client, ctx, custID, fkey)
+		for i, fkey := range featKeys {
+			if featMetered[i] {
+				grantMeteredEntitlement(b, client, ctx, custID, fkey)
+			} else {
+				grantBooleanEntitlement(b, client, ctx, custID, fkey)
+			}
 		}
 		custKeys = append(custKeys, ckey)
 	}
 
 	return custKeys, featKeys
 }
 
+// createGovernanceMeter creates a single SUM meter for the benchmark run and returns its
+// slug. All metered features in the run share this meter; the meter exists only so
+// metered entitlements resolve a usage query — no events are ingested against it.
+func createGovernanceMeter(b *testing.B, client *api.ClientWithResponses, ctx context.Context, run string) string {
+	b.Helper()
+
+	slug := run + "_meter"
+	resp, err := client.CreateMeterWithResponse(ctx, api.MeterCreate{
+		Slug:          slug,
+		Name:          convert.ToPointer(slug),
+		Aggregation:   api.MeterAggregationSum,
+		EventType:     run + "_event",
+		ValueProperty: convert.ToPointer("$.value"),
+	})
+	require.NoError(b, err)
+	// The meter API returns 200 (not 201) on create.
+	require.Equalf(b, http.StatusOK, resp.StatusCode(), "create meter: %s", resp.Body)
+
+	return slug
+}
+
 // grantBooleanEntitlement creates a boolean entitlement for the given customer and
 // feature key via the V2 customer-entitlement endpoint.
 func grantBooleanEntitlement(b *testing.B, client *api.ClientWithResponses, ctx context.Context, custID, featureKey string) {
@@ -159,3 +266,30 @@ func grantBooleanEntitlement(b *testing.B, client *api.ClientWithResponses, ctx
 	require.NoError(b, err)
 	require.Equalf(b, http.StatusCreated, resp.StatusCode(), "create boolean entitlement: %s", resp.Body)
 }
+
+// grantMeteredEntitlement creates a metered entitlement for the given customer and
+// feature key via the V2 customer-entitlement endpoint. The metered type is what forces
+// GetAccess to run a balance/usage query (ClickHouse) when governance resolves access.
+func grantMeteredEntitlement(b *testing.B, client *api.ClientWithResponses, ctx context.Context, custID, featureKey string) {
+	b.Helper()
+
+	month := &api.RecurringPeriodInterval{}
+	require.NoError(b, month.FromRecurringPeriodIntervalEnum(api.RecurringPeriodIntervalEnumMONTH))
+
+	var body api.CreateCustomerEntitlementV2JSONRequestBody
+	require.NoError(b, body.FromEntitlementMeteredV2CreateInputs(api.EntitlementMeteredV2CreateInputs{
+		Type:       "metered",
+		FeatureKey: lo.ToPtr(featureKey),
+		UsagePeriod: api.RecurringPeriodCreateInput{
+			// Anchor at seed time so only the current usage period is evaluated. A fixed past
+			// anchor would let the balance engine accrue a monthly reset period per elapsed
+			// month, drifting benchmark cost over time and breaking cross-run comparisons.
+			Anchor:   convert.ToPointer(time.Now().UTC()),
+			Interval: *month,
+		},
+	}))
+
+	resp, err := client.CreateCustomerEntitlementV2WithResponse(ctx, custID, body)
+	require.NoError(b, err)
+	require.Equalf(b, http.StatusCreated, resp.StatusCode(), "create metered entitlement: %s", resp.Body)
+}

openmeter/credit/helper.go

@@ -177,41 +177,51 @@ func (m *connector) snapshotEngineResult(ctx context.Context, snapParams snapsho
 	ctx, span := m.Tracer.Start(ctx, "credit.snapshotEngineResult", cTrace.WithOwner(snapParams.owner))
 	defer span.End()
 
-	if err := transaction.RunWithNoValue(ctx, m.GrantRepo, func(ctx context.Context) error {
-		return m.OwnerConnector.LockOwnerForTx(ctx, snapParams.owner, false)
-	}); err != nil {
-		// If we failed to acquire the lock we simply don't save the snapshot
-		return nil
-	}
-
 	// Skip snapshotting for LATEST type entitlements as the values fluctuate and snapshots can't be used
 	if snapParams.meter.Aggregation == meter.MeterAggregationLatest {
 		m.Logger.Debug("skipping snapshot for LATEST aggregation type entitlement", "owner", snapParams.owner, "meter", snapParams.meter.Key)
 		return nil
 	}
 
+	// Decide which segment (if any) to snapshot BEFORE acquiring the owner lock. This scan is
+	// pure (in-memory over the engine history), so when nothing qualifies we can skip the
+	// advisory-lock transaction entirely. That matters on the read path (GetBalanceAt runs this
+	// once per metered entitlement): without a qualifying segment the lock transaction would
+	// otherwise hold a pool connection for no persisted result.
 	segs := runRes.History.Segments()
 
+	snapshotIdx := -1
+
 	// i >= 1 because:
 	// The first segment starts with the last valid snapshot and we don't want to create another snapshot for that same time
 	for i := len(segs) - 1; i >= 1; i-- {
-		seg := segs[i]
-
 		// We can save a segment if its not after the current period start (this way backfilling, granting, resetting, etc... will work for the current UsagePeriod)
-		if !seg.From.After(snapParams.notAfter) {
-			snap, err := runRes.History.GetSnapshotAtStartOfSegment(i)
-			if err != nil {
-				return fmt.Errorf("failed to get snapshot at start of segment: %w", err)
-			}
-
-			if _, err := m.saveSnapshot(ctx, snapParams, snap); err != nil {
-				return fmt.Errorf("failed to save snapshot: %w", err)
-			}
-
+		if !segs[i].From.After(snapParams.notAfter) {
+			snapshotIdx = i
 			break
 		}
 	}
 
+	if snapshotIdx == -1 {
+		return nil
+	}
+
+	if err := transaction.RunWithNoValue(ctx, m.GrantRepo, func(ctx context.Context) error {
+		return m.OwnerConnector.LockOwnerForTx(ctx, snapParams.owner, false)
+	}); err != nil {
+		// If we failed to acquire the lock we simply don't save the snapshot
+		return nil
+	}
+
+	snap, err := runRes.History.GetSnapshotAtStartOfSegment(snapshotIdx)
+	if err != nil {
+		return fmt.Errorf("failed to get snapshot at start of segment: %w", err)
+	}
+
+	if _, err := m.saveSnapshot(ctx, snapParams, snap); err != nil {
+		return fmt.Errorf("failed to save snapshot: %w", err)
+	}
+
 	return nil
 }