diff --git a/AGENTS.md b/AGENTS.md index 6da3250109..0b44fdcedb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,6 +32,9 @@ - Name string enum constants ``, for example `InvoiceStatusDraft`. +- Prefer methods when behavior is intrinsic to an existing domain type. Use + freestanding functions when an operation has + no natural receiver. - Do not extract trivial or single-use helpers unless the name captures non-obvious domain intent. Inline pass-through wrappers. - Do not hide type switching, validation, persistence mapping, or meaningful diff --git a/openmeter/credit/README.md b/openmeter/credit/README.md new file mode 100644 index 0000000000..1ca903a4f8 --- /dev/null +++ b/openmeter/credit/README.md @@ -0,0 +1,126 @@ +# Credit + +Credit calculates how metered usage consumes grants over time. Its primary +owner is a metered entitlement, but the calculation model is generic: an owner +provides usage, grants, usage periods, and reset behavior. + +Credit balances are derived state. They are not money and they are not an +accounting ledger. See [Ledger](../ledger/README.md) for accounting state and +[Entitlements](../entitlement/README.md) for access semantics. + +## Grants + +A grant is an immutable, time-effective allocation of credit. Creation adds a +new grant; voiding ends its effective period without rewriting its earlier +meaning. + +Usage is burnt in deterministic order: + +1. lower numeric priority first; +2. earlier expiry first; +3. earlier creation first. + +Recurring grants replenish their own balance on their recurrence. Owner usage +resets are different: they apply each active grant's rollover limits and begin a +new usage period. + +Only active grants contribute to a balance. Expired and voided grants remain +part of historical calculations for times when they were active. + +## Time semantics + +Credit has one-minute temporal resolution. This is a domain constraint inherited +from minute-windowed metering, not merely a snapshot or database detail. + +Grant effective times, recurrence anchors, voids, and resets are normalized to +minute boundaries. A balance read at a sub-minute timestamp rounds up to include +usage from that partial minute. Callers must not rely on ordering or balance +changes within the same minute. + +History boundaries and persisted snapshots must use the same minute alignment. +Changing this resolution requires changing the metering and credit time model +together. + +## Balance calculation + +A balance at time `t` is calculated from: + +- the owner's metering and reset configuration; +- grants effective during the calculation; +- usage-period and explicit reset boundaries; +- metered usage up to `t`. + +The engine burns usage from active grants in priority order. Usage that cannot +be covered becomes overage. The result contains the remaining grant balances, +overage, and the usage-period state needed to continue the calculation. + +An engine run also produces a contiguous burn-down history anchored to its +starting balance. The anchor is part of the history's meaning: the segments +alone are not enough to reconstruct balances. + +## Resets + +A reset closes one usage period and starts another. Active grant balances are +first constrained by their rollover limits. If overage is preserved, it is then +burnt from those rolled-over balances in the new period. + +The burn-down history represents that reset transition explicitly. Rollover and +preserved overage burn must not be hidden by changing the meaning of the usage +segments on either side. + +Resetting changes both balance state and the owner's usage-period timeline. +Those changes are serialized for the owner so readers cannot observe only half +of the transition. + +## Balance snapshots + +A balance snapshot is a complete calculation checkpoint at a specific time. It +is a cache of derived state, not an event and not part of the source history. + +### Creation + +Balance calculations may save a snapshot at an eligible history breakpoint. +Snapshots are kept behind a configurable grace window so recent usage remains +recalculable when events arrive late. A reset saves the balance at the reset +boundary as part of the reset transaction. + +`LATEST` balances are not persisted as snapshots because a latest value is not +reusable cumulative state. + +### Use + +A balance calculation starts from the newest usable snapshot at or before the +query time and evaluates only the inputs after it. If no snapshot is usable, +calculation starts from the owner's measurement start. + +Snapshot presence affects cost, never meaning. Removing intermediate snapshots, +removing every snapshot before the latest one, or inserting an earlier snapshot +must not change the result obtained from a later checkpoint. + +Snapshots are independent checkpoints rather than a chain. Each one contains +all state required to continue the calculation, and incompatible snapshots are +ignored. + +### Invalidation + +A time-effective input can make saved derived state stale. The operation that +changes grants, resets, usage, or another calculation input is responsible for +invalidating affected snapshots. + +If a change affects the state represented by a snapshot, that snapshot must be +replaced or invalidated together with every later snapshot. Invalidated +snapshots are excluded from selection; the next read recomputes from an earlier +usable checkpoint or from measurement start. + +This rule is what makes snapshot pruning safe and prevents balance correctness +from depending on which checkpoints happen to exist. + +## Boundaries + +- The owner supplies metering, usage attribution, reset timing, and conversion + behavior. +- Streaming owns raw usage; credit consumes meter results. +- Credit owns grants, burn order, rollover, overage, calculation history, and + balance snapshots. +- Entitlement turns the calculated balance into an access decision. +- Billing and ledger consumers must not treat feature credit as money. diff --git a/openmeter/credit/adapter/balance_snapshot.go b/openmeter/credit/adapter/balance_snapshot.go index 9d4516bd18..7b4f494557 100644 --- a/openmeter/credit/adapter/balance_snapshot.go +++ b/openmeter/credit/adapter/balance_snapshot.go @@ -2,6 +2,7 @@ package adapter import ( "context" + "fmt" "time" "entgo.io/ent/dialect/sql" @@ -47,6 +48,7 @@ func (b *balanceSnapshotRepo) GetLatestValidAt(ctx context.Context, owner models db_balancesnapshot.Namespace(owner.Namespace), db_balancesnapshot.AtLTE(at), db_balancesnapshot.DeletedAtIsNil(), + db_balancesnapshot.UsageSnapshotNotNil(), ). // in case there were multiple snapshots for the same time return the newest one Order(db_balancesnapshot.ByAt(sql.OrderDesc()), db_balancesnapshot.ByUpdatedAt(sql.OrderDesc())). @@ -65,20 +67,27 @@ func (b *balanceSnapshotRepo) GetLatestValidAt(ctx context.Context, owner models func (b *balanceSnapshotRepo) Save(ctx context.Context, owner models.NamespacedID, balances []balance.Snapshot) error { return entutils.TransactingRepoWithNoValue(ctx, b, func(ctx context.Context, rep *balanceSnapshotRepo) error { commands := make([]*db.BalanceSnapshotCreate, 0, len(balances)) - for _, balance := range balances { + for _, snapshot := range balances { + if snapshot.UsageSnapshot == nil { + return fmt.Errorf("cannot save incomplete balance snapshot at %s", snapshot.At) + } + + // Keep writing the legacy usage representation for compatibility + // with old readers during the rolling migration. command := rep.db.BalanceSnapshot.Create(). SetNamespace(owner.Namespace). SetOwnerID(owner.ID). - SetBalance(balance.Balance()). - SetAt(balance.At). - SetGrantBalances(balance.Balances). - SetOverage(balance.Overage). - SetUsage(&balance.Usage) + SetBalance(snapshot.Balance()). + SetAt(snapshot.At). + SetGrantBalances(snapshot.Balances). + SetOverage(snapshot.Overage). + SetUsage(&snapshot.Usage). + SetUsageSnapshot(snapshot.UsageSnapshot) // Record the conversion regime this snapshot was computed under (OM-400) so // the resume path can refuse to reuse it under a different regime. Pointer // GoType has no SetNillable*, so nil-guard the set (nil = raw). - if balance.UnitConfig != nil { - command = command.SetUnitConfig(balance.UnitConfig) + if snapshot.UnitConfig != nil { + command = command.SetUnitConfig(snapshot.UnitConfig) } commands = append(commands, command) } @@ -93,7 +102,12 @@ func mapBalanceSnapshotEntity(entity *db.BalanceSnapshot) balance.Snapshot { Overage: entity.Overage, At: entity.At.In(time.UTC), } + if entity.UsageSnapshot != nil { + s.UsageSnapshot = entity.UsageSnapshot + } if entity.Usage != nil { + // Hydrate legacy usage only so subsequent snapshots remain readable by + // old binaries during the rolling migration. s.Usage = *entity.Usage } if entity.UnitConfig != nil { diff --git a/openmeter/credit/balance/balance.go b/openmeter/credit/balance/balance.go index 906ed4960e..82dcca97cf 100644 --- a/openmeter/credit/balance/balance.go +++ b/openmeter/credit/balance/balance.go @@ -7,7 +7,9 @@ import ( "github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig" ) -func NewStartingMap(grants []grant.Grant, at time.Time) Map { +// NewStartingSnapshot returns the complete zero-usage snapshot from which +// measurement starts. +func NewStartingSnapshot(grants []grant.Grant, at time.Time) Snapshot { balances := make(Map) for _, grant := range grants { if grant.ActiveAt(at) { @@ -16,7 +18,20 @@ func NewStartingMap(grants []grant.Grant, at time.Time) Map { balances.Set(grant.ID, 0.0) } } - return balances + + return Snapshot{ + Usage: SnapshottedUsage{ + Since: at, + Usage: 0, + }, + UsageSnapshot: &UsageSnapshot{ + Usage: 0, + TotalGrantUsage: 0, + }, + Balances: balances, + Overage: 0, + At: at, + } } // Represents a point in time balance of grants @@ -69,11 +84,18 @@ func (g Map) ExactlyForGrants(grants []grant.Grant) bool { return true } +// SnapshottedUsage is the legacy usage representation whose value is relative +// to an explicitly stored starting timestamp. +// +// Deprecated: use Snapshot.UsageSnapshot for complete usage-period state. type SnapshottedUsage struct { Usage float64 `json:"usage"` Since time.Time `json:"since"` } +// IsZero reports whether the legacy usage representation is unset. +// +// Deprecated: only use this while reading legacy snapshots. func (s SnapshottedUsage) IsZero() bool { return s.Usage == 0.0 && s.Since.IsZero() } @@ -86,7 +108,11 @@ type UsageSnapshot struct { } type Snapshot struct { + // Usage is retained for compatibility with the legacy persistence shape. + // + // Deprecated: use UsageSnapshot for engine calculations. Usage SnapshottedUsage + // UsageSnapshot is nil for snapshots created without complete usage-period // state. UsageSnapshot *UsageSnapshot @@ -103,6 +129,23 @@ type Snapshot struct { UnitConfig *unitconfig.UnitConfig } +// Clone returns a snapshot whose mutable balances and unit configuration are +// independent from the source. +func (s Snapshot) Clone() Snapshot { + cloned := s + cloned.Balances = s.Balances.Clone() + if s.UsageSnapshot != nil { + usageSnapshot := *s.UsageSnapshot + cloned.UsageSnapshot = &usageSnapshot + } + if s.UnitConfig != nil { + unitConfig := s.UnitConfig.Clone() + cloned.UnitConfig = &unitConfig + } + + return cloned +} + func (g Snapshot) Balance() float64 { return g.Balances.Balance() } diff --git a/openmeter/credit/balance/balance_test.go b/openmeter/credit/balance/balance_test.go index f928961b29..bd51c6c3f4 100644 --- a/openmeter/credit/balance/balance_test.go +++ b/openmeter/credit/balance/balance_test.go @@ -2,13 +2,44 @@ package balance_test import ( "testing" + "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" balance "github.com/openmeterio/openmeter/openmeter/credit/balance" "github.com/openmeterio/openmeter/openmeter/credit/grant" ) +func TestNewStartingSnapshot(t *testing.T) { + at := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + activeGrant := grant.Grant{ + ID: "active", + Amount: 100, + EffectiveAt: at, + } + futureGrant := grant.Grant{ + ID: "future", + Amount: 200, + EffectiveAt: at.Add(time.Hour), + } + + snapshot := balance.NewStartingSnapshot([]grant.Grant{activeGrant, futureGrant}, at) + + assert.Equal(t, at, snapshot.At) + assert.Equal(t, balance.Map{ + activeGrant.ID: activeGrant.Amount, + futureGrant.ID: 0, + }, snapshot.Balances) + assert.Zero(t, snapshot.Overage) + assert.Equal(t, balance.SnapshottedUsage{ + Since: at, + Usage: 0, + }, snapshot.Usage) + require.NotNil(t, snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{}, *snapshot.UsageSnapshot) +} + func TestGrantBalanceMap(t *testing.T) { makeGrant := func(id string) grant.Grant { return grant.Grant{ @@ -58,3 +89,26 @@ func TestGrantBalanceMap(t *testing.T) { })) }) } + +func TestSnapshotCloneCopiesUsageSnapshot(t *testing.T) { + snapshot := balance.Snapshot{ + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 10, + }, + } + + cloned := snapshot.Clone() + require.NotNil(t, cloned.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 10, + }, *cloned.UsageSnapshot) + + cloned.UsageSnapshot.Usage = 15 + cloned.UsageSnapshot.TotalGrantUsage = 20 + assert.Equal(t, balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 10, + }, *snapshot.UsageSnapshot) +} diff --git a/openmeter/credit/balance/repository.go b/openmeter/credit/balance/repository.go index 8f1415e6b6..bf407ec1c1 100644 --- a/openmeter/credit/balance/repository.go +++ b/openmeter/credit/balance/repository.go @@ -10,7 +10,7 @@ import ( type SnapshotRepo interface { InvalidateAfter(ctx context.Context, owner models.NamespacedID, at time.Time) error - // The returned Snapshot might not have usage data. + // GetLatestValidAt returns the latest complete snapshot. GetLatestValidAt(ctx context.Context, owner models.NamespacedID, at time.Time) (Snapshot, error) Save(ctx context.Context, owner models.NamespacedID, balances []Snapshot) error } diff --git a/openmeter/credit/balance/service.go b/openmeter/credit/balance/service.go index 4436a618cb..4a6484c0fb 100644 --- a/openmeter/credit/balance/service.go +++ b/openmeter/credit/balance/service.go @@ -4,10 +4,7 @@ import ( "context" "time" - "github.com/openmeterio/openmeter/openmeter/credit/grant" - "github.com/openmeterio/openmeter/openmeter/streaming" "github.com/openmeterio/openmeter/pkg/models" - "github.com/openmeterio/openmeter/pkg/timeutil" ) type SnapshotService interface { @@ -19,25 +16,16 @@ type SnapshotService interface { } type SnapshotServiceConfig struct { - OwnerConnector grant.OwnerConnector - StreamingConnector streaming.Connector - Repo SnapshotRepo + Repo SnapshotRepo } type service struct { - UsageQuerier UsageQuerier SnapshotServiceConfig } func NewSnapshotService(conf SnapshotServiceConfig) SnapshotService { return &service{ SnapshotServiceConfig: conf, - // We build a custom UsageQuerier for our usecase here - UsageQuerier: NewUsageQuerier(UsageQuerierConfig{ - StreamingConnector: conf.StreamingConnector, - DescribeOwner: conf.OwnerConnector.DescribeOwner, - GetUsagePeriodStartAt: conf.OwnerConnector.GetUsagePeriodStartAt, - }), } } @@ -48,33 +36,7 @@ func (s *service) InvalidateAfter(ctx context.Context, owner models.NamespacedID } func (s *service) GetLatestValidAt(ctx context.Context, owner models.NamespacedID, at time.Time) (Snapshot, error) { - res, err := s.Repo.GetLatestValidAt(ctx, owner, at) - if err != nil { - return Snapshot{}, err - } - - // We have to manually fill in the usage data if it wasn't saved - if res.Usage.IsZero() { - periodStart, err := s.OwnerConnector.GetUsagePeriodStartAt(ctx, owner, res.At) - if err != nil { - return Snapshot{}, err - } - - usage, err := s.UsageQuerier.QueryUsage(ctx, owner, timeutil.ClosedPeriod{ - From: periodStart, - To: res.At, - }) - if err != nil { - return Snapshot{}, err - } - - res.Usage = SnapshottedUsage{ - Usage: usage, - Since: periodStart, - } - } - - return res, nil + return s.Repo.GetLatestValidAt(ctx, owner, at) } func (s *service) Save(ctx context.Context, owner models.NamespacedID, balances []Snapshot) error { diff --git a/openmeter/credit/balance/service_test.go b/openmeter/credit/balance/service_test.go deleted file mode 100644 index b49eda1be6..0000000000 --- a/openmeter/credit/balance/service_test.go +++ /dev/null @@ -1,215 +0,0 @@ -package balance_test - -import ( - "context" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/openmeterio/openmeter/openmeter/credit/balance" - "github.com/openmeterio/openmeter/openmeter/credit/grant" - "github.com/openmeterio/openmeter/openmeter/meter" - "github.com/openmeterio/openmeter/openmeter/streaming" - "github.com/openmeterio/openmeter/openmeter/streaming/testutils" - "github.com/openmeterio/openmeter/pkg/models" - "github.com/openmeterio/openmeter/pkg/timeutil" -) - -// MockSnapshotRepo is a mock implementation of balance.SnapshotRepo -type MockSnapshotRepo struct { - snapshots map[string]balance.Snapshot -} - -func NewMockSnapshotRepo() *MockSnapshotRepo { - return &MockSnapshotRepo{ - snapshots: make(map[string]balance.Snapshot), - } -} - -func (m *MockSnapshotRepo) InvalidateAfter(ctx context.Context, owner models.NamespacedID, at time.Time) error { - return nil -} - -func (m *MockSnapshotRepo) GetLatestValidAt(ctx context.Context, owner models.NamespacedID, at time.Time) (balance.Snapshot, error) { - key := owner.Namespace + ":" + owner.ID - snapshot, ok := m.snapshots[key] - if !ok { - return balance.Snapshot{}, balance.NoSavedBalanceForOwnerError{ - Owner: owner, - Time: at, - } - } - return snapshot, nil -} - -func (m *MockSnapshotRepo) Save(ctx context.Context, owner models.NamespacedID, balances []balance.Snapshot) error { - key := owner.Namespace + ":" + owner.ID - if len(balances) > 0 { - m.snapshots[key] = balances[0] - } - return nil -} - -// MockOwnerConnector is a mock implementation of grant.OwnerConnector -type MockOwnerConnector struct { - usagePeriodStartAt time.Time - meterSlug string -} - -func NewMockOwnerConnector(usagePeriodStartAt time.Time, meterSlug string) *MockOwnerConnector { - return &MockOwnerConnector{ - usagePeriodStartAt: usagePeriodStartAt, - meterSlug: meterSlug, - } -} - -func (m *MockOwnerConnector) DescribeOwner(ctx context.Context, id models.NamespacedID) (grant.Owner, error) { - return grant.Owner{ - NamespacedID: id, - Meter: meter.Meter{ - Key: m.meterSlug, - Aggregation: meter.MeterAggregationSum, - }, - DefaultQueryParams: streaming.QueryParams{ - FilterSubject: []string{"subject1"}, - }, - }, nil -} - -func (m *MockOwnerConnector) GetResetTimelineInclusive(ctx context.Context, id models.NamespacedID, period timeutil.ClosedPeriod) (timeutil.SimpleTimeline, error) { - return timeutil.SimpleTimeline{}, nil -} - -func (m *MockOwnerConnector) GetUsagePeriodStartAt(ctx context.Context, id models.NamespacedID, at time.Time) (time.Time, error) { - return m.usagePeriodStartAt, nil -} - -func (m *MockOwnerConnector) GetStartOfMeasurement(ctx context.Context, id models.NamespacedID) (time.Time, error) { - return m.usagePeriodStartAt, nil -} - -func (m *MockOwnerConnector) EndCurrentUsagePeriod(ctx context.Context, id models.NamespacedID, params grant.EndCurrentUsagePeriodParams) error { - return nil -} - -func (m *MockOwnerConnector) LockOwnerForTx(ctx context.Context, id models.NamespacedID, wait bool) error { - return nil -} - -func TestGetLatestValidAt(t *testing.T) { - // Common setup - ctx := context.Background() - now := time.Now().UTC() - periodStart := now.Add(-24 * time.Hour) - meterSlug := "test-meter" - - // Create owner - owner := models.NamespacedID{ - Namespace: "test-namespace", - ID: "test-owner", - } - - t.Run("Should fill usage if snapshot has zero usage", func(t *testing.T) { - // Create mock streaming connector - streamingConnector := testutils.NewMockStreamingConnector(t) - - // Add usage data to the mock streaming connector - streamingConnector.AddSimpleEvent(meterSlug, 100.0, now.Add(-12*time.Hour)) - - // Create mock snapshot repo with a snapshot that has zero usage - mockRepo := NewMockSnapshotRepo() - mockRepo.snapshots[owner.Namespace+":"+owner.ID] = balance.Snapshot{ - Usage: balance.SnapshottedUsage{}, // Zero usage - Balances: balance.Map{"grant1": 1000.0}, - At: now.Add(-1 * time.Hour), - } - - // Create mock owner connector - mockOwnerConnector := NewMockOwnerConnector(periodStart, meterSlug) - - // Create the service - service := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: mockOwnerConnector, - StreamingConnector: streamingConnector, - Repo: mockRepo, - }) - - // Test - snapshot, err := service.GetLatestValidAt(ctx, owner, now) - - // Verify - require.NoError(t, err) - assert.Equal(t, 100.0, snapshot.Usage.Usage, "Usage should be filled with the value from streaming connector") - assert.Equal(t, periodStart, snapshot.Usage.Since, "Usage since should be set to the period start") - assert.Equal(t, 1000.0, snapshot.Balances["grant1"], "Balance should remain unchanged") - }) - - t.Run("Should preserve existing usage if snapshot already has usage data", func(t *testing.T) { - // Create mock streaming connector - streamingConnector := testutils.NewMockStreamingConnector(t) - - // Add usage data to the mock streaming connector - streamingConnector.AddSimpleEvent(meterSlug, 100.0, now.Add(-12*time.Hour)) - - // Create mock snapshot repo with a snapshot that already has usage data - mockRepo := NewMockSnapshotRepo() - mockRepo.snapshots[owner.Namespace+":"+owner.ID] = balance.Snapshot{ - Usage: balance.SnapshottedUsage{ - Usage: 50.0, - Since: periodStart, - }, - Balances: balance.Map{"grant1": 1000.0}, - At: now.Add(-1 * time.Hour), - } - - // Create mock owner connector - mockOwnerConnector := NewMockOwnerConnector(periodStart, meterSlug) - - // Create the service - service := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: mockOwnerConnector, - StreamingConnector: streamingConnector, - Repo: mockRepo, - }) - - // Test - snapshot, err := service.GetLatestValidAt(ctx, owner, now) - - // Verify - require.NoError(t, err) - assert.Equal(t, 50.0, snapshot.Usage.Usage, "Usage should remain unchanged") - assert.Equal(t, periodStart, snapshot.Usage.Since, "Usage since should remain unchanged") - assert.Equal(t, 1000.0, snapshot.Balances["grant1"], "Balance should remain unchanged") - }) - - t.Run("Should return error if no snapshot exists", func(t *testing.T) { - // Create mock streaming connector - streamingConnector := testutils.NewMockStreamingConnector(t) - - // Add usage data to the mock streaming connector - streamingConnector.AddSimpleEvent(meterSlug, 100.0, now.Add(-12*time.Hour)) - - // Create empty mock snapshot repo (no snapshots) - mockRepo := NewMockSnapshotRepo() - - // Create mock owner connector - mockOwnerConnector := NewMockOwnerConnector(periodStart, meterSlug) - - // Create the service - service := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: mockOwnerConnector, - StreamingConnector: streamingConnector, - Repo: mockRepo, - }) - - // Test - _, err := service.GetLatestValidAt(ctx, owner, now) - - // Verify - require.Error(t, err) - _, isNoSavedBalanceErr := err.(balance.NoSavedBalanceForOwnerError) - assert.True(t, isNoSavedBalanceErr, "Expected NoSavedBalanceForOwnerError") - }) -} diff --git a/openmeter/credit/engine/engine.go b/openmeter/credit/engine/engine.go index 513c2f2f23..cacb9ed7aa 100644 --- a/openmeter/credit/engine/engine.go +++ b/openmeter/credit/engine/engine.go @@ -4,9 +4,6 @@ import ( "context" "time" - "github.com/alpacahq/alpacadecimal" - "github.com/samber/lo" - "github.com/openmeterio/openmeter/openmeter/credit/balance" "github.com/openmeterio/openmeter/openmeter/credit/grant" "github.com/openmeterio/openmeter/openmeter/meter" @@ -39,7 +36,7 @@ func (p RunParams) Clone() RunParams { Meter: p.Meter, Grants: grants, Until: p.Until, - StartingSnapshot: p.StartingSnapshot, + StartingSnapshot: p.StartingSnapshot.Clone(), ResetBehavior: p.ResetBehavior, Resets: resets, } @@ -56,16 +53,6 @@ type RunResult struct { RunParams RunParams } -// TotalAvailableGrantAmountAtLastPeriod is the total grant amount available in the run period: -// grant-covered usage plus remaining balances still available at the end. -func (r RunResult) TotalAvailableGrantAmountAtLastPeriod() float64 { - lastUsagePeriodHistory, _ := lo.Last(r.History.ChunkByResets()) - - return lastUsagePeriodHistory.TotalGrantUsage(). - Add(alpacadecimal.NewFromFloat(r.Snapshot.Balance())). - InexactFloat64() -} - type Engine interface { // Burns down all grants in the defined period by the usage amounts. // diff --git a/openmeter/credit/engine/engine_test.go b/openmeter/credit/engine/engine_test.go index 37bac953b1..0c52661067 100644 --- a/openmeter/credit/engine/engine_test.go +++ b/openmeter/credit/engine/engine_test.go @@ -114,9 +114,10 @@ func Test_Fuzzing(t *testing.T) { engine.RunParams{ Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ - Balances: startingBalance, - Overage: 0, - At: start, + UsageSnapshot: zeroUsageSnapshot(), + Balances: startingBalance, + Overage: 0, + At: start, }, Meter: meter, Until: intermediate, @@ -212,9 +213,10 @@ func Test_Fuzzing(t *testing.T) { engine.RunParams{ Grants: gCp, StartingSnapshot: balance.Snapshot{ - Balances: balances, - Overage: 0, - At: start, + UsageSnapshot: zeroUsageSnapshot(), + Balances: balances, + Overage: 0, + At: start, }, Meter: meter, Until: end, @@ -299,9 +301,10 @@ func Test_Fuzzing(t *testing.T) { engine.RunParams{ Grants: gCp, StartingSnapshot: balance.Snapshot{ - Balances: startingBalances, - Overage: 0, - At: start, + UsageSnapshot: zeroUsageSnapshot(), + Balances: startingBalances, + Overage: 0, + At: start, }, Meter: meter, Until: end, @@ -337,9 +340,10 @@ func Test_Fuzzing(t *testing.T) { engine.RunParams{ Grants: gCp, StartingSnapshot: balance.Snapshot{ - Balances: balances, - Overage: overage, - At: pStart, + UsageSnapshot: zeroUsageSnapshot(), + Balances: balances, + Overage: overage, + At: pStart, }, Until: pEnd, Meter: meter, diff --git a/openmeter/credit/engine/history.go b/openmeter/credit/engine/history.go index 8bb3b15daf..c0ece5f0c0 100644 --- a/openmeter/credit/engine/history.go +++ b/openmeter/credit/engine/history.go @@ -3,6 +3,7 @@ package engine import ( "encoding/json" "fmt" + "math" "sort" "time" @@ -16,6 +17,8 @@ type SegmentTerminationReason struct { PriorityChange bool Recurrence []string // Grant IDs UsageReset bool + // Rollover marks grant balance rollover followed by preserved overage burn. + Rollover bool } type GrantUsageTerminationReason string @@ -43,13 +46,41 @@ type GrantUsage struct { TerminationReason GrantUsageTerminationReason } +type GrantUsages []GrantUsage + +func (u GrantUsages) Sum() alpacadecimal.Decimal { + total := alpacadecimal.NewFromFloat(0) + for _, usage := range u { + total = total.Add(alpacadecimal.NewFromFloat(usage.Usage)) + } + + return total +} + +func validateStartingSnapshot(snapshot balance.Snapshot) error { + if snapshot.UsageSnapshot == nil { + return fmt.Errorf("starting snapshot usage snapshot is missing") + } + + totalGrantUsage := snapshot.UsageSnapshot.TotalGrantUsage + if totalGrantUsage < 0 || math.IsNaN(totalGrantUsage) || math.IsInf(totalGrantUsage, 0) { + return fmt.Errorf("starting snapshot total grant usage must be a finite non-negative number, got %v", totalGrantUsage) + } + + return nil +} + // GrantBurnDownHistorySegment represents the smallest segment of grant usage which we store and calculate. // -// A segment represents a period of time in which: +// A non-rollover segment represents a period of time in which: // 1) The grant priority does not change // 2) Grants do not recurr // 3) There was no usage reset // +// A rollover segment is an instantaneous reset transition with no metered usage. +// Its starting balance is the rolled-over grant balance, and its grant usages +// capture preserved overage burnt from grants in the new usage period. +// // It is not necessarily the largest such segment. type GrantBurnDownHistorySegment struct { timeutil.ClosedPeriod @@ -58,7 +89,7 @@ type GrantBurnDownHistorySegment struct { TotalUsage float64 // Total usage of the feature in the Period OverageAtStart float64 // Usage beyond what could be burnt down from the grants in the previous segment (if any) Overage float64 // Usage beyond what cloud be burnt down from the grants - GrantUsages []GrantUsage // Grant usages in the segment order by grant priority + GrantUsages GrantUsages // Grant usages in the segment order by grant priority } // Returns GrantBalanceMap at the end of the segment @@ -70,32 +101,72 @@ func (s GrantBurnDownHistorySegment) ApplyUsage() balance.Map { return balance } -func NewGrantBurnDownHistory(segments []GrantBurnDownHistorySegment, usageAtStart balance.SnapshottedUsage) (GrantBurnDownHistory, error) { +// NewGrantBurnDownHistory creates a history anchored to startingSnapshot. +// Segments must continuously cover time beginning at the snapshot. +func NewGrantBurnDownHistory(segments []GrantBurnDownHistorySegment, startingSnapshot balance.Snapshot) (GrantBurnDownHistory, error) { + if err := validateStartingSnapshot(startingSnapshot); err != nil { + return GrantBurnDownHistory{}, err + } + s := make([]GrantBurnDownHistorySegment, len(segments)) copy(s, segments) - // sort segments by time - sort.Slice(s, func(i, j int) bool { + for i, segment := range s { + if segment.From.After(segment.To) { + return GrantBurnDownHistory{}, fmt.Errorf("segment %d starts after it ends", i) + } + } + + // Sort segments by time. Rollover transitions precede regular segments at + // the same timestamp so the reset transition is applied before new-period + // usage. + sort.SliceStable(s, func(i, j int) bool { + if s[i].ClosedPeriod.From.Equal(s[j].ClosedPeriod.From) { + return s[i].TerminationReasons.Rollover && !s[j].TerminationReasons.Rollover + } + return s[i].ClosedPeriod.From.Before(s[j].ClosedPeriod.From) }) - // validate no two segments overlap - for i := range s { - if i == 0 { - continue + if len(s) > 0 { + if !s[0].From.Equal(startingSnapshot.At) { + return GrantBurnDownHistory{}, fmt.Errorf( + "first segment starts at %s, expected starting snapshot at %s", + s[0].From, + startingSnapshot.At, + ) } - if s[i-1].To.After(s[i].From) { - return GrantBurnDownHistory{}, fmt.Errorf("segments %d and %d overlap", i-1, i) + if s[0].OverageAtStart != startingSnapshot.Overage { + return GrantBurnDownHistory{}, fmt.Errorf( + "first segment starts with overage %f, expected starting snapshot overage %f", + s[0].OverageAtStart, + startingSnapshot.Overage, + ) } } - return GrantBurnDownHistory{segments: s, usageAtStart: usageAtStart}, nil + for i := 1; i < len(s); i++ { + if !s[i-1].To.Equal(s[i].From) { + return GrantBurnDownHistory{}, fmt.Errorf( + "segments %d and %d are not contiguous: %s != %s", + i-1, + i, + s[i-1].To, + s[i].From, + ) + } + } + + return GrantBurnDownHistory{ + segments: s, + startingSnapshot: startingSnapshot.Clone(), + }, nil } type GrantBurnDownHistory struct { - segments []GrantBurnDownHistorySegment - usageAtStart balance.SnapshottedUsage + segments []GrantBurnDownHistorySegment + startingSnapshot balance.Snapshot } func (g GrantBurnDownHistory) MarshalJSON() ([]byte, error) { @@ -103,36 +174,48 @@ func (g GrantBurnDownHistory) MarshalJSON() ([]byte, error) { } func (g *GrantBurnDownHistory) GetSnapshotAtStartOfSegment(segmentIndex int) (balance.Snapshot, error) { - // Let's validate the segment index if segmentIndex < 0 || segmentIndex >= len(g.segments) { return balance.Snapshot{}, fmt.Errorf("segment index %d out of bounds", segmentIndex) } - // Let's get the segment + return g.getSnapshotAtStartOfSegment(segmentIndex), nil +} + +func (g *GrantBurnDownHistory) getSnapshotAtStartOfSegment(segmentIndex int) balance.Snapshot { segment := g.segments[segmentIndex] + snapshot := g.startingSnapshot.Clone() + snapshot.Usage = g.getUsageInPeriodUntilSegment(segmentIndex) + usageSnapshot := g.getUsageSnapshotAtStartOfSegment(segmentIndex) + snapshot.UsageSnapshot = &usageSnapshot + snapshot.Overage = segment.OverageAtStart + snapshot.Balances = segment.BalanceAtStart.Clone() + snapshot.At = segment.From + + return snapshot +} - // Let's get the usage in the period until the start of the segment - usage, err := g.GetUsageInPeriodUntilSegment(segmentIndex) - if err != nil { - return balance.Snapshot{}, fmt.Errorf("failed to get usage in period until segment: %w", err) - } +func (g *GrantBurnDownHistory) getUsageSnapshotAtStartOfSegment(segmentIndex int) balance.UsageSnapshot { + usageSnapshot := *g.startingSnapshot.UsageSnapshot - return balance.Snapshot{ - Usage: usage, - Overage: segment.OverageAtStart, - Balances: segment.BalanceAtStart, - At: segment.From, - }, nil -} + for i := 0; i < segmentIndex; i++ { + segment := g.segments[i] + if segment.TerminationReasons.UsageReset { + usageSnapshot = balance.UsageSnapshot{} + continue + } -// GetUsageInPeriodUntilSegment returns the SnapshottedUsage at the start of the given segment -func (g *GrantBurnDownHistory) GetUsageInPeriodUntilSegment(segmentIndex int) (balance.SnapshottedUsage, error) { - // Let's validate the segment index - if segmentIndex < 0 || segmentIndex >= len(g.segments) { - return balance.SnapshottedUsage{}, fmt.Errorf("segment index %d out of bounds", segmentIndex) + usageSnapshot.Usage += segment.TotalUsage + usageSnapshot.TotalGrantUsage = alpacadecimal.NewFromFloat(usageSnapshot.TotalGrantUsage). + Add(segment.GrantUsages.Sum()). + InexactFloat64() } - // Let's find the segment of the last reset before the provided segment + return usageSnapshot +} + +func (g *GrantBurnDownHistory) getUsageInPeriodUntilSegment(segmentIndex int) balance.SnapshottedUsage { + // Reconstruct the legacy, Since-relative usage representation while it is + // still required for persistence compatibility. lastResetSegmentIndex := -1 for i := 0; i < segmentIndex; i++ { if g.segments[i].TerminationReasons.UsageReset { @@ -140,8 +223,7 @@ func (g *GrantBurnDownHistory) GetUsageInPeriodUntilSegment(segmentIndex int) (b } } - // Now let's build a starting SnapshottedUsage - usage := g.usageAtStart + usage := g.startingSnapshot.Usage if lastResetSegmentIndex != -1 { // We need the segment right after the last reset @@ -151,12 +233,11 @@ func (g *GrantBurnDownHistory) GetUsageInPeriodUntilSegment(segmentIndex int) (b } } - // Now we need to add up the usage in all segments between the starting usage and the provided segment for i := lastResetSegmentIndex + 1; i < segmentIndex; i++ { usage.Usage += g.segments[i].TotalUsage } - return usage, nil + return usage } func (g *GrantBurnDownHistory) Segments() []GrantBurnDownHistorySegment { @@ -170,17 +251,23 @@ func (g GrantBurnDownHistory) ChunkByResets() []GrantBurnDownHistory { chunks := make([]GrantBurnDownHistory, 0, 1) current := GrantBurnDownHistory{ - usageAtStart: g.usageAtStart, - segments: make([]GrantBurnDownHistorySegment, 0, len(g.segments)), + startingSnapshot: g.startingSnapshot.Clone(), + segments: make([]GrantBurnDownHistorySegment, 0, len(g.segments)), } - for _, seg := range g.segments { + for i, seg := range g.segments { current.segments = append(current.segments, seg) if seg.TerminationReasons.UsageReset { chunks = append(chunks, current) + + var startingSnapshot balance.Snapshot + if i+1 < len(g.segments) { + startingSnapshot = g.getSnapshotAtStartOfSegment(i + 1) + } + current = GrantBurnDownHistory{ - usageAtStart: usageAtReset(seg.To), - segments: make([]GrantBurnDownHistorySegment, 0, len(g.segments)), + startingSnapshot: startingSnapshot, + segments: make([]GrantBurnDownHistorySegment, 0, len(g.segments)), } } } @@ -196,14 +283,13 @@ func (g GrantBurnDownHistory) TotalGrantUsage() alpacadecimal.Decimal { total := alpacadecimal.NewFromFloat(0) for _, seg := range g.segments { - for _, usage := range seg.GrantUsages { - total = total.Add(alpacadecimal.NewFromFloat(usage.Usage)) - } + total = total.Add(seg.GrantUsages.Sum()) } return total } +// usageAtReset creates the legacy usage representation for a reset boundary. func usageAtReset(at time.Time) balance.SnapshottedUsage { return balance.SnapshottedUsage{ Since: at, diff --git a/openmeter/credit/engine/history_test.go b/openmeter/credit/engine/history_test.go index 688fca6cb8..24beed66f4 100644 --- a/openmeter/credit/engine/history_test.go +++ b/openmeter/credit/engine/history_test.go @@ -4,11 +4,13 @@ import ( "testing" "time" + "github.com/alpacahq/alpacadecimal" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/openmeterio/openmeter/openmeter/credit/balance" "github.com/openmeterio/openmeter/openmeter/credit/engine" + "github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig" "github.com/openmeterio/openmeter/pkg/timeutil" ) @@ -18,6 +20,19 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { Since: start.Add(-time.Hour), Usage: 7, } + startingUnitConfig := &unitconfig.UnitConfig{ + Operation: unitconfig.UnitConfigOperationMultiply, + ConversionFactor: alpacadecimal.NewFromInt(2), + } + startingSnapshot := balance.Snapshot{ + At: start, + Usage: usageAtStart, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: usageAtStart.Usage, + TotalGrantUsage: 6, + }, + UnitConfig: startingUnitConfig, + } segment := func(idx int, grantUsage float64, reset bool) engine.GrantBurnDownHistorySegment { from := start.Add(time.Duration(idx) * time.Hour) @@ -41,7 +56,7 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { } t.Run("Empty history returns no chunks", func(t *testing.T) { - history, err := engine.NewGrantBurnDownHistory(nil, usageAtStart) + history, err := engine.NewGrantBurnDownHistory(nil, startingSnapshot) require.NoError(t, err) assert.Empty(t, history.ChunkByResets()) @@ -51,7 +66,7 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { history, err := engine.NewGrantBurnDownHistory([]engine.GrantBurnDownHistorySegment{ segment(0, 10, false), segment(1, 20, false), - }, usageAtStart) + }, startingSnapshot) require.NoError(t, err) chunks := history.ChunkByResets() @@ -59,9 +74,7 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { assert.Len(t, chunks[0].Segments(), 2) assert.Equal(t, 30.0, chunks[0].TotalGrantUsage().InexactFloat64()) - usage, err := chunks[0].GetUsageInPeriodUntilSegment(0) - require.NoError(t, err) - assert.Equal(t, usageAtStart, usage) + assertChunkSnapshotAtStart(t, chunks[0], usageAtStart, 6, startingUnitConfig) }) t.Run("History is chunked after reset segments", func(t *testing.T) { @@ -71,7 +84,7 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { segment(2, 30, false), segment(3, 40, true), segment(4, 50, false), - }, usageAtStart) + }, startingSnapshot) require.NoError(t, err) chunks := history.ChunkByResets() @@ -79,27 +92,27 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { assert.Len(t, chunks[0].Segments(), 2) assert.Equal(t, 30.0, chunks[0].TotalGrantUsage().InexactFloat64()) - assertChunkUsageAtStart(t, chunks[0], usageAtStart) + assertChunkSnapshotAtStart(t, chunks[0], usageAtStart, 6, startingUnitConfig) assert.Len(t, chunks[1].Segments(), 2) assert.Equal(t, 70.0, chunks[1].TotalGrantUsage().InexactFloat64()) - assertChunkUsageAtStart(t, chunks[1], balance.SnapshottedUsage{ + assertChunkSnapshotAtStart(t, chunks[1], balance.SnapshottedUsage{ Since: start.Add(2 * time.Hour), Usage: 0, - }) + }, 0, startingUnitConfig) assert.Len(t, chunks[2].Segments(), 1) assert.Equal(t, 50.0, chunks[2].TotalGrantUsage().InexactFloat64()) - assertChunkUsageAtStart(t, chunks[2], balance.SnapshottedUsage{ + assertChunkSnapshotAtStart(t, chunks[2], balance.SnapshottedUsage{ Since: start.Add(4 * time.Hour), Usage: 0, - }) + }, 0, startingUnitConfig) }) t.Run("Final reset does not create empty trailing chunk", func(t *testing.T) { history, err := engine.NewGrantBurnDownHistory([]engine.GrantBurnDownHistorySegment{ segment(0, 10, true), - }, usageAtStart) + }, startingSnapshot) require.NoError(t, err) chunks := history.ChunkByResets() @@ -107,12 +120,210 @@ func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) { assert.Len(t, chunks[0].Segments(), 1) assert.Equal(t, 10.0, chunks[0].TotalGrantUsage().InexactFloat64()) }) + + t.Run("Rollover starts the new usage period chunk", func(t *testing.T) { + resetAt := start.Add(time.Hour) + beforeReset := segment(0, 10, true) + rollover := engine.GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{ + From: resetAt, + To: resetAt, + }, + TerminationReasons: engine.SegmentTerminationReason{ + Rollover: true, + }, + GrantUsages: []engine.GrantUsage{ + { + GrantID: "grant-1", + Usage: 5, + }, + }, + } + afterReset := segment(1, 20, false) + + history, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{beforeReset, rollover, afterReset}, + startingSnapshot, + ) + require.NoError(t, err) + + chunks := history.ChunkByResets() + require.Len(t, chunks, 2) + require.Len(t, chunks[1].Segments(), 2) + assert.True(t, chunks[1].Segments()[0].TerminationReasons.Rollover) + assert.Equal(t, 25.0, chunks[1].TotalGrantUsage().InexactFloat64()) + assertChunkSnapshotAtStart(t, chunks[1], balance.SnapshottedUsage{ + Since: resetAt, + Usage: 0, + }, 0, startingUnitConfig) + + snapshotAfterRollover, err := history.GetSnapshotAtStartOfSegment(2) + require.NoError(t, err) + require.NotNil(t, snapshotAfterRollover.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + TotalGrantUsage: 5, + }, *snapshotAfterRollover.UsageSnapshot) + }) +} + +func TestGrantBurnDownHistory_RolloverPrecedesUsageAtSameTime(t *testing.T) { + at := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + usageSegment := engine.GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{ + From: at, + To: at.Add(time.Hour), + }, + } + rolloverSegment := engine.GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{ + From: at, + To: at, + }, + TerminationReasons: engine.SegmentTerminationReason{ + Rollover: true, + }, + } + + history, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{usageSegment, rolloverSegment}, + balance.Snapshot{At: at, UsageSnapshot: zeroUsageSnapshot()}, + ) + require.NoError(t, err) + require.Len(t, history.Segments(), 2) + assert.True(t, history.Segments()[0].TerminationReasons.Rollover) + assert.False(t, history.Segments()[1].TerminationReasons.Rollover) +} + +func TestNewGrantBurnDownHistory_ValidatesAnchorAndContinuity(t *testing.T) { + at := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + startingSnapshot := balance.Snapshot{ + At: at, + Balances: balance.Map{"grant-1": 100}, + Overage: 5, + UsageSnapshot: zeroUsageSnapshot(), + } + first := engine.GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{ + From: at, + To: at.Add(time.Hour), + }, + // Boundary changes such as recurrence can make this differ from the + // starting snapshot without breaking the history anchor. + BalanceAtStart: balance.Map{"grant-1": 200}, + OverageAtStart: 5, + } + second := engine.GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{ + From: at.Add(time.Hour), + To: at.Add(2 * time.Hour), + }, + } + + t.Run("aligned continuous history", func(t *testing.T) { + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{second, first}, + startingSnapshot, + ) + require.NoError(t, err) + }) + + t.Run("empty history", func(t *testing.T) { + _, err := engine.NewGrantBurnDownHistory(nil, startingSnapshot) + require.NoError(t, err) + }) + + t.Run("missing usage snapshot", func(t *testing.T) { + invalid := startingSnapshot + invalid.UsageSnapshot = nil + + _, err := engine.NewGrantBurnDownHistory(nil, invalid) + require.ErrorContains(t, err, "usage snapshot is missing") + }) + + t.Run("first segment does not start at snapshot", func(t *testing.T) { + misaligned := first + misaligned.ClosedPeriod = timeutil.ClosedPeriod{ + From: at.Add(time.Minute), + To: first.To, + } + + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{misaligned}, + startingSnapshot, + ) + require.ErrorContains(t, err, "expected starting snapshot") + }) + + t.Run("first segment overage does not match snapshot", func(t *testing.T) { + misaligned := first + misaligned.OverageAtStart = 6 + + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{misaligned}, + startingSnapshot, + ) + require.ErrorContains(t, err, "expected starting snapshot overage") + }) + + t.Run("segment starts after it ends", func(t *testing.T) { + invalid := first + invalid.ClosedPeriod = timeutil.ClosedPeriod{ + From: at, + To: at.Add(-time.Minute), + } + + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{invalid}, + startingSnapshot, + ) + require.ErrorContains(t, err, "starts after it ends") + }) + + t.Run("segments have a gap", func(t *testing.T) { + gapped := second + gapped.ClosedPeriod = timeutil.ClosedPeriod{ + From: second.From.Add(time.Minute), + To: second.To, + } + + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{first, gapped}, + startingSnapshot, + ) + require.ErrorContains(t, err, "are not contiguous") + }) + + t.Run("segments overlap", func(t *testing.T) { + overlapping := second + overlapping.ClosedPeriod = timeutil.ClosedPeriod{ + From: second.From.Add(-time.Minute), + To: second.To, + } + + _, err := engine.NewGrantBurnDownHistory( + []engine.GrantBurnDownHistorySegment{first, overlapping}, + startingSnapshot, + ) + require.ErrorContains(t, err, "are not contiguous") + }) } -func assertChunkUsageAtStart(t *testing.T, history engine.GrantBurnDownHistory, expected balance.SnapshottedUsage) { +func assertChunkSnapshotAtStart( + t *testing.T, + history engine.GrantBurnDownHistory, + expectedUsage balance.SnapshottedUsage, + expectedTotalGrantUsage float64, + expectedUnitConfig *unitconfig.UnitConfig, +) { t.Helper() - usage, err := history.GetUsageInPeriodUntilSegment(0) + snapshot, err := history.GetSnapshotAtStartOfSegment(0) require.NoError(t, err) - assert.Equal(t, expected, usage) + assert.Equal(t, expectedUsage, snapshot.Usage) + require.NotNil(t, snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: expectedUsage.Usage, + TotalGrantUsage: expectedTotalGrantUsage, + }, *snapshot.UsageSnapshot) + assert.True(t, expectedUnitConfig.Equal(snapshot.UnitConfig)) } diff --git a/openmeter/credit/engine/reset.go b/openmeter/credit/engine/reset.go index 6e0d3859a1..1a62962fd2 100644 --- a/openmeter/credit/engine/reset.go +++ b/openmeter/credit/engine/reset.go @@ -5,13 +5,16 @@ import ( "slices" "time" + "github.com/samber/lo" + "github.com/openmeterio/openmeter/openmeter/credit/balance" "github.com/openmeterio/openmeter/openmeter/credit/grant" + "github.com/openmeterio/openmeter/pkg/timeutil" ) // reset rolls over the grants and burns down the overage if needed. // It returns the new snapshot of the balances at the start of the next period. -func (e *engine) reset(grants []grant.Grant, snap balance.Snapshot, behavior grant.ResetBehavior, at time.Time) (balance.Snapshot, error) { +func (e *engine) reset(grants []grant.Grant, snap balance.Snapshot, behavior grant.ResetBehavior, at time.Time) (balance.Snapshot, GrantBurnDownHistorySegment, error) { // Let's build a grantMap from our grants for easier lookup grantMap := make(map[string]grant.Grant) for _, g := range grants { @@ -25,7 +28,7 @@ func (e *engine) reset(grants []grant.Grant, snap balance.Snapshot, behavior gra grant, ok := grantMap[grantID] // Inconsistency check, should never happen if !ok { - return balance.Snapshot{}, fmt.Errorf("grant %s not found", grantID) + return balance.Snapshot{}, GrantBurnDownHistorySegment{}, fmt.Errorf("grant %s not found", grantID) } // grants might become inactive at the reset time, in which case they're irrelevant for the next period @@ -44,14 +47,51 @@ func (e *engine) reset(grants []grant.Grant, snap balance.Snapshot, behavior gra prioritizedGrants := slices.Clone(grants) if err := PrioritizeGrants(prioritizedGrants); err != nil { - return balance.Snapshot{}, fmt.Errorf("failed to prioritize grants: %w", err) + return balance.Snapshot{}, GrantBurnDownHistorySegment{}, fmt.Errorf("failed to prioritize grants: %w", err) + } + + balances := rolledOver + overage := startingOverage + var grantUsages GrantUsages + if startingOverage != 0 { + balances, grantUsages, overage = e.burnDownGrants(rolledOver, prioritizedGrants, startingOverage) + } + + unitConfig := snap.UnitConfig + if unitConfig != nil { + unitConfig = lo.ToPtr(unitConfig.Clone()) + } + + // The reset snapshot is the point-in-time balance after grant balance + // rollover and preserved overage burn. + resetSnapshot := balance.Snapshot{ + At: at, + Balances: balances, + Overage: overage, + UnitConfig: unitConfig, + Usage: balance.SnapshottedUsage{ + Since: at, + Usage: 0, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 0, + TotalGrantUsage: grantUsages.Sum().InexactFloat64(), + }, } - rolledOver, _, startingOverage = e.burnDownGrants(rolledOver, prioritizedGrants, startingOverage) + // The rollover segment captures the instantaneous transition from rolled-over + // balances to the reset snapshot. + rolloverSegment := GrantBurnDownHistorySegment{ + ClosedPeriod: timeutil.ClosedPeriod{From: at, To: at}, + BalanceAtStart: rolledOver, + TerminationReasons: SegmentTerminationReason{ + Rollover: true, + }, + TotalUsage: 0, + OverageAtStart: startingOverage, + Overage: overage, + GrantUsages: grantUsages, + } - return balance.Snapshot{ - At: at, - Balances: rolledOver, - Overage: startingOverage, - }, nil + return resetSnapshot, rolloverSegment, nil } diff --git a/openmeter/credit/engine/reset_test.go b/openmeter/credit/engine/reset_test.go index 052558fde4..8067ac9333 100644 --- a/openmeter/credit/engine/reset_test.go +++ b/openmeter/credit/engine/reset_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/openmeterio/openmeter/openmeter/credit/balance" "github.com/openmeterio/openmeter/openmeter/credit/engine" @@ -79,6 +80,7 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Usage: balance.SnapshottedUsage{ Since: t1.AddDate(0, 0, -1), // Last "reset time", outside this period, arbitrary Usage: 0.0, @@ -107,8 +109,9 @@ func TestReset(t *testing.T) { assert.Equal(t, 0.0, res.Snapshot.Usage.Usage) // 0 usage after 5h mark assert.Equal(t, t1.Add(time.Hour*5), res.Snapshot.Usage.Since) // should mark since the last reset time - // History should have 2 segments, one before and one after the reset - assert.Equal(t, 2, len(res.History.Segments())) + // History should have a usage segment before the reset, the rollover + // transition, and a usage segment after the reset. + assert.Equal(t, 3, len(res.History.Segments())) // The first segment should have a balance of 100 with 10 usage assert.Equal(t, 100.0, res.History.Segments()[0].BalanceAtStart.Balance()) @@ -118,10 +121,21 @@ func TestReset(t *testing.T) { // It should end with a reset assert.True(t, res.History.Segments()[0].TerminationReasons.UsageReset) - // The second segment should have a balance of 50 with no usage + // The rollover transition applies no overage. + assert.True(t, res.History.Segments()[1].TerminationReasons.Rollover) + assert.Equal(t, t1.Add(time.Hour*5), res.History.Segments()[1].From) + assert.Equal(t, t1.Add(time.Hour*5), res.History.Segments()[1].To) assert.Equal(t, 50.0, res.History.Segments()[1].BalanceAtStart.Balance()) assert.Equal(t, 0.0, res.History.Segments()[1].OverageAtStart) assert.Equal(t, 0.0, res.History.Segments()[1].TotalUsage) + assert.Empty(t, res.History.Segments()[1].GrantUsages) + assert.Equal(t, 50.0, res.History.Segments()[1].ApplyUsage().Balance()) + + // The final segment starts after grant balance rollover and overage burn. + assert.False(t, res.History.Segments()[2].TerminationReasons.Rollover) + assert.Equal(t, 50.0, res.History.Segments()[2].BalanceAtStart.Balance()) + assert.Equal(t, 0.0, res.History.Segments()[2].OverageAtStart) + assert.Equal(t, 0.0, res.History.Segments()[2].TotalUsage) }) t.Run("Should carry over overage to next period", func(t *testing.T) { @@ -140,6 +154,7 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -157,9 +172,11 @@ func TestReset(t *testing.T) { // The grant should be rolled over: assert.Equal(t, 40.0, res.Snapshot.Balances[grant1.ID]) + assert.Equal(t, 50.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) - // History should have 2 segments, one before and one after the reset - assert.Equal(t, 2, len(res.History.Segments())) + // History should have a usage segment before the reset, the rollover + // transition, and a usage segment after the reset. + assert.Equal(t, 3, len(res.History.Segments())) // The first segment should have a balance of 0 with 10 overage assert.Equal(t, 100.0, res.History.Segments()[0].BalanceAtStart.Balance()) @@ -170,10 +187,29 @@ func TestReset(t *testing.T) { // It should end with a reset assert.True(t, res.History.Segments()[0].TerminationReasons.UsageReset) - // The second segment should have a balance of 50 - 10 with no usage (minRolloverAmount + overage) - assert.Equal(t, 40.0, res.History.Segments()[1].BalanceAtStart.Balance()) - assert.Equal(t, 0.0, res.History.Segments()[1].OverageAtStart) + // The rollover transition starts with the rolled balance and burns + // overage preserved from the previous usage period. + assert.True(t, res.History.Segments()[1].TerminationReasons.Rollover) + assert.Equal(t, t1.Add(time.Hour*5), res.History.Segments()[1].From) + assert.Equal(t, t1.Add(time.Hour*5), res.History.Segments()[1].To) + assert.Equal(t, 50.0, res.History.Segments()[1].BalanceAtStart.Balance()) + assert.Equal(t, 10.0, res.History.Segments()[1].OverageAtStart) assert.Equal(t, 0.0, res.History.Segments()[1].TotalUsage) + assert.Equal(t, 0.0, res.History.Segments()[1].Overage) + assert.Equal(t, engine.GrantUsages{ + { + GrantID: g1.ID, + Usage: 10.0, + TerminationReason: engine.GrantUsageTerminationReasonSegmentTermination, + }, + }, res.History.Segments()[1].GrantUsages) + assert.Equal(t, 40.0, res.History.Segments()[1].ApplyUsage().Balance()) + + // The final segment retains its existing post-reset semantics. + assert.False(t, res.History.Segments()[2].TerminationReasons.Rollover) + assert.Equal(t, 40.0, res.History.Segments()[2].BalanceAtStart.Balance()) + assert.Equal(t, 0.0, res.History.Segments()[2].OverageAtStart) + assert.Equal(t, 0.0, res.History.Segments()[2].TotalUsage) }) t.Run("No reset", func(t *testing.T) { @@ -194,6 +230,9 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{grant1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: &balance.UsageSnapshot{ + Usage: u.Usage, + }, Usage: u, Balances: balance.Map{ grant1.ID: 100.0, @@ -210,6 +249,8 @@ func TestReset(t *testing.T) { // If there was no reset, should extend the starting snapshot with the current usage data assert.Equal(t, 20.0, res.Snapshot.Usage.Usage) // 10 + 10 assert.Equal(t, u.Since, res.Snapshot.Usage.Since) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 20.0, res.Snapshot.UsageSnapshot.Usage) // Should have 2 periods, start - g2, g2 - end assert.Equal(t, 2, len(res.History.Segments())) @@ -233,6 +274,7 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -249,11 +291,13 @@ func TestReset(t *testing.T) { assert.Equal(t, 0.0, res.Snapshot.Usage.Usage) assert.Equal(t, resetTime, res.Snapshot.Usage.Since) - // Should have 2 periods, start - reset, reset - end where reset = end, 2nd period is 0 length - assert.Equal(t, 2, len(res.History.Segments()), "expected: %+v, got %+v, history: %+v", 2, len(res.History.Segments()), res.History.Segments()) + // The reset is represented by the ending usage segment, a rollover + // transition, and the existing zero-length new-period usage segment. + assert.Equal(t, 3, len(res.History.Segments()), "expected: %+v, got %+v, history: %+v", 3, len(res.History.Segments()), res.History.Segments()) assert.True(t, res.History.Segments()[0].TerminationReasons.UsageReset) - assert.False(t, res.History.Segments()[1].TerminationReasons.UsageReset) + assert.True(t, res.History.Segments()[1].TerminationReasons.Rollover) + assert.False(t, res.History.Segments()[2].TerminationReasons.UsageReset) }) t.Run("Should include grant recurrence in starting balance", func(t *testing.T) { @@ -274,6 +318,7 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -290,12 +335,14 @@ func TestReset(t *testing.T) { assert.Equal(t, 0.0, res.Snapshot.Usage.Usage) assert.Equal(t, resetTime, res.Snapshot.Usage.Since) - // Should have 2 periods, start - reset, reset - end where reset = end, 2nd period is 0 length - assert.Equal(t, 2, len(res.History.Segments()), "expected: %+v, got %+v, history: %+v", 2, len(res.History.Segments()), res.History.Segments()) + // The reset is represented by the ending usage segment, a rollover + // transition, and the existing zero-length new-period usage segment. + assert.Equal(t, 3, len(res.History.Segments()), "expected: %+v, got %+v, history: %+v", 3, len(res.History.Segments()), res.History.Segments()) assert.True(t, res.History.Segments()[0].TerminationReasons.UsageReset) - assert.False(t, res.History.Segments()[1].TerminationReasons.UsageReset) - assert.Equal(t, 100.0, res.History.Segments()[1].BalanceAtStart[g1.ID]) + assert.True(t, res.History.Segments()[1].TerminationReasons.Rollover) + assert.False(t, res.History.Segments()[2].TerminationReasons.UsageReset) + assert.Equal(t, 100.0, res.History.Segments()[2].BalanceAtStart[g1.ID]) // The starting balance should be the amount of the grant assert.Equal(t, 100.0, res.Snapshot.Balances[g1.ID]) @@ -314,9 +361,10 @@ func TestReset(t *testing.T) { Meter: meter, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ - Balances: balance.Map{g1.ID: 100.0}, - Overage: 0, - At: t1, + UsageSnapshot: zeroUsageSnapshot(), + Balances: balance.Map{g1.ID: 100.0}, + Overage: 0, + At: t1, }, Until: t1.AddDate(0, 0, 1), Resets: timeutil.NewSimpleTimeline([]time.Time{t1}), diff --git a/openmeter/credit/engine/run.go b/openmeter/credit/engine/run.go index 1d06fdaa71..9498cfaf99 100644 --- a/openmeter/credit/engine/run.go +++ b/openmeter/credit/engine/run.go @@ -15,6 +15,10 @@ import ( ) func (e *engine) Run(ctx context.Context, params RunParams) (RunResult, error) { + if err := validateStartingSnapshot(params.StartingSnapshot); err != nil { + return RunResult{}, err + } + resParams := params.Clone() // Let's build the timeline @@ -47,14 +51,6 @@ func (e *engine) Run(ctx context.Context, params RunParams) (RunResult, error) { periods := timeline.GetClosedPeriods() for idx, period := range periods { - // Let's reset the snapshot usage information as we're entering a new period (between resets) - if idx > 0 { - snapshot.Usage = balance.SnapshottedUsage{ - Since: period.From, - Usage: 0.0, - } - } - // We need to find the grants that are relevant for this period. // We do this filtering so that history isn't polluted with grants that are irrelevant. relevantGrants := e.filterRelevantGrants(params.Grants, snapshot.Balances, period) @@ -76,7 +72,7 @@ func (e *engine) Run(ctx context.Context, params RunParams) (RunResult, error) { if idx != len(periods)-1 { // We need to reset at each period, except the last one. // If the ending time is also a reset, there will be a 0 length period at the end. - snap, err := e.reset(relevantGrants, runRes.Snapshot, params.ResetBehavior, period.To) + snap, rolloverSegment, err := e.reset(relevantGrants, runRes.Snapshot, params.ResetBehavior, period.To) if err != nil { return RunResult{}, fmt.Errorf("failed to reset at end of period %s - %s: %w", period.From, period.To, err) } @@ -87,10 +83,12 @@ func (e *engine) Run(ctx context.Context, params RunParams) (RunResult, error) { if len(historySegments) > 0 { historySegments[len(historySegments)-1].TerminationReasons.UsageReset = true } + + historySegments = append(historySegments, rolloverSegment) } } - history, err := NewGrantBurnDownHistory(historySegments, params.StartingSnapshot.Usage) + history, err := NewGrantBurnDownHistory(historySegments, params.StartingSnapshot) if err != nil { return RunResult{}, fmt.Errorf("failed to create grant burn down history: %w", err) } @@ -250,20 +248,35 @@ func (e *engine) runBetweenResets(ctx context.Context, params inbetweenRunParams } } - history, err := NewGrantBurnDownHistory(segments, params.StartingSnapshot.Usage) + history, err := NewGrantBurnDownHistory(segments, params.StartingSnapshot) if err != nil { return RunResult{}, fmt.Errorf("failed to create grant burn down history: %w", err) } + totalUsage := history.TotalUsageInHistory() + usage := balance.SnapshottedUsage{ + Since: params.StartingSnapshot.Usage.Since, + Usage: params.StartingSnapshot.Usage.Usage + totalUsage, + } + usageSnapshot := balance.UsageSnapshot{ + Usage: params.StartingSnapshot.UsageSnapshot.Usage + totalUsage, + TotalGrantUsage: alpacadecimal.NewFromFloat(params.StartingSnapshot.UsageSnapshot.TotalGrantUsage). + Add(history.TotalGrantUsage()). + InexactFloat64(), + } + unitConfig := params.StartingSnapshot.UnitConfig + if unitConfig != nil { + unitConfig = lo.ToPtr(unitConfig.Clone()) + } + return RunResult{ Snapshot: balance.Snapshot{ - Balances: balancesAtPhaseStart, - Overage: overage, - At: period.To, - Usage: balance.SnapshottedUsage{ - Since: params.StartingSnapshot.Usage.Since, - Usage: params.StartingSnapshot.Usage.Usage + history.TotalUsageInHistory(), - }, + Usage: usage, + UsageSnapshot: &usageSnapshot, + Balances: balancesAtPhaseStart, + Overage: overage, + At: period.To, + UnitConfig: unitConfig, }, History: history, }, nil @@ -272,9 +285,9 @@ func (e *engine) runBetweenResets(ctx context.Context, params inbetweenRunParams // Burns down the grants of the priority sorted list. Manages overage. // // FIXME: calculations happen on inexact representations as float64, this can lead to rounding errors. -func (m *engine) burnDownGrants(startingBalances balance.Map, prioritized []grant.Grant, usage float64) (balance.Map, []GrantUsage, float64) { +func (m *engine) burnDownGrants(startingBalances balance.Map, prioritized []grant.Grant, usage float64) (balance.Map, GrantUsages, float64) { balances := startingBalances.Clone() - uses := make([]GrantUsage, 0, len(prioritized)) + uses := make(GrantUsages, 0, len(prioritized)) exactUsage := alpacadecimal.NewFromFloat(usage) getFloat := func(d alpacadecimal.Decimal) float64 { diff --git a/openmeter/credit/engine/run_test.go b/openmeter/credit/engine/run_test.go index 644495ce30..0e27f5c236 100644 --- a/openmeter/credit/engine/run_test.go +++ b/openmeter/credit/engine/run_test.go @@ -72,6 +72,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -88,6 +89,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -111,9 +113,10 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{}, StartingSnapshot: balance.Snapshot{ - Balances: balance.Map{}, - Overage: 0, - At: t1, + UsageSnapshot: zeroUsageSnapshot(), + Balances: balance.Map{}, + Overage: 0, + At: t1, }, Until: t1.AddDate(0, 0, 30), }) @@ -146,6 +149,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ grant1.ID: 100.0, }, @@ -172,6 +176,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ grant1.ID: 100.0, }, @@ -196,6 +201,9 @@ func TestEngine(t *testing.T) { Since: prevPeriodStart, Usage: 10.0, } + usageSnapshot := &balance.UsageSnapshot{ + Usage: u.Usage, + } res, err := eng.Run( context.Background(), @@ -203,7 +211,8 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{grant1}, StartingSnapshot: balance.Snapshot{ - Usage: u, + UsageSnapshot: usageSnapshot, + Usage: u, Balances: balance.Map{ grant1.ID: 100.0, }, @@ -215,7 +224,8 @@ func TestEngine(t *testing.T) { ) assert.NoError(t, err) assert.Equal(t, balance.Snapshot{ - Usage: u, // Should pass through the original usage info + Usage: u, // Should pass through the original usage info + UsageSnapshot: usageSnapshot, Balances: balance.Map{ grant1.ID: 100.0, }, @@ -234,6 +244,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{grant1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ grant1.ID: 100.0, }, @@ -261,6 +272,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -289,6 +301,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -319,6 +332,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -354,6 +368,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -388,6 +403,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -417,6 +433,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -449,6 +466,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 0.0, }, @@ -490,6 +508,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, g2.ID: 100.0, @@ -526,6 +545,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g2, g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, g2.ID: 100.0, @@ -562,6 +582,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g2, g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, g2.ID: 100.0, @@ -624,9 +645,10 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: grants, StartingSnapshot: balance.Snapshot{ - Balances: bm, - Overage: 0, - At: t1, + UsageSnapshot: zeroUsageSnapshot(), + Balances: bm, + Overage: 0, + At: t1, }, Until: t1.AddDate(0, 0, 1), }) @@ -660,6 +682,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, }, @@ -708,6 +731,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 80.0, // due to use before start g2.ID: 100.0, @@ -759,6 +783,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 80.0, // due to use before start g2.ID: 100.0, @@ -807,6 +832,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 100.0, g2.ID: 0.0, @@ -868,6 +894,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 80.0, // due to use before start g2.ID: 100.0, @@ -936,6 +963,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2, g3}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 80.0, // due to use before start g2.ID: 100.0, @@ -1036,6 +1064,7 @@ func TestEngine(t *testing.T) { Meter: mm, Grants: []grant.Grant{g1, g2, g3, g4}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g1.ID: 50.0, g2.ID: 100.0, @@ -1056,6 +1085,15 @@ func TestEngine(t *testing.T) { // 175 (active total at end) - 10 (last usage value) = 165 assert.Equal(t, 165.0, res.Snapshot.Balance(), "received following result %s", string(resJSON)) + // Preserve the existing usage-period accounting semantics: usage and + // grant usage accumulate across engine phases, including for LATEST. + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: 25, + TotalGrantUsage: 25, + }, *res.Snapshot.UsageSnapshot) + assert.Equal(t, 190.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) + // Now let's assert the history // We should have 2 segments: start -> 1h, 1h -> end assert.Equal(t, 2, len(res.History.Segments()), "received following result %s", string(resJSON)) diff --git a/openmeter/credit/engine/runresult_test.go b/openmeter/credit/engine/runresult_test.go index 2dd8398733..11e38877c2 100644 --- a/openmeter/credit/engine/runresult_test.go +++ b/openmeter/credit/engine/runresult_test.go @@ -19,7 +19,7 @@ import ( "github.com/openmeterio/openmeter/pkg/timeutil" ) -func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { +func TestRunResult_SnapshotUsage(t *testing.T) { t1, err := time.Parse(time.RFC3339, "2024-01-01T00:00:00Z") assert.NoError(t, err) meterSlug := "meter-1" @@ -78,6 +78,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -89,7 +90,9 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 100.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 0.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + assert.Equal(t, 100.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) }, }, { @@ -113,6 +116,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -124,7 +128,9 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 100.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 10.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + assert.Equal(t, 100.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) }, }, { @@ -139,16 +145,19 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{}, StartingSnapshot: balance.Snapshot{ - Balances: balance.Map{}, - Overage: 0, - At: t1, + UsageSnapshot: zeroUsageSnapshot(), + Balances: balance.Map{}, + Overage: 0, + At: t1, }, Until: t1.AddDate(0, 0, 1), }, ) require.NoError(t, err) - assert.Equal(t, 0.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 0.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + assert.Equal(t, 0.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) }, }, { @@ -177,6 +186,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -189,7 +199,9 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 0.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 0.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + assert.Equal(t, 0.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) }, }, { @@ -216,6 +228,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -227,9 +240,12 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 10.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 10.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + totalAvailableGrantAmount := res.Snapshot.UsageSnapshot.TotalGrantUsage + res.Snapshot.Balance() + assert.Equal(t, 10.0, totalAvailableGrantAmount) // When we have overage this holds true - assert.Equal(t, res.TotalAvailableGrantAmountAtLastPeriod(), res.Snapshot.Balance()+res.Snapshot.Usage.Usage-res.Snapshot.Overage, "balance %s, usage %s, overage %s", res.Snapshot.Balance(), res.Snapshot.Usage.Usage, res.Snapshot.Overage) + assert.Equal(t, totalAvailableGrantAmount, res.Snapshot.Balance()+res.Snapshot.Usage.Usage-res.Snapshot.Overage, "balance %s, usage %s, overage %s", res.Snapshot.Balance(), res.Snapshot.Usage.Usage, res.Snapshot.Overage) }, }, { @@ -261,6 +277,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -272,9 +289,12 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 370.0, res1.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res1.Snapshot.UsageSnapshot) + assert.Equal(t, 270.0, res1.Snapshot.UsageSnapshot.TotalGrantUsage) + totalAvailableGrantAmount := res1.Snapshot.UsageSnapshot.TotalGrantUsage + res1.Snapshot.Balance() + assert.Equal(t, 370.0, totalAvailableGrantAmount) // Should be true cause all usage was covered - assert.Equal(t, res1.TotalAvailableGrantAmountAtLastPeriod(), res1.Snapshot.Balance()+res1.Snapshot.Usage.Usage, "balance %s, usage %s", res1.Snapshot.Balance(), res1.Snapshot.Usage.Usage) + assert.Equal(t, totalAvailableGrantAmount, res1.Snapshot.Balance()+res1.Snapshot.Usage.Usage, "balance %s, usage %s", res1.Snapshot.Balance(), res1.Snapshot.Usage.Usage) res2, err := eng.Run( t.Context(), @@ -282,6 +302,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -294,9 +315,12 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - assert.Equal(t, 370.0, res2.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res2.Snapshot.UsageSnapshot) + assert.Equal(t, 270.0, res2.Snapshot.UsageSnapshot.TotalGrantUsage) + totalAvailableGrantAmount = res2.Snapshot.UsageSnapshot.TotalGrantUsage + res2.Snapshot.Balance() + assert.Equal(t, 370.0, totalAvailableGrantAmount) // Should be true cause all usage was covered - assert.Equal(t, res2.TotalAvailableGrantAmountAtLastPeriod(), res2.Snapshot.Balance()+res2.Snapshot.Usage.Usage, "balance %s, usage %s", res2.Snapshot.Balance(), res2.Snapshot.Usage.Usage) + assert.Equal(t, totalAvailableGrantAmount, res2.Snapshot.Balance()+res2.Snapshot.Usage.Usage, "balance %s, usage %s", res2.Snapshot.Balance(), res2.Snapshot.Usage.Usage) res3, err := eng.Run( t.Context(), @@ -304,6 +328,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -316,10 +341,13 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) + require.NotNil(t, res3.Snapshot.UsageSnapshot) + assert.Equal(t, 270.0, res3.Snapshot.UsageSnapshot.TotalGrantUsage) // So we have 270 used + 20 still available in the period - assert.Equal(t, 290.0, res3.TotalAvailableGrantAmountAtLastPeriod()) + totalAvailableGrantAmount = res3.Snapshot.UsageSnapshot.TotalGrantUsage + res3.Snapshot.Balance() + assert.Equal(t, 290.0, totalAvailableGrantAmount) // Should be true cause all usage was covered - assert.Equal(t, res3.TotalAvailableGrantAmountAtLastPeriod(), res3.Snapshot.Balance()+res3.Snapshot.Usage.Usage, "balance %s, usage %s", res3.Snapshot.Balance(), res3.Snapshot.Usage.Usage) + assert.Equal(t, totalAvailableGrantAmount, res3.Snapshot.Balance()+res3.Snapshot.Usage.Usage, "balance %s, usage %s", res3.Snapshot.Balance(), res3.Snapshot.Usage.Usage) res4, err := eng.Run( t.Context(), @@ -327,6 +355,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -338,10 +367,13 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { ) require.NoError(t, err) - // Going an extra day further it's still 280 used up + fresh 100 available as we're on the boundary - assert.Equal(t, 370.0, res4.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res4.Snapshot.UsageSnapshot) + assert.Equal(t, 270.0, res4.Snapshot.UsageSnapshot.TotalGrantUsage) + // Going an extra day further it's still 270 used up + fresh 100 available as we're on the boundary + totalAvailableGrantAmount = res4.Snapshot.UsageSnapshot.TotalGrantUsage + res4.Snapshot.Balance() + assert.Equal(t, 370.0, totalAvailableGrantAmount) // Should be true cause all usage was covered - assert.Equal(t, res4.TotalAvailableGrantAmountAtLastPeriod(), res4.Snapshot.Balance()+res4.Snapshot.Usage.Usage, "balance %s, usage %s", res4.Snapshot.Balance(), res4.Snapshot.Usage.Usage) + assert.Equal(t, totalAvailableGrantAmount, res4.Snapshot.Balance()+res4.Snapshot.Usage.Usage, "balance %s, usage %s", res4.Snapshot.Balance(), res4.Snapshot.Usage.Usage) }, }, { @@ -365,6 +397,7 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { Meter: mm, Grants: []grant.Grant{g}, StartingSnapshot: balance.Snapshot{ + UsageSnapshot: zeroUsageSnapshot(), Balances: balance.Map{ g.ID: 100.0, }, @@ -377,7 +410,9 @@ func TestRunResult_TotalAvailableGrantAmount(t *testing.T) { require.NoError(t, err) // We had 100 total grants in current period - assert.Equal(t, 100.0, res.TotalAvailableGrantAmountAtLastPeriod()) + require.NotNil(t, res.Snapshot.UsageSnapshot) + assert.Equal(t, 100.0, res.Snapshot.UsageSnapshot.TotalGrantUsage) + assert.Equal(t, 100.0, res.Snapshot.UsageSnapshot.TotalGrantUsage+res.Snapshot.Balance()) assert.Equal(t, 0.0, res.Snapshot.Balance()) assert.Equal(t, 0.0, res.Snapshot.Usage.Usage) assert.Equal(t, 0.0, res.Snapshot.Overage) diff --git a/openmeter/credit/engine/snapshot_test.go b/openmeter/credit/engine/snapshot_test.go new file mode 100644 index 0000000000..8c95a49e62 --- /dev/null +++ b/openmeter/credit/engine/snapshot_test.go @@ -0,0 +1,355 @@ +package engine_test + +import ( + "context" + "math" + "testing" + "time" + + "github.com/alpacahq/alpacadecimal" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/openmeterio/openmeter/openmeter/credit/balance" + "github.com/openmeterio/openmeter/openmeter/credit/engine" + "github.com/openmeterio/openmeter/openmeter/credit/grant" + "github.com/openmeterio/openmeter/openmeter/meter" + "github.com/openmeterio/openmeter/openmeter/productcatalog/unitconfig" + "github.com/openmeterio/openmeter/pkg/timeutil" +) + +func zeroUsageSnapshot() *balance.UsageSnapshot { + return lo.ToPtr(balance.UsageSnapshot{}) +} + +func TestEngineValidatesStartingUsageSnapshot(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + + tests := []struct { + name string + usageSnapshot *balance.UsageSnapshot + valid bool + }{ + { + name: "missing", + }, + { + name: "negative total grant usage", + usageSnapshot: &balance.UsageSnapshot{ + TotalGrantUsage: -1, + }, + }, + { + name: "total grant usage is not a number", + usageSnapshot: &balance.UsageSnapshot{ + TotalGrantUsage: math.NaN(), + }, + }, + { + name: "total grant usage is infinite", + usageSnapshot: &balance.UsageSnapshot{ + TotalGrantUsage: math.Inf(1), + }, + }, + { + name: "explicit zero", + usageSnapshot: &balance.UsageSnapshot{}, + valid: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + usageQueried := false + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, _, _ time.Time) (float64, error) { + usageQueried = true + return 0, nil + }, + }) + + _, err := eng.Run(t.Context(), engine.RunParams{ + StartingSnapshot: balance.Snapshot{ + At: start, + Balances: balance.Map{}, + UsageSnapshot: tt.usageSnapshot, + }, + Until: start.Add(time.Hour), + }) + + if tt.valid { + require.NoError(t, err) + assert.True(t, usageQueried) + return + } + + require.ErrorContains(t, err, "starting snapshot") + assert.False(t, usageQueried) + }) + } +} + +func TestEngineAccumulatesTotalGrantUsageFromSnapshot(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + startingTotalGrantUsage := 25.0 + g := grant.Grant{ + ID: "grant-1", + Amount: 100, + Priority: 1, + EffectiveAt: start, + } + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, _, _ time.Time) (float64, error) { + return 10, nil + }, + }) + + result, err := eng.Run(t.Context(), engine.RunParams{ + Meter: meter.Meter{ + Aggregation: meter.MeterAggregationSum, + }, + Grants: []grant.Grant{g}, + StartingSnapshot: balance.Snapshot{ + At: start, + Balances: balance.Map{g.ID: g.Amount}, + UsageSnapshot: &balance.UsageSnapshot{ + TotalGrantUsage: startingTotalGrantUsage, + }, + }, + Until: start.Add(time.Hour), + }) + require.NoError(t, err) + require.NotNil(t, result.Snapshot.UsageSnapshot) + assert.Equal(t, 10.0, result.Snapshot.UsageSnapshot.Usage) + assert.Equal(t, 35.0, result.Snapshot.UsageSnapshot.TotalGrantUsage) + + historyStart, err := result.History.GetSnapshotAtStartOfSegment(0) + require.NoError(t, err) + require.NotNil(t, historyStart.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + TotalGrantUsage: 25, + }, *historyStart.UsageSnapshot) +} + +func TestEngineResetsTotalGrantUsageBeforeRolloverOverageBurn(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + resetAt := start.Add(time.Hour) + end := resetAt.Add(time.Hour) + startingTotalGrantUsage := 25.0 + g := grant.Grant{ + ID: "grant-1", + Amount: 100, + Priority: 1, + EffectiveAt: start, + ResetMinRollover: 50, + ResetMaxRollover: 50, + } + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, from, _ time.Time) (float64, error) { + if from.Before(resetAt) { + return 110, nil + } + + return 5, nil + }, + }) + + result, err := eng.Run(t.Context(), engine.RunParams{ + Meter: meter.Meter{ + Aggregation: meter.MeterAggregationSum, + }, + Grants: []grant.Grant{g}, + StartingSnapshot: balance.Snapshot{ + At: start, + Balances: balance.Map{g.ID: g.Amount}, + UsageSnapshot: &balance.UsageSnapshot{ + TotalGrantUsage: startingTotalGrantUsage, + }, + }, + Until: end, + ResetBehavior: grant.ResetBehavior{ + PreserveOverage: true, + }, + Resets: timeutil.NewSimpleTimeline([]time.Time{resetAt}), + }) + require.NoError(t, err) + require.Len(t, result.History.Segments(), 3) + + rolloverStart, err := result.History.GetSnapshotAtStartOfSegment(1) + require.NoError(t, err) + require.NotNil(t, rolloverStart.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{}, *rolloverStart.UsageSnapshot) + + afterRollover, err := result.History.GetSnapshotAtStartOfSegment(2) + require.NoError(t, err) + require.NotNil(t, afterRollover.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + TotalGrantUsage: 10, + }, *afterRollover.UsageSnapshot) + + require.NotNil(t, result.Snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 15, + }, *result.Snapshot.UsageSnapshot) +} + +func TestEnginePreservesUsageSnapshotAcrossRuns(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + intermediate := start.Add(time.Hour) + end := intermediate.Add(time.Hour) + g := grant.Grant{ + ID: "grant-1", + Amount: 100, + Priority: 1, + EffectiveAt: start, + } + startingSnapshot := balance.Snapshot{ + At: start, + Balances: balance.Map{g.ID: 95}, + Usage: balance.SnapshottedUsage{ + Since: start, + Usage: 5, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 5, + }, + } + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, from, to time.Time) (float64, error) { + return to.Sub(from).Hours() * 10, nil + }, + }) + runParams := engine.RunParams{ + Meter: meter.Meter{ + Aggregation: meter.MeterAggregationSum, + }, + Grants: []grant.Grant{g}, + StartingSnapshot: startingSnapshot, + Until: end, + } + + singleRun, err := eng.Run(t.Context(), runParams) + require.NoError(t, err) + + firstRunParams := runParams + firstRunParams.Until = intermediate + firstRun, err := eng.Run(t.Context(), firstRunParams) + require.NoError(t, err) + + secondRunParams := runParams + secondRunParams.StartingSnapshot = firstRun.Snapshot + secondRun, err := eng.Run(t.Context(), secondRunParams) + require.NoError(t, err) + + require.NotNil(t, singleRun.Snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: 25, + TotalGrantUsage: 25, + }, *singleRun.Snapshot.UsageSnapshot) + assert.Equal(t, singleRun.Snapshot, secondRun.Snapshot) +} + +func TestEnginePreservesUsageSnapshotWhenResumingAfterRollover(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + resetAt := start.Add(time.Hour) + end := resetAt.Add(time.Hour) + g := grant.Grant{ + ID: "grant-1", + Amount: 100, + Priority: 1, + EffectiveAt: start, + ResetMinRollover: 50, + ResetMaxRollover: 50, + } + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, from, _ time.Time) (float64, error) { + if from.Before(resetAt) { + return 110, nil + } + + return 5, nil + }, + }) + runParams := engine.RunParams{ + Meter: meter.Meter{ + Aggregation: meter.MeterAggregationSum, + }, + Grants: []grant.Grant{g}, + StartingSnapshot: balance.NewStartingSnapshot([]grant.Grant{g}, start), + Until: end, + ResetBehavior: grant.ResetBehavior{ + PreserveOverage: true, + }, + Resets: timeutil.NewSimpleTimeline([]time.Time{resetAt}), + } + + singleRun, err := eng.Run(t.Context(), runParams) + require.NoError(t, err) + require.Len(t, singleRun.History.Segments(), 3) + + afterRollover, err := singleRun.History.GetSnapshotAtStartOfSegment(2) + require.NoError(t, err) + + resumedRunParams := runParams + resumedRunParams.StartingSnapshot = afterRollover + resumedRunParams.Resets = timeutil.SimpleTimeline{} + resumedRun, err := eng.Run(t.Context(), resumedRunParams) + require.NoError(t, err) + + require.NotNil(t, resumedRun.Snapshot.UsageSnapshot) + assert.Equal(t, balance.UsageSnapshot{ + Usage: 5, + TotalGrantUsage: 15, + }, *resumedRun.Snapshot.UsageSnapshot) + assert.Equal(t, singleRun.Snapshot, resumedRun.Snapshot) +} + +func TestEnginePreservesUnitConfig(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + resetAt := start.Add(time.Hour) + + tests := []struct { + name string + resets timeutil.SimpleTimeline + }{ + { + name: "within a usage period", + }, + { + name: "across a reset", + resets: timeutil.NewSimpleTimeline([]time.Time{resetAt}), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + displayUnit := "requests" + startingUnitConfig := &unitconfig.UnitConfig{ + Operation: unitconfig.UnitConfigOperationDivide, + ConversionFactor: alpacadecimal.NewFromInt(1000), + DisplayUnit: &displayUnit, + } + startingSnapshot := balance.NewStartingSnapshot(nil, start) + startingSnapshot.UnitConfig = startingUnitConfig + eng := engine.NewEngine(engine.EngineConfig{ + QueryUsage: func(_ context.Context, _, _ time.Time) (float64, error) { + return 0, nil + }, + }) + + result, err := eng.Run(t.Context(), engine.RunParams{ + StartingSnapshot: startingSnapshot, + Until: start.Add(2 * time.Hour), + Resets: tt.resets, + }) + require.NoError(t, err) + require.NotNil(t, result.Snapshot.UnitConfig) + assert.True(t, startingUnitConfig.Equal(result.Snapshot.UnitConfig)) + assert.NotSame(t, startingUnitConfig, result.Snapshot.UnitConfig) + assert.NotSame(t, startingUnitConfig.DisplayUnit, result.Snapshot.UnitConfig.DisplayUnit) + }) + } +} diff --git a/openmeter/credit/helper.go b/openmeter/credit/helper.go index 6bb9c2ac9a..02100f284e 100644 --- a/openmeter/credit/helper.go +++ b/openmeter/credit/helper.go @@ -76,11 +76,7 @@ func (m *connector) startOfMeasurementSnapshot(ctx context.Context, owner models return balance.Snapshot{}, err } - return balance.Snapshot{ - At: startOfMeasurement, - Balances: balance.NewStartingMap(grants, startOfMeasurement), - Overage: 0.0, // There cannot be overage at the start of measurement - }, nil + return balance.NewStartingSnapshot(grants, startOfMeasurement), nil } func (m *connector) runEngineInSpan(ctx context.Context, eng engine.Engine, runParams engine.RunParams) (engine.RunResult, error) { diff --git a/openmeter/entitlement/README.md b/openmeter/entitlement/README.md new file mode 100644 index 0000000000..5058308d77 --- /dev/null +++ b/openmeter/entitlement/README.md @@ -0,0 +1,93 @@ +# Entitlements + +Entitlements answer whether a customer may use a feature and, when relevant, +return the state behind that decision. They connect product configuration to a +customer's usage attribution, but they do not own raw usage or accounting +entries. + +## Domain model + +An entitlement belongs to a customer and a feature. At most one entitlement for +the same feature is active for a customer at a given time. + +| Type | Meaning | +| --- | --- | +| `boolean` | grants access while active | +| `static` | grants access and carries an immutable JSON configuration | +| `metered` | derives access from metered usage and grant balance | + +Entitlement values are time-bound. An entitlement is active from its configured +start, or creation time, until its configured end or deletion. The end is +exclusive. Outside that interval it provides no access. + +## Ownership and lifecycle + +Entitlements may be created directly or derived from subscriptions. +[Subscription](../subscription/README.md) owns the customer-specific schedule +that materializes and supersedes subscription-managed entitlements. +Entitlement owns their persisted lifecycle and value resolution. + +Definitions are historical. Deletion ends an entitlement at a timestamp rather +than erasing it, and overriding supersedes the previous definition rather than +changing its meaning in the past. + +An entitlement identifies the customer-feature relationship. Metered value +calculations resolve that customer's usage attribution when querying the +feature's meter. + +## Metered entitlements + +A metered entitlement connects a feature's meter to grants owned by the +entitlement. Metered usage consumes active grants in priority order. Uncovered +usage becomes overage. + +Access is allowed while balance remains. A soft limit continues to allow access +after the balance is exhausted; it does not prevent usage or overage from being +calculated. + +Measurement begins at the entitlement's measurement start. Usage before that +time is outside the entitlement calculation. + +Metered values inherit [Credit](../credit/README.md)'s one-minute resolution. +They must not be used to express sub-minute access transitions. + +## Usage periods and resets + +Usage periods partition a metered entitlement's usage over time. The initial +period begins at measurement start and later periods follow the configured +recurrence. + +An explicit reset ends the current period and starts a new one. The reset may +retain the original recurrence anchor or establish a new one. Grant rollover +and preserved-overage behavior are part of the credit reset, not separate +entitlement balances. + +See [Credit](../credit/README.md) for grant consumption, reset, history, and +snapshot semantics. + +## Value events + +Entitlement value reads calculate the value at the requested time. Separately, +the balance worker recalculates values when usage or entitlement state changes +and publishes value events for notifications and downstream consumers. + +Those event payloads are sometimes called snapshots, but they are not the +durable credit balance snapshots used to resume calculations. Event handlers +are retried, so consumers must process them idempotently. + +Crossing a future entitlement or grant activation or expiry time does not by +itself produce a value event. Reads reflect the new state, while event consumers +observe it after another recalculation trigger or a scheduled recalculation. + +## Boundaries + +- Features and meters define what is measured; entitlement selects the feature, + customer attribution, and access policy. +- Streaming owns raw usage and meter queries. +- [Credit](../credit/README.md) owns grants, balance calculation, resets, + history, and balance snapshots for metered entitlements. +- [Subscriptions](../subscription/README.md) may materialize and supersede + subscription-managed entitlements. +- Notifications consume entitlement value events; those events are not a + transactional source of balance truth. +- Billing and [Ledger](../ledger/README.md) own monetary and accounting state. diff --git a/openmeter/entitlement/metered/balance.go b/openmeter/entitlement/metered/balance.go index d4f88c20fa..0c3845d3cf 100644 --- a/openmeter/entitlement/metered/balance.go +++ b/openmeter/entitlement/metered/balance.go @@ -98,9 +98,9 @@ func (e *connector) GetEntitlementBalance(ctx context.Context, entitlementID mod return &EntitlementBalance{ EntitlementID: entitlementID.ID, Balance: res.Snapshot.Balance(), - UsageInPeriod: res.Snapshot.Usage.Usage, + UsageInPeriod: res.Snapshot.UsageSnapshot.Usage, Overage: res.Snapshot.Overage, - TotalAvailableGrantAmount: res.TotalAvailableGrantAmountAtLastPeriod(), + TotalAvailableGrantAmount: res.Snapshot.UsageSnapshot.TotalGrantUsage + res.Snapshot.Balance(), GrantBalances: grantBalances, StartOfPeriod: startOfPeriod, }, nil diff --git a/openmeter/entitlement/metered/balance_test.go b/openmeter/entitlement/metered/balance_test.go index d07bb00814..947c1e5395 100644 --- a/openmeter/entitlement/metered/balance_test.go +++ b/openmeter/entitlement/metered/balance_test.go @@ -287,6 +287,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -306,6 +307,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -331,6 +333,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: queryTime, // querytime is the start of a UsagePeriod, so this snapshot will be at the start of the usage period Usage: 0, // And at a reset time the usage is 0 }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 800, }, @@ -396,6 +399,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -415,6 +419,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -440,6 +445,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime.AddDate(0, 0, 9), // will create a snapshot at the start of the usage period Usage: 0, // And at a reset time the usage is 0 }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 800, }, @@ -524,6 +530,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: datetime.NewDateTime(anchor).AddDateNoOverflow(0, 2, 0).Time, // Will create a snapshot at the last history breakpoint outside 7 day grace period (which is the last reset time) Usage: 0, // And at a reset time the usage is 0 }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 800, }, @@ -659,6 +666,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -701,6 +709,7 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime, Usage: 0, }, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -726,6 +735,10 @@ func TestGetEntitlementBalance(t *testing.T) { Since: startTime.AddDate(0, 1, 0), // The programmatic reset time Usage: 200, // Total usage in second period so far }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 200, + TotalGrantUsage: 200, + }, Balances: balance.Map{ g1.ID: 600, g2.ID: 1000, @@ -785,6 +798,7 @@ func TestGetEntitlementBalance(t *testing.T) { ctx, owner, []balance.Snapshot{ { + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, @@ -1429,7 +1443,12 @@ func TestGetEntitlementHistory(t *testing.T) { // check returned burndownhistory segments := burndownHistory.Segments() - assert.Len(t, segments, 2) + require.Len(t, segments, 3) + assert.True(t, segments[0].TerminationReasons.UsageReset) + assert.True(t, segments[1].TerminationReasons.Rollover) + assert.Equal(t, resetTime, segments[1].From) + assert.Equal(t, resetTime, segments[1].To) + assert.False(t, segments[2].TerminationReasons.Rollover) assert.Len(t, windowedHistory, 2) diff --git a/openmeter/entitlement/metered/balance_total_available_test.go b/openmeter/entitlement/metered/balance_total_available_test.go new file mode 100644 index 0000000000..636571be86 --- /dev/null +++ b/openmeter/entitlement/metered/balance_total_available_test.go @@ -0,0 +1,312 @@ +package meteredentitlement_test + +import ( + "testing" + "time" + + "github.com/samber/lo" + "github.com/stretchr/testify/require" + + "github.com/openmeterio/openmeter/openmeter/credit/balance" + "github.com/openmeterio/openmeter/openmeter/credit/grant" + db_balancesnapshot "github.com/openmeterio/openmeter/openmeter/ent/db/balancesnapshot" + "github.com/openmeterio/openmeter/openmeter/entitlement" + "github.com/openmeterio/openmeter/openmeter/productcatalog/feature" + "github.com/openmeterio/openmeter/openmeter/testutils" + "github.com/openmeterio/openmeter/pkg/convert" + "github.com/openmeterio/openmeter/pkg/filter" + "github.com/openmeterio/openmeter/pkg/models" + "github.com/openmeterio/openmeter/pkg/timeutil" +) + +func TestGetEntitlementBalanceTotalAvailableGrantAmountAfterSnapshot(t *testing.T) { + connector, deps := setupConnector(t) + defer deps.Teardown() + + ctx := t.Context() + periodStart := getAnchor(t) + snapshotAt := periodStart.Add(time.Hour) + queryAt := snapshotAt.Add(time.Hour) + + feat, err := deps.featureRepo.CreateFeature(ctx, feature.CreateFeatureInputs{ + Namespace: namespace, + Name: "feature1", + Key: "feature-1", + MeterID: &deps.meterID, + MeterGroupByFilters: map[string]filter.FilterString{}, + }) + require.NoError(t, err) + + randName := testutils.NameGenerator.Generate() + cust := createCustomerAndSubject(t, deps.subjectService, deps.customerService, namespace, randName.Key, randName.Name) + + usagePeriod := entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{ + Anchor: periodStart, + Interval: timeutil.RecurrencePeriodYear, + }) + currentUsagePeriod, err := usagePeriod.GetValue().GetPeriodAt(queryAt) + require.NoError(t, err) + + ent, err := deps.entitlementRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{ + Namespace: namespace, + FeatureID: feat.ID, + FeatureKey: feat.Key, + UsageAttribution: cust.GetUsageAttribution(), + MeasureUsageFrom: &periodStart, + EntitlementType: entitlement.EntitlementTypeMetered, + IssueAfterReset: convert.ToPointer(0.0), + IsSoftLimit: convert.ToPointer(false), + UsagePeriod: &usagePeriod, + CurrentUsagePeriod: ¤tUsagePeriod, + }) + require.NoError(t, err) + + g, err := deps.grantRepo.CreateGrant(ctx, grant.RepoCreateInput{ + OwnerID: ent.ID, + Namespace: namespace, + Amount: 1000, + Priority: 1, + EffectiveAt: periodStart, + ExpiresAt: lo.ToPtr(periodStart.AddDate(1, 0, 0)), + }) + require.NoError(t, err) + + // given: + // - 200 usage already consumed from a 1000-unit grant in the current period + // - a persisted mid-period snapshot carrying the cumulative usage and remaining balance + deps.streamingConnector.AddSimpleEvent(meterSlug, 200, periodStart.Add(time.Minute)) + err = deps.balanceSnapshotService.Save(ctx, models.NamespacedID{ + Namespace: namespace, + ID: ent.ID, + }, []balance.Snapshot{ + { + At: snapshotAt, + Balances: balance.Map{g.ID: 800}, + Usage: balance.SnapshottedUsage{ + Since: periodStart, + Usage: 200, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 200, + TotalGrantUsage: 200, + }, + }, + }) + require.NoError(t, err) + + // when: another 100 usage is consumed after the snapshot + deps.streamingConnector.AddSimpleEvent(meterSlug, 100, snapshotAt.Add(time.Minute)) + entBalance, err := connector.GetEntitlementBalance(ctx, models.NamespacedID{ + Namespace: namespace, + ID: ent.ID, + }, queryAt) + require.NoError(t, err) + + // then: the period values must remain internally consistent regardless of snapshotting + require.Equal(t, 300.0, entBalance.UsageInPeriod) + require.Equal(t, 700.0, entBalance.Balance) + require.Equal(t, 0.0, entBalance.Overage) + require.Equal(t, 1000.0, entBalance.TotalAvailableGrantAmount) +} + +func TestBalanceSnapshotPersistenceRequiresUsageSnapshot(t *testing.T) { + _, deps := setupConnector(t) + defer deps.Teardown() + + err := deps.balanceSnapshotService.Save(t.Context(), models.NamespacedID{ + Namespace: namespace, + ID: "01ARZ3NDEKTSV4RRFFQ69G5FAV", + }, []balance.Snapshot{{ + At: getAnchor(t), + Balances: balance.Map{"grant-1": 800}, + }}) + require.ErrorContains(t, err, "cannot save incomplete balance snapshot") +} + +func TestBalanceSnapshotSelectionDuringUsageSnapshotMigration(t *testing.T) { + _, deps := setupConnector(t) + defer deps.Teardown() + + ctx := t.Context() + periodStart := getAnchor(t) + completeSnapshotAt := periodStart.Add(time.Hour) + legacySnapshotAt := completeSnapshotAt.Add(time.Hour) + owner := createBalanceSnapshotOwner(t, deps, periodStart) + completeSnapshot := balance.Snapshot{ + At: completeSnapshotAt, + Balances: balance.Map{"grant-1": 800}, + Usage: balance.SnapshottedUsage{ + Since: periodStart, + Usage: 200, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 200, + TotalGrantUsage: 200, + }, + } + + err := deps.balanceSnapshotService.Save(ctx, owner, []balance.Snapshot{completeSnapshot}) + require.NoError(t, err) + + _, err = deps.dbClient.BalanceSnapshot.Create(). + SetNamespace(owner.Namespace). + SetOwnerID(owner.ID). + SetAt(legacySnapshotAt). + SetBalance(700). + SetGrantBalances(balance.Map{"grant-1": 700}). + SetOverage(0). + SetUsage(&balance.SnapshottedUsage{ + Since: periodStart, + Usage: 300, + }). + Save(ctx) + require.NoError(t, err) + + selectedSnapshot, err := deps.balanceSnapshotService.GetLatestValidAt(ctx, owner, legacySnapshotAt) + require.NoError(t, err) + require.Equal(t, completeSnapshot, selectedSnapshot) + + _, err = deps.dbClient.BalanceSnapshot.Delete(). + Where( + db_balancesnapshot.Namespace(owner.Namespace), + db_balancesnapshot.OwnerID(owner.ID), + db_balancesnapshot.UsageSnapshotNotNil(), + ). + Exec(ctx) + require.NoError(t, err) + + _, err = deps.balanceSnapshotService.GetLatestValidAt(ctx, owner, legacySnapshotAt) + var noSavedSnapshot *balance.NoSavedBalanceForOwnerError + require.ErrorAs(t, err, &noSavedSnapshot) +} + +func TestBalanceSnapshotVersionsAreIndependent(t *testing.T) { + _, deps := setupConnector(t) + defer deps.Teardown() + + ctx := t.Context() + periodStart := getAnchor(t) + snapshotAt := periodStart.Add(time.Hour) + queryAt := snapshotAt.Add(time.Hour) + owner := createBalanceSnapshotOwner(t, deps, periodStart) + snapshotA := balance.Snapshot{ + At: snapshotAt, + Balances: balance.Map{"grant-1": 800}, + Usage: balance.SnapshottedUsage{ + Since: periodStart, + Usage: 200, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 200, + TotalGrantUsage: 200, + }, + } + snapshotB := balance.Snapshot{ + At: snapshotAt.Add(10 * time.Minute), + Balances: balance.Map{"grant-1": 700}, + Usage: balance.SnapshottedUsage{ + Since: periodStart, + Usage: 300, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 300, + TotalGrantUsage: 300, + }, + } + snapshotC := balance.Snapshot{ + At: snapshotAt.Add(20 * time.Minute), + Balances: balance.Map{"grant-1": 600}, + Usage: balance.SnapshottedUsage{ + Since: periodStart, + Usage: 400, + }, + UsageSnapshot: &balance.UsageSnapshot{ + Usage: 400, + TotalGrantUsage: 400, + }, + } + + err := deps.balanceSnapshotService.Save(ctx, owner, []balance.Snapshot{snapshotA, snapshotB, snapshotC}) + require.NoError(t, err) + + _, err = deps.dbClient.BalanceSnapshot.Delete(). + Where( + db_balancesnapshot.Namespace(owner.Namespace), + db_balancesnapshot.OwnerID(owner.ID), + db_balancesnapshot.AtEQ(snapshotB.At), + ). + Exec(ctx) + require.NoError(t, err) + + selectedSnapshot, err := deps.balanceSnapshotService.GetLatestValidAt(ctx, owner, queryAt) + require.NoError(t, err) + require.Equal(t, snapshotC, selectedSnapshot) + + _, err = deps.dbClient.BalanceSnapshot.Delete(). + Where( + db_balancesnapshot.Namespace(owner.Namespace), + db_balancesnapshot.OwnerID(owner.ID), + db_balancesnapshot.AtEQ(snapshotA.At), + ). + Exec(ctx) + require.NoError(t, err) + + selectedSnapshot, err = deps.balanceSnapshotService.GetLatestValidAt(ctx, owner, queryAt) + require.NoError(t, err) + require.Equal(t, snapshotC, selectedSnapshot) + + err = deps.balanceSnapshotService.Save(ctx, owner, []balance.Snapshot{snapshotB}) + require.NoError(t, err) + + selectedSnapshot, err = deps.balanceSnapshotService.GetLatestValidAt(ctx, owner, queryAt) + require.NoError(t, err) + require.Equal(t, snapshotC, selectedSnapshot) +} + +func createBalanceSnapshotOwner(t *testing.T, deps *dependencies, at time.Time) models.NamespacedID { + ctx := t.Context() + featureName := testutils.NameGenerator.Generate() + feat, err := deps.featureRepo.CreateFeature(ctx, feature.CreateFeatureInputs{ + Namespace: namespace, + Name: featureName.Name, + Key: featureName.Key, + MeterID: &deps.meterID, + MeterGroupByFilters: map[string]filter.FilterString{}, + }) + require.NoError(t, err) + + customerName := testutils.NameGenerator.Generate() + cust := createCustomerAndSubject( + t, + deps.subjectService, + deps.customerService, + namespace, + customerName.Key, + customerName.Name, + ) + usagePeriod := entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{ + Anchor: at, + Interval: timeutil.RecurrencePeriodYear, + }) + currentUsagePeriod, err := usagePeriod.GetValue().GetPeriodAt(at) + require.NoError(t, err) + + ent, err := deps.entitlementRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{ + Namespace: namespace, + FeatureID: feat.ID, + FeatureKey: feat.Key, + UsageAttribution: cust.GetUsageAttribution(), + MeasureUsageFrom: &at, + EntitlementType: entitlement.EntitlementTypeMetered, + IssueAfterReset: convert.ToPointer(0.0), + IsSoftLimit: convert.ToPointer(false), + UsagePeriod: &usagePeriod, + CurrentUsagePeriod: ¤tUsagePeriod, + }) + require.NoError(t, err) + + return models.NamespacedID{ + Namespace: namespace, + ID: ent.ID, + } +} diff --git a/openmeter/entitlement/metered/balance_unitconfig_test.go b/openmeter/entitlement/metered/balance_unitconfig_test.go index d261e7f673..fa740d1574 100644 --- a/openmeter/entitlement/metered/balance_unitconfig_test.go +++ b/openmeter/entitlement/metered/balance_unitconfig_test.go @@ -181,11 +181,12 @@ func TestBalanceSnapshotRegimeMismatchRecomputes(t *testing.T) { staleAt := startTime.Add(30 * time.Minute) err = deps.balanceSnapshotService.Save(ctx, models.NamespacedID{Namespace: namespace, ID: ent.ID}, []balance.Snapshot{ { - At: staleAt, - Balances: balance.Map{g.ID: 999}, - Overage: 0, - Usage: balance.SnapshottedUsage{Usage: 0, Since: startTime}, - UnitConfig: nil, + At: staleAt, + Balances: balance.Map{g.ID: 999}, + Overage: 0, + Usage: balance.SnapshottedUsage{Usage: 0, Since: startTime}, + UsageSnapshot: &balance.UsageSnapshot{}, + UnitConfig: nil, }, }) require.NoError(t, err) diff --git a/openmeter/entitlement/metered/lateevents_test.go b/openmeter/entitlement/metered/lateevents_test.go index 0ad2d9b962..56cab6e063 100644 --- a/openmeter/entitlement/metered/lateevents_test.go +++ b/openmeter/entitlement/metered/lateevents_test.go @@ -173,9 +173,7 @@ func TestGetEntitlementBalanceConsistency(t *testing.T) { ) balanceSnapshotService := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: ownerConnector, - StreamingConnector: streamingConnector, - Repo: balanceSnapshotRepo, + Repo: balanceSnapshotRepo, }) transactionManager := enttx.NewCreator(dbClient) diff --git a/openmeter/entitlement/metered/reset_test.go b/openmeter/entitlement/metered/reset_test.go index 4b04f72a9d..3d0f6fd17b 100644 --- a/openmeter/entitlement/metered/reset_test.go +++ b/openmeter/entitlement/metered/reset_test.go @@ -300,8 +300,9 @@ func TestResetEntitlementUsage(t *testing.T) { ID: ent.ID, }, []balance.Snapshot{ { - At: g1.EffectiveAt, - Overage: 0, + At: g1.EffectiveAt, + Overage: 0, + UsageSnapshot: &balance.UsageSnapshot{}, Balances: balance.Map{ g1.ID: 1000, }, diff --git a/openmeter/entitlement/metered/utils_test.go b/openmeter/entitlement/metered/utils_test.go index 1d31c80770..ac933fb537 100644 --- a/openmeter/entitlement/metered/utils_test.go +++ b/openmeter/entitlement/metered/utils_test.go @@ -156,9 +156,7 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies) transactionManager := enttx.NewCreator(dbClient) snapshotService := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: ownerConnector, - StreamingConnector: streamingConnector, - Repo: balanceSnapshotRepo, + Repo: balanceSnapshotRepo, }) creditConnector := credit.NewCreditConnector( diff --git a/openmeter/registry/builder/entitlement.go b/openmeter/registry/builder/entitlement.go index 7a03596958..0623d8154d 100644 --- a/openmeter/registry/builder/entitlement.go +++ b/openmeter/registry/builder/entitlement.go @@ -63,9 +63,7 @@ func GetEntitlementRegistry(opts EntitlementOptions) *registry.Entitlement { transactionManager := enttx.NewCreator(opts.DatabaseClient) balanceSnapshotService := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: entitlementOwnerConnector, - StreamingConnector: opts.StreamingConnector, - Repo: balanceSnashotDBAdapter, + Repo: balanceSnashotDBAdapter, }) creditConnector := credit.NewCreditConnector( diff --git a/openmeter/subscription/README.md b/openmeter/subscription/README.md index e7564b29e3..59d1522a2f 100644 --- a/openmeter/subscription/README.md +++ b/openmeter/subscription/README.md @@ -102,9 +102,9 @@ and phase; do not persist a derived absolute end as independent source truth. - The product catalog supplies the plan shape used to build the initial spec. The subscription owns the resulting customer-specific schedule, not the mutable plan definition. -- Entitlements are derived from entitlement-bearing rate cards and share the - effective cadence of their subscription item. Recreating an item may - recreate its entitlement. +- [Entitlements](../entitlement/README.md) are derived from + entitlement-bearing rate cards and share the effective cadence of their + subscription item. Recreating an item may recreate its entitlement. - Subscription commands publish lifecycle events. Successful subscription mutation means the desired schedule was committed; it does not mean billing artifacts have already been reconciled. diff --git a/test/billing/subscription_suite.go b/test/billing/subscription_suite.go index eeb4ababca..323cb16d08 100644 --- a/test/billing/subscription_suite.go +++ b/test/billing/subscription_suite.go @@ -269,9 +269,7 @@ func (s *SubscriptionMixin) SetupEntitlements(t *testing.T, deps SubscriptionMix ) balanceSnapshotService := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: owner, - StreamingConnector: deps.MockStreamingConnector, - Repo: balanceSnapshotRepo, + Repo: balanceSnapshotRepo, }) transactionManager := enttx.NewCreator(deps.DBClient) diff --git a/test/entitlement/regression/framework_test.go b/test/entitlement/regression/framework_test.go index 3641fe2304..451bb6132d 100644 --- a/test/entitlement/regression/framework_test.go +++ b/test/entitlement/regression/framework_test.go @@ -159,9 +159,7 @@ func setupDependencies(t *testing.T) Dependencies { ) balanceSnapshotService := balance.NewSnapshotService(balance.SnapshotServiceConfig{ - OwnerConnector: owner, - StreamingConnector: streaming, - Repo: balanceSnapshotRepo, + Repo: balanceSnapshotRepo, }) transactionManager := enttx.NewCreator(dbClient)