fix(credit): enforce transaction and snapshot consistency

openmeter/credit/adapter/balance_snapshot.go

@@ -8,10 +8,13 @@ import (
 	"entgo.io/ent/dialect/sql"
 
 	"github.com/openmeterio/openmeter/openmeter/credit/balance"
+	"github.com/openmeterio/openmeter/openmeter/credit/grant"
 	"github.com/openmeterio/openmeter/openmeter/ent/db"
 	db_balancesnapshot "github.com/openmeterio/openmeter/openmeter/ent/db/balancesnapshot"
+	db_entitlement "github.com/openmeterio/openmeter/openmeter/ent/db/entitlement"
 	"github.com/openmeterio/openmeter/pkg/clock"
 	"github.com/openmeterio/openmeter/pkg/framework/entutils"
+	"github.com/openmeterio/openmeter/pkg/framework/transaction"
 	"github.com/openmeterio/openmeter/pkg/models"
 )
 
@@ -27,16 +30,54 @@ func NewPostgresBalanceSnapshotRepo(db *db.Client) *balanceSnapshotRepo {
 }
 
 func (b *balanceSnapshotRepo) InvalidateAfter(ctx context.Context, owner models.NamespacedID, at time.Time) error {
-	return entutils.TransactingRepoWithNoValue(ctx, b, func(ctx context.Context, rep *balanceSnapshotRepo) error {
-		return rep.db.BalanceSnapshot.Update().
+	return transaction.RunWithNoValue(ctx, b, func(ctx context.Context) error {
+		return entutils.TransactingRepoWithNoValue(ctx, b, func(ctx context.Context, rep *balanceSnapshotRepo) error {
+			updated, err := rep.db.Entitlement.Update().
+				Where(
+					db_entitlement.ID(owner.ID),
+					db_entitlement.Namespace(owner.Namespace),
+				).
+				AddBalanceSnapshotInvalidationVersion(1).
+				Save(ctx)
+			if err != nil {
+				return err
+			}
+
+			if updated == 0 {
+				return grant.NewOwnerNotFoundError(owner, "entitlement")
+			}
+
+			return rep.db.BalanceSnapshot.Update().
+				Where(
+					db_balancesnapshot.OwnerID(owner.ID),
+					db_balancesnapshot.Namespace(owner.Namespace),
+					db_balancesnapshot.AtGT(at),
+					db_balancesnapshot.DeletedAtIsNil(),
+				).
+				SetDeletedAt(clock.Now()).
+				Exec(ctx)
+		})
+	})
+}
+
+func (b *balanceSnapshotRepo) GetInvalidationVersion(ctx context.Context, owner models.NamespacedID) (balance.SnapshotInvalidationVersion, error) {
+	return entutils.TransactingRepo(ctx, b, func(ctx context.Context, rep *balanceSnapshotRepo) (balance.SnapshotInvalidationVersion, error) {
+		entitlement, err := rep.db.Entitlement.Query().
 			Where(
-				db_balancesnapshot.OwnerID(owner.ID),
-				db_balancesnapshot.Namespace(owner.Namespace),
-				db_balancesnapshot.AtGT(at),
-				db_balancesnapshot.DeletedAtIsNil(),
+				db_entitlement.ID(owner.ID),
+				db_entitlement.Namespace(owner.Namespace),
 			).
-			SetDeletedAt(clock.Now()).
-			Exec(ctx)
+			Select(db_entitlement.FieldBalanceSnapshotInvalidationVersion).
+			Only(ctx)
+		if err != nil {
+			if db.IsNotFound(err) {
+				return 0, grant.NewOwnerNotFoundError(owner, "entitlement")
+			}
+
+			return 0, err
+		}
+
+		return balance.SnapshotInvalidationVersion(entitlement.BalanceSnapshotInvalidationVersion), nil
 	})
 }
 

openmeter/credit/adapter/balance_snapshot_test.go

@@ -0,0 +1,102 @@
+package adapter_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/oklog/ulid/v2"
+	"github.com/stretchr/testify/require"
+
+	creditadapter "github.com/openmeterio/openmeter/openmeter/credit/adapter"
+	"github.com/openmeterio/openmeter/openmeter/credit/balance"
+	"github.com/openmeterio/openmeter/openmeter/entitlement"
+	entitlementadapter "github.com/openmeterio/openmeter/openmeter/entitlement/adapter"
+	"github.com/openmeterio/openmeter/openmeter/streaming"
+	"github.com/openmeterio/openmeter/openmeter/testutils"
+	"github.com/openmeterio/openmeter/pkg/framework/transaction"
+	"github.com/openmeterio/openmeter/pkg/models"
+)
+
+func TestBalanceSnapshotInvalidationVersionIsAtomicWithInvalidation(t *testing.T) {
+	// given: an entitlement with a saved balance snapshot
+	testDB := testutils.InitPostgresDB(t, testutils.PostgresDBStateEntMigrated)
+	dbClient := testDB.EntDriver.Client()
+	t.Cleanup(func() {
+		_ = dbClient.Close()
+		testDB.Close(t)
+	})
+
+	namespace := ulid.Make().String()
+	customer, err := dbClient.Customer.Create().
+		SetNamespace(namespace).
+		SetName("Customer").
+		SetKey("customer-" + ulid.Make().String()).
+		Save(t.Context())
+	require.NoError(t, err)
+
+	feature, err := dbClient.Feature.Create().
+		SetNamespace(namespace).
+		SetName("Feature").
+		SetKey("feature-" + ulid.Make().String()).
+		Save(t.Context())
+	require.NoError(t, err)
+
+	entitlementRepo := entitlementadapter.NewPostgresEntitlementRepo(dbClient)
+	ent, err := entitlementRepo.CreateEntitlement(t.Context(), entitlement.CreateEntitlementRepoInputs{
+		Namespace:  namespace,
+		FeatureID:  feature.ID,
+		FeatureKey: feature.Key,
+		UsageAttribution: streaming.CustomerUsageAttribution{
+			ID:  customer.ID,
+			Key: &customer.Key,
+		},
+		EntitlementType: entitlement.EntitlementTypeMetered,
+	})
+	require.NoError(t, err)
+
+	owner := models.NamespacedID{Namespace: namespace, ID: ent.ID}
+	snapshotAt := time.Now().UTC().Truncate(time.Minute)
+	repo := creditadapter.NewPostgresBalanceSnapshotRepo(dbClient)
+	require.NoError(t, repo.Save(t.Context(), owner, []balance.Snapshot{
+		balance.NewStartingSnapshot(nil, snapshotAt),
+	}))
+
+	version, err := repo.GetInvalidationVersion(t.Context(), owner)
+	require.NoError(t, err)
+	require.Zero(t, version)
+
+	// when: the outer transaction fails after invalidation
+	rollbackErr := errors.New("rollback")
+	err = transaction.RunWithNoValue(t.Context(), repo, func(ctx context.Context) error {
+		if err := repo.InvalidateAfter(ctx, owner, snapshotAt.Add(-time.Minute)); err != nil {
+			return err
+		}
+
+		version, err := repo.GetInvalidationVersion(ctx, owner)
+		require.NoError(t, err)
+		require.Equal(t, balance.SnapshotInvalidationVersion(1), version)
+
+		return rollbackErr
+	})
+	require.ErrorIs(t, err, rollbackErr)
+
+	// then: the version increment and snapshot invalidation both roll back
+	version, err = repo.GetInvalidationVersion(t.Context(), owner)
+	require.NoError(t, err)
+	require.Zero(t, version)
+
+	_, err = repo.GetLatestValidAt(t.Context(), owner, snapshotAt)
+	require.NoError(t, err)
+
+	// and: a committed invalidation advances the durable fence
+	require.NoError(t, repo.InvalidateAfter(t.Context(), owner, snapshotAt.Add(-time.Minute)))
+	version, err = repo.GetInvalidationVersion(t.Context(), owner)
+	require.NoError(t, err)
+	require.Equal(t, balance.SnapshotInvalidationVersion(1), version)
+
+	_, err = repo.GetLatestValidAt(t.Context(), owner, snapshotAt)
+	var noSnapshot *balance.NoSavedBalanceForOwnerError
+	require.ErrorAs(t, err, &noSnapshot)
+}

openmeter/credit/adapter/grant.go

@@ -35,7 +35,6 @@ func NewPostgresGrantRepo(db *db.Client) grant.Repo {
 }
 
 func (g *grantDBADapter) CreateGrant(ctx context.Context, grant grant.RepoCreateInput) (*grant.Grant, error) {
-	// TODO: transaction and locking
 	if err := ensureEntitlementOwnerExists(ctx, g.db, models.NamespacedID{
 		Namespace: grant.Namespace,
 		ID:        grant.OwnerID,
@@ -82,7 +81,6 @@ func (g *grantDBADapter) DeleteOwnerGrants(ctx context.Context, ownerID models.N
 
 // translates to a delete
 func (g *grantDBADapter) VoidGrant(ctx context.Context, grantID models.NamespacedID, at time.Time) error {
-	// TODO: transaction and locking
 	command := g.db.Grant.Update().
 		SetVoidedAt(at).
 		Where(db_grant.ID(grantID.ID), db_grant.Namespace(grantID.Namespace))

openmeter/credit/balance.go

@@ -41,7 +41,13 @@ type BalanceConnector interface {
 var _ BalanceConnector = &connector{}
 
 // GetBalanceSinceSnapshot returns the result of the engine.Run since a given snapshot.
-func (m *connector) getBalanceSinceSnapshot(ctx context.Context, ownerID models.NamespacedID, snap balance.Snapshot, at time.Time) (engine.RunResult, error) {
+func (m *connector) getBalanceSinceSnapshot(
+	ctx context.Context,
+	ownerID models.NamespacedID,
+	snap balance.Snapshot,
+	at time.Time,
+	snapshotInvalidationVersion balance.SnapshotInvalidationVersion,
+) (engine.RunResult, error) {
 	ctx, span := m.Tracer.Start(ctx, "credit.GetBalanceSinceSnapshot", cTrace.WithOwner(ownerID), trace.WithAttributes(attribute.String("at", at.String())))
 	defer span.End()
 
@@ -110,11 +116,12 @@ func (m *connector) getBalanceSinceSnapshot(ctx context.Context, ownerID models.
 	// TODO: it might be the case that we don't save any snapshots as they require a history breakpoint. To solve this,
 	// we should introduce artificial history breakpoints in the engine, but that would result in more streaming.Query calls, so first lets improve the visibility of what's happening.
 	if err := m.snapshotEngineResult(ctx, snapshotParams{
-		grants:     grants,
-		owner:      ownerID,
-		notAfter:   m.getSnapshotNotAfter(periodStart, clock.Now()),
-		meter:      owner.Meter,
-		unitConfig: owner.UnitConfig,
+		grants:                      grants,
+		owner:                       ownerID,
+		notAfter:                    m.getSnapshotNotAfter(periodStart, clock.Now()),
+		meter:                       owner.Meter,
+		unitConfig:                  owner.UnitConfig,
+		snapshotInvalidationVersion: snapshotInvalidationVersion,
 	}, result); err != nil {
 		return def, fmt.Errorf("failed to snapshot engine result: %w", err)
 	}
@@ -137,13 +144,18 @@ func (m *connector) GetBalanceAt(ctx context.Context, ownerID models.NamespacedI
 		at = trunc.Add(time.Minute)
 	}
 
+	snapshotInvalidationVersion, err := m.BalanceSnapshotService.GetInvalidationVersion(ctx, ownerID)
+	if err != nil {
+		return def, fmt.Errorf("failed to get snapshot invalidation version for owner %s: %w", ownerID.ID, err)
+	}
+
 	// get last valid grantbalances
 	snap, err := m.GetLastValidSnapshotAt(ctx, ownerID, at)
 	if err != nil {
 		return def, err
 	}
 
-	return m.getBalanceSinceSnapshot(ctx, ownerID, snap, at)
+	return m.getBalanceSinceSnapshot(ctx, ownerID, snap, at, snapshotInvalidationVersion)
 }
 
 func (m *connector) GetBalanceForPeriod(ctx context.Context, ownerID models.NamespacedID, period timeutil.ClosedPeriod) (engine.RunResult, error) {

openmeter/credit/balance/repository.go

@@ -8,8 +8,11 @@ import (
 	"github.com/openmeterio/openmeter/pkg/models"
 )
 
+type SnapshotInvalidationVersion uint64
+
 type SnapshotRepo interface {
 	InvalidateAfter(ctx context.Context, owner models.NamespacedID, at time.Time) error
+	GetInvalidationVersion(ctx context.Context, owner models.NamespacedID) (SnapshotInvalidationVersion, error)
 	// 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

openmeter/credit/balance/service.go

@@ -9,6 +9,7 @@ import (
 
 type SnapshotService interface {
 	InvalidateAfter(ctx context.Context, owner models.NamespacedID, at time.Time) error
+	GetInvalidationVersion(ctx context.Context, owner models.NamespacedID) (SnapshotInvalidationVersion, error)
 	GetLatestValidAt(ctx context.Context, owner models.NamespacedID, at time.Time) (Snapshot, error)
 	Save(ctx context.Context, owner models.NamespacedID, balances []Snapshot) error
 	// To make sure repo doesn't implement the service interface
@@ -35,6 +36,10 @@ func (s *service) InvalidateAfter(ctx context.Context, owner models.NamespacedID
 	return s.Repo.InvalidateAfter(ctx, owner, at)
 }
 
+func (s *service) GetInvalidationVersion(ctx context.Context, owner models.NamespacedID) (SnapshotInvalidationVersion, error) {
+	return s.Repo.GetInvalidationVersion(ctx, owner)
+}
+
 func (s *service) GetLatestValidAt(ctx context.Context, owner models.NamespacedID, at time.Time) (Snapshot, error) {
 	return s.Repo.GetLatestValidAt(ctx, owner, at)
 }

openmeter/credit/grant.go

@@ -73,6 +73,11 @@ func (m *connector) CreateGrant(ctx context.Context, ownerID models.NamespacedID
 			input.Recurrence.Anchor = input.Recurrence.Anchor.Truncate(granularity)
 		}
 
+		err = m.OwnerConnector.LockOwnerForTx(ctx, ownerID, true)
+		if err != nil {
+			return nil, err
+		}
+
 		periodStart, err := m.OwnerConnector.GetUsagePeriodStartAt(ctx, ownerID, clock.Now())
 		if err != nil {
 			return nil, err
@@ -82,11 +87,6 @@ func (m *connector) CreateGrant(ctx context.Context, ownerID models.NamespacedID
 			return nil, models.NewGenericValidationError(fmt.Errorf("grant effective date %s is before the current usage period %s", input.EffectiveAt, periodStart))
 		}
 
-		err = m.OwnerConnector.LockOwnerForTx(ctx, ownerID, true)
-		if err != nil {
-			return nil, err
-		}
-
 		repoInp := grant.RepoCreateInput{
 			OwnerID:          ownerID.ID,
 			Namespace:        ownerID.Namespace,
@@ -137,10 +137,6 @@ func (m *connector) VoidGrant(ctx context.Context, grantID models.NamespacedID,
 		return err
 	}
 
-	if g.VoidedAt != nil {
-		return models.NewGenericValidationError(fmt.Errorf("grant already voided"))
-	}
-
 	// Resolve the void time: use provided `at` or default to now.
 	voidAt := clock.Now()
 	if at != nil {
@@ -156,33 +152,42 @@ func (m *connector) VoidGrant(ctx context.Context, grantID models.NamespacedID,
 
 	ownerID := models.NamespacedID{Namespace: grantID.Namespace, ID: g.OwnerID}
 
-	// Validate: at must be within the current usage period
-	periodStart, err := m.OwnerConnector.GetUsagePeriodStartAt(ctx, ownerID, clock.Now())
-	if err != nil {
-		return err
-	}
-
-	if voidAt.Before(periodStart) {
-		return models.NewGenericValidationError(fmt.Errorf("void time %s is before the current usage period start %s", voidAt, periodStart))
-	}
-
 	_, err = transaction.Run(ctx, m.GrantRepo, func(ctx context.Context) (*interface{}, error) {
 		tx, err := entutils.GetDriverFromContext(ctx)
 		if err != nil {
 			return nil, err
 		}
 
-		owner, err := m.OwnerConnector.DescribeOwner(ctx, ownerID)
+		err = m.OwnerConnector.LockOwnerForTx(ctx, ownerID, true)
 		if err != nil {
 			return nil, err
 		}
 
-		err = m.OwnerConnector.LockOwnerForTx(ctx, ownerID, true)
+		txGrantRepo := m.GrantRepo.WithTx(ctx, tx)
+		g, err = txGrantRepo.GetGrant(ctx, grantID)
+		if err != nil {
+			return nil, err
+		}
+
+		if g.VoidedAt != nil {
+			return nil, models.NewGenericValidationError(fmt.Errorf("grant already voided"))
+		}
+
+		periodStart, err := m.OwnerConnector.GetUsagePeriodStartAt(ctx, ownerID, clock.Now())
+		if err != nil {
+			return nil, err
+		}
+
+		if voidAt.Before(periodStart) {
+			return nil, models.NewGenericValidationError(fmt.Errorf("void time %s is before the current usage period start %s", voidAt, periodStart))
+		}
+
+		owner, err := m.OwnerConnector.DescribeOwner(ctx, ownerID)
 		if err != nil {
 			return nil, err
 		}
 
-		err = m.GrantRepo.WithTx(ctx, tx).VoidGrant(ctx, grantID, voidAt)
+		err = txGrantRepo.VoidGrant(ctx, grantID, voidAt)
 		if err != nil {
 			return nil, err
 		}