Update various balance call sites

ocp/balance/calculator.go

@@ -395,6 +395,93 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide
 	return res, nil
 }
 
+// BalanceWithUsdCostBasis holds a token account's quark balance and USD cost
+// basis, in balance.UsdQuarksPerUnit.
+type BalanceWithUsdCostBasis struct {
+	Quarks       uint64
+	UsdCostBasis int64
+}
+
+// BatchCalculateWithUsdCostBasisFromCache calculates balances and USD cost
+// bases for a set of account records. When ledger reads are enabled, both
+// values for an account come from the same balance record read, so they are
+// guaranteed consistent with each other. Accounts without a backfilled
+// record fall back to the legacy aggregates, which read the two values from
+// separate sources.
+//
+// Note: Use this method when calculating balances for accounts that are managed by
+// Code (ie. Timelock account) and operate within the L2 system.
+func BatchCalculateWithUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, accountRecordsBatch ...*common.AccountRecords) (map[string]*BalanceWithUsdCostBasis, error) {
+	tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateWithUsdCostBasisFromCache")
+	defer tracer.End()
+
+	var tokenAccounts []string
+	for _, accountRecords := range accountRecordsBatch {
+		if !accountRecords.IsTimelock() || !accountRecords.IsManagedByCode(ctx) {
+			tracer.OnError(ErrNotManagedByCode)
+			return nil, ErrNotManagedByCode
+		}
+		tokenAccounts = append(tokenAccounts, accountRecords.General.TokenAccount)
+	}
+
+	balanceRecords := make(map[string]*balance.Record)
+	if enableLedgerReads.Get(ctx) {
+		var err error
+		balanceRecords, err = data.GetBalanceBatch(ctx, tokenAccounts...)
+		if err != nil {
+			tracer.OnError(err)
+			return nil, err
+		}
+	}
+
+	res := make(map[string]*BalanceWithUsdCostBasis, len(tokenAccounts))
+	var remaining []string
+	for _, tokenAccount := range tokenAccounts {
+		balanceRecord, ok := balanceRecords[tokenAccount]
+		if !ok || !balanceRecord.IsBackfilled {
+			remaining = append(remaining, tokenAccount)
+			continue
+		}
+
+		quarks, err := quarksFromRecord(balanceRecord)
+		if err != nil {
+			tracer.OnError(err)
+			return nil, err
+		}
+		res[tokenAccount] = &BalanceWithUsdCostBasis{
+			Quarks:       quarks,
+			UsdCostBasis: balanceRecord.UsdCostBasis,
+		}
+	}
+
+	if len(remaining) == 0 {
+		return res, nil
+	}
+
+	legacyQuarks, err := CalculateBatch(
+		ctx,
+		remaining,
+		NetBalanceFromIntentActionsBatch(ctx, data),
+		FundingFromExternalDepositsBatch(ctx, data),
+	)
+	if err != nil {
+		tracer.OnError(err)
+		return nil, err
+	}
+	for _, tokenAccount := range remaining {
+		usdCostBasis, err := legacyUsdCostBasis(ctx, data, tokenAccount)
+		if err != nil {
+			tracer.OnError(err)
+			return nil, err
+		}
+		res[tokenAccount] = &BalanceWithUsdCostBasis{
+			Quarks:       legacyQuarks[tokenAccount],
+			UsdCostBasis: usdCostBasis,
+		}
+	}
+	return res, nil
+}
+
 // CalculateUsdCostBasisFromCache calculates a token account's USD cost basis,
 // in balance.UsdQuarksPerUnit, using cached values.
 //

ocp/balance/calculator_test.go

@@ -419,6 +419,73 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) {
 	assert.Equal(t, expected, balanceByAccount)
 }
 
