fix: totalAvailableGrantAmount (#4304)

Author: GAlexIHU

openmeter/credit/engine/burnphase.go

@@ -31,9 +31,9 @@ type phasePlan struct {
 //
 // Note that grant balance does not effect the burndown order if we simply ignore grants that don't
 // have balance while burning down.
-func (e *engine) getPhases(grants []grant.Grant, period timeutil.ClosedPeriod) (phasePlan, error) {
+func (e *engine) getPhases(grants []grant.Grant, period timeutil.ClosedPeriod, recurrenceEndBoundaryBehavior timeutil.Boundary) (phasePlan, error) {
 	activityChanges := e.getGrantActivityChanges(grants, period)
-	recurrenceTimes, err := e.getGrantRecurrenceTimes(grants, period)
+	recurrenceTimes, err := e.getGrantRecurrenceTimes(grants, period, recurrenceEndBoundaryBehavior)
 	if err != nil {
 		return phasePlan{}, fmt.Errorf("failed to get grant recurrence times: %w", err)
 	}
@@ -68,6 +68,10 @@ func (e *engine) getPhases(grants []grant.Grant, period timeutil.ClosedPeriod) (
 	phaseFrom := period.From
 
 	appendPhase := func(phase burnPhase) {
+		if !phase.to.After(phase.from) {
+			return
+		}
+
 		phases = append(phases, phase)
 		phaseFrom = phase.to
 	}
@@ -102,10 +106,7 @@ func (e *engine) getPhases(grants []grant.Grant, period timeutil.ClosedPeriod) (
 			rtI++
 		}
 
-		// If it's a valid phase (non-zero duration), we save it and break
-		if phase.to.After(phase.from) {
-			appendPhase(phase)
-		}
+		appendPhase(phase)
 	}
 
 	// order here doesn't matter as one or both of them is empty
@@ -133,6 +134,25 @@ func (e *engine) getPhases(grants []grant.Grant, period timeutil.ClosedPeriod) (
 		})
 	}
 
+	if len(phases) > 0 {
+		lastPhase := phases[len(phases)-1]
+		terminalEvent := lastPhase.to.Equal(period.To) &&
+			len(lastPhase.grantsRecurredAtEnd) > 0
+		if terminalEvent {
+			phases = append(phases, burnPhase{
+				from: period.To,
+				to:   period.To,
+			})
+		}
+	}
+
+	if len(phases) == 0 {
+		phases = append(phases, burnPhase{
+			from: period.From,
+			to:   period.To,
+		})
+	}
+
 	return phasePlan{
 		phases:                phases,
 		grantsRecurredAtStart: grantsRecurredAtStart,

openmeter/credit/engine/engine.go

@@ -56,36 +56,14 @@ type RunResult struct {
 	RunParams RunParams
 }
 
-// TotalAvailableGrantAmount is the total amount of grants either currently active + the used amount of currently inactive grants.
-func (r RunResult) TotalAvailableGrantAmount() float64 {
-	// First, let's calculate the total amount of grants active at the end of the period.
-	activeAmount := lo.Reduce(r.RunParams.Grants, func(agg alpacadecimal.Decimal, grant grant.Grant, _ int) alpacadecimal.Decimal {
-		if !grant.ActiveAt(r.RunParams.Until) {
-			return agg
-		}
-
-		return agg.Add(alpacadecimal.NewFromFloat(grant.Amount))
-	}, alpacadecimal.NewFromFloat(0))
-
-	// Second, let's calculate the used-up amount of since inactive grants.
-	inactiveGrants := lo.Filter(r.RunParams.Grants, func(grant grant.Grant, _ int) bool {
-		return !grant.ActiveAt(r.RunParams.Until)
-	})
-
-	usedInactive := alpacadecimal.NewFromFloat(0)
-	if len(inactiveGrants) > 0 {
-		for _, seg := range r.History.Segments() {
-			for _, usage := range seg.GrantUsages {
-				if lo.SomeBy(inactiveGrants, func(grant grant.Grant) bool {
-					return grant.ID == usage.GrantID
-				}) {
-					usedInactive = usedInactive.Add(alpacadecimal.NewFromFloat(usage.Usage))
-				}
-			}
-		}
-	}
-
-	return activeAmount.Add(usedInactive).InexactFloat64()
+// 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 {

openmeter/credit/engine/grant.go

@@ -16,26 +16,27 @@ import (
 // An activity change is a grant becoming active or a grant expiring.
 func (e *engine) getGrantActivityChanges(grants []grant.Grant, period timeutil.ClosedPeriod) []time.Time {
 	activityChanges := []time.Time{}
+
 	for _, grant := range grants {
 		// grants that take effect in the period
-		if grant.EffectiveAt.After(period.From) && (grant.EffectiveAt.Before(period.To)) {
+		if period.Contains(grant.EffectiveAt) {
 			activityChanges = append(activityChanges, grant.EffectiveAt)
 		}
 		// grants that expire in the period
 		if grant.ExpiresAt != nil {
-			if grant.ExpiresAt.After(period.From) && (grant.ExpiresAt.Before(period.To)) {
+			if period.Contains(*grant.ExpiresAt) {
 				activityChanges = append(activityChanges, *grant.ExpiresAt)
 			}
 		}
 		// grants that are deleted in the period
 		if grant.DeletedAt != nil {
-			if grant.DeletedAt.After(period.From) && (grant.DeletedAt.Before(period.To)) {
+			if period.Contains(*grant.DeletedAt) {
 				activityChanges = append(activityChanges, *grant.DeletedAt)
 			}
 		}
 		// grants that are voided in the period
 		if grant.VoidedAt != nil {
-			if grant.VoidedAt.After(period.From) && (grant.VoidedAt.Before(period.To)) {
+			if period.Contains(*grant.VoidedAt) {
 				activityChanges = append(activityChanges, *grant.VoidedAt)
 			}
 		}
@@ -62,7 +63,7 @@ func (e *engine) getGrantActivityChanges(grants []grant.Grant, period timeutil.C
 }
 
 // Get all times grants recurr in the period.
-func (e *engine) getGrantRecurrenceTimes(grants []grant.Grant, period timeutil.ClosedPeriod) ([]struct {
+func (e *engine) getGrantRecurrenceTimes(grants []grant.Grant, period timeutil.ClosedPeriod, endBoundaryBehavior timeutil.Boundary) ([]struct {
 	time     time.Time
 	grantIDs []string
 }, error,
@@ -86,12 +87,29 @@ func (e *engine) getGrantRecurrenceTimes(grants []grant.Grant, period timeutil.C
 		if err != nil {
 			return nil, err
 		}
-		// Write all recurrence times in [period.From, period.To).
+
+		// Write all recurrence times in [period.From, period.To]).
 		// For a zero-length period [T, T], include T so a recurrence at the
 		// start of that period can still be applied.
+		// Include period.To for the final run period so a run ending exactly on
+		// a recurrence produces a zero-length terminal phase where it is applied.
 		inPeriod := func(at time.Time) bool {
-			return at.Before(period.To) || (period.From.Equal(period.To) && at.Equal(period.From))
+			behavior := endBoundaryBehavior
+
+			if period.IsEmpty() && at.Equal(period.From) {
+				behavior = timeutil.Inclusive
+			}
+
+			switch behavior {
+			case timeutil.Inclusive:
+				return period.ContainsInclusive(at) // [from, to]
+			case timeutil.Exclusive:
+				return period.Contains(at) // [from, to)
+			default:
+				return false
+			}
 		}
+
 		for inPeriod(it.At) && grant.ActiveAt(it.At) {
 			times = append(times, struct {
 				time    time.Time

openmeter/credit/engine/history.go

@@ -4,6 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"sort"
+	"time"
+
+	"github.com/alpacahq/alpacadecimal"
 
 	"github.com/openmeterio/openmeter/openmeter/credit/balance"
 	"github.com/openmeterio/openmeter/pkg/timeutil"
@@ -144,10 +147,7 @@ func (g *GrantBurnDownHistory) GetUsageInPeriodUntilSegment(segmentIndex int) (b
 		// We need the segment right after the last reset
 		if lastResetSegmentIndex+1 < len(g.segments) {
 			firstSeg := g.segments[lastResetSegmentIndex+1]
-			usage = balance.SnapshottedUsage{
-				Since: firstSeg.From,
-				Usage: 0.0,
-			}
+			usage = usageAtReset(firstSeg.From)
 		}
 	}
 
@@ -163,6 +163,54 @@ func (g *GrantBurnDownHistory) Segments() []GrantBurnDownHistorySegment {
 	return g.segments
 }
 
+func (g GrantBurnDownHistory) ChunkByResets() []GrantBurnDownHistory {
+	if len(g.segments) == 0 {
+		return nil
+	}
+
+	chunks := make([]GrantBurnDownHistory, 0, 1)
+	current := GrantBurnDownHistory{
+		usageAtStart: g.usageAtStart,
+		segments:     make([]GrantBurnDownHistorySegment, 0, len(g.segments)),
+	}
+
+	for _, seg := range g.segments {
+		current.segments = append(current.segments, seg)
+		if seg.TerminationReasons.UsageReset {
+			chunks = append(chunks, current)
+			current = GrantBurnDownHistory{
+				usageAtStart: usageAtReset(seg.To),
+				segments:     make([]GrantBurnDownHistorySegment, 0, len(g.segments)),
+			}
+		}
+	}
+
+	if len(current.segments) > 0 {
+		chunks = append(chunks, current)
+	}
+
+	return chunks
+}
+
+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))
+		}
+	}
+
+	return total
+}
+
+func usageAtReset(at time.Time) balance.SnapshottedUsage {
+	return balance.SnapshottedUsage{
+		Since: at,
+		Usage: 0.0,
+	}
+}
+
 func (g *GrantBurnDownHistory) TotalUsageInHistory() float64 {
 	var total float64
 	for _, s := range g.segments {

openmeter/credit/engine/history_test.go

@@ -0,0 +1,118 @@
+package engine_test
+
+import (
+	"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/engine"
+	"github.com/openmeterio/openmeter/pkg/timeutil"
+)
+
+func TestGrantBurnDownHistory_ChunkByResets(t *testing.T) {
+	start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
+	usageAtStart := balance.SnapshottedUsage{
+		Since: start.Add(-time.Hour),
+		Usage: 7,
+	}
+
+	segment := func(idx int, grantUsage float64, reset bool) engine.GrantBurnDownHistorySegment {
+		from := start.Add(time.Duration(idx) * time.Hour)
+
+		return engine.GrantBurnDownHistorySegment{
+			ClosedPeriod: timeutil.ClosedPeriod{
+				From: from,
+				To:   from.Add(time.Hour),
+			},
+			TotalUsage: grantUsage + 1,
+			GrantUsages: []engine.GrantUsage{
+				{
+					GrantID: "grant-1",
+					Usage:   grantUsage,
+				},
+			},
+			TerminationReasons: engine.SegmentTerminationReason{
+				UsageReset: reset,
+			},
+		}
+	}
+
+	t.Run("Empty history returns no chunks", func(t *testing.T) {
+		history, err := engine.NewGrantBurnDownHistory(nil, usageAtStart)
+		require.NoError(t, err)
+
+		assert.Empty(t, history.ChunkByResets())
+	})
+
+	t.Run("History without resets stays as one chunk", func(t *testing.T) {
+		history, err := engine.NewGrantBurnDownHistory([]engine.GrantBurnDownHistorySegment{
+			segment(0, 10, false),
+			segment(1, 20, false),
+		}, usageAtStart)
+		require.NoError(t, err)
+
+		chunks := history.ChunkByResets()
+		require.Len(t, chunks, 1)
+		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)
+	})
+
+	t.Run("History is chunked after reset segments", func(t *testing.T) {
+		history, err := engine.NewGrantBurnDownHistory([]engine.GrantBurnDownHistorySegment{
+			segment(0, 10, false),
+			segment(1, 20, true),
+			segment(2, 30, false),
+			segment(3, 40, true),
+			segment(4, 50, false),
+		}, usageAtStart)
+		require.NoError(t, err)
+
+		chunks := history.ChunkByResets()
+		require.Len(t, chunks, 3)
+
+		assert.Len(t, chunks[0].Segments(), 2)
+		assert.Equal(t, 30.0, chunks[0].TotalGrantUsage().InexactFloat64())
+		assertChunkUsageAtStart(t, chunks[0], usageAtStart)
+
+		assert.Len(t, chunks[1].Segments(), 2)
+		assert.Equal(t, 70.0, chunks[1].TotalGrantUsage().InexactFloat64())
+		assertChunkUsageAtStart(t, chunks[1], balance.SnapshottedUsage{
+			Since: start.Add(2 * time.Hour),
+			Usage: 0,
+		})
+
+		assert.Len(t, chunks[2].Segments(), 1)
+		assert.Equal(t, 50.0, chunks[2].TotalGrantUsage().InexactFloat64())
+		assertChunkUsageAtStart(t, chunks[2], balance.SnapshottedUsage{
+			Since: start.Add(4 * time.Hour),
+			Usage: 0,
+		})
+	})
+
+	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)
+		require.NoError(t, err)
+
+		chunks := history.ChunkByResets()
+		require.Len(t, chunks, 1)
+		assert.Len(t, chunks[0].Segments(), 1)
+		assert.Equal(t, 10.0, chunks[0].TotalGrantUsage().InexactFloat64())
+	})
+}
+
+func assertChunkUsageAtStart(t *testing.T, history engine.GrantBurnDownHistory, expected balance.SnapshottedUsage) {
+	t.Helper()
+
+	usage, err := history.GetUsageInPeriodUntilSegment(0)
+	require.NoError(t, err)
+	assert.Equal(t, expected, usage)
+}