diff --git a/e2e/Makefile b/e2e/Makefile index 09b94b15ac..26de6256b8 100644 --- a/e2e/Makefile +++ b/e2e/Makefile @@ -22,6 +22,12 @@ bench-governance-matrix: ## Benchmark governance query latency (full 3x3 custome TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_FULL_MATRIX=1 \ go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./... +.PHONY: bench-governance-namespace +bench-governance-namespace: ## Benchmark governance query latency vs. namespace size (fixed 100 customers x 1 feature, 0/10k/100k decoys); requires direct Postgres access + $(call print-target) + TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) \ + go test -run='^$$' -bench='^BenchmarkGovernanceQueryNamespaceScale$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./... + .PHONY: test-local test-local: ## Run tests against local openmeter $(call print-target) diff --git a/e2e/governance_bench_namespace_test.go b/e2e/governance_bench_namespace_test.go new file mode 100644 index 0000000000..a743a81db2 --- /dev/null +++ b/e2e/governance_bench_namespace_test.go @@ -0,0 +1,173 @@ +package e2e + +import ( + "fmt" + "net/http" + "os" + "slices" + "sort" + "strconv" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/oklog/ulid/v2" + "github.com/stretchr/testify/require" + + v3sdk "github.com/openmeterio/openmeter/api/v3/client" +) + +const ( + defaultGovernanceBenchNamespaceDecoys = 100_000 + governanceBenchNamespaceDecoysEnv = "GOV_BENCH_NAMESPACE_DECOYS" +) + +// BenchmarkGovernanceQueryNamespaceScale isolates the effect of total namespace +// size on governance query latency, independent of the customers/features axes +// BenchmarkGovernanceQuery already covers. It fixes a small, constant query +// load (100 customers, matching the OAS customer-key cap, x 1 feature) so the +// entitlement GetAccess fan-out cost stays flat across sub-benchmarks, and +// varies only how many other (never-queried) customers/subjects exist in the +// namespace. +// +// This targets the customer usage-attribution resolution path: pre-UNION-ALL, +// a large decoy count made GetCustomersByUsageAttribution seq-scan the +// customers table (see #4684 and the follow-up bulk fix); this benchmark shows +// whether decoy count still moves total request latency post-fix. Decoy +// counts accumulate across sub-benchmarks (0 -> 10k -> N), so each step only +// seeds the incremental delta. +// +// Decoys are seeded directly via SQL (not HTTP, which would dominate setup +// time) and require direct Postgres access alongside OPENMETER_ADDRESS (see +// initE2EPostgresPool). Set GOV_BENCH_NAMESPACE_DECOYS to change the top +// decoy count from the default 100,000. +func BenchmarkGovernanceQueryNamespaceScale(b *testing.B) { + client := initClient(b) + v3 := newV3Client(b) + pool := initE2EPostgresPool(b) + + const ( + queryCustomers = 100 + queryFeatures = 1 + ) + + custKeys, featKeys := seedGovernanceFixture(b, client, queryCustomers, queryFeatures) + namespace := getCustomerNamespaceByKey(b, pool, custKeys[0]) + decoyRun := uniqueKey("gov_bench_ns_decoy") + + // Sorted and deduplicated so the cumulative seeding below stays monotonically + // increasing regardless of GOV_BENCH_NAMESPACE_DECOYS: an override below the + // hardcoded 10_000 tier would otherwise seed 10_000 first, then skip seeding + // for the smaller tier (seeded > target) while still reporting that smaller, + // now-incorrect decoy count as the benchmark's label. + decoyCounts := []int{0, 10_000, governanceBenchNamespaceDecoyCount(b)} + sort.Ints(decoyCounts) + decoyCounts = slices.Compact(decoyCounts) + + reqBody := v3sdk.GovernanceQueryRequest{ + Customer: v3sdk.GovernanceQueryRequestCustomers{Keys: custKeys}, + Feature: &v3sdk.GovernanceQueryRequestFeatures{Keys: featKeys}, + } + + seeded := 0 + for _, target := range decoyCounts { + b.Run(fmt.Sprintf("decoys=%d", target), func(b *testing.B) { + if target > seeded { + seedNamespaceDecoys(b, pool, namespace, decoyRun, seeded, target) + seeded = target + } + + // Warm-up + correctness gate: a wrong result (e.g. missing customers) + // would make the latency number meaningless. + resp, err := v3.Governance.QueryAccess(b.Context(), reqBody, v3sdk.GovernanceQueryResultListParams{}) + v3.requireStatus(http.StatusOK, err) + require.Lenf(b, resp.Data, queryCustomers, "expected %d resolved customers", queryCustomers) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := v3.Governance.QueryAccess(b.Context(), reqBody, v3sdk.GovernanceQueryResultListParams{}); err != nil { + b.Fatalf("governance query failed: %v", err) + } + if s := v3.statuses.last(); s != http.StatusOK { + b.Fatalf("governance query returned %d", s) + } + } + b.StopTimer() + b.ReportMetric(float64(target), "decoys") + }) + } +} + +func governanceBenchNamespaceDecoyCount(b *testing.B) int { + b.Helper() + + value := os.Getenv(governanceBenchNamespaceDecoysEnv) + if value == "" { + return defaultGovernanceBenchNamespaceDecoys + } + + count, err := strconv.Atoi(value) + if err != nil || count < 1 { + b.Fatalf("%s must be a positive integer, got %q", governanceBenchNamespaceDecoysEnv, value) + } + + return count +} + +func getCustomerNamespaceByKey(tb testing.TB, pool *pgxpool.Pool, key string) string { + tb.Helper() + + var namespace string + err := pool.QueryRow( + tb.Context(), + `SELECT namespace FROM customers WHERE key = $1`, + key, + ).Scan(&namespace) + require.NoError(tb, err) + + return namespace +} + +// seedNamespaceDecoys bulk-inserts decoy customers (and one subject key each) +// directly via SQL, numbered (from, to] under the given per-benchmark-run +// prefix, so repeated calls with a growing `to` only seed the incremental +// delta, and separate benchmark runs never collide on the same key even if +// the database was not reset in between. Decoys are never included in a +// governance query; they exist purely to grow the namespace the resolution +// query has to search. ANALYZE keeps planner stats current, matching the +// customer adapter's own usage-attribution benchmark. +func seedNamespaceDecoys(tb testing.TB, pool *pgxpool.Pool, namespace, runPrefix string, from, to int) { + tb.Helper() + + n := to - from + ids := make([]string, n) + keys := make([]string, n) + subjectKeys := make([]string, n) + + for i := range n { + ids[i] = ulid.Make().String() + keys[i] = fmt.Sprintf("%s_%d", runPrefix, from+i) + subjectKeys[i] = keys[i] + "_subj" + } + + rctx := tb.Context() + + _, err := pool.Exec(rctx, ` + INSERT INTO customers (id, namespace, created_at, updated_at, name, key) + SELECT id, $1, now(), now(), key, key + FROM unnest($2::text[], $3::text[]) AS t(id, key) + `, namespace, ids, keys) + require.NoError(tb, err) + + _, err = pool.Exec(rctx, ` + INSERT INTO customer_subjects (namespace, subject_key, created_at, customer_id) + SELECT $1, subject_key, now(), customer_id + FROM unnest($2::text[], $3::text[]) AS t(subject_key, customer_id) + `, namespace, subjectKeys, ids) + require.NoError(tb, err) + + _, err = pool.Exec(rctx, "ANALYZE customers") + require.NoError(tb, err) + _, err = pool.Exec(rctx, "ANALYZE customer_subjects") + require.NoError(tb, err) +} diff --git a/e2e/helpers.go b/e2e/helpers.go index 7940807114..d1dcf6e048 100644 --- a/e2e/helpers.go +++ b/e2e/helpers.go @@ -4,14 +4,44 @@ import ( "context" "fmt" "net/http" + "os" + "strings" "testing" + "github.com/jackc/pgx/v5/pgxpool" "github.com/samber/lo" "github.com/stretchr/testify/require" api "github.com/openmeterio/openmeter/api/client/go" ) +// initE2EPostgresPool opens a direct Postgres connection alongside OPENMETER_ADDRESS, for +// e2e checks and fixtures that need to read or seed data the public API doesn't expose +// (ledger account mappings, bulk namespace seeding). Requires OPENMETER_E2E_POSTGRES_URL, +// or the local compose stack's default exposed port. +func initE2EPostgresPool(t testing.TB) *pgxpool.Pool { + t.Helper() + + dsn := os.Getenv("OPENMETER_E2E_POSTGRES_URL") + if dsn == "" { + address := os.Getenv("OPENMETER_ADDRESS") + if !strings.Contains(address, "localhost:38888") && !strings.Contains(address, "127.0.0.1:38888") { + t.Skipf("this e2e check requires OPENMETER_E2E_POSTGRES_URL or local compose stack at localhost:38888, got %q", address) + } + + dsn = "postgres://postgres:postgres@127.0.0.1:35432/postgres?sslmode=disable" + } + + pool, err := pgxpool.New(t.Context(), dsn) + require.NoError(t, err) + + t.Cleanup(pool.Close) + + require.NoError(t, pool.Ping(t.Context())) + + return pool +} + // This will not be needed once we get rid of subjects. Cloud middleware already handles subject / customer creation. func CreateCustomerWithSubject(t *testing.T, client *api.ClientWithResponses, customerKey string, subjectKey string) *api.Customer { t.Helper() diff --git a/e2e/ledger_accounts_test.go b/e2e/ledger_accounts_test.go index 24ce7267fc..da5e370391 100644 --- a/e2e/ledger_accounts_test.go +++ b/e2e/ledger_accounts_test.go @@ -1,7 +1,6 @@ package e2e import ( - "os" "strings" "testing" @@ -38,29 +37,6 @@ func TestNewCustomerHasLedgerAccounts(t *testing.T) { } } -func initE2EPostgresPool(t *testing.T) *pgxpool.Pool { - t.Helper() - - dsn := os.Getenv("OPENMETER_E2E_POSTGRES_URL") - if dsn == "" { - address := os.Getenv("OPENMETER_ADDRESS") - if !strings.Contains(address, "localhost:38888") && !strings.Contains(address, "127.0.0.1:38888") { - t.Skipf("ledger account e2e requires OPENMETER_E2E_POSTGRES_URL or local compose stack at localhost:38888, got %q", address) - } - - dsn = "postgres://postgres:postgres@127.0.0.1:35432/postgres?sslmode=disable" - } - - pool, err := pgxpool.New(t.Context(), dsn) - require.NoError(t, err) - - t.Cleanup(pool.Close) - - require.NoError(t, pool.Ping(t.Context())) - - return pool -} - func getCustomerNamespace(t *testing.T, pool *pgxpool.Pool, customerID string) string { t.Helper() diff --git a/openmeter/customer/adapter.go b/openmeter/customer/adapter.go index e234655801..ad55f36763 100644 --- a/openmeter/customer/adapter.go +++ b/openmeter/customer/adapter.go @@ -20,7 +20,6 @@ type CustomerAdapter interface { CreateCustomer(ctx context.Context, params CreateCustomerInput) (*Customer, error) DeleteCustomer(ctx context.Context, customer DeleteCustomerInput) error GetCustomer(ctx context.Context, customer GetCustomerInput) (*Customer, error) - GetCustomerByUsageAttribution(ctx context.Context, input GetCustomerByUsageAttributionInput) (*Customer, error) GetCustomersByUsageAttribution(ctx context.Context, input GetCustomersByUsageAttributionInput) ([]Customer, error) UpdateCustomer(ctx context.Context, params UpdateCustomerInput) (*Customer, error) } diff --git a/openmeter/customer/adapter/customer.go b/openmeter/customer/adapter/customer.go index 865eb6416b..a20a589e72 100644 --- a/openmeter/customer/adapter/customer.go +++ b/openmeter/customer/adapter/customer.go @@ -500,90 +500,50 @@ func (a *adapter) GetCustomer(ctx context.Context, input customer.GetCustomerInp }) } -// GetCustomerByUsageAttribution gets a customer by usage attribution -func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input customer.GetCustomerByUsageAttributionInput) (*customer.Customer, error) { - if err := input.Validate(); err != nil { - return nil, models.NewGenericValidationError( - fmt.Errorf("error getting customer by usage attribution: %w", err), - ) - } - - return entutils.TransactingRepo(ctx, a, func(ctx context.Context, repo *adapter) (*customer.Customer, error) { - now := clock.Now().UTC() - - query := repo.db.Customer.Query(). - Where( - customerdb.Namespace(input.Namespace), - customerdb.DeletedAtIsNil(), - customerMatchesUsageAttributionKey(input.Namespace, input.Key, now), - ) - query = WithSubjects(query, now) - if slices.Contains(input.Expands, customer.ExpandSubscriptions) { - query = WithActiveSubscriptions(query, now) - } - - customerEntity, err := query.First(ctx) - if err != nil { - if entdb.IsNotFound(err) { - return nil, models.NewGenericNotFoundError( - fmt.Errorf("customer with subject key %s not found in %s namespace", input.Key, input.Namespace), - ) - } - - return nil, fmt.Errorf("failed to fetch customer: %w", err) - } - - if customerEntity == nil { - return nil, fmt.Errorf("invalid query result: nil customer received") - } - - return CustomerFromDBEntity(*customerEntity, input.Expands) - }) -} - -// customerMatchesUsageAttributionKey resolves a customer key before a subject -// key while keeping both lookup branches independently indexable. The returned -// candidate is applied to the outer customer query as: +// customersMatchUsageAttributionKeys resolves customers whose own key OR one of their subject keys +// is in the given key set, via two independently-indexable UNION ALL branches. It returns the full +// candidate set rather than one winner per key: a key that matches both a customer's own key and a +// different customer's subject key legitimately resolves to two distinct customers here, so the +// caller applies key-over-subject precedence per input key (see +// customer/service.resolveCustomersByKeyWithPrecedence). Both the single-key and bulk usage-attribution lookups +// use this one predicate. The generated SQL shape mirrors: // // WHERE customers.id IN ( -// SELECT matches.id -// FROM ( -// SELECT c.id, 0 AS lookup_priority -// FROM customers AS c -// WHERE c.namespace = $1 -// AND c.key = $2 -// AND c.deleted_at IS NULL +// SELECT c.id FROM customers AS c +// WHERE c.namespace = $1 +// AND c.key IN (...) +// AND (c.deleted_at IS NULL OR c.deleted_at > $2) // -// UNION ALL +// UNION ALL // -// SELECT c.id, 1 AS lookup_priority -// FROM customer_subjects AS cs -// JOIN customers AS c ON c.id = cs.customer_id -// WHERE cs.namespace = $1 -// AND cs.subject_key = $2 -// AND (cs.deleted_at IS NULL OR cs.deleted_at > $3) -// AND c.namespace = $1 -// AND c.deleted_at IS NULL -// ) AS matches -// ORDER BY matches.lookup_priority -// LIMIT 1 +// SELECT c.id +// FROM customer_subjects AS cs +// JOIN customers AS c ON c.id = cs.customer_id +// WHERE cs.namespace = $1 +// AND cs.subject_key IN (...) +// AND (cs.deleted_at IS NULL OR cs.deleted_at > $2) +// AND c.namespace = $1 +// AND (c.deleted_at IS NULL OR c.deleted_at > $2) // ) -func customerMatchesUsageAttributionKey(namespace, key string, at time.Time) predicate.Customer { +func customersMatchUsageAttributionKeys(namespace string, keys []string, at time.Time) predicate.Customer { return func(s *sql.Selector) { + keyValues := lo.ToAnySlice(keys) + keyCustomerTable := sql.Table(customerdb.Table).As("customer_by_key") customerKeyMatch := sql.Select(keyCustomerTable.C(customerdb.FieldID)). - AppendSelectExprAs(sql.Expr("0"), "lookup_priority"). From(keyCustomerTable). Where(sql.And( sql.EQ(keyCustomerTable.C(customerdb.FieldNamespace), namespace), - sql.EQ(keyCustomerTable.C(customerdb.FieldKey), key), - sql.IsNull(keyCustomerTable.C(customerdb.FieldDeletedAt)), + sql.In(keyCustomerTable.C(customerdb.FieldKey), keyValues...), + sql.Or( + sql.IsNull(keyCustomerTable.C(customerdb.FieldDeletedAt)), + sql.GT(keyCustomerTable.C(customerdb.FieldDeletedAt), at), + ), )) customerSubjectsTable := sql.Table(customersubjectsdb.Table).As("customer_subjects") subjectCustomerTable := sql.Table(customerdb.Table).As("customer_by_subject") subjectKeyMatch := sql.Select(subjectCustomerTable.C(customerdb.FieldID)). - AppendSelectExprAs(sql.Expr("1"), "lookup_priority"). From(customerSubjectsTable). Join(subjectCustomerTable). On( @@ -592,58 +552,44 @@ func customerMatchesUsageAttributionKey(namespace, key string, at time.Time) pre ). Where(sql.And( sql.EQ(customerSubjectsTable.C(customersubjectsdb.FieldNamespace), namespace), - sql.EQ(customerSubjectsTable.C(customersubjectsdb.FieldSubjectKey), key), + sql.In(customerSubjectsTable.C(customersubjectsdb.FieldSubjectKey), keyValues...), sql.Or( sql.IsNull(customerSubjectsTable.C(customersubjectsdb.FieldDeletedAt)), sql.GT(customerSubjectsTable.C(customersubjectsdb.FieldDeletedAt), at), ), sql.EQ(subjectCustomerTable.C(customerdb.FieldNamespace), namespace), - sql.IsNull(subjectCustomerTable.C(customerdb.FieldDeletedAt)), + sql.Or( + sql.IsNull(subjectCustomerTable.C(customerdb.FieldDeletedAt)), + sql.GT(subjectCustomerTable.C(customerdb.FieldDeletedAt), at), + ), )) - candidates := customerKeyMatch. - UnionAll(subjectKeyMatch). - As("usage_attribution_matches") - preferredCustomerID := sql.Select(candidates.C(customerdb.FieldID)). - From(candidates). - OrderBy(candidates.C("lookup_priority")). - Limit(1) + candidates := customerKeyMatch.UnionAll(subjectKeyMatch) - s.Where(sql.In(s.C(customerdb.FieldID), preferredCustomerID)) + s.Where(sql.In(s.C(customerdb.FieldID), candidates)) } } -// GetCustomersByUsageAttribution resolves multiple customers by usage attribution keys in a single query. -// A key matches a customer either by the customer's own key or by one of its subject keys, mirroring -// the single-key GetCustomerByUsageAttribution. Keys that match no customer are simply absent from the -// result; the caller derives which keys were not found. +// GetCustomersByUsageAttribution resolves customers by usage attribution keys in a single query. +// A key matches a customer either by the customer's own key or by one of its subject keys. It is the +// sole usage-attribution lookup for both the bulk and single-key service paths; keys that match no +// customer are simply absent from the result, and the caller derives which keys were not found. func (a *adapter) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) ([]customer.Customer, error) { return entutils.TransactingRepo(ctx, a, func(ctx context.Context, repo *adapter) ([]customer.Customer, error) { now := clock.Now().UTC() - // TODO: consider adding a CreatedAtLTE(now) guard to the customer/subject query, - // to defend against the edge case of having customers/subjects becoming active in the future + // Deliberately no CreatedAtLTE(now) guard here or in WithSubjects below: this keeps the + // candidate predicate and the eager-loaded subject list in agreement. created_at is + // server-assigned and immutable (see ResourceMixin/CustomerSubjects schema), so a + // future-created subject can't occur in production — guarding only one of the two paths + // would instead cause a real, batch-composition-dependent resolution bug for a case that + // can't happen. query := repo.db.Customer.Query(). - Where(customerdb.Namespace(input.Namespace)). Where( - customerdb.Or( - // We lookup customers by subject key in the subjects table - customerdb.HasSubjectsWith( - customersubjectsdb.SubjectKeyIn(input.Keys...), - customersubjectsdb.CreatedAtLTE(now), - customersubjectsdb.Or( - customersubjectsdb.DeletedAtIsNil(), - customersubjectsdb.DeletedAtGT(now), - ), - ), - // Or else we lookup customers by key in the customers table - customerdb.KeyIn(input.Keys...), - ), - ). - Where(customerdb.DeletedAtIsNil()) + customerdb.Namespace(input.Namespace), + customersMatchUsageAttributionKeys(input.Namespace, input.Keys, now), + ) - // TODO: consider adding a CreatedAtLTE(now) guard to the WithSubjects query, - // to defend against the edge case of having customers/subjects becoming active in the future query = WithSubjects(query, now) if slices.Contains(input.Expands, customer.ExpandSubscriptions) { diff --git a/openmeter/customer/adapter/customer_test.go b/openmeter/customer/adapter/customer_test.go index 97fae3b099..015839fe93 100644 --- a/openmeter/customer/adapter/customer_test.go +++ b/openmeter/customer/adapter/customer_test.go @@ -389,194 +389,6 @@ func customerIDs(customers []customer.Customer) []string { }) } -func TestGetCustomerByUsageAttribution(t *testing.T) { - t.Run("MatchByCustomerKey", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1") - - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "cust-key", - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, id, got.ID) - require.NotNil(t, got.UsageAttribution) - assert.Equal(t, []string{"subj-1"}, got.UsageAttribution.SubjectKeys) - }) - - t.Run("MatchBySubjectKey", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1", "subj-2") - - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "subj-2", - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, id, got.ID) - require.NotNil(t, got.UsageAttribution) - assert.Equal(t, []string{"subj-1", "subj-2"}, got.UsageAttribution.SubjectKeys) - }) - - t.Run("CustomerKeyTakesPriorityOverAnotherCustomerSubjectKey", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - customerKeyID := env.seedCustomerWithKey(ns, "shared-key", "customer-key-subject") - _ = env.seedCustomerWithKey(ns, "subject-owner", "shared-key") - - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "shared-key", - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, customerKeyID, got.ID, "a direct customer-key match must take priority over a subject-key match") - }) - - t.Run("CustomerMatchedByOwnKeyAndSubjectKeyReturnedOnce", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "shared-key", "shared-key") - - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "shared-key", - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, id, got.ID) - }) - - t.Run("SoftDeletedCustomerExcluded", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1") - _ = freezeTime(t, time.Now()) - - require.NoError(t, env.adapter.DeleteCustomer(t.Context(), customer.DeleteCustomerInput{ - Namespace: ns, - ID: id, - })) - - for _, key := range []string{"cust-key", "subj-1"} { - _, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: key, - }) - require.Error(t, err) - assert.True(t, models.IsGenericNotFoundError(err)) - } - }) - - t.Run("SubjectDeletedInPastExcluded", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1") - now := freezeTime(t, time.Now()) - - _, err := env.db.CustomerSubjects.Update(). - Where( - customersubjectsdb.Namespace(ns), - customersubjectsdb.CustomerID(id), - customersubjectsdb.SubjectKey("subj-1"), - ). - SetDeletedAt(now.Add(-time.Minute)). - Save(t.Context()) - require.NoError(t, err) - - _, err = env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "subj-1", - }) - require.Error(t, err) - assert.True(t, models.IsGenericNotFoundError(err)) - }) - - t.Run("SubjectDeletedAtLookupTimeExcluded", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1") - now := freezeTime(t, time.Now()) - - _, err := env.db.CustomerSubjects.Update(). - Where( - customersubjectsdb.Namespace(ns), - customersubjectsdb.CustomerID(id), - customersubjectsdb.SubjectKey("subj-1"), - ). - SetDeletedAt(now). - Save(t.Context()) - require.NoError(t, err) - - _, err = env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "subj-1", - }) - require.Error(t, err) - assert.True(t, models.IsGenericNotFoundError(err)) - }) - - t.Run("SubjectDeletedInFutureIncluded", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - id := env.seedCustomerWithKey(ns, "cust-key", "subj-1") - now := freezeTime(t, time.Now()) - - _, err := env.db.CustomerSubjects.Update(). - Where( - customersubjectsdb.Namespace(ns), - customersubjectsdb.CustomerID(id), - customersubjectsdb.SubjectKey("subj-1"), - ). - SetDeletedAt(now.Add(time.Minute)). - Save(t.Context()) - require.NoError(t, err) - - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "subj-1", - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, id, got.ID) - }) - - t.Run("NamespaceIsolation", func(t *testing.T) { - env := newTestEnv(t) - nsA := ulid.Make().String() - nsB := ulid.Make().String() - idA := env.seedCustomerWithKey(nsA, "shared-customer-key", "shared-subject-key") - _ = env.seedCustomerWithKey(nsB, "shared-customer-key", "shared-subject-key") - - for _, key := range []string{"shared-customer-key", "shared-subject-key"} { - got, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: nsA, - Key: key, - }) - require.NoError(t, err) - require.NotNil(t, got) - assert.Equal(t, idA, got.ID) - } - }) - - t.Run("UnknownKeyReturnsNotFound", func(t *testing.T) { - env := newTestEnv(t) - ns := ulid.Make().String() - _ = env.seedCustomerWithKey(ns, "cust-key", "subj-1") - - _, err := env.adapter.GetCustomerByUsageAttribution(t.Context(), customer.GetCustomerByUsageAttributionInput{ - Namespace: ns, - Key: "missing", - }) - require.Error(t, err) - assert.True(t, models.IsGenericNotFoundError(err)) - }) -} - func TestGetCustomersByUsageAttribution(t *testing.T) { t.Run("MatchByCustomerKey", func(t *testing.T) { env := newTestEnv(t) @@ -664,6 +476,35 @@ func TestGetCustomersByUsageAttribution(t *testing.T) { assert.Empty(t, customerIDs(got), "soft-deleted customer must not be returned") }) + t.Run("CustomerDeletedInFutureIncluded", func(t *testing.T) { + // given: + // - a customer whose own deleted_at is future-dated, set directly via Ent; the bulk analog + // of the single-key CustomerDeletedInFutureIncluded. + // then: + // - the customer is returned once for a key set covering both its own key and its subject key, + // because the deleted_at grace window now applies to the owning customer on both branches. + env := newTestEnv(t) + ns := ulid.Make().String() + id := env.seedCustomerWithKey(ns, "key-a", "subj-a") + now := freezeTime(t, time.Now()) + + _, err := env.db.Customer.Update(). + Where( + customerdb.Namespace(ns), + customerdb.ID(id), + ). + SetDeletedAt(now.Add(time.Minute)). + Save(t.Context()) + require.NoError(t, err) + + got, err := env.adapter.GetCustomersByUsageAttribution(t.Context(), customer.GetCustomersByUsageAttributionInput{ + Namespace: ns, + Keys: []string{"key-a", "subj-a"}, + }) + require.NoError(t, err) + assert.ElementsMatch(t, []string{id}, customerIDs(got)) + }) + t.Run("SoftDeletedSubjectExcludedButCustomerKeyStillMatches", func(t *testing.T) { env := newTestEnv(t) ns := ulid.Make().String() diff --git a/openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go b/openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go index aeaa8f7d9a..7a10a20a9a 100644 --- a/openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go +++ b/openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go @@ -3,36 +3,30 @@ package adapter import ( "context" "fmt" - "os" - "strconv" "testing" "time" - entdb "github.com/openmeterio/openmeter/openmeter/ent/db" + "github.com/openmeterio/openmeter/openmeter/customer/testutils/uabench" customerdb "github.com/openmeterio/openmeter/openmeter/ent/db/customer" customersubjectsdb "github.com/openmeterio/openmeter/openmeter/ent/db/customersubjects" "github.com/openmeterio/openmeter/openmeter/testutils" ) -const ( - defaultUsageAttributionBenchmarkCustomers = 100_000 - usageAttributionBenchmarkSizeEnv = "CUSTOMER_USAGE_ATTRIBUTION_BENCH_CUSTOMERS" -) - -// BenchmarkCustomerUsageAttributionLookup compares the cross-table OR with the -// UNION ALL candidate-ID lookup. Seeding is outside the timed sections. Set -// CUSTOMER_USAGE_ATTRIBUTION_BENCH_CUSTOMERS to change the namespace size from -// the default 100,000 customers. -func BenchmarkCustomerUsageAttributionLookup(b *testing.B) { +// BenchmarkCustomerUsageAttributionQuery compares the cross-table OR with the UNION ALL candidate-ID +// query for a single key. It measures only the adapter-level candidate query (no precedence, no +// hydration); the full service path incl. key-over-subject precedence is benchmarked separately in +// the customer service package. Seeding is outside the timed sections. Set +// CUSTOMER_USAGE_ATTRIBUTION_BENCH_CUSTOMERS to change the namespace size from the default 100,000. +func BenchmarkCustomerUsageAttributionQuery(b *testing.B) { b.StopTimer() db := testutils.InitPostgresDB(b, testutils.PostgresDBStateEntMigrated) b.Cleanup(func() { db.Close(b) }) client := db.EntDriver.Client() - customerCount := usageAttributionBenchmarkCustomerCount(b) + customerCount := uabench.CustomerCount(b) namespace := "customer-usage-attribution-benchmark" - seedUsageAttributionBenchmark(b, client, namespace, customerCount) + uabench.Seed(b, client, namespace, customerCount) targetCustomerID := fmt.Sprintf("%026d", customerCount) now := time.Now().UTC() @@ -84,8 +78,7 @@ func BenchmarkCustomerUsageAttributionLookup(b *testing.B) { return client.Customer.Query(). Where( customerdb.Namespace(namespace), - customerdb.DeletedAtIsNil(), - customerMatchesUsageAttributionKey(namespace, lookup.key, now), + customersMatchUsageAttributionKeys(namespace, []string{lookup.key}, now), ). FirstID(ctx) }, @@ -122,57 +115,94 @@ func BenchmarkCustomerUsageAttributionLookup(b *testing.B) { } } -func usageAttributionBenchmarkCustomerCount(b *testing.B) int { - b.Helper() +// BenchmarkCustomersUsageAttributionBulkQuery compares the cross-table OR with the UNION ALL +// candidate-ID query for the bulk (multi-key) path, mirroring BenchmarkCustomerUsageAttributionQuery +// for the single-key path. Adapter-level query only (no precedence). Seeding is outside the timed +// sections. Set CUSTOMER_USAGE_ATTRIBUTION_BENCH_CUSTOMERS to change the namespace size. +func BenchmarkCustomersUsageAttributionBulkQuery(b *testing.B) { + b.StopTimer() - value := os.Getenv(usageAttributionBenchmarkSizeEnv) - if value == "" { - return defaultUsageAttributionBenchmarkCustomers - } + db := testutils.InitPostgresDB(b, testutils.PostgresDBStateEntMigrated) + b.Cleanup(func() { db.Close(b) }) - count, err := strconv.Atoi(value) - if err != nil || count < 1 { - b.Fatalf("%s must be a positive integer, got %q", usageAttributionBenchmarkSizeEnv, value) + client := db.EntDriver.Client() + + customerCount := uabench.CustomerCount(b) + if customerCount < 2*uabench.BulkKeyCount { + b.Fatalf("%s must be at least %d to seed a non-overlapping bulk key set, got %d", + uabench.CustomersEnv, 2*uabench.BulkKeyCount, customerCount) } - return count -} + namespace := "customer-usage-attribution-bulk-benchmark" + uabench.Seed(b, client, namespace, customerCount) -func seedUsageAttributionBenchmark(b *testing.B, client *entdb.Client, namespace string, customerCount int) { - b.Helper() - - _, err := client.ExecContext(b.Context(), ` - INSERT INTO customers (id, namespace, created_at, updated_at, name, key) - SELECT - lpad(customer_number::text, 26, '0'), - $1, - now(), - now(), - 'customer-' || customer_number, - 'customer-' || customer_number - FROM generate_series(1, $2) AS customer_number - `, namespace, customerCount) - if err != nil { - b.Fatalf("seed customers: %v", err) - } + now := time.Now().UTC() + keys := uabench.BulkKeys(customerCount) - _, err = client.ExecContext(b.Context(), ` - INSERT INTO customer_subjects (namespace, subject_key, created_at, customer_id) - SELECT - $1, - 'subject-' || customer_number, - now(), - lpad(customer_number::text, 26, '0') - FROM generate_series(1, $2) AS customer_number - `, namespace, customerCount) - if err != nil { - b.Fatalf("seed customer subjects: %v", err) + variants := []struct { + name string + query func(context.Context) (int, error) + }{ + { + name: "cross_table_or", + query: func(ctx context.Context) (int, error) { + return client.Customer.Query(). + Where( + customerdb.Namespace(namespace), + customerdb.Or( + customerdb.HasSubjectsWith( + customersubjectsdb.SubjectKeyIn(keys...), + customersubjectsdb.Or( + customersubjectsdb.DeletedAtIsNil(), + customersubjectsdb.DeletedAtGT(now), + ), + ), + customerdb.KeyIn(keys...), + ), + customerdb.DeletedAtIsNil(), + ). + Count(ctx) + }, + }, + { + name: "union_all", + query: func(ctx context.Context) (int, error) { + return client.Customer.Query(). + Where( + customerdb.Namespace(namespace), + customerdb.DeletedAtIsNil(), + customersMatchUsageAttributionKeys(namespace, keys, now), + ). + Count(ctx) + }, + }, } - if _, err = client.ExecContext(b.Context(), "ANALYZE customers"); err != nil { - b.Fatalf("analyze customers: %v", err) - } - if _, err = client.ExecContext(b.Context(), "ANALYZE customer_subjects"); err != nil { - b.Fatalf("analyze customer subjects: %v", err) + for _, variant := range variants { + b.Run(variant.name, func(b *testing.B) { + count, err := variant.query(b.Context()) + if err != nil { + b.Fatalf("warm up lookup: %v", err) + } + if count != uabench.BulkKeyCount { + b.Fatalf("warm up lookup returned %d customers, expected %d", count, uabench.BulkKeyCount) + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + count, err = variant.query(b.Context()) + if err != nil { + b.Fatalf("lookup customers: %v", err) + } + } + b.StopTimer() + b.ReportMetric(float64(customerCount), "customers") + b.ReportMetric(float64(uabench.BulkKeyCount), "keys") + + if count != uabench.BulkKeyCount { + b.Fatalf("lookup returned %d customers, expected %d", count, uabench.BulkKeyCount) + } + }) } } diff --git a/openmeter/customer/service.go b/openmeter/customer/service.go index f39fad415d..63991071b1 100644 --- a/openmeter/customer/service.go +++ b/openmeter/customer/service.go @@ -26,6 +26,6 @@ type CustomerService interface { DeleteCustomer(ctx context.Context, customer DeleteCustomerInput) error GetCustomer(ctx context.Context, customer GetCustomerInput) (*Customer, error) GetCustomerByUsageAttribution(ctx context.Context, input GetCustomerByUsageAttributionInput) (*Customer, error) - GetCustomersByUsageAttribution(ctx context.Context, input GetCustomersByUsageAttributionInput) ([]Customer, error) + GetCustomersByUsageAttribution(ctx context.Context, input GetCustomersByUsageAttributionInput) (map[string]*Customer, error) UpdateCustomer(ctx context.Context, params UpdateCustomerInput) (*Customer, error) } diff --git a/openmeter/customer/service/customer.go b/openmeter/customer/service/customer.go index 9c0d1ebf4c..cf400f15bb 100644 --- a/openmeter/customer/service/customer.go +++ b/openmeter/customer/service/customer.go @@ -131,20 +131,103 @@ func (s *Service) GetCustomer(ctx context.Context, input customer.GetCustomerInp return s.adapter.GetCustomer(ctx, input) } -// GetCustomerByUsageAttribution gets a customer by usage attribution +// GetCustomerByUsageAttribution gets a customer by usage attribution. It resolves the single key +// through the same bulk candidate query and key-over-subject precedence as +// GetCustomersByUsageAttribution, so both paths share one predicate and one precedence mechanism. func (s *Service) GetCustomerByUsageAttribution(ctx context.Context, input customer.GetCustomerByUsageAttributionInput) (*customer.Customer, error) { - return s.adapter.GetCustomerByUsageAttribution(ctx, input) + if err := input.Validate(); err != nil { + return nil, models.NewGenericValidationError( + fmt.Errorf("error getting customer by usage attribution: %w", err), + ) + } + + resolved, err := s.resolveCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: input.Namespace, + Keys: []string{input.Key}, + Expands: input.Expands, + }) + if err != nil { + return nil, err + } + + c := resolved[input.Key] + if c == nil { + return nil, models.NewGenericNotFoundError( + fmt.Errorf("customer with subject key %s not found in %s namespace", input.Key, input.Namespace), + ) + } + + return c, nil } -// GetCustomersByUsageAttribution resolves multiple customers by usage attribution keys in a single query -func (s *Service) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) ([]customer.Customer, error) { +// GetCustomersByUsageAttribution resolves multiple customers by usage attribution keys in a single +// query, mapping each input key to the customer it matches with key-over-subject precedence applied. +// Every input key is present in the returned map; keys that match no customer map to a nil value. +func (s *Service) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) { if err := input.Validate(); err != nil { return nil, models.NewGenericValidationError( fmt.Errorf("error getting customers by usage attribution: %w", err), ) } - return s.adapter.GetCustomersByUsageAttribution(ctx, input) + return s.resolveCustomersByUsageAttribution(ctx, input) +} + +// resolveCustomersByUsageAttribution runs the shared usage-attribution lookup: fetch the raw +// candidate customers from the adapter, then apply key-over-subject precedence. Both the single-key +// and bulk service methods route through this so their core resolution stays identical. +func (s *Service) resolveCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) { + customers, err := s.adapter.GetCustomersByUsageAttribution(ctx, input) + if err != nil { + return nil, err + } + + return resolveCustomersByKeyWithPrecedence(customers, input.Keys), nil +} + +// resolveCustomersByKeyWithPrecedence maps each input key to the customer it matches. A key matches a customer +// either by the customer's own key or by one of its subject keys; when a key matches both a +// distinct key-owner and a distinct subject-owner, the key-owner takes precedence. Key-over-subject +// collisions are structurally rare and resolved deterministically here, so they are not surfaced as +// errors or logs; investigate the underlying data via the database if ever needed. Every input key +// is present in the returned map; keys that match no customer have a nil value. +func resolveCustomersByKeyWithPrecedence(customers []customer.Customer, keys []string) map[string]*customer.Customer { + byKey := make(map[string]*customer.Customer, len(customers)) + bySubject := make(map[string]*customer.Customer, len(customers)) + + for i := range customers { + c := &customers[i] + + if c.Key != nil { + byKey[*c.Key] = c + } + + if c.UsageAttribution != nil { + for _, sk := range c.UsageAttribution.SubjectKeys { + if _, ok := bySubject[sk]; !ok { + bySubject[sk] = c + } + } + } + } + + resolved := make(map[string]*customer.Customer, len(keys)) + + for _, k := range keys { + if keyOwner, ok := byKey[k]; ok { + resolved[k] = keyOwner + continue + } + + if subjectOwner, ok := bySubject[k]; ok { + resolved[k] = subjectOwner + continue + } + + resolved[k] = nil // present, but no matching customer + } + + return resolved } // UpdateCustomer updates a customer diff --git a/openmeter/customer/service/customer_test.go b/openmeter/customer/service/customer_test.go new file mode 100644 index 0000000000..a2b6703120 --- /dev/null +++ b/openmeter/customer/service/customer_test.go @@ -0,0 +1,104 @@ +package customerservice + +import ( + "testing" + + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + + "github.com/openmeterio/openmeter/openmeter/customer" + "github.com/openmeterio/openmeter/pkg/models" +) + +func Test_resolveCustomersByKeyWithPrecedence(t *testing.T) { + customerA := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-a"}, + Key: lo.ToPtr("key-a"), + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: []string{"subject-a"}, + }, + } + + t.Run("MatchByKey", func(t *testing.T) { + resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"key-a"}) + + assert.Equal(t, map[string]*customer.Customer{"key-a": &customerA}, resolved) + }) + + t.Run("MatchBySubject", func(t *testing.T) { + resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"subject-a"}) + + assert.Equal(t, map[string]*customer.Customer{"subject-a": &customerA}, resolved) + }) + + t.Run("UnmatchedKeyIsNil", func(t *testing.T) { + resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"no-such-key"}) + + assert.Equal(t, map[string]*customer.Customer{"no-such-key": nil}, resolved) + }) + + t.Run("KeyOwnerTakesPrecedenceOverDistinctSubjectOwner", func(t *testing.T) { + // "shared" is customerA's own key AND customerB's subject key: the key-owner (A) must win. + sharedKeyCustomer := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-a"}, + Key: lo.ToPtr("shared"), + } + sharedSubjectCustomer := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-b"}, + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: []string{"shared"}, + }, + } + + resolved := resolveCustomersByKeyWithPrecedence( + []customer.Customer{sharedKeyCustomer, sharedSubjectCustomer}, + []string{"shared"}, + ) + + assert.Equal(t, map[string]*customer.Customer{"shared": &sharedKeyCustomer}, resolved) + }) + + t.Run("SameCustomerMatchedByOwnKeyAndSubjectKey", func(t *testing.T) { + selfMatched := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-a"}, + Key: lo.ToPtr("dual"), + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: []string{"dual"}, + }, + } + + resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{selfMatched}, []string{"dual"}) + + assert.Equal(t, map[string]*customer.Customer{"dual": &selfMatched}, resolved) + }) + + t.Run("CrossCollisionResolvesEachKeyToItsOwnKeyOwner", func(t *testing.T) { + // customerA's key is customerB's subject key and vice versa: no single ordering of + // [A, B] can satisfy both keys under a naive single-pass first-match-wins map, but + // resolving key-owner and subject-owner independently per key does. + crossA := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-a"}, + Key: lo.ToPtr("key-1"), + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: []string{"key-2"}, + }, + } + crossB := customer.Customer{ + ManagedResource: models.ManagedResource{ID: "customer-b"}, + Key: lo.ToPtr("key-2"), + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: []string{"key-1"}, + }, + } + + resolved := resolveCustomersByKeyWithPrecedence( + []customer.Customer{crossA, crossB}, + []string{"key-1", "key-2"}, + ) + + assert.Equal(t, map[string]*customer.Customer{ + "key-1": &crossA, + "key-2": &crossB, + }, resolved) + }) +} diff --git a/openmeter/customer/service/customer_usage_attribution_benchmark_test.go b/openmeter/customer/service/customer_usage_attribution_benchmark_test.go new file mode 100644 index 0000000000..25d17e6bfd --- /dev/null +++ b/openmeter/customer/service/customer_usage_attribution_benchmark_test.go @@ -0,0 +1,169 @@ +package customerservice_test + +import ( + "context" + "fmt" + "testing" + + "github.com/openmeterio/openmeter/openmeter/customer" + customeradapter "github.com/openmeterio/openmeter/openmeter/customer/adapter" + customerservice "github.com/openmeterio/openmeter/openmeter/customer/service" + "github.com/openmeterio/openmeter/openmeter/customer/testutils/uabench" + "github.com/openmeterio/openmeter/openmeter/testutils" + noop "github.com/openmeterio/openmeter/openmeter/watermill/driver/noop" + "github.com/openmeterio/openmeter/openmeter/watermill/eventbus" +) + +// newUsageAttributionServiceBenchmark builds a real customer service over a freshly seeded namespace +// (100k customers by default) and returns it with the namespace and seeded size. It constructs the +// service directly (adapter + service, no-op event publisher) rather than via the *testing.T-only +// NewTestEnv, so it is usable from a *testing.B. +func newUsageAttributionServiceBenchmark(b *testing.B) (customer.Service, string, int) { + b.Helper() + + db := testutils.InitPostgresDB(b, testutils.PostgresDBStateEntMigrated) + b.Cleanup(func() { db.Close(b) }) + + client := db.EntDriver.Client() + + logger := testutils.NewDiscardLogger(b) + + publisher, err := eventbus.New(eventbus.Options{ + Publisher: &noop.Publisher{}, + TopicMapping: eventbus.TopicMapping{ + IngestEventsTopic: "bench-ingest-events", + SystemEventsTopic: "bench-system-events", + BalanceWorkerEventsTopic: "bench-balance-worker-events", + }, + Logger: logger, + }) + if err != nil { + b.Fatalf("build event publisher: %v", err) + } + + adapter, err := customeradapter.New(customeradapter.Config{Client: client, Logger: logger}) + if err != nil { + b.Fatalf("build customer adapter: %v", err) + } + + svc, err := customerservice.New(customerservice.Config{ + Adapter: adapter, + Publisher: publisher, + }) + if err != nil { + b.Fatalf("build customer service: %v", err) + } + + customerCount := uabench.CustomerCount(b) + namespace := "customer-usage-attribution-service-benchmark" + uabench.Seed(b, client, namespace, customerCount) + + return svc, namespace, customerCount +} + +// BenchmarkCustomerUsageAttributionLookup measures the full single-key service path — candidate +// query, subject hydration, and the real key-over-subject precedence in +// resolveCustomersByKeyWithPrecedence — via the public service method. The adapter-level query cost +// alone is benchmarked by BenchmarkCustomerUsageAttributionQuery in the adapter package. +func BenchmarkCustomerUsageAttributionLookup(b *testing.B) { + b.StopTimer() + + svc, namespace, customerCount := newUsageAttributionServiceBenchmark(b) + targetCustomerID := fmt.Sprintf("%026d", customerCount) + + lookups := []struct { + name string + key string + }{ + {name: "customer_key", key: fmt.Sprintf("customer-%d", customerCount)}, + {name: "subject_key", key: fmt.Sprintf("subject-%d", customerCount)}, + } + + for _, lookup := range lookups { + b.Run(lookup.name, func(b *testing.B) { + get := func(ctx context.Context) (string, error) { + c, err := svc.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{ + Namespace: namespace, + Key: lookup.key, + }) + if err != nil { + return "", err + } + return c.ID, nil + } + + customerID, err := get(b.Context()) + if err != nil { + b.Fatalf("warm up lookup: %v", err) + } + if customerID != targetCustomerID { + b.Fatalf("warm up lookup returned customer %q, expected %q", customerID, targetCustomerID) + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := get(b.Context()); err != nil { + b.Fatalf("lookup customer: %v", err) + } + } + b.StopTimer() + b.ReportMetric(float64(customerCount), "customers") + }) + } +} + +// BenchmarkCustomersUsageAttributionBulkLookup measures the full bulk service path — one candidate +// query for the whole key set plus per-key precedence resolution into the returned map — via the +// public service method. +func BenchmarkCustomersUsageAttributionBulkLookup(b *testing.B) { + svc, namespace, customerCount := newUsageAttributionServiceBenchmark(b) + if customerCount < 2*uabench.BulkKeyCount { + b.Fatalf("%s must be at least %d to build a non-overlapping bulk key set, got %d", + uabench.CustomersEnv, 2*uabench.BulkKeyCount, customerCount) + } + + keys := uabench.BulkKeys(customerCount) + + // Count non-nil entries: every input key is present in the result map, with a nil value for + // keys that resolved to no customer. + get := func(ctx context.Context) (int, error) { + resolved, err := svc.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: namespace, + Keys: keys, + }) + if err != nil { + return 0, err + } + matched := 0 + for _, c := range resolved { + if c != nil { + matched++ + } + } + return matched, nil + } + + count, err := get(b.Context()) + if err != nil { + b.Fatalf("warm up lookup: %v", err) + } + if count != uabench.BulkKeyCount { + b.Fatalf("warm up lookup resolved %d customers, expected %d", count, uabench.BulkKeyCount) + } + + // The timed loop runs inside b.Run so it gets a fresh, running timer: setup above (which seeds + // 100k rows) is excluded from the measurement without relying on StopTimer/ResetTimer ordering. + b.Run(fmt.Sprintf("keys_%d", uabench.BulkKeyCount), func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := get(b.Context()); err != nil { + b.Fatalf("lookup customers: %v", err) + } + } + b.StopTimer() + b.ReportMetric(float64(customerCount), "customers") + b.ReportMetric(float64(uabench.BulkKeyCount), "keys") + }) +} diff --git a/openmeter/customer/service/service_test.go b/openmeter/customer/service/service_test.go index 2fe58607fc..e6ce871541 100644 --- a/openmeter/customer/service/service_test.go +++ b/openmeter/customer/service/service_test.go @@ -19,8 +19,6 @@ func Test_CustomerService(t *testing.T) { env.Close(t) }) - // Run database migrations - // Get new namespace ID namespace := customertestutils.NewTestNamespace(t) @@ -220,24 +218,171 @@ func Test_CustomerService(t *testing.T) { }) } -func Test_GetCustomersByUsageAttribution_Validation(t *testing.T) { +// Test_GetCustomersByUsageAttribution protects the bulk method's contract with the same rigor as +// Test_GetCustomerByUsageAttribution does the single-key one: each input key maps to its customer +// with key-over-subject precedence, and unmatched keys are present with a nil value (no error). The predicate +// itself is covered by the adapter's TestGetCustomersByUsageAttribution, and the fully symmetric +// {K1,K2} cross-collision by Test_resolveCustomersByKeyWithPrecedence — that state cannot be built +// through the service here, because CreateCustomer rejects a customer whose key overlaps another +// customer's subject key (see the overlap guard in the customer adapter). +func Test_GetCustomersByUsageAttribution(t *testing.T) { env := customertestutils.NewTestEnv(t) t.Cleanup(func() { env.Close(t) }) namespace := customertestutils.NewTestNamespace(t) + ctx := t.Context() + + create := func(t *testing.T, key string, subjectKeys ...string) *customer.Customer { + t.Helper() + cus, err := env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{ + Namespace: namespace, + CustomerMutate: customer.CustomerMutate{ + Key: lo.ToPtr(key), + Name: key, + UsageAttribution: &customer.CustomerUsageAttribution{SubjectKeys: subjectKeys}, + }, + }) + require.NoError(t, err) + return cus + } + + t.Run("ResolvesByKeyAndSubject", func(t *testing.T) { + // given: + // - customer A matched by its own key, customer B matched by one of its subject keys + // then: + // - each input key maps to the correct distinct customer + a := create(t, "a-key", "a-subject") + b := create(t, "b-key", "b-subject") + + got, err := env.CustomerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: namespace, + Keys: []string{"a-key", "b-subject"}, + }) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, a.ID, got["a-key"].ID) + assert.Equal(t, b.ID, got["b-subject"].ID) + }) + + t.Run("KeyOverSubjectPrecedence", func(t *testing.T) { + // given: + // - "shared" is customer P's own key AND customer Q's subject key + // then: + // - the key owner (P) wins, resolved in the service + p := create(t, "shared", "p-subject") + _ = create(t, "q-key", "shared") + + got, err := env.CustomerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: namespace, + Keys: []string{"shared"}, + }) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, p.ID, got["shared"].ID, "a direct customer-key match must win over a subject-key match") + }) + + t.Run("UnmatchedKeysNil", func(t *testing.T) { + // given: + // - a mix of a known key and an unknown one + // then: + // - every input key is present in the map; the unknown key has a nil value (unlike the + // single-key method, the bulk method does NOT return a not-found error) + known := create(t, "known-key") + + got, err := env.CustomerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: namespace, + Keys: []string{"known-key", "totally-unknown"}, + }) + require.NoError(t, err) + require.Len(t, got, 2) + require.NotNil(t, got["known-key"]) + assert.Equal(t, known.ID, got["known-key"].ID) + unknown, ok := got["totally-unknown"] + assert.True(t, ok, "every input key must be present in the result map") + assert.Nil(t, unknown, "unmatched key must have a nil value") + }) + + t.Run("EmptyKeysFailsValidation", func(t *testing.T) { + // given: + // - an empty key set + // then: + // - the service rejects it with a validation error before any DB query is built + _, err := env.CustomerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + Namespace: namespace, + Keys: nil, + }) + require.Error(t, err, "empty key set must fail validation") + assert.True(t, models.IsGenericValidationError(err), "error must be a validation error") + }) +} - // given: - // - an empty key set - // when: - // - resolving customers by usage attribution - // then: - // - the service rejects it with a validation error before any DB query is built - _, err := env.CustomerService.GetCustomersByUsageAttribution(t.Context(), customer.GetCustomersByUsageAttributionInput{ - Namespace: namespace, - Keys: nil, +// Test_GetCustomerByUsageAttribution covers the single-key wiring that Test_CustomerService's +// happy-path lookups don't reach: the key-over-subject precedence now resolved in the service (not +// in SQL), the not-found mapping, and input validation. The predicate itself is covered by the +// adapter's TestGetCustomersByUsageAttribution. +func Test_GetCustomerByUsageAttribution(t *testing.T) { + env := customertestutils.NewTestEnv(t) + t.Cleanup(func() { + env.Close(t) + }) + + namespace := customertestutils.NewTestNamespace(t) + ctx := t.Context() + + t.Run("KeyOverSubjectPrecedence", func(t *testing.T) { + // given: + // - customer A owns the key "shared-key" + // - customer B carries "shared-key" as one of its subject keys + // when: + // - resolving "shared-key" by usage attribution + // then: + // - the direct key owner (A) wins over the subject-key match (B), resolved in the service + keyOwner, err := env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{ + Namespace: namespace, + CustomerMutate: customer.CustomerMutate{ + Key: lo.ToPtr("shared-key"), + Name: "Key Owner", + UsageAttribution: &customer.CustomerUsageAttribution{SubjectKeys: []string{"key-owner-subject"}}, + }, + }) + require.NoError(t, err) + + _, err = env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{ + Namespace: namespace, + CustomerMutate: customer.CustomerMutate{ + Key: lo.ToPtr("subject-owner"), + Name: "Subject Owner", + UsageAttribution: &customer.CustomerUsageAttribution{SubjectKeys: []string{"shared-key"}}, + }, + }) + require.NoError(t, err) + + got, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{ + Namespace: namespace, + Key: "shared-key", + }) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, keyOwner.ID, got.ID, "a direct customer-key match must win over a subject-key match") + }) + + t.Run("UnknownKeyReturnsNotFound", func(t *testing.T) { + _, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{ + Namespace: namespace, + Key: "no-such-key", + }) + require.Error(t, err) + assert.True(t, models.IsGenericNotFoundError(err), "unmatched key must map to a not-found error") + }) + + t.Run("EmptyKeyFailsValidation", func(t *testing.T) { + _, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{ + Namespace: namespace, + Key: "", + }) + require.Error(t, err) + assert.True(t, models.IsGenericValidationError(err), "empty key must fail validation") }) - require.Error(t, err, "empty key set must fail validation") - assert.True(t, models.IsGenericValidationError(err), "error must be a validation error") } diff --git a/openmeter/customer/testutils/uabench/uabench.go b/openmeter/customer/testutils/uabench/uabench.go new file mode 100644 index 0000000000..c551fcf052 --- /dev/null +++ b/openmeter/customer/testutils/uabench/uabench.go @@ -0,0 +1,105 @@ +// Package uabench holds shared fixtures for the customer usage-attribution benchmarks. It imports +// only the generated ent client (not the customer adapter or service packages), so both the +// adapter-layer benchmark (package adapter) and the service-layer benchmark can seed from one place +// without an import cycle. +package uabench + +import ( + "fmt" + "os" + "strconv" + "testing" + + entdb "github.com/openmeterio/openmeter/openmeter/ent/db" +) + +const ( + // CustomersEnv overrides the seeded namespace size. + CustomersEnv = "CUSTOMER_USAGE_ATTRIBUTION_BENCH_CUSTOMERS" + + defaultCustomers = 100_000 + + // BulkKeyCount mirrors the governance QueryAccess OAS cap (@maxItems(100) on customer.keys), + // the primary bulk consumer's realistic upper bound. + BulkKeyCount = 100 +) + +// CustomerCount returns the seeded namespace size, overridable via CustomersEnv (default 100,000). +func CustomerCount(tb testing.TB) int { + tb.Helper() + + value := os.Getenv(CustomersEnv) + if value == "" { + return defaultCustomers + } + + count, err := strconv.Atoi(value) + if err != nil || count < 1 { + tb.Fatalf("%s must be a positive integer, got %q", CustomersEnv, value) + } + + return count +} + +// BulkKeys builds a realistic mixed bulk key set of BulkKeyCount keys: half resolve via a direct +// customer-key match, half via a subject-key match on a disjoint customer range, mirroring how +// governance resolves customer keys and subject keys through the same call. Requires +// customerCount >= 2*BulkKeyCount. +func BulkKeys(customerCount int) []string { + half := BulkKeyCount / 2 + + keys := make([]string, 0, BulkKeyCount) + for i := 1; i <= half; i++ { + keys = append(keys, fmt.Sprintf("customer-%d", i)) + } + for i := half + 1; i <= 2*half; i++ { + keys = append(keys, fmt.Sprintf("subject-%d", i)) + } + + return keys +} + +// Seed bulk-inserts customerCount customers (key "customer-N") each with one subject ("subject-N") +// directly via SQL — orders of magnitude faster than the service create path — then ANALYZEs both +// tables so the planner has stats. Customer ids are the zero-padded ordinal, so callers can compute +// the target id for row N as lpad(N, 26, '0'). +func Seed(tb testing.TB, client *entdb.Client, namespace string, customerCount int) { + tb.Helper() + + ctx := tb.Context() + + _, err := client.ExecContext(ctx, ` + INSERT INTO customers (id, namespace, created_at, updated_at, name, key) + SELECT + lpad(customer_number::text, 26, '0'), + $1, + now(), + now(), + 'customer-' || customer_number, + 'customer-' || customer_number + FROM generate_series(1, $2) AS customer_number + `, namespace, customerCount) + if err != nil { + tb.Fatalf("seed customers: %v", err) + } + + _, err = client.ExecContext(ctx, ` + INSERT INTO customer_subjects (namespace, subject_key, created_at, customer_id) + SELECT + $1, + 'subject-' || customer_number, + now(), + lpad(customer_number::text, 26, '0') + FROM generate_series(1, $2) AS customer_number + `, namespace, customerCount) + if err != nil { + tb.Fatalf("seed customer subjects: %v", err) + } + + if _, err = client.ExecContext(ctx, "ANALYZE customers"); err != nil { + tb.Fatalf("analyze customers: %v", err) + } + if _, err = client.ExecContext(ctx, "ANALYZE customer_subjects"); err != nil { + tb.Fatalf("analyze customer subjects: %v", err) + } +} diff --git a/openmeter/governance/service/service.go b/openmeter/governance/service/service.go index 209e3868c8..ca5c7d5501 100644 --- a/openmeter/governance/service/service.go +++ b/openmeter/governance/service/service.go @@ -287,7 +287,10 @@ func (s *service) resolveCustomers(ctx context.Context, input governance.QueryAc attribute.Int("requested", len(input.CustomerKeys)), ) - customers, err := s.customerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ + // GetCustomersByUsageAttribution already resolves each key to its customer with + // key-over-subject precedence applied, so there is no local mapping/collision handling + // left to do here. + keyToCustomer, err := s.customerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{ Namespace: input.Namespace, Keys: input.CustomerKeys, }) @@ -295,37 +298,14 @@ func (s *service) resolveCustomers(ctx context.Context, input governance.QueryAc return resolveCustomersResult{}, fmt.Errorf("failed to resolve customer keys: %w", err) } - // Map each lookup key to the customer it resolved to. A key matches a customer either by - // the customer's own key or by one of its subject keys. First match wins on the rare - // collision, mirroring the single-key First() lookup. - keyToCustomer := make(map[string]*customer.Customer, len(customers)) - - for i := range customers { - cus := &customers[i] - - if cus.Key != nil { - if _, ok := keyToCustomer[*cus.Key]; !ok { - keyToCustomer[*cus.Key] = cus - } - } - - if cus.UsageAttribution != nil { - for _, sk := range cus.UsageAttribution.SubjectKeys { - if _, ok := keyToCustomer[sk]; !ok { - keyToCustomer[sk] = cus - } - } - } - } - customerMap := make(map[string]*resolvedCustomer) var queryErrors []governance.QueryError // Iterate the input keys in order so matched keys and not-found errors keep input ordering. for _, key := range input.CustomerKeys { - cus, ok := keyToCustomer[key] + cus := keyToCustomer[key] - if !ok { + if cus == nil { queryErrors = append(queryErrors, governance.QueryError{ CustomerKey: key, Code: governance.QueryErrorCustomerNotFound, diff --git a/openmeter/server/server_test.go b/openmeter/server/server_test.go index 2438a81961..77ea96527e 100644 --- a/openmeter/server/server_test.go +++ b/openmeter/server/server_test.go @@ -1458,7 +1458,7 @@ func (n NoopCustomerService) GetCustomerByUsageAttribution(ctx context.Context, return &customer.Customer{}, nil } -func (n NoopCustomerService) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) ([]customer.Customer, error) { +func (n NoopCustomerService) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) { return nil, nil }