From 3bd0fe7cf8507f0f6b3de55a1a1f74ff9f7c0afc Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 11:49:24 -0400 Subject: [PATCH 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 51f5de75442d7d48961e83d674ac6b3b6abaedd9 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 28 Aug 2026 13:41:55 -0400 Subject: [PATCH 9/9] 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) + }) }