From 3bd0fe7cf8507f0f6b3de55a1a1f74ff9f7c0afc Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 11:49:24 -0400 Subject: [PATCH 01/26] Implement new balance table --- ocp/data/balance/memory/store.go | 278 ++++++++---- ocp/data/balance/memory/store_legacy.go | 131 ++++++ ocp/data/balance/postgres/model.go | 385 ++++++++++------- ocp/data/balance/postgres/model_legacy.go | 199 +++++++++ ocp/data/balance/postgres/store.go | 102 +++-- ocp/data/balance/postgres/store_legacy.go | 53 +++ ocp/data/balance/postgres/store_test.go | 23 + ocp/data/balance/record.go | 170 ++++++++ ocp/data/balance/store.go | 102 +++++ ocp/data/balance/tests/tests.go | 489 +++++++++++++++++++++- ocp/data/internal.go | 32 ++ 11 files changed, 1719 insertions(+), 245 deletions(-) create mode 100644 ocp/data/balance/memory/store_legacy.go create mode 100644 ocp/data/balance/postgres/model_legacy.go create mode 100644 ocp/data/balance/postgres/store_legacy.go create mode 100644 ocp/data/balance/record.go diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index bfe25f9..00cec4a 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -5,152 +5,286 @@ import ( "sync" "time" + "github.com/code-payments/ocp-server/database/query" "github.com/code-payments/ocp-server/ocp/data/balance" ) type store struct { - mu sync.Mutex + mu sync.Mutex + balanceRecords []*balance.Record + balanceRecordsByTokenAccount map[string]*balance.Record + cachedBalanceVersionsByAccount map[string]uint64 closedAccounts map[string]any externalCheckpointRecords []*balance.ExternalCheckpointRecord - last uint64 + + 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), } } -// GetCachedVersion implements balance.Store.GetCachedVersion -func (s *store) GetCachedVersion(_ context.Context, account string) (uint64, error) { +// Create implements balance.Store.Create +func (s *store) Create(_ context.Context, record *balance.Record) error { + if err := record.Validate(); err != nil { + return err + } + s.mu.Lock() defer s.mu.Unlock() - current, ok := s.cachedBalanceVersionsByAccount[account] - if !ok { - return 0, nil + if _, ok := s.balanceRecordsByTokenAccount[record.TokenAccount]; ok { + return balance.ErrRecordExists } - return current, nil + + s.last++ + record.Id = s.last + record.UpdatedAt = time.Now() + + cloned := record.Clone() + s.balanceRecordsByTokenAccount[record.TokenAccount] = &cloned + s.balanceRecords = append(s.balanceRecords, &cloned) + + return nil } -// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion -func (s *store) AdvanceCachedVersion(_ context.Context, account string, currentVersion uint64) error { +// Get implements balance.Store.Get +func (s *store) Get(_ context.Context, tokenAccount string) (*balance.Record, error) { s.mu.Lock() defer s.mu.Unlock() - actualVersion, ok := s.cachedBalanceVersionsByAccount[account] + item, ok := s.balanceRecordsByTokenAccount[tokenAccount] if !ok { - if currentVersion != 0 { - return balance.ErrStaleCachedBalanceVersion - } - - s.cachedBalanceVersionsByAccount[account] = 1 - - return nil + return nil, balance.ErrRecordNotFound } + cloned := item.Clone() + return &cloned, nil +} + +// GetBatch implements balance.Store.GetBatch +func (s *store) GetBatch(_ context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) { + s.mu.Lock() + defer s.mu.Unlock() - if actualVersion != currentVersion { - return balance.ErrStaleCachedBalanceVersion + res := make(map[string]*balance.Record) + for _, tokenAccount := range tokenAccounts { + item, ok := s.balanceRecordsByTokenAccount[tokenAccount] + if !ok { + continue + } + cloned := item.Clone() + res[tokenAccount] = &cloned } + return res, nil +} - s.cachedBalanceVersionsByAccount[account]++ +// GetAllByOwner implements balance.Store.GetAllByOwner +func (s *store) GetAllByOwner(_ context.Context, owner string) ([]*balance.Record, error) { + s.mu.Lock() + defer s.mu.Unlock() - return nil + return s.filter(func(item *balance.Record) bool { + return item.OwnerAccount == owner + }) } -// CheckNotClosed implements balance.Store.CheckNotClosed -func (s *store) CheckNotClosed(ctx context.Context, account string) error { +// GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint +func (s *store) GetAllByOwnerAndMint(_ context.Context, owner, mint string) ([]*balance.Record, error) { s.mu.Lock() defer s.mu.Unlock() - if _, ok := s.closedAccounts[account]; ok { - return balance.ErrAccountClosed - } - - return nil + return s.filter(func(item *balance.Record) bool { + return item.OwnerAccount == owner && item.MintAccount == mint + }) } -// MarkAsClosed implements balance.Store.MarkAsClosed -func (s *store) MarkAsClosed(ctx context.Context, account string) error { +// GetAllByMint implements balance.Store.GetAllByMint +func (s *store) GetAllByMint(_ context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { s.mu.Lock() defer s.mu.Unlock() - s.closedAccounts[account] = true + res, err := s.filter(func(item *balance.Record) bool { + if item.MintAccount != mint || item.Quarks < minQuarks { + return false + } + if len(cursor) > 0 { + if direction == query.Ascending && item.Id <= cursor.ToUint64() { + return false + } + if direction == query.Descending && item.Id >= cursor.ToUint64() { + return false + } + } + return true + }) + if err != nil { + return nil, err + } - return nil + if direction == query.Descending { + for i, j := 0, len(res)-1; i < j; i, j = i+1, j-1 { + res[i], res[j] = res[j], res[i] + } + } + + if limit > 0 && uint64(len(res)) > limit { + res = res[:limit] + } + return res, nil } -// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint -func (s *store) SaveExternalCheckpoint(_ context.Context, data *balance.ExternalCheckpointRecord) error { - if err := data.Validate(); err != nil { - return err +// ApplyDeltas implements balance.Store.ApplyDeltas +func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { + for _, delta := range deltas { + if err := delta.Validate(); err != nil { + return err + } } + sorted := make([]*balance.Delta, len(deltas)) + copy(sorted, deltas) + balance.SortDeltas(sorted) + s.mu.Lock() defer s.mu.Unlock() - s.last++ - if item := s.findExternalCheckpoint(data); item != nil { - if data.SlotCheckpoint <= item.SlotCheckpoint { - return balance.ErrStaleCheckpoint + // Apply to copies first so a failure part way through leaves the store + // untouched, matching the transactional behaviour of the DB store. + updated := make(map[string]*balance.Record) + for _, delta := range sorted { + item, ok := updated[delta.TokenAccount] + if !ok { + original, ok := s.balanceRecordsByTokenAccount[delta.TokenAccount] + if !ok { + continue // Not an account we track + } + cloned := original.Clone() + item = &cloned + updated[delta.TokenAccount] = item } - item.SlotCheckpoint = data.SlotCheckpoint - item.Quarks = data.Quarks - item.LastUpdatedAt = time.Now() - item.CopyTo(data) - } else { - if data.Id == 0 { - data.Id = s.last + if err := applyDelta(item, delta); err != nil { + return err } - data.LastUpdatedAt = time.Now() - c := data.Clone() - s.externalCheckpointRecords = append(s.externalCheckpointRecords, &c) } + now := time.Now() + for tokenAccount, item := range updated { + item.UpdatedAt = now + item.CopyTo(s.balanceRecordsByTokenAccount[tokenAccount]) + } return nil } -// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint -func (s *store) GetExternalCheckpoint(_ context.Context, account string) (*balance.ExternalCheckpointRecord, error) { - s.mu.Lock() - defer s.mu.Unlock() +func applyDelta(item *balance.Record, delta *balance.Delta) error { + enforce := item.IsBackfilled - if item := s.findExternalCheckpointByTokenAccount(account); item != nil { - cloned := item.Clone() - return &cloned, nil + switch delta.Kind { + case balance.DeltaCredit: + if enforce && !item.IsOpen { + return balance.ErrAccountClosed + } + item.Quarks += int64(delta.Quarks) + item.UsdCostBasis += delta.UsdCostBasis + case balance.DeltaDebit: + if enforce && item.Quarks < int64(delta.Quarks) { + return balance.ErrInsufficientBalance + } + item.Quarks -= int64(delta.Quarks) + item.UsdCostBasis -= delta.UsdCostBasis + case balance.DeltaDrain: + if enforce { + if !item.IsOpen { + return balance.ErrAccountClosed + } + if item.Quarks != int64(delta.Quarks) { + return balance.ErrBalanceChanged + } + item.Quarks = 0 + item.UsdCostBasis = 0 + } else { + item.Quarks -= int64(delta.Quarks) + item.UsdCostBasis -= delta.UsdCostBasis + } + item.IsOpen = false + case balance.DeltaClose: + if enforce { + if !item.IsOpen { + return balance.ErrAccountClosed + } + if item.Quarks != 0 { + return balance.ErrBalanceChanged + } + } + item.IsOpen = false } - return nil, balance.ErrCheckpointNotFound + return nil } -func (s *store) findExternalCheckpoint(data *balance.ExternalCheckpointRecord) *balance.ExternalCheckpointRecord { - for _, item := range s.externalCheckpointRecords { - if item.Id == data.Id { - return item - } - if data.TokenAccount == item.TokenAccount { - return item - } +// Backfill implements balance.Store.Backfill +// +// Note: The lock is released while fn runs, since fn reads from other stores +// sharing the provider. Unlike the DB store, this doesn't block concurrent +// deltas, which tests don't exercise against a backfill. +func (s *store) Backfill(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { + s.mu.Lock() + item, ok := s.balanceRecordsByTokenAccount[tokenAccount] + if !ok { + s.mu.Unlock() + return balance.ErrRecordNotFound + } + if item.IsBackfilled { + s.mu.Unlock() + return balance.ErrAlreadyBackfilled + } + s.mu.Unlock() + + result, err := fn(ctx) + if err != nil { + return err } + if result.Quarks < 0 { + return balance.ErrNegativeBalance + } + + s.mu.Lock() + defer s.mu.Unlock() + + item.Quarks = result.Quarks + item.UsdCostBasis = result.UsdCostBasis + item.IsOpen = result.IsOpen + item.IsBackfilled = true + item.UpdatedAt = time.Now() return nil } -func (s *store) findExternalCheckpointByTokenAccount(account string) *balance.ExternalCheckpointRecord { - for _, item := range s.externalCheckpointRecords { - if account == item.TokenAccount { - return item +func (s *store) filter(fn func(*balance.Record) bool) ([]*balance.Record, error) { + var res []*balance.Record + for _, item := range s.balanceRecords { + if !fn(item) { + continue } + cloned := item.Clone() + res = append(res, &cloned) } - return nil + if len(res) == 0 { + return nil, balance.ErrRecordNotFound + } + return res, nil } func (s *store) reset() { s.mu.Lock() defer s.mu.Unlock() + 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 diff --git a/ocp/data/balance/memory/store_legacy.go b/ocp/data/balance/memory/store_legacy.go new file mode 100644 index 0000000..580ffd4 --- /dev/null +++ b/ocp/data/balance/memory/store_legacy.go @@ -0,0 +1,131 @@ +package memory + +import ( + "context" + "time" + + "github.com/code-payments/ocp-server/ocp/data/balance" +) + +// GetCachedVersion implements balance.Store.GetCachedVersion +func (s *store) GetCachedVersion(_ context.Context, account string) (uint64, error) { + s.mu.Lock() + defer s.mu.Unlock() + + current, ok := s.cachedBalanceVersionsByAccount[account] + if !ok { + return 0, nil + } + return current, nil +} + +// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion +func (s *store) AdvanceCachedVersion(_ context.Context, account string, currentVersion uint64) error { + s.mu.Lock() + defer s.mu.Unlock() + + actualVersion, ok := s.cachedBalanceVersionsByAccount[account] + if !ok { + if currentVersion != 0 { + return balance.ErrStaleCachedBalanceVersion + } + + s.cachedBalanceVersionsByAccount[account] = 1 + + return nil + } + + if actualVersion != currentVersion { + return balance.ErrStaleCachedBalanceVersion + } + + s.cachedBalanceVersionsByAccount[account]++ + + return nil +} + +// CheckNotClosed implements balance.Store.CheckNotClosed +func (s *store) CheckNotClosed(ctx context.Context, account string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if _, ok := s.closedAccounts[account]; ok { + return balance.ErrAccountClosed + } + + return nil +} + +// MarkAsClosed implements balance.Store.MarkAsClosed +func (s *store) MarkAsClosed(ctx context.Context, account string) error { + s.mu.Lock() + defer s.mu.Unlock() + + s.closedAccounts[account] = true + + return nil +} + +// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint +func (s *store) SaveExternalCheckpoint(_ context.Context, data *balance.ExternalCheckpointRecord) error { + if err := data.Validate(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + + s.last++ + if item := s.findExternalCheckpoint(data); item != nil { + if data.SlotCheckpoint <= item.SlotCheckpoint { + return balance.ErrStaleCheckpoint + } + + item.SlotCheckpoint = data.SlotCheckpoint + item.Quarks = data.Quarks + item.LastUpdatedAt = time.Now() + item.CopyTo(data) + } else { + if data.Id == 0 { + data.Id = s.last + } + data.LastUpdatedAt = time.Now() + c := data.Clone() + s.externalCheckpointRecords = append(s.externalCheckpointRecords, &c) + } + + return nil +} + +// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint +func (s *store) GetExternalCheckpoint(_ context.Context, account string) (*balance.ExternalCheckpointRecord, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if item := s.findExternalCheckpointByTokenAccount(account); item != nil { + cloned := item.Clone() + return &cloned, nil + } + return nil, balance.ErrCheckpointNotFound +} + +func (s *store) findExternalCheckpoint(data *balance.ExternalCheckpointRecord) *balance.ExternalCheckpointRecord { + for _, item := range s.externalCheckpointRecords { + if item.Id == data.Id { + return item + } + if data.TokenAccount == item.TokenAccount { + return item + } + } + return nil +} + +func (s *store) findExternalCheckpointByTokenAccount(account string) *balance.ExternalCheckpointRecord { + for _, item := range s.externalCheckpointRecords { + if account == item.TokenAccount { + return item + } + } + return nil +} diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index b387ffc..75366d7 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -3,197 +3,292 @@ package postgres import ( "context" "database/sql" - "errors" + "fmt" "time" "github.com/jmoiron/sqlx" - "github.com/code-payments/ocp-server/ocp/data/balance" pgutil "github.com/code-payments/ocp-server/database/postgres" + q "github.com/code-payments/ocp-server/database/query" + "github.com/code-payments/ocp-server/ocp/data/balance" ) const ( - cachedBalanceVersionTableName = "ocp__core_cachedbalanceversion" - openCloseLocksTableName = "ocp__core_opencloselocks" - externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" + tableName = "ocp__core_balance" + + allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_backfilled, updated_at" ) -type externalCheckpointModel struct { +type model struct { Id sql.NullInt64 `db:"id"` - TokenAccount string `db:"token_account"` - Quarks uint64 `db:"quarks"` - SlotCheckpoint uint64 `db:"slot_checkpoint"` + TokenAccount string `db:"token_account"` + OwnerAccount string `db:"owner_account"` + MintAccount string `db:"mint_account"` + + Quarks int64 `db:"quarks"` + UsdCostBasis int64 `db:"usd_cost_basis"` - LastUpdatedAt time.Time `db:"last_updated_at"` + IsOpen bool `db:"is_open"` + IsBackfilled bool `db:"is_backfilled"` + + UpdatedAt time.Time `db:"updated_at"` } -func dbGetCachedVersion(ctx context.Context, db *sqlx.DB, account string) (uint64, error) { - var res uint64 - err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - insertQuery := `INSERT INTO ` + cachedBalanceVersionTableName + ` - (token_account, version) - VALUES($1, 0) - ON CONFLICT DO NOTHING - ` - sqlResult, err := tx.ExecContext(ctx, insertQuery, account) - if err != nil { - return err - } - rowsAffected, err := sqlResult.RowsAffected() - if err != nil { - return err - } - if rowsAffected == 1 { - res = 0 - return nil - } +func toModel(obj *balance.Record) (*model, error) { + if err := obj.Validate(); err != nil { + return nil, err + } - selectQuery := `SELECT version FROM ` + cachedBalanceVersionTableName + ` - WHERE token_account = $1 - FOR UPDATE` - return db.GetContext(ctx, &res, selectQuery, account) - }) - return res, err + return &model{ + TokenAccount: obj.TokenAccount, + OwnerAccount: obj.OwnerAccount, + MintAccount: obj.MintAccount, + + Quarks: obj.Quarks, + UsdCostBasis: obj.UsdCostBasis, + IsOpen: obj.IsOpen, + IsBackfilled: obj.IsBackfilled, + + UpdatedAt: obj.UpdatedAt, + }, nil } -func dbAdvanceCachedVersion(ctx context.Context, db *sqlx.DB, account string, currentVersion uint64) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - var res uint64 - query := `UPDATE ` + cachedBalanceVersionTableName + ` - SET version = version + 1 - WHERE token_account = $1 AND version = $2 - RETURNING version - ` - err := tx.GetContext(ctx, &res, query, account, currentVersion) - if pgutil.IsNoRows(err) || pgutil.IsUniqueViolation(err) { - return balance.ErrStaleCachedBalanceVersion - } - if err != nil { - return err - } - return nil - }) +func fromModel(obj *model) *balance.Record { + return &balance.Record{ + Id: uint64(obj.Id.Int64), + + TokenAccount: obj.TokenAccount, + OwnerAccount: obj.OwnerAccount, + MintAccount: obj.MintAccount, + Quarks: obj.Quarks, + UsdCostBasis: obj.UsdCostBasis, + + IsOpen: obj.IsOpen, + IsBackfilled: obj.IsBackfilled, + + UpdatedAt: obj.UpdatedAt, + } } -func dbCheckNotClosed(ctx context.Context, db *sqlx.DB, account string) error { +func (m *model) dbCreate(ctx context.Context, db *sqlx.DB) error { return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - insertQuery := `INSERT INTO ` + openCloseLocksTableName + ` - (token_account, is_open) - VALUES ($1, TRUE) - ON CONFLICT DO NOTHING - ` - - _, err := tx.ExecContext(ctx, insertQuery, account) - if err != nil { - return err - } + query := `INSERT INTO ` + tableName + ` + (token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_backfilled, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + RETURNING ` + allColumns - selectQuery := `SELECT is_open FROM ` + openCloseLocksTableName + ` - WHERE token_account = $1 - FOR UPDATE - ` - var isOpen bool - err = tx.GetContext(ctx, &isOpen, selectQuery, account) - if err != nil { - return err - } - if !isOpen { - return balance.ErrAccountClosed - } - return nil + m.UpdatedAt = time.Now() + + err := tx.QueryRowxContext( + ctx, + query, + m.TokenAccount, + m.OwnerAccount, + m.MintAccount, + m.Quarks, + m.UsdCostBasis, + m.IsOpen, + m.IsBackfilled, + m.UpdatedAt.UTC(), + ).StructScan(m) + + return pgutil.CheckUniqueViolation(err, balance.ErrRecordExists) }) } -func dbMarkAsClosed(ctx context.Context, db *sqlx.DB, account string) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - query := `INSERT INTO ` + openCloseLocksTableName + ` - (token_account, is_open) - VALUES ($1, FALSE) - - ON CONFLICT (token_account) - DO UPDATE - SET is_open = FALSE - WHERE ` + openCloseLocksTableName + `.token_account = $1 - - RETURNING is_open - ` - var isOpen bool - err := tx.GetContext(ctx, &isOpen, query, account) - if err != nil { - return err - } - if isOpen { - return errors.New("unexpected state transition") - } - return nil +func dbGet(ctx context.Context, db *sqlx.DB, tokenAccount string) (*model, error) { + res := &model{} + + query := `SELECT ` + allColumns + ` FROM ` + tableName + ` + WHERE token_account = $1` + + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + return tx.GetContext(ctx, res, query, tokenAccount) }) + if err != nil { + return nil, pgutil.CheckNoRows(err, balance.ErrRecordNotFound) + } + return res, nil } -func toExternalCheckpointModel(obj *balance.ExternalCheckpointRecord) (*externalCheckpointModel, error) { - if err := obj.Validate(); err != nil { - return nil, err +func dbGetBatch(ctx context.Context, db *sqlx.DB, tokenAccounts ...string) ([]*model, error) { + res := []*model{} + if len(tokenAccounts) == 0 { + return res, nil } - return &externalCheckpointModel{ - TokenAccount: obj.TokenAccount, - Quarks: obj.Quarks, - SlotCheckpoint: obj.SlotCheckpoint, - LastUpdatedAt: obj.LastUpdatedAt, - }, nil + query := `SELECT ` + allColumns + ` FROM ` + tableName + ` + WHERE token_account = ANY($1)` + + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + return tx.SelectContext(ctx, &res, query, tokenAccounts) + }) + if err != nil && !pgutil.IsNoRows(err) { + return nil, err + } + return res, nil } -func fromExternalCheckpoingModel(obj *externalCheckpointModel) *balance.ExternalCheckpointRecord { - return &balance.ExternalCheckpointRecord{ - Id: uint64(obj.Id.Int64), - TokenAccount: obj.TokenAccount, - Quarks: obj.Quarks, - SlotCheckpoint: obj.SlotCheckpoint, - LastUpdatedAt: obj.LastUpdatedAt, +func dbGetAllByOwner(ctx context.Context, db *sqlx.DB, owner string, mint *string) ([]*model, error) { + res := []*model{} + + query := `SELECT ` + allColumns + ` FROM ` + tableName + ` + WHERE owner_account = $1` + args := []any{owner} + if mint != nil { + query += ` AND mint_account = $2` + args = append(args, *mint) + } + query += ` ORDER BY id ASC` + + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + return tx.SelectContext(ctx, &res, query, args...) + }) + if err != nil { + return nil, pgutil.CheckNoRows(err, balance.ErrRecordNotFound) } + if len(res) == 0 { + return nil, balance.ErrRecordNotFound + } + return res, nil } -func (m *externalCheckpointModel) dbSave(ctx context.Context, db *sqlx.DB) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - query := `INSERT INTO ` + externalCheckpointTableName + ` - (token_account, quarks, slot_checkpoint, last_updated_at) - VALUES ($1, $2, $3, $4) +func dbGetAllByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64, cursor q.Cursor, limit uint64, direction q.Ordering) ([]*model, error) { + res := []*model{} - ON CONFLICT (token_account) - DO UPDATE - SET quarks = $2, slot_checkpoint = $3, last_updated_at = $4 - WHERE ` + externalCheckpointTableName + `.token_account = $1 AND ` + externalCheckpointTableName + `.slot_checkpoint < $3 + query := `SELECT ` + allColumns + ` FROM ` + tableName + ` + WHERE (mint_account = $1 AND quarks >= $2)` + query, args := q.PaginateQuery(query, []any{mint, minQuarks}, cursor, limit, direction) - RETURNING - id, token_account, quarks, slot_checkpoint, last_updated_at` + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + return tx.SelectContext(ctx, &res, query, args...) + }) + if err != nil { + return nil, pgutil.CheckNoRows(err, balance.ErrRecordNotFound) + } + if len(res) == 0 { + return nil, balance.ErrRecordNotFound + } + return res, nil +} - m.LastUpdatedAt = time.Now() +// dbApplyDeltas applies every delta in a single transaction. Each delta is one +// conditional UPDATE, so its predicate is evaluated against the row after the +// row lock is acquired. Predicates only apply to backfilled rows. +func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + for _, delta := range deltas { + var query string + var args []any + switch delta.Kind { + case balance.DeltaCredit: + query = `UPDATE ` + tableName + ` + SET quarks = quarks + $2, usd_cost_basis = usd_cost_basis + $3, updated_at = $4 + WHERE token_account = $1 AND (NOT is_backfilled OR is_open)` + args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} + case balance.DeltaDebit: + query = `UPDATE ` + tableName + ` + SET quarks = quarks - $2, usd_cost_basis = usd_cost_basis - $3, updated_at = $4 + WHERE token_account = $1 AND (NOT is_backfilled OR quarks >= $2)` + args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} + case balance.DeltaDrain: + query = `UPDATE ` + tableName + ` + SET quarks = CASE WHEN is_backfilled THEN 0 ELSE quarks - $2 END, + usd_cost_basis = CASE WHEN is_backfilled THEN 0 ELSE usd_cost_basis - $3 END, + is_open = FALSE, + updated_at = $4 + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND quarks = $2))` + args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} + case balance.DeltaClose: + query = `UPDATE ` + tableName + ` + SET is_open = FALSE, updated_at = $2 + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND quarks = 0))` + args = []any{delta.TokenAccount, time.Now().UTC()} + default: + return fmt.Errorf("unsupported delta kind: %s", delta.Kind) + } - err := tx.QueryRowxContext( - ctx, - query, - m.TokenAccount, - m.Quarks, - m.SlotCheckpoint, - m.LastUpdatedAt.UTC(), - ).StructScan(m) + sqlResult, err := tx.ExecContext(ctx, query, args...) + if err != nil { + return err + } + rowsAffected, err := sqlResult.RowsAffected() + if err != nil { + return err + } + if rowsAffected == 1 { + continue + } - return pgutil.CheckNoRows(err, balance.ErrStaleCheckpoint) + // Either the predicate failed or there is no record. Classify which. + var current model + err = tx.GetContext(ctx, ¤t, `SELECT `+allColumns+` FROM `+tableName+` WHERE token_account = $1`, delta.TokenAccount) + if pgutil.IsNoRows(err) { + continue // Not an account we track + } else if err != nil { + return err + } + return classifyFailedDelta(delta, fromModel(¤t)) + } + return nil }) } -func dbGetExternalCheckpoint(ctx context.Context, db *sqlx.DB, account string) (*externalCheckpointModel, error) { - res := &externalCheckpointModel{} +func classifyFailedDelta(delta *balance.Delta, current *balance.Record) error { + switch delta.Kind { + case balance.DeltaCredit: + return balance.ErrAccountClosed + case balance.DeltaDebit: + return balance.ErrInsufficientBalance + case balance.DeltaDrain, balance.DeltaClose: + if !current.IsOpen { + return balance.ErrAccountClosed + } + return balance.ErrBalanceChanged + } + return fmt.Errorf("unsupported delta kind: %s", delta.Kind) +} - query := `SELECT id, token_account, quarks, slot_checkpoint, last_updated_at FROM ` + externalCheckpointTableName + ` - WHERE token_account = $1 - LIMIT 1` +func dbBackfill(ctx context.Context, db *sqlx.DB, tokenAccount string, fn balance.BackfillFunc) error { + return executeTxWithinCtxOrJoin(ctx, db, func(ctx context.Context) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + var current model + err := tx.GetContext(ctx, ¤t, `SELECT `+allColumns+` FROM `+tableName+` WHERE token_account = $1 FOR UPDATE`, tokenAccount) + if err != nil { + return pgutil.CheckNoRows(err, balance.ErrRecordNotFound) + } + if current.IsBackfilled { + return balance.ErrAlreadyBackfilled + } - err := db.GetContext(ctx, res, query, account) - if err != nil { - return nil, pgutil.CheckNoRows(err, balance.ErrCheckpointNotFound) + // The row lock is held across fn, so it observes every committed + // write to the account and blocks every in-flight one. + result, err := fn(ctx) + if err != nil { + return err + } + if result.Quarks < 0 { + return balance.ErrNegativeBalance + } + + query := `UPDATE ` + tableName + ` + SET quarks = $2, usd_cost_basis = $3, is_open = $4, is_backfilled = TRUE, updated_at = $5 + WHERE token_account = $1` + _, err = tx.ExecContext(ctx, query, tokenAccount, result.Quarks, result.UsdCostBasis, result.IsOpen, time.Now().UTC()) + return err + }) + }) +} + +// executeTxWithinCtxOrJoin runs fn with a context carrying a DB transaction, +// starting one if the context doesn't already have one. +func executeTxWithinCtxOrJoin(ctx context.Context, db *sqlx.DB, fn func(ctx context.Context) error) error { + err := pgutil.ExecuteTxWithinCtx(ctx, db, sql.LevelDefault, fn) + if err == pgutil.ErrAlreadyInTx { + return fn(ctx) } - return res, nil + return err } diff --git a/ocp/data/balance/postgres/model_legacy.go b/ocp/data/balance/postgres/model_legacy.go new file mode 100644 index 0000000..dbefe71 --- /dev/null +++ b/ocp/data/balance/postgres/model_legacy.go @@ -0,0 +1,199 @@ +package postgres + +import ( + "context" + "database/sql" + "errors" + "time" + + "github.com/jmoiron/sqlx" + + pgutil "github.com/code-payments/ocp-server/database/postgres" + "github.com/code-payments/ocp-server/ocp/data/balance" +) + +const ( + cachedBalanceVersionTableName = "ocp__core_cachedbalanceversion" + openCloseLocksTableName = "ocp__core_opencloselocks" + externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" +) + +type externalCheckpointModel struct { + Id sql.NullInt64 `db:"id"` + + TokenAccount string `db:"token_account"` + Quarks uint64 `db:"quarks"` + SlotCheckpoint uint64 `db:"slot_checkpoint"` + + LastUpdatedAt time.Time `db:"last_updated_at"` +} + +func dbGetCachedVersion(ctx context.Context, db *sqlx.DB, account string) (uint64, error) { + var res uint64 + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + insertQuery := `INSERT INTO ` + cachedBalanceVersionTableName + ` + (token_account, version) + VALUES($1, 0) + ON CONFLICT DO NOTHING + ` + sqlResult, err := tx.ExecContext(ctx, insertQuery, account) + if err != nil { + return err + } + rowsAffected, err := sqlResult.RowsAffected() + if err != nil { + return err + } + if rowsAffected == 1 { + res = 0 + return nil + } + + selectQuery := `SELECT version FROM ` + cachedBalanceVersionTableName + ` + WHERE token_account = $1 + FOR UPDATE` + return db.GetContext(ctx, &res, selectQuery, account) + }) + return res, err + +} + +func dbAdvanceCachedVersion(ctx context.Context, db *sqlx.DB, account string, currentVersion uint64) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + var res uint64 + query := `UPDATE ` + cachedBalanceVersionTableName + ` + SET version = version + 1 + WHERE token_account = $1 AND version = $2 + RETURNING version + ` + err := tx.GetContext(ctx, &res, query, account, currentVersion) + if pgutil.IsNoRows(err) || pgutil.IsUniqueViolation(err) { + return balance.ErrStaleCachedBalanceVersion + } + if err != nil { + return err + } + return nil + }) + +} + +func dbCheckNotClosed(ctx context.Context, db *sqlx.DB, account string) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + insertQuery := `INSERT INTO ` + openCloseLocksTableName + ` + (token_account, is_open) + VALUES ($1, TRUE) + ON CONFLICT DO NOTHING + ` + + _, err := tx.ExecContext(ctx, insertQuery, account) + if err != nil { + return err + } + + selectQuery := `SELECT is_open FROM ` + openCloseLocksTableName + ` + WHERE token_account = $1 + FOR UPDATE + ` + var isOpen bool + err = tx.GetContext(ctx, &isOpen, selectQuery, account) + if err != nil { + return err + } + if !isOpen { + return balance.ErrAccountClosed + } + return nil + }) +} + +func dbMarkAsClosed(ctx context.Context, db *sqlx.DB, account string) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + query := `INSERT INTO ` + openCloseLocksTableName + ` + (token_account, is_open) + VALUES ($1, FALSE) + + ON CONFLICT (token_account) + DO UPDATE + SET is_open = FALSE + WHERE ` + openCloseLocksTableName + `.token_account = $1 + + RETURNING is_open + ` + var isOpen bool + err := tx.GetContext(ctx, &isOpen, query, account) + if err != nil { + return err + } + if isOpen { + return errors.New("unexpected state transition") + } + return nil + }) +} + +func toExternalCheckpointModel(obj *balance.ExternalCheckpointRecord) (*externalCheckpointModel, error) { + if err := obj.Validate(); err != nil { + return nil, err + } + + return &externalCheckpointModel{ + TokenAccount: obj.TokenAccount, + Quarks: obj.Quarks, + SlotCheckpoint: obj.SlotCheckpoint, + LastUpdatedAt: obj.LastUpdatedAt, + }, nil +} + +func fromExternalCheckpoingModel(obj *externalCheckpointModel) *balance.ExternalCheckpointRecord { + return &balance.ExternalCheckpointRecord{ + Id: uint64(obj.Id.Int64), + TokenAccount: obj.TokenAccount, + Quarks: obj.Quarks, + SlotCheckpoint: obj.SlotCheckpoint, + LastUpdatedAt: obj.LastUpdatedAt, + } +} + +func (m *externalCheckpointModel) dbSave(ctx context.Context, db *sqlx.DB) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + query := `INSERT INTO ` + externalCheckpointTableName + ` + (token_account, quarks, slot_checkpoint, last_updated_at) + VALUES ($1, $2, $3, $4) + + ON CONFLICT (token_account) + DO UPDATE + SET quarks = $2, slot_checkpoint = $3, last_updated_at = $4 + WHERE ` + externalCheckpointTableName + `.token_account = $1 AND ` + externalCheckpointTableName + `.slot_checkpoint < $3 + + RETURNING + id, token_account, quarks, slot_checkpoint, last_updated_at` + + m.LastUpdatedAt = time.Now() + + err := tx.QueryRowxContext( + ctx, + query, + m.TokenAccount, + m.Quarks, + m.SlotCheckpoint, + m.LastUpdatedAt.UTC(), + ).StructScan(m) + + return pgutil.CheckNoRows(err, balance.ErrStaleCheckpoint) + }) +} + +func dbGetExternalCheckpoint(ctx context.Context, db *sqlx.DB, account string) (*externalCheckpointModel, error) { + res := &externalCheckpointModel{} + + query := `SELECT id, token_account, quarks, slot_checkpoint, last_updated_at FROM ` + externalCheckpointTableName + ` + WHERE token_account = $1 + LIMIT 1` + + err := db.GetContext(ctx, res, query, account) + if err != nil { + return nil, pgutil.CheckNoRows(err, balance.ErrCheckpointNotFound) + } + return res, nil +} diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index bac477b..1bdbdde 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -6,6 +6,7 @@ import ( "github.com/jmoiron/sqlx" + "github.com/code-payments/ocp-server/database/query" "github.com/code-payments/ocp-server/ocp/data/balance" ) @@ -20,48 +21,95 @@ func New(db *sql.DB) balance.Store { } } -// GetCachedVersion implements balance.Store.GetCachedVersion -func (s *store) GetCachedVersion(ctx context.Context, account string) (uint64, error) { - return dbGetCachedVersion(ctx, s.db, account) -} +// Create implements balance.Store.Create +func (s *store) Create(ctx context.Context, record *balance.Record) error { + model, err := toModel(record) + if err != nil { + return err + } -// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion -func (s *store) AdvanceCachedVersion(ctx context.Context, account string, currentVersion uint64) error { - return dbAdvanceCachedVersion(ctx, s.db, account, currentVersion) -} + if err := model.dbCreate(ctx, s.db); err != nil { + return err + } -// CheckNotClosed implements balance.Store.CheckNotClosed -func (s *store) CheckNotClosed(ctx context.Context, account string) error { - return dbCheckNotClosed(ctx, s.db, account) + fromModel(model).CopyTo(record) + return nil } -// MarkAsClosed implements balance.Store.MarkAsClosed -func (s *store) MarkAsClosed(ctx context.Context, account string) error { - return dbMarkAsClosed(ctx, s.db, account) +// Get implements balance.Store.Get +func (s *store) Get(ctx context.Context, tokenAccount string) (*balance.Record, error) { + model, err := dbGet(ctx, s.db, tokenAccount) + if err != nil { + return nil, err + } + return fromModel(model), nil } -// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint -func (s *store) SaveExternalCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { - model, err := toExternalCheckpointModel(record) +// GetBatch implements balance.Store.GetBatch +func (s *store) GetBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) { + models, err := dbGetBatch(ctx, s.db, tokenAccounts...) if err != nil { - return err + return nil, err } - if err := model.dbSave(ctx, s.db); err != nil { - return err + res := make(map[string]*balance.Record, len(models)) + for _, model := range models { + res[model.TokenAccount] = fromModel(model) } + return res, nil +} - res := fromExternalCheckpoingModel(model) - res.CopyTo(record) +// GetAllByOwner implements balance.Store.GetAllByOwner +func (s *store) GetAllByOwner(ctx context.Context, owner string) ([]*balance.Record, error) { + models, err := dbGetAllByOwner(ctx, s.db, owner, nil) + if err != nil { + return nil, err + } + return fromModels(models), nil +} - return nil +// GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint +func (s *store) GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { + models, err := dbGetAllByOwner(ctx, s.db, owner, &mint) + if err != nil { + return nil, err + } + return fromModels(models), nil } -// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint -func (s *store) GetExternalCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) { - model, err := dbGetExternalCheckpoint(ctx, s.db, account) +// GetAllByMint implements balance.Store.GetAllByMint +func (s *store) GetAllByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { + models, err := dbGetAllByMint(ctx, s.db, mint, minQuarks, cursor, limit, direction) if err != nil { return nil, err } - return fromExternalCheckpoingModel(model), nil + return fromModels(models), nil +} + +// ApplyDeltas implements balance.Store.ApplyDeltas +func (s *store) ApplyDeltas(ctx context.Context, deltas ...*balance.Delta) error { + for _, delta := range deltas { + if err := delta.Validate(); err != nil { + return err + } + } + + sorted := make([]*balance.Delta, len(deltas)) + copy(sorted, deltas) + balance.SortDeltas(sorted) + + return dbApplyDeltas(ctx, s.db, sorted) +} + +// Backfill implements balance.Store.Backfill +func (s *store) Backfill(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { + return dbBackfill(ctx, s.db, tokenAccount, fn) +} + +func fromModels(models []*model) []*balance.Record { + res := make([]*balance.Record, len(models)) + for i, model := range models { + res[i] = fromModel(model) + } + return res } diff --git a/ocp/data/balance/postgres/store_legacy.go b/ocp/data/balance/postgres/store_legacy.go new file mode 100644 index 0000000..c66bede --- /dev/null +++ b/ocp/data/balance/postgres/store_legacy.go @@ -0,0 +1,53 @@ +package postgres + +import ( + "context" + + "github.com/code-payments/ocp-server/ocp/data/balance" +) + +// GetCachedVersion implements balance.Store.GetCachedVersion +func (s *store) GetCachedVersion(ctx context.Context, account string) (uint64, error) { + return dbGetCachedVersion(ctx, s.db, account) +} + +// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion +func (s *store) AdvanceCachedVersion(ctx context.Context, account string, currentVersion uint64) error { + return dbAdvanceCachedVersion(ctx, s.db, account, currentVersion) +} + +// CheckNotClosed implements balance.Store.CheckNotClosed +func (s *store) CheckNotClosed(ctx context.Context, account string) error { + return dbCheckNotClosed(ctx, s.db, account) +} + +// MarkAsClosed implements balance.Store.MarkAsClosed +func (s *store) MarkAsClosed(ctx context.Context, account string) error { + return dbMarkAsClosed(ctx, s.db, account) +} + +// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint +func (s *store) SaveExternalCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { + model, err := toExternalCheckpointModel(record) + if err != nil { + return err + } + + if err := model.dbSave(ctx, s.db); err != nil { + return err + } + + res := fromExternalCheckpoingModel(model) + res.CopyTo(record) + + return nil +} + +// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint +func (s *store) GetExternalCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) { + model, err := dbGetExternalCheckpoint(ctx, s.db, account) + if err != nil { + return nil, err + } + return fromExternalCheckpoingModel(model), nil +} diff --git a/ocp/data/balance/postgres/store_test.go b/ocp/data/balance/postgres/store_test.go index 6980945..77c0146 100644 --- a/ocp/data/balance/postgres/store_test.go +++ b/ocp/data/balance/postgres/store_test.go @@ -24,6 +24,28 @@ var ( const ( // Used for testing ONLY, the table and migrations are external to this repository tableCreate = ` + CREATE TABLE ocp__core_balance ( + id SERIAL NOT NULL PRIMARY KEY, + + token_account TEXT NOT NULL, + owner_account TEXT NOT NULL, + mint_account TEXT NOT NULL, + + quarks BIGINT NOT NULL DEFAULT 0, + usd_cost_basis BIGINT NOT NULL DEFAULT 0, + + is_open BOOL NOT NULL DEFAULT TRUE, + is_backfilled BOOL NOT NULL DEFAULT FALSE, + + updated_at TIMESTAMP WITH TIME ZONE NOT NULL, + + CONSTRAINT ocp__core_balance__uniq__token_account UNIQUE (token_account), + CONSTRAINT ocp__core_balance__check__nonnegative CHECK (NOT is_backfilled OR quarks >= 0) + ) WITH (fillfactor = 90); + + CREATE INDEX ocp__core_balance__idx__owner_account__mint_account ON ocp__core_balance (owner_account, mint_account); + CREATE INDEX ocp__core_balance__idx__mint_account ON ocp__core_balance (mint_account); + CREATE TABLE ocp__core_cachedbalanceversion ( id SERIAL NOT NULL PRIMARY KEY, @@ -57,6 +79,7 @@ const ( // Used for testing ONLY, the table and migrations are external to this repository tableDestroy = ` + DROP TABLE ocp__core_balance; DROP TABLE ocp__core_cachedbalanceversion; DROP TABLE ocp__core_opencloselocks; DROP TABLE ocp__core_externalbalancecheckpoint; diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go new file mode 100644 index 0000000..d656ed4 --- /dev/null +++ b/ocp/data/balance/record.go @@ -0,0 +1,170 @@ +package balance + +import ( + "errors" + "sort" + "time" +) + +// UsdQuarksPerUnit is the scale of UsdCostBasis: 1 unit is $0.000001. It is +// deliberately equal to the core mint's quarks per unit, so a core mint account's +// USD cost basis is exactly its quark balance. +const UsdQuarksPerUnit = 1_000_000 + +// Record is the materialized balance of a token account managed by Code. +type Record struct { + Id uint64 + + TokenAccount string + OwnerAccount string + MintAccount string + + // Quarks is signed because a record that has not been backfilled only + // accumulates deltas, which may temporarily net negative. Backfilled + // records are guaranteed to be non-negative. + Quarks int64 + + // UsdCostBasis is the account's USD cost basis, in UsdQuarksPerUnit. + // A cost basis may legitimately be negative. + UsdCostBasis int64 + + IsOpen bool + + // IsBackfilled indicates the record reflects the full history of the + // account. Until it does, deltas are recorded without enforcing any + // balance predicates. + IsBackfilled bool + + UpdatedAt time.Time +} + +func (r *Record) Validate() error { + if len(r.TokenAccount) == 0 { + return errors.New("token account is required") + } + + if len(r.OwnerAccount) == 0 { + return errors.New("owner account is required") + } + + if len(r.MintAccount) == 0 { + return errors.New("mint account is required") + } + + if r.IsBackfilled && r.Quarks < 0 { + return errors.New("backfilled quarks cannot be negative") + } + + return nil +} + +func (r *Record) Clone() Record { + return Record{ + Id: r.Id, + + TokenAccount: r.TokenAccount, + OwnerAccount: r.OwnerAccount, + MintAccount: r.MintAccount, + + Quarks: r.Quarks, + UsdCostBasis: r.UsdCostBasis, + + IsOpen: r.IsOpen, + IsBackfilled: r.IsBackfilled, + + UpdatedAt: r.UpdatedAt, + } +} + +func (r *Record) CopyTo(dst *Record) { + dst.Id = r.Id + + dst.TokenAccount = r.TokenAccount + dst.OwnerAccount = r.OwnerAccount + dst.MintAccount = r.MintAccount + + dst.Quarks = r.Quarks + dst.UsdCostBasis = r.UsdCostBasis + + dst.IsOpen = r.IsOpen + dst.IsBackfilled = r.IsBackfilled + + dst.UpdatedAt = r.UpdatedAt +} + +// DeltaKind selects the predicate a Delta is applied under. Predicates are +// only enforced on backfilled records. +type DeltaKind uint8 + +const ( + // DeltaCredit adds funds to an open account. + DeltaCredit DeltaKind = iota + 1 + + // DeltaDebit removes funds from an account with sufficient balance. + DeltaDebit + + // DeltaDrain removes exactly the account's full balance and closes it. + DeltaDrain + + // DeltaClose closes an account with a zero balance. + DeltaClose +) + +// Delta is a single balance change to apply to a token account. +type Delta struct { + TokenAccount string + Kind DeltaKind + + // Quarks is the amount credited, debited or drained. Ignored for DeltaClose. + Quarks uint64 + + // UsdCostBasis is added on credit and subtracted on debit. It is signed + // so that a credit can also carry a downward reconciliation. Ignored for + // DeltaDrain and DeltaClose on backfilled records, where the basis is + // zeroed along with the balance. + UsdCostBasis int64 +} + +func (d *Delta) Validate() error { + if len(d.TokenAccount) == 0 { + return errors.New("token account is required") + } + + switch d.Kind { + case DeltaCredit, DeltaDebit, DeltaDrain: + if d.Quarks == 0 && d.UsdCostBasis == 0 { + return errors.New("delta is a no-op") + } + case DeltaClose: + default: + return errors.New("invalid delta kind") + } + + return nil +} + +// SortDeltas orders deltas by token account, then by kind, so every store +// implementation acquires row locks in the same order and cannot deadlock +// against another transaction applying deltas to the same accounts. +func SortDeltas(deltas []*Delta) { + sort.SliceStable(deltas, func(i, j int) bool { + if deltas[i].TokenAccount != deltas[j].TokenAccount { + return deltas[i].TokenAccount < deltas[j].TokenAccount + } + return deltas[i].Kind < deltas[j].Kind + }) +} + +func (k DeltaKind) String() string { + switch k { + case DeltaCredit: + return "credit" + case DeltaDebit: + return "debit" + case DeltaDrain: + return "drain" + case DeltaClose: + return "close" + } + return "unknown" +} diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index cd80428..da38824 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -3,9 +3,28 @@ package balance import ( "context" "errors" + + "github.com/code-payments/ocp-server/database/query" ) var ( + ErrRecordNotFound = errors.New("balance record not found") + ErrRecordExists = errors.New("balance record already exists") + + // ErrInsufficientBalance is returned when a debit exceeds the balance. + ErrInsufficientBalance = errors.New("insufficient balance") + + // ErrBalanceChanged is returned when a drain or close expected a different + // balance than the one on record. + ErrBalanceChanged = errors.New("balance is not the expected value") + + ErrAlreadyBackfilled = errors.New("balance record is already backfilled") + + // ErrNegativeBalance is returned when a backfill computes a negative + // balance, which indicates inconsistent historical data that must be + // reviewed rather than recorded. + ErrNegativeBalance = errors.New("backfilled balance is negative") + ErrStaleCachedBalanceVersion = errors.New("cached balance version is stale") ErrAccountClosed = errors.New("account open state is stale") @@ -14,14 +33,91 @@ var ( ErrStaleCheckpoint = errors.New("checkpoint is stale") ) +// BackfillResult is the full historical state of a token account. +type BackfillResult struct { + Quarks int64 + UsdCostBasis int64 + + // IsOpen is false for accounts that can no longer receive funds, such as + // claimed gift cards and distributed pools. + IsOpen bool +} + +// BackfillFunc computes the full historical state of a token account. It +// is called while the record is locked, with a context that is part of the +// same DB transaction, so any store reads made through it observe every +// committed change and block every in-flight one. +type BackfillFunc func(ctx context.Context) (*BackfillResult, error) + type Store interface { + // Create creates a new balance record. + // + // ErrRecordExists is returned if the token account already has a record. + Create(ctx context.Context, record *Record) error + + // Get gets the balance record for a token account. + // + // ErrRecordNotFound is returned if no record exists. + Get(ctx context.Context, tokenAccount string) (*Record, error) + + // GetBatch gets balance records for a set of token accounts. Accounts + // without a record are omitted from the result. + GetBatch(ctx context.Context, tokenAccounts ...string) (map[string]*Record, error) + + // GetAllByOwner gets all balance records for an owner. + // + // ErrRecordNotFound is returned if no records exist. + GetAllByOwner(ctx context.Context, owner string) ([]*Record, error) + + // GetAllByOwnerAndMint gets all balance records for an owner and mint. + // + // ErrRecordNotFound is returned if no records exist. + GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([]*Record, error) + + // GetAllByMint gets balance records for a mint with at least minQuarks, + // paged by record ID. + // + // ErrRecordNotFound is returned if no records exist. + GetAllByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) + + // ApplyDeltas atomically applies a set of deltas. Either every delta is + // applied or none are. Deltas are applied in SortDeltas order. + // + // Predicates are enforced only on backfilled records; records that are + // not backfilled simply accumulate the change. Deltas for token accounts + // without a record are skipped, since only accounts managed by Code are + // tracked. + // + // ErrInsufficientBalance is returned when a debit exceeds the balance. + // ErrBalanceChanged is returned when a drain or close doesn't match the + // balance. ErrAccountClosed is returned when a credit, drain or close + // targets a closed account. + ApplyDeltas(ctx context.Context, deltas ...*Delta) error + + // Backfill locks a record that is not yet backfilled, calls fn to compute + // its full historical balance, and overwrites the record with the result, + // marking it as backfilled. Deltas recorded before the backfill are + // intentionally discarded, since fn observes them. + // + // ErrRecordNotFound is returned if no record exists. ErrAlreadyBackfilled + // is returned if the record is already backfilled, in which case fn is + // not called. ErrNegativeBalance is returned if fn computes a negative + // balance, leaving the record untouched. + Backfill(ctx context.Context, tokenAccount string, fn BackfillFunc) error + // GetCachedVersion gets the current cached balance version, which can be used // for optimistic locking cached balances for operations with outgoing transfers. + // + // Note: Use ApplyDeltas, whose predicates replace the version check. + // Retained for accounts that are not yet backfilled. GetCachedVersion(ctx context.Context, account string) (uint64, error) // AdvanceCachedVersion advances an account's cached balance version. // // ErrStaleCachedBalanceVersion is returned if the currentVersion is out of date. + // + // Note: Use ApplyDeltas, whose predicates replace the version check. + // Retained for accounts that are not yet backfilled. AdvanceCachedVersion(ctx context.Context, account string, currentVersion uint64) error // CheckNotClosed checks whether an account is closed under a lock to guarantee @@ -29,10 +125,16 @@ type Store interface { // account. // // ErrAccountClosed is returned if the account has been closed. + // + // Note: Use ApplyDeltas with DeltaCredit. Retained for accounts that + // are not yet backfilled. CheckNotClosed(ctx context.Context, account string) error // MarkAsClosed marks an account as being closed and unable to receive payments // as a destination. + // + // Note: Use ApplyDeltas with DeltaDrain or DeltaClose. Retained for + // accounts that are not yet backfilled. MarkAsClosed(ctx context.Context, account string) error // SaveExternalCheckpoint saves an external balance at a checkpoint. diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 30f80ba..e9e654d 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -2,17 +2,26 @@ package tests import ( "context" + "sync" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/code-payments/ocp-server/database/query" "github.com/code-payments/ocp-server/ocp/data/balance" ) func RunTests(t *testing.T, s balance.Store, teardown func()) { for _, tf := range []func(t *testing.T, s balance.Store){ + testRecordHappyPath, + testGetAllByMint, + testApplyDeltasBackfilled, + testApplyDeltasNotBackfilled, + testApplyDeltasAtomicity, + testApplyDeltasConcurrency, + testBackfill, testCachedBalanceVersionHappyPath, testClosedAccountHappyPath, testExternalCheckpointHappyPath, @@ -22,12 +31,472 @@ func RunTests(t *testing.T, s balance.Store, teardown func()) { } } +func testRecordHappyPath(t *testing.T, s balance.Store) { + t.Run("testRecordHappyPath", func(t *testing.T) { + ctx := context.Background() + + _, err := s.Get(ctx, "token_account_1") + assert.Equal(t, balance.ErrRecordNotFound, err) + + _, err = s.GetAllByOwner(ctx, "owner_1") + assert.Equal(t, balance.ErrRecordNotFound, err) + + _, err = s.GetAllByOwnerAndMint(ctx, "owner_1", "mint_1") + assert.Equal(t, balance.ErrRecordNotFound, err) + + batch, err := s.GetBatch(ctx, "token_account_1", "token_account_2") + require.NoError(t, err) + assert.Empty(t, batch) + + start := time.Now() + + expected := &balance.Record{ + TokenAccount: "token_account_1", + OwnerAccount: "owner_1", + MintAccount: "mint_1", + Quarks: 100, + UsdCostBasis: 200, + IsOpen: true, + IsBackfilled: true, + } + cloned := expected.Clone() + + require.NoError(t, s.Create(ctx, expected)) + assert.EqualValues(t, 1, expected.Id) + assert.True(t, expected.UpdatedAt.After(start)) + + assert.Equal(t, balance.ErrRecordExists, s.Create(ctx, &cloned)) + + actual, err := s.Get(ctx, "token_account_1") + require.NoError(t, err) + assertEquivalentRecords(t, &cloned, actual) + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_2", + OwnerAccount: "owner_1", + MintAccount: "mint_2", + IsOpen: true, + })) + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_3", + OwnerAccount: "owner_2", + MintAccount: "mint_1", + IsOpen: true, + })) + + batch, err = s.GetBatch(ctx, "token_account_1", "token_account_3", "token_account_4") + require.NoError(t, err) + require.Len(t, batch, 2) + assertEquivalentRecords(t, &cloned, batch["token_account_1"]) + assert.Equal(t, "token_account_3", batch["token_account_3"].TokenAccount) + + byOwner, err := s.GetAllByOwner(ctx, "owner_1") + require.NoError(t, err) + require.Len(t, byOwner, 2) + assert.Equal(t, "token_account_1", byOwner[0].TokenAccount) + assert.Equal(t, "token_account_2", byOwner[1].TokenAccount) + + byOwnerAndMint, err := s.GetAllByOwnerAndMint(ctx, "owner_1", "mint_2") + require.NoError(t, err) + require.Len(t, byOwnerAndMint, 1) + assert.Equal(t, "token_account_2", byOwnerAndMint[0].TokenAccount) + + _, err = s.GetAllByOwnerAndMint(ctx, "owner_2", "mint_2") + assert.Equal(t, balance.ErrRecordNotFound, err) + + assert.Error(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_5", + OwnerAccount: "owner_1", + MintAccount: "mint_1", + Quarks: -1, + IsBackfilled: true, + })) + }) +} + +func testGetAllByMint(t *testing.T, s balance.Store) { + t.Run("testGetAllByMint", func(t *testing.T) { + ctx := context.Background() + + _, err := s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) + assert.Equal(t, balance.ErrRecordNotFound, err) + + for i := range 5 { + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_" + string(rune('a'+i)), + OwnerAccount: "owner", + MintAccount: "mint_1", + Quarks: int64(i * 10), + IsOpen: true, + IsBackfilled: true, + })) + } + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_other", + OwnerAccount: "owner", + MintAccount: "mint_2", + Quarks: 1000, + IsOpen: true, + IsBackfilled: true, + })) + + records, err := s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) + require.NoError(t, err) + require.Len(t, records, 5) + for i, record := range records { + assert.EqualValues(t, i+1, record.Id) + } + + records, err = s.GetAllByMint(ctx, "mint_1", 20, query.EmptyCursor, 10, query.Ascending) + require.NoError(t, err) + require.Len(t, records, 3) + assert.EqualValues(t, 20, records[0].Quarks) + + records, err = s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Ascending) + require.NoError(t, err) + require.Len(t, records, 2) + assert.EqualValues(t, 1, records[0].Id) + assert.EqualValues(t, 2, records[1].Id) + + records, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(2), 2, query.Ascending) + require.NoError(t, err) + require.Len(t, records, 2) + assert.EqualValues(t, 3, records[0].Id) + assert.EqualValues(t, 4, records[1].Id) + + records, err = s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Descending) + require.NoError(t, err) + require.Len(t, records, 2) + assert.EqualValues(t, 5, records[0].Id) + assert.EqualValues(t, 4, records[1].Id) + + records, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(4), 10, query.Descending) + require.NoError(t, err) + require.Len(t, records, 3) + assert.EqualValues(t, 3, records[0].Id) + + _, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(5), 10, query.Ascending) + assert.Equal(t, balance.ErrRecordNotFound, err) + }) +} + +func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { + t.Run("testApplyDeltasBackfilled", func(t *testing.T) { + ctx := context.Background() + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_1", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + IsBackfilled: true, + })) + + // Deltas for accounts without a record are skipped + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDebit, Quarks: 1})) + + // Invalid deltas are rejected + assert.Error(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit})) + assert.Error(t, s.ApplyDeltas(ctx, &balance.Delta{Kind: balance.DeltaCredit, Quarks: 1})) + + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 100, UsdCostBasis: 50})) + assertBalance(t, s, "token_account_1", 100, 50, true) + + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 101, UsdCostBasis: 1})) + assertBalance(t, s, "token_account_1", 100, 50, true) + + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 30, UsdCostBasis: 60})) + assertBalance(t, s, "token_account_1", 70, -10, true) + + // A credit can carry a signed USD-only reconciliation + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, UsdCostBasis: -5})) + assertBalance(t, s, "token_account_1", 70, -15, true) + + assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) + assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 69})) + assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 71})) + assertBalance(t, s, "token_account_1", 70, -15, true) + + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 70, UsdCostBasis: 12345})) + assertBalance(t, s, "token_account_1", 0, 0, false) + + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1})) + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 0, UsdCostBasis: 1})) + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) + assertBalance(t, s, "token_account_1", 0, 0, false) + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_2", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + IsBackfilled: true, + })) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaClose})) + assertBalance(t, s, "token_account_2", 0, 0, false) + }) +} + +func testApplyDeltasNotBackfilled(t *testing.T, s balance.Store) { + t.Run("testApplyDeltasNotBackfilled", func(t *testing.T) { + ctx := context.Background() + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_1", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + })) + + // No predicates are enforced, and the balance can go negative + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 30, UsdCostBasis: 10})) + assertBalance(t, s, "token_account_1", -30, -10, true) + + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 100, UsdCostBasis: 50})) + assertBalance(t, s, "token_account_1", 70, 40, true) + + // A drain records the delta rather than zeroing, but still closes + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 100, UsdCostBasis: 60})) + assertBalance(t, s, "token_account_1", -30, -20, false) + + // Closed accounts still accumulate + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 5, UsdCostBasis: 5})) + assertBalance(t, s, "token_account_1", -25, -15, false) + + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) + assertBalance(t, s, "token_account_1", -25, -15, false) + }) +} + +func testApplyDeltasAtomicity(t *testing.T, s balance.Store) { + t.Run("testApplyDeltasAtomicity", func(t *testing.T) { + ctx := context.Background() + + for _, tokenAccount := range []string{"token_account_1", "token_account_2"} { + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: tokenAccount, + OwnerAccount: "owner", + MintAccount: "mint", + Quarks: 100, + IsOpen: true, + IsBackfilled: true, + })) + } + + // A transfer applies both sides + require.NoError(t, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 40, UsdCostBasis: 4}, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 40, UsdCostBasis: 4}, + )) + assertBalance(t, s, "token_account_1", 60, -4, true) + assertBalance(t, s, "token_account_2", 140, 4, true) + + // A failure on either side rolls back the other, regardless of order + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 500}, + &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaDebit, Quarks: 141}, + )) + assertBalance(t, s, "token_account_1", 60, -4, true) + assertBalance(t, s, "token_account_2", 140, 4, true) + + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 61}, + &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 500}, + )) + assertBalance(t, s, "token_account_1", 60, -4, true) + assertBalance(t, s, "token_account_2", 140, 4, true) + + // Multiple deltas to the same account apply in kind order: credit, then debit + require.NoError(t, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 100}, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 50}, + )) + assertBalance(t, s, "token_account_1", 10, -4, true) + }) +} + +func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { + t.Run("testApplyDeltasConcurrency", func(t *testing.T) { + ctx := context.Background() + + const initialBalance = 20 + const attempts = 50 + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "sender", + OwnerAccount: "owner_1", + MintAccount: "mint", + Quarks: initialBalance, + IsOpen: true, + IsBackfilled: true, + })) + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "receiver", + OwnerAccount: "owner_2", + MintAccount: "mint", + IsOpen: true, + IsBackfilled: true, + })) + + // Concurrent sends of 1 quark each: exactly initialBalance succeed, and + // the rest fail with an insufficient balance. Every send credits the + // receiver in the same batch. + var wg sync.WaitGroup + results := make(chan error, attempts) + for range attempts { + wg.Go(func() { + results <- s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "sender", Kind: balance.DeltaDebit, Quarks: 1, UsdCostBasis: 1}, + &balance.Delta{TokenAccount: "receiver", Kind: balance.DeltaCredit, Quarks: 1, UsdCostBasis: 1}, + ) + }) + } + wg.Wait() + close(results) + + var succeeded, insufficient int + for err := range results { + switch err { + case nil: + succeeded++ + case balance.ErrInsufficientBalance: + insufficient++ + default: + require.NoError(t, err) + } + } + assert.Equal(t, initialBalance, succeeded) + assert.Equal(t, attempts-initialBalance, insufficient) + + assertBalance(t, s, "sender", 0, -initialBalance, true) + assertBalance(t, s, "receiver", initialBalance, initialBalance, true) + + // Concurrent credits never fail + wg = sync.WaitGroup{} + results = make(chan error, attempts) + for range attempts { + wg.Go(func() { + results <- s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "receiver", Kind: balance.DeltaCredit, Quarks: 1}) + }) + } + wg.Wait() + close(results) + for err := range results { + require.NoError(t, err) + } + assertBalance(t, s, "receiver", initialBalance+attempts, initialBalance, true) + + // Concurrent drains: exactly one wins + wg = sync.WaitGroup{} + results = make(chan error, attempts) + for range attempts { + wg.Go(func() { + results <- s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "receiver", Kind: balance.DeltaDrain, Quarks: initialBalance + attempts}) + }) + } + wg.Wait() + close(results) + + var drained, closed int + for err := range results { + switch err { + case nil: + drained++ + case balance.ErrAccountClosed: + closed++ + default: + require.NoError(t, err) + } + } + assert.Equal(t, 1, drained) + assert.Equal(t, attempts-1, closed) + assertBalance(t, s, "receiver", 0, 0, false) + }) +} + +func testBackfill(t *testing.T, s balance.Store) { + t.Run("testBackfill", func(t *testing.T) { + ctx := context.Background() + + fn := func(quarks, usdCostBasis int64, isOpen bool) balance.BackfillFunc { + return func(ctx context.Context) (*balance.BackfillResult, error) { + return &balance.BackfillResult{Quarks: quarks, UsdCostBasis: usdCostBasis, IsOpen: isOpen}, nil + } + } + + assert.Equal(t, balance.ErrRecordNotFound, s.Backfill(ctx, "token_account_1", fn(1, 1, true))) + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_1", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + })) + + // Deltas recorded before the backfill are discarded by it + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 30, UsdCostBasis: 10})) + assertBalance(t, s, "token_account_1", -30, -10, true) + + // A failing computation leaves the record untouched + assert.Error(t, s.Backfill(ctx, "token_account_1", func(ctx context.Context) (*balance.BackfillResult, error) { + return nil, assert.AnError + })) + record, err := s.Get(ctx, "token_account_1") + require.NoError(t, err) + assert.False(t, record.IsBackfilled) + assert.EqualValues(t, -30, record.Quarks) + + // So does a negative computed balance + assert.Equal(t, balance.ErrNegativeBalance, s.Backfill(ctx, "token_account_1", fn(-1, 0, true))) + record, err = s.Get(ctx, "token_account_1") + require.NoError(t, err) + assert.False(t, record.IsBackfilled) + assert.EqualValues(t, -30, record.Quarks) + + require.NoError(t, s.Backfill(ctx, "token_account_1", fn(500, 250, true))) + record, err = s.Get(ctx, "token_account_1") + require.NoError(t, err) + assert.True(t, record.IsBackfilled) + assertBalance(t, s, "token_account_1", 500, 250, true) + + // Predicates are enforced from now on + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 501})) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 500})) + assertBalance(t, s, "token_account_1", 0, 250, true) + + called := false + assert.Equal(t, balance.ErrAlreadyBackfilled, s.Backfill(ctx, "token_account_1", func(ctx context.Context) (*balance.BackfillResult, error) { + called = true + return &balance.BackfillResult{}, nil + })) + assert.False(t, called) + assertBalance(t, s, "token_account_1", 0, 250, true) + + // A backfill can establish the account as closed, e.g. a claimed gift card + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_2", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + })) + require.NoError(t, s.Backfill(ctx, "token_account_2", fn(0, 0, false))) + assertBalance(t, s, "token_account_2", 0, 0, false) + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 1})) + }) +} + func testCachedBalanceVersionHappyPath(t *testing.T, s balance.Store) { t.Run("testCachedBalanceVersionHappyPath", func(t *testing.T) { ctx := context.Background() for i := range 100 { - for j := 0; j < 10; j++ { + for range 10 { currentVersion, err := s.GetCachedVersion(ctx, "token_account_1") require.NoError(t, err) assert.EqualValues(t, i, currentVersion) @@ -114,6 +583,24 @@ func testExternalCheckpointHappyPath(t *testing.T, s balance.Store) { }) } +func assertBalance(t *testing.T, s balance.Store, tokenAccount string, quarks, usdCostBasis int64, isOpen bool) { + record, err := s.Get(context.Background(), tokenAccount) + require.NoError(t, err) + assert.EqualValues(t, quarks, record.Quarks, "quarks") + assert.EqualValues(t, usdCostBasis, record.UsdCostBasis, "usd market value") + assert.Equal(t, isOpen, record.IsOpen, "is open") +} + +func assertEquivalentRecords(t *testing.T, obj1, obj2 *balance.Record) { + assert.Equal(t, obj1.TokenAccount, obj2.TokenAccount) + assert.Equal(t, obj1.OwnerAccount, obj2.OwnerAccount) + assert.Equal(t, obj1.MintAccount, obj2.MintAccount) + assert.Equal(t, obj1.Quarks, obj2.Quarks) + assert.Equal(t, obj1.UsdCostBasis, obj2.UsdCostBasis) + assert.Equal(t, obj1.IsOpen, obj2.IsOpen) + assert.Equal(t, obj1.IsBackfilled, obj2.IsBackfilled) +} + func assertEquivalentExternalCheckpoingRecords(t *testing.T, obj1, obj2 *balance.ExternalCheckpointRecord) { assert.Equal(t, obj1.TokenAccount, obj2.TokenAccount) assert.Equal(t, obj1.Quarks, obj2.Quarks) diff --git a/ocp/data/internal.go b/ocp/data/internal.go index 1304fa4..a3111f3 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -121,6 +121,14 @@ type DatabaseData interface { // Balance // -------------------------------------------------------------------------------- + CreateBalance(ctx context.Context, record *balance.Record) error + GetBalance(ctx context.Context, tokenAccount string) (*balance.Record, error) + GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) + GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) + GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) + GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) + ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error + BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) AdvanceCachedBalanceVersion(ctx context.Context, account string, currentVersion uint64) error CheckNotClosedForBalanceUpdate(ctx context.Context, account string) error @@ -460,6 +468,30 @@ func (dp *DatabaseProvider) HasFeeAction(ctx context.Context, intent string, fee // Balance // -------------------------------------------------------------------------------- +func (dp *DatabaseProvider) CreateBalance(ctx context.Context, record *balance.Record) error { + return dp.balance.Create(ctx, record) +} +func (dp *DatabaseProvider) GetBalance(ctx context.Context, tokenAccount string) (*balance.Record, error) { + return dp.balance.Get(ctx, tokenAccount) +} +func (dp *DatabaseProvider) GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) { + return dp.balance.GetBatch(ctx, tokenAccounts...) +} +func (dp *DatabaseProvider) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) { + return dp.balance.GetAllByOwner(ctx, owner) +} +func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { + return dp.balance.GetAllByOwnerAndMint(ctx, owner, mint) +} +func (dp *DatabaseProvider) GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { + return dp.balance.GetAllByMint(ctx, mint, minQuarks, cursor, limit, direction) +} +func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error { + return dp.balance.ApplyDeltas(ctx, deltas...) +} +func (dp *DatabaseProvider) BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { + return dp.balance.Backfill(ctx, tokenAccount, fn) +} func (dp *DatabaseProvider) GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) { return dp.balance.GetCachedVersion(ctx, account) } From 28e37d597380c6781a76651ed748beeef28794c7 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 11:57:37 -0400 Subject: [PATCH 02/26] Update balance calculators --- ocp/balance/calculator.go | 158 ++++++++++++++++++++++++++++++--- ocp/balance/calculator_test.go | 148 ++++++++++++++++++++++++++++++ 2 files changed, 296 insertions(+), 10 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 606445c..bedd437 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -2,13 +2,17 @@ package balance import ( "context" + "math" "time" "github.com/pkg/errors" + commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" + "github.com/code-payments/ocp-server/metrics" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" + "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/data/timelock" "github.com/code-payments/ocp-server/solana" @@ -103,14 +107,22 @@ func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccoun return 0, ErrNotManagedByCode } - // Pick a set of strategies relevant for the type of account, so we can optimize - // the number of DB calls. - // - // Overall, we're using a simple strategy that iterates over an account's history - // to unblock a scheduler implementation optimized for privacy. - // - // todo: Come up with a heurisitc that enables some form of checkpointing, so - // we're not iterating over all records every time. + // Prefer the materialized balance record, when the account has one that + // reflects its full history. + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err == nil && balanceRecord.IsBackfilled { + quarks, err := quarksFromRecord(balanceRecord) + if err != nil { + tracer.OnError(err) + return 0, err + } + return quarks, nil + } else if err != nil && err != balance.ErrRecordNotFound { + tracer.OnError(err) + return 0, err + } + + // Otherwise, fall back to iterating over the account's history. strategies := []Strategy{ NetBalanceFromIntentActions(ctx, data), FundingFromExternalDeposits(ctx, data), @@ -335,12 +347,138 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide tokenAccounts = append(tokenAccounts, timelockRecord.VaultAddress) } - return CalculateBatch( + // Prefer materialized balance records, and only iterate over history for + // accounts that don't yet have a fully backfilled one. + balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccounts...) + if err != nil { + return nil, err + } + + res := make(map[string]uint64, 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 { + return nil, err + } + res[tokenAccount] = quarks + } + + if len(remaining) == 0 { + return res, nil + } + + legacyRes, err := CalculateBatch( ctx, - tokenAccounts, + remaining, NetBalanceFromIntentActionsBatch(ctx, data), FundingFromExternalDepositsBatch(ctx, data), ) + if err != nil { + return nil, err + } + for tokenAccount, quarks := range legacyRes { + res[tokenAccount] = quarks + } + return res, nil +} + +// CalculateUsdCostBasisFromCache calculates a token account's USD cost basis, +// in balance.UsdQuarksPerUnit, using cached values. +// +// Note: Unlike quark balances, a cost basis for an account not managed by Code +// is still meaningful, so no timelock check is performed. +func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (int64, error) { + tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "CalculateUsdCostBasisFromCache") + tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) + defer tracer.End() + + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err == nil && balanceRecord.IsBackfilled { + return balanceRecord.UsdCostBasis, nil + } else if err != nil && err != balance.ErrRecordNotFound { + tracer.OnError(err) + return 0, err + } + + res, err := legacyUsdCostBasis(ctx, data, tokenAccount.PublicKey().ToBase58()) + if err != nil { + tracer.OnError(err) + return 0, err + } + return res, nil +} + +// BatchCalculateUsdCostBasisFromCache is like CalculateUsdCostBasisFromCache, +// but for a set of token accounts. +func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccounts ...*common.Account) (map[string]int64, error) { + tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateUsdCostBasisFromCache") + defer tracer.End() + + tokenAccountStrings := make([]string, len(tokenAccounts)) + for i, tokenAccount := range tokenAccounts { + tokenAccountStrings[i] = tokenAccount.PublicKey().ToBase58() + } + + balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccountStrings...) + if err != nil { + tracer.OnError(err) + return nil, err + } + + res := make(map[string]int64, len(tokenAccounts)) + for _, tokenAccount := range tokenAccountStrings { + balanceRecord, ok := balanceRecords[tokenAccount] + if ok && balanceRecord.IsBackfilled { + res[tokenAccount] = balanceRecord.UsdCostBasis + continue + } + + // todo: The legacy calculation has a batch variant by owner, but the + // fallback is temporary and per-owner batching doesn't map onto + // token accounts without an extra lookup anyway. + usdCostBasis, err := legacyUsdCostBasis(ctx, data, tokenAccount) + if err != nil { + tracer.OnError(err) + return nil, err + } + res[tokenAccount] = usdCostBasis + } + return res, nil +} + +// legacyUsdCostBasis derives a token account's cost basis from the owner-level +// intent aggregate, which is only defined for primary accounts. +func legacyUsdCostBasis(ctx context.Context, data ocp_data.Provider, tokenAccount string) (int64, error) { + accountInfoRecord, err := data.GetAccountInfoByTokenAddress(ctx, tokenAccount) + if err == account.ErrAccountInfoNotFound { + return 0, nil + } else if err != nil { + return 0, err + } + + if accountInfoRecord.AccountType != commonpb.AccountType_PRIMARY { + return 0, nil + } + + usdCostBasis, err := data.GetUsdCostBasis(ctx, accountInfoRecord.OwnerAccount, accountInfoRecord.MintAccount) + if err != nil { + return 0, err + } + return int64(math.Round(usdCostBasis * balance.UsdQuarksPerUnit)), nil +} + +func quarksFromRecord(record *balance.Record) (uint64, error) { + if record.Quarks < 0 { + return 0, ErrNegativeBalance + } + return uint64(record.Quarks), nil } // NetBalanceFromIntentActionsBatch is a balance calculation strategy that incorporates diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 93fba32..5d7076e 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -16,6 +16,7 @@ import ( ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" + "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/data/deposit" "github.com/code-payments/ocp-server/ocp/data/intent" "github.com/code-payments/ocp-server/ocp/data/transaction" @@ -337,6 +338,153 @@ func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { assert.Equal(t, ErrNotManagedByCode, err) } +func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { + env := setupBalanceTestEnv(t) + + vmConfig := testutil.NewRandomVmConfig(t, true) + backfilledOwner := testutil.NewRandomAccount(t) + backfilledAccount, err := backfilledOwner.ToTimelockVault(vmConfig) + require.NoError(t, err) + pendingOwner := testutil.NewRandomAccount(t) + pendingAccount, err := pendingOwner.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, pendingOwner, legacyOwner}, + transactions: []balanceTestTransaction{ + {source: externalAccount, destination: backfilledAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, + {source: externalAccount, destination: pendingAccount, quantity: 22, transactionState: transaction.ConfirmationFinalized}, + {source: externalAccount, destination: legacyAccount, quantity: 33, transactionState: transaction.ConfirmationFinalized}, + }, + } + + setupBalanceTestData(t, env, data) + + // A backfilled record is authoritative, 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, + IsOpen: true, + IsBackfilled: true, + })) + + // A record that isn't backfilled is ignored in favour of history + require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ + TokenAccount: pendingAccount.PublicKey().ToBase58(), + OwnerAccount: pendingOwner.PublicKey().ToBase58(), + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + Quarks: -5, + IsOpen: true, + })) + + expected := map[string]uint64{ + backfilledAccount.PublicKey().ToBase58(): 42, + pendingAccount.PublicKey().ToBase58(): 22, + legacyAccount.PublicKey().ToBase58(): 33, + } + + for tokenAccount, expectedQuarks := range expected { + account, err := common.NewAccountFromPublicKeyString(tokenAccount) + require.NoError(t, err) + + actual, err := CalculateFromCache(env.ctx, env.data, account) + require.NoError(t, err) + assert.EqualValues(t, expectedQuarks, actual, tokenAccount) + } + + balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, backfilledAccount, pendingAccount, legacyAccount) + require.NoError(t, err) + assert.Equal(t, expected, balanceByAccount) + + var allAccountRecords []*common.AccountRecords + for _, owner := range data.codeUsers { + accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner) + require.NoError(t, err) + allAccountRecords = append(allAccountRecords, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) + } + balanceByAccount, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, allAccountRecords...) + require.NoError(t, err) + assert.Equal(t, expected, balanceByAccount) +} + +func TestUsdCostBasisCalculationMethods(t *testing.T) { + env := setupBalanceTestEnv(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) + + data := &balanceTestData{ + vmConfig: vmConfig, + codeUsers: []*common.Account{backfilledOwner, legacyOwner}, + } + + setupBalanceTestData(t, env, data) + + require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ + TokenAccount: backfilledAccount.PublicKey().ToBase58(), + OwnerAccount: backfilledOwner.PublicKey().ToBase58(), + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + UsdCostBasis: -123456, + IsOpen: true, + IsBackfilled: true, + })) + + // The legacy calculation is owner-level, and derived from intents + require.NoError(t, env.data.SaveIntent(env.ctx, &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.ExternalDeposit, + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + InitiatorOwnerAccount: legacyOwner.PublicKey().ToBase58(), + ExternalDepositMetadata: &intent.ExternalDepositMetadata{ + DestinationTokenAccount: legacyAccount.PublicKey().ToBase58(), + Quantity: 1, + ExchangeCurrency: currency_lib.USD, + ExchangeRate: 1.0, + NativeAmount: 1.5, + UsdMarketValue: 1.5, + }, + State: intent.StateConfirmed, + CreatedAt: time.Now(), + })) + + expected := map[string]int64{ + backfilledAccount.PublicKey().ToBase58(): -123456, + legacyAccount.PublicKey().ToBase58(): 1_500_000, + } + + for tokenAccount, expectedUsdCostBasis := range expected { + account, err := common.NewAccountFromPublicKeyString(tokenAccount) + require.NoError(t, err) + + actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, account) + require.NoError(t, err) + assert.EqualValues(t, expectedUsdCostBasis, actual, tokenAccount) + } + + usdCostBasisByAccount, err := BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, backfilledAccount, legacyAccount) + require.NoError(t, err) + assert.Equal(t, expected, usdCostBasisByAccount) + + // Accounts unknown to the system have no cost basis + unknownAccount := testutil.NewRandomAccount(t) + actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, unknownAccount) + require.NoError(t, err) + assert.EqualValues(t, 0, actual) +} + func TestDefaultCalculation_ExternalAccount(t *testing.T) { env := setupBalanceTestEnv(t) externalAccount := testutil.NewRandomAccount(t) From 823bcf49c4a6cbe5560d50884af32961b2149277 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 13:17:34 -0400 Subject: [PATCH 03/26] Updates to balance store implementations --- ocp/data/balance/memory/store.go | 69 ++++++++++++++++++- ocp/data/balance/memory/store_legacy.go | 65 ------------------ ocp/data/balance/postgres/model.go | 84 ++++++++++++++++++++++- ocp/data/balance/postgres/model_legacy.go | 78 --------------------- ocp/data/balance/postgres/store.go | 26 +++++++ ocp/data/balance/postgres/store_legacy.go | 28 -------- ocp/data/balance/record.go | 14 ++++ ocp/data/balance/store.go | 31 +++++---- ocp/data/balance/tests/tests.go | 21 +++++- 9 files changed, 226 insertions(+), 190 deletions(-) diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 00cec4a..b50b0d4 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -162,7 +162,10 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { if !ok { original, ok := s.balanceRecordsByTokenAccount[delta.TokenAccount] if !ok { - continue // Not an account we track + if delta.Kind == balance.DeltaCredit { + continue // Credits to accounts we don't track, like external wallets, are expected + } + return balance.ErrRecordNotFound // Everything else only ever targets accounts we track } cloned := original.Clone() item = &cloned @@ -290,3 +293,67 @@ func (s *store) reset() { s.externalCheckpointRecords = nil s.last = 0 } + +// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint +func (s *store) SaveExternalCheckpoint(_ context.Context, data *balance.ExternalCheckpointRecord) error { + if err := data.Validate(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + + s.last++ + if item := s.findExternalCheckpoint(data); item != nil { + if data.SlotCheckpoint <= item.SlotCheckpoint { + return balance.ErrStaleCheckpoint + } + + item.SlotCheckpoint = data.SlotCheckpoint + item.Quarks = data.Quarks + item.LastUpdatedAt = time.Now() + item.CopyTo(data) + } else { + if data.Id == 0 { + data.Id = s.last + } + data.LastUpdatedAt = time.Now() + c := data.Clone() + s.externalCheckpointRecords = append(s.externalCheckpointRecords, &c) + } + + return nil +} + +// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint +func (s *store) GetExternalCheckpoint(_ context.Context, account string) (*balance.ExternalCheckpointRecord, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if item := s.findExternalCheckpointByTokenAccount(account); item != nil { + cloned := item.Clone() + return &cloned, nil + } + return nil, balance.ErrCheckpointNotFound +} + +func (s *store) findExternalCheckpoint(data *balance.ExternalCheckpointRecord) *balance.ExternalCheckpointRecord { + for _, item := range s.externalCheckpointRecords { + if item.Id == data.Id { + return item + } + if data.TokenAccount == item.TokenAccount { + return item + } + } + return nil +} + +func (s *store) findExternalCheckpointByTokenAccount(account string) *balance.ExternalCheckpointRecord { + for _, item := range s.externalCheckpointRecords { + if account == item.TokenAccount { + return item + } + } + return nil +} diff --git a/ocp/data/balance/memory/store_legacy.go b/ocp/data/balance/memory/store_legacy.go index 580ffd4..9e062f9 100644 --- a/ocp/data/balance/memory/store_legacy.go +++ b/ocp/data/balance/memory/store_legacy.go @@ -2,7 +2,6 @@ package memory import ( "context" - "time" "github.com/code-payments/ocp-server/ocp/data/balance" ) @@ -65,67 +64,3 @@ func (s *store) MarkAsClosed(ctx context.Context, account string) error { return nil } - -// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint -func (s *store) SaveExternalCheckpoint(_ context.Context, data *balance.ExternalCheckpointRecord) error { - if err := data.Validate(); err != nil { - return err - } - - s.mu.Lock() - defer s.mu.Unlock() - - s.last++ - if item := s.findExternalCheckpoint(data); item != nil { - if data.SlotCheckpoint <= item.SlotCheckpoint { - return balance.ErrStaleCheckpoint - } - - item.SlotCheckpoint = data.SlotCheckpoint - item.Quarks = data.Quarks - item.LastUpdatedAt = time.Now() - item.CopyTo(data) - } else { - if data.Id == 0 { - data.Id = s.last - } - data.LastUpdatedAt = time.Now() - c := data.Clone() - s.externalCheckpointRecords = append(s.externalCheckpointRecords, &c) - } - - return nil -} - -// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint -func (s *store) GetExternalCheckpoint(_ context.Context, account string) (*balance.ExternalCheckpointRecord, error) { - s.mu.Lock() - defer s.mu.Unlock() - - if item := s.findExternalCheckpointByTokenAccount(account); item != nil { - cloned := item.Clone() - return &cloned, nil - } - return nil, balance.ErrCheckpointNotFound -} - -func (s *store) findExternalCheckpoint(data *balance.ExternalCheckpointRecord) *balance.ExternalCheckpointRecord { - for _, item := range s.externalCheckpointRecords { - if item.Id == data.Id { - return item - } - if data.TokenAccount == item.TokenAccount { - return item - } - } - return nil -} - -func (s *store) findExternalCheckpointByTokenAccount(account string) *balance.ExternalCheckpointRecord { - for _, item := range s.externalCheckpointRecords { - if account == item.TokenAccount { - return item - } - } - return nil -} diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index 75366d7..7a9b7c0 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -14,7 +14,8 @@ import ( ) const ( - tableName = "ocp__core_balance" + tableName = "ocp__core_balance" + externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_backfilled, updated_at" ) @@ -227,7 +228,10 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er var current model err = tx.GetContext(ctx, ¤t, `SELECT `+allColumns+` FROM `+tableName+` WHERE token_account = $1`, delta.TokenAccount) if pgutil.IsNoRows(err) { - continue // Not an account we track + if delta.Kind == balance.DeltaCredit { + continue // Credits to accounts we don't track, like external wallets, are expected + } + return balance.ErrRecordNotFound // Everything else only ever targets accounts we track } else if err != nil { return err } @@ -292,3 +296,79 @@ func executeTxWithinCtxOrJoin(ctx context.Context, db *sqlx.DB, fn func(ctx cont } return err } + +type externalCheckpointModel struct { + Id sql.NullInt64 `db:"id"` + + TokenAccount string `db:"token_account"` + Quarks uint64 `db:"quarks"` + SlotCheckpoint uint64 `db:"slot_checkpoint"` + + LastUpdatedAt time.Time `db:"last_updated_at"` +} + +func toExternalCheckpointModel(obj *balance.ExternalCheckpointRecord) (*externalCheckpointModel, error) { + if err := obj.Validate(); err != nil { + return nil, err + } + + return &externalCheckpointModel{ + TokenAccount: obj.TokenAccount, + Quarks: obj.Quarks, + SlotCheckpoint: obj.SlotCheckpoint, + LastUpdatedAt: obj.LastUpdatedAt, + }, nil +} + +func fromExternalCheckpoingModel(obj *externalCheckpointModel) *balance.ExternalCheckpointRecord { + return &balance.ExternalCheckpointRecord{ + Id: uint64(obj.Id.Int64), + TokenAccount: obj.TokenAccount, + Quarks: obj.Quarks, + SlotCheckpoint: obj.SlotCheckpoint, + LastUpdatedAt: obj.LastUpdatedAt, + } +} + +func (m *externalCheckpointModel) dbSave(ctx context.Context, db *sqlx.DB) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + query := `INSERT INTO ` + externalCheckpointTableName + ` + (token_account, quarks, slot_checkpoint, last_updated_at) + VALUES ($1, $2, $3, $4) + + ON CONFLICT (token_account) + DO UPDATE + SET quarks = $2, slot_checkpoint = $3, last_updated_at = $4 + WHERE ` + externalCheckpointTableName + `.token_account = $1 AND ` + externalCheckpointTableName + `.slot_checkpoint < $3 + + RETURNING + id, token_account, quarks, slot_checkpoint, last_updated_at` + + m.LastUpdatedAt = time.Now() + + err := tx.QueryRowxContext( + ctx, + query, + m.TokenAccount, + m.Quarks, + m.SlotCheckpoint, + m.LastUpdatedAt.UTC(), + ).StructScan(m) + + return pgutil.CheckNoRows(err, balance.ErrStaleCheckpoint) + }) +} + +func dbGetExternalCheckpoint(ctx context.Context, db *sqlx.DB, account string) (*externalCheckpointModel, error) { + res := &externalCheckpointModel{} + + query := `SELECT id, token_account, quarks, slot_checkpoint, last_updated_at FROM ` + externalCheckpointTableName + ` + WHERE token_account = $1 + LIMIT 1` + + err := db.GetContext(ctx, res, query, account) + if err != nil { + return nil, pgutil.CheckNoRows(err, balance.ErrCheckpointNotFound) + } + return res, nil +} diff --git a/ocp/data/balance/postgres/model_legacy.go b/ocp/data/balance/postgres/model_legacy.go index dbefe71..247c9a4 100644 --- a/ocp/data/balance/postgres/model_legacy.go +++ b/ocp/data/balance/postgres/model_legacy.go @@ -4,7 +4,6 @@ import ( "context" "database/sql" "errors" - "time" "github.com/jmoiron/sqlx" @@ -15,19 +14,8 @@ import ( const ( cachedBalanceVersionTableName = "ocp__core_cachedbalanceversion" openCloseLocksTableName = "ocp__core_opencloselocks" - externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" ) -type externalCheckpointModel struct { - Id sql.NullInt64 `db:"id"` - - TokenAccount string `db:"token_account"` - Quarks uint64 `db:"quarks"` - SlotCheckpoint uint64 `db:"slot_checkpoint"` - - LastUpdatedAt time.Time `db:"last_updated_at"` -} - func dbGetCachedVersion(ctx context.Context, db *sqlx.DB, account string) (uint64, error) { var res uint64 err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { @@ -131,69 +119,3 @@ func dbMarkAsClosed(ctx context.Context, db *sqlx.DB, account string) error { return nil }) } - -func toExternalCheckpointModel(obj *balance.ExternalCheckpointRecord) (*externalCheckpointModel, error) { - if err := obj.Validate(); err != nil { - return nil, err - } - - return &externalCheckpointModel{ - TokenAccount: obj.TokenAccount, - Quarks: obj.Quarks, - SlotCheckpoint: obj.SlotCheckpoint, - LastUpdatedAt: obj.LastUpdatedAt, - }, nil -} - -func fromExternalCheckpoingModel(obj *externalCheckpointModel) *balance.ExternalCheckpointRecord { - return &balance.ExternalCheckpointRecord{ - Id: uint64(obj.Id.Int64), - TokenAccount: obj.TokenAccount, - Quarks: obj.Quarks, - SlotCheckpoint: obj.SlotCheckpoint, - LastUpdatedAt: obj.LastUpdatedAt, - } -} - -func (m *externalCheckpointModel) dbSave(ctx context.Context, db *sqlx.DB) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - query := `INSERT INTO ` + externalCheckpointTableName + ` - (token_account, quarks, slot_checkpoint, last_updated_at) - VALUES ($1, $2, $3, $4) - - ON CONFLICT (token_account) - DO UPDATE - SET quarks = $2, slot_checkpoint = $3, last_updated_at = $4 - WHERE ` + externalCheckpointTableName + `.token_account = $1 AND ` + externalCheckpointTableName + `.slot_checkpoint < $3 - - RETURNING - id, token_account, quarks, slot_checkpoint, last_updated_at` - - m.LastUpdatedAt = time.Now() - - err := tx.QueryRowxContext( - ctx, - query, - m.TokenAccount, - m.Quarks, - m.SlotCheckpoint, - m.LastUpdatedAt.UTC(), - ).StructScan(m) - - return pgutil.CheckNoRows(err, balance.ErrStaleCheckpoint) - }) -} - -func dbGetExternalCheckpoint(ctx context.Context, db *sqlx.DB, account string) (*externalCheckpointModel, error) { - res := &externalCheckpointModel{} - - query := `SELECT id, token_account, quarks, slot_checkpoint, last_updated_at FROM ` + externalCheckpointTableName + ` - WHERE token_account = $1 - LIMIT 1` - - err := db.GetContext(ctx, res, query, account) - if err != nil { - return nil, pgutil.CheckNoRows(err, balance.ErrCheckpointNotFound) - } - return res, nil -} diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index 1bdbdde..def96ff 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -113,3 +113,29 @@ func fromModels(models []*model) []*balance.Record { } return res } + +// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint +func (s *store) SaveExternalCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { + model, err := toExternalCheckpointModel(record) + if err != nil { + return err + } + + if err := model.dbSave(ctx, s.db); err != nil { + return err + } + + res := fromExternalCheckpoingModel(model) + res.CopyTo(record) + + return nil +} + +// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint +func (s *store) GetExternalCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) { + model, err := dbGetExternalCheckpoint(ctx, s.db, account) + if err != nil { + return nil, err + } + return fromExternalCheckpoingModel(model), nil +} diff --git a/ocp/data/balance/postgres/store_legacy.go b/ocp/data/balance/postgres/store_legacy.go index c66bede..f947bf1 100644 --- a/ocp/data/balance/postgres/store_legacy.go +++ b/ocp/data/balance/postgres/store_legacy.go @@ -2,8 +2,6 @@ package postgres import ( "context" - - "github.com/code-payments/ocp-server/ocp/data/balance" ) // GetCachedVersion implements balance.Store.GetCachedVersion @@ -25,29 +23,3 @@ func (s *store) CheckNotClosed(ctx context.Context, account string) error { func (s *store) MarkAsClosed(ctx context.Context, account string) error { return dbMarkAsClosed(ctx, s.db, account) } - -// SaveExternalCheckpoint implements balance.Store.SaveExternalCheckpoint -func (s *store) SaveExternalCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { - model, err := toExternalCheckpointModel(record) - if err != nil { - return err - } - - if err := model.dbSave(ctx, s.db); err != nil { - return err - } - - res := fromExternalCheckpoingModel(model) - res.CopyTo(record) - - return nil -} - -// GetExternalCheckpoint implements balance.Store.GetExternalCheckpoint -func (s *store) GetExternalCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) { - model, err := dbGetExternalCheckpoint(ctx, s.db, account) - if err != nil { - return nil, err - } - return fromExternalCheckpoingModel(model), nil -} diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index d656ed4..1a32f4b 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -2,6 +2,7 @@ package balance import ( "errors" + "math" "sort" "time" ) @@ -11,6 +12,19 @@ import ( // USD cost basis is exactly its quark balance. const UsdQuarksPerUnit = 1_000_000 +// UsdCostBasisFromFloat converts a USD value into UsdQuarksPerUnit, rounding +// to the nearest unit. This is the single conversion point, so every caller +// rounds identically. +func UsdCostBasisFromFloat(usd float64) int64 { + return int64(math.Round(usd * UsdQuarksPerUnit)) +} + +// UsdCostBasisToFloat converts a value in UsdQuarksPerUnit back into USD. Use +// it only at the edge, e.g. when populating a client-facing response. +func UsdCostBasisToFloat(usdCostBasis int64) float64 { + return float64(usdCostBasis) / UsdQuarksPerUnit +} + // Record is the materialized balance of a token account managed by Code. type Record struct { Id uint64 diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index da38824..aae4c3e 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -84,9 +84,12 @@ type Store interface { // applied or none are. Deltas are applied in SortDeltas order. // // Predicates are enforced only on backfilled records; records that are - // not backfilled simply accumulate the change. Deltas for token accounts - // without a record are skipped, since only accounts managed by Code are - // tracked. + // not backfilled simply accumulate the change. + // + // Only accounts managed by OCP have records. A credit to an account + // without one is skipped, since external destinations are routinely paid. + // Any other kind targeting an account without a record is + // ErrRecordNotFound, since funds only ever leave accounts managed by OCP. // // ErrInsufficientBalance is returned when a debit exceeds the balance. // ErrBalanceChanged is returned when a drain or close doesn't match the @@ -105,6 +108,17 @@ type Store interface { // balance, leaving the record untouched. Backfill(ctx context.Context, tokenAccount string, fn BackfillFunc) error + // SaveExternalCheckpoint saves an external balance at a checkpoint. + // + // ErrStaleCheckpoint is returned if the checkpoint is outdated + SaveExternalCheckpoint(ctx context.Context, record *ExternalCheckpointRecord) error + + // GetExternalCheckpoint gets an exeternal balance checkpoint for a + // given account. + // + // ErrCheckpointNotFound is returend if no DB record exists. + GetExternalCheckpoint(ctx context.Context, account string) (*ExternalCheckpointRecord, error) + // GetCachedVersion gets the current cached balance version, which can be used // for optimistic locking cached balances for operations with outgoing transfers. // @@ -136,15 +150,4 @@ type Store interface { // Note: Use ApplyDeltas with DeltaDrain or DeltaClose. Retained for // accounts that are not yet backfilled. MarkAsClosed(ctx context.Context, account string) error - - // SaveExternalCheckpoint saves an external balance at a checkpoint. - // - // ErrStaleCheckpoint is returned if the checkpoint is outdated - SaveExternalCheckpoint(ctx context.Context, record *ExternalCheckpointRecord) error - - // GetExternalCheckpoint gets an exeternal balance checkpoint for a - // given account. - // - // ErrCheckpointNotFound is returend if no DB record exists. - GetExternalCheckpoint(ctx context.Context, account string) (*ExternalCheckpointRecord, error) } diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index e9e654d..2b9a6d5 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -192,8 +192,25 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { IsBackfilled: true, })) - // Deltas for accounts without a record are skipped - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDebit, Quarks: 1})) + // Credits to accounts without a record are skipped, but nothing else is + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1})) + assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDebit, Quarks: 1})) + assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDrain, Quarks: 1})) + assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaClose})) + + // A batch mixing a tracked account with an untracked credit, like a + // withdrawal to an external wallet, applies the tracked side + require.NoError(t, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1}, + &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1}, + )) + require.NoError(t, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1}, + &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1}, + )) + assertBalance(t, s, "token_account_1", 0, 0, true) // Invalid deltas are rejected assert.Error(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit})) From 8231ed67df1f2704dad100099b8c8b56eded6b7c Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 13:17:49 -0400 Subject: [PATCH 04/26] Gate new balance table reads on a config --- ocp/balance/calculator.go | 54 +++++++++++++++++++------------- ocp/balance/calculator_test.go | 56 ++++++++++++++++++++++++++++++++++ ocp/balance/config.go | 17 +++++++++++ 3 files changed, 106 insertions(+), 21 deletions(-) create mode 100644 ocp/balance/config.go diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index bedd437..9a10f42 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -109,17 +109,19 @@ func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccoun // Prefer the materialized balance record, when the account has one that // reflects its full history. - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err == nil && balanceRecord.IsBackfilled { - quarks, err := quarksFromRecord(balanceRecord) - if err != nil { + if enableLedgerReads.Get(ctx) { + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err == nil && balanceRecord.IsBackfilled { + quarks, err := quarksFromRecord(balanceRecord) + if err != nil { + tracer.OnError(err) + return 0, err + } + return quarks, nil + } else if err != nil && err != balance.ErrRecordNotFound { tracer.OnError(err) return 0, err } - return quarks, nil - } else if err != nil && err != balance.ErrRecordNotFound { - tracer.OnError(err) - return 0, err } // Otherwise, fall back to iterating over the account's history. @@ -349,9 +351,13 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide // Prefer materialized balance records, and only iterate over history for // accounts that don't yet have a fully backfilled one. - balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccounts...) - if err != nil { - return nil, err + balanceRecords := make(map[string]*balance.Record) + if enableLedgerReads.Get(ctx) { + var err error + balanceRecords, err = data.GetBalanceBatch(ctx, tokenAccounts...) + if err != nil { + return nil, err + } } res := make(map[string]uint64, len(tokenAccounts)) @@ -399,12 +405,14 @@ func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) defer tracer.End() - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err == nil && balanceRecord.IsBackfilled { - return balanceRecord.UsdCostBasis, nil - } else if err != nil && err != balance.ErrRecordNotFound { - tracer.OnError(err) - return 0, err + if enableLedgerReads.Get(ctx) { + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err == nil && balanceRecord.IsBackfilled { + return balanceRecord.UsdCostBasis, nil + } else if err != nil && err != balance.ErrRecordNotFound { + tracer.OnError(err) + return 0, err + } } res, err := legacyUsdCostBasis(ctx, data, tokenAccount.PublicKey().ToBase58()) @@ -426,10 +434,14 @@ func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Prov tokenAccountStrings[i] = tokenAccount.PublicKey().ToBase58() } - balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccountStrings...) - if err != nil { - tracer.OnError(err) - return nil, err + balanceRecords := make(map[string]*balance.Record) + if enableLedgerReads.Get(ctx) { + var err error + balanceRecords, err = data.GetBalanceBatch(ctx, tokenAccountStrings...) + if err != nil { + tracer.OnError(err) + return nil, err + } } res := make(map[string]int64, len(tokenAccounts)) diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 5d7076e..c662202 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -11,6 +11,8 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" + "github.com/code-payments/ocp-server/config/memory" + "github.com/code-payments/ocp-server/config/wrapper" currency_lib "github.com/code-payments/ocp-server/currency" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" @@ -340,6 +342,7 @@ func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) + enableLedgerReadsForTest(t) vmConfig := testutil.NewRandomVmConfig(t, true) backfilledOwner := testutil.NewRandomAccount(t) @@ -415,8 +418,53 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { assert.Equal(t, expected, balanceByAccount) } +func TestDefaultCalculationMethods_BalanceRecordReadsDisabled(t *testing.T) { + env := setupBalanceTestEnv(t) + + vmConfig := testutil.NewRandomVmConfig(t, true) + owner := testutil.NewRandomAccount(t) + tokenAccount, err := owner.ToTimelockVault(vmConfig) + require.NoError(t, err) + + externalAccount := testutil.NewRandomAccount(t) + + data := &balanceTestData{ + vmConfig: vmConfig, + codeUsers: []*common.Account{owner}, + transactions: []balanceTestTransaction{ + {source: externalAccount, destination: tokenAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, + }, + } + + setupBalanceTestData(t, env, data) + + // A backfilled record exists, but reads are disabled, so history wins + require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ + TokenAccount: tokenAccount.PublicKey().ToBase58(), + OwnerAccount: owner.PublicKey().ToBase58(), + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + Quarks: 42, + UsdCostBasis: 123, + IsOpen: true, + IsBackfilled: true, + })) + + actual, err := CalculateFromCache(env.ctx, env.data, tokenAccount) + require.NoError(t, err) + assert.EqualValues(t, 11, actual) + + balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) + require.NoError(t, err) + assert.EqualValues(t, 11, balanceByAccount[tokenAccount.PublicKey().ToBase58()]) + + usdCostBasis, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) + require.NoError(t, err) + assert.EqualValues(t, 0, usdCostBasis) +} + func TestUsdCostBasisCalculationMethods(t *testing.T) { env := setupBalanceTestEnv(t) + enableLedgerReadsForTest(t) vmConfig := testutil.NewRandomVmConfig(t, true) backfilledOwner := testutil.NewRandomAccount(t) @@ -494,6 +542,14 @@ func TestDefaultCalculation_ExternalAccount(t *testing.T) { // Note: not possible with batch method, since we wouldn't have account records } +func enableLedgerReadsForTest(t *testing.T) { + previous := enableLedgerReads + enableLedgerReads = wrapper.NewBoolConfig(memory.NewConfig(true), defaultEnableLedgerReads) + t.Cleanup(func() { + enableLedgerReads = previous + }) +} + type balanceTestEnv struct { ctx context.Context data ocp_data.Provider diff --git a/ocp/balance/config.go b/ocp/balance/config.go new file mode 100644 index 0000000..a85e1e4 --- /dev/null +++ b/ocp/balance/config.go @@ -0,0 +1,17 @@ +package balance + +import ( + "github.com/code-payments/ocp-server/config" + "github.com/code-payments/ocp-server/config/env" +) + +const ( + // EnableLedgerReadsConfigEnvName gates whether balance calculators read + // from the new ocp__core_balance ledger. When disabled, calculators use + // the legacy strategies exclusively. + EnableLedgerReadsConfigEnvName = "BALANCE_ENABLE_LEDGER_READS" + + defaultEnableLedgerReads = false +) + +var enableLedgerReads config.Bool = env.NewBoolConfig(EnableLedgerReadsConfigEnvName, defaultEnableLedgerReads) From 9d26aaca31d2a0181a86131e7e128e853d71f2da Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Thu, 27 Aug 2026 10:33:47 -0400 Subject: [PATCH 05/26] Add more balance utilities --- ocp/balance/config.go | 12 +- ocp/balance/delta.go | 327 ++++++++++++++++++++++ ocp/balance/delta_test.go | 431 +++++++++++++++++++++++++++++ ocp/balance/ledger.go | 140 ++++++++++ ocp/balance/ledger_test.go | 194 +++++++++++++ ocp/data/balance/memory/store.go | 11 +- ocp/data/balance/postgres/model.go | 5 +- ocp/data/balance/postgres/store.go | 6 +- ocp/data/balance/record.go | 27 ++ ocp/data/balance/record_test.go | 44 +++ ocp/data/balance/store.go | 8 +- ocp/data/balance/tests/tests.go | 34 ++- 12 files changed, 1206 insertions(+), 33 deletions(-) create mode 100644 ocp/balance/delta.go create mode 100644 ocp/balance/delta_test.go create mode 100644 ocp/balance/ledger.go create mode 100644 ocp/balance/ledger_test.go create mode 100644 ocp/data/balance/record_test.go diff --git a/ocp/balance/config.go b/ocp/balance/config.go index a85e1e4..78be8a4 100644 --- a/ocp/balance/config.go +++ b/ocp/balance/config.go @@ -11,7 +11,15 @@ const ( // the legacy strategies exclusively. EnableLedgerReadsConfigEnvName = "BALANCE_ENABLE_LEDGER_READS" - defaultEnableLedgerReads = false + // EnableLedgerWritesConfigEnvName gates whether ApplyDeltasInTx writes + // to the ledger at all. When disabled, it is a no-op. + EnableLedgerWritesConfigEnvName = "BALANCE_ENABLE_LEDGER_WRITES" + + defaultEnableLedgerReads = false + defaultEnableLedgerWrites = false ) -var enableLedgerReads config.Bool = env.NewBoolConfig(EnableLedgerReadsConfigEnvName, defaultEnableLedgerReads) +var ( + enableLedgerReads config.Bool = env.NewBoolConfig(EnableLedgerReadsConfigEnvName, defaultEnableLedgerReads) + enableLedgerWrites config.Bool = env.NewBoolConfig(EnableLedgerWritesConfigEnvName, defaultEnableLedgerWrites) +) diff --git a/ocp/balance/delta.go b/ocp/balance/delta.go new file mode 100644 index 0000000..9007224 --- /dev/null +++ b/ocp/balance/delta.go @@ -0,0 +1,327 @@ +package balance + +import ( + "errors" + "fmt" + + transactionpb "github.com/code-payments/ocp-protobuf-api/generated/go/transaction/v1" + + "github.com/code-payments/ocp-server/ocp/config" + "github.com/code-payments/ocp-server/ocp/data/action" + "github.com/code-payments/ocp-server/ocp/data/balance" + "github.com/code-payments/ocp-server/ocp/data/intent" +) + +// ErrUnsupportedBalanceChange is returned when records describe a balance +// change the ledger has no rule for. It's a bug to commit such records with +// ledger writes enabled, so callers should fail the transaction. +var ErrUnsupportedBalanceChange = errors.New("unsupported balance change") + +// DeltasForSubmittedIntent returns the ledger deltas for an intent and its +// actions as committed by SubmitIntent. Every account that funds move +// between gets a delta, and USD cost basis moves with the funds. +// +// Actions without a quantity are deferred (eg. a gift card auto-return) and +// contribute nothing until the quantity is set, at which point the flow +// setting it is responsible for the delta. +// +// USD cost basis is attributed per intent: the intent's USD market value is +// the gross amount leaving the source. A fee action carries the configured +// fee's USD value and the principal carries the remainder, so an intent may +// have at most one quantified principal action and at most one fee action. +// withdrawalFeeQuarks is the configured create-on-send withdrawal fee, in +// core mint quarks. +func DeltasForSubmittedIntent(intentRecord *intent.Record, actionRecords []*action.Record, withdrawalFeeQuarks uint64) ([]*balance.Delta, error) { + if err := requireSupported(intentRecord, actionRecords); err != nil { + return nil, err + } + + switch intentRecord.IntentType { + case intent.OpenAccounts, intent.SendPublicPayment, intent.ReceivePaymentsPublicly: + default: + return nil, fmt.Errorf("%w: %d intent is not submitted", ErrUnsupportedBalanceChange, intentRecord.IntentType) + } + + usdByAction, err := usdCostBasisByAction(intentRecord, actionRecords, withdrawalFeeQuarks) + if err != nil { + return nil, err + } + + var deltas []*balance.Delta + for _, actionRecord := range actionRecords { + if actionRecord.Intent != intentRecord.IntentId { + return nil, errors.New("action does not belong to intent") + } + + actionDeltas, err := deltasForAction(actionRecord, usdByAction[actionRecord.ActionId]) + if err != nil { + return nil, err + } + deltas = append(deltas, actionDeltas...) + } + + balance.SortDeltas(deltas) + return deltas, nil +} + +// DeltasForExternalDeposit returns the ledger deltas for an external deposit +// intent, which is created by workers once funds are observed on chain. Only +// confirmed deposits are supported, since that's the only state workers +// commit; the funds are credited to the destination in full. +func DeltasForExternalDeposit(intentRecord *intent.Record) ([]*balance.Delta, error) { + if intentRecord.IntentType != intent.ExternalDeposit { + return nil, fmt.Errorf("%w: %d intent is not an external deposit", ErrUnsupportedBalanceChange, intentRecord.IntentType) + } + if intentRecord.State != intent.StateConfirmed { + return nil, fmt.Errorf("%w: external deposit is not confirmed", ErrUnsupportedBalanceChange) + } + + usdCostBasis, err := UsdCostBasisForIntent(intentRecord) + if err != nil { + return nil, err + } + + metadata := intentRecord.ExternalDepositMetadata + return []*balance.Delta{{ + TokenAccount: metadata.DestinationTokenAccount, + Kind: balance.DeltaCredit, + Quarks: metadata.Quantity, + UsdCostBasis: usdCostBasis, + }}, nil +} + +// DeltasForGiftCardAutoReturn returns the ledger deltas for returning a gift +// card's funds to its issuer. The auto-return action is deferred at issuance +// and contributes nothing until the worker sets its quantity and commits the +// synthetic return intent, which is when this applies. The gift card is +// drained and closed, and the issued value is returned to the issuer. +func DeltasForGiftCardAutoReturn(autoReturnIntent *intent.Record, autoReturnAction *action.Record) ([]*balance.Delta, error) { + if err := requireSupported(autoReturnIntent, []*action.Record{autoReturnAction}); err != nil { + return nil, err + } + + if autoReturnIntent.IntentType != intent.ReceivePaymentsPublicly { + return nil, fmt.Errorf("%w: %d intent is not a gift card return", ErrUnsupportedBalanceChange, autoReturnIntent.IntentType) + } + metadata := autoReturnIntent.ReceivePaymentsPubliclyMetadata + if !metadata.IsIndirectSend || (!metadata.IsReturned && !metadata.IsIssuerVoidingGiftCard) { + return nil, fmt.Errorf("%w: intent is not a gift card return", ErrUnsupportedBalanceChange) + } + + if autoReturnAction.ActionType != action.NoPrivacyWithdraw { + return nil, fmt.Errorf("%w: auto-return is not a withdraw", ErrUnsupportedBalanceChange) + } + if autoReturnAction.Quantity == nil { + return nil, fmt.Errorf("%w: auto-return quantity is not set", ErrUnsupportedBalanceChange) + } + if autoReturnAction.Destination == nil { + return nil, errors.New("destination is required for a withdraw") + } + if autoReturnAction.Source != metadata.Source { + return nil, errors.New("auto-return action does not match intent") + } + + usdCostBasis, err := UsdCostBasisForIntent(autoReturnIntent) + if err != nil { + return nil, err + } + + deltas := []*balance.Delta{ + { + TokenAccount: autoReturnAction.Source, + Kind: balance.DeltaDrain, + Quarks: *autoReturnAction.Quantity, + UsdCostBasis: usdCostBasis, + }, + { + TokenAccount: *autoReturnAction.Destination, + Kind: balance.DeltaCredit, + Quarks: *autoReturnAction.Quantity, + UsdCostBasis: usdCostBasis, + }, + } + balance.SortDeltas(deltas) + return deltas, nil +} + +// DeltasForSwapSellReconciliation returns the ledger deltas for reconciling a +// swap sell's funding payment to the value the sell actually realized. The +// funding payment was committed with an estimated USD market value that the +// swap worker later overwrites, so the source's cost basis is adjusted by +// the difference. No quarks move, and the swap destination isn't tracked by +// the ledger, so only the source is adjusted. +// +// previous and updated are the funding intent before and after the worker +// reconciles its value. The actions are those of the funding intent, which +// identify the source. +func DeltasForSwapSellReconciliation(previous, updated *intent.Record, actionRecords []*action.Record) ([]*balance.Delta, error) { + if previous.IntentId != updated.IntentId { + return nil, errors.New("intent records do not match") + } + if err := requireSupported(updated, actionRecords); err != nil { + return nil, err + } + if updated.IntentType != intent.SendPublicPayment || !updated.SendPublicPaymentMetadata.IsSwapSell { + return nil, fmt.Errorf("%w: intent is not a swap sell", ErrUnsupportedBalanceChange) + } + if previous.IntentType != intent.SendPublicPayment || !previous.SendPublicPaymentMetadata.IsSwapSell { + return nil, fmt.Errorf("%w: previous intent is not a swap sell", ErrUnsupportedBalanceChange) + } + + var funding *action.Record + for _, actionRecord := range actionRecords { + if actionRecord.Intent != updated.IntentId { + return nil, errors.New("action does not belong to intent") + } + if actionRecord.Quantity == nil { + continue + } + if actionRecord.FeeType != nil { + return nil, fmt.Errorf("%w: swap sell pays a fee", ErrUnsupportedBalanceChange) + } + if funding != nil { + return nil, fmt.Errorf("%w: intent pays more than one account", ErrUnsupportedBalanceChange) + } + funding = actionRecord + } + if funding == nil { + return nil, fmt.Errorf("%w: swap sell has no funding action", ErrUnsupportedBalanceChange) + } + + adjustment := balance.UsdCostBasisFromFloat(updated.SendPublicPaymentMetadata.UsdMarketValue) - balance.UsdCostBasisFromFloat(previous.SendPublicPaymentMetadata.UsdMarketValue) + if adjustment == 0 { + return nil, nil + } + + // A debit subtracts the signed basis, so a higher realized value removes + // more basis from the source and a lower one gives some back + return []*balance.Delta{{ + TokenAccount: funding.Source, + Kind: balance.DeltaDebit, + UsdCostBasis: adjustment, + }}, nil +} + +// UsdCostBasisForIntent is the gross USD cost basis moved by an intent, in +// balance.UsdQuarksPerUnit. +func UsdCostBasisForIntent(intentRecord *intent.Record) (int64, error) { + switch intentRecord.IntentType { + case intent.OpenAccounts: + return 0, nil + case intent.ExternalDeposit: + return balance.UsdCostBasisFromFloat(intentRecord.ExternalDepositMetadata.UsdMarketValue), nil + case intent.SendPublicPayment: + return balance.UsdCostBasisFromFloat(intentRecord.SendPublicPaymentMetadata.UsdMarketValue), nil + case intent.ReceivePaymentsPublicly: + return balance.UsdCostBasisFromFloat(intentRecord.ReceivePaymentsPubliclyMetadata.UsdMarketValue), nil + default: + return 0, fmt.Errorf("%w: %d intent", ErrUnsupportedBalanceChange, intentRecord.IntentType) + } +} + +func requireSupported(intentRecord *intent.Record, actionRecords []*action.Record) error { + if intentRecord.IntentType == intent.PublicDistribution { + return fmt.Errorf("%w: public distribution", ErrUnsupportedBalanceChange) + } + if intentRecord.State == intent.StateRevoked { + return fmt.Errorf("%w: revoked intent", ErrUnsupportedBalanceChange) + } + for _, actionRecord := range actionRecords { + if actionRecord.State == action.StateRevoked { + return fmt.Errorf("%w: revoked action", ErrUnsupportedBalanceChange) + } + } + return nil +} + +// usdCostBasisByAction splits an intent's USD cost basis across its +// quantified actions, keyed by action ID. +func usdCostBasisByAction(intentRecord *intent.Record, actionRecords []*action.Record, withdrawalFeeQuarks uint64) (map[uint32]int64, error) { + gross, err := UsdCostBasisForIntent(intentRecord) + if err != nil { + return nil, err + } + + var principal, fee *action.Record + for _, actionRecord := range actionRecords { + if actionRecord.Quantity == nil { + continue + } + + if actionRecord.FeeType != nil { + if fee != nil { + return nil, fmt.Errorf("%w: intent pays more than one fee", ErrUnsupportedBalanceChange) + } + fee = actionRecord + continue + } + + if principal != nil { + return nil, fmt.Errorf("%w: intent pays more than one account", ErrUnsupportedBalanceChange) + } + principal = actionRecord + } + + res := make(map[uint32]int64) + if fee != nil { + switch *fee.FeeType { + case transactionpb.FeePaymentAction_CREATE_ON_SEND_WITHDRAWAL: + // The fee is a fixed core mint amount, so its USD value is fixed + // regardless of how the intent's mint is valued + feeUsd := balance.UsdCostBasisFromFloat(float64(withdrawalFeeQuarks) / float64(config.CoreMintQuarksPerUnit)) + res[fee.ActionId] = feeUsd + gross -= feeUsd + default: + return nil, fmt.Errorf("%w: %s fee", ErrUnsupportedBalanceChange, fee.FeeType.String()) + } + } + if principal != nil { + res[principal.ActionId] = gross + } else if gross != 0 { + return nil, fmt.Errorf("%w: intent has value but no quantified action", ErrUnsupportedBalanceChange) + } + return res, nil +} + +func deltasForAction(actionRecord *action.Record, usdCostBasis int64) ([]*balance.Delta, error) { + switch actionRecord.ActionType { + case action.OpenAccount: + return nil, nil + + case action.CloseEmptyAccount: + return []*balance.Delta{{ + TokenAccount: actionRecord.Source, + Kind: balance.DeltaClose, + }}, nil + + case action.NoPrivacyTransfer, action.NoPrivacyWithdraw: + if actionRecord.Quantity == nil { + return nil, nil + } + if actionRecord.Destination == nil { + return nil, errors.New("destination is required for a transfer") + } + + outgoingKind := balance.DeltaDebit + if actionRecord.ActionType == action.NoPrivacyWithdraw { + outgoingKind = balance.DeltaDrain + } + return []*balance.Delta{ + { + TokenAccount: actionRecord.Source, + Kind: outgoingKind, + Quarks: *actionRecord.Quantity, + UsdCostBasis: usdCostBasis, + }, + { + TokenAccount: *actionRecord.Destination, + Kind: balance.DeltaCredit, + Quarks: *actionRecord.Quantity, + UsdCostBasis: usdCostBasis, + }, + }, nil + + default: + return nil, fmt.Errorf("%w: %d action", ErrUnsupportedBalanceChange, actionRecord.ActionType) + } +} diff --git a/ocp/balance/delta_test.go b/ocp/balance/delta_test.go new file mode 100644 index 0000000..cf30f7c --- /dev/null +++ b/ocp/balance/delta_test.go @@ -0,0 +1,431 @@ +package balance + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + transactionpb "github.com/code-payments/ocp-protobuf-api/generated/go/transaction/v1" + + currency_lib "github.com/code-payments/ocp-server/currency" + "github.com/code-payments/ocp-server/ocp/config" + "github.com/code-payments/ocp-server/ocp/data/action" + "github.com/code-payments/ocp-server/ocp/data/balance" + "github.com/code-payments/ocp-server/ocp/data/intent" + "github.com/code-payments/ocp-server/pointer" + "github.com/code-payments/ocp-server/testutil" +) + +const testWithdrawalFeeQuarks = config.CoreMintQuarksPerUnit / 4 // $0.25 + +func TestDeltasForSubmittedIntent_SendPublicPayment(t *testing.T) { + intentRecord := newDeltaTestSendPublicPaymentIntent(t, 150_000, 1.5) + actionRecords := []*action.Record{ + newDeltaTestTransferAction(intentRecord, 0, "source", "destination", 150_000), + } + + deltas, err := DeltasForSubmittedIntent(intentRecord, actionRecords, testWithdrawalFeeQuarks) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "destination", Kind: balance.DeltaCredit, Quarks: 150_000, UsdCostBasis: 1_500_000}, + {TokenAccount: "source", Kind: balance.DeltaDebit, Quarks: 150_000, UsdCostBasis: 1_500_000}, + }, deltas) +} + +func TestDeltasForSubmittedIntent_WithdrawalWithFee(t *testing.T) { + intentRecord := newDeltaTestSendPublicPaymentIntent(t, 150_000, 1.5) + intentRecord.SendPublicPaymentMetadata.IsWithdrawal = true + + feeAction := newDeltaTestTransferAction(intentRecord, 0, "source", "fee_collector", testWithdrawalFeeQuarks) + feeType := transactionpb.FeePaymentAction_CREATE_ON_SEND_WITHDRAWAL + feeAction.FeeType = &feeType + actionRecords := []*action.Record{ + feeAction, + newDeltaTestTransferAction(intentRecord, 1, "source", "destination", 125_000), + } + + deltas, err := DeltasForSubmittedIntent(intentRecord, actionRecords, testWithdrawalFeeQuarks) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "destination", Kind: balance.DeltaCredit, Quarks: 125_000, UsdCostBasis: 1_250_000}, + {TokenAccount: "fee_collector", Kind: balance.DeltaCredit, Quarks: testWithdrawalFeeQuarks, UsdCostBasis: 250_000}, + {TokenAccount: "source", Kind: balance.DeltaDebit, Quarks: testWithdrawalFeeQuarks, UsdCostBasis: 250_000}, + {TokenAccount: "source", Kind: balance.DeltaDebit, Quarks: 125_000, UsdCostBasis: 1_250_000}, + }, deltas) +} + +func TestDeltasForSubmittedIntent_GiftCardIssuanceAndClaim(t *testing.T) { + // Issuance: a transfer to the gift card plus a deferred auto-return that + // contributes nothing until its quantity is set + issueIntent := newDeltaTestSendPublicPaymentIntent(t, 100_000, 1.0) + issueIntent.SendPublicPaymentMetadata.IsIndirectSend = true + autoReturn := newDeltaTestTransferAction(issueIntent, 1, "gift_card", "source", 0) + autoReturn.ActionType = action.NoPrivacyWithdraw + autoReturn.Quantity = nil + autoReturn.State = action.StateUnknown + actionRecords := []*action.Record{ + newDeltaTestTransferAction(issueIntent, 0, "source", "gift_card", 100_000), + autoReturn, + } + + deltas, err := DeltasForSubmittedIntent(issueIntent, actionRecords, testWithdrawalFeeQuarks) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "gift_card", Kind: balance.DeltaCredit, Quarks: 100_000, UsdCostBasis: 1_000_000}, + {TokenAccount: "source", Kind: balance.DeltaDebit, Quarks: 100_000, UsdCostBasis: 1_000_000}, + }, deltas) + + // Claim: a withdrawal drains the gift card and closes it + claimIntent := &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.ReceivePaymentsPublicly, + MintAccount: "mint", + InitiatorOwnerAccount: "claimer", + ReceivePaymentsPubliclyMetadata: &intent.ReceivePaymentsPubliclyMetadata{ + Source: "gift_card", + Quantity: 100_000, + IsIndirectSend: true, + OriginalExchangeCurrency: currency_lib.USD, + OriginalExchangeRate: 1.0, + OriginalNativeAmount: 1.0, + UsdMarketValue: 1.0, + }, + State: intent.StatePending, + } + claim := newDeltaTestTransferAction(claimIntent, 0, "gift_card", "claimer_primary", 100_000) + claim.ActionType = action.NoPrivacyWithdraw + + deltas, err = DeltasForSubmittedIntent(claimIntent, []*action.Record{claim}, testWithdrawalFeeQuarks) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "claimer_primary", Kind: balance.DeltaCredit, Quarks: 100_000, UsdCostBasis: 1_000_000}, + {TokenAccount: "gift_card", Kind: balance.DeltaDrain, Quarks: 100_000, UsdCostBasis: 1_000_000}, + }, deltas) +} + +func TestDeltasForSubmittedIntent_OpenAccounts(t *testing.T) { + intentRecord := &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.OpenAccounts, + MintAccount: "mint", + InitiatorOwnerAccount: "owner", + OpenAccountsMetadata: &intent.OpenAccountsMetadata{}, + State: intent.StatePending, + } + actionRecords := []*action.Record{{ + Intent: intentRecord.IntentId, + IntentType: intentRecord.IntentType, + ActionId: 0, + ActionType: action.OpenAccount, + Source: "primary", + State: action.StatePending, + }} + + deltas, err := DeltasForSubmittedIntent(intentRecord, actionRecords, testWithdrawalFeeQuarks) + require.NoError(t, err) + assert.Empty(t, deltas) +} + +func TestDeltasForSubmittedIntent_Unsupported(t *testing.T) { + for _, tc := range []struct { + name string + mutate func(*intent.Record, []*action.Record) (*intent.Record, []*action.Record) + }{ + { + name: "revoked intent", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + i.State = intent.StateRevoked + return i, a + }, + }, + { + name: "revoked action", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + a[0].State = action.StateRevoked + return i, a + }, + }, + { + name: "public distribution", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + i.IntentType = intent.PublicDistribution + return i, a + }, + }, + { + name: "external deposit", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + i.IntentType = intent.ExternalDeposit + i.ExternalDepositMetadata = &intent.ExternalDepositMetadata{UsdMarketValue: 1.5} + return i, a + }, + }, + { + name: "more than one payment", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + return i, append(a, newDeltaTestTransferAction(i, 1, "source", "other", 1)) + }, + }, + { + name: "more than one fee", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + feeType := transactionpb.FeePaymentAction_CREATE_ON_SEND_WITHDRAWAL + fee1 := newDeltaTestTransferAction(i, 1, "source", "fee_collector", 1) + fee1.FeeType = &feeType + fee2 := newDeltaTestTransferAction(i, 2, "source", "fee_collector", 1) + fee2.FeeType = &feeType + return i, append(a, fee1, fee2) + }, + }, + { + name: "unknown fee type", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + feeType := transactionpb.FeePaymentAction_FeeType(99) + fee := newDeltaTestTransferAction(i, 1, "source", "fee_collector", 1) + fee.FeeType = &feeType + return i, append(a, fee) + }, + }, + { + name: "action from another intent", + mutate: func(i *intent.Record, a []*action.Record) (*intent.Record, []*action.Record) { + a[0].Intent = "other" + return i, a + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + intentRecord := newDeltaTestSendPublicPaymentIntent(t, 150_000, 1.5) + actionRecords := []*action.Record{ + newDeltaTestTransferAction(intentRecord, 0, "source", "destination", 150_000), + } + intentRecord, actionRecords = tc.mutate(intentRecord, actionRecords) + + _, err := DeltasForSubmittedIntent(intentRecord, actionRecords, testWithdrawalFeeQuarks) + assert.Error(t, err) + }) + } +} + +func TestDeltasForExternalDeposit(t *testing.T) { + intentRecord := &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.ExternalDeposit, + MintAccount: "mint", + InitiatorOwnerAccount: "owner", + ExternalDepositMetadata: &intent.ExternalDepositMetadata{ + DestinationTokenAccount: "destination", + Quantity: 150_000, + ExchangeCurrency: currency_lib.USD, + ExchangeRate: 1.0, + NativeAmount: 1.5, + UsdMarketValue: 1.5, + }, + State: intent.StateConfirmed, + } + + deltas, err := DeltasForExternalDeposit(intentRecord) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "destination", Kind: balance.DeltaCredit, Quarks: 150_000, UsdCostBasis: 1_500_000}, + }, deltas) + + // Only confirmed deposits are committed by workers + for _, state := range []intent.State{intent.StateUnknown, intent.StatePending, intent.StateFailed, intent.StateRevoked} { + intentRecord.State = state + _, err = DeltasForExternalDeposit(intentRecord) + assert.ErrorIs(t, err, ErrUnsupportedBalanceChange) + } + + _, err = DeltasForExternalDeposit(newDeltaTestSendPublicPaymentIntent(t, 1, 1.0)) + assert.ErrorIs(t, err, ErrUnsupportedBalanceChange) +} + +func TestDeltasForGiftCardAutoReturn(t *testing.T) { + newFixtures := func() (*intent.Record, *action.Record) { + intentRecord := &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.ReceivePaymentsPublicly, + MintAccount: "mint", + InitiatorOwnerAccount: "issuer", + ReceivePaymentsPubliclyMetadata: &intent.ReceivePaymentsPubliclyMetadata{ + Source: "gift_card", + Quantity: 100_000, + IsIndirectSend: true, + IsReturned: true, + OriginalExchangeCurrency: currency_lib.USD, + OriginalExchangeRate: 1.0, + OriginalNativeAmount: 1.0, + UsdMarketValue: 1.0, + }, + State: intent.StateConfirmed, + } + actionRecord := &action.Record{ + Intent: "issued_intent", + IntentType: intent.SendPublicPayment, + ActionId: 1, + ActionType: action.NoPrivacyWithdraw, + Source: "gift_card", + Destination: pointer.String("issuer_primary"), + Quantity: pointer.Uint64(100_000), + State: action.StatePending, + } + return intentRecord, actionRecord + } + + intentRecord, actionRecord := newFixtures() + deltas, err := DeltasForGiftCardAutoReturn(intentRecord, actionRecord) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "gift_card", Kind: balance.DeltaDrain, Quarks: 100_000, UsdCostBasis: 1_000_000}, + {TokenAccount: "issuer_primary", Kind: balance.DeltaCredit, Quarks: 100_000, UsdCostBasis: 1_000_000}, + }, deltas) + + // Voiding by the issuer is the same movement + intentRecord, actionRecord = newFixtures() + intentRecord.ReceivePaymentsPubliclyMetadata.IsReturned = false + intentRecord.ReceivePaymentsPubliclyMetadata.IsIssuerVoidingGiftCard = true + _, err = DeltasForGiftCardAutoReturn(intentRecord, actionRecord) + require.NoError(t, err) + + for _, tc := range []struct { + name string + mutate func(*intent.Record, *action.Record) + }{ + {"claim rather than return", func(i *intent.Record, a *action.Record) { + i.ReceivePaymentsPubliclyMetadata.IsReturned = false + }}, + {"not a gift card", func(i *intent.Record, a *action.Record) { + i.ReceivePaymentsPubliclyMetadata.IsIndirectSend = false + }}, + {"deferred action", func(i *intent.Record, a *action.Record) { + a.Quantity = nil + a.State = action.StateUnknown + }}, + {"revoked action", func(i *intent.Record, a *action.Record) { + a.State = action.StateRevoked + }}, + {"not a withdraw", func(i *intent.Record, a *action.Record) { + a.ActionType = action.NoPrivacyTransfer + }}, + {"wrong source", func(i *intent.Record, a *action.Record) { + a.Source = "other_gift_card" + }}, + } { + t.Run(tc.name, func(t *testing.T) { + intentRecord, actionRecord := newFixtures() + tc.mutate(intentRecord, actionRecord) + _, err := DeltasForGiftCardAutoReturn(intentRecord, actionRecord) + assert.Error(t, err) + }) + } +} + +func TestDeltasForSwapSellReconciliation(t *testing.T) { + newFixtures := func(previousUsd, updatedUsd float64) (*intent.Record, *intent.Record, []*action.Record) { + previous := newDeltaTestSendPublicPaymentIntent(t, 150_000, previousUsd) + previous.SendPublicPaymentMetadata.IsSwapSell = true + previous.State = intent.StateConfirmed + + updatedClone := previous.Clone() + updated := &updatedClone + updated.SendPublicPaymentMetadata.UsdMarketValue = updatedUsd + + actionRecords := []*action.Record{ + newDeltaTestTransferAction(previous, 0, "source", "swap", 150_000), + } + return previous, updated, actionRecords + } + + // Realized more than estimated: more basis leaves the source + previous, updated, actionRecords := newFixtures(1.5, 1.75) + deltas, err := DeltasForSwapSellReconciliation(previous, updated, actionRecords) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "source", Kind: balance.DeltaDebit, UsdCostBasis: 250_000}, + }, deltas) + + // Realized less than estimated: basis is returned to the source + previous, updated, actionRecords = newFixtures(1.5, 1.25) + deltas, err = DeltasForSwapSellReconciliation(previous, updated, actionRecords) + require.NoError(t, err) + assert.Equal(t, []*balance.Delta{ + {TokenAccount: "source", Kind: balance.DeltaDebit, UsdCostBasis: -250_000}, + }, deltas) + + // No change is a no-op + previous, updated, actionRecords = newFixtures(1.5, 1.5) + deltas, err = DeltasForSwapSellReconciliation(previous, updated, actionRecords) + require.NoError(t, err) + assert.Empty(t, deltas) + + for _, tc := range []struct { + name string + mutate func(previous, updated *intent.Record, a []*action.Record) []*action.Record + }{ + {"not a swap sell", func(p, u *intent.Record, a []*action.Record) []*action.Record { + u.SendPublicPaymentMetadata.IsSwapSell = false + return a + }}, + {"different intents", func(p, u *intent.Record, a []*action.Record) []*action.Record { + u.IntentId = "other" + return a + }}, + {"revoked", func(p, u *intent.Record, a []*action.Record) []*action.Record { + u.State = intent.StateRevoked + return a + }}, + {"no funding action", func(p, u *intent.Record, a []*action.Record) []*action.Record { + return nil + }}, + {"pays a fee", func(p, u *intent.Record, a []*action.Record) []*action.Record { + feeType := transactionpb.FeePaymentAction_CREATE_ON_SEND_WITHDRAWAL + fee := newDeltaTestTransferAction(u, 1, "source", "fee_collector", 1) + fee.FeeType = &feeType + return append(a, fee) + }}, + {"more than one payment", func(p, u *intent.Record, a []*action.Record) []*action.Record { + return append(a, newDeltaTestTransferAction(u, 1, "source", "other", 1)) + }}, + } { + t.Run(tc.name, func(t *testing.T) { + previous, updated, actionRecords := newFixtures(1.5, 1.75) + actionRecords = tc.mutate(previous, updated, actionRecords) + _, err := DeltasForSwapSellReconciliation(previous, updated, actionRecords) + assert.Error(t, err) + }) + } +} + +func newDeltaTestSendPublicPaymentIntent(t *testing.T, quantity uint64, usd float64) *intent.Record { + return &intent.Record{ + IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + IntentType: intent.SendPublicPayment, + MintAccount: "mint", + InitiatorOwnerAccount: "owner", + SendPublicPaymentMetadata: &intent.SendPublicPaymentMetadata{ + DestinationOwnerAccount: "destination_owner", + DestinationTokenAccount: "destination", + Quantity: quantity, + ExchangeCurrency: currency_lib.USD, + ExchangeRate: 1.0, + NativeAmount: usd, + UsdMarketValue: usd, + }, + State: intent.StatePending, + } +} + +func newDeltaTestTransferAction(intentRecord *intent.Record, actionId uint32, source, destination string, quantity uint64) *action.Record { + return &action.Record{ + Intent: intentRecord.IntentId, + IntentType: intentRecord.IntentType, + ActionId: actionId, + ActionType: action.NoPrivacyTransfer, + Source: source, + Destination: pointer.String(destination), + Quantity: pointer.Uint64(quantity), + State: action.StatePending, + } +} diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go new file mode 100644 index 0000000..631f197 --- /dev/null +++ b/ocp/balance/ledger.go @@ -0,0 +1,140 @@ +package balance + +import ( + "context" + "errors" + "fmt" + + ocp_data "github.com/code-payments/ocp-server/ocp/data" + "github.com/code-payments/ocp-server/ocp/data/account" + "github.com/code-payments/ocp-server/ocp/data/balance" +) + +// ErrUntrackedAccount is returned when funds would leave an account the +// ledger doesn't track. +var ErrUntrackedAccount = errors.New("account is not tracked by the balance ledger") + +// ApplyDeltasInTx applies balance deltas to the ledger. It must be called +// within the DB transaction that commits the records the deltas are derived +// from, so the ledger can never disagree with them. +// +// It is a no-op while ledger writes are disabled. +// +// The ledger only tracks timelock accounts. Credits to any other account, +// like an external wallet or the fee collector, are dropped, since delta +// builders don't know which destinations OCP manages. Outgoing deltas from +// an account the ledger doesn't track are ErrUntrackedAccount, since funds +// only ever leave accounts OCP manages. +// +// Any timelock account in the delta set that has no ledger record yet lazily +// gets one that is not backfilled, so accounts that predate the ledger start +// accumulating deltas on first touch regardless of direction. +// +// Store predicate failures (balance.ErrInsufficientBalance, +// balance.ErrBalanceChanged, balance.ErrAccountClosed) are returned as is +// for the caller to map. +func ApplyDeltasInTx(ctx context.Context, data ocp_data.Provider, deltas ...*balance.Delta) error { + if !enableLedgerWrites.Get(ctx) || len(deltas) == 0 { + return nil + } + + for _, delta := range deltas { + if err := delta.Validate(); err != nil { + return err + } + } + + tracked, err := resolveRecords(ctx, data, deltas) + if err != nil { + return err + } + + var applicable []*balance.Delta + for _, delta := range deltas { + if tracked[delta.TokenAccount] { + applicable = append(applicable, delta) + } else if delta.Kind != balance.DeltaCredit { + return fmt.Errorf("%w: %s", ErrUntrackedAccount, delta.TokenAccount) + } + } + if len(applicable) == 0 { + return nil + } + + return data.ApplyBalanceDeltas(ctx, applicable...) +} + +// CreateRecordInTx creates the ledger record for a newly opened account. It +// must be called within the DB transaction that creates the account info +// record. A new account has no history, so its record is created backfilled +// at zero and predicates are enforced from the start. +// +// It is a no-op while ledger writes are disabled, and for accounts that +// aren't timelock accounts, which the ledger doesn't track. +func CreateRecordInTx(ctx context.Context, data ocp_data.Provider, accountInfoRecord *account.Record) error { + if !enableLedgerWrites.Get(ctx) || !accountInfoRecord.IsTimelock() { + return nil + } + + err := data.CreateBalance(ctx, &balance.Record{ + TokenAccount: accountInfoRecord.TokenAccount, + OwnerAccount: accountInfoRecord.OwnerAccount, + MintAccount: accountInfoRecord.MintAccount, + IsOpen: true, + IsBackfilled: true, + }) + if errors.Is(err, balance.ErrRecordExists) { + return nil + } + return err +} + +// resolveRecords reports which accounts in the delta set the ledger tracks, +// creating a non-backfilled record for every timelock account that doesn't +// have one yet. +func resolveRecords(ctx context.Context, data ocp_data.Provider, deltas []*balance.Delta) (map[string]bool, error) { + tracked := make(map[string]bool) + var tokenAccounts []string + for _, delta := range deltas { + if _, ok := tracked[delta.TokenAccount]; ok { + continue + } + tracked[delta.TokenAccount] = false + tokenAccounts = append(tokenAccounts, delta.TokenAccount) + } + + existing, err := data.GetBalanceBatch(ctx, tokenAccounts...) + if err != nil { + return nil, err + } + + for _, tokenAccount := range tokenAccounts { + if _, ok := existing[tokenAccount]; ok { + tracked[tokenAccount] = true + continue + } + + accountInfoRecord, err := data.GetAccountInfoByTokenAddress(ctx, tokenAccount) + if errors.Is(err, account.ErrAccountInfoNotFound) { + continue + } else if err != nil { + return nil, err + } + if !accountInfoRecord.IsTimelock() { + continue + } + + err = data.CreateBalance(ctx, &balance.Record{ + TokenAccount: accountInfoRecord.TokenAccount, + OwnerAccount: accountInfoRecord.OwnerAccount, + MintAccount: accountInfoRecord.MintAccount, + IsOpen: true, + IsBackfilled: false, + }) + if err != nil && !errors.Is(err, balance.ErrRecordExists) { + return nil, err + } + tracked[tokenAccount] = true + } + return tracked, nil +} diff --git a/ocp/balance/ledger_test.go b/ocp/balance/ledger_test.go new file mode 100644 index 0000000..34d4c37 --- /dev/null +++ b/ocp/balance/ledger_test.go @@ -0,0 +1,194 @@ +package balance + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" + + "github.com/code-payments/ocp-server/config/memory" + "github.com/code-payments/ocp-server/config/wrapper" + ocp_data "github.com/code-payments/ocp-server/ocp/data" + "github.com/code-payments/ocp-server/ocp/data/account" + "github.com/code-payments/ocp-server/ocp/data/balance" + "github.com/code-payments/ocp-server/testutil" +) + +func TestApplyDeltasInTx_WritesDisabled(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + + source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) + + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{ + TokenAccount: source, + Kind: balance.DeltaDebit, + Quarks: 100, + })) + + _, err := data.GetBalance(ctx, source) + assert.Equal(t, balance.ErrRecordNotFound, err) +} + +func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + enableLedgerWritesForTest(t) + + source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) + destination := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_REMOTE_SEND_GIFT_CARD) + swap := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_SWAP) + external := testutil.NewRandomAccount(t).PublicKey().ToBase58() + + require.NoError(t, ApplyDeltasInTx(ctx, data, + &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 100, UsdCostBasis: 1_000_000}, + &balance.Delta{TokenAccount: destination, Kind: balance.DeltaCredit, Quarks: 60, UsdCostBasis: 600_000}, + &balance.Delta{TokenAccount: swap, Kind: balance.DeltaCredit, Quarks: 20, UsdCostBasis: 200_000}, + &balance.Delta{TokenAccount: external, Kind: balance.DeltaCredit, Quarks: 20, UsdCostBasis: 200_000}, + )) + + // Existing accounts get a non-backfilled row that accumulates freely, + // including a negative balance for a source that predates the ledger + record, err := data.GetBalance(ctx, source) + require.NoError(t, err) + assert.EqualValues(t, -100, record.Quarks) + assert.EqualValues(t, -1_000_000, record.UsdCostBasis) + assert.False(t, record.IsBackfilled) + assert.True(t, record.IsOpen) + + record, err = data.GetBalance(ctx, destination) + require.NoError(t, err) + assert.EqualValues(t, 60, record.Quarks) + assert.EqualValues(t, 600_000, record.UsdCostBasis) + assert.False(t, record.IsBackfilled) + + // Credits to accounts OCP doesn't hold a timelock for are dropped, and + // those accounts never get a row + _, err = data.GetBalance(ctx, swap) + assert.Equal(t, balance.ErrRecordNotFound, err) + _, err = data.GetBalance(ctx, external) + assert.Equal(t, balance.ErrRecordNotFound, err) + + // Once backfilled, predicates are enforced + require.NoError(t, data.BackfillBalance(ctx, source, func(context.Context) (*balance.BackfillResult, error) { + return &balance.BackfillResult{Quarks: 400, UsdCostBasis: 4_000_000, IsOpen: true}, nil + })) + err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 401}) + assert.Equal(t, balance.ErrInsufficientBalance, err) + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 400, UsdCostBasis: 4_000_000})) + + record, err = data.GetBalance(ctx, source) + require.NoError(t, err) + assert.EqualValues(t, 0, record.Quarks) + assert.True(t, record.IsBackfilled) +} + +func TestApplyDeltasInTx_UnknownSource(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + enableLedgerWritesForTest(t) + + external := testutil.NewRandomAccount(t).PublicKey().ToBase58() + swap := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_SWAP) + for _, source := range []string{external, swap} { + err := ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 1}) + assert.ErrorIs(t, err, ErrUntrackedAccount) + } +} + +func TestApplyDeltasInTx_OnlyUntrackedCredits(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + enableLedgerWritesForTest(t) + + external := testutil.NewRandomAccount(t).PublicKey().ToBase58() + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: external, Kind: balance.DeltaCredit, Quarks: 1})) + _, err := data.GetBalance(ctx, external) + assert.Equal(t, balance.ErrRecordNotFound, err) +} + +func TestApplyDeltasInTx_InvalidDelta(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + enableLedgerWritesForTest(t) + + source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) + assert.Error(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit})) + + _, err := data.GetBalance(ctx, source) + assert.Equal(t, balance.ErrRecordNotFound, err) +} + +func TestCreateRecordInTx(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + + // 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) + + enableLedgerWritesForTest(t) + + // A new timelock account starts backfilled at zero, so predicates are + // enforced immediately + require.NoError(t, CreateRecordInTx(ctx, data, primary)) + record, err := data.GetBalance(ctx, primary.TokenAccount) + require.NoError(t, err) + assert.Equal(t, primary.TokenAccount, record.TokenAccount) + assert.Equal(t, primary.OwnerAccount, record.OwnerAccount) + assert.Equal(t, primary.MintAccount, record.MintAccount) + assert.EqualValues(t, 0, record.Quarks) + assert.EqualValues(t, 0, record.UsdCostBasis) + assert.True(t, record.IsOpen) + assert.True(t, record.IsBackfilled) + + err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: primary.TokenAccount, Kind: balance.DeltaDebit, Quarks: 1}) + assert.Equal(t, balance.ErrInsufficientBalance, err) + + // Re-creating is idempotent and doesn't reset the record + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: primary.TokenAccount, Kind: balance.DeltaCredit, Quarks: 10})) + require.NoError(t, CreateRecordInTx(ctx, data, primary)) + record, err = data.GetBalance(ctx, primary.TokenAccount) + require.NoError(t, err) + assert.EqualValues(t, 10, record.Quarks) + + // Non-timelock accounts are never tracked + swap := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_SWAP) + require.NoError(t, CreateRecordInTx(ctx, data, swap)) + _, err = data.GetBalance(ctx, swap.TokenAccount) + assert.Equal(t, balance.ErrRecordNotFound, err) +} + +func newLedgerTestAccount(t *testing.T, ctx context.Context, data ocp_data.Provider, accountType commonpb.AccountType) string { + return newLedgerTestAccountInfo(t, ctx, data, accountType).TokenAccount +} + +func newLedgerTestAccountInfo(t *testing.T, ctx context.Context, data ocp_data.Provider, accountType commonpb.AccountType) *account.Record { + owner := testutil.NewRandomAccount(t) + authority := owner + if accountType == commonpb.AccountType_SWAP { + authority = testutil.NewRandomAccount(t) + } + record := &account.Record{ + OwnerAccount: owner.PublicKey().ToBase58(), + AuthorityAccount: authority.PublicKey().ToBase58(), + TokenAccount: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + MintAccount: testutil.NewRandomAccount(t).PublicKey().ToBase58(), + AccountType: accountType, + } + require.NoError(t, data.CreateAccountInfo(ctx, record)) + return record +} + +func enableLedgerWritesForTest(t *testing.T) { + previous := enableLedgerWrites + enableLedgerWrites = wrapper.NewBoolConfig(memory.NewConfig(true), defaultEnableLedgerWrites) + t.Cleanup(func() { + enableLedgerWrites = previous + }) +} diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index b50b0d4..911b8fd 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -147,9 +147,7 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { } } - sorted := make([]*balance.Delta, len(deltas)) - copy(sorted, deltas) - balance.SortDeltas(sorted) + merged := balance.MergeDeltas(deltas) s.mu.Lock() defer s.mu.Unlock() @@ -157,15 +155,12 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { // Apply to copies first so a failure part way through leaves the store // untouched, matching the transactional behaviour of the DB store. updated := make(map[string]*balance.Record) - for _, delta := range sorted { + for _, delta := range merged { item, ok := updated[delta.TokenAccount] if !ok { original, ok := s.balanceRecordsByTokenAccount[delta.TokenAccount] if !ok { - if delta.Kind == balance.DeltaCredit { - continue // Credits to accounts we don't track, like external wallets, are expected - } - return balance.ErrRecordNotFound // Everything else only ever targets accounts we track + return balance.ErrRecordNotFound } cloned := original.Clone() item = &cloned diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index 7a9b7c0..b5e2c93 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -228,10 +228,7 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er var current model err = tx.GetContext(ctx, ¤t, `SELECT `+allColumns+` FROM `+tableName+` WHERE token_account = $1`, delta.TokenAccount) if pgutil.IsNoRows(err) { - if delta.Kind == balance.DeltaCredit { - continue // Credits to accounts we don't track, like external wallets, are expected - } - return balance.ErrRecordNotFound // Everything else only ever targets accounts we track + return balance.ErrRecordNotFound } else if err != nil { return err } diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index def96ff..8b8edcc 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -94,11 +94,7 @@ func (s *store) ApplyDeltas(ctx context.Context, deltas ...*balance.Delta) error } } - sorted := make([]*balance.Delta, len(deltas)) - copy(sorted, deltas) - balance.SortDeltas(sorted) - - return dbApplyDeltas(ctx, s.db, sorted) + return dbApplyDeltas(ctx, s.db, balance.MergeDeltas(deltas)) } // Backfill implements balance.Store.Backfill diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index 1a32f4b..5d6e8e8 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -169,6 +169,33 @@ func SortDeltas(deltas []*Delta) { }) } +// MergeDeltas returns a copy of deltas in SortDeltas order with consecutive +// credits and debits to the same account combined into one. Applying one +// combined delta is equivalent to applying the parts in sequence, since +// both kinds are additive and their predicates are monotonic in the amount, +// but it touches the row once. Drains and closes are never merged, since an +// account can only legitimately be drained or closed once. +func MergeDeltas(deltas []*Delta) []*Delta { + sorted := make([]*Delta, len(deltas)) + copy(sorted, deltas) + SortDeltas(sorted) + + merged := make([]*Delta, 0, len(sorted)) + for _, delta := range sorted { + if len(merged) > 0 { + last := merged[len(merged)-1] + if last.TokenAccount == delta.TokenAccount && last.Kind == delta.Kind && (delta.Kind == DeltaCredit || delta.Kind == DeltaDebit) { + last.Quarks += delta.Quarks + last.UsdCostBasis += delta.UsdCostBasis + continue + } + } + cloned := *delta + merged = append(merged, &cloned) + } + return merged +} + func (k DeltaKind) String() string { switch k { case DeltaCredit: diff --git a/ocp/data/balance/record_test.go b/ocp/data/balance/record_test.go new file mode 100644 index 0000000..87dc868 --- /dev/null +++ b/ocp/data/balance/record_test.go @@ -0,0 +1,44 @@ +package balance + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestMergeDeltas(t *testing.T) { + input := []*Delta{ + {TokenAccount: "b", Kind: DeltaDebit, Quarks: 5, UsdCostBasis: 1}, + {TokenAccount: "a", Kind: DeltaCredit, Quarks: 1, UsdCostBasis: 10}, + {TokenAccount: "b", Kind: DeltaCredit, Quarks: 2, UsdCostBasis: 20}, + {TokenAccount: "b", Kind: DeltaDebit, Quarks: 7, UsdCostBasis: -3}, + {TokenAccount: "a", Kind: DeltaCredit, Quarks: 3, UsdCostBasis: 30}, + {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, + {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, + {TokenAccount: "d", Kind: DeltaClose}, + {TokenAccount: "d", Kind: DeltaClose}, + } + original := make([]Delta, len(input)) + for i, delta := range input { + original[i] = *delta + } + + merged := MergeDeltas(input) + + assert.Equal(t, []*Delta{ + {TokenAccount: "a", Kind: DeltaCredit, Quarks: 4, UsdCostBasis: 40}, + {TokenAccount: "b", Kind: DeltaCredit, Quarks: 2, UsdCostBasis: 20}, + {TokenAccount: "b", Kind: DeltaDebit, Quarks: 12, UsdCostBasis: -2}, + {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, + {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, + {TokenAccount: "d", Kind: DeltaClose}, + {TokenAccount: "d", Kind: DeltaClose}, + }, merged) + + // The input is left untouched + for i, delta := range input { + assert.Equal(t, original[i], *delta) + } + + assert.Empty(t, MergeDeltas(nil)) +} diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index aae4c3e..8d212cf 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -86,10 +86,10 @@ type Store interface { // Predicates are enforced only on backfilled records; records that are // not backfilled simply accumulate the change. // - // Only accounts managed by OCP have records. A credit to an account - // without one is skipped, since external destinations are routinely paid. - // Any other kind targeting an account without a record is - // ErrRecordNotFound, since funds only ever leave accounts managed by OCP. + // Every delta must target an account with a record, otherwise + // ErrRecordNotFound is returned and nothing is applied. Callers are + // responsible for not producing deltas for accounts the ledger doesn't + // track, like external wallets. // // ErrInsufficientBalance is returned when a debit exceeds the balance. // ErrBalanceChanged is returned when a drain or close doesn't match the diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 2b9a6d5..4944bd3 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -192,24 +192,19 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { IsBackfilled: true, })) - // Credits to accounts without a record are skipped, but nothing else is - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1})) + // Every kind of delta requires a record + assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1})) assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDebit, Quarks: 1})) assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaDrain, Quarks: 1})) assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaClose})) - // A batch mixing a tracked account with an untracked credit, like a - // withdrawal to an external wallet, applies the tracked side - require.NoError(t, s.ApplyDeltas( + // A batch mixing a tracked account with an untracked credit is rejected + // as a whole, so the tracked side is untouched + assert.Equal(t, balance.ErrRecordNotFound, s.ApplyDeltas( ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1}, &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1}, )) - require.NoError(t, s.ApplyDeltas( - ctx, - &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1}, - &balance.Delta{TokenAccount: "untracked", Kind: balance.DeltaCredit, Quarks: 1}, - )) assertBalance(t, s, "token_account_1", 0, 0, true) // Invalid deltas are rejected @@ -334,6 +329,25 @@ func testApplyDeltasAtomicity(t *testing.T, s balance.Store) { &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 50}, )) assertBalance(t, s, "token_account_1", 10, -4, true) + + // Same-kind deltas to the same account are checked as one, so a pair of + // debits that together exceed the balance fails even though each fits + assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 6, UsdCostBasis: 1}, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 6, UsdCostBasis: 1}, + )) + assertBalance(t, s, "token_account_1", 10, -4, true) + + require.NoError(t, s.ApplyDeltas( + ctx, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 4, UsdCostBasis: 1}, + &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 6, UsdCostBasis: 2}, + &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 4, UsdCostBasis: 1}, + &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 6, UsdCostBasis: 2}, + )) + assertBalance(t, s, "token_account_1", 0, -7, true) + assertBalance(t, s, "token_account_2", 150, 7, true) }) } From ff19335ab46844f237adce36f3bf08956e668e9b Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Thu, 27 Aug 2026 14:38:39 -0400 Subject: [PATCH 06/26] Integrate new balance table into all call sites --- ocp/balance/ledger.go | 7 ++++ ocp/rpc/transaction/action_handler.go | 8 +++- ocp/rpc/transaction/intent.go | 22 ++++++++++ ocp/rpc/transaction/intent_handler.go | 3 ++ ocp/worker/account/gift_card.go | 23 ++++++++-- ocp/worker/currency/launcher/util.go | 5 +++ ocp/worker/geyser/external_deposit.go | 12 ++++++ ocp/worker/swap/util.go | 60 ++++++++++++++++++++++++--- 8 files changed, 129 insertions(+), 11 deletions(-) diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index 631f197..276bed8 100644 --- a/ocp/balance/ledger.go +++ b/ocp/balance/ledger.go @@ -14,6 +14,13 @@ import ( // ledger doesn't track. var ErrUntrackedAccount = errors.New("account is not tracked by the balance ledger") +// 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. +func LedgerWritesEnabled(ctx context.Context) bool { + return enableLedgerWrites.Get(ctx) +} + // ApplyDeltasInTx applies balance deltas to the ledger. It must be called // within the DB transaction that commits the records the deltas are derived // from, so the ledger can never disagree with them. diff --git a/ocp/rpc/transaction/action_handler.go b/ocp/rpc/transaction/action_handler.go index eecfad3..e40aacb 100644 --- a/ocp/rpc/transaction/action_handler.go +++ b/ocp/rpc/transaction/action_handler.go @@ -9,6 +9,7 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" transactionpb "github.com/code-payments/ocp-protobuf-api/generated/go/transaction/v1" + "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" @@ -190,7 +191,12 @@ func (h *OpenAccountActionHandler) OnCommitToDB(ctx context.Context) error { return err } - return h.data.CreateAccountInfo(ctx, h.unsavedAccountInfoRecord) + err = h.data.CreateAccountInfo(ctx, h.unsavedAccountInfoRecord) + if err != nil { + return err + } + + return balance.CreateRecordInTx(ctx, h.data, h.unsavedAccountInfoRecord) } type NoPrivacyTransferActionHandler struct { diff --git a/ocp/rpc/transaction/intent.go b/ocp/rpc/transaction/intent.go index b11fbf6..49e1d33 100644 --- a/ocp/rpc/transaction/intent.go +++ b/ocp/rpc/transaction/intent.go @@ -22,9 +22,11 @@ import ( transactionpb "github.com/code-payments/ocp-protobuf-api/generated/go/transaction/v1" "github.com/code-payments/ocp-server/grpc/client" + "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" + balance_store "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/data/fulfillment" "github.com/code-payments/ocp-server/ocp/data/intent" "github.com/code-payments/ocp-server/ocp/data/nonce" @@ -723,6 +725,22 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm } } + // Reflect the intent in the balance ledger. Applied last, so the + // ledger row locks are held for as little of the transaction as + // possible. + if balance.LedgerWritesEnabled(ctx) { + balanceDeltas, err := balance.DeltasForSubmittedIntent(intentRecord, actionRecords, s.conf.createOnSendWithdrawalFeeQuarks.Get(ctx)) + if err != nil { + log.With(zap.Error(err)).Warn("failure building balance deltas") + return err + } + err = balance.ApplyDeltasInTx(ctx, s.data, balanceDeltas...) + if err != nil { + log.With(zap.Error(err)).Warn("failure applying balance deltas") + return err + } + } + // Schedule app-defined tasks atomically with the intent, so their // execution is guaranteed once the intent is committed tasksToSchedule, err = s.submitIntentIntegration.GetTasksToSchedule(ctx, intentRecord) @@ -739,6 +757,10 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm return nil }) if err != nil { + if errors.Is(err, balance_store.ErrInsufficientBalance) || errors.Is(err, balance_store.ErrBalanceChanged) || errors.Is(err, balance_store.ErrAccountClosed) { + log.With(zap.Error(err)).Info("balance ledger rejected intent") + return handleSubmitIntentError(ctx, streamer, intentRecord, NewStaleStateErrorf("race detected: %s", err.Error())) + } if strings.Contains(err.Error(), "stale") || strings.Contains(err.Error(), "exist") { log.With(zap.Error(err)).Info("race condition detected") return handleSubmitIntentError(ctx, streamer, intentRecord, NewStaleStateErrorf("race detected: %s", err.Error())) diff --git a/ocp/rpc/transaction/intent_handler.go b/ocp/rpc/transaction/intent_handler.go index 88dceb5..9cbce69 100644 --- a/ocp/rpc/transaction/intent_handler.go +++ b/ocp/rpc/transaction/intent_handler.go @@ -1979,6 +1979,9 @@ func saveAutoOpenPrimaryAccountIntent(ctx context.Context, data ocp_data.Provide if err := data.CreateAccountInfo(ctx, req.accountInfo); err != nil { return err } + if err := balance.CreateRecordInTx(ctx, data, req.accountInfo); err != nil { + return err + } openFulfillmentRecord := &fulfillment.Record{ Intent: openIntentRecord.IntentId, diff --git a/ocp/worker/account/gift_card.go b/ocp/worker/account/gift_card.go index 41035ca..e2cbb6d 100644 --- a/ocp/worker/account/gift_card.go +++ b/ocp/worker/account/gift_card.go @@ -170,7 +170,7 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Prov } // Add a intent record to show the funds being returned back to the issuer - err = insertAutoReturnIntentRecord(ctx, data, giftCardIssuedIntent, isVoidedByUser) + autoReturnIntent, err := insertAutoReturnIntentRecord(ctx, data, giftCardIssuedIntent, isVoidedByUser) if err != nil { return err } @@ -219,6 +219,17 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Prov return err } + if balance.LedgerWritesEnabled(ctx) { + balanceDeltas, err := balance.DeltasForGiftCardAutoReturn(autoReturnIntent, autoReturnAction) + if err != nil { + return err + } + err = balance.ApplyDeltasInTx(ctx, data, balanceDeltas...) + if err != nil { + return err + } + } + // This will trigger the fulfillment worker to poll for the fulfillment. This // should be the very last DB update called. err = markFulfillmentAsActivelyScheduled(ctx, data, autoReturnFulfillment[0]) @@ -314,10 +325,10 @@ func updateAutoReturnFulfillmentPreSorting( return data.UpdateFulfillment(ctx, fulfillmentRecord) } -func insertAutoReturnIntentRecord(ctx context.Context, data ocp_data.Provider, giftCardIssuedIntent *intent.Record, isVoidedByUser bool) error { +func insertAutoReturnIntentRecord(ctx context.Context, data ocp_data.Provider, giftCardIssuedIntent *intent.Record, isVoidedByUser bool) (*intent.Record, error) { mintAccount, err := common.NewAccountFromPublicKeyString(giftCardIssuedIntent.MintAccount) if err != nil { - return err + return nil, err } // We need to insert a faked completed public receive intent so it can appear @@ -350,7 +361,11 @@ func insertAutoReturnIntentRecord(ctx context.Context, data ocp_data.Provider, g CreatedAt: time.Now(), } - return data.SaveIntent(ctx, intentRecord) + err = data.SaveIntent(ctx, intentRecord) + if err != nil { + return nil, err + } + return intentRecord, nil } func markActionAsRevoked(ctx context.Context, data ocp_data.Provider, actionRecord *action.Record) error { diff --git a/ocp/worker/currency/launcher/util.go b/ocp/worker/currency/launcher/util.go index 401602d..63dbe92 100644 --- a/ocp/worker/currency/launcher/util.go +++ b/ocp/worker/currency/launcher/util.go @@ -12,6 +12,7 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" + "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" @@ -1176,6 +1177,10 @@ func (p *runtime) initializeCreatorAcccount(ctx context.Context, currencyMetadat if err != nil { return errors.Wrap(err, "error saving creator account info record") } + err = balance.CreateRecordInTx(ctx, p.data, accountInfoRecord) + if err != nil { + return errors.Wrap(err, "error creating creator balance record") + } // Create the fulfillment record for initializing the timelock account fulfillmentRecord := &fulfillment.Record{ diff --git a/ocp/worker/geyser/external_deposit.go b/ocp/worker/geyser/external_deposit.go index 803ec2a..9e2356d 100644 --- a/ocp/worker/geyser/external_deposit.go +++ b/ocp/worker/geyser/external_deposit.go @@ -16,6 +16,7 @@ import ( "github.com/code-payments/ocp-server/cache" currency_lib "github.com/code-payments/ocp-server/currency" "github.com/code-payments/ocp-server/database/query" + balance_util "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" currency_util "github.com/code-payments/ocp-server/ocp/currency" ocp_data "github.com/code-payments/ocp-server/ocp/data" @@ -405,6 +406,17 @@ func processPotentialExternalDepositIntoVm(ctx context.Context, data ocp_data.Pr return errors.Wrap(err, "error saving external deposit record") } + if balance_util.LedgerWritesEnabled(ctx) { + balanceDeltas, err := balance_util.DeltasForExternalDeposit(intentRecord) + if err != nil { + return errors.Wrap(err, "error building balance deltas") + } + err = balance_util.ApplyDeltasInTx(ctx, data, balanceDeltas...) + if err != nil { + return errors.Wrap(err, "error applying balance deltas") + } + } + return nil }) if err != nil { diff --git a/ocp/worker/swap/util.go b/ocp/worker/swap/util.go index e07946a..85e764a 100644 --- a/ocp/worker/swap/util.go +++ b/ocp/worker/swap/util.go @@ -12,6 +12,7 @@ import ( "github.com/pkg/errors" currency_lib "github.com/code-payments/ocp-server/currency" + "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" currency_util "github.com/code-payments/ocp-server/ocp/currency" "github.com/code-payments/ocp-server/ocp/data/currency" @@ -464,6 +465,17 @@ func (p *runtime) markSwapCancelled(ctx context.Context, swapRecord *swap.Record return err } + if balance.LedgerWritesEnabled(ctx) { + balanceDeltas, err := balance.DeltasForExternalDeposit(refundIntentRecord) + if err != nil { + return err + } + err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) + if err != nil { + return err + } + } + // The swap was funded and entered transaction history, so its // record must reflect the refund err = history_util.MarkSwapAsFailed(ctx, p.data, swapRecord.SwapId) @@ -703,6 +715,7 @@ func (p *runtime) maybeUpdateBalancesForFinalizedReserveSwap(ctx context.Context var exchangeCurrency currency_lib.Code var nativeAmountWithoutFees float64 var usdMarketValueWithoutFees float64 + var previousFundingIntentRecord, reconciledFundingIntentRecord *intent.Record switch swapRecord.FundingSource { case swap.FundingSourceSubmitIntent: fundingIntentRecord, err := p.data.GetIntent(ctx, swapRecord.FundingId) @@ -760,12 +773,11 @@ func (p *runtime) maybeUpdateBalancesForFinalizedReserveSwap(ctx context.Context } // Reconcile the source funding payment's cost basis to the core mint actually - // realized by the sell. + // realized by the sell. Saved in the settlement transaction below. + cloned := fundingIntentRecord.Clone() + previousFundingIntentRecord = &cloned fundingIntentRecord.SendPublicPaymentMetadata.UsdMarketValue = usdMarketValueWithoutFees - err = p.data.SaveIntent(ctx, fundingIntentRecord) - if err != nil { - return 0, false, err - } + reconciledFundingIntentRecord = fundingIntentRecord } case swap.FundingSourceExternalWallet, swap.FundingSourceCoinbaseOnramp: if !common.IsCoreMint(fromMint) { @@ -830,7 +842,43 @@ func (p *runtime) maybeUpdateBalancesForFinalizedReserveSwap(ctx context.Context CreatedAt: time.Now(), } - return p.data.SaveExternalDeposit(ctx, externalDepositRecord) + err = p.data.SaveExternalDeposit(ctx, externalDepositRecord) + if err != nil { + return err + } + + if reconciledFundingIntentRecord != nil { + err = p.data.SaveIntent(ctx, reconciledFundingIntentRecord) + if err != nil { + return err + } + } + + if balance.LedgerWritesEnabled(ctx) { + balanceDeltas, err := balance.DeltasForExternalDeposit(intentRecord) + if err != nil { + return err + } + + if reconciledFundingIntentRecord != nil { + fundingActionRecords, err := p.data.GetAllActionsByIntent(ctx, reconciledFundingIntentRecord.IntentId) + if err != nil { + return err + } + reconciliationDeltas, err := balance.DeltasForSwapSellReconciliation(previousFundingIntentRecord, reconciledFundingIntentRecord, fundingActionRecords) + if err != nil { + return err + } + balanceDeltas = append(balanceDeltas, reconciliationDeltas...) + } + + err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) + if err != nil { + return err + } + } + + return nil }) if err != nil { return 0, false, err From 6144714b35332f9e02b1631e21d2eaba143c623e Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Thu, 27 Aug 2026 15:35:35 -0400 Subject: [PATCH 07/26] Bump Postgres test docker image to 14.24 --- database/postgres/test/util.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/database/postgres/test/util.go b/database/postgres/test/util.go index 6257c89..e3cb875 100644 --- a/database/postgres/test/util.go +++ b/database/postgres/test/util.go @@ -18,7 +18,7 @@ import ( const ( containerName = "postgres" - containerVersion = "10.4" + containerVersion = "14.24" containerAutoKill = 120 // seconds port = 5432 From 83773c8fe74fdef07332640904afe1d022c459c2 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 09:38:36 -0400 Subject: [PATCH 08/26] Update mint index --- ocp/data/balance/postgres/store_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ocp/data/balance/postgres/store_test.go b/ocp/data/balance/postgres/store_test.go index 77c0146..e322a0f 100644 --- a/ocp/data/balance/postgres/store_test.go +++ b/ocp/data/balance/postgres/store_test.go @@ -44,7 +44,7 @@ const ( ) WITH (fillfactor = 90); CREATE INDEX ocp__core_balance__idx__owner_account__mint_account ON ocp__core_balance (owner_account, mint_account); - CREATE INDEX ocp__core_balance__idx__mint_account ON ocp__core_balance (mint_account); + CREATE INDEX ocp__core_balance__idx__mint_account__id ON ocp__core_balance (mint_account, id); CREATE TABLE ocp__core_cachedbalanceversion ( id SERIAL NOT NULL PRIMARY KEY, From a59a3053b625347d170ad9161870ac52f996c923 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 12:02:21 -0400 Subject: [PATCH 09/26] Update holder count worker to use new balance table --- ocp/data/balance/memory/store.go | 14 ++++++++++++++ ocp/data/balance/postgres/model.go | 11 +++++++++++ ocp/data/balance/postgres/store.go | 5 +++++ ocp/data/balance/store.go | 5 +++++ ocp/data/balance/tests/tests.go | 27 +++++++++++++++++++++++++++ ocp/data/internal.go | 4 ++++ ocp/worker/currency/holder/runtime.go | 4 ++++ 7 files changed, 70 insertions(+) diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 911b8fd..9f62878 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -139,6 +139,20 @@ func (s *store) GetAllByMint(_ context.Context, mint string, minQuarks int64, cu return res, nil } +// CountByMint implements balance.Store.CountByMint +func (s *store) CountByMint(_ 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.IsBackfilled { + res++ + } + } + return res, nil +} + // ApplyDeltas implements balance.Store.ApplyDeltas func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { for _, delta := range deltas { diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index b5e2c93..dcaadcd 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -176,6 +176,17 @@ func dbGetAllByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int return res, nil } +func dbCountByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64) (uint64, error) { + var res uint64 + query := `SELECT COUNT(*) FROM ` + tableName + ` + WHERE mint_account = $1 AND quarks >= $2 AND is_backfilled` + err := db.GetContext(ctx, &res, query, mint, minQuarks) + if err != nil { + return 0, err + } + return res, nil +} + // dbApplyDeltas applies every delta in a single transaction. Each delta is one // conditional UPDATE, so its predicate is evaluated against the row after the // row lock is acquired. Predicates only apply to backfilled rows. diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index 8b8edcc..599b528 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -86,6 +86,11 @@ func (s *store) GetAllByMint(ctx context.Context, mint string, minQuarks int64, return fromModels(models), nil } +// CountByMint implements balance.Store.CountByMint +func (s *store) CountByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) { + return dbCountByMint(ctx, s.db, mint, minQuarks) +} + // ApplyDeltas implements balance.Store.ApplyDeltas func (s *store) ApplyDeltas(ctx context.Context, deltas ...*balance.Delta) error { for _, delta := range deltas { diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index 8d212cf..e5580e4 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -80,6 +80,11 @@ type Store interface { // ErrRecordNotFound is returned if no records exist. GetAllByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) + // CountByMint counts backfilled records for a mint with at least + // minQuarks. Records that are not backfilled are excluded, since their + // balances are partial sums that can't be compared against a threshold. + CountByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) + // ApplyDeltas atomically applies a set of deltas. Either every delta is // applied or none are. Deltas are applied in SortDeltas order. // diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 4944bd3..b9125c4 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -177,6 +177,33 @@ func testGetAllByMint(t *testing.T, s balance.Store) { _, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(5), 10, query.Ascending) assert.Equal(t, balance.ErrRecordNotFound, err) + + // Counting by mint uses the same threshold semantics, but only over + // backfilled records + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_not_backfilled", + OwnerAccount: "owner", + MintAccount: "mint_1", + Quarks: 1000, + IsOpen: true, + IsBackfilled: false, + })) + + count, err := s.CountByMint(ctx, "mint_1", 0) + require.NoError(t, err) + assert.EqualValues(t, 5, count) + + count, err = s.CountByMint(ctx, "mint_1", 20) + require.NoError(t, err) + assert.EqualValues(t, 3, count) + + count, err = s.CountByMint(ctx, "mint_2", 0) + require.NoError(t, err) + assert.EqualValues(t, 1, count) + + count, err = s.CountByMint(ctx, "mint_3", 0) + require.NoError(t, err) + assert.EqualValues(t, 0, count) }) } diff --git a/ocp/data/internal.go b/ocp/data/internal.go index a3111f3..09c1f8e 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -127,6 +127,7 @@ type DatabaseData interface { GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) + CountBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) @@ -486,6 +487,9 @@ func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, ow func (dp *DatabaseProvider) GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { return dp.balance.GetAllByMint(ctx, mint, minQuarks, cursor, limit, direction) } +func (dp *DatabaseProvider) CountBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) { + return dp.balance.CountByMint(ctx, mint, minQuarks) +} func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error { return dp.balance.ApplyDeltas(ctx, deltas...) } diff --git a/ocp/worker/currency/holder/runtime.go b/ocp/worker/currency/holder/runtime.go index bf6f84b..6ead637 100644 --- a/ocp/worker/currency/holder/runtime.go +++ b/ocp/worker/currency/holder/runtime.go @@ -115,6 +115,10 @@ func (p *holderRuntime) countHoldersForMint(ctx context.Context, mint string, cu return 0, nil } + if balance.LedgerReadsEnabled(ctx) { + return p.data.CountBalancesByMint(ctx, mint, int64(minHoldings)) + } + accountRecords, err := p.data.GetAccountInfosByMintAndType(ctx, mint, commonpb.AccountType_PRIMARY) if err == account.ErrAccountInfoNotFound { return 0, nil From f4ad70c96a31c01d07338d6afa2bcff46ab5815d Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 12:05:25 -0400 Subject: [PATCH 10/26] GetTokenAccountInfos now fully utilizes new balance table --- ocp/balance/calculator.go | 87 ++++++++++++++++++++++++++++++++++ ocp/balance/calculator_test.go | 66 ++++++++++++++++++++++++++ ocp/balance/ledger.go | 6 +++ ocp/rpc/account/server.go | 47 +++++++++--------- 4 files changed, 182 insertions(+), 24 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 9a10f42..6da1442 100644 --- a/ocp/balance/calculator.go +++ b/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. // diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index c662202..345a364 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -418,6 +418,72 @@ 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, + 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_BalanceRecordReadsDisabled(t *testing.T) { env := setupBalanceTestEnv(t) diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index 276bed8..9c3143f 100644 --- a/ocp/balance/ledger.go +++ b/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. diff --git a/ocp/rpc/account/server.go b/ocp/rpc/account/server.go index 8ea416d..b5ed848 100644 --- a/ocp/rpc/account/server.go +++ b/ocp/rpc/account/server.go @@ -25,6 +25,7 @@ import ( ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" + balance_store "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/rpc" account_worker "github.com/code-payments/ocp-server/ocp/worker/account" timelock_token_v1 "github.com/code-payments/ocp-server/solana/timelock/v1" @@ -36,8 +37,9 @@ var ( ) type balanceMetadata struct { - value uint64 - source accountpb.TokenAccountInfo_BalanceSource + quarks uint64 + usdCostBasis int64 + source accountpb.TokenAccountInfo_BalanceSource } type server struct { @@ -321,19 +323,21 @@ func (s *server) fetchBalances(ctx context.Context, allAccountRecords []*common. // Don't calculate a balance for now, since the caching strategy // is not possible. balanceMetadataByTokenAccount[accountRecords.General.TokenAccount] = &balanceMetadata{ - value: 0, - source: accountpb.TokenAccountInfo_BALANCE_SOURCE_UNKNOWN, + quarks: 0, + usdCostBasis: 0, + source: accountpb.TokenAccountInfo_BALANCE_SOURCE_UNKNOWN, } } } - balancesByTokenAccount, err := balance.BatchCalculateFromCacheWithAccountRecords(ctx, s.data, mangedByCodeRecords...) + balancesByTokenAccount, err := balance.BatchCalculateWithUsdCostBasisFromCache(ctx, s.data, mangedByCodeRecords...) if err != nil { return nil, err } - for tokenAccount, quarks := range balancesByTokenAccount { + for tokenAccount, cached := range balancesByTokenAccount { balanceMetadataByTokenAccount[tokenAccount] = &balanceMetadata{ - value: quarks, - source: accountpb.TokenAccountInfo_BALANCE_SOURCE_CACHE, + quarks: cached.Quarks, + usdCostBasis: cached.UsdCostBasis, + source: accountpb.TokenAccountInfo_BALANCE_SOURCE_CACHE, } } @@ -372,8 +376,9 @@ func (s *server) fetchBalances(ctx context.Context, allAccountRecords []*common. protoBalanceSource = accountpb.TokenAccountInfo_BALANCE_SOURCE_UNKNOWN } balanceMetadataByTokenAccount[tokenAccount.PublicKey().ToBase58()] = &balanceMetadata{ - value: quarks, - source: protoBalanceSource, + quarks: quarks, + usdCostBasis: 0, + source: protoBalanceSource, } } @@ -453,7 +458,7 @@ func (s *server) getProtoAccountInfo(ctx context.Context, records *common.Accoun } // Otherwise, check whether it looks like the gift card was claimed. - if prefetchedBalanceMetadata.source == accountpb.TokenAccountInfo_BALANCE_SOURCE_CACHE && prefetchedBalanceMetadata.value == 0 { + if prefetchedBalanceMetadata.source == accountpb.TokenAccountInfo_BALANCE_SOURCE_CACHE && prefetchedBalanceMetadata.quarks == 0 { claimState = accountpb.TokenAccountInfo_CLAIM_STATE_CLAIMED } else if records.Timelock.IsClosed() { claimState = accountpb.TokenAccountInfo_CLAIM_STATE_CLAIMED @@ -481,7 +486,8 @@ func (s *server) getProtoAccountInfo(ctx context.Context, records *common.Accoun // If the gift card account is claimed or expired, force the balance to zero. if claimState == accountpb.TokenAccountInfo_CLAIM_STATE_CLAIMED || claimState == accountpb.TokenAccountInfo_CLAIM_STATE_EXPIRED { prefetchedBalanceMetadata.source = accountpb.TokenAccountInfo_BALANCE_SOURCE_CACHE - prefetchedBalanceMetadata.value = 0 + prefetchedBalanceMetadata.quarks = 0 + prefetchedBalanceMetadata.usdCostBasis = 0 } } @@ -495,18 +501,11 @@ func (s *server) getProtoAccountInfo(ctx context.Context, records *common.Accoun var usdCostBasis float64 if common.IsCoreMint(mintAccount) && common.IsCoreMintUsdStableCoin() { - usdCostBasis = float64(prefetchedBalanceMetadata.value) / float64(common.CoreMintQuarksPerUnit) + usdCostBasis = float64(prefetchedBalanceMetadata.quarks) / float64(common.CoreMintQuarksPerUnit) } else { - switch records.General.AccountType { - case commonpb.AccountType_PRIMARY: - // todo: Assumes the structure that each user has exactly one primary account per mint - usdCostBasis, err = s.data.GetUsdCostBasis(ctx, ownerAccount.PublicKey().ToBase58(), mintAccount.PublicKey().ToBase58()) - if err != nil { - return nil, err - } - default: - usdCostBasis = 0 // Account type not supported - } + // Prefetched alongside the balance, from the same record where the + // ledger serves both + usdCostBasis = balance_store.UsdCostBasisToFloat(prefetchedBalanceMetadata.usdCostBasis) } var mintMetadata *currencypb.Mint var liveReserveState *currencypb.VerifiedLaunchpadCurrencyReserveState @@ -542,7 +541,7 @@ func (s *server) getProtoAccountInfo(ctx context.Context, records *common.Accoun AccountType: records.General.AccountType, Index: records.General.Index, BalanceSource: prefetchedBalanceMetadata.source, - Balance: prefetchedBalanceMetadata.value, + Balance: prefetchedBalanceMetadata.quarks, UsdCostBasis: usdCostBasis, ManagementState: managementState, BlockchainState: blockchainState, From 175d9c98891c989266c2038e5d3c69a29f758557 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 12:16:48 -0400 Subject: [PATCH 11/26] Remove balance locks --- ocp/balance/lock.go | 75 -------------- ocp/data/balance/memory/store.go | 10 +- ocp/data/balance/memory/store_legacy.go | 66 ------------ ocp/data/balance/postgres/model_legacy.go | 121 ---------------------- ocp/data/balance/postgres/store_legacy.go | 25 ----- ocp/data/balance/postgres/store_test.go | 20 ---- ocp/data/balance/store.go | 34 ------ ocp/data/balance/tests/tests.go | 40 ------- ocp/data/internal.go | 16 --- ocp/rpc/transaction/gift_card.go | 9 +- ocp/rpc/transaction/intent.go | 30 ++---- ocp/rpc/transaction/intent_handler.go | 107 +++---------------- ocp/worker/account/gift_card.go | 17 +-- 13 files changed, 34 insertions(+), 536 deletions(-) delete mode 100644 ocp/balance/lock.go delete mode 100644 ocp/data/balance/memory/store_legacy.go delete mode 100644 ocp/data/balance/postgres/model_legacy.go delete mode 100644 ocp/data/balance/postgres/store_legacy.go diff --git a/ocp/balance/lock.go b/ocp/balance/lock.go deleted file mode 100644 index c38e156..0000000 --- a/ocp/balance/lock.go +++ /dev/null @@ -1,75 +0,0 @@ -package balance - -import ( - "context" - "errors" - - "github.com/code-payments/ocp-server/ocp/common" - ocp_data "github.com/code-payments/ocp-server/ocp/data" - "github.com/code-payments/ocp-server/ocp/data/balance" -) - -// OptimisticVersionLock is an optimistic version lock on an account's cached -// balance, which can be paired with DB updates against balances that need to -// be protected against race conditions. -type OptimisticVersionLock struct { - vault *common.Account - currentVersion uint64 -} - -// GetOptimisticVersionLock gets an optimistic version lock for the vault account's -// cached balance -func GetOptimisticVersionLock(ctx context.Context, data ocp_data.Provider, vault *common.Account) (*OptimisticVersionLock, error) { - version, err := data.GetCachedBalanceVersion(ctx, vault.PublicKey().ToBase58()) - if err != nil { - return nil, err - } - return &OptimisticVersionLock{ - vault: vault, - currentVersion: version, - }, nil -} - -// OnNewBalanceVersion is called in the DB transaction updating the account's -// cached balance -func (l *OptimisticVersionLock) OnNewBalanceVersion(ctx context.Context, data ocp_data.Provider) error { - return data.AdvanceCachedBalanceVersion(ctx, l.vault.PublicKey().ToBase58(), l.currentVersion) -} - -// RequireSameBalanceVerion is called in the DB transaction requireing the -// account's cached balance not be changed -func (l *OptimisticVersionLock) RequireSameBalanceVerion(ctx context.Context, data ocp_data.Provider) error { - latestVersion, err := data.GetCachedBalanceVersion(ctx, l.vault.PublicKey().ToBase58()) - if err != nil { - return err - } - if latestVersion < l.currentVersion { - return errors.New("unexpected balance version detected") - } - if l.currentVersion != latestVersion { - return balance.ErrStaleCachedBalanceVersion - } - return nil -} - -// OpenCloseStatusLock is a lock on an account's open/close status -type OpenCloseStatusLock struct { - vault *common.Account -} - -func NewOpenCloseStatusLock(vault *common.Account) *OpenCloseStatusLock { - return &OpenCloseStatusLock{ - vault: vault, - } -} - -// OnPaymentToAccount is called in the DB transaction making a payment to the -// account that may be closed -func (l *OpenCloseStatusLock) OnPaymentToAccount(ctx context.Context, data ocp_data.Provider) error { - return data.CheckNotClosedForBalanceUpdate(ctx, l.vault.PublicKey().ToBase58()) -} - -// OnClose is called in the DB transaction closing the account -func (l *OpenCloseStatusLock) OnClose(ctx context.Context, data ocp_data.Provider) error { - return data.MarkAsClosedForBalanceUpdate(ctx, l.vault.PublicKey().ToBase58()) -} diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 9f62878..abb2e7d 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -14,9 +14,7 @@ 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 } @@ -24,9 +22,7 @@ type store struct { // 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), } } @@ -297,8 +293,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 } diff --git a/ocp/data/balance/memory/store_legacy.go b/ocp/data/balance/memory/store_legacy.go deleted file mode 100644 index 9e062f9..0000000 --- a/ocp/data/balance/memory/store_legacy.go +++ /dev/null @@ -1,66 +0,0 @@ -package memory - -import ( - "context" - - "github.com/code-payments/ocp-server/ocp/data/balance" -) - -// GetCachedVersion implements balance.Store.GetCachedVersion -func (s *store) GetCachedVersion(_ context.Context, account string) (uint64, error) { - s.mu.Lock() - defer s.mu.Unlock() - - current, ok := s.cachedBalanceVersionsByAccount[account] - if !ok { - return 0, nil - } - return current, nil -} - -// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion -func (s *store) AdvanceCachedVersion(_ context.Context, account string, currentVersion uint64) error { - s.mu.Lock() - defer s.mu.Unlock() - - actualVersion, ok := s.cachedBalanceVersionsByAccount[account] - if !ok { - if currentVersion != 0 { - return balance.ErrStaleCachedBalanceVersion - } - - s.cachedBalanceVersionsByAccount[account] = 1 - - return nil - } - - if actualVersion != currentVersion { - return balance.ErrStaleCachedBalanceVersion - } - - s.cachedBalanceVersionsByAccount[account]++ - - return nil -} - -// CheckNotClosed implements balance.Store.CheckNotClosed -func (s *store) CheckNotClosed(ctx context.Context, account string) error { - s.mu.Lock() - defer s.mu.Unlock() - - if _, ok := s.closedAccounts[account]; ok { - return balance.ErrAccountClosed - } - - return nil -} - -// MarkAsClosed implements balance.Store.MarkAsClosed -func (s *store) MarkAsClosed(ctx context.Context, account string) error { - s.mu.Lock() - defer s.mu.Unlock() - - s.closedAccounts[account] = true - - return nil -} diff --git a/ocp/data/balance/postgres/model_legacy.go b/ocp/data/balance/postgres/model_legacy.go deleted file mode 100644 index 247c9a4..0000000 --- a/ocp/data/balance/postgres/model_legacy.go +++ /dev/null @@ -1,121 +0,0 @@ -package postgres - -import ( - "context" - "database/sql" - "errors" - - "github.com/jmoiron/sqlx" - - pgutil "github.com/code-payments/ocp-server/database/postgres" - "github.com/code-payments/ocp-server/ocp/data/balance" -) - -const ( - cachedBalanceVersionTableName = "ocp__core_cachedbalanceversion" - openCloseLocksTableName = "ocp__core_opencloselocks" -) - -func dbGetCachedVersion(ctx context.Context, db *sqlx.DB, account string) (uint64, error) { - var res uint64 - err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - insertQuery := `INSERT INTO ` + cachedBalanceVersionTableName + ` - (token_account, version) - VALUES($1, 0) - ON CONFLICT DO NOTHING - ` - sqlResult, err := tx.ExecContext(ctx, insertQuery, account) - if err != nil { - return err - } - rowsAffected, err := sqlResult.RowsAffected() - if err != nil { - return err - } - if rowsAffected == 1 { - res = 0 - return nil - } - - selectQuery := `SELECT version FROM ` + cachedBalanceVersionTableName + ` - WHERE token_account = $1 - FOR UPDATE` - return db.GetContext(ctx, &res, selectQuery, account) - }) - return res, err - -} - -func dbAdvanceCachedVersion(ctx context.Context, db *sqlx.DB, account string, currentVersion uint64) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - var res uint64 - query := `UPDATE ` + cachedBalanceVersionTableName + ` - SET version = version + 1 - WHERE token_account = $1 AND version = $2 - RETURNING version - ` - err := tx.GetContext(ctx, &res, query, account, currentVersion) - if pgutil.IsNoRows(err) || pgutil.IsUniqueViolation(err) { - return balance.ErrStaleCachedBalanceVersion - } - if err != nil { - return err - } - return nil - }) - -} - -func dbCheckNotClosed(ctx context.Context, db *sqlx.DB, account string) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - insertQuery := `INSERT INTO ` + openCloseLocksTableName + ` - (token_account, is_open) - VALUES ($1, TRUE) - ON CONFLICT DO NOTHING - ` - - _, err := tx.ExecContext(ctx, insertQuery, account) - if err != nil { - return err - } - - selectQuery := `SELECT is_open FROM ` + openCloseLocksTableName + ` - WHERE token_account = $1 - FOR UPDATE - ` - var isOpen bool - err = tx.GetContext(ctx, &isOpen, selectQuery, account) - if err != nil { - return err - } - if !isOpen { - return balance.ErrAccountClosed - } - return nil - }) -} - -func dbMarkAsClosed(ctx context.Context, db *sqlx.DB, account string) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - query := `INSERT INTO ` + openCloseLocksTableName + ` - (token_account, is_open) - VALUES ($1, FALSE) - - ON CONFLICT (token_account) - DO UPDATE - SET is_open = FALSE - WHERE ` + openCloseLocksTableName + `.token_account = $1 - - RETURNING is_open - ` - var isOpen bool - err := tx.GetContext(ctx, &isOpen, query, account) - if err != nil { - return err - } - if isOpen { - return errors.New("unexpected state transition") - } - return nil - }) -} diff --git a/ocp/data/balance/postgres/store_legacy.go b/ocp/data/balance/postgres/store_legacy.go deleted file mode 100644 index f947bf1..0000000 --- a/ocp/data/balance/postgres/store_legacy.go +++ /dev/null @@ -1,25 +0,0 @@ -package postgres - -import ( - "context" -) - -// GetCachedVersion implements balance.Store.GetCachedVersion -func (s *store) GetCachedVersion(ctx context.Context, account string) (uint64, error) { - return dbGetCachedVersion(ctx, s.db, account) -} - -// AdvanceCachedVersion implements balance.Store.AdvanceCachedVersion -func (s *store) AdvanceCachedVersion(ctx context.Context, account string, currentVersion uint64) error { - return dbAdvanceCachedVersion(ctx, s.db, account, currentVersion) -} - -// CheckNotClosed implements balance.Store.CheckNotClosed -func (s *store) CheckNotClosed(ctx context.Context, account string) error { - return dbCheckNotClosed(ctx, s.db, account) -} - -// MarkAsClosed implements balance.Store.MarkAsClosed -func (s *store) MarkAsClosed(ctx context.Context, account string) error { - return dbMarkAsClosed(ctx, s.db, account) -} diff --git a/ocp/data/balance/postgres/store_test.go b/ocp/data/balance/postgres/store_test.go index e322a0f..2002cba 100644 --- a/ocp/data/balance/postgres/store_test.go +++ b/ocp/data/balance/postgres/store_test.go @@ -46,24 +46,6 @@ const ( CREATE INDEX ocp__core_balance__idx__owner_account__mint_account ON ocp__core_balance (owner_account, mint_account); CREATE INDEX ocp__core_balance__idx__mint_account__id ON ocp__core_balance (mint_account, id); - CREATE TABLE ocp__core_cachedbalanceversion ( - id SERIAL NOT NULL PRIMARY KEY, - - token_account TEXT NOT NULL, - version INTEGER NOT NULL, - - CONSTRAINT ocp__core_cachedbalanceversion__unique__token_account UNIQUE (token_account) - ); - - CREATE TABLE ocp__core_opencloselocks ( - id SERIAL NOT NULL PRIMARY KEY, - - token_account TEXT NOT NULL, - is_open BOOL NOT NULL, - - CONSTRAINT ocp__core_opencloselocks__unique__token_account UNIQUE (token_account) - ); - CREATE TABLE ocp__core_externalbalancecheckpoint ( id SERIAL NOT NULL PRIMARY KEY, @@ -80,8 +62,6 @@ const ( // Used for testing ONLY, the table and migrations are external to this repository tableDestroy = ` DROP TABLE ocp__core_balance; - DROP TABLE ocp__core_cachedbalanceversion; - DROP TABLE ocp__core_opencloselocks; DROP TABLE ocp__core_externalbalancecheckpoint; ` ) diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index e5580e4..bd86fe5 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -25,8 +25,6 @@ var ( // reviewed rather than recorded. ErrNegativeBalance = errors.New("backfilled balance is negative") - ErrStaleCachedBalanceVersion = errors.New("cached balance version is stale") - ErrAccountClosed = errors.New("account open state is stale") ErrCheckpointNotFound = errors.New("checkpoint not found") @@ -123,36 +121,4 @@ type Store interface { // // ErrCheckpointNotFound is returend if no DB record exists. GetExternalCheckpoint(ctx context.Context, account string) (*ExternalCheckpointRecord, error) - - // GetCachedVersion gets the current cached balance version, which can be used - // for optimistic locking cached balances for operations with outgoing transfers. - // - // Note: Use ApplyDeltas, whose predicates replace the version check. - // Retained for accounts that are not yet backfilled. - GetCachedVersion(ctx context.Context, account string) (uint64, error) - - // AdvanceCachedVersion advances an account's cached balance version. - // - // ErrStaleCachedBalanceVersion is returned if the currentVersion is out of date. - // - // Note: Use ApplyDeltas, whose predicates replace the version check. - // Retained for accounts that are not yet backfilled. - AdvanceCachedVersion(ctx context.Context, account string, currentVersion uint64) error - - // CheckNotClosed checks whether an account is closed under a lock to guarantee - // payments to a closeable destination with cached balances are made to an open - // account. - // - // ErrAccountClosed is returned if the account has been closed. - // - // Note: Use ApplyDeltas with DeltaCredit. Retained for accounts that - // are not yet backfilled. - CheckNotClosed(ctx context.Context, account string) error - - // MarkAsClosed marks an account as being closed and unable to receive payments - // as a destination. - // - // Note: Use ApplyDeltas with DeltaDrain or DeltaClose. Retained for - // accounts that are not yet backfilled. - MarkAsClosed(ctx context.Context, account string) error } diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index b9125c4..893ecc0 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -22,8 +22,6 @@ func RunTests(t *testing.T, s balance.Store, teardown func()) { testApplyDeltasAtomicity, testApplyDeltasConcurrency, testBackfill, - testCachedBalanceVersionHappyPath, - testClosedAccountHappyPath, testExternalCheckpointHappyPath, } { tf(t, s) @@ -549,44 +547,6 @@ func testBackfill(t *testing.T, s balance.Store) { }) } -func testCachedBalanceVersionHappyPath(t *testing.T, s balance.Store) { - t.Run("testCachedBalanceVersionHappyPath", func(t *testing.T) { - ctx := context.Background() - - for i := range 100 { - for range 10 { - currentVersion, err := s.GetCachedVersion(ctx, "token_account_1") - require.NoError(t, err) - assert.EqualValues(t, i, currentVersion) - } - - if i > 0 { - assert.Equal(t, balance.ErrStaleCachedBalanceVersion, s.AdvanceCachedVersion(ctx, "token_account_1", uint64(i-1))) - } - assert.Equal(t, balance.ErrStaleCachedBalanceVersion, s.AdvanceCachedVersion(ctx, "token_account_1", uint64(i+1))) - - require.NoError(t, s.AdvanceCachedVersion(ctx, "token_account_1", uint64(i))) - } - - currentVersion, err := s.GetCachedVersion(ctx, "token_account_2") - require.NoError(t, err) - assert.EqualValues(t, 0, currentVersion) - }) -} - -func testClosedAccountHappyPath(t *testing.T, s balance.Store) { - t.Run("testClosedAccountHappyPath", func(t *testing.T) { - ctx := context.Background() - - require.NoError(t, s.CheckNotClosed(ctx, "token_account_1")) - - require.NoError(t, s.MarkAsClosed(ctx, "token_account_1")) - - assert.Equal(t, balance.ErrAccountClosed, s.CheckNotClosed(ctx, "token_account_1")) - require.NoError(t, s.CheckNotClosed(ctx, "token_account_2s")) - }) -} - func testExternalCheckpointHappyPath(t *testing.T, s balance.Store) { t.Run("testExternalCheckpointHappyPath", func(t *testing.T) { ctx := context.Background() diff --git a/ocp/data/internal.go b/ocp/data/internal.go index 09c1f8e..f381450 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -130,10 +130,6 @@ type DatabaseData interface { CountBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error - GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) - AdvanceCachedBalanceVersion(ctx context.Context, account string, currentVersion uint64) error - CheckNotClosedForBalanceUpdate(ctx context.Context, account string) error - MarkAsClosedForBalanceUpdate(ctx context.Context, account string) error SaveExternalBalanceCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error GetExternalBalanceCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) @@ -496,18 +492,6 @@ func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*b func (dp *DatabaseProvider) BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { return dp.balance.Backfill(ctx, tokenAccount, fn) } -func (dp *DatabaseProvider) GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) { - return dp.balance.GetCachedVersion(ctx, account) -} -func (dp *DatabaseProvider) AdvanceCachedBalanceVersion(ctx context.Context, account string, currentVersion uint64) error { - return dp.balance.AdvanceCachedVersion(ctx, account, currentVersion) -} -func (dp *DatabaseProvider) CheckNotClosedForBalanceUpdate(ctx context.Context, account string) error { - return dp.balance.CheckNotClosed(ctx, account) -} -func (dp *DatabaseProvider) MarkAsClosedForBalanceUpdate(ctx context.Context, account string) error { - return dp.balance.MarkAsClosed(ctx, account) -} func (dp *DatabaseProvider) SaveExternalBalanceCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { return dp.balance.SaveExternalCheckpoint(ctx, record) } diff --git a/ocp/rpc/transaction/gift_card.go b/ocp/rpc/transaction/gift_card.go index bc1fbe4..06da3ee 100644 --- a/ocp/rpc/transaction/gift_card.go +++ b/ocp/rpc/transaction/gift_card.go @@ -12,7 +12,6 @@ import ( transactionpb "github.com/code-payments/ocp-protobuf-api/generated/go/transaction/v1" "github.com/code-payments/ocp-server/grpc/client" - "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" @@ -77,12 +76,6 @@ func (s *transactionServer) VoidGiftCard(ctx context.Context, req *transactionpb }, nil } - globalBalanceLock, err := balance.GetOptimisticVersionLock(ctx, s.data, giftCardVault) - if err != nil { - log.With(zap.Error(err)).Warn("failure getting balance lock") - return nil, status.Error(codes.Internal, "") - } - localAccountLock := s.getLocalAccountLock(giftCardVault) localAccountLock.Lock() defer localAccountLock.Unlock() @@ -120,7 +113,7 @@ func (s *transactionServer) VoidGiftCard(ctx context.Context, req *transactionpb return nil, status.Error(codes.Internal, "") } - err = account_worker.InitiateProcessToAutoReturnGiftCard(ctx, s.data, giftCardVault, true, globalBalanceLock) + err = account_worker.InitiateProcessToAutoReturnGiftCard(ctx, s.data, giftCardVault, true) if err != nil { log.With(zap.Error(err)).Warn("failure scheduling auto-return action") return nil, status.Error(codes.Internal, "") diff --git a/ocp/rpc/transaction/intent.go b/ocp/rpc/transaction/intent.go index 49e1d33..b2d1e2f 100644 --- a/ocp/rpc/transaction/intent.go +++ b/ocp/rpc/transaction/intent.go @@ -292,24 +292,24 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm return nil } - // Lock any acccounts with fund movement that is not resistent to race conditions - // 1. Global DB layer lock to guarantee balance consistency in a mult-server environment - // 2. Local in memory lock to avoid over consumption of local resources (eg. - // nonces) when we're likely to encounter a race resulting in DB txn rollback - // (eg. mass attempt to claim gift card). - globalBalanceLocks, err := intentHandler.GetBalanceLocks(ctx, intentRecord, submitActionsReq.Metadata) + // Locally lock any acccounts with fund movement to avoid over consumption + // of local resources (eg. nonces) when we're likely to encounter a race + // resulting in DB txn rollback (eg. mass attempt to claim gift card). + // Cross-server race resistance comes from the balance ledger's predicates + // in the intent DB transaction. + accountsToLock, err := intentHandler.GetAccountsToLock(ctx, intentRecord, submitActionsReq.Metadata) if err != nil { - log.With(zap.Error(err)).Warn("failure getting accounts with balances to lock") + log.With(zap.Error(err)).Warn("failure getting accounts to lock") return handleSubmitIntentError(ctx, streamer, intentRecord, err) } localAccountLocks := make([]*sync.Mutex, 0) locallyLockedAccounts := make(map[string]any) - for _, globalBalanceLock := range globalBalanceLocks { - _, ok := locallyLockedAccounts[globalBalanceLock.Account.PublicKey().ToBase58()] + for _, accountToLock := range accountsToLock { + _, ok := locallyLockedAccounts[accountToLock.PublicKey().ToBase58()] if !ok { - localAccountLocks = append(localAccountLocks, s.getLocalAccountLock(globalBalanceLock.Account)) + localAccountLocks = append(localAccountLocks, s.getLocalAccountLock(accountToLock)) } - locallyLockedAccounts[globalBalanceLock.Account.PublicKey().ToBase58()] = true + locallyLockedAccounts[accountToLock.PublicKey().ToBase58()] = true } for _, localAccountLock := range localAccountLocks { localAccountLock.Lock() @@ -717,14 +717,6 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm return err } - for _, globalBalanceLock := range globalBalanceLocks { - err = globalBalanceLock.CommitFn(ctx, s.data) - if err != nil { - log.With(zap.Error(err)).Warn("failure commiting balance update") - return err - } - } - // Reflect the intent in the balance ledger. Applied last, so the // ledger row locks are held for as little of the transaction as // possible. diff --git a/ocp/rpc/transaction/intent_handler.go b/ocp/rpc/transaction/intent_handler.go index 9cbce69..6b6d835 100644 --- a/ocp/rpc/transaction/intent_handler.go +++ b/ocp/rpc/transaction/intent_handler.go @@ -64,9 +64,14 @@ type CreateIntentHandler interface { // error. IsNoop(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata, actions []*transactionpb.Action) (bool, error) - // GetBalanceLocks gets a set of global balance locks to prevent race conditions + // GetAccountsToLock gets the set of accounts to lock in-process before + // validation, to avoid over consumption of local resources (eg. nonces) + // on requests that are likely to lose a race and roll back. Only expected + // mass races are worth locking, like a gift card being claimed from many + // devices at once. Race resistance itself comes from the balance ledger's + // predicates in the intent transaction. // against invalid balance updates that would result in intent fulfillment failure - GetBalanceLocks(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*intentBalanceLock, error) + GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) // AllowCreation determines whether the new intent creation should be allowed. AllowCreation(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata, actions []*transactionpb.Action) error @@ -192,7 +197,7 @@ func (h *OpenAccountsIntentHandler) IsNoop(ctx context.Context, intentRecord *in return accountInfoRecord.TokenAccount == tokenAccountToCheck.PublicKey().ToBase58() && accountInfoRecord.AccountType == expectedAccountType, nil } -func (h *OpenAccountsIntentHandler) GetBalanceLocks(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*intentBalanceLock, error) { +func (h *OpenAccountsIntentHandler) GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) { return nil, nil } @@ -508,58 +513,8 @@ func (h *SendPublicPaymentIntentHandler) IsNoop(ctx context.Context, intentRecor return false, nil } -func (h *SendPublicPaymentIntentHandler) GetBalanceLocks(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*intentBalanceLock, error) { - typedMetadata := metadata.GetSendPublicPayment() - if typedMetadata == nil { - return nil, errors.New("unexpected metadata proto message") - } - - sourceVault, err := common.NewAccountFromProto(typedMetadata.Source) - if err != nil { - return nil, err - } - - outgoingSourceBalanceLock, err := balance.GetOptimisticVersionLock(ctx, h.data, sourceVault) - if err != nil { - return nil, err - } - - intentBalanceLocks := []*intentBalanceLock{ - { - Account: sourceVault, - CommitFn: outgoingSourceBalanceLock.OnNewBalanceVersion, - }, - } - - if h.cachedDestinationAccountInfoRecord != nil { - switch h.cachedDestinationAccountInfoRecord.AccountType { - case commonpb.AccountType_POOL: - closeableDestinationVault, err := common.NewAccountFromProto(typedMetadata.Destination) - if err != nil { - return nil, err - } - - incomingDestinationBalanceLock, err := balance.GetOptimisticVersionLock(ctx, h.data, closeableDestinationVault) - if err != nil { - return nil, err - } - incomingDestinationOpenCloseLock := balance.NewOpenCloseStatusLock(closeableDestinationVault) - - intentBalanceLocks = append( - intentBalanceLocks, - &intentBalanceLock{ - Account: closeableDestinationVault, - CommitFn: incomingDestinationBalanceLock.RequireSameBalanceVerion, - }, - &intentBalanceLock{ - Account: closeableDestinationVault, - CommitFn: incomingDestinationOpenCloseLock.OnPaymentToAccount, - }, - ) - } - } - - return intentBalanceLocks, nil +func (h *SendPublicPaymentIntentHandler) GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) { + return nil, nil } func (h *SendPublicPaymentIntentHandler) AllowCreation(ctx context.Context, intentRecord *intent.Record, untypedMetadata *transactionpb.Metadata, actions []*transactionpb.Action) error { @@ -1088,23 +1043,16 @@ func (h *ReceivePaymentsPubliclyIntentHandler) IsNoop(ctx context.Context, inten return false, nil } -func (h *ReceivePaymentsPubliclyIntentHandler) GetBalanceLocks(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*intentBalanceLock, error) { +func (h *ReceivePaymentsPubliclyIntentHandler) GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) { + // A gift card claim is an expected mass race across devices, so requests + // on this server queue behind one lock instead of burning nonces on + // attempts that are doomed to lose the gift card's drain giftCardVault, err := common.NewAccountFromPublicKeyString(intentRecord.ReceivePaymentsPubliclyMetadata.Source) if err != nil { return nil, err } - outgoingGiftCardBalanceLock, err := balance.GetOptimisticVersionLock(ctx, h.data, giftCardVault) - if err != nil { - return nil, err - } - - return []*intentBalanceLock{ - { - Account: giftCardVault, - CommitFn: outgoingGiftCardBalanceLock.OnNewBalanceVersion, - }, - }, nil + return []*common.Account{giftCardVault}, nil } func (h *ReceivePaymentsPubliclyIntentHandler) AllowCreation(ctx context.Context, intentRecord *intent.Record, untypedMetadata *transactionpb.Metadata, actions []*transactionpb.Action) error { @@ -1437,29 +1385,8 @@ func (h *PublicDistributionIntentHandler) IsNoop(ctx context.Context, intentReco return false, nil } -func (h *PublicDistributionIntentHandler) GetBalanceLocks(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*intentBalanceLock, error) { - poolVault, err := common.NewAccountFromPublicKeyString(intentRecord.PublicDistributionMetadata.Source) - if err != nil { - return nil, err - } - - outgoingPoolBalanceLock, err := balance.GetOptimisticVersionLock(ctx, h.data, poolVault) - if err != nil { - return nil, err - } - - incomingPoolBalanceLock := balance.NewOpenCloseStatusLock(poolVault) - - return []*intentBalanceLock{ - { - Account: poolVault, - CommitFn: outgoingPoolBalanceLock.OnNewBalanceVersion, - }, - { - Account: poolVault, - CommitFn: incomingPoolBalanceLock.OnClose, - }, - }, nil +func (h *PublicDistributionIntentHandler) GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) { + return nil, nil } // todo: Not all multi-mint validation checks are implemented diff --git a/ocp/worker/account/gift_card.go b/ocp/worker/account/gift_card.go index e2cbb6d..e82989b 100644 --- a/ocp/worker/account/gift_card.go +++ b/ocp/worker/account/gift_card.go @@ -93,12 +93,6 @@ func (p *runtime) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn return err } - balanceLock, err := balance.GetOptimisticVersionLock(ctx, p.data, giftCardVaultAccount) - if err != nil { - log.With(zap.Error(err)).Warn("failure getting balance lock") - return err - } - _, err = p.data.GetGiftCardClaimedAction(ctx, giftCardVaultAccount.PublicKey().ToBase58()) if err == nil { log.Debug("gift card is claimed and will be removed from worker queue") @@ -130,7 +124,7 @@ func (p *runtime) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn // There's no action to claim the gift card and the expiry window has been met. // It's time to initiate the process of auto-returning the funds back to the // issuer. - err = InitiateProcessToAutoReturnGiftCard(ctx, p.data, giftCardVaultAccount, false, balanceLock) + err = InitiateProcessToAutoReturnGiftCard(ctx, p.data, giftCardVaultAccount, false) if err != nil { log.With(zap.Error(err)).Warn("failure initiating process to return gift card balance to issuer") return err @@ -149,7 +143,7 @@ func (p *runtime) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn // a good guide for similar actions in the future. // // todo: This probably belongs somewhere more common -func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Provider, giftCardVaultAccount *common.Account, isVoidedByUser bool, balanceLock *balance.OptimisticVersionLock) error { +func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Provider, giftCardVaultAccount *common.Account, isVoidedByUser bool) error { return data.ExecuteInTx(ctx, sql.LevelDefault, func(ctx context.Context) error { giftCardIssuedIntent, err := data.GetOriginalGiftCardIssuedIntent(ctx, giftCardVaultAccount.PublicKey().ToBase58()) if err != nil { @@ -232,12 +226,7 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Prov // This will trigger the fulfillment worker to poll for the fulfillment. This // should be the very last DB update called. - err = markFulfillmentAsActivelyScheduled(ctx, data, autoReturnFulfillment[0]) - if err != nil { - return err - } - - return balanceLock.OnNewBalanceVersion(ctx, data) + return markFulfillmentAsActivelyScheduled(ctx, data, autoReturnFulfillment[0]) }) } From da32ee2ea52b587c3b9e2b3b93abf45c323e312e Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 12:25:06 -0400 Subject: [PATCH 12/26] Tighten delta debit check against open status --- ocp/data/balance/memory/store.go | 3 +++ ocp/data/balance/postgres/model.go | 5 ++++- ocp/data/balance/store.go | 5 +++-- ocp/data/balance/tests/tests.go | 12 +++++++++++- 4 files changed, 21 insertions(+), 4 deletions(-) diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index abb2e7d..78f274f 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -201,6 +201,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 } diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index dcaadcd..2d3813d 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -204,7 +204,7 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er case balance.DeltaDebit: query = `UPDATE ` + tableName + ` SET quarks = quarks - $2, usd_cost_basis = usd_cost_basis - $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR quarks >= $2)` + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND quarks >= $2))` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDrain: query = `UPDATE ` + tableName + ` @@ -254,6 +254,9 @@ func classifyFailedDelta(delta *balance.Delta, current *balance.Record) error { case balance.DeltaCredit: return balance.ErrAccountClosed case balance.DeltaDebit: + if !current.IsOpen { + return balance.ErrAccountClosed + } return balance.ErrInsufficientBalance case balance.DeltaDrain, balance.DeltaClose: if !current.IsOpen { diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index bd86fe5..e87e265 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -96,8 +96,9 @@ type Store interface { // // ErrInsufficientBalance is returned when a debit exceeds the balance. // ErrBalanceChanged is returned when a drain or close doesn't match the - // balance. ErrAccountClosed is returned when a credit, drain or close - // targets a closed account. + // balance. ErrAccountClosed is returned when any delta targets a closed + // account: a closed account is frozen, so nothing may enter or leave it, + // including a zero-quark cost basis adjustment. ApplyDeltas(ctx context.Context, deltas ...*Delta) error // Backfill locks a record that is not yet backfilled, calls fn to compute diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 893ecc0..96df3d4 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -249,6 +249,12 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, UsdCostBasis: -5})) assertBalance(t, s, "token_account_1", 70, -15, true) + // So can a debit, while the account is open + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, UsdCostBasis: -5})) + assertBalance(t, s, "token_account_1", 70, -10, true) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, UsdCostBasis: 5})) + assertBalance(t, s, "token_account_1", 70, -15, true) + assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 69})) assert.Equal(t, balance.ErrBalanceChanged, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 71})) @@ -260,7 +266,11 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1})) assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 0, UsdCostBasis: 1})) assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) - assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) + + // A closed account is frozen: even a zero-quark cost basis adjustment + // cannot leave it + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, UsdCostBasis: 1})) assertBalance(t, s, "token_account_1", 0, 0, false) require.NoError(t, s.Create(ctx, &balance.Record{ From 7cd3d4583b3354afd5219ae63fd7db5593278c0c Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 12:39:32 -0400 Subject: [PATCH 13/26] Default balance ledger read and writes to true --- ocp/balance/config.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ocp/balance/config.go b/ocp/balance/config.go index 78be8a4..6925291 100644 --- a/ocp/balance/config.go +++ b/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 ( From 51f5de75442d7d48961e83d674ac6b3b6abaedd9 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 13:41:55 -0400 Subject: [PATCH 14/26] Balance rows now have lock state --- ocp/balance/calculator.go | 18 +++- ocp/balance/calculator_test.go | 53 ++++++++++++ ocp/balance/ledger.go | 11 ++- ocp/balance/ledger_test.go | 25 +++++- ocp/data/balance/memory/store.go | 26 +++++- ocp/data/balance/postgres/model.go | 49 ++++++++--- ocp/data/balance/postgres/store.go | 11 ++- ocp/data/balance/postgres/store_test.go | 3 +- ocp/data/balance/record.go | 9 ++ ocp/data/balance/store.go | 29 ++++++- ocp/data/balance/tests/tests.go | 104 +++++++++++++++++++++--- ocp/data/internal.go | 10 ++- ocp/worker/geyser/timelock.go | 9 +- 13 files changed, 314 insertions(+), 43 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 9a10f42..0d13e25 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -399,7 +399,10 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide // in balance.UsdQuarksPerUnit, using cached values. // // Note: Unlike quark balances, a cost basis for an account not managed by Code -// is still meaningful, so no timelock check is performed. +// is still meaningful, so no timelock check is performed and such accounts fall +// back to the legacy calculation. The materialized record is the exception: +// once a vault unlocks it holds the last managed state rather than a live cost +// basis, so reading one returns ErrNotManagedByCode. func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (int64, error) { tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "CalculateUsdCostBasisFromCache") tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) @@ -408,6 +411,10 @@ func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, if enableLedgerReads.Get(ctx) { balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) if err == nil && balanceRecord.IsBackfilled { + if !balanceRecord.IsLocked { + tracer.OnError(ErrNotManagedByCode) + return 0, ErrNotManagedByCode + } return balanceRecord.UsdCostBasis, nil } else if err != nil && err != balance.ErrRecordNotFound { tracer.OnError(err) @@ -448,6 +455,10 @@ func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Prov for _, tokenAccount := range tokenAccountStrings { balanceRecord, ok := balanceRecords[tokenAccount] if ok && balanceRecord.IsBackfilled { + if !balanceRecord.IsLocked { + tracer.OnError(ErrNotManagedByCode) + return nil, ErrNotManagedByCode + } res[tokenAccount] = balanceRecord.UsdCostBasis continue } @@ -487,6 +498,11 @@ func legacyUsdCostBasis(ctx context.Context, data ocp_data.Provider, tokenAccoun } func quarksFromRecord(record *balance.Record) (uint64, error) { + // Callers reject unlocked vaults on the timelock record before reaching + // here, so this only guards against a record that disagrees with it. + if !record.IsLocked { + return 0, ErrNotManagedByCode + } if record.Quarks < 0 { return 0, ErrNegativeBalance } diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index c662202..fc6ff0a 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -376,6 +376,7 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { MintAccount: vmConfig.Mint.PublicKey().ToBase58(), Quarks: 42, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) @@ -418,6 +419,56 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { assert.Equal(t, expected, balanceByAccount) } +func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) { + env := setupBalanceTestEnv(t) + enableLedgerReadsForTest(t) + + vmConfig := testutil.NewRandomVmConfig(t, true) + owner := testutil.NewRandomAccount(t) + tokenAccount, err := owner.ToTimelockVault(vmConfig) + require.NoError(t, err) + + externalAccount := testutil.NewRandomAccount(t) + + data := &balanceTestData{ + vmConfig: vmConfig, + codeUsers: []*common.Account{owner}, + transactions: []balanceTestTransaction{ + {source: externalAccount, destination: tokenAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, + }, + } + + setupBalanceTestData(t, env, data) + + // A backfilled record for an unlocked vault holds the last managed state, + // not a live balance, so it is refused even though the timelock record + // still passes the managed check. That pairing is inconsistent by + // construction: the timelock check normally rejects first, so the fixture + // exists to exercise the record's own guard. + require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ + TokenAccount: tokenAccount.PublicKey().ToBase58(), + OwnerAccount: owner.PublicKey().ToBase58(), + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + Quarks: 42, + UsdCostBasis: 4_200_000, + IsOpen: true, + IsLocked: false, + IsBackfilled: true, + })) + + _, err = CalculateFromCache(env.ctx, env.data, tokenAccount) + assert.Equal(t, ErrNotManagedByCode, err) + + _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) + assert.Equal(t, ErrNotManagedByCode, err) + + _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) + assert.Equal(t, ErrNotManagedByCode, err) + + _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) + assert.Equal(t, ErrNotManagedByCode, err) +} + func TestDefaultCalculationMethods_BalanceRecordReadsDisabled(t *testing.T) { env := setupBalanceTestEnv(t) @@ -446,6 +497,7 @@ func TestDefaultCalculationMethods_BalanceRecordReadsDisabled(t *testing.T) { Quarks: 42, UsdCostBasis: 123, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) @@ -487,6 +539,7 @@ func TestUsdCostBasisCalculationMethods(t *testing.T) { MintAccount: vmConfig.Mint.PublicKey().ToBase58(), UsdCostBasis: -123456, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index 276bed8..653cd33 100644 --- a/ocp/balance/ledger.go +++ b/ocp/balance/ledger.go @@ -31,15 +31,18 @@ func LedgerWritesEnabled(ctx context.Context) bool { // like an external wallet or the fee collector, are dropped, since delta // builders don't know which destinations OCP manages. Outgoing deltas from // an account the ledger doesn't track are ErrUntrackedAccount, since funds -// only ever leave accounts OCP manages. +// only ever leave accounts OCP manages. Any delta against a tracked account +// whose vault has unlocked fails loudly with balance.ErrAccountUnlocked: +// the record is no longer maintained, and a flow still moving funds through +// it is a bug to surface, never to paper over. // // Any timelock account in the delta set that has no ledger record yet lazily // gets one that is not backfilled, so accounts that predate the ledger start // accumulating deltas on first touch regardless of direction. // // Store predicate failures (balance.ErrInsufficientBalance, -// balance.ErrBalanceChanged, balance.ErrAccountClosed) are returned as is -// for the caller to map. +// balance.ErrBalanceChanged, balance.ErrAccountClosed, +// balance.ErrAccountUnlocked) are returned as is for the caller to map. func ApplyDeltasInTx(ctx context.Context, data ocp_data.Provider, deltas ...*balance.Delta) error { if !enableLedgerWrites.Get(ctx) || len(deltas) == 0 { return nil @@ -88,6 +91,7 @@ func CreateRecordInTx(ctx context.Context, data ocp_data.Provider, accountInfoRe OwnerAccount: accountInfoRecord.OwnerAccount, MintAccount: accountInfoRecord.MintAccount, IsOpen: true, + IsLocked: true, IsBackfilled: true, }) if errors.Is(err, balance.ErrRecordExists) { @@ -136,6 +140,7 @@ func resolveRecords(ctx context.Context, data ocp_data.Provider, deltas []*balan OwnerAccount: accountInfoRecord.OwnerAccount, MintAccount: accountInfoRecord.MintAccount, IsOpen: true, + IsLocked: true, IsBackfilled: false, }) if err != nil && !errors.Is(err, balance.ErrRecordExists) { diff --git a/ocp/balance/ledger_test.go b/ocp/balance/ledger_test.go index 34d4c37..09c90aa 100644 --- a/ocp/balance/ledger_test.go +++ b/ocp/balance/ledger_test.go @@ -74,7 +74,7 @@ func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { // Once backfilled, predicates are enforced require.NoError(t, data.BackfillBalance(ctx, source, func(context.Context) (*balance.BackfillResult, error) { - return &balance.BackfillResult{Quarks: 400, UsdCostBasis: 4_000_000, IsOpen: true}, nil + return &balance.BackfillResult{Quarks: 400, UsdCostBasis: 4_000_000, IsOpen: true, IsLocked: true}, nil })) err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 401}) assert.Equal(t, balance.ErrInsufficientBalance, err) @@ -110,6 +110,29 @@ func TestApplyDeltasInTx_OnlyUntrackedCredits(t *testing.T) { assert.Equal(t, balance.ErrRecordNotFound, err) } +func TestApplyDeltasInTx_UnlockedAccount(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + enableLedgerWritesForTest(t) + + unlocked := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY) + require.NoError(t, CreateRecordInTx(ctx, data, unlocked)) + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaCredit, Quarks: 100})) + require.NoError(t, data.MarkBalanceAsUnlocked(ctx, unlocked.TokenAccount)) + + // Any delta against an unlocked account fails loudly, so a flow still + // moving funds through it surfaces as a DB error + err := ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaCredit, Quarks: 1}) + assert.Equal(t, balance.ErrAccountUnlocked, err) + err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaDebit, Quarks: 1}) + assert.Equal(t, balance.ErrAccountUnlocked, err) + + record, err := data.GetBalance(ctx, unlocked.TokenAccount) + require.NoError(t, err) + assert.EqualValues(t, 100, record.Quarks) + assert.False(t, record.IsLocked) +} + func TestApplyDeltasInTx_InvalidDelta(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 911b8fd..549696c 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -104,13 +104,13 @@ func (s *store) GetAllByOwnerAndMint(_ context.Context, owner, mint string) ([]* }) } -// GetAllByMint implements balance.Store.GetAllByMint -func (s *store) GetAllByMint(_ context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { +// GetAllLockedByMint implements balance.Store.GetAllLockedByMint +func (s *store) GetAllLockedByMint(_ context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { s.mu.Lock() defer s.mu.Unlock() res, err := s.filter(func(item *balance.Record) bool { - if item.MintAccount != mint || item.Quarks < minQuarks { + if item.MintAccount != mint || item.Quarks < minQuarks || !item.IsLocked { return false } if len(cursor) > 0 { @@ -139,6 +139,21 @@ func (s *store) GetAllByMint(_ context.Context, mint string, minQuarks int64, cu return res, nil } +// MarkAsUnlocked implements balance.Store.MarkAsUnlocked +func (s *store) MarkAsUnlocked(_ context.Context, tokenAccount string) error { + s.mu.Lock() + defer s.mu.Unlock() + + item, ok := s.balanceRecordsByTokenAccount[tokenAccount] + if !ok { + return balance.ErrRecordNotFound + } + + item.IsLocked = false + item.UpdatedAt = time.Now() + return nil +} + // ApplyDeltas implements balance.Store.ApplyDeltas func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { for _, delta := range deltas { @@ -183,6 +198,10 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { func applyDelta(item *balance.Record, delta *balance.Delta) error { enforce := item.IsBackfilled + if enforce && !item.IsLocked { + return balance.ErrAccountUnlocked + } + switch delta.Kind { case balance.DeltaCredit: if enforce && !item.IsOpen { @@ -257,6 +276,7 @@ func (s *store) Backfill(ctx context.Context, tokenAccount string, fn balance.Ba item.Quarks = result.Quarks item.UsdCostBasis = result.UsdCostBasis item.IsOpen = result.IsOpen + item.IsLocked = result.IsLocked item.IsBackfilled = true item.UpdatedAt = time.Now() return nil diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index b5e2c93..30dae31 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -17,7 +17,7 @@ const ( tableName = "ocp__core_balance" externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" - allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_backfilled, updated_at" + allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_locked, is_backfilled, updated_at" ) type model struct { @@ -31,6 +31,7 @@ type model struct { UsdCostBasis int64 `db:"usd_cost_basis"` IsOpen bool `db:"is_open"` + IsLocked bool `db:"is_locked"` IsBackfilled bool `db:"is_backfilled"` UpdatedAt time.Time `db:"updated_at"` @@ -50,6 +51,7 @@ func toModel(obj *balance.Record) (*model, error) { UsdCostBasis: obj.UsdCostBasis, IsOpen: obj.IsOpen, + IsLocked: obj.IsLocked, IsBackfilled: obj.IsBackfilled, UpdatedAt: obj.UpdatedAt, @@ -68,6 +70,7 @@ func fromModel(obj *model) *balance.Record { UsdCostBasis: obj.UsdCostBasis, IsOpen: obj.IsOpen, + IsLocked: obj.IsLocked, IsBackfilled: obj.IsBackfilled, UpdatedAt: obj.UpdatedAt, @@ -77,8 +80,8 @@ func fromModel(obj *model) *balance.Record { func (m *model) dbCreate(ctx context.Context, db *sqlx.DB) error { return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { query := `INSERT INTO ` + tableName + ` - (token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_backfilled, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + (token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_locked, is_backfilled, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING ` + allColumns m.UpdatedAt = time.Now() @@ -92,6 +95,7 @@ func (m *model) dbCreate(ctx context.Context, db *sqlx.DB) error { m.Quarks, m.UsdCostBasis, m.IsOpen, + m.IsLocked, m.IsBackfilled, m.UpdatedAt.UTC(), ).StructScan(m) @@ -157,11 +161,11 @@ func dbGetAllByOwner(ctx context.Context, db *sqlx.DB, owner string, mint *strin return res, nil } -func dbGetAllByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64, cursor q.Cursor, limit uint64, direction q.Ordering) ([]*model, error) { +func dbGetAllLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64, cursor q.Cursor, limit uint64, direction q.Ordering) ([]*model, error) { res := []*model{} query := `SELECT ` + allColumns + ` FROM ` + tableName + ` - WHERE (mint_account = $1 AND quarks >= $2)` + WHERE (mint_account = $1 AND quarks >= $2 AND is_locked)` query, args := q.PaginateQuery(query, []any{mint, minQuarks}, cursor, limit, direction) err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { @@ -176,6 +180,26 @@ func dbGetAllByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int return res, nil } +func dbMarkAsUnlocked(ctx context.Context, db *sqlx.DB, tokenAccount string) error { + return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + query := `UPDATE ` + tableName + ` + SET is_locked = FALSE, updated_at = $2 + WHERE token_account = $1` + sqlResult, err := tx.ExecContext(ctx, query, tokenAccount, time.Now().UTC()) + if err != nil { + return err + } + rowsAffected, err := sqlResult.RowsAffected() + if err != nil { + return err + } + if rowsAffected == 0 { + return balance.ErrRecordNotFound + } + return nil + }) +} + // dbApplyDeltas applies every delta in a single transaction. Each delta is one // conditional UPDATE, so its predicate is evaluated against the row after the // row lock is acquired. Predicates only apply to backfilled rows. @@ -188,12 +212,12 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er case balance.DeltaCredit: query = `UPDATE ` + tableName + ` SET quarks = quarks + $2, usd_cost_basis = usd_cost_basis + $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR is_open)` + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked))` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDebit: query = `UPDATE ` + tableName + ` SET quarks = quarks - $2, usd_cost_basis = usd_cost_basis - $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR quarks >= $2)` + WHERE token_account = $1 AND (NOT is_backfilled OR (is_locked AND quarks >= $2))` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDrain: query = `UPDATE ` + tableName + ` @@ -201,12 +225,12 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er usd_cost_basis = CASE WHEN is_backfilled THEN 0 ELSE usd_cost_basis - $3 END, is_open = FALSE, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND quarks = $2))` + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks = $2))` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaClose: query = `UPDATE ` + tableName + ` SET is_open = FALSE, updated_at = $2 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND quarks = 0))` + WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks = 0))` args = []any{delta.TokenAccount, time.Now().UTC()} default: return fmt.Errorf("unsupported delta kind: %s", delta.Kind) @@ -239,6 +263,9 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er } func classifyFailedDelta(delta *balance.Delta, current *balance.Record) error { + if !current.IsLocked { + return balance.ErrAccountUnlocked + } switch delta.Kind { case balance.DeltaCredit: return balance.ErrAccountClosed @@ -276,9 +303,9 @@ func dbBackfill(ctx context.Context, db *sqlx.DB, tokenAccount string, fn balanc } query := `UPDATE ` + tableName + ` - SET quarks = $2, usd_cost_basis = $3, is_open = $4, is_backfilled = TRUE, updated_at = $5 + SET quarks = $2, usd_cost_basis = $3, is_open = $4, is_locked = $5, is_backfilled = TRUE, updated_at = $6 WHERE token_account = $1` - _, err = tx.ExecContext(ctx, query, tokenAccount, result.Quarks, result.UsdCostBasis, result.IsOpen, time.Now().UTC()) + _, err = tx.ExecContext(ctx, query, tokenAccount, result.Quarks, result.UsdCostBasis, result.IsOpen, result.IsLocked, time.Now().UTC()) return err }) }) diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index 8b8edcc..a47a362 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -77,15 +77,20 @@ func (s *store) GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([ return fromModels(models), nil } -// GetAllByMint implements balance.Store.GetAllByMint -func (s *store) GetAllByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { - models, err := dbGetAllByMint(ctx, s.db, mint, minQuarks, cursor, limit, direction) +// GetAllLockedByMint implements balance.Store.GetAllLockedByMint +func (s *store) GetAllLockedByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { + models, err := dbGetAllLockedByMint(ctx, s.db, mint, minQuarks, cursor, limit, direction) if err != nil { return nil, err } return fromModels(models), nil } +// MarkAsUnlocked implements balance.Store.MarkAsUnlocked +func (s *store) MarkAsUnlocked(ctx context.Context, tokenAccount string) error { + return dbMarkAsUnlocked(ctx, s.db, tokenAccount) +} + // ApplyDeltas implements balance.Store.ApplyDeltas func (s *store) ApplyDeltas(ctx context.Context, deltas ...*balance.Delta) error { for _, delta := range deltas { diff --git a/ocp/data/balance/postgres/store_test.go b/ocp/data/balance/postgres/store_test.go index e322a0f..50d1fa0 100644 --- a/ocp/data/balance/postgres/store_test.go +++ b/ocp/data/balance/postgres/store_test.go @@ -35,6 +35,7 @@ const ( usd_cost_basis BIGINT NOT NULL DEFAULT 0, is_open BOOL NOT NULL DEFAULT TRUE, + is_locked BOOL NOT NULL DEFAULT TRUE, is_backfilled BOOL NOT NULL DEFAULT FALSE, updated_at TIMESTAMP WITH TIME ZONE NOT NULL, @@ -44,7 +45,7 @@ const ( ) WITH (fillfactor = 90); CREATE INDEX ocp__core_balance__idx__owner_account__mint_account ON ocp__core_balance (owner_account, mint_account); - CREATE INDEX ocp__core_balance__idx__mint_account__id ON ocp__core_balance (mint_account, id); + CREATE INDEX ocp__core_balance__idx__mint_account__id ON ocp__core_balance (mint_account, id) WHERE is_locked; CREATE TABLE ocp__core_cachedbalanceversion ( id SERIAL NOT NULL PRIMARY KEY, diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index 5d6e8e8..769f095 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -44,6 +44,13 @@ type Record struct { IsOpen bool + // IsLocked indicates the timelock vault is still locked, so the account + // is managed by OCP and every balance change flows through the ledger. + // Once a vault unlocks, funds can move on chain without an intent, so + // the record's values are the last managed state and must not be + // trusted or aggregated. Unlocking is one-way. + IsLocked bool + // IsBackfilled indicates the record reflects the full history of the // account. Until it does, deltas are recorded without enforcing any // balance predicates. @@ -84,6 +91,7 @@ func (r *Record) Clone() Record { UsdCostBasis: r.UsdCostBasis, IsOpen: r.IsOpen, + IsLocked: r.IsLocked, IsBackfilled: r.IsBackfilled, UpdatedAt: r.UpdatedAt, @@ -101,6 +109,7 @@ func (r *Record) CopyTo(dst *Record) { dst.UsdCostBasis = r.UsdCostBasis dst.IsOpen = r.IsOpen + dst.IsLocked = r.IsLocked dst.IsBackfilled = r.IsBackfilled dst.UpdatedAt = r.UpdatedAt diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index 8d212cf..f721974 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -29,6 +29,11 @@ var ( ErrAccountClosed = errors.New("account open state is stale") + // ErrAccountUnlocked is returned when a delta targets an account whose + // timelock vault has unlocked. The ledger stops maintaining the record at + // unlock, so nothing may enter or leave it. + ErrAccountUnlocked = errors.New("account is unlocked") + ErrCheckpointNotFound = errors.New("checkpoint not found") ErrStaleCheckpoint = errors.New("checkpoint is stale") ) @@ -41,6 +46,10 @@ type BackfillResult struct { // IsOpen is false for accounts that can no longer receive funds, such as // claimed gift cards and distributed pools. IsOpen bool + + // IsLocked is false for accounts whose timelock vault has unlocked, so + // the balance is the last managed state rather than a live value. + IsLocked bool } // BackfillFunc computes the full historical state of a token account. It @@ -74,11 +83,13 @@ type Store interface { // ErrRecordNotFound is returned if no records exist. GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([]*Record, error) - // GetAllByMint gets balance records for a mint with at least minQuarks, - // paged by record ID. + // GetAllLockedByMint gets locked balance records for a mint with at + // least minQuarks, paged by record ID. Unlocked records are excluded, + // since funds can move on chain without an intent once a vault unlocks, + // making their balances stale. // // ErrRecordNotFound is returned if no records exist. - GetAllByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) + GetAllLockedByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) // ApplyDeltas atomically applies a set of deltas. Either every delta is // applied or none are. Deltas are applied in SortDeltas order. @@ -94,9 +105,19 @@ type Store interface { // ErrInsufficientBalance is returned when a debit exceeds the balance. // ErrBalanceChanged is returned when a drain or close doesn't match the // balance. ErrAccountClosed is returned when a credit, drain or close - // targets a closed account. + // targets a closed account. ErrAccountUnlocked is returned when any + // delta targets an unlocked account, whose record is no longer + // maintained. ApplyDeltas(ctx context.Context, deltas ...*Delta) error + // MarkAsUnlocked marks an account's timelock vault as unlocked, which is + // one-way and idempotent. It is called in the same transaction that + // commits the timelock record's transition out of the locked state, so + // the flag cannot disagree with the timelock record it mirrors. + // + // ErrRecordNotFound is returned if no record exists. + MarkAsUnlocked(ctx context.Context, tokenAccount string) error + // Backfill locks a record that is not yet backfilled, calls fn to compute // its full historical balance, and overwrites the record with the result, // marking it as backfilled. Deltas recorded before the backfill are diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 4944bd3..85195fd 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -16,12 +16,13 @@ import ( func RunTests(t *testing.T, s balance.Store, teardown func()) { for _, tf := range []func(t *testing.T, s balance.Store){ testRecordHappyPath, - testGetAllByMint, + testGetAllLockedByMint, testApplyDeltasBackfilled, testApplyDeltasNotBackfilled, testApplyDeltasAtomicity, testApplyDeltasConcurrency, testBackfill, + testMarkAsUnlocked, testCachedBalanceVersionHappyPath, testClosedAccountHappyPath, testExternalCheckpointHappyPath, @@ -57,6 +58,7 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { Quarks: 100, UsdCostBasis: 200, IsOpen: true, + IsLocked: true, IsBackfilled: true, } cloned := expected.Clone() @@ -76,12 +78,14 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { OwnerAccount: "owner_1", MintAccount: "mint_2", IsOpen: true, + IsLocked: true, })) require.NoError(t, s.Create(ctx, &balance.Record{ TokenAccount: "token_account_3", OwnerAccount: "owner_2", MintAccount: "mint_1", IsOpen: true, + IsLocked: true, })) batch, err = s.GetBatch(ctx, "token_account_1", "token_account_3", "token_account_4") @@ -114,11 +118,11 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { }) } -func testGetAllByMint(t *testing.T, s balance.Store) { - t.Run("testGetAllByMint", func(t *testing.T) { +func testGetAllLockedByMint(t *testing.T, s balance.Store) { + t.Run("testGetAllLockedByMint", func(t *testing.T) { ctx := context.Background() - _, err := s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) + _, err := s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) assert.Equal(t, balance.ErrRecordNotFound, err) for i := range 5 { @@ -128,6 +132,7 @@ func testGetAllByMint(t *testing.T, s balance.Store) { MintAccount: "mint_1", Quarks: int64(i * 10), IsOpen: true, + IsLocked: true, IsBackfilled: true, })) } @@ -137,45 +142,59 @@ func testGetAllByMint(t *testing.T, s balance.Store) { MintAccount: "mint_2", Quarks: 1000, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) - records, err := s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) + records, err := s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) require.NoError(t, err) require.Len(t, records, 5) for i, record := range records { assert.EqualValues(t, i+1, record.Id) } - records, err = s.GetAllByMint(ctx, "mint_1", 20, query.EmptyCursor, 10, query.Ascending) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 20, query.EmptyCursor, 10, query.Ascending) require.NoError(t, err) require.Len(t, records, 3) assert.EqualValues(t, 20, records[0].Quarks) - records, err = s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Ascending) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Ascending) require.NoError(t, err) require.Len(t, records, 2) assert.EqualValues(t, 1, records[0].Id) assert.EqualValues(t, 2, records[1].Id) - records, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(2), 2, query.Ascending) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.ToCursor(2), 2, query.Ascending) require.NoError(t, err) require.Len(t, records, 2) assert.EqualValues(t, 3, records[0].Id) assert.EqualValues(t, 4, records[1].Id) - records, err = s.GetAllByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Descending) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 2, query.Descending) require.NoError(t, err) require.Len(t, records, 2) assert.EqualValues(t, 5, records[0].Id) assert.EqualValues(t, 4, records[1].Id) - records, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(4), 10, query.Descending) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.ToCursor(4), 10, query.Descending) require.NoError(t, err) require.Len(t, records, 3) assert.EqualValues(t, 3, records[0].Id) - _, err = s.GetAllByMint(ctx, "mint_1", 0, query.ToCursor(5), 10, query.Ascending) + // Unlocked records are excluded, since their balances are stale + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_unlocked", + OwnerAccount: "owner", + MintAccount: "mint_1", + Quarks: 1000, + IsOpen: true, + IsBackfilled: true, + })) + records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) + require.NoError(t, err) + require.Len(t, records, 5) + + _, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.ToCursor(5), 10, query.Ascending) assert.Equal(t, balance.ErrRecordNotFound, err) }) } @@ -189,6 +208,7 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { OwnerAccount: "owner", MintAccount: "mint", IsOpen: true, + IsLocked: true, IsBackfilled: true, })) @@ -243,6 +263,7 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { OwnerAccount: "owner", MintAccount: "mint", IsOpen: true, + IsLocked: true, IsBackfilled: true, })) require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaClose})) @@ -259,6 +280,7 @@ func testApplyDeltasNotBackfilled(t *testing.T, s balance.Store) { OwnerAccount: "owner", MintAccount: "mint", IsOpen: true, + IsLocked: true, })) // No predicates are enforced, and the balance can go negative @@ -292,6 +314,7 @@ func testApplyDeltasAtomicity(t *testing.T, s balance.Store) { MintAccount: "mint", Quarks: 100, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) } @@ -364,6 +387,7 @@ func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { MintAccount: "mint", Quarks: initialBalance, IsOpen: true, + IsLocked: true, IsBackfilled: true, })) require.NoError(t, s.Create(ctx, &balance.Record{ @@ -371,6 +395,7 @@ func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { OwnerAccount: "owner_2", MintAccount: "mint", IsOpen: true, + IsLocked: true, IsBackfilled: true, })) @@ -457,7 +482,7 @@ func testBackfill(t *testing.T, s balance.Store) { fn := func(quarks, usdCostBasis int64, isOpen bool) balance.BackfillFunc { return func(ctx context.Context) (*balance.BackfillResult, error) { - return &balance.BackfillResult{Quarks: quarks, UsdCostBasis: usdCostBasis, IsOpen: isOpen}, nil + return &balance.BackfillResult{Quarks: quarks, UsdCostBasis: usdCostBasis, IsOpen: isOpen, IsLocked: true}, nil } } @@ -468,6 +493,7 @@ func testBackfill(t *testing.T, s balance.Store) { OwnerAccount: "owner", MintAccount: "mint", IsOpen: true, + IsLocked: true, })) // Deltas recorded before the backfill are discarded by it @@ -494,6 +520,7 @@ func testBackfill(t *testing.T, s balance.Store) { record, err = s.Get(ctx, "token_account_1") require.NoError(t, err) assert.True(t, record.IsBackfilled) + assert.True(t, record.IsLocked) assertBalance(t, s, "token_account_1", 500, 250, true) // Predicates are enforced from now on @@ -515,6 +542,7 @@ func testBackfill(t *testing.T, s balance.Store) { OwnerAccount: "owner", MintAccount: "mint", IsOpen: true, + IsLocked: true, })) require.NoError(t, s.Backfill(ctx, "token_account_2", fn(0, 0, false))) assertBalance(t, s, "token_account_2", 0, 0, false) @@ -522,6 +550,57 @@ func testBackfill(t *testing.T, s balance.Store) { }) } +func testMarkAsUnlocked(t *testing.T, s balance.Store) { + t.Run("testMarkAsUnlocked", func(t *testing.T) { + ctx := context.Background() + + assert.Equal(t, balance.ErrRecordNotFound, s.MarkAsUnlocked(ctx, "token_account_1")) + + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_1", + OwnerAccount: "owner", + MintAccount: "mint", + Quarks: 100, + IsOpen: true, + IsLocked: true, + IsBackfilled: true, + })) + + require.NoError(t, s.MarkAsUnlocked(ctx, "token_account_1")) + record, err := s.Get(ctx, "token_account_1") + require.NoError(t, err) + assert.False(t, record.IsLocked) + assert.EqualValues(t, 100, record.Quarks) + assert.True(t, record.IsOpen) + + // Unlocking is idempotent + require.NoError(t, s.MarkAsUnlocked(ctx, "token_account_1")) + record, err = s.Get(ctx, "token_account_1") + require.NoError(t, err) + assert.False(t, record.IsLocked) + + // An unlocked record is no longer maintained: nothing may enter or + // leave it + assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1})) + assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) + assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 100})) + assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) + assertBalance(t, s, "token_account_1", 100, 0, true) + + // A record that is not backfilled accumulates without predicates, + // unlocked included; the backfill computes the truth for it + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_2", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + IsLocked: false, + })) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 5})) + assertBalance(t, s, "token_account_2", 5, 0, true) + }) +} + func testCachedBalanceVersionHappyPath(t *testing.T, s balance.Store) { t.Run("testCachedBalanceVersionHappyPath", func(t *testing.T) { ctx := context.Background() @@ -627,6 +706,7 @@ func assertEquivalentRecords(t *testing.T, obj1, obj2 *balance.Record) { assert.Equal(t, obj1.OwnerAccount, obj2.OwnerAccount) assert.Equal(t, obj1.MintAccount, obj2.MintAccount) assert.Equal(t, obj1.Quarks, obj2.Quarks) + assert.Equal(t, obj1.IsLocked, obj2.IsLocked) assert.Equal(t, obj1.UsdCostBasis, obj2.UsdCostBasis) assert.Equal(t, obj1.IsOpen, obj2.IsOpen) assert.Equal(t, obj1.IsBackfilled, obj2.IsBackfilled) diff --git a/ocp/data/internal.go b/ocp/data/internal.go index a3111f3..4dc7b75 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -126,8 +126,9 @@ type DatabaseData interface { GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) - GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) + GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error + MarkBalanceAsUnlocked(ctx context.Context, tokenAccount string) error BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error GetCachedBalanceVersion(ctx context.Context, account string) (uint64, error) AdvanceCachedBalanceVersion(ctx context.Context, account string, currentVersion uint64) error @@ -483,8 +484,11 @@ func (dp *DatabaseProvider) GetAllBalancesByOwner(ctx context.Context, owner str func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { return dp.balance.GetAllByOwnerAndMint(ctx, owner, mint) } -func (dp *DatabaseProvider) GetAllBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { - return dp.balance.GetAllByMint(ctx, mint, minQuarks, cursor, limit, direction) +func (dp *DatabaseProvider) GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { + return dp.balance.GetAllLockedByMint(ctx, mint, minQuarks, cursor, limit, direction) +} +func (dp *DatabaseProvider) MarkBalanceAsUnlocked(ctx context.Context, tokenAccount string) error { + return dp.balance.MarkAsUnlocked(ctx, tokenAccount) } func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error { return dp.balance.ApplyDeltas(ctx, deltas...) diff --git a/ocp/worker/geyser/timelock.go b/ocp/worker/geyser/timelock.go index 9c2b3b0..66ec7ff 100644 --- a/ocp/worker/geyser/timelock.go +++ b/ocp/worker/geyser/timelock.go @@ -2,6 +2,7 @@ package geyser import ( "context" + "database/sql" "time" ocp_data "github.com/code-payments/ocp-server/ocp/data" @@ -29,5 +30,11 @@ func updateTimelockAccountRecord(ctx context.Context, data ocp_data.Provider, ti timelockRecord.UnlockAt = pointer.Uint64(uint64(unlockState.UnlockAt)) timelockRecord.Block = slot timelockRecord.LastUpdatedAt = time.Now() - return data.SaveTimelock(ctx, timelockRecord) + return data.ExecuteInTx(ctx, sql.LevelDefault, func(ctx context.Context) error { + err := data.SaveTimelock(ctx, timelockRecord) + if err != nil { + return err + } + return data.MarkBalanceAsUnlocked(ctx, timelockRecord.VaultAddress) + }) } From 2bf207192771d639e059cf7f1688bbffe3baccef Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 14:03:19 -0400 Subject: [PATCH 15/26] Remove intentBalanceLock --- ocp/rpc/transaction/intent_handler.go | 9 --------- 1 file changed, 9 deletions(-) diff --git a/ocp/rpc/transaction/intent_handler.go b/ocp/rpc/transaction/intent_handler.go index 6b6d835..25d42c2 100644 --- a/ocp/rpc/transaction/intent_handler.go +++ b/ocp/rpc/transaction/intent_handler.go @@ -34,15 +34,6 @@ import ( "github.com/code-payments/ocp-server/solana" ) -type intentBalanceLock struct { - // The account that's being locked - Account *common.Account - - // The function executed on intent DB commit that is guaranteed to prevent - // race conditions against invalid balance updates - CommitFn func(ctx context.Context, data ocp_data.Provider) error -} - // CreateIntentHandler is an interface for handling new intent creations type CreateIntentHandler interface { // PopulateMetadata adds intent metadata to the provided intent record From 095a2902be7050db6243be616256601a84406490 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 14:03:57 -0400 Subject: [PATCH 16/26] Fix comment --- ocp/rpc/transaction/intent_handler.go | 1 - 1 file changed, 1 deletion(-) diff --git a/ocp/rpc/transaction/intent_handler.go b/ocp/rpc/transaction/intent_handler.go index 25d42c2..4585d38 100644 --- a/ocp/rpc/transaction/intent_handler.go +++ b/ocp/rpc/transaction/intent_handler.go @@ -61,7 +61,6 @@ type CreateIntentHandler interface { // mass races are worth locking, like a gift card being claimed from many // devices at once. Race resistance itself comes from the balance ledger's // predicates in the intent transaction. - // against invalid balance updates that would result in intent fulfillment failure GetAccountsToLock(ctx context.Context, intentRecord *intent.Record, metadata *transactionpb.Metadata) ([]*common.Account, error) // AllowCreation determines whether the new intent creation should be allowed. From ae285817b057c994eedfb04c420403d2e68bafd4 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 14:11:29 -0400 Subject: [PATCH 17/26] Update testGetAllLockedByMint --- ocp/data/balance/tests/tests.go | 1 + 1 file changed, 1 insertion(+) diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index ac01c3e..b70d3c1 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -203,6 +203,7 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { MintAccount: "mint_1", Quarks: 1000, IsOpen: true, + IsLocked: true, IsBackfilled: false, })) From 928a0c2f4d62a0be862fc84c9ea178e3df6e9086 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 09:20:33 -0400 Subject: [PATCH 18/26] Unlocked gift card accounts are now cleaned up vs auto-returned --- ocp/worker/account/gift_card.go | 20 ++++++++++++++++++++ ocp/worker/account/gift_card_test.go | 16 ++++++++++++++++ ocp/worker/account/testutil.go | 13 +++++++++++++ 3 files changed, 49 insertions(+) diff --git a/ocp/worker/account/gift_card.go b/ocp/worker/account/gift_card.go index e2cbb6d..ac95600 100644 --- a/ocp/worker/account/gift_card.go +++ b/ocp/worker/account/gift_card.go @@ -116,6 +116,26 @@ func (p *runtime) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn return err } + timelockRecord, err := p.data.GetTimelockByVault(ctx, giftCardVaultAccount.PublicKey().ToBase58()) + if err != nil { + log.With(zap.Error(err)).Warn("failure getting timelock record") + return err + } + if !common.IsManagedByCode(ctx, timelockRecord) { + log.Debug("gift card is no longer managed by code and will be removed from worker queue") + + // The vault has unlocked, so the funds can no longer be moved by an + // auto-return withdraw. Clean it up like a claimed gift card, rather + // than scheduling a fulfillment that can't be executed. + err = InitiateProcessToCleanupGiftCardAutoReturn(ctx, p.data, giftCardVaultAccount) + if err != nil { + log.With(zap.Error(err)).Warn("failure cleaning up auto-return action") + return err + } + + return MarkAutoReturnCheckComplete(ctx, p.data, accountInfoRecord) + } + // Expiration window hasn't been met // // Note: Without distributed locks, we assume SubmitIntent uses expiry - delta diff --git a/ocp/worker/account/gift_card_test.go b/ocp/worker/account/gift_card_test.go index c6f8d19..7376712 100644 --- a/ocp/worker/account/gift_card_test.go +++ b/ocp/worker/account/gift_card_test.go @@ -63,6 +63,22 @@ func TestGiftCardAutoReturn_AlreadyClaimed(t *testing.T) { env.assertGiftCardNotAutoReturned(t, giftCard2, true) } +func TestGiftCardAutoReturn_VaultUnlocked(t *testing.T) { + env := setup(t) + + giftCard1 := env.generateRandomGiftCard(t, time.Now()) + env.simulateGiftCardVaultBeingUnlocked(t, giftCard1) + + giftCard2 := env.generateRandomGiftCard(t, time.Now().Add(-GiftCardExpiry-24*time.Hour)) + env.simulateGiftCardVaultBeingUnlocked(t, giftCard2) + + require.NoError(t, env.runtime.maybeInitiateGiftCardAutoReturn(env.ctx, giftCard1.accountInfoRecord)) + env.assertGiftCardNotAutoReturned(t, giftCard1, true) + + require.NoError(t, env.runtime.maybeInitiateGiftCardAutoReturn(env.ctx, giftCard2.accountInfoRecord)) + env.assertGiftCardNotAutoReturned(t, giftCard2, true) +} + func TestGiftCardAutoReturn_IntentId(t *testing.T) { intentId1 := testutil.NewRandomAccount(t).PublicKey().ToBase58() intentId2 := testutil.NewRandomAccount(t).PublicKey().ToBase58() diff --git a/ocp/worker/account/testutil.go b/ocp/worker/account/testutil.go index af6aa33..18827dc 100644 --- a/ocp/worker/account/testutil.go +++ b/ocp/worker/account/testutil.go @@ -19,7 +19,9 @@ import ( "github.com/code-payments/ocp-server/ocp/data/action" "github.com/code-payments/ocp-server/ocp/data/fulfillment" "github.com/code-payments/ocp-server/ocp/data/intent" + "github.com/code-payments/ocp-server/ocp/data/timelock" "github.com/code-payments/ocp-server/pointer" + timelock_token "github.com/code-payments/ocp-server/solana/timelock/v1" "github.com/code-payments/ocp-server/testutil" ) @@ -31,6 +33,7 @@ type testEnv struct { type testGiftCard struct { accountInfoRecord *account.Record + timelockRecord *timelock.Record issuedIntentRecord *intent.Record claimedActionRecord *action.Record @@ -74,6 +77,9 @@ func (e *testEnv) generateRandomGiftCard(t *testing.T, creationTs time.Time) *te } require.NoError(t, e.data.CreateAccountInfo(e.ctx, accountInfoRecord)) + timelockRecord := timelockAccounts.ToDBRecord() + require.NoError(t, e.data.SaveTimelock(e.ctx, timelockRecord)) + intentRecord := &intent.Record{ IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), IntentType: intent.SendPublicPayment, @@ -148,6 +154,7 @@ func (e *testEnv) generateRandomGiftCard(t *testing.T, creationTs time.Time) *te return &testGiftCard{ accountInfoRecord: accountInfoRecord, + timelockRecord: timelockRecord, issuedIntentRecord: intentRecord, autoReturnActionRecord: autoReturnActionRecord, @@ -174,6 +181,12 @@ func (e *testEnv) simulateGiftCardBeingClaimed(t *testing.T, giftCard *testGiftC require.NoError(t, e.data.PutAllActions(e.ctx, giftCard.claimedActionRecord)) } +func (e *testEnv) simulateGiftCardVaultBeingUnlocked(t *testing.T, giftCard *testGiftCard) { + giftCard.timelockRecord.VaultState = timelock_token.StateUnlocked + giftCard.timelockRecord.Block += 1 + require.NoError(t, e.data.SaveTimelock(e.ctx, giftCard.timelockRecord)) +} + func (e *testEnv) assertGiftCardAutoReturned(t *testing.T, giftCard *testGiftCard) { accountInfoRecord, err := e.data.GetAccountInfoByTokenAddress(e.ctx, giftCard.accountInfoRecord.TokenAccount) require.NoError(t, err) From 56dd2518a1603206a60b1903f31bf05b8f6db2b9 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 09:26:16 -0400 Subject: [PATCH 19/26] VoidGiftCard now handles unlocked gift card accounts --- ocp/rpc/transaction/gift_card.go | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ocp/rpc/transaction/gift_card.go b/ocp/rpc/transaction/gift_card.go index bc1fbe4..21d8d72 100644 --- a/ocp/rpc/transaction/gift_card.go +++ b/ocp/rpc/transaction/gift_card.go @@ -120,6 +120,17 @@ func (s *transactionServer) VoidGiftCard(ctx context.Context, req *transactionpb return nil, status.Error(codes.Internal, "") } + timelockRecord, err := s.data.GetTimelockByVault(ctx, giftCardVault.PublicKey().ToBase58()) + if err != nil { + log.With(zap.Error(err)).Warn("failure getting timelock record") + return nil, status.Error(codes.Internal, "") + } + if !common.IsManagedByCode(ctx, timelockRecord) { + return &transactionpb.VoidGiftCardResponse{ + Result: transactionpb.VoidGiftCardResponse_CLAIMED_BY_OTHER_USER, + }, nil + } + err = account_worker.InitiateProcessToAutoReturnGiftCard(ctx, s.data, giftCardVault, true, globalBalanceLock) if err != nil { log.With(zap.Error(err)).Warn("failure scheduling auto-return action") From d095032797bfdf3807bfc7237f82383ef8bd9f9e Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 10:30:05 -0400 Subject: [PATCH 20/26] Allow credits to unlocked accounts and introduce a cost basis adjustment delta --- ocp/balance/delta.go | 9 +++--- ocp/balance/delta_test.go | 4 +-- ocp/balance/ledger.go | 10 ++++--- ocp/balance/ledger_test.go | 14 ++++++---- ocp/data/balance/memory/store.go | 14 ++++++++-- ocp/data/balance/postgres/model.go | 17 ++++++++++-- ocp/data/balance/record.go | 44 ++++++++++++++++++++++-------- ocp/data/balance/record_test.go | 3 ++ ocp/data/balance/store.go | 15 ++++++---- ocp/data/balance/tests/tests.go | 38 ++++++++++++++++++++++++-- 10 files changed, 127 insertions(+), 41 deletions(-) diff --git a/ocp/balance/delta.go b/ocp/balance/delta.go index 9007224..e3fda2d 100644 --- a/ocp/balance/delta.go +++ b/ocp/balance/delta.go @@ -188,16 +188,17 @@ func DeltasForSwapSellReconciliation(previous, updated *intent.Record, actionRec return nil, fmt.Errorf("%w: swap sell has no funding action", ErrUnsupportedBalanceChange) } - adjustment := balance.UsdCostBasisFromFloat(updated.SendPublicPaymentMetadata.UsdMarketValue) - balance.UsdCostBasisFromFloat(previous.SendPublicPaymentMetadata.UsdMarketValue) + // The funding payment already removed the basis it was committed with, so + // the correction added back is what it overcharged: negative when the sell + // realized more than estimated, positive when it realized less. + adjustment := balance.UsdCostBasisFromFloat(previous.SendPublicPaymentMetadata.UsdMarketValue) - balance.UsdCostBasisFromFloat(updated.SendPublicPaymentMetadata.UsdMarketValue) if adjustment == 0 { return nil, nil } - // A debit subtracts the signed basis, so a higher realized value removes - // more basis from the source and a lower one gives some back return []*balance.Delta{{ TokenAccount: funding.Source, - Kind: balance.DeltaDebit, + Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: adjustment, }}, nil } diff --git a/ocp/balance/delta_test.go b/ocp/balance/delta_test.go index cf30f7c..7e7bcd1 100644 --- a/ocp/balance/delta_test.go +++ b/ocp/balance/delta_test.go @@ -343,7 +343,7 @@ func TestDeltasForSwapSellReconciliation(t *testing.T) { deltas, err := DeltasForSwapSellReconciliation(previous, updated, actionRecords) require.NoError(t, err) assert.Equal(t, []*balance.Delta{ - {TokenAccount: "source", Kind: balance.DeltaDebit, UsdCostBasis: 250_000}, + {TokenAccount: "source", Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: -250_000}, }, deltas) // Realized less than estimated: basis is returned to the source @@ -351,7 +351,7 @@ func TestDeltasForSwapSellReconciliation(t *testing.T) { deltas, err = DeltasForSwapSellReconciliation(previous, updated, actionRecords) require.NoError(t, err) assert.Equal(t, []*balance.Delta{ - {TokenAccount: "source", Kind: balance.DeltaDebit, UsdCostBasis: -250_000}, + {TokenAccount: "source", Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: 250_000}, }, deltas) // No change is a no-op diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index 653cd33..99f7fa0 100644 --- a/ocp/balance/ledger.go +++ b/ocp/balance/ledger.go @@ -31,10 +31,12 @@ func LedgerWritesEnabled(ctx context.Context) bool { // like an external wallet or the fee collector, are dropped, since delta // builders don't know which destinations OCP manages. Outgoing deltas from // an account the ledger doesn't track are ErrUntrackedAccount, since funds -// only ever leave accounts OCP manages. Any delta against a tracked account -// whose vault has unlocked fails loudly with balance.ErrAccountUnlocked: -// the record is no longer maintained, and a flow still moving funds through -// it is a bug to surface, never to paper over. +// only ever leave accounts OCP manages. An outgoing delta against a tracked +// account whose vault has unlocked fails loudly with +// balance.ErrAccountUnlocked: the record is no longer maintained, and a flow +// still taking funds out of it is a bug to surface, never to paper over. +// Credits are still applied, since an unlocked record is excluded from every +// read and turning one away only blocks the flow recording it. // // Any timelock account in the delta set that has no ledger record yet lazily // gets one that is not backfilled, so accounts that predate the ledger start diff --git a/ocp/balance/ledger_test.go b/ocp/balance/ledger_test.go index 09c90aa..5f72e57 100644 --- a/ocp/balance/ledger_test.go +++ b/ocp/balance/ledger_test.go @@ -120,16 +120,18 @@ func TestApplyDeltasInTx_UnlockedAccount(t *testing.T) { require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaCredit, Quarks: 100})) require.NoError(t, data.MarkBalanceAsUnlocked(ctx, unlocked.TokenAccount)) - // Any delta against an unlocked account fails loudly, so a flow still - // moving funds through it surfaces as a DB error - err := ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaCredit, Quarks: 1}) - assert.Equal(t, balance.ErrAccountUnlocked, err) - err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaDebit, Quarks: 1}) + // An outgoing delta against an unlocked account fails loudly, so a flow + // still taking funds out of it surfaces as a DB error + err := ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaDebit, Quarks: 1}) assert.Equal(t, balance.ErrAccountUnlocked, err) + // A credit still applies, so a flow recording funds that have already + // moved isn't blocked by an unlock it raced against + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: unlocked.TokenAccount, Kind: balance.DeltaCredit, Quarks: 1})) + record, err := data.GetBalance(ctx, unlocked.TokenAccount) require.NoError(t, err) - assert.EqualValues(t, 100, record.Quarks) + assert.EqualValues(t, 101, record.Quarks) assert.False(t, record.IsLocked) } diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 549696c..f31e676 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -198,8 +198,16 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { func applyDelta(item *balance.Record, delta *balance.Delta) error { enforce := item.IsBackfilled - if enforce && !item.IsLocked { - return balance.ErrAccountUnlocked + // A credit doesn't require the vault to be locked, since an unlocked + // record is excluded from every read anyway and turning the credit away + // only blocks the flow recording it. A cost basis adjustment carries no + // predicate at all. + switch delta.Kind { + case balance.DeltaCredit, balance.DeltaAdjustUsdCostBasis: + default: + if enforce && !item.IsLocked { + return balance.ErrAccountUnlocked + } } switch delta.Kind { @@ -240,6 +248,8 @@ func applyDelta(item *balance.Record, delta *balance.Delta) error { } } item.IsOpen = false + case balance.DeltaAdjustUsdCostBasis: + item.UsdCostBasis += delta.UsdCostBasis } return nil } diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index 30dae31..dbe59dc 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -212,7 +212,7 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er case balance.DeltaCredit: query = `UPDATE ` + tableName + ` SET quarks = quarks + $2, usd_cost_basis = usd_cost_basis + $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked))` + WHERE token_account = $1 AND (NOT is_backfilled OR is_open)` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDebit: query = `UPDATE ` + tableName + ` @@ -232,6 +232,13 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er SET is_open = FALSE, updated_at = $2 WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks = 0))` args = []any{delta.TokenAccount, time.Now().UTC()} + case balance.DeltaAdjustUsdCostBasis: + // No predicate: no quarks move, so nothing the other kinds + // guard against can happen here + query = `UPDATE ` + tableName + ` + SET usd_cost_basis = usd_cost_basis + $2, updated_at = $3 + WHERE token_account = $1` + args = []any{delta.TokenAccount, delta.UsdCostBasis, time.Now().UTC()} default: return fmt.Errorf("unsupported delta kind: %s", delta.Kind) } @@ -263,12 +270,16 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er } func classifyFailedDelta(delta *balance.Delta, current *balance.Record) error { + // A credit doesn't require the vault to be locked, so a closed account is + // the only thing that turns one away. + if delta.Kind == balance.DeltaCredit { + return balance.ErrAccountClosed + } + if !current.IsLocked { return balance.ErrAccountUnlocked } switch delta.Kind { - case balance.DeltaCredit: - return balance.ErrAccountClosed case balance.DeltaDebit: return balance.ErrInsufficientBalance case balance.DeltaDrain, balance.DeltaClose: diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index 769f095..c82f580 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -46,9 +46,11 @@ type Record struct { // IsLocked indicates the timelock vault is still locked, so the account // is managed by OCP and every balance change flows through the ledger. - // Once a vault unlocks, funds can move on chain without an intent, so - // the record's values are the last managed state and must not be - // trusted or aggregated. Unlocking is one-way. + // Once a vault unlocks, funds can move on chain without an intent, so the + // record's values must not be trusted or aggregated: nothing may leave + // the account through the ledger, while credits keep being recorded + // against a balance that no longer reflects the chain. Unlocking is + // one-way. IsLocked bool // IsBackfilled indicates the record reflects the full history of the @@ -131,6 +133,14 @@ const ( // DeltaClose closes an account with a zero balance. DeltaClose + + // DeltaAdjustUsdCostBasis adds a signed correction to an account's USD + // cost basis without moving quarks. It carries no predicate: every + // predicate on the other kinds protects a quark invariant, and none of + // them apply when no quarks move. A correction is only ever issued for a + // period the ledger was already tracking, so refusing one would leave the + // basis wrong rather than protect anything. + DeltaAdjustUsdCostBasis ) // Delta is a single balance change to apply to a token account. @@ -138,13 +148,15 @@ type Delta struct { TokenAccount string Kind DeltaKind - // Quarks is the amount credited, debited or drained. Ignored for DeltaClose. + // Quarks is the amount credited, debited or drained. Ignored for + // DeltaClose, and must be zero for DeltaAdjustUsdCostBasis. Quarks uint64 // UsdCostBasis is added on credit and subtracted on debit. It is signed // so that a credit can also carry a downward reconciliation. Ignored for // DeltaDrain and DeltaClose on backfilled records, where the basis is - // zeroed along with the balance. + // zeroed along with the balance. For DeltaAdjustUsdCostBasis it is the + // signed correction, added as is. UsdCostBasis int64 } @@ -158,6 +170,13 @@ func (d *Delta) Validate() error { if d.Quarks == 0 && d.UsdCostBasis == 0 { return errors.New("delta is a no-op") } + case DeltaAdjustUsdCostBasis: + if d.Quarks != 0 { + return errors.New("cost basis adjustment cannot move quarks") + } + if d.UsdCostBasis == 0 { + return errors.New("delta is a no-op") + } case DeltaClose: default: return errors.New("invalid delta kind") @@ -179,11 +198,12 @@ func SortDeltas(deltas []*Delta) { } // MergeDeltas returns a copy of deltas in SortDeltas order with consecutive -// credits and debits to the same account combined into one. Applying one -// combined delta is equivalent to applying the parts in sequence, since -// both kinds are additive and their predicates are monotonic in the amount, -// but it touches the row once. Drains and closes are never merged, since an -// account can only legitimately be drained or closed once. +// credits, debits and cost basis adjustments to the same account combined +// into one. Applying one combined delta is equivalent to applying the parts +// in sequence, since those kinds are additive and their predicates are +// monotonic in the amount, but it touches the row once. Drains and closes are +// never merged, since an account can only legitimately be drained or closed +// once. func MergeDeltas(deltas []*Delta) []*Delta { sorted := make([]*Delta, len(deltas)) copy(sorted, deltas) @@ -193,7 +213,7 @@ func MergeDeltas(deltas []*Delta) []*Delta { for _, delta := range sorted { if len(merged) > 0 { last := merged[len(merged)-1] - if last.TokenAccount == delta.TokenAccount && last.Kind == delta.Kind && (delta.Kind == DeltaCredit || delta.Kind == DeltaDebit) { + if last.TokenAccount == delta.TokenAccount && last.Kind == delta.Kind && (delta.Kind == DeltaCredit || delta.Kind == DeltaDebit || delta.Kind == DeltaAdjustUsdCostBasis) { last.Quarks += delta.Quarks last.UsdCostBasis += delta.UsdCostBasis continue @@ -215,6 +235,8 @@ func (k DeltaKind) String() string { return "drain" case DeltaClose: return "close" + case DeltaAdjustUsdCostBasis: + return "adjust usd cost basis" } return "unknown" } diff --git a/ocp/data/balance/record_test.go b/ocp/data/balance/record_test.go index 87dc868..e0b7bbf 100644 --- a/ocp/data/balance/record_test.go +++ b/ocp/data/balance/record_test.go @@ -17,6 +17,8 @@ func TestMergeDeltas(t *testing.T) { {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, {TokenAccount: "d", Kind: DeltaClose}, {TokenAccount: "d", Kind: DeltaClose}, + {TokenAccount: "a", Kind: DeltaAdjustUsdCostBasis, UsdCostBasis: 5}, + {TokenAccount: "a", Kind: DeltaAdjustUsdCostBasis, UsdCostBasis: -8}, } original := make([]Delta, len(input)) for i, delta := range input { @@ -27,6 +29,7 @@ func TestMergeDeltas(t *testing.T) { assert.Equal(t, []*Delta{ {TokenAccount: "a", Kind: DeltaCredit, Quarks: 4, UsdCostBasis: 40}, + {TokenAccount: "a", Kind: DeltaAdjustUsdCostBasis, UsdCostBasis: -3}, {TokenAccount: "b", Kind: DeltaCredit, Quarks: 2, UsdCostBasis: 20}, {TokenAccount: "b", Kind: DeltaDebit, Quarks: 12, UsdCostBasis: -2}, {TokenAccount: "c", Kind: DeltaDrain, Quarks: 9, UsdCostBasis: 9}, diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index f721974..bf3d693 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -29,9 +29,11 @@ var ( ErrAccountClosed = errors.New("account open state is stale") - // ErrAccountUnlocked is returned when a delta targets an account whose - // timelock vault has unlocked. The ledger stops maintaining the record at - // unlock, so nothing may enter or leave it. + // ErrAccountUnlocked is returned when a delta other than a credit targets + // an account whose timelock vault has unlocked. The ledger stops + // maintaining the record at unlock, so nothing may leave it. Credits are + // still applied, since an unlocked record is excluded from every read and + // turning one away only blocks the flow recording it. ErrAccountUnlocked = errors.New("account is unlocked") ErrCheckpointNotFound = errors.New("checkpoint not found") @@ -105,9 +107,10 @@ type Store interface { // ErrInsufficientBalance is returned when a debit exceeds the balance. // ErrBalanceChanged is returned when a drain or close doesn't match the // balance. ErrAccountClosed is returned when a credit, drain or close - // targets a closed account. ErrAccountUnlocked is returned when any - // delta targets an unlocked account, whose record is no longer - // maintained. + // targets a closed account. ErrAccountUnlocked is returned when a delta + // other than a credit targets an unlocked account, whose record is no + // longer maintained. DeltaAdjustUsdCostBasis carries no predicate and + // only ever fails when the record is missing. ApplyDeltas(ctx context.Context, deltas ...*Delta) error // MarkAsUnlocked marks an account's timelock vault as unlocked, which is diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 85195fd..9823fd0 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -258,6 +258,18 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) assertBalance(t, s, "token_account_1", 0, 0, false) + // A cost basis adjustment carries no predicate, so it still applies to + // a closed account and never moves quarks + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: 7})) + assertBalance(t, s, "token_account_1", 0, 7, false) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: -20})) + assertBalance(t, s, "token_account_1", 0, -13, false) + + // It must not carry quarks, and a zero adjustment is a no-op + assert.Error(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaAdjustUsdCostBasis, Quarks: 1, UsdCostBasis: 1})) + assert.Error(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaAdjustUsdCostBasis})) + assertBalance(t, s, "token_account_1", 0, -13, false) + require.NoError(t, s.Create(ctx, &balance.Record{ TokenAccount: "token_account_2", OwnerAccount: "owner", @@ -579,14 +591,34 @@ func testMarkAsUnlocked(t *testing.T, s balance.Store) { require.NoError(t, err) assert.False(t, record.IsLocked) - // An unlocked record is no longer maintained: nothing may enter or - // leave it - assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1})) + // Nothing may leave an unlocked record, which is no longer maintained assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 1})) assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 100})) assert.Equal(t, balance.ErrAccountUnlocked, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) assertBalance(t, s, "token_account_1", 100, 0, true) + // Credits still apply, so a flow recording funds that have already + // moved is never blocked by an unlock it raced against + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 1, UsdCostBasis: 2})) + assertBalance(t, s, "token_account_1", 101, 2, true) + + // As do cost basis adjustments, which carry no predicate at all + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaAdjustUsdCostBasis, UsdCostBasis: -5})) + assertBalance(t, s, "token_account_1", 101, -3, true) + + // Unless the account is also closed + require.NoError(t, s.Create(ctx, &balance.Record{ + TokenAccount: "token_account_3", + OwnerAccount: "owner", + MintAccount: "mint", + IsOpen: true, + IsLocked: true, + IsBackfilled: true, + })) + require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_3", Kind: balance.DeltaClose})) + require.NoError(t, s.MarkAsUnlocked(ctx, "token_account_3")) + assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_3", Kind: balance.DeltaCredit, Quarks: 1})) + // A record that is not backfilled accumulates without predicates, // unlocked included; the backfill computes the truth for it require.NoError(t, s.Create(ctx, &balance.Record{ From ff2cb2cb1bd1bdd962b121ceb331f9676714d6eb Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 13:42:45 -0400 Subject: [PATCH 21/26] Remove legacy balance calculations and ledger configs --- ocp/balance/calculator.go | 386 +++---------- ocp/balance/calculator_test.go | 780 +++++--------------------- ocp/balance/config.go | 25 - ocp/balance/ledger.go | 23 +- ocp/balance/ledger_test.go | 49 -- ocp/rpc/account/server_test.go | 25 +- ocp/rpc/balance/server_test.go | 23 +- ocp/rpc/transaction/intent.go | 20 +- ocp/worker/account/gift_card.go | 16 +- ocp/worker/currency/holder/runtime.go | 66 +-- ocp/worker/geyser/external_deposit.go | 16 +- ocp/worker/swap/util.go | 46 +- 12 files changed, 271 insertions(+), 1204 deletions(-) delete mode 100644 ocp/balance/config.go diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 73343cc..f12dbdd 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -2,17 +2,13 @@ package balance import ( "context" - "math" "time" "github.com/pkg/errors" - commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" - "github.com/code-payments/ocp-server/metrics" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" - "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/data/timelock" "github.com/code-payments/ocp-server/solana" @@ -31,8 +27,8 @@ const ( ) var ( - // ErrNegativeBalance indicates that a balance calculation resulted in a - // negative value. + // ErrNegativeBalance indicates that a ledger record holds a negative + // value, which can only happen before it's backfilled. ErrNegativeBalance = errors.New("balance calculation resulted in negative value") // ErrNotManagedByCode indicates that an account is not owned by Code. @@ -40,46 +36,14 @@ var ( // the context of a balance. ErrNotManagedByCode = errors.New("explicitly not handling account not managed by code") - // ErrUnhandledAccount indicates that the balance calculator does not - // have strategies to handle the provided account. - ErrUnhandledAccount = errors.New("unhandled account") + // ErrIncompleteRecord indicates a ledger record doesn't yet reflect the + // full history of the account, so it isn't an authoritative balance. + // + // todo: Remove alongside balance.Record.IsBackfilled once the ledger + // backfill is retired. + ErrIncompleteRecord = errors.New("balance record is not backfilled") ) -// Calculator is a function that calculates a token account's balance -type Calculator func(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (uint64, error) - -type Strategy func(ctx context.Context, tokenAccount *common.Account, state *State) (*State, error) - -type State struct { - // We allow for negative balances in intermediary steps. This is to simplify - // coordination between strategies. In the end, the sum of all strategies must - // reflect an accurate picture of the balance, at which point we'll enforce this - // is positive. - current int64 -} - -// Calculate calculates a token account's balance using a starting point and a set -// of strategies. Each may be incomplete individually, but in total must form a -// complete balance calculation. -func Calculate(ctx context.Context, tokenAccount *common.Account, initialBalance uint64, strategies ...Strategy) (balance uint64, err error) { - balanceState := &State{ - current: int64(initialBalance), - } - - for _, strategy := range strategies { - balanceState, err = strategy(ctx, tokenAccount, balanceState) - if err != nil { - return 0, err - } - } - - if balanceState.current < 0 { - return 0, ErrNegativeBalance - } - - return uint64(balanceState.current), nil -} - // CalculateFromCache is the default and recommended strategy for reliably estimating // a token account's balance using cached values. // @@ -99,48 +63,26 @@ func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccoun return 0, err } - // The strategy uses cached values from the intents system. The account must - // be managed by Code in order to return accurate values. + // The balance ledger is only maintained for accounts managed by Code. The + // account must be managed in order to return accurate values. isManagedByCode := common.IsManagedByCode(ctx, timelockRecord) if !isManagedByCode { tracer.OnError(ErrNotManagedByCode) return 0, ErrNotManagedByCode } - // Prefer the materialized balance record, when the account has one that - // reflects its full history. - if enableLedgerReads.Get(ctx) { - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err == nil && balanceRecord.IsBackfilled { - quarks, err := quarksFromRecord(balanceRecord) - if err != nil { - tracer.OnError(err) - return 0, err - } - return quarks, nil - } else if err != nil && err != balance.ErrRecordNotFound { - tracer.OnError(err) - return 0, err - } - } - - // Otherwise, fall back to iterating over the account's history. - strategies := []Strategy{ - NetBalanceFromIntentActions(ctx, data), - FundingFromExternalDeposits(ctx, data), + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err != nil { + tracer.OnError(err) + return 0, err } - balance, err := Calculate( - ctx, - tokenAccount, - 0, - strategies..., - ) + quarks, err := quarksFromRecord(balanceRecord) if err != nil { tracer.OnError(err) - return 0, errors.Wrap(err, "error calculating token account balance") + return 0, err } - return balance, nil + return quarks, nil } // CalculateFromBlockchain is the default and recommended strategy for reliably @@ -203,74 +145,6 @@ func CalculateFromBlockchain(ctx context.Context, data ocp_data.Provider, tokenA return quarks, BlockchainSource, nil } -// NetBalanceFromIntentActions is a balance calculation strategy that incorporates -// the net balance by applying payment intents to the current balance. -func NetBalanceFromIntentActions(ctx context.Context, data ocp_data.Provider) Strategy { - return func(ctx context.Context, tokenAccount *common.Account, state *State) (*State, error) { - netBalance, err := data.GetNetBalanceFromActions(ctx, tokenAccount.PublicKey().ToBase58()) - if err != nil { - return nil, errors.Wrap(err, "error getting net balance from intent actions") - } - - state.current += netBalance - return state, nil - } -} - -// FundingFromExternalDeposits is a balance calculation strategy that adds funding -// from deposits from external accounts. -func FundingFromExternalDeposits(ctx context.Context, data ocp_data.Provider) Strategy { - return func(ctx context.Context, tokenAccount *common.Account, state *State) (*State, error) { - amount, err := data.GetTotalExternalDepositedAmountInQuarks(ctx, tokenAccount.PublicKey().ToBase58()) - if err != nil { - return nil, errors.Wrap(err, "error getting external deposit amount") - } - state.current += int64(amount) - - return state, nil - } -} - -// BatchCalculator is a functiona that calculates a batch of accounts' balances -type BatchCalculator func(ctx context.Context, data ocp_data.Provider, accountRecordsBatch []*common.AccountRecords) (map[string]uint64, error) - -type BatchStrategy func(ctx context.Context, tokenAccounts []string, state *BatchState) (*BatchState, error) - -type BatchState struct { - // We allow for negative balances in intermediary steps. This is to simplify - // coordination between strategies. In the end, the sum of all strategies must - // reflect an accurate picture of the balance, at which point we'll enforce this - // is positive. - current map[string]int64 -} - -// CalculateBatch calculates a set of token accounts' balance using a starting point -// and a set of strategies. Each may be incomplete individually, but in total must -// form a complete balance calculation. -func CalculateBatch(ctx context.Context, tokenAccounts []string, strategies ...BatchStrategy) (balanceByTokenAccount map[string]uint64, err error) { - balanceState := &BatchState{ - current: make(map[string]int64), - } - - for _, strategy := range strategies { - balanceState, err = strategy(ctx, tokenAccounts, balanceState) - if err != nil { - return nil, err - } - } - - res := make(map[string]uint64) - for tokenAccount, balance := range balanceState.current { - if balance < 0 { - return nil, ErrNegativeBalance - } - - res[tokenAccount] = uint64(balance) - } - - return res, nil -} - // BatchCalculateFromCacheWithAccountRecords is the default and recommended batch strategy // or reliably estimating a set of token accounts' balance when common.AccountRecords are // available. @@ -337,10 +211,10 @@ func BatchCalculateFromCacheWithTokenAccounts(ctx context.Context, data ocp_data } func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provider, timelockRecords []*timelock.Record) (map[string]uint64, error) { - var tokenAccounts []string + tokenAccounts := make([]string, 0, len(timelockRecords)) for _, timelockRecord := range timelockRecords { - // The strategy uses cached values from the intents system. The account must - // be managed by Code in order to return accurate values. + // The balance ledger is only maintained for accounts managed by Code. + // The account must be managed in order to return accurate values. isManagedByCode := common.IsManagedByCode(ctx, timelockRecord) if !isManagedByCode { return nil, ErrNotManagedByCode @@ -349,24 +223,16 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide tokenAccounts = append(tokenAccounts, timelockRecord.VaultAddress) } - // Prefer materialized balance records, and only iterate over history for - // accounts that don't yet have a fully backfilled one. - balanceRecords := make(map[string]*balance.Record) - if enableLedgerReads.Get(ctx) { - var err error - balanceRecords, err = data.GetBalanceBatch(ctx, tokenAccounts...) - if err != nil { - return nil, err - } + balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccounts...) + if err != nil { + return nil, err } res := make(map[string]uint64, len(tokenAccounts)) - var remaining []string for _, tokenAccount := range tokenAccounts { balanceRecord, ok := balanceRecords[tokenAccount] - if !ok || !balanceRecord.IsBackfilled { - remaining = append(remaining, tokenAccount) - continue + if !ok { + return nil, balance.ErrRecordNotFound } quarks, err := quarksFromRecord(balanceRecord) @@ -375,23 +241,6 @@ func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provide } res[tokenAccount] = quarks } - - if len(remaining) == 0 { - return res, nil - } - - legacyRes, err := CalculateBatch( - ctx, - remaining, - NetBalanceFromIntentActionsBatch(ctx, data), - FundingFromExternalDepositsBatch(ctx, data), - ) - if err != nil { - return nil, err - } - for tokenAccount, quarks := range legacyRes { - res[tokenAccount] = quarks - } return res, nil } @@ -403,11 +252,9 @@ type BalanceWithUsdCostBasis struct { } // 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. +// bases for a set of account records. Both values for an account come from +// the same balance record read, so they are guaranteed consistent with each +// other. // // Note: Use this method when calculating balances for accounts that are managed by // Code (ie. Timelock account) and operate within the L2 system. @@ -415,7 +262,7 @@ func BatchCalculateWithUsdCostBasisFromCache(ctx context.Context, data ocp_data. tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateWithUsdCostBasisFromCache") defer tracer.End() - var tokenAccounts []string + tokenAccounts := make([]string, 0, len(accountRecordsBatch)) for _, accountRecords := range accountRecordsBatch { if !accountRecords.IsTimelock() || !accountRecords.IsManagedByCode(ctx) { tracer.OnError(ErrNotManagedByCode) @@ -424,23 +271,18 @@ func BatchCalculateWithUsdCostBasisFromCache(ctx context.Context, data ocp_data. 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 - } + 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 + if !ok { + tracer.OnError(balance.ErrRecordNotFound) + return nil, balance.ErrRecordNotFound } quarks, err := quarksFromRecord(balanceRecord) @@ -453,68 +295,33 @@ func BatchCalculateWithUsdCostBasisFromCache(ctx context.Context, data ocp_data. 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. +// in balance.UsdQuarksPerUnit, from its ledger record. // -// Note: Unlike quark balances, a cost basis for an account not managed by Code -// is still meaningful, so no timelock check is performed and such accounts fall -// back to the legacy calculation. The materialized record is the exception: -// once a vault unlocks it holds the last managed state rather than a live cost -// basis, so reading one returns ErrNotManagedByCode. +// Note: Unlike the quark balance calculators, no timelock check is performed, +// since the ledger record carries the lock state itself. Once a vault unlocks +// its record holds the last managed state rather than a live cost basis, so +// reading one returns ErrNotManagedByCode. An account the ledger doesn't +// track returns balance.ErrRecordNotFound. func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (int64, error) { tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "CalculateUsdCostBasisFromCache") tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) defer tracer.End() - if enableLedgerReads.Get(ctx) { - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err == nil && balanceRecord.IsBackfilled { - if !balanceRecord.IsLocked { - tracer.OnError(ErrNotManagedByCode) - return 0, ErrNotManagedByCode - } - return balanceRecord.UsdCostBasis, nil - } else if err != nil && err != balance.ErrRecordNotFound { - tracer.OnError(err) - return 0, err - } + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err != nil { + tracer.OnError(err) + return 0, err } - res, err := legacyUsdCostBasis(ctx, data, tokenAccount.PublicKey().ToBase58()) - if err != nil { + if err := checkRecord(balanceRecord); err != nil { tracer.OnError(err) return 0, err } - return res, nil + return balanceRecord.UsdCostBasis, nil } // BatchCalculateUsdCostBasisFromCache is like CalculateUsdCostBasisFromCache, @@ -528,106 +335,53 @@ func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Prov tokenAccountStrings[i] = tokenAccount.PublicKey().ToBase58() } - balanceRecords := make(map[string]*balance.Record) - if enableLedgerReads.Get(ctx) { - var err error - balanceRecords, err = data.GetBalanceBatch(ctx, tokenAccountStrings...) - if err != nil { - tracer.OnError(err) - return nil, err - } + balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccountStrings...) + if err != nil { + tracer.OnError(err) + return nil, err } res := make(map[string]int64, len(tokenAccounts)) for _, tokenAccount := range tokenAccountStrings { balanceRecord, ok := balanceRecords[tokenAccount] - if ok && balanceRecord.IsBackfilled { - if !balanceRecord.IsLocked { - tracer.OnError(ErrNotManagedByCode) - return nil, ErrNotManagedByCode - } - res[tokenAccount] = balanceRecord.UsdCostBasis - continue + if !ok { + tracer.OnError(balance.ErrRecordNotFound) + return nil, balance.ErrRecordNotFound } - // todo: The legacy calculation has a batch variant by owner, but the - // fallback is temporary and per-owner batching doesn't map onto - // token accounts without an extra lookup anyway. - usdCostBasis, err := legacyUsdCostBasis(ctx, data, tokenAccount) - if err != nil { + if err := checkRecord(balanceRecord); err != nil { tracer.OnError(err) return nil, err } - res[tokenAccount] = usdCostBasis + res[tokenAccount] = balanceRecord.UsdCostBasis } return res, nil } -// legacyUsdCostBasis derives a token account's cost basis from the owner-level -// intent aggregate, which is only defined for primary accounts. -func legacyUsdCostBasis(ctx context.Context, data ocp_data.Provider, tokenAccount string) (int64, error) { - accountInfoRecord, err := data.GetAccountInfoByTokenAddress(ctx, tokenAccount) - if err == account.ErrAccountInfoNotFound { - return 0, nil - } else if err != nil { - return 0, err +// checkRecord verifies a ledger record is an authoritative view of an account +// that Code still manages. +func checkRecord(record *balance.Record) error { + if !record.IsBackfilled { + return ErrIncompleteRecord } - if accountInfoRecord.AccountType != commonpb.AccountType_PRIMARY { - return 0, nil - } - - usdCostBasis, err := data.GetUsdCostBasis(ctx, accountInfoRecord.OwnerAccount, accountInfoRecord.MintAccount) - if err != nil { - return 0, err - } - return int64(math.Round(usdCostBasis * balance.UsdQuarksPerUnit)), nil -} - -func quarksFromRecord(record *balance.Record) (uint64, error) { // Callers reject unlocked vaults on the timelock record before reaching // here, so this only guards against a record that disagrees with it. if !record.IsLocked { - return 0, ErrNotManagedByCode + return ErrNotManagedByCode } - if record.Quarks < 0 { - return 0, ErrNegativeBalance - } - return uint64(record.Quarks), nil + return nil } -// NetBalanceFromIntentActionsBatch is a balance calculation strategy that incorporates -// the net balance by applying payment intents to the current balance. -func NetBalanceFromIntentActionsBatch(ctx context.Context, data ocp_data.Provider) BatchStrategy { - return func(ctx context.Context, tokenAccounts []string, state *BatchState) (*BatchState, error) { - netBalanceByAccount, err := data.GetNetBalanceFromActionsBatch(ctx, tokenAccounts...) - if err != nil { - return nil, errors.Wrap(err, "error getting net balance from intent actions") - } - - for tokenAccount, netBalance := range netBalanceByAccount { - state.current[tokenAccount] += netBalance - } - - return state, nil +func quarksFromRecord(record *balance.Record) (uint64, error) { + if err := checkRecord(record); err != nil { + return 0, err } -} -// FundingFromExternalDepositsBatch is a balance calculation strategy that adds -// funding from deposits from external accounts. -func FundingFromExternalDepositsBatch(ctx context.Context, data ocp_data.Provider) BatchStrategy { - return func(ctx context.Context, tokenAccounts []string, state *BatchState) (*BatchState, error) { - amountByAccount, err := data.GetTotalExternalDepositedAmountInQuarksBatch(ctx, tokenAccounts...) - if err != nil { - return nil, errors.Wrap(err, "error getting external deposit amount") - } - - for tokenAccount, amount := range amountByAccount { - state.current[tokenAccount] += int64(amount) - } - - return state, nil + if record.Quarks < 0 { + return 0, ErrNegativeBalance } + return uint64(record.Quarks), nil } func (s Source) String() string { diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index a03223d..13b8f1d 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -2,656 +2,191 @@ package balance import ( "context" - "fmt" "testing" - "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" - "github.com/code-payments/ocp-server/config/memory" - "github.com/code-payments/ocp-server/config/wrapper" - currency_lib "github.com/code-payments/ocp-server/currency" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" - "github.com/code-payments/ocp-server/ocp/data/action" "github.com/code-payments/ocp-server/ocp/data/balance" - "github.com/code-payments/ocp-server/ocp/data/deposit" - "github.com/code-payments/ocp-server/ocp/data/intent" - "github.com/code-payments/ocp-server/ocp/data/transaction" timelock_token_v1 "github.com/code-payments/ocp-server/solana/timelock/v1" "github.com/code-payments/ocp-server/testutil" ) -func TestDefaultCalculationMethods_NewCodeAccount(t *testing.T) { - env := setupBalanceTestEnv(t) - - vmConfig := testutil.NewRandomVmConfig(t, true) - newOwnerAccount := testutil.NewRandomAccount(t) - newTokenAccount, err := newOwnerAccount.ToTimelockVault(vmConfig) - require.NoError(t, err) - - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{newOwnerAccount}, - } - - setupBalanceTestData(t, env, data) - - accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, newOwnerAccount) - require.NoError(t, err) - - balance, err := CalculateFromCache(env.ctx, env.data, newTokenAccount) - require.NoError(t, err) - assert.EqualValues(t, 0, balance) - - balanceByAccount, err := BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 0, balanceByAccount[newTokenAccount.PublicKey().ToBase58()]) - - balanceByAccount, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, newTokenAccount) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 0, balanceByAccount[newTokenAccount.PublicKey().ToBase58()]) -} - -func TestDefaultCalculationMethods_DepositFromExternalWallet(t *testing.T) { +func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) - vmConfig := testutil.NewRandomVmConfig(t, true) - owner := testutil.NewRandomAccount(t) - depositAccount, err := owner.ToTimelockVault(vmConfig) - require.NoError(t, err) + first := newBalanceTestAccount(t, env) + second := newBalanceTestAccount(t, env) - externalAccount := testutil.NewRandomAccount(t) + saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 0, IsOpen: true, IsLocked: true, IsBackfilled: true}) - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{owner}, - transactions: []balanceTestTransaction{ - // The following entries are added to the balance - {source: externalAccount, destination: depositAccount, quantity: 1, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: depositAccount, quantity: 10, transactionState: transaction.ConfirmationFinalized}, - // The following entries aren't added to the balance because they aren't finalized - {source: externalAccount, destination: depositAccount, quantity: 100, transactionState: transaction.ConfirmationFailed}, - {source: externalAccount, destination: depositAccount, quantity: 1000, transactionState: transaction.ConfirmationPending}, - {source: externalAccount, destination: depositAccount, quantity: 10000, transactionState: transaction.ConfirmationUnknown}, - }, + expected := map[string]uint64{ + first.tokenAccount.PublicKey().ToBase58(): 42, + second.tokenAccount.PublicKey().ToBase58(): 0, } - setupBalanceTestData(t, env, data) - - balance, err := CalculateFromCache(env.ctx, env.data, depositAccount) - require.NoError(t, err) - assert.EqualValues(t, 11, balance) - - accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner) - require.NoError(t, err) - - balanceByAccount, err := BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 11, balanceByAccount[depositAccount.PublicKey().ToBase58()]) - balanceByAccount, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, depositAccount) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 11, balanceByAccount[depositAccount.PublicKey().ToBase58()]) -} - -func TestDefaultCalculationMethods_MultipleIntents(t *testing.T) { - env := setupBalanceTestEnv(t) - - vmConfig := testutil.NewRandomVmConfig(t, true) - - owner1 := testutil.NewRandomAccount(t) - a1, err := owner1.ToTimelockVault(vmConfig) - require.NoError(t, err) - - owner2 := testutil.NewRandomAccount(t) - a2, err := owner2.ToTimelockVault(vmConfig) - require.NoError(t, err) - - owner3 := testutil.NewRandomAccount(t) - a3, err := owner3.ToTimelockVault(vmConfig) - require.NoError(t, err) - - owner4 := testutil.NewRandomAccount(t) - a4, err := owner4.ToTimelockVault(vmConfig) - require.NoError(t, err) - - externalAccount := testutil.NewRandomAccount(t) - - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{owner1, owner2, owner3, owner4}, - transactions: []balanceTestTransaction{ - // Fund account a1 through a4 with an external deposit - {source: externalAccount, destination: a1, quantity: 1, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: a2, quantity: 10, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: a3, quantity: 100, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: a4, quantity: 1000, transactionState: transaction.ConfirmationFinalized}, - // Confirmed intents are incorporated into balance calculations - {source: a4, destination: a1, quantity: 1, intentID: "i1", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - {source: a4, destination: a1, quantity: 2, intentID: "i2", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - // Pending intents are incorporated into balance calculations - {source: a4, destination: a2, quantity: 3, intentID: "i3", intentState: intent.StatePending, actionState: action.StatePending}, - {source: a4, destination: a2, quantity: 4, intentID: "i4", intentState: intent.StatePending, actionState: action.StatePending}, - // Failed intents are incorporated into balance calculations. We'll - // always make the user whole. - {source: a4, destination: a3, quantity: 5, intentID: "i5", intentState: intent.StateFailed, actionState: action.StateFailed}, - {source: a4, destination: a3, quantity: 6, intentID: "i6", intentState: intent.StateFailed, actionState: action.StateFailed}, - // Intents in the unknown state are incorporated differently depending - // on the intent type, since it infers which intent system it came from. - // Legacy intents are not incorporated, as the intent is not committed by - // the client. Intents could theoretically by in the unknown state under - // the new system, but we should limit this as much as possible. - {source: a4, destination: a1, quantity: 7, intentID: "i7", intentState: intent.StateUnknown, actionState: action.StateUnknown}, - // Revoked intents are not incorporated into balance calculations. - {source: a4, destination: a2, quantity: 8, intentID: "i8", intentState: intent.StateRevoked, actionState: action.StateRevoked}, - }, + for _, testAccount := range []*balanceTestAccount{first, second} { + actual, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) + require.NoError(t, err) + assert.EqualValues(t, expected[testAccount.tokenAccount.PublicKey().ToBase58()], actual) } - setupBalanceTestData(t, env, data) - - balance, err := CalculateFromCache(env.ctx, env.data, a1) - require.NoError(t, err) - assert.EqualValues(t, 11, balance) - - balance, err = CalculateFromCache(env.ctx, env.data, a2) + balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, first.tokenAccount, second.tokenAccount) require.NoError(t, err) - assert.EqualValues(t, 17, balance) - - balance, err = CalculateFromCache(env.ctx, env.data, a3) - require.NoError(t, err) - assert.EqualValues(t, 111, balance) - - balance, err = CalculateFromCache(env.ctx, env.data, a4) - require.NoError(t, err) - assert.EqualValues(t, 972, balance) - - accountRecords1, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner1) - require.NoError(t, err) - - accountRecords2, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner2) - require.NoError(t, err) - - accountRecords3, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner3) - require.NoError(t, err) - - accountRecords4, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner4) - require.NoError(t, err) - - balanceByAccount, err := BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords1[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0], accountRecords2[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0], accountRecords3[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0], accountRecords4[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - require.NoError(t, err) - require.Len(t, balanceByAccount, 4) - assert.EqualValues(t, 11, balanceByAccount[a1.PublicKey().ToBase58()]) - assert.EqualValues(t, 17, balanceByAccount[a2.PublicKey().ToBase58()]) - assert.EqualValues(t, 111, balanceByAccount[a3.PublicKey().ToBase58()]) - assert.EqualValues(t, 972, balanceByAccount[a4.PublicKey().ToBase58()]) + assert.Equal(t, expected, balanceByAccount) - balanceByAccount, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, a1, a2, a3, a4) + balanceByAccount, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, first.accountRecords(t, env), second.accountRecords(t, env)) require.NoError(t, err) - require.Len(t, balanceByAccount, 4) - assert.EqualValues(t, 11, balanceByAccount[a1.PublicKey().ToBase58()]) - assert.EqualValues(t, 17, balanceByAccount[a2.PublicKey().ToBase58()]) - assert.EqualValues(t, 111, balanceByAccount[a3.PublicKey().ToBase58()]) - assert.EqualValues(t, 972, balanceByAccount[a4.PublicKey().ToBase58()]) + assert.Equal(t, expected, balanceByAccount) } -func TestDefaultCalculationMethods_BackAndForth(t *testing.T) { +func TestDefaultCalculationMethods_MissingBalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) - vmConfig := testutil.NewRandomVmConfig(t, true) - - owner1 := testutil.NewRandomAccount(t) - a1, err := owner1.ToTimelockVault(vmConfig) - require.NoError(t, err) - - owner2 := testutil.NewRandomAccount(t) - a2, err := owner2.ToTimelockVault(vmConfig) - require.NoError(t, err) - - externalAccount := testutil.NewRandomAccount(t) - - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{owner1, owner2}, - transactions: []balanceTestTransaction{ - // Fund account a1 through an external deposit - {source: externalAccount, destination: a1, quantity: 1, transactionState: transaction.ConfirmationFinalized}, - // Setup a set of intents that result in back and forth movement of the Kin - {source: a1, destination: a2, quantity: 1, intentID: "i1", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - {source: a2, destination: a1, quantity: 1, intentID: "i2", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - {source: a1, destination: a2, quantity: 1, intentID: "i3", intentState: intent.StatePending, actionState: action.StatePending}, - {source: a2, destination: a1, quantity: 1, intentID: "i4", intentState: intent.StatePending, actionState: action.StatePending}, - {source: a1, destination: a2, quantity: 1, intentID: "i5", intentState: intent.StatePending, actionState: action.StatePending}, - }, - } - - setupBalanceTestData(t, env, data) - - balance, err := CalculateFromCache(env.ctx, env.data, a1) - require.NoError(t, err) - assert.EqualValues(t, 0, balance) + // The account is managed, but the ledger has no record for it, which is a + // broken invariant rather than a balance of zero + testAccount := newBalanceTestAccount(t, env) - balance, err = CalculateFromCache(env.ctx, env.data, a2) - require.NoError(t, err) - assert.EqualValues(t, 1, balance) + _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, balance.ErrRecordNotFound, err) - accountRecords1, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner1) - require.NoError(t, err) + _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, balance.ErrRecordNotFound, err) - accountRecords2, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner2) - require.NoError(t, err) + _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) + assert.Equal(t, balance.ErrRecordNotFound, err) - balanceByAccount, err := BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords1[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0], accountRecords2[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - require.NoError(t, err) - require.Len(t, balanceByAccount, 2) - assert.EqualValues(t, 0, balanceByAccount[a1.PublicKey().ToBase58()]) - assert.EqualValues(t, 1, balanceByAccount[a2.PublicKey().ToBase58()]) + _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, balance.ErrRecordNotFound, err) - balanceByAccount, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, a1, a2) - require.NoError(t, err) - require.Len(t, balanceByAccount, 2) - assert.EqualValues(t, 0, balanceByAccount[a1.PublicKey().ToBase58()]) - assert.EqualValues(t, 1, balanceByAccount[a2.PublicKey().ToBase58()]) + _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, balance.ErrRecordNotFound, err) } -func TestDefaultCalculationMethods_SelfPayments(t *testing.T) { +func TestDefaultCalculationMethods_IncompleteBalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) - vmConfig := testutil.NewRandomVmConfig(t, true) - ownerAccount := testutil.NewRandomAccount(t) - tokenAccount, err := ownerAccount.ToTimelockVault(vmConfig) - require.NoError(t, err) - - externalAccount := testutil.NewRandomAccount(t) + // A record that isn't backfilled only holds deltas observed since it was + // seeded, so it can't answer for the account + testAccount := newBalanceTestAccount(t, env) + saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{ownerAccount}, - transactions: []balanceTestTransaction{ - // Fund account the token account through an external deposit - {source: externalAccount, destination: tokenAccount, quantity: 1, transactionState: transaction.ConfirmationFinalized}, - // Setup a set of intents that result in self-payments and no-ops to - // the balance calculation - {source: tokenAccount, destination: tokenAccount, quantity: 1, intentID: "i1", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - {source: tokenAccount, destination: tokenAccount, quantity: 1, intentID: "i2", intentState: intent.StateConfirmed, actionState: action.StateConfirmed, transactionState: transaction.ConfirmationFinalized}, - {source: tokenAccount, destination: tokenAccount, quantity: 1, intentID: "i3", intentState: intent.StatePending, actionState: action.StatePending}, - {source: tokenAccount, destination: tokenAccount, quantity: 1, intentID: "i4", intentState: intent.StatePending, actionState: action.StatePending}, - }, - } + _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, ErrIncompleteRecord, err) - setupBalanceTestData(t, env, data) + _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, ErrIncompleteRecord, err) - balance, err := CalculateFromCache(env.ctx, env.data, tokenAccount) - require.NoError(t, err) - assert.EqualValues(t, 1, balance) + _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) + assert.Equal(t, ErrIncompleteRecord, err) - accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, ownerAccount) - require.NoError(t, err) - - balanceByAccount, err := BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 1, balanceByAccount[tokenAccount.PublicKey().ToBase58()]) + _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, ErrIncompleteRecord, err) - balanceByAccount, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) - require.NoError(t, err) - require.Len(t, balanceByAccount, 1) - assert.EqualValues(t, 1, balanceByAccount[tokenAccount.PublicKey().ToBase58()]) + _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) + assert.Equal(t, ErrIncompleteRecord, err) } func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { env := setupBalanceTestEnv(t) - vmConfig := testutil.NewRandomVmConfig(t, true) - ownerAccount := testutil.NewRandomAccount(t) - tokenAccount, err := ownerAccount.ToTimelockVault(vmConfig) - require.NoError(t, err) + testAccount := newBalanceTestAccount(t, env) + saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true, IsBackfilled: true}) - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{ownerAccount}, - } - - setupBalanceTestData(t, env, data) - - timelockRecord, err := env.data.GetTimelockByVault(env.ctx, tokenAccount.PublicKey().ToBase58()) + timelockRecord, err := env.data.GetTimelockByVault(env.ctx, testAccount.tokenAccount.PublicKey().ToBase58()) require.NoError(t, err) timelockRecord.VaultState = timelock_token_v1.StateWaitingForTimeout timelockRecord.Block += 1 require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord)) - accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, ownerAccount) - require.NoError(t, err) - - _, err = CalculateFromCache(env.ctx, env.data, tokenAccount) + _, err = CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) + _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) + _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) assert.Equal(t, ErrNotManagedByCode, err) } -func TestDefaultCalculationMethods_BalanceRecord(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) - pendingOwner := testutil.NewRandomAccount(t) - pendingAccount, err := pendingOwner.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, pendingOwner, legacyOwner}, - transactions: []balanceTestTransaction{ - {source: externalAccount, destination: backfilledAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: pendingAccount, quantity: 22, transactionState: transaction.ConfirmationFinalized}, - {source: externalAccount, destination: legacyAccount, quantity: 33, transactionState: transaction.ConfirmationFinalized}, - }, - } - - setupBalanceTestData(t, env, data) - - // A backfilled record is authoritative, 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, - IsOpen: true, - IsLocked: true, - IsBackfilled: true, - })) - - // A record that isn't backfilled is ignored in favour of history - require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ - TokenAccount: pendingAccount.PublicKey().ToBase58(), - OwnerAccount: pendingOwner.PublicKey().ToBase58(), - MintAccount: vmConfig.Mint.PublicKey().ToBase58(), - Quarks: -5, - IsOpen: true, - })) - - expected := map[string]uint64{ - backfilledAccount.PublicKey().ToBase58(): 42, - pendingAccount.PublicKey().ToBase58(): 22, - legacyAccount.PublicKey().ToBase58(): 33, - } - - for tokenAccount, expectedQuarks := range expected { - account, err := common.NewAccountFromPublicKeyString(tokenAccount) - require.NoError(t, err) - - actual, err := CalculateFromCache(env.ctx, env.data, account) - require.NoError(t, err) - assert.EqualValues(t, expectedQuarks, actual, tokenAccount) - } - - balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, backfilledAccount, pendingAccount, legacyAccount) - require.NoError(t, err) - assert.Equal(t, expected, balanceByAccount) - - var allAccountRecords []*common.AccountRecords - for _, owner := range data.codeUsers { - accountRecords, err := common.GetLatestTokenAccountRecordsForOwner(env.ctx, env.data, owner) - require.NoError(t, err) - allAccountRecords = append(allAccountRecords, accountRecords[vmConfig.Mint.PublicKey().ToBase58()][commonpb.AccountType_PRIMARY][0]) - } - balanceByAccount, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, allAccountRecords...) - require.NoError(t, err) - 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) - - vmConfig := testutil.NewRandomVmConfig(t, true) - owner := testutil.NewRandomAccount(t) - tokenAccount, err := owner.ToTimelockVault(vmConfig) - require.NoError(t, err) - - externalAccount := testutil.NewRandomAccount(t) - - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{owner}, - transactions: []balanceTestTransaction{ - {source: externalAccount, destination: tokenAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, - }, - } - - setupBalanceTestData(t, env, data) // A backfilled record for an unlocked vault holds the last managed state, // not a live balance, so it is refused even though the timelock record // still passes the managed check. That pairing is inconsistent by // construction: the timelock check normally rejects first, so the fixture // exists to exercise the record's own guard. - require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ - TokenAccount: tokenAccount.PublicKey().ToBase58(), - OwnerAccount: owner.PublicKey().ToBase58(), - MintAccount: vmConfig.Mint.PublicKey().ToBase58(), - Quarks: 42, - UsdCostBasis: 4_200_000, - IsOpen: true, - IsLocked: false, - IsBackfilled: true, - })) + testAccount := newBalanceTestAccount(t, env) + saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsBackfilled: true}) - _, err = CalculateFromCache(env.ctx, env.data, tokenAccount) + _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) + _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) + _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) + _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) } -func TestDefaultCalculationMethods_BalanceRecordReadsDisabled(t *testing.T) { +func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) { env := setupBalanceTestEnv(t) - disableLedgerReadsForTest(t) - - vmConfig := testutil.NewRandomVmConfig(t, true) - owner := testutil.NewRandomAccount(t) - tokenAccount, err := owner.ToTimelockVault(vmConfig) - require.NoError(t, err) - - externalAccount := testutil.NewRandomAccount(t) - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{owner}, - transactions: []balanceTestTransaction{ - {source: externalAccount, destination: tokenAccount, quantity: 11, transactionState: transaction.ConfirmationFinalized}, - }, - } + first := newBalanceTestAccount(t, env) + second := newBalanceTestAccount(t, env) - setupBalanceTestData(t, env, data) - - // A backfilled record exists, but reads are disabled, so history wins - require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ - TokenAccount: tokenAccount.PublicKey().ToBase58(), - OwnerAccount: owner.PublicKey().ToBase58(), - MintAccount: vmConfig.Mint.PublicKey().ToBase58(), - Quarks: 42, - UsdCostBasis: 123, - IsOpen: true, - IsLocked: true, - IsBackfilled: true, - })) + // Both values come from the same record, so they can't disagree + saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 33, UsdCostBasis: -123456, IsOpen: true, IsLocked: true, IsBackfilled: true}) - actual, err := CalculateFromCache(env.ctx, env.data, tokenAccount) + res, err := BatchCalculateWithUsdCostBasisFromCache(env.ctx, env.data, first.accountRecords(t, env), second.accountRecords(t, env)) require.NoError(t, err) - assert.EqualValues(t, 11, actual) + require.Len(t, res, 2) - balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, tokenAccount) - require.NoError(t, err) - assert.EqualValues(t, 11, balanceByAccount[tokenAccount.PublicKey().ToBase58()]) + cached := res[first.tokenAccount.PublicKey().ToBase58()] + require.NotNil(t, cached) + assert.EqualValues(t, 42, cached.Quarks) + assert.EqualValues(t, 4_200_000, cached.UsdCostBasis) - usdCostBasis, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, tokenAccount) - require.NoError(t, err) - assert.EqualValues(t, 0, usdCostBasis) + cached = res[second.tokenAccount.PublicKey().ToBase58()] + require.NotNil(t, cached) + assert.EqualValues(t, 33, cached.Quarks) + assert.EqualValues(t, -123456, cached.UsdCostBasis) } func TestUsdCostBasisCalculationMethods(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) - - data := &balanceTestData{ - vmConfig: vmConfig, - codeUsers: []*common.Account{backfilledOwner, legacyOwner}, - } - setupBalanceTestData(t, env, data) + first := newBalanceTestAccount(t, env) + second := newBalanceTestAccount(t, env) - require.NoError(t, env.data.CreateBalance(env.ctx, &balance.Record{ - TokenAccount: backfilledAccount.PublicKey().ToBase58(), - OwnerAccount: backfilledOwner.PublicKey().ToBase58(), - MintAccount: vmConfig.Mint.PublicKey().ToBase58(), - UsdCostBasis: -123456, - IsOpen: true, - IsLocked: true, - IsBackfilled: true, - })) - - // The legacy calculation is owner-level, and derived from intents - require.NoError(t, env.data.SaveIntent(env.ctx, &intent.Record{ - IntentId: testutil.NewRandomAccount(t).PublicKey().ToBase58(), - IntentType: intent.ExternalDeposit, - MintAccount: vmConfig.Mint.PublicKey().ToBase58(), - InitiatorOwnerAccount: legacyOwner.PublicKey().ToBase58(), - ExternalDepositMetadata: &intent.ExternalDepositMetadata{ - DestinationTokenAccount: legacyAccount.PublicKey().ToBase58(), - Quantity: 1, - ExchangeCurrency: currency_lib.USD, - ExchangeRate: 1.0, - NativeAmount: 1.5, - UsdMarketValue: 1.5, - }, - State: intent.StateConfirmed, - CreatedAt: time.Now(), - })) + saveBalanceTestRecord(t, env, first, &balance.Record{UsdCostBasis: -123456, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 1, UsdCostBasis: 1_500_000, IsOpen: true, IsLocked: true, IsBackfilled: true}) expected := map[string]int64{ - backfilledAccount.PublicKey().ToBase58(): -123456, - legacyAccount.PublicKey().ToBase58(): 1_500_000, + first.tokenAccount.PublicKey().ToBase58(): -123456, + second.tokenAccount.PublicKey().ToBase58(): 1_500_000, } - for tokenAccount, expectedUsdCostBasis := range expected { - account, err := common.NewAccountFromPublicKeyString(tokenAccount) - require.NoError(t, err) - - actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, account) + for _, testAccount := range []*balanceTestAccount{first, second} { + actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) require.NoError(t, err) - assert.EqualValues(t, expectedUsdCostBasis, actual, tokenAccount) + assert.EqualValues(t, expected[testAccount.tokenAccount.PublicKey().ToBase58()], actual) } - usdCostBasisByAccount, err := BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, backfilledAccount, legacyAccount) + usdCostBasisByAccount, err := BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, first.tokenAccount, second.tokenAccount) require.NoError(t, err) assert.Equal(t, expected, usdCostBasisByAccount) - - // Accounts unknown to the system have no cost basis - unknownAccount := testutil.NewRandomAccount(t) - actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, unknownAccount) - require.NoError(t, err) - assert.EqualValues(t, 0, actual) } func TestDefaultCalculation_ExternalAccount(t *testing.T) { @@ -663,42 +198,15 @@ func TestDefaultCalculation_ExternalAccount(t *testing.T) { // Note: not possible with batch method, since we wouldn't have account records } -func enableLedgerReadsForTest(t *testing.T) { - previous := enableLedgerReads - enableLedgerReads = wrapper.NewBoolConfig(memory.NewConfig(true), defaultEnableLedgerReads) - t.Cleanup(func() { - enableLedgerReads = previous - }) -} - -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 } -type balanceTestData struct { +type balanceTestAccount struct { vmConfig *common.VmConfig - codeUsers []*common.Account - transactions []balanceTestTransaction -} - -type balanceTestTransaction struct { - source, destination *common.Account - quantity uint64 - - intentID string - intentState intent.State - actionState action.State - - transactionState transaction.Confirmation + owner *common.Account + tokenAccount *common.Account } func setupBalanceTestEnv(t *testing.T) (env balanceTestEnv) { @@ -708,77 +216,53 @@ func setupBalanceTestEnv(t *testing.T) (env balanceTestEnv) { return env } -func setupBalanceTestData(t *testing.T, env balanceTestEnv, data *balanceTestData) { - for _, owner := range data.codeUsers { - timelockAccounts, err := owner.GetTimelockAccounts(data.vmConfig) - require.NoError(t, err) - timelockRecord := timelockAccounts.ToDBRecord() - timelockRecord.VaultState = timelock_token_v1.StateLocked - timelockRecord.Block += 1 - require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord)) - - accountInfoRecord := &account.Record{ - OwnerAccount: owner.PublicKey().ToBase58(), - AuthorityAccount: owner.PublicKey().ToBase58(), - TokenAccount: timelockRecord.VaultAddress, - MintAccount: data.vmConfig.Mint.PublicKey().ToBase58(), - AccountType: commonpb.AccountType_PRIMARY, - } - require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountInfoRecord)) +// newBalanceTestAccount creates a locked timelock account, with an account +// info record but no ledger record. +func newBalanceTestAccount(t *testing.T, env balanceTestEnv) *balanceTestAccount { + vmConfig := testutil.NewRandomVmConfig(t, true) + owner := testutil.NewRandomAccount(t) + + timelockAccounts, err := owner.GetTimelockAccounts(vmConfig) + require.NoError(t, err) + timelockRecord := timelockAccounts.ToDBRecord() + timelockRecord.VaultState = timelock_token_v1.StateLocked + timelockRecord.Block += 1 + require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord)) + + require.NoError(t, env.data.CreateAccountInfo(env.ctx, &account.Record{ + OwnerAccount: owner.PublicKey().ToBase58(), + AuthorityAccount: owner.PublicKey().ToBase58(), + TokenAccount: timelockRecord.VaultAddress, + MintAccount: vmConfig.Mint.PublicKey().ToBase58(), + AccountType: commonpb.AccountType_PRIMARY, + })) + + tokenAccount, err := common.NewAccountFromPublicKeyString(timelockRecord.VaultAddress) + require.NoError(t, err) + + return &balanceTestAccount{ + vmConfig: vmConfig, + owner: owner, + tokenAccount: tokenAccount, } +} - for i, txn := range data.transactions { - // Setup the intent record with an equivalent action record - if len(txn.intentID) > 0 { - intentRecord := &intent.Record{ - IntentId: txn.intentID, - IntentType: intent.SendPublicPayment, - MintAccount: data.vmConfig.Mint.PublicKey().ToBase58(), - InitiatorOwnerAccount: "owner", - SendPublicPaymentMetadata: &intent.SendPublicPaymentMetadata{ - DestinationOwnerAccount: testutil.NewRandomAccount(t).PublicKey().ToBase58(), - DestinationTokenAccount: txn.destination.PublicKey().ToBase58(), - Quantity: txn.quantity, - - ExchangeCurrency: currency_lib.USD, - ExchangeRate: 1.0, - NativeAmount: 1.0, - UsdMarketValue: 1.0, - }, - State: txn.intentState, - CreatedAt: time.Now(), - } - require.NoError(t, env.data.SaveIntent(env.ctx, intentRecord)) - - actionRecord := &action.Record{ - Intent: txn.intentID, - IntentType: intent.SendPublicPayment, - - ActionId: 0, - ActionType: action.NoPrivacyTransfer, - - Source: txn.source.PublicKey().ToBase58(), - Destination: &intentRecord.SendPublicPaymentMetadata.DestinationTokenAccount, - Quantity: &intentRecord.SendPublicPaymentMetadata.Quantity, - - State: txn.actionState, - } - require.NoError(t, env.data.PutAllActions(env.ctx, actionRecord)) - } - - // There's no intent, so we have an external deposit - if len(txn.intentID) == 0 && txn.transactionState != transaction.ConfirmationUnknown { - depositRecord := &deposit.Record{ - Signature: fmt.Sprintf("txn%d", i), - Destination: txn.destination.PublicKey().ToBase58(), - Amount: txn.quantity, - - Slot: 12345, - ConfirmationState: txn.transactionState, - - CreatedAt: time.Now(), - } - require.NoError(t, env.data.SaveExternalDeposit(env.ctx, depositRecord)) - } +// saveBalanceTestRecord creates the account's ledger record, filling in the +// identifying fields from the account. +func saveBalanceTestRecord(t *testing.T, env balanceTestEnv, testAccount *balanceTestAccount, record *balance.Record) { + record.TokenAccount = testAccount.tokenAccount.PublicKey().ToBase58() + record.OwnerAccount = testAccount.owner.PublicKey().ToBase58() + record.MintAccount = testAccount.vmConfig.Mint.PublicKey().ToBase58() + require.NoError(t, env.data.CreateBalance(env.ctx, record)) +} + +func (a *balanceTestAccount) accountRecords(t *testing.T, env balanceTestEnv) *common.AccountRecords { + generalRecord, err := env.data.GetAccountInfoByTokenAddress(env.ctx, a.tokenAccount.PublicKey().ToBase58()) + require.NoError(t, err) + timelockRecord, err := env.data.GetTimelockByVault(env.ctx, a.tokenAccount.PublicKey().ToBase58()) + require.NoError(t, err) + return &common.AccountRecords{ + General: generalRecord, + Timelock: timelockRecord, } } diff --git a/ocp/balance/config.go b/ocp/balance/config.go deleted file mode 100644 index 6925291..0000000 --- a/ocp/balance/config.go +++ /dev/null @@ -1,25 +0,0 @@ -package balance - -import ( - "github.com/code-payments/ocp-server/config" - "github.com/code-payments/ocp-server/config/env" -) - -const ( - // EnableLedgerReadsConfigEnvName gates whether balance calculators read - // from the new ocp__core_balance ledger. When disabled, calculators use - // the legacy strategies exclusively. - EnableLedgerReadsConfigEnvName = "BALANCE_ENABLE_LEDGER_READS" - - // EnableLedgerWritesConfigEnvName gates whether ApplyDeltasInTx writes - // to the ledger at all. When disabled, it is a no-op. - EnableLedgerWritesConfigEnvName = "BALANCE_ENABLE_LEDGER_WRITES" - - defaultEnableLedgerReads = true - defaultEnableLedgerWrites = true -) - -var ( - enableLedgerReads config.Bool = env.NewBoolConfig(EnableLedgerReadsConfigEnvName, defaultEnableLedgerReads) - enableLedgerWrites config.Bool = env.NewBoolConfig(EnableLedgerWritesConfigEnvName, defaultEnableLedgerWrites) -) diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index 4ad81c5..aec6fd0 100644 --- a/ocp/balance/ledger.go +++ b/ocp/balance/ledger.go @@ -14,25 +14,10 @@ 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. -func LedgerWritesEnabled(ctx context.Context) bool { - return enableLedgerWrites.Get(ctx) -} - // ApplyDeltasInTx applies balance deltas to the ledger. It must be called // within the DB transaction that commits the records the deltas are derived // from, so the ledger can never disagree with them. // -// It is a no-op while ledger writes are disabled. -// // The ledger only tracks timelock accounts. Credits to any other account, // like an external wallet or the fee collector, are dropped, since delta // builders don't know which destinations OCP manages. Outgoing deltas from @@ -52,7 +37,7 @@ func LedgerWritesEnabled(ctx context.Context) bool { // balance.ErrBalanceChanged, balance.ErrAccountClosed, // balance.ErrAccountUnlocked) are returned as is for the caller to map. func ApplyDeltasInTx(ctx context.Context, data ocp_data.Provider, deltas ...*balance.Delta) error { - if !enableLedgerWrites.Get(ctx) || len(deltas) == 0 { + if len(deltas) == 0 { return nil } @@ -87,10 +72,10 @@ func ApplyDeltasInTx(ctx context.Context, data ocp_data.Provider, deltas ...*bal // record. A new account has no history, so its record is created backfilled // at zero and predicates are enforced from the start. // -// It is a no-op while ledger writes are disabled, and for accounts that -// aren't timelock accounts, which the ledger doesn't track. +// It is a no-op for accounts that aren't timelock accounts, which the ledger +// doesn't track. func CreateRecordInTx(ctx context.Context, data ocp_data.Provider, accountInfoRecord *account.Record) error { - if !enableLedgerWrites.Get(ctx) || !accountInfoRecord.IsTimelock() { + if !accountInfoRecord.IsTimelock() { return nil } diff --git a/ocp/balance/ledger_test.go b/ocp/balance/ledger_test.go index 5bf6caa..c842e46 100644 --- a/ocp/balance/ledger_test.go +++ b/ocp/balance/ledger_test.go @@ -9,35 +9,15 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" - "github.com/code-payments/ocp-server/config/memory" - "github.com/code-payments/ocp-server/config/wrapper" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/testutil" ) -func TestApplyDeltasInTx_WritesDisabled(t *testing.T) { - ctx := context.Background() - data := ocp_data.NewTestDataProvider() - disableLedgerWritesForTest(t) - - source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) - - require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{ - TokenAccount: source, - Kind: balance.DeltaDebit, - Quarks: 100, - })) - - _, err := data.GetBalance(ctx, source) - assert.Equal(t, balance.ErrRecordNotFound, err) -} - func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - enableLedgerWritesForTest(t) source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) destination := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_REMOTE_SEND_GIFT_CARD) @@ -90,7 +70,6 @@ func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { func TestApplyDeltasInTx_UnknownSource(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - enableLedgerWritesForTest(t) external := testutil.NewRandomAccount(t).PublicKey().ToBase58() swap := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_SWAP) @@ -103,7 +82,6 @@ func TestApplyDeltasInTx_UnknownSource(t *testing.T) { func TestApplyDeltasInTx_OnlyUntrackedCredits(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - enableLedgerWritesForTest(t) external := testutil.NewRandomAccount(t).PublicKey().ToBase58() require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: external, Kind: balance.DeltaCredit, Quarks: 1})) @@ -114,7 +92,6 @@ func TestApplyDeltasInTx_OnlyUntrackedCredits(t *testing.T) { func TestApplyDeltasInTx_UnlockedAccount(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - enableLedgerWritesForTest(t) unlocked := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY) require.NoError(t, CreateRecordInTx(ctx, data, unlocked)) @@ -139,7 +116,6 @@ func TestApplyDeltasInTx_UnlockedAccount(t *testing.T) { func TestApplyDeltasInTx_InvalidDelta(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - enableLedgerWritesForTest(t) source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) assert.Error(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit})) @@ -152,16 +128,7 @@ func TestCreateRecordInTx(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - // Disabled writes are a no-op primary := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY) - func() { - disableLedgerWritesForTest(t) - require.NoError(t, CreateRecordInTx(ctx, data, primary)) - _, err := data.GetBalance(ctx, primary.TokenAccount) - assert.Equal(t, balance.ErrRecordNotFound, err) - }() - - enableLedgerWritesForTest(t) // A new timelock account starts backfilled at zero, so predicates are // enforced immediately @@ -213,19 +180,3 @@ func newLedgerTestAccountInfo(t *testing.T, ctx context.Context, data ocp_data.P require.NoError(t, data.CreateAccountInfo(ctx, record)) 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) - t.Cleanup(func() { - enableLedgerWrites = previous - }) -} diff --git a/ocp/rpc/account/server_test.go b/ocp/rpc/account/server_test.go index 98eba93..fad9efb 100644 --- a/ocp/rpc/account/server_test.go +++ b/ocp/rpc/account/server_test.go @@ -3,8 +3,6 @@ package account import ( "context" "crypto/ed25519" - "fmt" - "math/rand" "testing" "time" @@ -19,20 +17,20 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" "github.com/code-payments/ocp-server/currency" + balance_util "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" currency_util "github.com/code-payments/ocp-server/ocp/currency" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" + "github.com/code-payments/ocp-server/ocp/data/balance" exchange_memory "github.com/code-payments/ocp-server/ocp/data/currency/exchange/memory" "github.com/code-payments/ocp-server/ocp/data/currency/holder" holder_memory "github.com/code-payments/ocp-server/ocp/data/currency/holder/memory" "github.com/code-payments/ocp-server/ocp/data/currency/reserve" reserve_memory "github.com/code-payments/ocp-server/ocp/data/currency/reserve/memory" - "github.com/code-payments/ocp-server/ocp/data/deposit" "github.com/code-payments/ocp-server/ocp/data/intent" "github.com/code-payments/ocp-server/ocp/data/timelock" - "github.com/code-payments/ocp-server/ocp/data/transaction" "github.com/code-payments/ocp-server/pointer" "github.com/code-payments/ocp-server/solana/currencycreator" timelock_token_v1 "github.com/code-payments/ocp-server/solana/timelock/v1" @@ -634,6 +632,7 @@ func TestGetTokenAccountInfos_BlockchainState(t *testing.T) { accountRecords.Timelock.VaultState = tc.timelockState accountRecords.Timelock.Block += 1 require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountRecords.General)) + require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountRecords.General)) require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) resp, err := env.client.GetTokenAccountInfos(env.ctx, req) @@ -702,6 +701,7 @@ func TestGetTokenAccountInfos_ManagementState(t *testing.T) { accountRecords.Timelock.VaultState = tc.timelockState accountRecords.Timelock.Block = tc.block require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountRecords.General)) + require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountRecords.General)) require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) resp, err := env.client.GetTokenAccountInfos(env.ctx, req) @@ -804,6 +804,7 @@ func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccou accountRecords := getDefaultTestAccountRecords(t, ownerAccount, authorityAccount, vmConfig, index, accountType) require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountRecords.General)) + require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountRecords.General)) if accountRecords.IsTimelock() { accountRecords.Timelock.VaultState = timelock_token_v1.StateLocked @@ -847,14 +848,10 @@ func getDefaultTestAccountRecords(t *testing.T, ownerAccount, authorityAccount * } } -func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, balance uint64) { - depositRecord := &deposit.Record{ - Signature: fmt.Sprintf("txn%d", rand.Uint64()), - Destination: accountRecords.General.TokenAccount, - Amount: balance, - - ConfirmationState: transaction.ConfirmationFinalized, - Slot: 12345, - } - require.NoError(t, env.data.SaveExternalDeposit(env.ctx, depositRecord)) +func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, quarks uint64) { + require.NoError(t, balance_util.ApplyDeltasInTx(env.ctx, env.data, &balance.Delta{ + TokenAccount: accountRecords.General.TokenAccount, + Kind: balance.DeltaCredit, + Quarks: quarks, + })) } diff --git a/ocp/rpc/balance/server_test.go b/ocp/rpc/balance/server_test.go index 27d90eb..823f66c 100644 --- a/ocp/rpc/balance/server_test.go +++ b/ocp/rpc/balance/server_test.go @@ -2,8 +2,6 @@ package balance import ( "context" - "fmt" - "math/rand" "testing" "time" @@ -15,17 +13,17 @@ import ( balancepb "github.com/code-payments/ocp-protobuf-api/generated/go/balance/v1" commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" + balance_util "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" currency_util "github.com/code-payments/ocp-server/ocp/currency" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" + "github.com/code-payments/ocp-server/ocp/data/balance" exchange_memory "github.com/code-payments/ocp-server/ocp/data/currency/exchange/memory" "github.com/code-payments/ocp-server/ocp/data/currency/holder" holder_memory "github.com/code-payments/ocp-server/ocp/data/currency/holder/memory" "github.com/code-payments/ocp-server/ocp/data/currency/reserve" reserve_memory "github.com/code-payments/ocp-server/ocp/data/currency/reserve/memory" - "github.com/code-payments/ocp-server/ocp/data/deposit" - "github.com/code-payments/ocp-server/ocp/data/transaction" "github.com/code-payments/ocp-server/solana/currencycreator" timelock_token_v1 "github.com/code-payments/ocp-server/solana/timelock/v1" "github.com/code-payments/ocp-server/testutil" @@ -197,6 +195,7 @@ func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccou require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountInfoRecord)) require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord)) + require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountInfoRecord)) return &common.AccountRecords{ General: accountInfoRecord, @@ -204,14 +203,10 @@ func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccou } } -func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, balance uint64) { - depositRecord := &deposit.Record{ - Signature: fmt.Sprintf("txn%d", rand.Uint64()), - Destination: accountRecords.General.TokenAccount, - Amount: balance, - - ConfirmationState: transaction.ConfirmationFinalized, - Slot: 12345, - } - require.NoError(t, env.data.SaveExternalDeposit(env.ctx, depositRecord)) +func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, quarks uint64) { + require.NoError(t, balance_util.ApplyDeltasInTx(env.ctx, env.data, &balance.Delta{ + TokenAccount: accountRecords.General.TokenAccount, + Kind: balance.DeltaCredit, + Quarks: quarks, + })) } diff --git a/ocp/rpc/transaction/intent.go b/ocp/rpc/transaction/intent.go index b2d1e2f..e6f6668 100644 --- a/ocp/rpc/transaction/intent.go +++ b/ocp/rpc/transaction/intent.go @@ -720,17 +720,15 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm // Reflect the intent in the balance ledger. Applied last, so the // ledger row locks are held for as little of the transaction as // possible. - if balance.LedgerWritesEnabled(ctx) { - balanceDeltas, err := balance.DeltasForSubmittedIntent(intentRecord, actionRecords, s.conf.createOnSendWithdrawalFeeQuarks.Get(ctx)) - if err != nil { - log.With(zap.Error(err)).Warn("failure building balance deltas") - return err - } - err = balance.ApplyDeltasInTx(ctx, s.data, balanceDeltas...) - if err != nil { - log.With(zap.Error(err)).Warn("failure applying balance deltas") - return err - } + balanceDeltas, err := balance.DeltasForSubmittedIntent(intentRecord, actionRecords, s.conf.createOnSendWithdrawalFeeQuarks.Get(ctx)) + if err != nil { + log.With(zap.Error(err)).Warn("failure building balance deltas") + return err + } + err = balance.ApplyDeltasInTx(ctx, s.data, balanceDeltas...) + if err != nil { + log.With(zap.Error(err)).Warn("failure applying balance deltas") + return err } // Schedule app-defined tasks atomically with the intent, so their diff --git a/ocp/worker/account/gift_card.go b/ocp/worker/account/gift_card.go index 031c15b..7c32442 100644 --- a/ocp/worker/account/gift_card.go +++ b/ocp/worker/account/gift_card.go @@ -233,15 +233,13 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data ocp_data.Prov return err } - if balance.LedgerWritesEnabled(ctx) { - balanceDeltas, err := balance.DeltasForGiftCardAutoReturn(autoReturnIntent, autoReturnAction) - if err != nil { - return err - } - err = balance.ApplyDeltasInTx(ctx, data, balanceDeltas...) - if err != nil { - return err - } + balanceDeltas, err := balance.DeltasForGiftCardAutoReturn(autoReturnIntent, autoReturnAction) + if err != nil { + return err + } + err = balance.ApplyDeltasInTx(ctx, data, balanceDeltas...) + if err != nil { + return err } // This will trigger the fulfillment worker to poll for the fulfillment. This diff --git a/ocp/worker/currency/holder/runtime.go b/ocp/worker/currency/holder/runtime.go index 72d78e4..3ac018a 100644 --- a/ocp/worker/currency/holder/runtime.go +++ b/ocp/worker/currency/holder/runtime.go @@ -4,16 +4,11 @@ import ( "context" "time" - "github.com/pkg/errors" "go.uber.org/zap" - commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" - "github.com/code-payments/ocp-server/metrics" - "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" - "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/currency" currency_holder "github.com/code-payments/ocp-server/ocp/data/currency/holder" "github.com/code-payments/ocp-server/ocp/data/currency/reserve" @@ -115,64 +110,5 @@ func (p *holderRuntime) countHoldersForMint(ctx context.Context, mint string, cu return 0, nil } - if balance.LedgerReadsEnabled(ctx) { - return p.data.CountLockedBalancesByMint(ctx, mint, int64(minHoldings)) - } - - accountRecords, err := p.data.GetAccountInfosByMintAndType(ctx, mint, commonpb.AccountType_PRIMARY) - if err == account.ErrAccountInfoNotFound { - return 0, nil - } else if err != nil { - return 0, errors.Wrap(err, "error getting primary accounts") - } - - if len(accountRecords) == 0 { - return 0, nil - } - - vaultAddresses := make([]string, len(accountRecords)) - for i, record := range accountRecords { - vaultAddresses[i] = record.TokenAccount - } - - timelockRecordsByVault, err := p.data.GetTimelockByVaultBatch(ctx, vaultAddresses...) - if err != nil { - return 0, errors.Wrap(err, "error getting timelock records") - } - - tokenAccounts := make([]*common.Account, 0) - for _, record := range accountRecords { - timelockRecord, ok := timelockRecordsByVault[record.TokenAccount] - if !ok { - return 0, errors.Errorf("timelock record unexpectedly missing for vault %s", record.TokenAccount) - } - - if !timelockRecord.IsLocked() { - continue - } - - tokenAccount, err := common.NewAccountFromPublicKeyString(record.TokenAccount) - if err != nil { - return 0, errors.Wrap(err, "invalid token account public key") - } - tokenAccounts = append(tokenAccounts, tokenAccount) - } - - if len(tokenAccounts) == 0 { - return 0, nil - } - - balances, err := balance.BatchCalculateFromCacheWithTokenAccounts(ctx, p.data, tokenAccounts...) - if err != nil { - return 0, errors.Wrap(err, "error calculating balances batch") - } - - var count uint64 - for _, bal := range balances { - if bal >= minHoldings { - count++ - } - } - - return count, nil + return p.data.CountLockedBalancesByMint(ctx, mint, int64(minHoldings)) } diff --git a/ocp/worker/geyser/external_deposit.go b/ocp/worker/geyser/external_deposit.go index 9e2356d..cc4d553 100644 --- a/ocp/worker/geyser/external_deposit.go +++ b/ocp/worker/geyser/external_deposit.go @@ -406,15 +406,13 @@ func processPotentialExternalDepositIntoVm(ctx context.Context, data ocp_data.Pr return errors.Wrap(err, "error saving external deposit record") } - if balance_util.LedgerWritesEnabled(ctx) { - balanceDeltas, err := balance_util.DeltasForExternalDeposit(intentRecord) - if err != nil { - return errors.Wrap(err, "error building balance deltas") - } - err = balance_util.ApplyDeltasInTx(ctx, data, balanceDeltas...) - if err != nil { - return errors.Wrap(err, "error applying balance deltas") - } + balanceDeltas, err := balance_util.DeltasForExternalDeposit(intentRecord) + if err != nil { + return errors.Wrap(err, "error building balance deltas") + } + err = balance_util.ApplyDeltasInTx(ctx, data, balanceDeltas...) + if err != nil { + return errors.Wrap(err, "error applying balance deltas") } return nil diff --git a/ocp/worker/swap/util.go b/ocp/worker/swap/util.go index 85e764a..e154c5f 100644 --- a/ocp/worker/swap/util.go +++ b/ocp/worker/swap/util.go @@ -465,15 +465,13 @@ func (p *runtime) markSwapCancelled(ctx context.Context, swapRecord *swap.Record return err } - if balance.LedgerWritesEnabled(ctx) { - balanceDeltas, err := balance.DeltasForExternalDeposit(refundIntentRecord) - if err != nil { - return err - } - err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) - if err != nil { - return err - } + balanceDeltas, err := balance.DeltasForExternalDeposit(refundIntentRecord) + if err != nil { + return err + } + err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) + if err != nil { + return err } // The swap was funded and entered transaction history, so its @@ -854,28 +852,26 @@ func (p *runtime) maybeUpdateBalancesForFinalizedReserveSwap(ctx context.Context } } - if balance.LedgerWritesEnabled(ctx) { - balanceDeltas, err := balance.DeltasForExternalDeposit(intentRecord) + balanceDeltas, err := balance.DeltasForExternalDeposit(intentRecord) + if err != nil { + return err + } + + if reconciledFundingIntentRecord != nil { + fundingActionRecords, err := p.data.GetAllActionsByIntent(ctx, reconciledFundingIntentRecord.IntentId) if err != nil { return err } - - if reconciledFundingIntentRecord != nil { - fundingActionRecords, err := p.data.GetAllActionsByIntent(ctx, reconciledFundingIntentRecord.IntentId) - if err != nil { - return err - } - reconciliationDeltas, err := balance.DeltasForSwapSellReconciliation(previousFundingIntentRecord, reconciledFundingIntentRecord, fundingActionRecords) - if err != nil { - return err - } - balanceDeltas = append(balanceDeltas, reconciliationDeltas...) - } - - err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) + reconciliationDeltas, err := balance.DeltasForSwapSellReconciliation(previousFundingIntentRecord, reconciledFundingIntentRecord, fundingActionRecords) if err != nil { return err } + balanceDeltas = append(balanceDeltas, reconciliationDeltas...) + } + + err = balance.ApplyDeltasInTx(ctx, p.data, balanceDeltas...) + if err != nil { + return err } return nil From e742c1998b39c1ab04255cce08f82526ae2e4d85 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 13:56:25 -0400 Subject: [PATCH 22/26] Remove backfill --- ocp/balance/calculator.go | 19 +--- ocp/balance/calculator_test.go | 44 ++------- ocp/balance/ledger.go | 30 ++---- ocp/balance/ledger_test.go | 56 ++++++----- ocp/data/balance/memory/store.go | 81 ++++----------- ocp/data/balance/postgres/model.go | 83 ++++------------ ocp/data/balance/postgres/store.go | 5 - ocp/data/balance/record.go | 29 ++---- ocp/data/balance/store.go | 47 +-------- ocp/data/balance/tests/tests.go | 152 +---------------------------- ocp/data/internal.go | 4 - ocp/worker/account/testutil.go | 11 +++ 12 files changed, 116 insertions(+), 445 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index f12dbdd..6a7e4ac 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -28,20 +28,13 @@ const ( var ( // ErrNegativeBalance indicates that a ledger record holds a negative - // value, which can only happen before it's backfilled. + // value, which the store's delta predicates should make impossible. ErrNegativeBalance = errors.New("balance calculation resulted in negative value") // ErrNotManagedByCode indicates that an account is not owned by Code. // It's up to callers to determine how to handle this situation within // the context of a balance. ErrNotManagedByCode = errors.New("explicitly not handling account not managed by code") - - // ErrIncompleteRecord indicates a ledger record doesn't yet reflect the - // full history of the account, so it isn't an authoritative balance. - // - // todo: Remove alongside balance.Record.IsBackfilled once the ledger - // backfill is retired. - ErrIncompleteRecord = errors.New("balance record is not backfilled") ) // CalculateFromCache is the default and recommended strategy for reliably estimating @@ -359,14 +352,10 @@ func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Prov } // checkRecord verifies a ledger record is an authoritative view of an account -// that Code still manages. +// that Code still manages. Quark balance callers reject unlocked vaults on the +// timelock record before reaching here, so for them this only guards against a +// record that disagrees with it. func checkRecord(record *balance.Record) error { - if !record.IsBackfilled { - return ErrIncompleteRecord - } - - // Callers reject unlocked vaults on the timelock record before reaching - // here, so this only guards against a record that disagrees with it. if !record.IsLocked { return ErrNotManagedByCode } diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 13b8f1d..7dd9141 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -23,8 +23,8 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { first := newBalanceTestAccount(t, env) second := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true, IsBackfilled: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 0, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 0, IsOpen: true, IsLocked: true}) expected := map[string]uint64{ first.tokenAccount.PublicKey().ToBase58(): 42, @@ -69,35 +69,11 @@ func TestDefaultCalculationMethods_MissingBalanceRecord(t *testing.T) { assert.Equal(t, balance.ErrRecordNotFound, err) } -func TestDefaultCalculationMethods_IncompleteBalanceRecord(t *testing.T) { - env := setupBalanceTestEnv(t) - - // A record that isn't backfilled only holds deltas observed since it was - // seeded, so it can't answer for the account - testAccount := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) - - _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrIncompleteRecord, err) - - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrIncompleteRecord, err) - - _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) - assert.Equal(t, ErrIncompleteRecord, err) - - _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrIncompleteRecord, err) - - _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrIncompleteRecord, err) -} - func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { env := setupBalanceTestEnv(t) testAccount := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) timelockRecord, err := env.data.GetTimelockByVault(env.ctx, testAccount.tokenAccount.PublicKey().ToBase58()) require.NoError(t, err) @@ -118,13 +94,13 @@ func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) - // A backfilled record for an unlocked vault holds the last managed state, - // not a live balance, so it is refused even though the timelock record + // A record for an unlocked vault holds the last managed state, not a live + // balance, so it is refused even though the timelock record // still passes the managed check. That pairing is inconsistent by // construction: the timelock check normally rejects first, so the fixture // exists to exercise the record's own guard. testAccount := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true}) _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) @@ -146,8 +122,8 @@ func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) { second := newBalanceTestAccount(t, env) // Both values come from the same record, so they can't disagree - saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true, IsBackfilled: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 33, UsdCostBasis: -123456, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 33, UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) res, err := BatchCalculateWithUsdCostBasisFromCache(env.ctx, env.data, first.accountRecords(t, env), second.accountRecords(t, env)) require.NoError(t, err) @@ -170,8 +146,8 @@ func TestUsdCostBasisCalculationMethods(t *testing.T) { first := newBalanceTestAccount(t, env) second := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, first, &balance.Record{UsdCostBasis: -123456, IsOpen: true, IsLocked: true, IsBackfilled: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 1, UsdCostBasis: 1_500_000, IsOpen: true, IsLocked: true, IsBackfilled: true}) + saveBalanceTestRecord(t, env, first, &balance.Record{UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 1, UsdCostBasis: 1_500_000, IsOpen: true, IsLocked: true}) expected := map[string]int64{ first.tokenAccount.PublicKey().ToBase58(): -123456, diff --git a/ocp/balance/ledger.go b/ocp/balance/ledger.go index aec6fd0..84eb7cc 100644 --- a/ocp/balance/ledger.go +++ b/ocp/balance/ledger.go @@ -29,9 +29,8 @@ var ErrUntrackedAccount = errors.New("account is not tracked by the balance ledg // Credits are still applied, since an unlocked record is excluded from every // read and turning one away only blocks the flow recording it. // -// Any timelock account in the delta set that has no ledger record yet lazily -// gets one that is not backfilled, so accounts that predate the ledger start -// accumulating deltas on first touch regardless of direction. +// A timelock account with no ledger record at all is balance.ErrRecordNotFound +// in either direction, since every one gets a record when it's opened. // // Store predicate failures (balance.ErrInsufficientBalance, // balance.ErrBalanceChanged, balance.ErrAccountClosed, @@ -69,8 +68,7 @@ func ApplyDeltasInTx(ctx context.Context, data ocp_data.Provider, deltas ...*bal // CreateRecordInTx creates the ledger record for a newly opened account. It // must be called within the DB transaction that creates the account info -// record. A new account has no history, so its record is created backfilled -// at zero and predicates are enforced from the start. +// record, so every timelock account has a record from the moment it exists. // // It is a no-op for accounts that aren't timelock accounts, which the ledger // doesn't track. @@ -85,7 +83,6 @@ func CreateRecordInTx(ctx context.Context, data ocp_data.Provider, accountInfoRe MintAccount: accountInfoRecord.MintAccount, IsOpen: true, IsLocked: true, - IsBackfilled: true, }) if errors.Is(err, balance.ErrRecordExists) { return nil @@ -93,9 +90,11 @@ func CreateRecordInTx(ctx context.Context, data ocp_data.Provider, accountInfoRe return err } -// resolveRecords reports which accounts in the delta set the ledger tracks, -// creating a non-backfilled record for every timelock account that doesn't -// have one yet. +// resolveRecords reports which accounts in the delta set the ledger tracks. A +// timelock account without a record is a broken invariant rather than an +// untracked account, since CreateRecordInTx gives every one a record when it's +// opened, so it fails with balance.ErrRecordNotFound instead of being seeded +// with a balance that has no relationship to the account's history. func resolveRecords(ctx context.Context, data ocp_data.Provider, deltas []*balance.Delta) (map[string]bool, error) { tracked := make(map[string]bool) var tokenAccounts []string @@ -128,18 +127,7 @@ func resolveRecords(ctx context.Context, data ocp_data.Provider, deltas []*balan continue } - err = data.CreateBalance(ctx, &balance.Record{ - TokenAccount: accountInfoRecord.TokenAccount, - OwnerAccount: accountInfoRecord.OwnerAccount, - MintAccount: accountInfoRecord.MintAccount, - IsOpen: true, - IsLocked: true, - IsBackfilled: false, - }) - if err != nil && !errors.Is(err, balance.ErrRecordExists) { - return nil, err - } - tracked[tokenAccount] = true + return nil, fmt.Errorf("%w: %s", balance.ErrRecordNotFound, tokenAccount) } return tracked, nil } diff --git a/ocp/balance/ledger_test.go b/ocp/balance/ledger_test.go index c842e46..a1eeefe 100644 --- a/ocp/balance/ledger_test.go +++ b/ocp/balance/ledger_test.go @@ -15,36 +15,36 @@ import ( "github.com/code-payments/ocp-server/testutil" ) -func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { +func TestApplyDeltasInTx_TrackedAccounts(t *testing.T) { ctx := context.Background() data := ocp_data.NewTestDataProvider() - source := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) - destination := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_REMOTE_SEND_GIFT_CARD) + source := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY) + destination := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_REMOTE_SEND_GIFT_CARD) swap := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_SWAP) external := testutil.NewRandomAccount(t).PublicKey().ToBase58() + require.NoError(t, CreateRecordInTx(ctx, data, source)) + require.NoError(t, CreateRecordInTx(ctx, data, destination)) + require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source.TokenAccount, Kind: balance.DeltaCredit, Quarks: 100, UsdCostBasis: 1_000_000})) + require.NoError(t, ApplyDeltasInTx(ctx, data, - &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 100, UsdCostBasis: 1_000_000}, - &balance.Delta{TokenAccount: destination, Kind: balance.DeltaCredit, Quarks: 60, UsdCostBasis: 600_000}, + &balance.Delta{TokenAccount: source.TokenAccount, Kind: balance.DeltaDebit, Quarks: 100, UsdCostBasis: 1_000_000}, + &balance.Delta{TokenAccount: destination.TokenAccount, Kind: balance.DeltaCredit, Quarks: 60, UsdCostBasis: 600_000}, &balance.Delta{TokenAccount: swap, Kind: balance.DeltaCredit, Quarks: 20, UsdCostBasis: 200_000}, &balance.Delta{TokenAccount: external, Kind: balance.DeltaCredit, Quarks: 20, UsdCostBasis: 200_000}, )) - // Existing accounts get a non-backfilled row that accumulates freely, - // including a negative balance for a source that predates the ledger - record, err := data.GetBalance(ctx, source) + record, err := data.GetBalance(ctx, source.TokenAccount) require.NoError(t, err) - assert.EqualValues(t, -100, record.Quarks) - assert.EqualValues(t, -1_000_000, record.UsdCostBasis) - assert.False(t, record.IsBackfilled) + assert.EqualValues(t, 0, record.Quarks) + assert.EqualValues(t, 0, record.UsdCostBasis) assert.True(t, record.IsOpen) - record, err = data.GetBalance(ctx, destination) + record, err = data.GetBalance(ctx, destination.TokenAccount) require.NoError(t, err) assert.EqualValues(t, 60, record.Quarks) assert.EqualValues(t, 600_000, record.UsdCostBasis) - assert.False(t, record.IsBackfilled) // Credits to accounts OCP doesn't hold a timelock for are dropped, and // those accounts never get a row @@ -53,18 +53,24 @@ func TestApplyDeltasInTx_SeedsTimelockAccounts(t *testing.T) { _, err = data.GetBalance(ctx, external) assert.Equal(t, balance.ErrRecordNotFound, err) - // Once backfilled, predicates are enforced - require.NoError(t, data.BackfillBalance(ctx, source, func(context.Context) (*balance.BackfillResult, error) { - return &balance.BackfillResult{Quarks: 400, UsdCostBasis: 4_000_000, IsOpen: true, IsLocked: true}, nil - })) - err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 401}) + // Predicates are enforced against every record + err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: destination.TokenAccount, Kind: balance.DeltaDebit, Quarks: 61}) assert.Equal(t, balance.ErrInsufficientBalance, err) - require.NoError(t, ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: source, Kind: balance.DeltaDebit, Quarks: 400, UsdCostBasis: 4_000_000})) +} - record, err = data.GetBalance(ctx, source) - require.NoError(t, err) - assert.EqualValues(t, 0, record.Quarks) - assert.True(t, record.IsBackfilled) +func TestApplyDeltasInTx_MissingRecord(t *testing.T) { + ctx := context.Background() + data := ocp_data.NewTestDataProvider() + + // A timelock account with no record is a broken invariant in either + // direction, not an account the ledger doesn't track + tokenAccount := newLedgerTestAccount(t, ctx, data, commonpb.AccountType_PRIMARY) + for _, delta := range []*balance.Delta{ + {TokenAccount: tokenAccount, Kind: balance.DeltaCredit, Quarks: 1}, + {TokenAccount: tokenAccount, Kind: balance.DeltaDebit, Quarks: 1}, + } { + assert.ErrorIs(t, ApplyDeltasInTx(ctx, data, delta), balance.ErrRecordNotFound) + } } func TestApplyDeltasInTx_UnknownSource(t *testing.T) { @@ -130,8 +136,7 @@ func TestCreateRecordInTx(t *testing.T) { primary := newLedgerTestAccountInfo(t, ctx, data, commonpb.AccountType_PRIMARY) - // A new timelock account starts backfilled at zero, so predicates are - // enforced immediately + // A new timelock account starts at zero, with predicates enforced require.NoError(t, CreateRecordInTx(ctx, data, primary)) record, err := data.GetBalance(ctx, primary.TokenAccount) require.NoError(t, err) @@ -141,7 +146,6 @@ func TestCreateRecordInTx(t *testing.T) { assert.EqualValues(t, 0, record.Quarks) assert.EqualValues(t, 0, record.UsdCostBasis) assert.True(t, record.IsOpen) - assert.True(t, record.IsBackfilled) err = ApplyDeltasInTx(ctx, data, &balance.Delta{TokenAccount: primary.TokenAccount, Kind: balance.DeltaDebit, Quarks: 1}) assert.Equal(t, balance.ErrInsufficientBalance, err) diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 179e20c..0035070 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -142,7 +142,7 @@ func (s *store) CountLockedByMint(_ context.Context, mint string, minQuarks int6 var res uint64 for _, item := range s.balanceRecordsByTokenAccount { - if item.MintAccount == mint && item.Quarks >= minQuarks && item.IsLocked && item.IsBackfilled { + if item.MintAccount == mint && item.Quarks >= minQuarks && item.IsLocked { res++ } } @@ -206,8 +206,6 @@ func (s *store) ApplyDeltas(_ context.Context, deltas ...*balance.Delta) error { } func applyDelta(item *balance.Record, delta *balance.Delta) error { - enforce := item.IsBackfilled - // A credit doesn't require the vault to be locked, since an unlocked // record is excluded from every read anyway and turning the credit away // only blocks the flow recording it. A cost basis adjustment carries no @@ -215,50 +213,43 @@ func applyDelta(item *balance.Record, delta *balance.Delta) error { switch delta.Kind { case balance.DeltaCredit, balance.DeltaAdjustUsdCostBasis: default: - if enforce && !item.IsLocked { + if !item.IsLocked { return balance.ErrAccountUnlocked } } switch delta.Kind { case balance.DeltaCredit: - if enforce && !item.IsOpen { + if !item.IsOpen { return balance.ErrAccountClosed } item.Quarks += int64(delta.Quarks) item.UsdCostBasis += delta.UsdCostBasis case balance.DeltaDebit: - if enforce && !item.IsOpen { + if !item.IsOpen { return balance.ErrAccountClosed } - if enforce && item.Quarks < int64(delta.Quarks) { + if item.Quarks < int64(delta.Quarks) { return balance.ErrInsufficientBalance } item.Quarks -= int64(delta.Quarks) item.UsdCostBasis -= delta.UsdCostBasis case balance.DeltaDrain: - if enforce { - if !item.IsOpen { - return balance.ErrAccountClosed - } - if item.Quarks != int64(delta.Quarks) { - return balance.ErrBalanceChanged - } - item.Quarks = 0 - item.UsdCostBasis = 0 - } else { - item.Quarks -= int64(delta.Quarks) - item.UsdCostBasis -= delta.UsdCostBasis + if !item.IsOpen { + return balance.ErrAccountClosed } + if item.Quarks != int64(delta.Quarks) { + return balance.ErrBalanceChanged + } + item.Quarks = 0 + item.UsdCostBasis = 0 item.IsOpen = false case balance.DeltaClose: - if enforce { - if !item.IsOpen { - return balance.ErrAccountClosed - } - if item.Quarks != 0 { - return balance.ErrBalanceChanged - } + if !item.IsOpen { + return balance.ErrAccountClosed + } + if item.Quarks != 0 { + return balance.ErrBalanceChanged } item.IsOpen = false case balance.DeltaAdjustUsdCostBasis: @@ -267,44 +258,6 @@ func applyDelta(item *balance.Record, delta *balance.Delta) error { return nil } -// Backfill implements balance.Store.Backfill -// -// Note: The lock is released while fn runs, since fn reads from other stores -// sharing the provider. Unlike the DB store, this doesn't block concurrent -// deltas, which tests don't exercise against a backfill. -func (s *store) Backfill(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { - s.mu.Lock() - item, ok := s.balanceRecordsByTokenAccount[tokenAccount] - if !ok { - s.mu.Unlock() - return balance.ErrRecordNotFound - } - if item.IsBackfilled { - s.mu.Unlock() - return balance.ErrAlreadyBackfilled - } - s.mu.Unlock() - - result, err := fn(ctx) - if err != nil { - return err - } - if result.Quarks < 0 { - return balance.ErrNegativeBalance - } - - s.mu.Lock() - defer s.mu.Unlock() - - item.Quarks = result.Quarks - item.UsdCostBasis = result.UsdCostBasis - item.IsOpen = result.IsOpen - item.IsLocked = result.IsLocked - item.IsBackfilled = true - item.UpdatedAt = time.Now() - return nil -} - func (s *store) filter(fn func(*balance.Record) bool) ([]*balance.Record, error) { var res []*balance.Record for _, item := range s.balanceRecords { diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index 9a29133..6890ffa 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -17,7 +17,10 @@ const ( tableName = "ocp__core_balance" externalCheckpointTableName = "ocp__core_externalbalancecheckpoint" - allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_locked, is_backfilled, updated_at" + // Note: is_backfilled is deliberately absent. The column still exists in + // prod, where every row is backfilled, so it is written as TRUE on insert + // and never read. + allColumns = "id, token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_locked, updated_at" ) type model struct { @@ -30,9 +33,8 @@ type model struct { Quarks int64 `db:"quarks"` UsdCostBasis int64 `db:"usd_cost_basis"` - IsOpen bool `db:"is_open"` - IsLocked bool `db:"is_locked"` - IsBackfilled bool `db:"is_backfilled"` + IsOpen bool `db:"is_open"` + IsLocked bool `db:"is_locked"` UpdatedAt time.Time `db:"updated_at"` } @@ -50,9 +52,8 @@ func toModel(obj *balance.Record) (*model, error) { Quarks: obj.Quarks, UsdCostBasis: obj.UsdCostBasis, - IsOpen: obj.IsOpen, - IsLocked: obj.IsLocked, - IsBackfilled: obj.IsBackfilled, + IsOpen: obj.IsOpen, + IsLocked: obj.IsLocked, UpdatedAt: obj.UpdatedAt, }, nil @@ -69,9 +70,8 @@ func fromModel(obj *model) *balance.Record { Quarks: obj.Quarks, UsdCostBasis: obj.UsdCostBasis, - IsOpen: obj.IsOpen, - IsLocked: obj.IsLocked, - IsBackfilled: obj.IsBackfilled, + IsOpen: obj.IsOpen, + IsLocked: obj.IsLocked, UpdatedAt: obj.UpdatedAt, } @@ -81,7 +81,7 @@ func (m *model) dbCreate(ctx context.Context, db *sqlx.DB) error { return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { query := `INSERT INTO ` + tableName + ` (token_account, owner_account, mint_account, quarks, usd_cost_basis, is_open, is_locked, is_backfilled, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + VALUES ($1, $2, $3, $4, $5, $6, $7, TRUE, $8) RETURNING ` + allColumns m.UpdatedAt = time.Now() @@ -96,7 +96,6 @@ func (m *model) dbCreate(ctx context.Context, db *sqlx.DB) error { m.UsdCostBasis, m.IsOpen, m.IsLocked, - m.IsBackfilled, m.UpdatedAt.UTC(), ).StructScan(m) @@ -183,7 +182,7 @@ func dbGetAllLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuar func dbCountLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64) (uint64, error) { var res uint64 query := `SELECT COUNT(*) FROM ` + tableName + ` - WHERE mint_account = $1 AND quarks >= $2 AND is_locked AND is_backfilled` + WHERE mint_account = $1 AND quarks >= $2 AND is_locked` err := db.GetContext(ctx, &res, query, mint, minQuarks) if err != nil { return 0, err @@ -213,7 +212,7 @@ func dbMarkAsUnlocked(ctx context.Context, db *sqlx.DB, tokenAccount string) err // dbApplyDeltas applies every delta in a single transaction. Each delta is one // conditional UPDATE, so its predicate is evaluated against the row after the -// row lock is acquired. Predicates only apply to backfilled rows. +// row lock is acquired. func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) error { return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { for _, delta := range deltas { @@ -223,25 +222,22 @@ func dbApplyDeltas(ctx context.Context, db *sqlx.DB, deltas []*balance.Delta) er case balance.DeltaCredit: query = `UPDATE ` + tableName + ` SET quarks = quarks + $2, usd_cost_basis = usd_cost_basis + $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR is_open)` + WHERE token_account = $1 AND is_open` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDebit: query = `UPDATE ` + tableName + ` SET quarks = quarks - $2, usd_cost_basis = usd_cost_basis - $3, updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks >= $2))` + WHERE token_account = $1 AND is_open AND is_locked AND quarks >= $2` args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} case balance.DeltaDrain: query = `UPDATE ` + tableName + ` - SET quarks = CASE WHEN is_backfilled THEN 0 ELSE quarks - $2 END, - usd_cost_basis = CASE WHEN is_backfilled THEN 0 ELSE usd_cost_basis - $3 END, - is_open = FALSE, - updated_at = $4 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks = $2))` - args = []any{delta.TokenAccount, int64(delta.Quarks), delta.UsdCostBasis, time.Now().UTC()} + SET quarks = 0, usd_cost_basis = 0, is_open = FALSE, updated_at = $3 + WHERE token_account = $1 AND is_open AND is_locked AND quarks = $2` + args = []any{delta.TokenAccount, int64(delta.Quarks), time.Now().UTC()} case balance.DeltaClose: query = `UPDATE ` + tableName + ` SET is_open = FALSE, updated_at = $2 - WHERE token_account = $1 AND (NOT is_backfilled OR (is_open AND is_locked AND quarks = 0))` + WHERE token_account = $1 AND is_open AND is_locked AND quarks = 0` args = []any{delta.TokenAccount, time.Now().UTC()} case balance.DeltaAdjustUsdCostBasis: // No predicate: no quarks move, so nothing the other kinds @@ -305,47 +301,6 @@ func classifyFailedDelta(delta *balance.Delta, current *balance.Record) error { return fmt.Errorf("unsupported delta kind: %s", delta.Kind) } -func dbBackfill(ctx context.Context, db *sqlx.DB, tokenAccount string, fn balance.BackfillFunc) error { - return executeTxWithinCtxOrJoin(ctx, db, func(ctx context.Context) error { - return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { - var current model - err := tx.GetContext(ctx, ¤t, `SELECT `+allColumns+` FROM `+tableName+` WHERE token_account = $1 FOR UPDATE`, tokenAccount) - if err != nil { - return pgutil.CheckNoRows(err, balance.ErrRecordNotFound) - } - if current.IsBackfilled { - return balance.ErrAlreadyBackfilled - } - - // The row lock is held across fn, so it observes every committed - // write to the account and blocks every in-flight one. - result, err := fn(ctx) - if err != nil { - return err - } - if result.Quarks < 0 { - return balance.ErrNegativeBalance - } - - query := `UPDATE ` + tableName + ` - SET quarks = $2, usd_cost_basis = $3, is_open = $4, is_locked = $5, is_backfilled = TRUE, updated_at = $6 - WHERE token_account = $1` - _, err = tx.ExecContext(ctx, query, tokenAccount, result.Quarks, result.UsdCostBasis, result.IsOpen, result.IsLocked, time.Now().UTC()) - return err - }) - }) -} - -// executeTxWithinCtxOrJoin runs fn with a context carrying a DB transaction, -// starting one if the context doesn't already have one. -func executeTxWithinCtxOrJoin(ctx context.Context, db *sqlx.DB, fn func(ctx context.Context) error) error { - err := pgutil.ExecuteTxWithinCtx(ctx, db, sql.LevelDefault, fn) - if err == pgutil.ErrAlreadyInTx { - return fn(ctx) - } - return err -} - type externalCheckpointModel struct { Id sql.NullInt64 `db:"id"` diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index 67d5b96..d2c6178 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -107,11 +107,6 @@ func (s *store) ApplyDeltas(ctx context.Context, deltas ...*balance.Delta) error return dbApplyDeltas(ctx, s.db, balance.MergeDeltas(deltas)) } -// Backfill implements balance.Store.Backfill -func (s *store) Backfill(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { - return dbBackfill(ctx, s.db, tokenAccount, fn) -} - func fromModels(models []*model) []*balance.Record { res := make([]*balance.Record, len(models)) for i, model := range models { diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index c82f580..fe3f4a4 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -33,9 +33,8 @@ type Record struct { OwnerAccount string MintAccount string - // Quarks is signed because a record that has not been backfilled only - // accumulates deltas, which may temporarily net negative. Backfilled - // records are guaranteed to be non-negative. + // Quarks is signed to match the column's type. Delta predicates keep it + // non-negative. Quarks int64 // UsdCostBasis is the account's USD cost basis, in UsdQuarksPerUnit. @@ -53,11 +52,6 @@ type Record struct { // one-way. IsLocked bool - // IsBackfilled indicates the record reflects the full history of the - // account. Until it does, deltas are recorded without enforcing any - // balance predicates. - IsBackfilled bool - UpdatedAt time.Time } @@ -74,8 +68,8 @@ func (r *Record) Validate() error { return errors.New("mint account is required") } - if r.IsBackfilled && r.Quarks < 0 { - return errors.New("backfilled quarks cannot be negative") + if r.Quarks < 0 { + return errors.New("quarks cannot be negative") } return nil @@ -92,9 +86,8 @@ func (r *Record) Clone() Record { Quarks: r.Quarks, UsdCostBasis: r.UsdCostBasis, - IsOpen: r.IsOpen, - IsLocked: r.IsLocked, - IsBackfilled: r.IsBackfilled, + IsOpen: r.IsOpen, + IsLocked: r.IsLocked, UpdatedAt: r.UpdatedAt, } @@ -112,13 +105,11 @@ func (r *Record) CopyTo(dst *Record) { dst.IsOpen = r.IsOpen dst.IsLocked = r.IsLocked - dst.IsBackfilled = r.IsBackfilled dst.UpdatedAt = r.UpdatedAt } -// DeltaKind selects the predicate a Delta is applied under. Predicates are -// only enforced on backfilled records. +// DeltaKind selects the predicate a Delta is applied under. type DeltaKind uint8 const ( @@ -154,9 +145,9 @@ type Delta struct { // UsdCostBasis is added on credit and subtracted on debit. It is signed // so that a credit can also carry a downward reconciliation. Ignored for - // DeltaDrain and DeltaClose on backfilled records, where the basis is - // zeroed along with the balance. For DeltaAdjustUsdCostBasis it is the - // signed correction, added as is. + // DeltaDrain and DeltaClose, where the basis is zeroed along with the + // balance. For DeltaAdjustUsdCostBasis it is the signed correction, added + // as is. UsdCostBasis int64 } diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index b427b17..b904fb4 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -18,13 +18,6 @@ var ( // balance than the one on record. ErrBalanceChanged = errors.New("balance is not the expected value") - ErrAlreadyBackfilled = errors.New("balance record is already backfilled") - - // ErrNegativeBalance is returned when a backfill computes a negative - // balance, which indicates inconsistent historical data that must be - // reviewed rather than recorded. - ErrNegativeBalance = errors.New("backfilled balance is negative") - ErrAccountClosed = errors.New("account open state is stale") // ErrAccountUnlocked is returned when a delta other than a credit targets @@ -38,26 +31,6 @@ var ( ErrStaleCheckpoint = errors.New("checkpoint is stale") ) -// BackfillResult is the full historical state of a token account. -type BackfillResult struct { - Quarks int64 - UsdCostBasis int64 - - // IsOpen is false for accounts that can no longer receive funds, such as - // claimed gift cards and distributed pools. - IsOpen bool - - // IsLocked is false for accounts whose timelock vault has unlocked, so - // the balance is the last managed state rather than a live value. - IsLocked bool -} - -// BackfillFunc computes the full historical state of a token account. It -// is called while the record is locked, with a context that is part of the -// same DB transaction, so any store reads made through it observe every -// committed change and block every in-flight one. -type BackfillFunc func(ctx context.Context) (*BackfillResult, error) - type Store interface { // Create creates a new balance record. // @@ -91,19 +64,14 @@ type Store interface { // ErrRecordNotFound is returned if no records exist. GetAllLockedByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) - // CountLockedByMint counts locked, backfilled records for a mint with at - // least minQuarks. Records that are not backfilled are excluded, since - // their balances are partial sums that can't be compared against a - // threshold. Unlocked records are excluded, since their balances are + // CountLockedByMint counts locked records for a mint with at least + // minQuarks. Unlocked records are excluded, since their balances are // stale. CountLockedByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) // ApplyDeltas atomically applies a set of deltas. Either every delta is // applied or none are. Deltas are applied in SortDeltas order. // - // Predicates are enforced only on backfilled records; records that are - // not backfilled simply accumulate the change. - // // Every delta must target an account with a record, otherwise // ErrRecordNotFound is returned and nothing is applied. Callers are // responsible for not producing deltas for accounts the ledger doesn't @@ -127,17 +95,6 @@ type Store interface { // ErrRecordNotFound is returned if no record exists. MarkAsUnlocked(ctx context.Context, tokenAccount string) error - // Backfill locks a record that is not yet backfilled, calls fn to compute - // its full historical balance, and overwrites the record with the result, - // marking it as backfilled. Deltas recorded before the backfill are - // intentionally discarded, since fn observes them. - // - // ErrRecordNotFound is returned if no record exists. ErrAlreadyBackfilled - // is returned if the record is already backfilled, in which case fn is - // not called. ErrNegativeBalance is returned if fn computes a negative - // balance, leaving the record untouched. - Backfill(ctx context.Context, tokenAccount string, fn BackfillFunc) error - // SaveExternalCheckpoint saves an external balance at a checkpoint. // // ErrStaleCheckpoint is returned if the checkpoint is outdated diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index bb569dd..dbcf6ed 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -17,11 +17,9 @@ func RunTests(t *testing.T, s balance.Store, teardown func()) { for _, tf := range []func(t *testing.T, s balance.Store){ testRecordHappyPath, testGetAllLockedByMint, - testApplyDeltasBackfilled, - testApplyDeltasNotBackfilled, + testApplyDeltas, testApplyDeltasAtomicity, testApplyDeltasConcurrency, - testBackfill, testMarkAsUnlocked, testExternalCheckpointHappyPath, } { @@ -57,7 +55,6 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { UsdCostBasis: 200, IsOpen: true, IsLocked: true, - IsBackfilled: true, } cloned := expected.Clone() @@ -111,7 +108,6 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { OwnerAccount: "owner_1", MintAccount: "mint_1", Quarks: -1, - IsBackfilled: true, })) }) } @@ -131,7 +127,6 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { Quarks: int64(i * 10), IsOpen: true, IsLocked: true, - IsBackfilled: true, })) } require.NoError(t, s.Create(ctx, &balance.Record{ @@ -141,7 +136,6 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { Quarks: 1000, IsOpen: true, IsLocked: true, - IsBackfilled: true, })) records, err := s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) @@ -186,7 +180,6 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { MintAccount: "mint_1", Quarks: 1000, IsOpen: true, - IsBackfilled: true, })) records, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.EmptyCursor, 10, query.Ascending) require.NoError(t, err) @@ -195,18 +188,7 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { _, err = s.GetAllLockedByMint(ctx, "mint_1", 0, query.ToCursor(5), 10, query.Ascending) assert.Equal(t, balance.ErrRecordNotFound, err) - // Counting by mint uses the same threshold semantics, but only over - // backfilled records - require.NoError(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_not_backfilled", - OwnerAccount: "owner", - MintAccount: "mint_1", - Quarks: 1000, - IsOpen: true, - IsLocked: true, - IsBackfilled: false, - })) - + // Counting by mint uses the same threshold semantics count, err := s.CountLockedByMint(ctx, "mint_1", 0) require.NoError(t, err) assert.EqualValues(t, 5, count) @@ -225,8 +207,8 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { }) } -func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { - t.Run("testApplyDeltasBackfilled", func(t *testing.T) { +func testApplyDeltas(t *testing.T, s balance.Store) { + t.Run("testApplyDeltas", func(t *testing.T) { ctx := context.Background() require.NoError(t, s.Create(ctx, &balance.Record{ @@ -235,7 +217,6 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { MintAccount: "mint", IsOpen: true, IsLocked: true, - IsBackfilled: true, })) // Every kind of delta requires a record @@ -312,45 +293,12 @@ func testApplyDeltasBackfilled(t *testing.T, s balance.Store) { MintAccount: "mint", IsOpen: true, IsLocked: true, - IsBackfilled: true, })) require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaClose})) assertBalance(t, s, "token_account_2", 0, 0, false) }) } -func testApplyDeltasNotBackfilled(t *testing.T, s balance.Store) { - t.Run("testApplyDeltasNotBackfilled", func(t *testing.T) { - ctx := context.Background() - - require.NoError(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_1", - OwnerAccount: "owner", - MintAccount: "mint", - IsOpen: true, - IsLocked: true, - })) - - // No predicates are enforced, and the balance can go negative - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 30, UsdCostBasis: 10})) - assertBalance(t, s, "token_account_1", -30, -10, true) - - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 100, UsdCostBasis: 50})) - assertBalance(t, s, "token_account_1", 70, 40, true) - - // A drain records the delta rather than zeroing, but still closes - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDrain, Quarks: 100, UsdCostBasis: 60})) - assertBalance(t, s, "token_account_1", -30, -20, false) - - // Closed accounts still accumulate - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaCredit, Quarks: 5, UsdCostBasis: 5})) - assertBalance(t, s, "token_account_1", -25, -15, false) - - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaClose})) - assertBalance(t, s, "token_account_1", -25, -15, false) - }) -} - func testApplyDeltasAtomicity(t *testing.T, s balance.Store) { t.Run("testApplyDeltasAtomicity", func(t *testing.T) { ctx := context.Background() @@ -363,7 +311,6 @@ func testApplyDeltasAtomicity(t *testing.T, s balance.Store) { Quarks: 100, IsOpen: true, IsLocked: true, - IsBackfilled: true, })) } @@ -436,7 +383,6 @@ func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { Quarks: initialBalance, IsOpen: true, IsLocked: true, - IsBackfilled: true, })) require.NoError(t, s.Create(ctx, &balance.Record{ TokenAccount: "receiver", @@ -444,7 +390,6 @@ func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { MintAccount: "mint", IsOpen: true, IsLocked: true, - IsBackfilled: true, })) // Concurrent sends of 1 quark each: exactly initialBalance succeed, and @@ -524,80 +469,6 @@ func testApplyDeltasConcurrency(t *testing.T, s balance.Store) { }) } -func testBackfill(t *testing.T, s balance.Store) { - t.Run("testBackfill", func(t *testing.T) { - ctx := context.Background() - - fn := func(quarks, usdCostBasis int64, isOpen bool) balance.BackfillFunc { - return func(ctx context.Context) (*balance.BackfillResult, error) { - return &balance.BackfillResult{Quarks: quarks, UsdCostBasis: usdCostBasis, IsOpen: isOpen, IsLocked: true}, nil - } - } - - assert.Equal(t, balance.ErrRecordNotFound, s.Backfill(ctx, "token_account_1", fn(1, 1, true))) - - require.NoError(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_1", - OwnerAccount: "owner", - MintAccount: "mint", - IsOpen: true, - IsLocked: true, - })) - - // Deltas recorded before the backfill are discarded by it - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 30, UsdCostBasis: 10})) - assertBalance(t, s, "token_account_1", -30, -10, true) - - // A failing computation leaves the record untouched - assert.Error(t, s.Backfill(ctx, "token_account_1", func(ctx context.Context) (*balance.BackfillResult, error) { - return nil, assert.AnError - })) - record, err := s.Get(ctx, "token_account_1") - require.NoError(t, err) - assert.False(t, record.IsBackfilled) - assert.EqualValues(t, -30, record.Quarks) - - // So does a negative computed balance - assert.Equal(t, balance.ErrNegativeBalance, s.Backfill(ctx, "token_account_1", fn(-1, 0, true))) - record, err = s.Get(ctx, "token_account_1") - require.NoError(t, err) - assert.False(t, record.IsBackfilled) - assert.EqualValues(t, -30, record.Quarks) - - require.NoError(t, s.Backfill(ctx, "token_account_1", fn(500, 250, true))) - record, err = s.Get(ctx, "token_account_1") - require.NoError(t, err) - assert.True(t, record.IsBackfilled) - assert.True(t, record.IsLocked) - assertBalance(t, s, "token_account_1", 500, 250, true) - - // Predicates are enforced from now on - assert.Equal(t, balance.ErrInsufficientBalance, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 501})) - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_1", Kind: balance.DeltaDebit, Quarks: 500})) - assertBalance(t, s, "token_account_1", 0, 250, true) - - called := false - assert.Equal(t, balance.ErrAlreadyBackfilled, s.Backfill(ctx, "token_account_1", func(ctx context.Context) (*balance.BackfillResult, error) { - called = true - return &balance.BackfillResult{}, nil - })) - assert.False(t, called) - assertBalance(t, s, "token_account_1", 0, 250, true) - - // A backfill can establish the account as closed, e.g. a claimed gift card - require.NoError(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_2", - OwnerAccount: "owner", - MintAccount: "mint", - IsOpen: true, - IsLocked: true, - })) - require.NoError(t, s.Backfill(ctx, "token_account_2", fn(0, 0, false))) - assertBalance(t, s, "token_account_2", 0, 0, false) - assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 1})) - }) -} - func testMarkAsUnlocked(t *testing.T, s balance.Store) { t.Run("testMarkAsUnlocked", func(t *testing.T) { ctx := context.Background() @@ -611,7 +482,6 @@ func testMarkAsUnlocked(t *testing.T, s balance.Store) { Quarks: 100, IsOpen: true, IsLocked: true, - IsBackfilled: true, })) require.NoError(t, s.MarkAsUnlocked(ctx, "token_account_1")) @@ -649,23 +519,10 @@ func testMarkAsUnlocked(t *testing.T, s balance.Store) { MintAccount: "mint", IsOpen: true, IsLocked: true, - IsBackfilled: true, })) require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_3", Kind: balance.DeltaClose})) require.NoError(t, s.MarkAsUnlocked(ctx, "token_account_3")) assert.Equal(t, balance.ErrAccountClosed, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_3", Kind: balance.DeltaCredit, Quarks: 1})) - - // A record that is not backfilled accumulates without predicates, - // unlocked included; the backfill computes the truth for it - require.NoError(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_2", - OwnerAccount: "owner", - MintAccount: "mint", - IsOpen: true, - IsLocked: false, - })) - require.NoError(t, s.ApplyDeltas(ctx, &balance.Delta{TokenAccount: "token_account_2", Kind: balance.DeltaCredit, Quarks: 5})) - assertBalance(t, s, "token_account_2", 5, 0, true) }) } @@ -739,7 +596,6 @@ func assertEquivalentRecords(t *testing.T, obj1, obj2 *balance.Record) { assert.Equal(t, obj1.IsLocked, obj2.IsLocked) assert.Equal(t, obj1.UsdCostBasis, obj2.UsdCostBasis) assert.Equal(t, obj1.IsOpen, obj2.IsOpen) - assert.Equal(t, obj1.IsBackfilled, obj2.IsBackfilled) } func assertEquivalentExternalCheckpoingRecords(t *testing.T, obj1, obj2 *balance.ExternalCheckpointRecord) { diff --git a/ocp/data/internal.go b/ocp/data/internal.go index af4a9e0..841ea1c 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -130,7 +130,6 @@ type DatabaseData interface { CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error MarkBalanceAsUnlocked(ctx context.Context, tokenAccount string) error - BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error SaveExternalBalanceCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error GetExternalBalanceCheckpoint(ctx context.Context, account string) (*balance.ExternalCheckpointRecord, error) @@ -493,9 +492,6 @@ func (dp *DatabaseProvider) CountLockedBalancesByMint(ctx context.Context, mint func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error { return dp.balance.ApplyDeltas(ctx, deltas...) } -func (dp *DatabaseProvider) BackfillBalance(ctx context.Context, tokenAccount string, fn balance.BackfillFunc) error { - return dp.balance.Backfill(ctx, tokenAccount, fn) -} func (dp *DatabaseProvider) SaveExternalBalanceCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error { return dp.balance.SaveExternalCheckpoint(ctx, record) } diff --git a/ocp/worker/account/testutil.go b/ocp/worker/account/testutil.go index 18827dc..e9031ae 100644 --- a/ocp/worker/account/testutil.go +++ b/ocp/worker/account/testutil.go @@ -13,10 +13,12 @@ import ( commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" currency_lib "github.com/code-payments/ocp-server/currency" + "github.com/code-payments/ocp-server/ocp/balance" "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/account" "github.com/code-payments/ocp-server/ocp/data/action" + balance_data "github.com/code-payments/ocp-server/ocp/data/balance" "github.com/code-payments/ocp-server/ocp/data/fulfillment" "github.com/code-payments/ocp-server/ocp/data/intent" "github.com/code-payments/ocp-server/ocp/data/timelock" @@ -106,6 +108,15 @@ func (e *testEnv) generateRandomGiftCard(t *testing.T, creationTs time.Time) *te } require.NoError(t, e.data.SaveIntent(e.ctx, intentRecord)) + // The issuance funded the gift card, which the ledger recorded + require.NoError(t, balance.CreateRecordInTx(e.ctx, e.data, accountInfoRecord)) + require.NoError(t, balance.ApplyDeltasInTx(e.ctx, e.data, &balance_data.Delta{ + TokenAccount: accountInfoRecord.TokenAccount, + Kind: balance_data.DeltaCredit, + Quarks: intentRecord.SendPublicPaymentMetadata.Quantity, + UsdCostBasis: balance_data.UsdCostBasisFromFloat(intentRecord.SendPublicPaymentMetadata.UsdMarketValue), + })) + autoReturnActionRecord := &action.Record{ Intent: intentRecord.IntentId, IntentType: intentRecord.IntentType, From 0f9336112193b33ebd8b85df73b1596d6155703d Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 14:01:12 -0400 Subject: [PATCH 23/26] Quarks can now be uint64 --- ocp/balance/calculator.go | 10 +--------- ocp/data/balance/memory/store.go | 12 ++++++------ ocp/data/balance/postgres/model.go | 8 ++++---- ocp/data/balance/postgres/store.go | 4 ++-- ocp/data/balance/record.go | 8 +------- ocp/data/balance/store.go | 4 ++-- ocp/data/balance/tests/tests.go | 11 ++--------- ocp/data/internal.go | 8 ++++---- ocp/worker/currency/holder/runtime.go | 2 +- 9 files changed, 23 insertions(+), 44 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 6a7e4ac..4ca7580 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -27,10 +27,6 @@ const ( ) var ( - // ErrNegativeBalance indicates that a ledger record holds a negative - // value, which the store's delta predicates should make impossible. - ErrNegativeBalance = errors.New("balance calculation resulted in negative value") - // ErrNotManagedByCode indicates that an account is not owned by Code. // It's up to callers to determine how to handle this situation within // the context of a balance. @@ -366,11 +362,7 @@ func quarksFromRecord(record *balance.Record) (uint64, error) { if err := checkRecord(record); err != nil { return 0, err } - - if record.Quarks < 0 { - return 0, ErrNegativeBalance - } - return uint64(record.Quarks), nil + return record.Quarks, nil } func (s Source) String() string { diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index 0035070..c7c5794 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -101,7 +101,7 @@ func (s *store) GetAllByOwnerAndMint(_ context.Context, owner, mint string) ([]* } // GetAllLockedByMint implements balance.Store.GetAllLockedByMint -func (s *store) GetAllLockedByMint(_ context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { +func (s *store) GetAllLockedByMint(_ context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { s.mu.Lock() defer s.mu.Unlock() @@ -136,7 +136,7 @@ func (s *store) GetAllLockedByMint(_ context.Context, mint string, minQuarks int } // CountLockedByMint implements balance.Store.CountLockedByMint -func (s *store) CountLockedByMint(_ context.Context, mint string, minQuarks int64) (uint64, error) { +func (s *store) CountLockedByMint(_ context.Context, mint string, minQuarks uint64) (uint64, error) { s.mu.Lock() defer s.mu.Unlock() @@ -223,22 +223,22 @@ func applyDelta(item *balance.Record, delta *balance.Delta) error { if !item.IsOpen { return balance.ErrAccountClosed } - item.Quarks += int64(delta.Quarks) + item.Quarks += delta.Quarks item.UsdCostBasis += delta.UsdCostBasis case balance.DeltaDebit: if !item.IsOpen { return balance.ErrAccountClosed } - if item.Quarks < int64(delta.Quarks) { + if item.Quarks < delta.Quarks { return balance.ErrInsufficientBalance } - item.Quarks -= int64(delta.Quarks) + item.Quarks -= delta.Quarks item.UsdCostBasis -= delta.UsdCostBasis case balance.DeltaDrain: if !item.IsOpen { return balance.ErrAccountClosed } - if item.Quarks != int64(delta.Quarks) { + if item.Quarks != delta.Quarks { return balance.ErrBalanceChanged } item.Quarks = 0 diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index 6890ffa..eea77b8 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -30,8 +30,8 @@ type model struct { OwnerAccount string `db:"owner_account"` MintAccount string `db:"mint_account"` - Quarks int64 `db:"quarks"` - UsdCostBasis int64 `db:"usd_cost_basis"` + Quarks uint64 `db:"quarks"` + UsdCostBasis int64 `db:"usd_cost_basis"` IsOpen bool `db:"is_open"` IsLocked bool `db:"is_locked"` @@ -160,7 +160,7 @@ func dbGetAllByOwner(ctx context.Context, db *sqlx.DB, owner string, mint *strin return res, nil } -func dbGetAllLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64, cursor q.Cursor, limit uint64, direction q.Ordering) ([]*model, error) { +func dbGetAllLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks uint64, cursor q.Cursor, limit uint64, direction q.Ordering) ([]*model, error) { res := []*model{} query := `SELECT ` + allColumns + ` FROM ` + tableName + ` @@ -179,7 +179,7 @@ func dbGetAllLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuar return res, nil } -func dbCountLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks int64) (uint64, error) { +func dbCountLockedByMint(ctx context.Context, db *sqlx.DB, mint string, minQuarks uint64) (uint64, error) { var res uint64 query := `SELECT COUNT(*) FROM ` + tableName + ` WHERE mint_account = $1 AND quarks >= $2 AND is_locked` diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index d2c6178..00bbf3d 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -78,7 +78,7 @@ func (s *store) GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([ } // GetAllLockedByMint implements balance.Store.GetAllLockedByMint -func (s *store) GetAllLockedByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { +func (s *store) GetAllLockedByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { models, err := dbGetAllLockedByMint(ctx, s.db, mint, minQuarks, cursor, limit, direction) if err != nil { return nil, err @@ -87,7 +87,7 @@ func (s *store) GetAllLockedByMint(ctx context.Context, mint string, minQuarks i } // CountLockedByMint implements balance.Store.CountLockedByMint -func (s *store) CountLockedByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) { +func (s *store) CountLockedByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error) { return dbCountLockedByMint(ctx, s.db, mint, minQuarks) } diff --git a/ocp/data/balance/record.go b/ocp/data/balance/record.go index fe3f4a4..a471344 100644 --- a/ocp/data/balance/record.go +++ b/ocp/data/balance/record.go @@ -33,9 +33,7 @@ type Record struct { OwnerAccount string MintAccount string - // Quarks is signed to match the column's type. Delta predicates keep it - // non-negative. - Quarks int64 + Quarks uint64 // UsdCostBasis is the account's USD cost basis, in UsdQuarksPerUnit. // A cost basis may legitimately be negative. @@ -68,10 +66,6 @@ func (r *Record) Validate() error { return errors.New("mint account is required") } - if r.Quarks < 0 { - return errors.New("quarks cannot be negative") - } - return nil } diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index b904fb4..5b6a8db 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -62,12 +62,12 @@ type Store interface { // making their balances stale. // // ErrRecordNotFound is returned if no records exist. - GetAllLockedByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) + GetAllLockedByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) // CountLockedByMint counts locked records for a mint with at least // minQuarks. Unlocked records are excluded, since their balances are // stale. - CountLockedByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) + CountLockedByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error) // ApplyDeltas atomically applies a set of deltas. Either every delta is // applied or none are. Deltas are applied in SortDeltas order. diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index dbcf6ed..0ee5985 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -102,13 +102,6 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { _, err = s.GetAllByOwnerAndMint(ctx, "owner_2", "mint_2") assert.Equal(t, balance.ErrRecordNotFound, err) - - assert.Error(t, s.Create(ctx, &balance.Record{ - TokenAccount: "token_account_5", - OwnerAccount: "owner_1", - MintAccount: "mint_1", - Quarks: -1, - })) }) } @@ -124,7 +117,7 @@ func testGetAllLockedByMint(t *testing.T, s balance.Store) { TokenAccount: "token_account_" + string(rune('a'+i)), OwnerAccount: "owner", MintAccount: "mint_1", - Quarks: int64(i * 10), + Quarks: uint64(i * 10), IsOpen: true, IsLocked: true, })) @@ -580,7 +573,7 @@ func testExternalCheckpointHappyPath(t *testing.T, s balance.Store) { }) } -func assertBalance(t *testing.T, s balance.Store, tokenAccount string, quarks, usdCostBasis int64, isOpen bool) { +func assertBalance(t *testing.T, s balance.Store, tokenAccount string, quarks uint64, usdCostBasis int64, isOpen bool) { record, err := s.Get(context.Background(), tokenAccount) require.NoError(t, err) assert.EqualValues(t, quarks, record.Quarks, "quarks") diff --git a/ocp/data/internal.go b/ocp/data/internal.go index 841ea1c..64bc59e 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -126,8 +126,8 @@ type DatabaseData interface { GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) - GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) - CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) + GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) + CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error MarkBalanceAsUnlocked(ctx context.Context, tokenAccount string) error SaveExternalBalanceCheckpoint(ctx context.Context, record *balance.ExternalCheckpointRecord) error @@ -480,13 +480,13 @@ func (dp *DatabaseProvider) GetAllBalancesByOwner(ctx context.Context, owner str func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { return dp.balance.GetAllByOwnerAndMint(ctx, owner, mint) } -func (dp *DatabaseProvider) GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { +func (dp *DatabaseProvider) GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) { return dp.balance.GetAllLockedByMint(ctx, mint, minQuarks, cursor, limit, direction) } func (dp *DatabaseProvider) MarkBalanceAsUnlocked(ctx context.Context, tokenAccount string) error { return dp.balance.MarkAsUnlocked(ctx, tokenAccount) } -func (dp *DatabaseProvider) CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks int64) (uint64, error) { +func (dp *DatabaseProvider) CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error) { return dp.balance.CountLockedByMint(ctx, mint, minQuarks) } func (dp *DatabaseProvider) ApplyBalanceDeltas(ctx context.Context, deltas ...*balance.Delta) error { diff --git a/ocp/worker/currency/holder/runtime.go b/ocp/worker/currency/holder/runtime.go index 3ac018a..c9e3f66 100644 --- a/ocp/worker/currency/holder/runtime.go +++ b/ocp/worker/currency/holder/runtime.go @@ -110,5 +110,5 @@ func (p *holderRuntime) countHoldersForMint(ctx context.Context, mint string, cu return 0, nil } - return p.data.CountLockedBalancesByMint(ctx, mint, int64(minHoldings)) + return p.data.CountLockedBalancesByMint(ctx, mint, minHoldings) } From a87b69a62e9110a84dc6569eecdf3da1f49ad1ca Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 14:04:01 -0400 Subject: [PATCH 24/26] Remove unused USD cost basis helpers --- ocp/balance/calculator.go | 77 ++-------------------------------- ocp/balance/calculator_test.go | 45 ++------------------ 2 files changed, 8 insertions(+), 114 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 4ca7580..4d30cb8 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -287,80 +287,11 @@ func BatchCalculateWithUsdCostBasisFromCache(ctx context.Context, data ocp_data. return res, nil } -// CalculateUsdCostBasisFromCache calculates a token account's USD cost basis, -// in balance.UsdQuarksPerUnit, from its ledger record. -// -// Note: Unlike the quark balance calculators, no timelock check is performed, -// since the ledger record carries the lock state itself. Once a vault unlocks -// its record holds the last managed state rather than a live cost basis, so -// reading one returns ErrNotManagedByCode. An account the ledger doesn't -// track returns balance.ErrRecordNotFound. -func CalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (int64, error) { - tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "CalculateUsdCostBasisFromCache") - tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) - defer tracer.End() - - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err != nil { - tracer.OnError(err) - return 0, err - } - - if err := checkRecord(balanceRecord); err != nil { - tracer.OnError(err) - return 0, err - } - return balanceRecord.UsdCostBasis, nil -} - -// BatchCalculateUsdCostBasisFromCache is like CalculateUsdCostBasisFromCache, -// but for a set of token accounts. -func BatchCalculateUsdCostBasisFromCache(ctx context.Context, data ocp_data.Provider, tokenAccounts ...*common.Account) (map[string]int64, error) { - tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateUsdCostBasisFromCache") - defer tracer.End() - - tokenAccountStrings := make([]string, len(tokenAccounts)) - for i, tokenAccount := range tokenAccounts { - tokenAccountStrings[i] = tokenAccount.PublicKey().ToBase58() - } - - balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccountStrings...) - if err != nil { - tracer.OnError(err) - return nil, err - } - - res := make(map[string]int64, len(tokenAccounts)) - for _, tokenAccount := range tokenAccountStrings { - balanceRecord, ok := balanceRecords[tokenAccount] - if !ok { - tracer.OnError(balance.ErrRecordNotFound) - return nil, balance.ErrRecordNotFound - } - - if err := checkRecord(balanceRecord); err != nil { - tracer.OnError(err) - return nil, err - } - res[tokenAccount] = balanceRecord.UsdCostBasis - } - return res, nil -} - -// checkRecord verifies a ledger record is an authoritative view of an account -// that Code still manages. Quark balance callers reject unlocked vaults on the -// timelock record before reaching here, so for them this only guards against a -// record that disagrees with it. -func checkRecord(record *balance.Record) error { - if !record.IsLocked { - return ErrNotManagedByCode - } - return nil -} - func quarksFromRecord(record *balance.Record) (uint64, error) { - if err := checkRecord(record); err != nil { - return 0, err + // Callers reject unlocked vaults on the timelock record before reaching + // here, so this only guards against a record that disagrees with it. + if !record.IsLocked { + return 0, ErrNotManagedByCode } return record.Quarks, nil } diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 7dd9141..25db7ae 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -61,12 +61,6 @@ func TestDefaultCalculationMethods_MissingBalanceRecord(t *testing.T) { _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) assert.Equal(t, balance.ErrRecordNotFound, err) - - _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, balance.ErrRecordNotFound, err) - - _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, balance.ErrRecordNotFound, err) } func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { @@ -95,10 +89,10 @@ func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) // A record for an unlocked vault holds the last managed state, not a live - // balance, so it is refused even though the timelock record - // still passes the managed check. That pairing is inconsistent by - // construction: the timelock check normally rejects first, so the fixture - // exists to exercise the record's own guard. + // balance, so it is refused even though the timelock record still passes + // the managed check. That pairing is inconsistent by construction: the + // timelock check normally rejects first, so the fixture exists to exercise + // the record's own guard. testAccount := newBalanceTestAccount(t, env) saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true}) @@ -107,12 +101,6 @@ func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) { _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - - _, err = CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrNotManagedByCode, err) - - _, err = BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrNotManagedByCode, err) } func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) { @@ -140,31 +128,6 @@ func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) { assert.EqualValues(t, -123456, cached.UsdCostBasis) } -func TestUsdCostBasisCalculationMethods(t *testing.T) { - env := setupBalanceTestEnv(t) - - first := newBalanceTestAccount(t, env) - second := newBalanceTestAccount(t, env) - - saveBalanceTestRecord(t, env, first, &balance.Record{UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 1, UsdCostBasis: 1_500_000, IsOpen: true, IsLocked: true}) - - expected := map[string]int64{ - first.tokenAccount.PublicKey().ToBase58(): -123456, - second.tokenAccount.PublicKey().ToBase58(): 1_500_000, - } - - for _, testAccount := range []*balanceTestAccount{first, second} { - actual, err := CalculateUsdCostBasisFromCache(env.ctx, env.data, testAccount.tokenAccount) - require.NoError(t, err) - assert.EqualValues(t, expected[testAccount.tokenAccount.PublicKey().ToBase58()], actual) - } - - usdCostBasisByAccount, err := BatchCalculateUsdCostBasisFromCache(env.ctx, env.data, first.tokenAccount, second.tokenAccount) - require.NoError(t, err) - assert.Equal(t, expected, usdCostBasisByAccount) -} - func TestDefaultCalculation_ExternalAccount(t *testing.T) { env := setupBalanceTestEnv(t) externalAccount := testutil.NewRandomAccount(t) From 3fa809a64438024ac338ed41ac9a7327c804cdb4 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 14:17:11 -0400 Subject: [PATCH 25/26] Further simplify balance calculators --- ocp/balance/calculator.go | 194 +++++------------------- ocp/balance/calculator_test.go | 124 +++++---------- ocp/rpc/account/server.go | 32 ++-- ocp/rpc/account/server_test.go | 23 ++- ocp/rpc/balance/server.go | 23 +-- ocp/rpc/balance/server_test.go | 3 + ocp/rpc/transaction/local_simulation.go | 22 +-- 7 files changed, 143 insertions(+), 278 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 4d30cb8..e1213f5 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -10,7 +10,6 @@ import ( "github.com/code-payments/ocp-server/ocp/common" ocp_data "github.com/code-payments/ocp-server/ocp/data" "github.com/code-payments/ocp-server/ocp/data/balance" - "github.com/code-payments/ocp-server/ocp/data/timelock" "github.com/code-payments/ocp-server/solana" ) @@ -36,6 +35,11 @@ var ( // CalculateFromCache is the default and recommended strategy for reliably estimating // a token account's balance using cached values. // +// The ledger record is the whole answer: it exists for exactly the accounts +// Code manages, and carries the lock state that says whether it still does. +// An account with no record, or one whose vault has unlocked, is +// ErrNotManagedByCode. +// // Note: Use this method when calculating balances for accounts that are managed by // Code (ie. Timelock account) and operate within the L2 system. func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccount *common.Account) (uint64, error) { @@ -43,8 +47,8 @@ func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccoun tracer.AddAttribute("account", tokenAccount.PublicKey().ToBase58()) defer tracer.End() - timelockRecord, err := data.GetTimelockByVault(ctx, tokenAccount.PublicKey().ToBase58()) - if err == timelock.ErrTimelockNotFound { + balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) + if err == balance.ErrRecordNotFound { tracer.OnError(ErrNotManagedByCode) return 0, ErrNotManagedByCode } else if err != nil { @@ -52,26 +56,13 @@ func CalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccoun return 0, err } - // The balance ledger is only maintained for accounts managed by Code. The - // account must be managed in order to return accurate values. - isManagedByCode := common.IsManagedByCode(ctx, timelockRecord) - if !isManagedByCode { + // Once a vault unlocks, funds can move on chain without an intent, so the + // record stops being maintained and its balance must not be trusted. + if !balanceRecord.IsLocked { tracer.OnError(ErrNotManagedByCode) return 0, ErrNotManagedByCode } - - balanceRecord, err := data.GetBalance(ctx, tokenAccount.PublicKey().ToBase58()) - if err != nil { - tracer.OnError(err) - return 0, err - } - - quarks, err := quarksFromRecord(balanceRecord) - if err != nil { - tracer.OnError(err) - return 0, err - } - return quarks, nil + return balanceRecord.Quarks, nil } // CalculateFromBlockchain is the default and recommended strategy for reliably @@ -134,42 +125,27 @@ func CalculateFromBlockchain(ctx context.Context, data ocp_data.Provider, tokenA return quarks, BlockchainSource, nil } -// BatchCalculateFromCacheWithAccountRecords is the default and recommended batch strategy -// or reliably estimating a set of token accounts' balance when common.AccountRecords are -// available. -// -// Note: Use this method when calculating balances for accounts that are managed by -// Code (ie. Timelock account) and operate within the L2 system. -func BatchCalculateFromCacheWithAccountRecords(ctx context.Context, data ocp_data.Provider, accountRecordsBatch ...*common.AccountRecords) (map[string]uint64, error) { - tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheWithAccountRecords") - defer tracer.End() - - timelockRecords := make([]*timelock.Record, 0) - for _, accountRecords := range accountRecordsBatch { - if !accountRecords.IsTimelock() { - tracer.OnError(ErrNotManagedByCode) - return nil, ErrNotManagedByCode - } - - timelockRecords = append(timelockRecords, accountRecords.Timelock) - } - - balanceByTokenAccount, err := defaultBatchCalculationFromCache(ctx, data, timelockRecords) - if err != nil { - tracer.OnError(err) - return nil, err - } - return balanceByTokenAccount, nil +// Balance is a token account's quark balance and USD cost basis, in +// balance.UsdQuarksPerUnit. +type Balance struct { + Quarks uint64 + UsdCostBasis int64 } -// BatchCalculateFromCacheWithTokenAccounts is the default and recommended batch strategy -// or reliably estimating a set of token accounts' balance when common.Account are -// available. +// BatchCalculateFromCache is the default and recommended batch strategy for +// reliably estimating a set of token accounts' balances using cached values. +// Both values for an account come from the same ledger record read, so they +// are guaranteed consistent with each other. +// +// Accounts the ledger doesn't manage are omitted from the result: those with +// no record, and those whose vault has unlocked. A caller that requires every +// account it asked for to be managed compares the result's length against its +// input. // // Note: Use this method when calculating balances for accounts that are managed by // Code (ie. Timelock account) and operate within the L2 system. -func BatchCalculateFromCacheWithTokenAccounts(ctx context.Context, data ocp_data.Provider, tokenAccounts ...*common.Account) (map[string]uint64, error) { - tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheWithTokenAccounts") +func BatchCalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenAccounts ...*common.Account) (map[string]*Balance, error) { + tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCache") defer tracer.End() tokenAccountStrings := make([]string, len(tokenAccounts)) @@ -177,125 +153,29 @@ func BatchCalculateFromCacheWithTokenAccounts(ctx context.Context, data ocp_data tokenAccountStrings[i] = tokenAccount.PublicKey().ToBase58() } - timelockRecordsByVault, err := data.GetTimelockByVaultBatch(ctx, tokenAccountStrings...) - if err == timelock.ErrTimelockNotFound { - tracer.OnError(ErrNotManagedByCode) - return nil, ErrNotManagedByCode - } else if err != nil { - tracer.OnError(err) - return nil, err - } - - timelockRecords := make([]*timelock.Record, 0, len(timelockRecordsByVault)) - for _, timelockRecord := range timelockRecordsByVault { - timelockRecords = append(timelockRecords, timelockRecord) - } - - balanceByTokenAccount, err := defaultBatchCalculationFromCache(ctx, data, timelockRecords) + balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccountStrings...) if err != nil { tracer.OnError(err) return nil, err } - return balanceByTokenAccount, nil -} -func defaultBatchCalculationFromCache(ctx context.Context, data ocp_data.Provider, timelockRecords []*timelock.Record) (map[string]uint64, error) { - tokenAccounts := make([]string, 0, len(timelockRecords)) - for _, timelockRecord := range timelockRecords { - // The balance ledger is only maintained for accounts managed by Code. - // The account must be managed in order to return accurate values. - isManagedByCode := common.IsManagedByCode(ctx, timelockRecord) - if !isManagedByCode { - return nil, ErrNotManagedByCode + res := make(map[string]*Balance, len(balanceRecords)) + for tokenAccount, balanceRecord := range balanceRecords { + // Once a vault unlocks, funds can move on chain without an intent, so + // the record stops being maintained and its balance must not be + // trusted or aggregated. + if !balanceRecord.IsLocked { + continue } - tokenAccounts = append(tokenAccounts, timelockRecord.VaultAddress) - } - - balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccounts...) - if err != nil { - return nil, err - } - - res := make(map[string]uint64, len(tokenAccounts)) - for _, tokenAccount := range tokenAccounts { - balanceRecord, ok := balanceRecords[tokenAccount] - if !ok { - return nil, balance.ErrRecordNotFound - } - - quarks, err := quarksFromRecord(balanceRecord) - if err != nil { - return nil, err - } - res[tokenAccount] = quarks - } - 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. Both values for an account come from -// the same balance record read, so they are guaranteed consistent with each -// other. -// -// 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() - - tokenAccounts := make([]string, 0, len(accountRecordsBatch)) - for _, accountRecords := range accountRecordsBatch { - if !accountRecords.IsTimelock() || !accountRecords.IsManagedByCode(ctx) { - tracer.OnError(ErrNotManagedByCode) - return nil, ErrNotManagedByCode - } - tokenAccounts = append(tokenAccounts, accountRecords.General.TokenAccount) - } - - balanceRecords, err := data.GetBalanceBatch(ctx, tokenAccounts...) - if err != nil { - tracer.OnError(err) - return nil, err - } - - res := make(map[string]*BalanceWithUsdCostBasis, len(tokenAccounts)) - for _, tokenAccount := range tokenAccounts { - balanceRecord, ok := balanceRecords[tokenAccount] - if !ok { - tracer.OnError(balance.ErrRecordNotFound) - return nil, balance.ErrRecordNotFound - } - - quarks, err := quarksFromRecord(balanceRecord) - if err != nil { - tracer.OnError(err) - return nil, err - } - res[tokenAccount] = &BalanceWithUsdCostBasis{ - Quarks: quarks, + res[tokenAccount] = &Balance{ + Quarks: balanceRecord.Quarks, UsdCostBasis: balanceRecord.UsdCostBasis, } } return res, nil } -func quarksFromRecord(record *balance.Record) (uint64, error) { - // Callers reject unlocked vaults on the timelock record before reaching - // here, so this only guards against a record that disagrees with it. - if !record.IsLocked { - return 0, ErrNotManagedByCode - } - return record.Quarks, nil -} - func (s Source) String() string { switch s { case UnknownSource: diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 25db7ae..ab20f84 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -23,44 +23,49 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { first := newBalanceTestAccount(t, env) second := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 0, IsOpen: true, IsLocked: true}) - - expected := map[string]uint64{ - first.tokenAccount.PublicKey().ToBase58(): 42, - second.tokenAccount.PublicKey().ToBase58(): 0, - } - - for _, testAccount := range []*balanceTestAccount{first, second} { - actual, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) + // Quarks and cost basis come from the same record, so they can't disagree + saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 0, UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) + + for _, tc := range []struct { + testAccount *balanceTestAccount + quarks uint64 + usdCostBasis int64 + }{ + {first, 42, 4_200_000}, + {second, 0, -123456}, + } { + actual, err := CalculateFromCache(env.ctx, env.data, tc.testAccount.tokenAccount) require.NoError(t, err) - assert.EqualValues(t, expected[testAccount.tokenAccount.PublicKey().ToBase58()], actual) + assert.EqualValues(t, tc.quarks, actual) } - balanceByAccount, err := BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, first.tokenAccount, second.tokenAccount) - require.NoError(t, err) - assert.Equal(t, expected, balanceByAccount) - - balanceByAccount, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, first.accountRecords(t, env), second.accountRecords(t, env)) + balanceByAccount, err := BatchCalculateFromCache(env.ctx, env.data, first.tokenAccount, second.tokenAccount) require.NoError(t, err) - assert.Equal(t, expected, balanceByAccount) + assert.Equal(t, map[string]*Balance{ + first.tokenAccount.PublicKey().ToBase58(): {Quarks: 42, UsdCostBasis: 4_200_000}, + second.tokenAccount.PublicKey().ToBase58(): {Quarks: 0, UsdCostBasis: -123456}, + }, balanceByAccount) } func TestDefaultCalculationMethods_MissingBalanceRecord(t *testing.T) { env := setupBalanceTestEnv(t) - // The account is managed, but the ledger has no record for it, which is a - // broken invariant rather than a balance of zero - testAccount := newBalanceTestAccount(t, env) - - _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, balance.ErrRecordNotFound, err) + tracked := newBalanceTestAccount(t, env) + saveBalanceTestRecord(t, env, tracked, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) + untracked := newBalanceTestAccount(t, env) - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, balance.ErrRecordNotFound, err) + // The ledger record is the whole answer, so an account without one simply + // isn't managed by Code + _, err := CalculateFromCache(env.ctx, env.data, untracked.tokenAccount) + assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) - assert.Equal(t, balance.ErrRecordNotFound, err) + // The batch variant says so by omission, and the rest of the batch still + // resolves + balanceByAccount, err := BatchCalculateFromCache(env.ctx, env.data, tracked.tokenAccount, untracked.tokenAccount) + require.NoError(t, err) + require.Len(t, balanceByAccount, 1) + assert.EqualValues(t, 42, balanceByAccount[tracked.tokenAccount.PublicKey().ToBase58()].Quarks) } func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { @@ -69,63 +74,23 @@ func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { testAccount := newBalanceTestAccount(t, env) saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) + // The vault unlocks. In production the timelock and ledger records move in + // the same transaction, so they can't disagree. timelockRecord, err := env.data.GetTimelockByVault(env.ctx, testAccount.tokenAccount.PublicKey().ToBase58()) require.NoError(t, err) timelockRecord.VaultState = timelock_token_v1.StateWaitingForTimeout timelockRecord.Block += 1 require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord)) + require.NoError(t, env.data.MarkBalanceAsUnlocked(env.ctx, testAccount.tokenAccount.PublicKey().ToBase58())) + // A record for an unlocked vault holds the last managed state rather than + // a live balance, so it is refused outright and omitted from a batch _, err = CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) assert.Equal(t, ErrNotManagedByCode, err) - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrNotManagedByCode, err) - - _, err = BatchCalculateFromCacheWithAccountRecords(env.ctx, env.data, testAccount.accountRecords(t, env)) - assert.Equal(t, ErrNotManagedByCode, err) -} - -func TestDefaultCalculationMethods_UnlockedBalanceRecord(t *testing.T) { - env := setupBalanceTestEnv(t) - - // A record for an unlocked vault holds the last managed state, not a live - // balance, so it is refused even though the timelock record still passes - // the managed check. That pairing is inconsistent by construction: the - // timelock check normally rejects first, so the fixture exists to exercise - // the record's own guard. - testAccount := newBalanceTestAccount(t, env) - saveBalanceTestRecord(t, env, testAccount, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true}) - - _, err := CalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrNotManagedByCode, err) - - _, err = BatchCalculateFromCacheWithTokenAccounts(env.ctx, env.data, testAccount.tokenAccount) - assert.Equal(t, ErrNotManagedByCode, err) -} - -func TestDefaultCalculationMethods_BalanceWithUsdCostBasis(t *testing.T) { - env := setupBalanceTestEnv(t) - - first := newBalanceTestAccount(t, env) - second := newBalanceTestAccount(t, env) - - // Both values come from the same record, so they can't disagree - saveBalanceTestRecord(t, env, first, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) - saveBalanceTestRecord(t, env, second, &balance.Record{Quarks: 33, UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) - - res, err := BatchCalculateWithUsdCostBasisFromCache(env.ctx, env.data, first.accountRecords(t, env), second.accountRecords(t, env)) + balanceByAccount, err := BatchCalculateFromCache(env.ctx, env.data, testAccount.tokenAccount) require.NoError(t, err) - require.Len(t, res, 2) - - cached := res[first.tokenAccount.PublicKey().ToBase58()] - require.NotNil(t, cached) - assert.EqualValues(t, 42, cached.Quarks) - assert.EqualValues(t, 4_200_000, cached.UsdCostBasis) - - cached = res[second.tokenAccount.PublicKey().ToBase58()] - require.NotNil(t, cached) - assert.EqualValues(t, 33, cached.Quarks) - assert.EqualValues(t, -123456, cached.UsdCostBasis) + assert.Empty(t, balanceByAccount) } func TestDefaultCalculation_ExternalAccount(t *testing.T) { @@ -194,14 +159,3 @@ func saveBalanceTestRecord(t *testing.T, env balanceTestEnv, testAccount *balanc record.MintAccount = testAccount.vmConfig.Mint.PublicKey().ToBase58() require.NoError(t, env.data.CreateBalance(env.ctx, record)) } - -func (a *balanceTestAccount) accountRecords(t *testing.T, env balanceTestEnv) *common.AccountRecords { - generalRecord, err := env.data.GetAccountInfoByTokenAddress(env.ctx, a.tokenAccount.PublicKey().ToBase58()) - require.NoError(t, err) - timelockRecord, err := env.data.GetTimelockByVault(env.ctx, a.tokenAccount.PublicKey().ToBase58()) - require.NoError(t, err) - return &common.AccountRecords{ - General: generalRecord, - Timelock: timelockRecord, - } -} diff --git a/ocp/rpc/account/server.go b/ocp/rpc/account/server.go index b5ed848..e1b5922 100644 --- a/ocp/rpc/account/server.go +++ b/ocp/rpc/account/server.go @@ -315,21 +315,29 @@ func (s *server) GetTokenAccountInfos(ctx context.Context, req *accountpb.GetTok func (s *server) fetchBalances(ctx context.Context, allAccountRecords []*common.AccountRecords) (map[string]*balanceMetadata, error) { balanceMetadataByTokenAccount := make(map[string]*balanceMetadata) - var mangedByCodeRecords []*common.AccountRecords + // Every Timelock account starts out without a balance, since the caching + // strategy isn't possible for the ones that have left the L2 system. The + // ledger omits those, so they keep this value. + var tokenAccounts []*common.Account for _, accountRecords := range allAccountRecords { - if accountRecords.IsManagedByCode(ctx) { - mangedByCodeRecords = append(mangedByCodeRecords, accountRecords) - } else { - // Don't calculate a balance for now, since the caching strategy - // is not possible. - balanceMetadataByTokenAccount[accountRecords.General.TokenAccount] = &balanceMetadata{ - quarks: 0, - usdCostBasis: 0, - source: accountpb.TokenAccountInfo_BALANCE_SOURCE_UNKNOWN, - } + if !accountRecords.IsTimelock() { + continue + } + + balanceMetadataByTokenAccount[accountRecords.General.TokenAccount] = &balanceMetadata{ + quarks: 0, + usdCostBasis: 0, + source: accountpb.TokenAccountInfo_BALANCE_SOURCE_UNKNOWN, } + + tokenAccount, err := common.NewAccountFromPublicKeyString(accountRecords.General.TokenAccount) + if err != nil { + return nil, err + } + tokenAccounts = append(tokenAccounts, tokenAccount) } - balancesByTokenAccount, err := balance.BatchCalculateWithUsdCostBasisFromCache(ctx, s.data, mangedByCodeRecords...) + + balancesByTokenAccount, err := balance.BatchCalculateFromCache(ctx, s.data, tokenAccounts...) if err != nil { return nil, err } diff --git a/ocp/rpc/account/server_test.go b/ocp/rpc/account/server_test.go index fad9efb..e2a504c 100644 --- a/ocp/rpc/account/server_test.go +++ b/ocp/rpc/account/server_test.go @@ -517,9 +517,7 @@ func TestGetTokenAccountInfos_RemoteSendGiftCard_HappyPath(t *testing.T) { require.NoError(t, env.data.PutAllActions(env.ctx, claimActionRecord)) } - accountRecords.Timelock.VaultState = tc.timelockState - accountRecords.Timelock.Block += 1 - require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) + setTimelockState(t, env, accountRecords, tc.timelockState) for _, requestingOwnerAccount := range []*common.Account{ nil, @@ -629,11 +627,9 @@ func TestGetTokenAccountInfos_BlockchainState(t *testing.T) { } accountRecords := getDefaultTestAccountRecords(t, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY) - accountRecords.Timelock.VaultState = tc.timelockState - accountRecords.Timelock.Block += 1 require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountRecords.General)) require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountRecords.General)) - require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) + setTimelockState(t, env, accountRecords, tc.timelockState) resp, err := env.client.GetTokenAccountInfos(env.ctx, req) require.NoError(t, err) @@ -703,6 +699,9 @@ func TestGetTokenAccountInfos_ManagementState(t *testing.T) { require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountRecords.General)) require.NoError(t, balance_util.CreateRecordInTx(env.ctx, env.data, accountRecords.General)) require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) + if !accountRecords.Timelock.IsLocked() { + require.NoError(t, env.data.MarkBalanceAsUnlocked(env.ctx, accountRecords.General.TokenAccount)) + } resp, err := env.client.GetTokenAccountInfos(env.ctx, req) require.NoError(t, err) @@ -848,6 +847,18 @@ func getDefaultTestAccountRecords(t *testing.T, ownerAccount, authorityAccount * } } +// setTimelockState transitions a vault the way the geyser worker does, moving +// the ledger record's lock state with it so the two can't disagree. +func setTimelockState(t *testing.T, env testEnv, accountRecords *common.AccountRecords, state timelock_token_v1.TimelockState) { + accountRecords.Timelock.VaultState = state + accountRecords.Timelock.Block += 1 + require.NoError(t, env.data.SaveTimelock(env.ctx, accountRecords.Timelock)) + + if !accountRecords.Timelock.IsLocked() { + require.NoError(t, env.data.MarkBalanceAsUnlocked(env.ctx, accountRecords.General.TokenAccount)) + } +} + func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, quarks uint64) { require.NoError(t, balance_util.ApplyDeltasInTx(env.ctx, env.data, &balance.Delta{ TokenAccount: accountRecords.General.TokenAccount, diff --git a/ocp/rpc/balance/server.go b/ocp/rpc/balance/server.go index d83413c..fd2dc30 100644 --- a/ocp/rpc/balance/server.go +++ b/ocp/rpc/balance/server.go @@ -79,35 +79,40 @@ func (s *server) calculateCoreMintValue(ctx context.Context, owner *common.Accou return 0, err } - // Accounts that have left the L2 system don't have a cached balance that can - // be trusted, so they're excluded from the calculation. mintByTokenAccount := make(map[string]string) - var managedByCodeRecords []*common.AccountRecords + var tokenAccounts []*common.Account for mint, recordsByType := range recordsByMintAndType { for _, recordsList := range recordsByType { for _, records := range recordsList { - if !records.IsManagedByCode(ctx) { + if !records.IsTimelock() { continue } + tokenAccount, err := common.NewAccountFromPublicKeyString(records.General.TokenAccount) + if err != nil { + return 0, err + } + mintByTokenAccount[records.General.TokenAccount] = mint - managedByCodeRecords = append(managedByCodeRecords, records) + tokenAccounts = append(tokenAccounts, tokenAccount) } } } - if len(managedByCodeRecords) == 0 { + if len(tokenAccounts) == 0 { return 0, nil } - balanceByTokenAccount, err := balance.BatchCalculateFromCacheWithAccountRecords(ctx, s.data, managedByCodeRecords...) + // Accounts that have left the L2 system don't have a cached balance that can + // be trusted, so the ledger omits them from the result. + balanceByTokenAccount, err := balance.BatchCalculateFromCache(ctx, s.data, tokenAccounts...) if err != nil { return 0, err } quarksByMint := make(map[string]uint64) - for tokenAccount, quarks := range balanceByTokenAccount { - quarksByMint[mintByTokenAccount[tokenAccount]] += quarks + for tokenAccount, cached := range balanceByTokenAccount { + quarksByMint[mintByTokenAccount[tokenAccount]] += cached.Quarks } var coreMintValue uint64 diff --git a/ocp/rpc/balance/server_test.go b/ocp/rpc/balance/server_test.go index 823f66c..26abe2a 100644 --- a/ocp/rpc/balance/server_test.go +++ b/ocp/rpc/balance/server_test.go @@ -131,9 +131,12 @@ func TestGetBalance_UnmanagedAccountsExcluded(t *testing.T) { // The pool account has left the L2 system, so there isn't a cached balance that // can be trusted for it. + // The geyser worker moves both records in the same transaction, so the + // ledger record's lock state can't disagree with the timelock record's primaryCoreMintAccountRecords.Timelock.VaultState = timelock_token_v1.StateUnlocked primaryCoreMintAccountRecords.Timelock.Block += 1 require.NoError(t, env.data.SaveTimelock(env.ctx, primaryCoreMintAccountRecords.Timelock)) + require.NoError(t, env.data.MarkBalanceAsUnlocked(env.ctx, primaryCoreMintAccountRecords.General.TokenAccount)) resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ Owner: ownerAccount.ToProto(), diff --git a/ocp/rpc/transaction/local_simulation.go b/ocp/rpc/transaction/local_simulation.go index 4e47635..b1c5fb1 100644 --- a/ocp/rpc/transaction/local_simulation.go +++ b/ocp/rpc/transaction/local_simulation.go @@ -340,14 +340,18 @@ func LocalSimulation(ctx context.Context, data ocp_data.Provider, actions []*tra tokenAccountsToFetchBalance = append(tokenAccountsToFetchBalance, sim.TokenAccount) } } - prefetchedBalances := make(map[string]uint64) + prefetchedBalances := make(map[string]*balance.Balance) if len(tokenAccountsToFetchBalance) > 0 { - prefetchedBalances, err = balance.BatchCalculateFromCacheWithTokenAccounts(ctx, data, tokenAccountsToFetchBalance...) - if err == balance.ErrNotManagedByCode { - return nil, ErrSourceNotManagedByCode - } else if err != nil { + prefetchedBalances, err = balance.BatchCalculateFromCache(ctx, data, tokenAccountsToFetchBalance...) + if err != nil { return nil, err } + + // The ledger omits accounts it doesn't manage, and every account we + // simulate against must be managed + if len(prefetchedBalances) != len(tokenAccountsToFetchBalance) { + return nil, ErrSourceNotManagedByCode + } } // Do more complex simulation validation on each involved account using all combined actions @@ -356,16 +360,16 @@ func LocalSimulation(ctx context.Context, data ocp_data.Provider, actions []*tra continue } - var ok bool - var balance uint64 + var quarks uint64 if sim.RequiresBalanceFetch() { - balance, ok = prefetchedBalances[sim.TokenAccount.PublicKey().ToBase58()] + cached, ok := prefetchedBalances[sim.TokenAccount.PublicKey().ToBase58()] if !ok { return nil, errors.New("prefetched balance is unavailable") } + quarks = cached.Quarks } - err := sim.EnforceBalances(ctx, data, balance) + err := sim.EnforceBalances(ctx, data, quarks) if err != nil { return nil, err } From 90396bda20d47928cccefce5f10bb4198aedfd19 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 31 Aug 2026 14:24:01 -0400 Subject: [PATCH 26/26] Add BatchCalculateFromCacheByOwner utility and use it in the Balance RPC service --- ocp/balance/calculator.go | 59 ++++++++++++++++++++++++++++------ ocp/balance/calculator_test.go | 42 +++++++++++++++++++++--- ocp/rpc/balance/server.go | 40 ++++------------------- 3 files changed, 94 insertions(+), 47 deletions(-) diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index e1213f5..3f515b4 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -126,8 +126,9 @@ func CalculateFromBlockchain(ctx context.Context, data ocp_data.Provider, tokenA } // Balance is a token account's quark balance and USD cost basis, in -// balance.UsdQuarksPerUnit. +// balance.UsdQuarksPerUnit, alongside the mint it holds. type Balance struct { + MintAccount string Quarks uint64 UsdCostBasis int64 } @@ -161,21 +162,61 @@ func BatchCalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenA res := make(map[string]*Balance, len(balanceRecords)) for tokenAccount, balanceRecord := range balanceRecords { - // Once a vault unlocks, funds can move on chain without an intent, so - // the record stops being maintained and its balance must not be - // trusted or aggregated. - if !balanceRecord.IsLocked { - continue + if cached, ok := balanceFromRecord(balanceRecord); ok { + res[tokenAccount] = cached } + } + return res, nil +} + +// BatchCalculateFromCacheByOwner is the default and recommended strategy for +// reliably estimating the balance of every account an owner holds, keyed by +// token account. Each balance carries the mint its account holds, so a caller +// aggregating across mints doesn't need the account records to say which is +// which. +// +// Accounts the ledger doesn't manage are omitted, as are owners it holds +// nothing for, so an owner outside the L2 system is an empty result rather +// than an error. +// +// Note: Use this method when calculating balances for accounts that are managed by +// Code (ie. Timelock account) and operate within the L2 system. +func BatchCalculateFromCacheByOwner(ctx context.Context, data ocp_data.Provider, owner *common.Account) (map[string]*Balance, error) { + tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheByOwner") + tracer.AddAttribute("owner", owner.PublicKey().ToBase58()) + defer tracer.End() - res[tokenAccount] = &Balance{ - Quarks: balanceRecord.Quarks, - UsdCostBasis: balanceRecord.UsdCostBasis, + balanceRecords, err := data.GetAllBalancesByOwner(ctx, owner.PublicKey().ToBase58()) + if err != nil && err != balance.ErrRecordNotFound { + tracer.OnError(err) + return nil, err + } + + res := make(map[string]*Balance, len(balanceRecords)) + for _, balanceRecord := range balanceRecords { + if cached, ok := balanceFromRecord(balanceRecord); ok { + res[balanceRecord.TokenAccount] = cached } } return res, nil } +// balanceFromRecord reports a ledger record as a Balance, and whether the +// ledger still maintains it. Once a vault unlocks, funds can move on chain +// without an intent, so the record's balance must not be trusted or +// aggregated. +func balanceFromRecord(record *balance.Record) (*Balance, bool) { + if !record.IsLocked { + return nil, false + } + + return &Balance{ + MintAccount: record.MintAccount, + Quarks: record.Quarks, + UsdCostBasis: record.UsdCostBasis, + }, true +} + func (s Source) String() string { switch s { case UnknownSource: diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index ab20f84..54ad38d 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -43,8 +43,8 @@ func TestDefaultCalculationMethods_BalanceRecord(t *testing.T) { balanceByAccount, err := BatchCalculateFromCache(env.ctx, env.data, first.tokenAccount, second.tokenAccount) require.NoError(t, err) assert.Equal(t, map[string]*Balance{ - first.tokenAccount.PublicKey().ToBase58(): {Quarks: 42, UsdCostBasis: 4_200_000}, - second.tokenAccount.PublicKey().ToBase58(): {Quarks: 0, UsdCostBasis: -123456}, + first.tokenAccount.PublicKey().ToBase58(): {MintAccount: first.mint(), Quarks: 42, UsdCostBasis: 4_200_000}, + second.tokenAccount.PublicKey().ToBase58(): {MintAccount: second.mint(), Quarks: 0, UsdCostBasis: -123456}, }, balanceByAccount) } @@ -93,6 +93,33 @@ func TestDefaultCalculationMethods_NotManagedByCode(t *testing.T) { assert.Empty(t, balanceByAccount) } +func TestDefaultCalculationMethods_ByOwner(t *testing.T) { + env := setupBalanceTestEnv(t) + + owner := testutil.NewRandomAccount(t) + coreMint := newBalanceTestAccountForOwner(t, env, owner, testutil.NewRandomVmConfig(t, true)) + otherMint := newBalanceTestAccountForOwner(t, env, owner, testutil.NewRandomVmConfig(t, false)) + unlocked := newBalanceTestAccountForOwner(t, env, owner, testutil.NewRandomVmConfig(t, false)) + + saveBalanceTestRecord(t, env, coreMint, &balance.Record{Quarks: 42, UsdCostBasis: 4_200_000, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, otherMint, &balance.Record{Quarks: 33, UsdCostBasis: -123456, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, unlocked, &balance.Record{Quarks: 99, IsOpen: true}) + + // Every account the ledger manages for the owner is reported with the mint + // it holds. The unlocked one is omitted, since its balance is stale. + balanceByAccount, err := BatchCalculateFromCacheByOwner(env.ctx, env.data, owner) + require.NoError(t, err) + assert.Equal(t, map[string]*Balance{ + coreMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: coreMint.mint(), Quarks: 42, UsdCostBasis: 4_200_000}, + otherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: otherMint.mint(), Quarks: 33, UsdCostBasis: -123456}, + }, balanceByAccount) + + // An owner the ledger holds nothing for is empty rather than an error + balanceByAccount, err = BatchCalculateFromCacheByOwner(env.ctx, env.data, testutil.NewRandomAccount(t)) + require.NoError(t, err) + assert.Empty(t, balanceByAccount) +} + func TestDefaultCalculation_ExternalAccount(t *testing.T) { env := setupBalanceTestEnv(t) externalAccount := testutil.NewRandomAccount(t) @@ -123,9 +150,12 @@ func setupBalanceTestEnv(t *testing.T) (env balanceTestEnv) { // newBalanceTestAccount creates a locked timelock account, with an account // info record but no ledger record. func newBalanceTestAccount(t *testing.T, env balanceTestEnv) *balanceTestAccount { - vmConfig := testutil.NewRandomVmConfig(t, true) - owner := testutil.NewRandomAccount(t) + return newBalanceTestAccountForOwner(t, env, testutil.NewRandomAccount(t), testutil.NewRandomVmConfig(t, true)) +} +// newBalanceTestAccountForOwner is like newBalanceTestAccount, for an owner +// that holds accounts across several mints. +func newBalanceTestAccountForOwner(t *testing.T, env balanceTestEnv, owner *common.Account, vmConfig *common.VmConfig) *balanceTestAccount { timelockAccounts, err := owner.GetTimelockAccounts(vmConfig) require.NoError(t, err) timelockRecord := timelockAccounts.ToDBRecord() @@ -151,6 +181,10 @@ func newBalanceTestAccount(t *testing.T, env balanceTestEnv) *balanceTestAccount } } +func (a *balanceTestAccount) mint() string { + return a.vmConfig.Mint.PublicKey().ToBase58() +} + // saveBalanceTestRecord creates the account's ledger record, filling in the // identifying fields from the account. func saveBalanceTestRecord(t *testing.T, env balanceTestEnv, testAccount *balanceTestAccount, record *balance.Record) { diff --git a/ocp/rpc/balance/server.go b/ocp/rpc/balance/server.go index fd2dc30..b383252 100644 --- a/ocp/rpc/balance/server.go +++ b/ocp/rpc/balance/server.go @@ -74,45 +74,17 @@ func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceReques } func (s *server) calculateCoreMintValue(ctx context.Context, owner *common.Account) (uint64, error) { - recordsByMintAndType, err := common.GetLatestCodeTimelockAccountRecordsForOwner(ctx, s.data, owner) - if err != nil { - return 0, err - } - - mintByTokenAccount := make(map[string]string) - var tokenAccounts []*common.Account - for mint, recordsByType := range recordsByMintAndType { - for _, recordsList := range recordsByType { - for _, records := range recordsList { - if !records.IsTimelock() { - continue - } - - tokenAccount, err := common.NewAccountFromPublicKeyString(records.General.TokenAccount) - if err != nil { - return 0, err - } - - mintByTokenAccount[records.General.TokenAccount] = mint - tokenAccounts = append(tokenAccounts, tokenAccount) - } - } - } - - if len(tokenAccounts) == 0 { - return 0, nil - } - - // Accounts that have left the L2 system don't have a cached balance that can - // be trusted, so the ledger omits them from the result. - balanceByTokenAccount, err := balance.BatchCalculateFromCache(ctx, s.data, tokenAccounts...) + // The ledger holds a record for every account Code manages for the owner, + // and each carries the mint it holds. Accounts that have left the L2 system + // don't have a cached balance that can be trusted, so it omits them. + balanceByTokenAccount, err := balance.BatchCalculateFromCacheByOwner(ctx, s.data, owner) if err != nil { return 0, err } quarksByMint := make(map[string]uint64) - for tokenAccount, cached := range balanceByTokenAccount { - quarksByMint[mintByTokenAccount[tokenAccount]] += cached.Quarks + for _, cached := range balanceByTokenAccount { + quarksByMint[cached.MintAccount] += cached.Quarks } var coreMintValue uint64