refactor: balance worker partition affinity #3603
Conversation
📝 WalkthroughWalkthroughThis PR refactors the balance worker's event handling by introducing a new Changes
Sequence DiagramsequenceDiagram
participant EventSource as Event Handler
participant EventBus as EventBus
participant Worker as Balance Worker
participant Handler as handleEntitlementEvent
rect rgb(200, 220, 240)
Note over EventSource,Handler: Previous Flow
EventSource->>Handler: Call directly with options
Handler->>Handler: Process event
end
rect rgb(220, 240, 200)
Note over EventSource,Handler: New Flow (This PR)
EventSource->>EventBus: Publish RecalculateEvent<br/>(SourceOperation, OriginalEventSource, etc.)
EventBus->>Worker: Consume RecalculateEvent
Worker->>Handler: Call with mapped operation<br/>& options from event
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Areas requiring extra attention:
Possibly related PRs
Suggested labels
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 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 |
turip
changed the title
…tions This allows us to not to use redis for any caches in balance worker.
turip
force-pushed
the
refactor/balance-worker-events
branch
from
d21436b to
946813e
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
openmeter/entitlement/balanceworker/ingesthandler.go (1)
54-61: RecalculateEvent publish path looks good; consider cleaning up error handling and double‑checkAsOfchoiceNice move pushing ingest handling through
events.RecalculateEvent– the payload looks consistent with the struct (entitlement ID, original source,SourceOperationIngest, andRawIngestedEvents).Two small follow‑ups worth considering:
- You now return on the first
Publisherror, but theerrsslice anderrors.Joinat the end are never populated anymore. Either:
- drop
errsand the finalJoinentirely, or- if you still want partial progress, keep aggregating per‑entitlement publish errors and only return after processing all entitlements.
AsOfis currentlyevent.StoredAt. If the downstream snapshot logic expects “when the underlying usage happened” rather than “when we stored the batch”, you might want to feed it something derived fromeventTimestamps(e.g. the max timestamp) instead. If “storage time” is the intended semantics, then this is fine as‑is.I’d lean toward cleaning up the unused
errslogic so the function’s intent is obvious at a glance.Also applies to: 62-63
openmeter/entitlement/balanceworker/worker.go (2)
17-17: Nice unification of event handling viaRecalculateEventThis refactor to have all the entitlement/grant/reset paths publish a single
events.RecalculateEventshape looks solid:
- Entitlement identity is consistently
pkgmodels.NamespacedID{Namespace: ..., ID: ...}derived from the event’s namespace + owner/entitlement ID.OriginalEventSourceusesmetadata.ComposeResourcePathwith appropriate entities and IDs so you can still trace back to the concrete entitlement/grant event.AsOfis taken fromCreatedAt/DeletedAt/UpdatedAt/ResetAt, which lines up with how you’d expect the snapshot to be cut.SourceOperationis set explicitly for each case, which lets the downstream handler switch cleanly on intent.If you find yourself adding more producers later, it might be worth a tiny helper like
newRecalculateEvent(entitlementID, originalSource string, asOf time.Time, op events.OperationType)to avoid repeating the struct literal boilerplate, but that’s purely a readability nit.Also applies to: 211-218, 223-229, 234-240, 245-251, 256-262, 267-273, 278-284
302-315: RecalculateEvent consumer and snapshot operation mapping look correctThe new handler for
events.RecalculateEventlooks coherent with the rest of the design:
- Defaulting
snapshotOperationtosnapshot.ValueOperationUpdateand only switching toResetforOperationTypeMeteredEntitlementResetandDeleteforOperationTypeEntitlementDeletedmatches how those upstream events are modeled.- Passing
OriginalEventSourceandAsOfintoWithSource/WithEventAtpreserves the original context all the way intohandleEntitlementEvent.- Conditionally adding
WithRawIngestedEventskeeps ingest‑specific payloads attached without impacting other operation types.- Wrapping
handleEntitlementEventinEventBus.WithContext(ctx).PublishIfNoError(...)is consistent with the publishing pattern elsewhere and avoids extra ceremony here.Only tiny thought: if
RawIngestedEventsslices can be large, keep an eye on allocations as they flow through the pipeline, but that’s more of an operational concern than a blocker.Also applies to: 317-319, 321-327
openmeter/entitlement/balanceworker/events/recalculate.go (1)
6-6: OperationType enum + validation are clear; possible micro‑tweakDefining
OperationTypeas a string with explicit constants and aValidatemethod is a nice way to keep the event schema self‑describing. Usingslices.Contains(o.Values(), o)is totally fine for this scale.If you ever care about shaving a few allocations here (since
Values()builds a slice on each call), you could switchValidateto a simpleswitchover the known constants and drop theslicesimport. Not urgent, just a small optimization knob if validation ends up in a hot path.Also applies to: 10-10, 20-49
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
openmeter/entitlement/balanceworker/events/recalculate.go(4 hunks)openmeter/entitlement/balanceworker/ingesthandler.go(2 hunks)openmeter/entitlement/balanceworker/worker.go(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.go
⚙️ CodeRabbit configuration file
**/*.go: In general when reviewing the Golang code make readability and maintainability a priority, even potentially suggest restructuring the code to improve them.Performance should be a priority in critical code paths. Anything related to event ingestion, message processing, database operations (regardless of database) should be vetted for potential performance bottlenecks.
Files:
openmeter/entitlement/balanceworker/ingesthandler.goopenmeter/entitlement/balanceworker/events/recalculate.goopenmeter/entitlement/balanceworker/worker.go
🧠 Learnings (1)
📚 Learning: 2025-03-07T12:17:43.129Z
Learnt from: GAlexIHU
Repo: openmeterio/openmeter PR: 2383
File: openmeter/entitlement/metered/lateevents_test.go:37-45
Timestamp: 2025-03-07T12:17:43.129Z
Learning: In the OpenMeter codebase, test files like `openmeter/entitlement/metered/lateevents_test.go` may use variables like `meterSlug` and `namespace` without explicit declarations visible in the same file. This appears to be an accepted pattern in their test structure.
Applied to files:
openmeter/entitlement/balanceworker/ingesthandler.go
🧬 Code graph analysis (3)
openmeter/entitlement/balanceworker/ingesthandler.go (3)
openmeter/entitlement/balanceworker/events/recalculate.go (2)
RecalculateEvent(63-69)OperationTypeIngest(29-29)pkg/models/id.go (1)
NamespacedID(7-10)openmeter/event/metadata/resourcepath.go (2)
ComposeResourcePath(29-31)EntityEvent(26-26)
openmeter/entitlement/balanceworker/events/recalculate.go (3)
pkg/models/validator.go (1)
Validate(16-26)pkg/models/id.go (1)
NamespacedID(7-10)openmeter/ingest/kafkaingest/serializer/serializer.go (1)
CloudEventsKafkaPayload(20-28)
openmeter/entitlement/balanceworker/worker.go (7)
openmeter/entitlement/balanceworker/events/recalculate.go (6)
RecalculateEvent(63-69)OperationTypeEntitlementCreated(23-23)OperationTypeEntitlementDeleted(24-24)OperationTypeGrantCreated(25-25)OperationTypeGrantVoided(27-27)OperationTypeMeteredEntitlementReset(28-28)pkg/models/id.go (1)
NamespacedID(7-10)openmeter/event/metadata/resourcepath.go (3)
ComposeResourcePath(29-31)EntityEntitlement(10-10)EntityGrant(16-16)pkg/models/model.go (1)
ManagedModel(119-125)openmeter/entitlement/metered/events.go (1)
EntitlementResetEvent(17-24)openmeter/entitlement/snapshot/event.go (3)
ValueOperationUpdate(23-23)ValueOperationReset(22-22)ValueOperationDelete(24-24)openmeter/entitlement/balanceworker/entitlementhandler.go (4)
WithSource(49-53)WithEventAt(55-59)WithSourceOperation(61-65)WithRawIngestedEvents(67-71)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)
- GitHub Check: Artifacts / Container image
- GitHub Check: Artifacts / Benthos Collector Container image
- GitHub Check: Lint
- GitHub Check: Code Generators
- GitHub Check: Migration Checks
- GitHub Check: Build
- GitHub Check: Test
- GitHub Check: Secret Scanning
- GitHub Check: Scan GitHub Workflows
- GitHub Check: Repository Scan
- GitHub Check: SAST (semgrep)
- GitHub Check: Analyze (go)
🔇 Additional comments (2)
openmeter/entitlement/balanceworker/ingesthandler.go (1)
12-12: Importing balanceworker events looks goodPulling in the typed
eventspackage here fits the new RecalculateEvent flow and keeps the handler nicely aligned with the worker logic. No issues from my side.openmeter/entitlement/balanceworker/events/recalculate.go (1)
57-57: RecalculateEvent v2 shape and validation look consistent with usageThe v2
RecalculateEventstruct and its helpers line up well with how the worker is using it:
Entitlement+AsOfcarry the core recalculation coordinates.OriginalEventSourceis now the source forEventMetadata.Source, whileSubjectcontinues to be the entitlement resource path, giving you both “where it came from” and “what it’s about”.SourceOperationandRawIngestedEventsgive the consumer enough context to decide how to treat the recalculation (update/reset/delete / ingest details).Validate()now requires a non‑zeroSourceOperation, which should help catch mis‑constructed events early.The version bump to
"v2"makes sense given the schema change; just make sure your producers/consumers are all aligned on the new version and that any legacy v1 messages in the wild are either drained or handled gracefully.Also applies to: 64-68, 76-79, 93-95
Overview
This patch makes sure that for each incoming event balance worker yields another event, that is used to process the entitlement recalculation. This is useful as going forward we will not need redis based caches as the entitlements end up on the same worker instance.
Summary by CodeRabbit
New Features
Refactor