+func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) {
+	env := setupBalanceTestEnv(t)
+	enableLedgerReadsForTest(t)
+
+	vmConfig := testutil.NewRandomVmConfig(t, true)
+	backfilledOwner := testutil.NewRandomAccount(t)
+	backfilledAccount, err := backfilledOwner.ToTimelockVault(vmConfig)
+	require.NoError(t, err)
+	legacyOwner := testutil.NewRandomAccount(t)
+	legacyAccount, err := legacyOwner.ToTimelockVault(vmConfig)
+	require.NoError(t, err)
+
+	externalAccount := testutil.NewRandomAccount(t)
+
+	data := &balanceTestData{
+		vmConfig:  vmConfig,
+		codeUsers: []*common.Account{backfilledOwner, legacyOwner},
+		transactions: []balanceTestTransaction{
+			{source: externalAccount, destination: backfilledAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized},
+			{source: externalAccount, destination: legacyAccount, quantity: 33, transactionState: transaction.ConfirmationFinalized},
+		},
+	}
+
+	setupBalanceTestData(t, env, data)
+
+	// Both values come from the same record for a backfilled account, even
+	// where it disagrees with history
+	require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{
+		TokenAccount: backfilledAccount.PublicKey().ToBase58(),
+		OwnerAccount: backfilledOwner.PublicKey().ToBase58(),
+		MintAccount:  vmConfig.Mint.PublicKey().ToBase58(),
+		Quarks:       42,
+		UsdCostBasis: 4_200_000,
+		IsOpen:       true,
+		IsLocked:     true,
+		IsBackfilled: true,
+	}))
+
+	accountRecordsBatch := make([]*common.AccountRecords, 0)
+	for _, tokenAccount := range []*common.Account{backfilledAccount, legacyAccount} {
+		generalRecord, err := env.data.GetAccountInfoByTokenAddress(env.ctx, tokenAccount.PublicKey().ToBase58())
+		require.NoError(t, err)
+		timelockRecord, err := env.data.GetTimelockByVault(env.ctx, tokenAccount.PublicKey().ToBase58())
+		require.NoError(t, err)
+		accountRecordsBatch = append(accountRecordsBatch, &common.AccountRecords{
+			General:  generalRecord,
+			Timelock: timelockRecord,
+		})
+	}
+
+	res, err := BatchCalculateWithUsdCostBasisFromCache(env.ctx, env.data, accountRecordsBatch...)
+	require.NoError(t, err)
+	require.Len(t, res, 2)
+
+	cached := res[backfilledAccount.PublicKey().ToBase58()]
+	require.NotNil(t, cached)
+	assert.EqualValues(t, 42, cached.Quarks)
+	assert.EqualValues(t, 4_200_000, cached.UsdCostBasis)
+
+	// An account without a backfilled record falls back to the legacy
+	// aggregates for both values
+	cached = res[legacyAccount.PublicKey().ToBase58()]
+	require.NotNil(t, cached)
+	assert.EqualValues(t, 33, cached.Quarks)
+	assert.EqualValues(t, 0, cached.UsdCostBasis) // deposits aren't primary-owner intents in this fixture
+}
+
 func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) {
 	env := setupBalanceTestEnv(t)
 	enableLedgerReadsForTest(t)
@@ -471,6 +538,7 @@ func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) {
 
 func TestDefaultCalculationMethods_BalanceRecordReadsDisabled(t *testing.T) {
 	env := setupBalanceTestEnv(t)
+	disableLedgerReadsForTest(t)
 
 	vmConfig := testutil.NewRandomVmConfig(t, true)
 	owner := testutil.NewRandomAccount(t)
@@ -603,6 +671,14 @@ func enableLedgerReadsForTest(t *testing.T) {
 	})
 }
 
