feat: bulk customer usage attribution lookup performance and correctness#4686
Merged
gergely-kurucz-konghq merged 13 commits intoJul 23, 2026
Merged
Conversation
gergely-kurucz-konghq
added
release-note/misc
Contributor
📝 WalkthroughWalkthroughUsage-attribution resolution now returns key-to-customer maps with key-over-subject precedence. Governance consumes the resolved map directly, while adapter, service, and E2E benchmarks measure lookup and namespace-scaling behavior. ChangesUsage attribution resolution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant GovernanceService
participant CustomerService
participant CustomerAdapter
participant Postgres
GovernanceService->>CustomerService: Resolve usage-attribution keys
CustomerService->>CustomerAdapter: Fetch matching customers
CustomerAdapter->>Postgres: Query customer and subject-key candidates
Postgres-->>CustomerAdapter: Return candidate customer data
CustomerAdapter-->>CustomerService: Return candidates
CustomerService->>CustomerService: Apply key-over-subject precedence
CustomerService-->>GovernanceService: Return key-to-customer map
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@e2e/governance_bench_namespace_test.go`:
- Line 57: Validate the custom value returned by
governanceBenchNamespaceDecoyCount before constructing decoyCounts, requiring it
to be at least 10,000 (and fail clearly if not). Update the setup around the
decoy-count benchmark cases so decoyCounts remains sorted and the incremental
seeding logic reports accurate metrics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 8a2301f5-8bfa-48bc-af01-381cec6cb90a
📒 Files selected for processing (27)
app/common/customer.goe2e/Makefilee2e/governance_bench_namespace_test.goe2e/helpers.goe2e/ledger_accounts_test.goopenmeter/customer/adapter/customer.goopenmeter/customer/adapter/customer_usage_attribution_benchmark_test.goopenmeter/customer/service.goopenmeter/customer/service/customer.goopenmeter/customer/service/customer_test.goopenmeter/customer/service/hooks/subjectcustomer_test.goopenmeter/customer/service/service.goopenmeter/customer/testutils/env.goopenmeter/entitlement/metered/lateevents_test.goopenmeter/entitlement/metered/utils_test.goopenmeter/entitlement/service/utils_test.goopenmeter/governance/service/service.goopenmeter/governance/service/service_test.goopenmeter/server/server_test.goopenmeter/subject/service/hooks/customersubject_test.goopenmeter/subject/testutils/env.goopenmeter/subscription/testutils/customer.gotest/app/stripe/testenv.gotest/app/testenv.gotest/billing/suite.gotest/customer/testenv.gotest/entitlement/regression/framework_test.go
💤 Files with no reviewable changes (1)
- e2e/ledger_accounts_test.go
Contributor
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
GAlexIHU
reviewed
Jul 16, 2026
GetCustomersByUsageAttribution had the same cross-table OR anti-pattern PR #4684 fixed for the single-key lookup, causing a Postgres seq scan on large namespaces. Split into two independently-indexable UNION ALL branches, same as the single-key fix, without lookup_priority (bulk still returns every match). Benchmark: 222ms -> 1ms (220x) at 100 keys / 100k customers.
…ally GetCustomersByUsageAttribution returned an unordered []Customer, so callers (governance + a private-repo consumer) each built their own map[key]Customer with first-match-wins semantics that depended on DB row order — unfixable by reordering when two customers collide on the same key from opposite sides (one's own key vs. the other's subject key). Service now returns map[string]Customer with key-over-subject precedence resolved once, centrally, mirroring the single-key lookup's UNION ALL priority. Each ambiguous collision is logged with the key and both customer IDs for operator follow-up, replacing the duplicated (and now provably buggy) mapping loops in every consumer.
Isolates namespace size from the customers/features axes BenchmarkGovernanceQuery already covers: fixes query load at 100 customers x 1 feature, varies decoy customers/subjects seeded directly via SQL (HTTP creation would dominate setup time at 100k+). Confirmed against the pre-fix bulk lookup: latency scales with namespace size (38ms -> 160ms at 0 -> 100k decoys); post-fix it stays flat (48ms -> 54ms), matching the adapter-level benchmark's win. Moves initE2EPostgresPool into helpers.go so both this and the existing ledger e2e check can share it.
… size decoyCounts hardcoded 10_000 as the middle tier; GOV_BENCH_NAMESPACE_DECOYS below that seeded up to 10_000 first, then skipped seeding for the smaller tier (seeded > target) while still reporting the smaller, now-incorrect count as the benchmark label. Sort + dedup keeps the sequence monotonic regardless of the env override.
…ibution lookup Bulk subject-key matching added a CreatedAtLTE guard the single-key path never had, diverging from its own eager-loaded subject list and forcing a second query to reconcile them. created_at is server-assigned and immutable, so a future-created subject can't occur in production. Remove the guard to restore parity with the single-key predicate.
Trim redundant clauses in two comments (governance namespace-scale benchmark doc, bulk usage-attribution created_at-guard note) that restated the same point twice or over-qualified it. No behavior change.
…bution lookups Give the subject-branch owning-customer check the same Or(IsNull, GT(at)) grace window as the key branch in both GetCustomerByUsageAttribution and GetCustomersByUsageAttribution, so a customer resolves identically via its own key or a subject key. Behavior-equivalent in production (deletion always stamps clock.Now().UTC()); adds CustomerDeletedInFutureIncluded tests.
…resolve precedence in service Delete the single-key adapter GetCustomerByUsageAttribution and its SQL-precedence helper (lookup_priority / ORDER BY / LIMIT). Both single-key and bulk lookups now use the one customersMatchUsageAttributionKeys predicate; the single-key service method calls the bulk adapter with a one-element key set and applies key-over-subject precedence in Go via resolveCustomersByKey. The service signature and error contract are unchanged (validation / not-found-same-message / never-conflict), so no OSS or downstream consumer changes. Drop the ambiguous-match logging from both single and bulk paths: key-over-subject collisions are structurally rare and resolved deterministically, so resolveCustomersByKey now returns only the resolved map. A benchmark of the single-key path (SQL-side vs Go-side precedence) confirms it is perf-neutral, with slightly fewer allocations.
…ution core
Route the single-key and bulk GetCustomer(s)ByUsageAttribution service methods
through one resolveCustomersByUsageAttribution helper (adapter fetch + precedence)
so their core resolution can't drift. Rename resolveCustomersByKey to
resolveCustomersByKeyWithPrecedence to name its key-over-subject behavior.
Add Test_GetCustomersByUsageAttribution covering key/subject resolution,
key-over-subject precedence, unmatched-key-absent (no error), and empty-key
validation, mirroring Test_GetCustomerByUsageAttribution. The fully symmetric
{K1,K2} cross-collision is left to the unit and adapter tests: CreateCustomer's
overlap guard rejects a customer key that overlaps another customer's subject,
so that state can't be constructed through the service.
The adapter benchmark inlined a copy of the service's key-over-subject precedence, mixing layers and risking drift from resolveCustomersByKeyWithPrecedence. Split into layer-honest benchmarks: - adapter package: BenchmarkCustomer(s)UsageAttribution*Query measure only the candidate query (cross_table_or vs union_all), no precedence. - service package: BenchmarkCustomer(s)UsageAttribution*Lookup drive the real service.Get(s)ByUsageAttribution end-to-end, exercising the real precedence. Shared seed/count/bulk-key fixtures move to a new leaf package customer/testutils/uabench (imports only the ent client, so both benchmark packages use it without an import cycle). The inline precedence copy is deleted; the seed SQL is unchanged. Run the timed loop of the bulk service benchmark inside b.Run so it gets a fresh, running timer: the previous flat form called b.StopTimer() during setup and b.ResetTimer() (which does not restart a stopped timer), so the loop was measured with the timer off and reported no ns/op. Test-only, no production change.
Rebase onto main changed the harness: NewTestEnv migrates internally and InitPostgresDB takes PostgresDBStateEntMigrated. Drop the leftover DBSchemaMigrate call and stale comment, and the now-redundant Schema.Create in the adapter and service benchmarks.
gergely-kurucz-konghq
force-pushed
the
feat/bulk-customer-usage-attribution-lookup-perf
branch
from
July 17, 2026 15:15
e8ddca3 to
84e266c
Compare
The service Logger was added on this branch only for the ambiguous-match logging, which was removed; s.logger had no remaining uses. Remove the field from customerservice.Config and the Logger argument from all call sites.
Contributor
Author
|
@coderabbitai resume |
Contributor
✅ Action performedReviews resumed. |
turip
previously approved these changes
Jul 21, 2026
…ageAttribution Align the bulk resolver's semantics with the single-key lookup and the repo's featureresolver.BatchResolve convention: every input key is present in the returned map, with a nil value marking not-found (was: value map, absent = not-found). Callers can distinguish requested-but-missing from not-requested, and the map holds pointers instead of copied Customer structs.
chrisgacsal
approved these changes
Jul 21, 2026
gergely-kurucz-konghq
deleted the
feat/bulk-customer-usage-attribution-lookup-perf
branch
July 23, 2026 09:54
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Optimizes and unifies the customer usage-attribution lookups. The bulk
GetCustomersByUsageAttributionstill had the cross-tableORanti-pattern that #4684 fixed for the single-key path. This PR fixes the bulk perf, collapses the single-key and bulk paths onto one predicate and one precedence mechanism, and adds layer-separated benchmarks.Why
OR. Callers can pass an uncapped key set, so this is real production exposure, not just a consistency fix.map[key]customerthat depended on DB row order — provably unfixable by reordering when two customers collide on the same key from opposite sides (one's own key vs. the other's subject key).ORDER BY lookup_priority LIMIT 1vs. an in-memory pass) and two nearly-identical predicates — a drift hazard.How
UNION ALLbranches (same shape as fix: optimize customer usage attribution lookup #4684), returning the full candidate set.customer.Service:GetCustomersByUsageAttributionnow returnsmap[string]Customerwith precedence resolved once, centrally — deleting the duplicated (and buggy) mapping logic from callers (governance).resolveCustomersByKeyWithPrecedence. The single-key public signature and error contract are unchanged (validation / not-found / never-conflict), so no consumer changes.deleted_atgrace window (IS NULL OR > at) uniformly across both predicate branches, matching theWithSubjectsconvention. Thecreated_atactive-time guard is intentionally omitted —created_atis server-assigned and immutable, so a future-dated subject can't occur through the service; guarding only one path would introduce a batch-composition-dependent bug for a state that can't happen.openmeter/customer/testutils/uabench.customer.Service.GetCustomersByUsageAttribution(now returns a map).Benchmarks
Adapter — candidate query only, cross-table
ORvsUNION ALL(100k customers, median of 3)Service — full lookup path (query + subject hydration + key-over-subject precedence, 100k customers, median of 3)
Governance endpoint — namespace-scale (fixed 100 customers × 1 feature)
Pre-fix: 4.2x jump (38 ms → 160 ms) as namespace grows. Post-fix: 1.11x (flat, within noise).
Governance endpoint — customers × features diagonal (baseline, no regression)
Summary by CodeRabbit
Summary by CodeRabbit
Improvements
key → customermapping.nilentries.Performance
Developer Experience
Greptile Summary
This PR improves customer usage-attribution lookup performance and centralizes resolution behavior. The main changes are:
UNION ALLcandidate queries.Confidence Score: 5/5
This looks safe to merge.
Important Files Changed
UNION ALLcandidate branches.Reviews (6): Last reviewed commit: "refactor(customer): return map[string]*C..." | Re-trigger Greptile
Context used (3)