+func disableLedgerReadsForTest(t *testing.T) {
+	previous := enableLedgerReads
+	enableLedgerReads = wrapper.NewBoolConfig(memory.NewConfig(false), defaultEnableLedgerReads)
+	t.Cleanup(func() {
+		enableLedgerReads = previous
+	})
+}
+
 type balanceTestEnv struct {
 	ctx  context.Context
 	data ocp_data.Provider

ocp/balance/config.go

@@ -15,8 +15,8 @@ const (
 	// to the ledger at all. When disabled, it is a no-op.
 	EnableLedgerWritesConfigEnvName = "BALANCE_ENABLE_LEDGER_WRITES"
 
-	defaultEnableLedgerReads  = false
-	defaultEnableLedgerWrites = false
+	defaultEnableLedgerReads  = true
+	defaultEnableLedgerWrites = true
 )
 
 var (

ocp/balance/ledger.go

@@ -14,6 +14,12 @@ import (
 // ledger doesn't track.
 var ErrUntrackedAccount = errors.New("account is not tracked by the balance ledger")
 
+// LedgerReadsEnabled reports whether backfilled ledger records are the
+// authoritative source for balance reads.
+func LedgerReadsEnabled(ctx context.Context) bool {
+	return enableLedgerReads.Get(ctx)
+}
+
 // LedgerWritesEnabled reports whether the ledger is being written to.
 // Callers use it to skip building deltas entirely when writes are disabled,
 // since builders reject flows the ledger doesn't support.

ocp/balance/ledger_test.go

@@ -20,6 +20,7 @@ import (
 func TestApplyDeltasInTx_WritesDisabled(t *testing.T) {
 	ctx := context.Background()
 	data := ocp_data.NewTestDataProvider()
+	disableLedgerWritesForTest(t)
 
 	source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY)
 
@@ -153,9 +154,12 @@ func TestCreateRecordInTx(t *testing.T) {
 
 	// Disabled writes are a no-op
 	primary := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY)
-	require.NoError(t, CreateRecordInTx(ctx, data, primary))
-	_, err := data.GetBalance(ctx, primary.TokenAccount)
-	assert.Equal(t, balance.ErrRecordNotFound, err)
+	func() {
+		disableLedgerWritesForTest(t)
+		require.NoError(t, CreateRecordInTx(ctx, data, primary))
+		_, err := data.GetBalance(ctx, primary.TokenAccount)
+		assert.Equal(t, balance.ErrRecordNotFound, err)
+	}()
 
 	enableLedgerWritesForTest(t)
 
@@ -210,6 +214,14 @@ func newLedgerTestAccountInfo(t *testing.T, ctx context.Context, data ocp_data.P
 	return record
 }
 
+func disableLedgerWritesForTest(t *testing.T) {
+	previous := enableLedgerWrites
+	enableLedgerWrites = wrapper.NewBoolConfig(memory.NewConfig(false), defaultEnableLedgerWrites)
+	t.Cleanup(func() {
+		enableLedgerWrites = previous
+	})
+}
+
 func enableLedgerWritesForTest(t *testing.T) {
 	previous := enableLedgerWrites
 	enableLedgerWrites = wrapper.NewBoolConfig(memory.NewConfig(true), defaultEnableLedgerWrites)

ocp/data/balance/memory/store.go

@@ -14,19 +14,15 @@ type store struct {
 	balanceRecords               []*balance.Record
 	balanceRecordsByTokenAccount map[string]*balance.Record
 
-	cachedBalanceVersionsByAccount map[string]uint64
-	closedAccounts                 map[string]any
-	externalCheckpointRecords      []*balance.ExternalCheckpointRecord
+	externalCheckpointRecords []*balance.ExternalCheckpointRecord
 
 	last uint64
 }
 
 // New returns a new in memory balance.Store
 func New() balance.Store {
 	return &store{
-		balanceRecordsByTokenAccount:   make(map[string]*balance.Record),
-		cachedBalanceVersionsByAccount: make(map[string]uint64),
-		closedAccounts:                 make(map[string]any),
+		balanceRecordsByTokenAccount: make(map[string]*balance.Record),
 	}
 }
 
@@ -139,6 +135,20 @@ func (s *store) GetAllLockedByMint(_ context.Context, mint string, minQuarks int
 	return res, nil
 }
 
+// CountLockedByMint implements balance.Store.CountLockedByMint
+func (s *store) CountLockedByMint(_ context.Context, mint string, minQuarks int64) (uint64, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var res uint64
+	for _, item := range s.balanceRecordsByTokenAccount {
+		if item.MintAccount == mint && item.Quarks >= minQuarks && item.IsLocked && item.IsBackfilled {
+			res++
+		}
+	}
+	return res, nil
+}
+
 // MarkAsUnlocked implements balance.Store.MarkAsUnlocked
 func (s *store) MarkAsUnlocked(_ context.Context, tokenAccount string) error {
 	s.mu.Lock()
@@ -218,6 +228,9 @@ func applyDelta(item *balance.Record, delta *balance.Delta) error {
 		item.Quarks += int64(delta.Quarks)
 		item.UsdCostBasis += delta.UsdCostBasis
 	case balance.DeltaDebit:
+		if enforce && !item.IsOpen {
+			return balance.ErrAccountClosed
+		}
 		if enforce && item.Quarks < int64(delta.Quarks) {
 			return balance.ErrInsufficientBalance
 		}
@@ -313,8 +326,6 @@ func (s *store) reset() {
 
 	s.balanceRecords = nil
 	s.balanceRecordsByTokenAccount = make(map[string]*balance.Record)
-	s.cachedBalanceVersionsByAccount = make(map[string]uint64)
-	s.closedAccounts = make(map[string]any)
 	s.externalCheckpointRecords = nil
 	s.last = 0
 }