From 822d55a2bef41349fe3e020987f5f5f71e7e22b5 Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Tue, 27 May 2025 14:45:46 -0400 Subject: [PATCH 1/8] DB-level changes to support distributed nonce pools --- pkg/code/async/nonce/metrics.go | 1 + pkg/code/async/nonce/pool.go | 30 +- pkg/code/data/internal.go | 4 + pkg/code/data/nonce/memory/store.go | 76 ++- pkg/code/data/nonce/nonce.go | 76 ++- pkg/code/data/nonce/postgres/model.go | 135 +++- pkg/code/data/nonce/postgres/store.go | 15 + pkg/code/data/nonce/postgres/store_test.go | 21 +- pkg/code/data/nonce/store.go | 14 + pkg/code/data/nonce/tests/tests.go | 722 ++++++++++++++------- pkg/code/transaction/nonce.go | 12 +- pkg/code/transaction/nonce_test.go | 73 ++- 12 files changed, 894 insertions(+), 285 deletions(-) diff --git a/pkg/code/async/nonce/metrics.go b/pkg/code/async/nonce/metrics.go index b6dada11..f1d4ee89 100644 --- a/pkg/code/async/nonce/metrics.go +++ b/pkg/code/async/nonce/metrics.go @@ -31,6 +31,7 @@ func (p *service) metricsGaugeWorker(ctx context.Context) error { nonce.StateAvailable, nonce.StateReserved, nonce.StateInvalid, + nonce.StateClaimed, } { count, err := p.data.GetNonceCountByStateAndPurpose(ctx, nonce.EnvironmentCvm, common.CodeVmAccount.PublicKey().ToBase58(), state, nonce.PurposeClientTransaction) if err != nil { diff --git a/pkg/code/async/nonce/pool.go b/pkg/code/async/nonce/pool.go index 548957c6..53e2cd39 100644 --- a/pkg/code/async/nonce/pool.go +++ b/pkg/code/async/nonce/pool.go @@ -95,16 +95,20 @@ func (p *service) handle(ctx context.Context, record *nonce.Record) error { blockhash yet. StateAvailable: - Available to be used by a payment intent, subscription, or - other nonce-related transaction/instruction. + Available to be by a fulfillment for a virtual instruction + or transaction. StateReserved: - Reserved by a payment intent, subscription, or other - nonce-related transaction/instruction. + Reserved by a by a fulfillment for a virtual instruction + or transaction. StateInvalid: The nonce account is invalid (e.g. insufficient funds, etc). + StateClaimed: + The nonce is claimed by a process for future use (identified + by a node ID) + Transitions: StateUnknown -> StateInvalid @@ -112,14 +116,15 @@ func (p *service) handle(ctx context.Context, record *nonce.Record) error { StateReleased -> StateAvailable StateAvailable - -> [externally] StateReserved (nonce used in a new transaction) + -> [externally] StateClaimed (nonce claimed by a process for future use) + StateClaimed + -> [externally] StateAvailable (nonce released by the process that claimed it) + -> [externally] StateReserved (nonce reserved for future use in a virtual instruction or transaction) StateReserved - -> [externally] StateReleased (nonce used in a submitted transaction) - -> [externally] StateAvailable (nonce will never be submitted in the transaction - eg. it became revoked) + -> [externally] StateReleased (nonce used in a submitted virtual instruction or transaction) + -> [externally] StateAvailable (nonce will never be submitted in the virtual instruction or transaction - eg. it became revoked) */ - // todo: distributed lock on the nonce - switch record.State { case nonce.StateUnknown: return p.handleUnknown(ctx, record) @@ -131,6 +136,8 @@ func (p *service) handle(ctx context.Context, record *nonce.Record) error { return p.handleReserved(ctx, record) case nonce.StateInvalid: return p.handleInvalid(ctx, record) + case nonce.StateClaimed: + return p.handleClaimed(ctx, record) default: return nil } @@ -259,3 +266,8 @@ func (p *service) handleInvalid(ctx context.Context, record *nonce.Record) error // as is for further investigation. return nil } + +func (p *service) handleClaimed(ctx context.Context, record *nonce.Record) error { + // Nothing to do here + return nil +} diff --git a/pkg/code/data/internal.go b/pkg/code/data/internal.go index 9c88d881..e8e10677 100644 --- a/pkg/code/data/internal.go +++ b/pkg/code/data/internal.go @@ -136,6 +136,7 @@ type DatabaseData interface { GetNonceCountByStateAndPurpose(ctx context.Context, env nonce.Environment, instance string, state nonce.State, purpose nonce.Purpose) (uint64, error) GetAllNonceByState(ctx context.Context, env nonce.Environment, instance string, state nonce.State, opts ...query.Option) ([]*nonce.Record, error) GetRandomAvailableNonceByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) + BatchClaimAvailableNoncesByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) SaveNonce(ctx context.Context, record *nonce.Record) error // Fulfillment @@ -532,6 +533,9 @@ func (dp *DatabaseProvider) GetAllNonceByState(ctx context.Context, env nonce.En func (dp *DatabaseProvider) GetRandomAvailableNonceByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) { return dp.nonces.GetRandomAvailableByPurpose(ctx, env, instance, purpose) } +func (dp *DatabaseProvider) BatchClaimAvailableNoncesByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { + return dp.nonces.BatchClaimAvailableByPurpose(ctx, env, instance, purpose, limit, nodeID, minExpireAt, maxExpireAt) +} func (dp *DatabaseProvider) SaveNonce(ctx context.Context, record *nonce.Record) error { return dp.nonces.Save(ctx, record) } diff --git a/pkg/code/data/nonce/memory/store.go b/pkg/code/data/nonce/memory/store.go index f89f877c..e657b658 100644 --- a/pkg/code/data/nonce/memory/store.go +++ b/pkg/code/data/nonce/memory/store.go @@ -5,9 +5,11 @@ import ( "math/rand" "sort" "sync" + "time" "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/database/query" + "github.com/code-payments/code-server/pkg/pointer" ) type store struct { @@ -131,6 +133,16 @@ func (s *store) filter(items []*nonce.Record, cursor query.Cursor, limit uint64, return res } +func (s *store) filterAvailable(items []*nonce.Record) []*nonce.Record { + var res []*nonce.Record + for _, item := range items { + if item.IsAvailable() { + res = append(res, item) + } + } + return res +} + func (s *store) Count(ctx context.Context, env nonce.Environment, instance string) (uint64, error) { s.mu.Lock() defer s.mu.Unlock() @@ -165,13 +177,25 @@ func (s *store) Save(ctx context.Context, data *nonce.Record) error { s.last++ if item := s.find(data); item != nil { + if item.Version != data.Version { + return nonce.ErrStaleVersion + } + + data.Version++ + item.Blockhash = data.Blockhash - item.State = data.State item.Signature = data.Signature + item.State = data.State + item.ClaimNodeID = pointer.StringCopy(data.ClaimNodeID) + item.ClaimExpiresAt = pointer.TimeCopy(data.ClaimExpiresAt) + item.Version = data.Version } else { if data.Id == 0 { data.Id = s.last } + + data.Version++ + c := data.Clone() s.records = append(s.records, &c) } @@ -184,7 +208,8 @@ func (s *store) Get(ctx context.Context, address string) (*nonce.Record, error) defer s.mu.Unlock() if item := s.findAddress(address); item != nil { - return item, nil + cloned := item.Clone() + return &cloned, nil } return nil, nonce.ErrNonceNotFound @@ -201,7 +226,7 @@ func (s *store) GetAllByState(ctx context.Context, env nonce.Environment, instan return nil, nonce.ErrNonceNotFound } - return res, nil + return clonedRecords(res), nil } return nil, nonce.ErrNonceNotFound @@ -212,10 +237,53 @@ func (s *store) GetRandomAvailableByPurpose(ctx context.Context, env nonce.Envir defer s.mu.Unlock() items := s.findByStateAndPurpose(env, instance, nonce.StateAvailable, purpose) + items = append(items, s.findByStateAndPurpose(env, instance, nonce.StateClaimed, purpose)...) + items = s.filterAvailable(items) if len(items) == 0 { return nil, nonce.ErrNonceNotFound } index := rand.Intn(len(items)) - return items[index], nil + cloned := items[index].Clone() + return &cloned, nil +} + +func (s *store) BatchClaimAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { + s.mu.Lock() + defer s.mu.Unlock() + + items := s.findByStateAndPurpose(env, instance, nonce.StateAvailable, purpose) + items = append(items, s.findByStateAndPurpose(env, instance, nonce.StateClaimed, purpose)...) + items = s.filterAvailable(items) + if len(items) == 0 { + return nil, nonce.ErrNonceNotFound + } + if len(items) > limit { + items = items[:limit] + } + + for i, l := 0, len(items); i < l; i++ { + j := rand.Intn(l) + items[i], items[j] = items[j], items[i] + } + for i := 0; i < len(items); i++ { + window := maxExpireAt.Sub(minExpireAt) + expiry := minExpireAt.Add(time.Duration(rand.Intn(int(window)))) + + items[i].State = nonce.StateClaimed + items[i].ClaimNodeID = pointer.String(nodeID) + items[i].ClaimExpiresAt = pointer.Time(expiry) + items[i].Version++ + } + + return clonedRecords(items), nil +} + +func clonedRecords(items []*nonce.Record) []*nonce.Record { + res := make([]*nonce.Record, len(items)) + for i, item := range items { + cloned := item.Clone() + res[i] = &cloned + } + return res } diff --git a/pkg/code/data/nonce/nonce.go b/pkg/code/data/nonce/nonce.go index f50b1bca..ba43b167 100644 --- a/pkg/code/data/nonce/nonce.go +++ b/pkg/code/data/nonce/nonce.go @@ -3,7 +3,9 @@ package nonce import ( "crypto/ed25519" "errors" + "time" + "github.com/code-payments/code-server/pkg/pointer" "github.com/mr-tron/base58" ) @@ -22,8 +24,8 @@ const ( ) var ( + ErrStaleVersion = errors.New("nonce version is stale") ErrNonceNotFound = errors.New("no records could be found") - ErrInvalidNonce = errors.New("invalid nonce") ) type State uint8 @@ -31,13 +33,13 @@ type State uint8 const ( StateUnknown State = iota StateReleased // The nonce is almost ready but we don't know its blockhash yet. - StateAvailable // The nonce is available to be used by a payment intent, subscription, or other nonce-related transaction/instruction. - StateReserved // The nonce is reserved by a payment intent, subscription, or other nonce-related transaction/instruction. + StateAvailable // The nonce is available to be used by a fulfillment for a virtual instruction or transaction. + StateReserved // The nonce is reserved by a fulfillment for a virtual instruction or transaction. StateInvalid // The nonce account is invalid (e.g. insufficient funds, etc). + StateClaimed // The nonce is claimed by a process for future use (identified by a node ID) ) -// Split nonce pool across different use cases. This has an added benefit of: -// - Solving for race conditions without distributed locks. +// Split nonce pool across different use cases. This has an added benefit of:. // - Avoiding different use cases from starving each other and ending up in a // deadlocked state. Concretely, it would be really bad if clients could starve // internal processes from creating transactions that would allow us to progress @@ -65,10 +67,11 @@ type Record struct { State State Signature string -} -func (r *Record) GetPublicKey() (ed25519.PublicKey, error) { - return base58.Decode(r.Address) + ClaimNodeID *string + ClaimExpiresAt *time.Time + + Version uint64 } func (r *Record) Clone() Record { @@ -82,6 +85,9 @@ func (r *Record) Clone() Record { Purpose: r.Purpose, State: r.State, Signature: r.Signature, + ClaimNodeID: pointer.StringCopy(r.ClaimNodeID), + ClaimExpiresAt: pointer.TimeCopy(r.ClaimExpiresAt), + Version: r.Version, } } @@ -95,32 +101,72 @@ func (r *Record) CopyTo(dst *Record) { dst.Purpose = r.Purpose dst.State = r.State dst.Signature = r.Signature + dst.ClaimNodeID = pointer.StringCopy(r.ClaimNodeID) + dst.ClaimExpiresAt = pointer.TimeCopy(r.ClaimExpiresAt) + dst.Version = r.Version } -func (v *Record) Validate() error { - if len(v.Address) == 0 { +func (r *Record) Validate() error { + if len(r.Address) == 0 { return errors.New("nonce account address is required") } - if len(v.Authority) == 0 { + if len(r.Authority) == 0 { return errors.New("authority address is required") } - if v.Environment == EnvironmentUnknown { + if r.Environment == EnvironmentUnknown { return errors.New("nonce environment must be set") } - if len(v.EnvironmentInstance) == 0 { + if len(r.EnvironmentInstance) == 0 { return errors.New("nonce environment instance must be set") } - if v.Purpose == PurposeUnknown { + if r.Purpose == PurposeUnknown { return errors.New("nonce purpose must be set") } + switch r.State { + case StateClaimed: + if r.ClaimNodeID == nil { + return errors.New("claim node id is required") + } + + if r.ClaimExpiresAt == nil { + return errors.New("claim expiration timestamp is required") + } + default: + if r.ClaimNodeID != nil { + return errors.New("claim node id cannot be set") + } + + if r.ClaimExpiresAt != nil { + return errors.New("claim expiration timestamp cannot be set") + } + } + return nil } +func (r *Record) IsAvailable() bool { + if r.State == StateAvailable { + return true + } + if r.State != StateClaimed { + return false + } + if r.ClaimExpiresAt == nil { + return false + } + + return time.Now().After(*r.ClaimExpiresAt) +} + +func (r *Record) GetPublicKey() (ed25519.PublicKey, error) { + return base58.Decode(r.Address) +} + func (e Environment) String() string { switch e { case EnvironmentUnknown: @@ -145,6 +191,8 @@ func (s State) String() string { return "reserved" case StateInvalid: return "invalid" + case StateClaimed: + return "claimed" } return "unknown" diff --git a/pkg/code/data/nonce/postgres/model.go b/pkg/code/data/nonce/postgres/model.go index 119e1e40..a97b5be8 100644 --- a/pkg/code/data/nonce/postgres/model.go +++ b/pkg/code/data/nonce/postgres/model.go @@ -3,10 +3,12 @@ package postgres import ( "context" "database/sql" + "time" "github.com/jmoiron/sqlx" "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/pointer" pgutil "github.com/code-payments/code-server/pkg/database/postgres" q "github.com/code-payments/code-server/pkg/database/query" @@ -17,15 +19,18 @@ const ( ) type nonceModel struct { - Id sql.NullInt64 `db:"id"` - Address string `db:"address"` - Authority string `db:"authority"` - Blockhash string `db:"blockhash"` - Environment uint `db:"environment"` - EnvironmentInstance string `db:"environment_instance"` - Purpose uint `db:"purpose"` - State uint `db:"state"` - Signature string `db:"signature"` + Id sql.NullInt64 `db:"id"` + Address string `db:"address"` + Authority string `db:"authority"` + Blockhash string `db:"blockhash"` + Environment uint `db:"environment"` + EnvironmentInstance string `db:"environment_instance"` + Purpose uint `db:"purpose"` + State uint `db:"state"` + Signature string `db:"signature"` + ClaimNodeId sql.NullString `db:"claim_node_id"` + ClaimExpiresAtMs sql.NullInt64 `db:"claim_expires_at"` + Version int64 `db:"version"` } func toNonceModel(obj *nonce.Record) (*nonceModel, error) { @@ -33,6 +38,11 @@ func toNonceModel(obj *nonce.Record) (*nonceModel, error) { return nil, err } + var claimExpiresAtMs int64 + if obj.ClaimExpiresAt != nil { + claimExpiresAtMs = int64(obj.ClaimExpiresAt.UnixMilli()) + } + return &nonceModel{ Id: sql.NullInt64{Int64: int64(obj.Id), Valid: true}, Address: obj.Address, @@ -43,10 +53,25 @@ func toNonceModel(obj *nonce.Record) (*nonceModel, error) { Purpose: uint(obj.Purpose), State: uint(obj.State), Signature: obj.Signature, + ClaimNodeId: sql.NullString{ + Valid: obj.ClaimNodeID != nil, + String: *pointer.StringOrDefault(obj.ClaimNodeID, ""), + }, + ClaimExpiresAtMs: sql.NullInt64{ + Valid: claimExpiresAtMs > 0, + Int64: claimExpiresAtMs, + }, + Version: int64(obj.Version), }, nil } func fromNonceModel(obj *nonceModel) *nonce.Record { + var claimExpiresAt *time.Time + if obj.ClaimExpiresAtMs.Valid { + ts := time.UnixMilli(obj.ClaimExpiresAtMs.Int64) + claimExpiresAt = &ts + } + return &nonce.Record{ Id: uint64(obj.Id.Int64), Address: obj.Address, @@ -57,20 +82,23 @@ func fromNonceModel(obj *nonceModel) *nonce.Record { Purpose: nonce.Purpose(obj.Purpose), State: nonce.State(obj.State), Signature: obj.Signature, + ClaimNodeID: pointer.StringIfValid(obj.ClaimNodeId.Valid, obj.ClaimNodeId.String), + ClaimExpiresAt: claimExpiresAt, + Version: uint64(obj.Version), } } func (m *nonceModel) dbSave(ctx context.Context, db *sqlx.DB) error { return pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { query := `INSERT INTO ` + nonceTableName + ` - (address, authority, blockhash, environment, environment_instance, purpose, state, signature) - VALUES ($1,$2,$3,$4,$5,$6,$7,$8) + (address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11 + 1) ON CONFLICT (address) DO UPDATE - SET blockhash = $3, state = $7, signature = $8 - WHERE ` + nonceTableName + `.address = $1 + SET blockhash = $3, state = $7, signature = $8, claim_node_id = $9, claim_expires_at = $10, version = ` + nonceTableName + `.version + 1 + WHERE ` + nonceTableName + `.address = $1 AND ` + nonceTableName + `.version = $11 RETURNING - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature` + id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version` err := tx.QueryRowxContext( ctx, @@ -83,9 +111,12 @@ func (m *nonceModel) dbSave(ctx context.Context, db *sqlx.DB) error { m.Purpose, m.State, m.Signature, + m.ClaimNodeId, + m.ClaimExpiresAtMs, + m.Version, ).StructScan(m) - return pgutil.CheckNoRows(err, nonce.ErrInvalidNonce) + return pgutil.CheckNoRows(err, nonce.ErrStaleVersion) }) } @@ -129,7 +160,7 @@ func dbGetNonce(ctx context.Context, db *sqlx.DB, address string) (*nonceModel, res := &nonceModel{} query := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature + id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version FROM ` + nonceTableName + ` WHERE address = $1 ` @@ -150,7 +181,7 @@ func dbGetAllByState(ctx context.Context, db *sqlx.DB, env nonce.Environment, in // // todo: Fix said nonce records query := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature + id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version FROM ` + nonceTableName + ` WHERE environment = $1 AND environment_instance = $2 AND state = $3 AND signature IS NOT NULL ` @@ -183,20 +214,21 @@ func dbGetRandomAvailableByPurpose(ctx context.Context, db *sqlx.DB, env nonce.E // // todo: Fix said nonce records query := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature + id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version FROM ` + nonceTableName + ` - WHERE environment = $1 AND environment_instance = $2 AND state = $3 AND purpose = $4 AND signature IS NOT NULL + WHERE environment = $1 AND environment_instance = $2 AND ((state = $3) OR (state = $4 AND claim_expires_at < $5)) AND purpose = $6 AND signature IS NOT NULL OFFSET FLOOR(RANDOM() * 100) LIMIT 1 ` fallbackQuery := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature + id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version FROM ` + nonceTableName + ` - WHERE environment = $1 AND environment_instance = $2 AND state = $3 AND purpose = $4 AND signature IS NOT NULL + WHERE environment = $1 AND environment_instance = $2 AND ((state = $3) OR (state = $4 AND claim_expires_at < $5)) AND purpose = $6 AND signature IS NOT NULL LIMIT 1 ` - err := db.GetContext(ctx, res, query, env, instance, nonce.StateAvailable, purpose) + nowMs := time.Now().UnixMilli() + err := db.GetContext(ctx, res, query, env, instance, nonce.StateAvailable, nonce.StateClaimed, nowMs, purpose) if err != nil { err = pgutil.CheckNoRows(err, nonce.ErrNonceNotFound) @@ -204,7 +236,7 @@ func dbGetRandomAvailableByPurpose(ctx context.Context, db *sqlx.DB, env nonce.E // strategy that will guarantee to select something if an available // nonce exists. if err == nonce.ErrNonceNotFound { - err := db.GetContext(ctx, res, fallbackQuery, env, instance, nonce.StateAvailable, purpose) + err := db.GetContext(ctx, res, fallbackQuery, env, instance, nonce.StateAvailable, nonce.StateClaimed, nowMs, purpose) if err != nil { return nil, pgutil.CheckNoRows(err, nonce.ErrNonceNotFound) } @@ -215,3 +247,60 @@ func dbGetRandomAvailableByPurpose(ctx context.Context, db *sqlx.DB, env nonce.E } return res, nil } + +func dbBatchClaimAvailableByPurpose( + ctx context.Context, + db *sqlx.DB, + env nonce.Environment, + instance string, + purpose nonce.Purpose, + limit int, + nodeID string, + minExpireAt time.Time, + maxExpireAt time.Time, +) ([]*nonceModel, error) { + // Signature null check is required because some legacy records didn't have this + // set and causes this call to fail. This is a result of the field not being + // defined at the time of record creation. + // + // todo: Fix said nonce records + res := []*nonceModel{} + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + query := `WITH selected_nonces AS ( + SELECT id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version + FROM ` + nonceTableName + ` + WHERE environment = $1 AND environment_instance = $2 AND ((state = $3) OR (state = $4 AND claim_expires_at < $5)) AND purpose = $6 AND signature IS NOT NULL + LIMIT $7 + FOR UPDATE + ) + UPDATE ` + nonceTableName + ` + SET state = $8, claim_node_id = $9, claim_expires_at = $10 + FLOOR(RANDOM() * $11), version = ` + nonceTableName + `.version + 1 + FROM selected_nonces + WHERE ` + nonceTableName + `.id = selected_nonces.id + RETURNING ` + nonceTableName + `.*` + + return tx.SelectContext( + ctx, + &res, + query, + env, + instance, + nonce.StateAvailable, + nonce.StateClaimed, + time.Now().UnixMilli(), + purpose, + limit, + nonce.StateClaimed, + nodeID, + minExpireAt.UnixMilli(), + maxExpireAt.Sub(minExpireAt).Milliseconds(), + ) + }) + if err != nil { + return nil, pgutil.CheckNoRows(err, nonce.ErrNonceNotFound) + } + if len(res) == 0 { + return nil, nonce.ErrNonceNotFound + } + return res, nil +} diff --git a/pkg/code/data/nonce/postgres/store.go b/pkg/code/data/nonce/postgres/store.go index adf2b98e..f523c05b 100644 --- a/pkg/code/data/nonce/postgres/store.go +++ b/pkg/code/data/nonce/postgres/store.go @@ -3,6 +3,7 @@ package postgres import ( "context" "database/sql" + "time" "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/database/query" @@ -78,3 +79,17 @@ func (s *store) GetRandomAvailableByPurpose(ctx context.Context, env nonce.Envir } return fromNonceModel(model), nil } + +func (s *store) BatchClaimAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { + models, err := dbBatchClaimAvailableByPurpose(ctx, s.db, env, instance, purpose, limit, nodeID, minExpireAt, maxExpireAt) + if err != nil { + return nil, err + } + + nonces := make([]*nonce.Record, len(models)) + for i, model := range models { + nonces[i] = fromNonceModel(model) + } + + return nonces, nil +} diff --git a/pkg/code/data/nonce/postgres/store_test.go b/pkg/code/data/nonce/postgres/store_test.go index fe1fb4a5..7db68368 100644 --- a/pkg/code/data/nonce/postgres/store_test.go +++ b/pkg/code/data/nonce/postgres/store_test.go @@ -22,17 +22,22 @@ const ( CREATE TABLE codewallet__core_nonce( id SERIAL NOT NULL PRIMARY KEY, - address text NOT NULL UNIQUE, - authority text NOT NULL, - blockhash text NULL, + address TEXT NOT NULL UNIQUE, + authority TEXT NOT NULL, + blockhash TEXT NULL, - environment integer NOT NULL, - environment_instance text NOT NULL, + environment INTEGER NOT NULL, + environment_instance TEXT NOT NULL, - purpose integer NOT NULL, - state integer NOT NULL, + purpose INTEGER NOT NULL, + state INTEGER NOT NULL, - signature text NULL + signature TEXT NULL, + + claim_node_id TEXT NULL, + claim_expires_at BIGINT NULL, + + version BIGINT NOT NULL ); ` diff --git a/pkg/code/data/nonce/store.go b/pkg/code/data/nonce/store.go index 4bc8e01c..35cb23d2 100644 --- a/pkg/code/data/nonce/store.go +++ b/pkg/code/data/nonce/store.go @@ -2,6 +2,7 @@ package nonce import ( "context" + "time" "github.com/code-payments/code-server/pkg/database/query" ) @@ -36,5 +37,18 @@ type Store interface { // an environment instance. // // Returns ErrNotFound if no records are found. + // + // Deprecated in favour of BatchClaimAvailableByPurpose GetRandomAvailableByPurpose(ctx context.Context, env Environment, instance string, purpose Purpose) (*Record, error) + + // BatchClaimAvailableByPurpose batch claims up to the specified limit. + // + // The returned nonces will be marked as claimed by the current node, with + // the specified expiry date. + // + // Note: Implementations need not randomize the results/selection. + // The transactional nature of the call means that any contention exists + // on the tx level (which always occurs), and not around fighting over + // individual nonces. + BatchClaimAvailableByPurpose(ctx context.Context, env Environment, instance string, purpose Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*Record, error) } diff --git a/pkg/code/data/nonce/tests/tests.go b/pkg/code/data/nonce/tests/tests.go index 6d90a2d0..15b7afc1 100644 --- a/pkg/code/data/nonce/tests/tests.go +++ b/pkg/code/data/nonce/tests/tests.go @@ -3,13 +3,18 @@ package tests import ( "context" "fmt" + "math" + "slices" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/database/query" + "github.com/code-payments/code-server/pkg/pointer" ) func RunTests(t *testing.T, s nonce.Store, teardown func()) { @@ -19,6 +24,8 @@ func RunTests(t *testing.T, s nonce.Store, teardown func()) { testGetAllByState, testGetCount, testGetRandomAvailableByPurpose, + testBatchClaimAvailableByPurpose, + testBatchClaimAvailableByPurposeExpirationRandomness, } { tf(t, s) teardown() @@ -26,238 +33,256 @@ func RunTests(t *testing.T, s nonce.Store, teardown func()) { } func testRoundTrip(t *testing.T, s nonce.Store) { - ctx := context.Background() - - actual, err := s.Get(ctx, "test_address") - require.Error(t, err) - assert.Equal(t, nonce.ErrNonceNotFound, err) - assert.Nil(t, actual) - - expected := nonce.Record{ - Address: "test_address", - Authority: "test_authority", - Blockhash: "test_blockhash", - Environment: nonce.EnvironmentSolana, - EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, - Purpose: nonce.PurposeClientTransaction, - State: nonce.StateReserved, - } - cloned := expected.Clone() - err = s.Save(ctx, &expected) - require.NoError(t, err) - - actual, err = s.Get(ctx, "test_address") - require.NoError(t, err) - assert.Equal(t, cloned.Address, actual.Address) - assert.Equal(t, cloned.Authority, actual.Authority) - assert.Equal(t, cloned.Blockhash, actual.Blockhash) - assert.Equal(t, cloned.Environment, actual.Environment) - assert.Equal(t, cloned.EnvironmentInstance, actual.EnvironmentInstance) - assert.Equal(t, cloned.Purpose, actual.Purpose) - assert.Equal(t, cloned.State, actual.State) - assert.EqualValues(t, 1, actual.Id) + t.Run("testRoundTrip", func(t *testing.T) { + ctx := context.Background() + + actual, err := s.Get(ctx, "test_address") + require.Error(t, err) + assert.Equal(t, nonce.ErrNonceNotFound, err) + assert.Nil(t, actual) + + expected := nonce.Record{ + Address: "test_address", + Authority: "test_authority", + Blockhash: "test_blockhash", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, + Purpose: nonce.PurposeClientTransaction, + State: nonce.StateClaimed, + ClaimNodeID: pointer.String("test_claim_node_id"), + ClaimExpiresAt: pointer.Time(time.Now().Add(time.Hour)), + } + cloned := expected.Clone() + err = s.Save(ctx, &expected) + require.NoError(t, err) + assert.EqualValues(t, 1, expected.Id) + assert.EqualValues(t, 1, expected.Version) + + actual, err = s.Get(ctx, "test_address") + require.NoError(t, err) + assertEquivalentRecords(t, &cloned, actual) + assert.EqualValues(t, 1, actual.Id) + assert.EqualValues(t, 1, actual.Version) + }) } func testUpdate(t *testing.T, s nonce.Store) { - ctx := context.Background() - - expected := nonce.Record{ - Address: "test_address", - Authority: "test_authority", - Blockhash: "test_blockhash", - Environment: nonce.EnvironmentSolana, - EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, - Purpose: nonce.PurposeInternalServerProcess, - } - err := s.Save(ctx, &expected) - require.NoError(t, err) - assert.EqualValues(t, 1, expected.Id) - - expected.State = nonce.StateUnknown - expected.Signature = "test_signature" - - err = s.Save(ctx, &expected) - require.NoError(t, err) - - actual, err := s.Get(ctx, "test_address") - require.NoError(t, err) - assert.Equal(t, expected.Address, actual.Address) - assert.Equal(t, expected.Authority, actual.Authority) - assert.Equal(t, expected.Blockhash, actual.Blockhash) - assert.Equal(t, expected.Purpose, actual.Purpose) - assert.Equal(t, expected.State, actual.State) - assert.Equal(t, expected.Signature, actual.Signature) - assert.EqualValues(t, 1, actual.Id) + t.Run("testUpdate", func(t *testing.T) { + ctx := context.Background() + + expected := nonce.Record{ + Address: "test_address", + Authority: "test_authority", + Blockhash: "test_blockhash", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, + Purpose: nonce.PurposeInternalServerProcess, + } + cloned := expected.Clone() + err := s.Save(ctx, &expected) + require.NoError(t, err) + assert.EqualValues(t, 1, expected.Id) + assert.EqualValues(t, 1, expected.Version) + + actual, err := s.Get(ctx, "test_address") + require.NoError(t, err) + assertEquivalentRecords(t, &cloned, actual) + assert.EqualValues(t, 1, actual.Id) + assert.EqualValues(t, 1, actual.Version) + + expected = actual.Clone() + expected.State = nonce.StateClaimed + expected.Blockhash = "test_blockhash2" + expected.Signature = "test_signature" + expected.ClaimNodeID = pointer.String("test_claim_node_id") + expected.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) + cloned = expected.Clone() + + err = s.Save(ctx, &expected) + require.NoError(t, err) + assert.EqualValues(t, 1, expected.Id) + assert.EqualValues(t, 2, expected.Version) + + actual, err = s.Get(ctx, "test_address") + require.NoError(t, err) + assertEquivalentRecords(t, &cloned, actual) + assert.EqualValues(t, 1, actual.Id) + assert.EqualValues(t, 2, actual.Version) + }) } func testGetAllByState(t *testing.T, s nonce.Store) { - ctx := context.Background() - - expected := []nonce.Record{ - {Address: "t1", Authority: "a1", Blockhash: "b1", State: nonce.StateUnknown, Signature: "s1"}, - {Address: "t2", Authority: "a2", Blockhash: "b1", State: nonce.StateInvalid, Signature: "s2"}, - {Address: "t3", Authority: "a3", Blockhash: "b1", State: nonce.StateReserved, Signature: "s3"}, - {Address: "t4", Authority: "a1", Blockhash: "b2", State: nonce.StateReserved, Signature: "s4"}, - {Address: "t5", Authority: "a2", Blockhash: "b2", State: nonce.StateReserved, Signature: "s5"}, - {Address: "t6", Authority: "a3", Blockhash: "b2", State: nonce.StateInvalid, Signature: "s6"}, - } + t.Run("testGetAllByState", func(t *testing.T) { + ctx := context.Background() - for _, item := range expected { - item.Environment = nonce.EnvironmentSolana - item.EnvironmentInstance = nonce.EnvironmentInstanceSolanaMainnet - item.Purpose = nonce.PurposeInternalServerProcess + expected := []nonce.Record{ + {Address: "t1", Authority: "a1", Blockhash: "b1", State: nonce.StateUnknown, Signature: "s1"}, + {Address: "t2", Authority: "a2", Blockhash: "b1", State: nonce.StateInvalid, Signature: "s2"}, + {Address: "t3", Authority: "a3", Blockhash: "b1", State: nonce.StateReserved, Signature: "s3"}, + {Address: "t4", Authority: "a1", Blockhash: "b2", State: nonce.StateReserved, Signature: "s4"}, + {Address: "t5", Authority: "a2", Blockhash: "b2", State: nonce.StateReserved, Signature: "s5"}, + {Address: "t6", Authority: "a3", Blockhash: "b2", State: nonce.StateInvalid, Signature: "s6"}, + } + + for _, item := range expected { + item.Environment = nonce.EnvironmentSolana + item.EnvironmentInstance = nonce.EnvironmentInstanceSolanaMainnet + item.Purpose = nonce.PurposeInternalServerProcess + + err := s.Save(ctx, &item) + require.NoError(t, err) + } - err := s.Save(ctx, &item) + // Simple get all by state + actual, err := s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Ascending) require.NoError(t, err) - } + assert.Equal(t, 3, len(actual)) - // Simple get all by state - actual, err := s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, query.EmptyCursor, 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 1, len(actual)) - - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid, query.EmptyCursor, 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 2, len(actual)) - - // Simple get all by state (reverse) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, query.EmptyCursor, 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 1, len(actual)) - - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid, query.EmptyCursor, 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 2, len(actual)) - - // Check items (asc) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - assert.Equal(t, "t3", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - assert.Equal(t, "t5", actual[2].Address) - - // Check items (desc) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - assert.Equal(t, "t5", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - assert.Equal(t, "t3", actual[2].Address) - - // Check items (asc + limit) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 2, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 2, len(actual)) - assert.Equal(t, "t3", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - - // Check items (desc + limit) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 2, query.Descending) - require.NoError(t, err) - assert.Equal(t, 2, len(actual)) - assert.Equal(t, "t5", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - - // Check items (asc + cursor) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(1), 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - assert.Equal(t, "t3", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - assert.Equal(t, "t5", actual[2].Address) - - // Check items (desc + cursor) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(6), 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 3, len(actual)) - assert.Equal(t, "t5", actual[0].Address) - assert.Equal(t, "t4", actual[1].Address) - assert.Equal(t, "t3", actual[2].Address) - - // Check items (asc + cursor) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(3), 5, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 2, len(actual)) - assert.Equal(t, "t4", actual[0].Address) - assert.Equal(t, "t5", actual[1].Address) - - // Check items (desc + cursor) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(4), 5, query.Descending) - require.NoError(t, err) - assert.Equal(t, 1, len(actual)) - assert.Equal(t, "t3", actual[0].Address) - - // Check items (asc + cursor + limit) - actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(3), 1, query.Ascending) - require.NoError(t, err) - assert.Equal(t, 1, len(actual)) - assert.Equal(t, "t4", actual[0].Address) -} + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, query.EmptyCursor, 5, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 1, len(actual)) -func testGetCount(t *testing.T, s nonce.Store) { - ctx := context.Background() - - expected := []nonce.Record{ - {Address: "t1", Authority: "a1", Blockhash: "b1", State: nonce.StateUnknown, Purpose: nonce.PurposeClientTransaction, Signature: "s1"}, - {Address: "t2", Authority: "a2", Blockhash: "b1", State: nonce.StateInvalid, Purpose: nonce.PurposeClientTransaction, Signature: "s2"}, - {Address: "t3", Authority: "a3", Blockhash: "b1", State: nonce.StateReserved, Purpose: nonce.PurposeClientTransaction, Signature: "s3"}, - {Address: "t4", Authority: "a1", Blockhash: "b2", State: nonce.StateReserved, Purpose: nonce.PurposeClientTransaction, Signature: "s4"}, - {Address: "t5", Authority: "a2", Blockhash: "b2", State: nonce.StateReserved, Purpose: nonce.PurposeInternalServerProcess, Signature: "s5"}, - {Address: "t6", Authority: "a3", Blockhash: "b2", State: nonce.StateInvalid, Purpose: nonce.PurposeClientTransaction, Signature: "s6"}, - } + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid, query.EmptyCursor, 5, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 2, len(actual)) - for index, item := range expected { - item.Environment = nonce.EnvironmentSolana - item.EnvironmentInstance = nonce.EnvironmentInstanceSolanaMainnet + // Simple get all by state (reverse) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Descending) + require.NoError(t, err) + assert.Equal(t, 3, len(actual)) - count, err := s.Count(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, query.EmptyCursor, 5, query.Descending) require.NoError(t, err) - assert.EqualValues(t, index, count) + assert.Equal(t, 1, len(actual)) - err = s.Save(ctx, &item) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid, query.EmptyCursor, 5, query.Descending) require.NoError(t, err) - } + assert.Equal(t, 2, len(actual)) + + // Check items (asc) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 3, len(actual)) + assert.Equal(t, "t3", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) + assert.Equal(t, "t5", actual[2].Address) - count, err := s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateAvailable) - require.NoError(t, err) - assert.EqualValues(t, 0, count) + // Check items (desc) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 5, query.Descending) + require.NoError(t, err) + assert.Equal(t, 3, len(actual)) + assert.Equal(t, "t5", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) + assert.Equal(t, "t3", actual[2].Address) - count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown) - require.NoError(t, err) - assert.EqualValues(t, 1, count) + // Check items (asc + limit) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 2, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 2, len(actual)) + assert.Equal(t, "t3", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) - count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid) - require.NoError(t, err) - assert.EqualValues(t, 2, count) + // Check items (desc + limit) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.EmptyCursor, 2, query.Descending) + require.NoError(t, err) + assert.Equal(t, 2, len(actual)) + assert.Equal(t, "t5", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) - count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved) - require.NoError(t, err) - assert.EqualValues(t, 3, count) + // Check items (asc + cursor) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(1), 5, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 3, len(actual)) + assert.Equal(t, "t3", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) + assert.Equal(t, "t5", actual[2].Address) - count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, nonce.PurposeClientTransaction) - require.NoError(t, err) - assert.EqualValues(t, 2, count) + // Check items (desc + cursor) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(6), 5, query.Descending) + require.NoError(t, err) + assert.Equal(t, 3, len(actual)) + assert.Equal(t, "t5", actual[0].Address) + assert.Equal(t, "t4", actual[1].Address) + assert.Equal(t, "t3", actual[2].Address) - count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, nonce.PurposeInternalServerProcess) - require.NoError(t, err) - assert.EqualValues(t, 1, count) + // Check items (asc + cursor) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(3), 5, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 2, len(actual)) + assert.Equal(t, "t4", actual[0].Address) + assert.Equal(t, "t5", actual[1].Address) - count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, nonce.PurposeClientTransaction) - require.NoError(t, err) - assert.EqualValues(t, 1, count) + // Check items (desc + cursor) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(4), 5, query.Descending) + require.NoError(t, err) + assert.Equal(t, 1, len(actual)) + assert.Equal(t, "t3", actual[0].Address) - count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, nonce.PurposeInternalServerProcess) - require.NoError(t, err) - assert.EqualValues(t, 0, count) + // Check items (asc + cursor + limit) + actual, err = s.GetAllByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, query.ToCursor(3), 1, query.Ascending) + require.NoError(t, err) + assert.Equal(t, 1, len(actual)) + assert.Equal(t, "t4", actual[0].Address) + }) +} + +func testGetCount(t *testing.T, s nonce.Store) { + t.Run("testGetCount", func(t *testing.T) { + ctx := context.Background() + + expected := []nonce.Record{ + {Address: "t1", Authority: "a1", Blockhash: "b1", State: nonce.StateUnknown, Purpose: nonce.PurposeClientTransaction, Signature: "s1"}, + {Address: "t2", Authority: "a2", Blockhash: "b1", State: nonce.StateInvalid, Purpose: nonce.PurposeClientTransaction, Signature: "s2"}, + {Address: "t3", Authority: "a3", Blockhash: "b1", State: nonce.StateReserved, Purpose: nonce.PurposeClientTransaction, Signature: "s3"}, + {Address: "t4", Authority: "a1", Blockhash: "b2", State: nonce.StateReserved, Purpose: nonce.PurposeClientTransaction, Signature: "s4"}, + {Address: "t5", Authority: "a2", Blockhash: "b2", State: nonce.StateReserved, Purpose: nonce.PurposeInternalServerProcess, Signature: "s5"}, + {Address: "t6", Authority: "a3", Blockhash: "b2", State: nonce.StateInvalid, Purpose: nonce.PurposeClientTransaction, Signature: "s6"}, + } + + for index, item := range expected { + item.Environment = nonce.EnvironmentSolana + item.EnvironmentInstance = nonce.EnvironmentInstanceSolanaMainnet + + count, err := s.Count(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet) + require.NoError(t, err) + assert.EqualValues(t, index, count) + + err = s.Save(ctx, &item) + require.NoError(t, err) + } + + count, err := s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateAvailable) + require.NoError(t, err) + assert.EqualValues(t, 0, count) + + count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown) + require.NoError(t, err) + assert.EqualValues(t, 1, count) + + count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateInvalid) + require.NoError(t, err) + assert.EqualValues(t, 2, count) + + count, err = s.CountByState(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved) + require.NoError(t, err) + assert.EqualValues(t, 3, count) + + count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, nonce.PurposeClientTransaction) + require.NoError(t, err) + assert.EqualValues(t, 2, count) + + count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateReserved, nonce.PurposeInternalServerProcess) + require.NoError(t, err) + assert.EqualValues(t, 1, count) + + count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, nonce.PurposeClientTransaction) + require.NoError(t, err) + assert.EqualValues(t, 1, count) + + count, err = s.CountByStateAndPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.StateUnknown, nonce.PurposeInternalServerProcess) + require.NoError(t, err) + assert.EqualValues(t, 0, count) + }) } func testGetRandomAvailableByPurpose(t *testing.T, s nonce.Store) { @@ -275,8 +300,10 @@ func testGetRandomAvailableByPurpose(t *testing.T, s nonce.Store) { nonce.StateUnknown, nonce.StateAvailable, nonce.StateReserved, + nonce.StateReleased, + nonce.StateClaimed, } { - for i := 0; i < 500; i++ { + for i := 0; i < 50; i++ { record := &nonce.Record{ Address: fmt.Sprintf("nonce_%s_%s_%d", purpose, state, i), Authority: "authority", @@ -287,29 +314,280 @@ func testGetRandomAvailableByPurpose(t *testing.T, s nonce.Store) { State: state, Signature: "", } + if state == nonce.StateClaimed { + record.ClaimNodeID = pointer.String("node_id") + + if i < 25 { + record.ClaimExpiresAt = pointer.Time(time.Now().Add(-time.Hour)) + } else { + record.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) + } + } require.NoError(t, s.Save(ctx, record)) } } } - selectedByAddress := make(map[string]struct{}) for i := 0; i < 100; i++ { - actual, err := s.GetRandomAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.NoError(t, err) - assert.Equal(t, nonce.PurposeClientTransaction, actual.Purpose) - assert.Equal(t, nonce.StateAvailable, actual.State) - selectedByAddress[actual.Address] = struct{}{} + record := &nonce.Record{ + Address: fmt.Sprintf("nonce_devnet_%d", i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaDevnet, + Purpose: nonce.PurposeInternalServerProcess, + State: nonce.StateAvailable, + Signature: "", + } + require.NoError(t, s.Save(ctx, record)) + + record = &nonce.Record{ + Address: fmt.Sprintf("nonce_cvm_%d", i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentCvm, + EnvironmentInstance: "pubkey", + Purpose: nonce.PurposeClientTransaction, + State: nonce.StateClaimed, + Signature: "", + ClaimNodeID: pointer.String("node_id"), + ClaimExpiresAt: pointer.Time(time.Now().Add(-time.Hour)), + } + require.NoError(t, s.Save(ctx, record)) + } + + for _, purpose := range []nonce.Purpose{nonce.PurposeClientTransaction, nonce.PurposeInternalServerProcess} { + availableState, claimedState := 0, 0 + selectedByAddress := make(map[string]struct{}) + for i := 0; i < 100; i++ { + actual, err := s.GetRandomAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, purpose) + require.NoError(t, err) + assert.Equal(t, purpose, actual.Purpose) + assert.Equal(t, nonce.EnvironmentSolana, actual.Environment) + assert.Equal(t, nonce.EnvironmentInstanceSolanaMainnet, actual.EnvironmentInstance) + assert.True(t, actual.IsAvailable()) + switch actual.State { + case nonce.StateAvailable: + availableState++ + case nonce.StateClaimed: + claimedState++ + assert.True(t, time.Now().After(*actual.ClaimExpiresAt)) + default: + } + selectedByAddress[actual.Address] = struct{}{} + } + assert.True(t, len(selectedByAddress) > 10) + assert.NotZero(t, availableState) + assert.NotZero(t, claimedState) + } + }) +} + +func testBatchClaimAvailableByPurpose(t *testing.T, s nonce.Store) { + t.Run("testBatchClaimAvailableByPurpose", func(t *testing.T) { + ctx := context.Background() + + minExpiry := time.Now().Add(time.Hour).Truncate(time.Millisecond) + maxExpiry := time.Now().Add(2 * time.Hour).Truncate(time.Millisecond) + + nonces, err := s.BatchClaimAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 100, "my_node_id", minExpiry, maxExpiry) + require.Equal(t, nonce.ErrNonceNotFound, err) + require.Empty(t, nonces) + + for _, purpose := range []nonce.Purpose{ + nonce.PurposeClientTransaction, + nonce.PurposeInternalServerProcess, + } { + for _, state := range []nonce.State{ + nonce.StateUnknown, + nonce.StateAvailable, + nonce.StateReserved, + nonce.StateClaimed, + } { + for i := 0; i < 50; i++ { + record := &nonce.Record{ + Address: fmt.Sprintf("nonce_%s_%s_%d", purpose, state, i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, + Purpose: purpose, + State: state, + Signature: "", + } + if state == nonce.StateClaimed { + record.ClaimNodeID = pointer.String("other_node_id") + + if i < 25 { + record.ClaimExpiresAt = pointer.Time(time.Now().Add(-time.Hour)) + } else { + record.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) + } + } + + require.NoError(t, s.Save(ctx, record)) + } + } } - assert.True(t, len(selectedByAddress) > 10) - selectedByAddress = make(map[string]struct{}) for i := 0; i < 100; i++ { - actual, err := s.GetRandomAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeInternalServerProcess) + record := &nonce.Record{ + Address: fmt.Sprintf("nonce_devnet_%d", i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaDevnet, + Purpose: nonce.PurposeInternalServerProcess, + State: nonce.StateAvailable, + Signature: "", + } + require.NoError(t, s.Save(ctx, record)) + + record = &nonce.Record{ + Address: fmt.Sprintf("nonce_cvm_%d", i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentCvm, + EnvironmentInstance: "pubkey", + Purpose: nonce.PurposeClientTransaction, + State: nonce.StateClaimed, + Signature: "", + ClaimNodeID: pointer.String("other_node_id"), + ClaimExpiresAt: pointer.Time(time.Now().Add(-time.Hour)), + } + require.NoError(t, s.Save(ctx, record)) + } + + var claimed []*nonce.Record + for remaining := 75; remaining > 0; { + nonces, err = s.BatchClaimAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 10, "my_node_d", minExpiry, maxExpiry) require.NoError(t, err) - assert.Equal(t, nonce.PurposeInternalServerProcess, actual.Purpose) - assert.Equal(t, nonce.StateAvailable, actual.State) - selectedByAddress[actual.Address] = struct{}{} + require.Len(t, nonces, min(remaining, 10)) + + remaining -= len(nonces) + + for _, n := range nonces { + actual, err := s.Get(ctx, n.Address) + require.NoError(t, err) + require.Equal(t, nonce.StateClaimed, actual.State) + require.NotNil(t, actual.ClaimNodeID) + require.NotNil(t, actual.ClaimExpiresAt) + require.Equal(t, "my_node_d", *actual.ClaimNodeID) + require.GreaterOrEqual(t, *actual.ClaimExpiresAt, minExpiry) + require.LessOrEqual(t, *actual.ClaimExpiresAt, maxExpiry) + require.Equal(t, nonce.EnvironmentSolana, actual.Environment) + require.Equal(t, nonce.EnvironmentInstanceSolanaMainnet, actual.EnvironmentInstance) + require.Equal(t, nonce.PurposeClientTransaction, actual.Purpose) + require.EqualValues(t, 2, actual.Version) + + claimed = append(claimed, actual) + } + } + + nonces, err = s.BatchClaimAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 10, "my_node_id", minExpiry, maxExpiry) + require.Equal(t, nonce.ErrNonceNotFound, err) + require.Empty(t, nonces) + + for i := range claimed[:20] { + claimed[i].State = nonce.StateAvailable + claimed[i].ClaimNodeID = nil + claimed[i].ClaimExpiresAt = nil + s.Save(ctx, claimed[i]) + } + + nonces, err = s.BatchClaimAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 30, "my_node_id2", minExpiry, maxExpiry) + require.NoError(t, err) + require.Len(t, nonces, 20) + + slices.SortFunc(claimed[:20], func(a, b *nonce.Record) int { + return strings.Compare(a.Address, b.Address) + }) + slices.SortFunc(nonces, func(a, b *nonce.Record) int { + return strings.Compare(a.Address, b.Address) + }) + + for i, actual := range nonces { + require.Equal(t, nonce.StateClaimed, actual.State) + require.NotNil(t, actual.ClaimNodeID) + require.NotNil(t, actual.ClaimExpiresAt) + require.Equal(t, "my_node_id2", *actual.ClaimNodeID) + require.GreaterOrEqual(t, *actual.ClaimExpiresAt, minExpiry) + require.LessOrEqual(t, *actual.ClaimExpiresAt, maxExpiry) + require.Equal(t, nonce.EnvironmentSolana, actual.Environment) + require.Equal(t, nonce.EnvironmentInstanceSolanaMainnet, actual.EnvironmentInstance) + require.Equal(t, nonce.PurposeClientTransaction, actual.Purpose) + require.Equal(t, claimed[i].Address, actual.Address) + require.EqualValues(t, 4, actual.Version) } - assert.True(t, len(selectedByAddress) > 10) }) } + +func testBatchClaimAvailableByPurposeExpirationRandomness(t *testing.T, s nonce.Store) { + t.Run("testBatchClaimAvailableByPurposeExpirationRandomness", func(t *testing.T) { + ctx := context.Background() + + min := time.Now().Add(time.Hour).Truncate(time.Millisecond) + max := time.Now().Add(2 * time.Hour).Truncate(time.Millisecond) + + for i := 0; i < 1000; i++ { + record := &nonce.Record{ + Address: fmt.Sprintf("nonce_%s_%s_%d", nonce.PurposeClientTransaction, nonce.StateAvailable, i), + Authority: "authority", + Blockhash: "bh", + Environment: nonce.EnvironmentSolana, + EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, + Purpose: nonce.PurposeClientTransaction, + State: nonce.StateAvailable, + Signature: "", + } + + require.NoError(t, s.Save(ctx, record)) + } + + nonces, err := s.BatchClaimAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 1000, "my_node_id", min, max) + require.NoError(t, err) + require.Len(t, nonces, 1000) + + // To verify that we have a rough random distribution of expirations, + // we bucket the expiration space, and compute the standard deviation. + // + // We then compare against the expected value with a tolerance. + // Specifically, we know there should be 50 nonces per bucket in + // an ideal world, and we allow for a 15% deviation on this. + bins := make([]int64, 20) + expected := float64(len(nonces)) / float64(len(bins)) + for _, n := range nonces { + // Formula: bin = k(val - min) / (max-min+1) + // + // We use '+1' in the divisor to ensure we don't divide by zero. + // In practive, this should produce pretty much no bias since our + // testing ranges are large. + bin := int(n.ClaimExpiresAt.Sub(min).Milliseconds()) * len(bins) / int(max.Sub(min).Milliseconds()+1) + bins[bin]++ + } + + sum := 0.0 + for _, count := range bins { + diff := float64(count) - expected + sum += diff * diff + } + + stdDev := math.Sqrt(sum / float64(len(bins))) + assert.LessOrEqual(t, stdDev, 0.15*expected, "expected: %v, bins %v:", expected, bins) + }) +} + +func assertEquivalentRecords(t *testing.T, obj1, obj2 *nonce.Record) { + assert.Equal(t, obj1.Address, obj2.Address) + assert.Equal(t, obj1.Authority, obj2.Authority) + assert.Equal(t, obj1.Blockhash, obj2.Blockhash) + assert.Equal(t, obj1.Environment, obj2.Environment) + assert.Equal(t, obj1.EnvironmentInstance, obj2.EnvironmentInstance) + assert.Equal(t, obj1.Purpose, obj2.Purpose) + assert.Equal(t, obj1.State, obj2.State) + assert.Equal(t, obj1.ClaimNodeID, obj2.ClaimNodeID) + assert.Equal(t, obj1.ClaimExpiresAt == nil, obj2.ClaimExpiresAt == nil) + if obj1.ClaimExpiresAt != nil { + assert.Equal(t, obj1.ClaimExpiresAt.UnixMilli(), obj2.ClaimExpiresAt.UnixMilli()) + } +} diff --git a/pkg/code/transaction/nonce.go b/pkg/code/transaction/nonce.go index 50710fae..fdc5fe56 100644 --- a/pkg/code/transaction/nonce.go +++ b/pkg/code/transaction/nonce.go @@ -83,7 +83,7 @@ func SelectAvailableNonce(ctx context.Context, data code_data.Provider, env nonc return err } - if record.State != nonce.StateAvailable { + if !record.IsAvailable() { // Unlock and try again lock.Unlock() return errors.New("selected nonce that became unavailable") @@ -104,6 +104,8 @@ func SelectAvailableNonce(ctx context.Context, data code_data.Provider, env nonc // Reserve the nonce for use with a fulfillment record.State = nonce.StateReserved + record.ClaimNodeID = nil + record.ClaimExpiresAt = nil err = data.SaveNonce(ctx, record) if err != nil { lock.Unlock() @@ -152,12 +154,14 @@ func (n *SelectedNonce) MarkReservedWithSignature(ctx context.Context, sig strin return n.data.SaveNonce(ctx, n.record) } - if n.record.State != nonce.StateAvailable { + if !n.record.IsAvailable() { return errors.New("nonce must be available to reserve") } n.record.State = nonce.StateReserved n.record.Signature = sig + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil return n.data.SaveNonce(ctx, n.record) } @@ -173,7 +177,7 @@ func (n *SelectedNonce) ReleaseIfNotReserved() error { return errors.New("nonce is unlocked") } - if n.record.State == nonce.StateAvailable { + if n.record.IsAvailable() { return nil } @@ -184,6 +188,8 @@ func (n *SelectedNonce) ReleaseIfNotReserved() error { defer cancel() n.record.State = nonce.StateAvailable + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil return n.data.SaveNonce(ctx, n.record) } diff --git a/pkg/code/transaction/nonce_test.go b/pkg/code/transaction/nonce_test.go index 766ea9eb..8b487ec9 100644 --- a/pkg/code/transaction/nonce_test.go +++ b/pkg/code/transaction/nonce_test.go @@ -14,6 +14,7 @@ import ( code_data "github.com/code-payments/code-server/pkg/code/data" "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/code/data/vault" + "github.com/code-payments/code-server/pkg/pointer" "github.com/code-payments/code-server/pkg/solana" "github.com/code-payments/code-server/pkg/testutil" ) @@ -58,10 +59,45 @@ func TestNonce_SelectAvailableNonce(t *testing.T) { _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) assert.Equal(t, ErrNoAvailableNonces, err) +} + +func TestNonce_SelectAvailableNonceClaimed(t *testing.T) { + env := setupNonceTestEnv(t) + + expiredNonces := map[string]*nonce.Record{} + for i := 0; i < 10; i++ { + n := generateClaimedNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, true) + expiredNonces[n.Address] = n - _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeInternalServerProcess) + generateClaimedNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, false) + } + + _, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeInternalServerProcess) assert.Equal(t, ErrNoAvailableNonces, err) + _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentCvm, testutil.NewRandomAccount(t).PublicKey().ToBase58(), nonce.PurposeClientTransaction) + assert.Equal(t, ErrNoAvailableNonces, err) + + for i := 0; i < 10; i++ { + selectedNonce, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + require.NoError(t, err) + + nonceRecord, ok := expiredNonces[selectedNonce.Account.PublicKey().ToBase58()] + require.True(t, ok) + require.True(t, nonceRecord.ClaimExpiresAt.Before(time.Now())) + assert.Equal(t, nonceRecord.Address, selectedNonce.record.Address) + assert.Equal(t, nonceRecord.Address, selectedNonce.Account.PublicKey().ToBase58()) + assert.Equal(t, nonceRecord.Blockhash, base58.Encode(selectedNonce.Blockhash[:])) + delete(expiredNonces, selectedNonce.Account.PublicKey().ToBase58()) + + updatedRecord, err := env.data.GetNonce(env.ctx, selectedNonce.Account.PublicKey().ToBase58()) + require.NoError(t, err) + assert.Equal(t, nonce.StateReserved, updatedRecord.State) + assert.Empty(t, updatedRecord.Signature) + } + + _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + require.Equal(t, ErrNoAvailableNonces, err) } func TestNonce_MarkReservedWithSignature(t *testing.T) { @@ -132,7 +168,7 @@ func generateAvailableNonce(t *testing.T, env nonceTestEnv, nonceEnv nonce.Envir nonceKey := &vault.Record{ PublicKey: nonceAccount.PublicKey().ToBase58(), PrivateKey: nonceAccount.PrivateKey().ToBase58(), - State: vault.StateAvailable, + State: vault.StateReserved, CreatedAt: time.Now(), } nonceRecord := &nonce.Record{ @@ -149,6 +185,39 @@ func generateAvailableNonce(t *testing.T, env nonceTestEnv, nonceEnv nonce.Envir return nonceRecord } +func generateClaimedNonce(t *testing.T, env nonceTestEnv, nonceEnv nonce.Environment, instance string, expired bool) *nonce.Record { + nonceAccount := testutil.NewRandomAccount(t) + + var bh solana.Blockhash + rand.Read(bh[:]) + + nonceKey := &vault.Record{ + PublicKey: nonceAccount.PublicKey().ToBase58(), + PrivateKey: nonceAccount.PrivateKey().ToBase58(), + State: vault.StateReserved, + CreatedAt: time.Now(), + } + nonceRecord := &nonce.Record{ + Address: nonceAccount.PublicKey().ToBase58(), + Authority: common.GetSubsidizer().PublicKey().ToBase58(), + Blockhash: base58.Encode(bh[:]), + Environment: nonceEnv, + EnvironmentInstance: instance, + Purpose: nonce.PurposeClientTransaction, + State: nonce.StateClaimed, + ClaimNodeID: pointer.String("my_node_id"), + } + if expired { + nonceRecord.ClaimExpiresAt = pointer.Time(time.Now().Add(-time.Hour)) + } else { + nonceRecord.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) + } + + require.NoError(t, env.data.SaveKey(env.ctx, nonceKey)) + require.NoError(t, env.data.SaveNonce(env.ctx, nonceRecord)) + return nonceRecord +} + func generateAvailableNonces(t *testing.T, env nonceTestEnv, nonceEnv nonce.Environment, instance string, useCase nonce.Purpose, count int) []*nonce.Record { var nonces []*nonce.Record for i := 0; i < count; i++ { From e2347fff2316fe4ac1779297c9893f0f0d03036e Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 09:14:11 -0400 Subject: [PATCH 2/8] Initial LocalNoncePool implementation --- go.mod | 11 - go.sum | 24 -- pkg/code/data/nonce/memory/store.go | 8 +- pkg/code/data/nonce/nonce.go | 15 +- pkg/code/data/nonce/tests/tests.go | 3 +- pkg/code/transaction/nonce.go | 10 +- pkg/code/transaction/nonce_pool.go | 522 ++++++++++++++++++++++++ pkg/code/transaction/nonce_pool_test.go | 270 ++++++++++++ 8 files changed, 812 insertions(+), 51 deletions(-) create mode 100644 pkg/code/transaction/nonce_pool.go create mode 100644 pkg/code/transaction/nonce_pool_test.go diff --git a/go.mod b/go.mod index 1f7447f2..d51e168a 100644 --- a/go.mod +++ b/go.mod @@ -31,14 +31,9 @@ require ( github.com/stretchr/testify v1.8.4 github.com/vence722/base122-go v0.0.2 github.com/ybbus/jsonrpc v2.1.2+incompatible - go.etcd.io/etcd/api/v3 v3.5.13 - go.etcd.io/etcd/client/v3 v3.5.13 - go.uber.org/zap v1.17.0 golang.org/x/crypto v0.32.0 - golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 golang.org/x/net v0.34.0 golang.org/x/text v0.21.0 - golang.org/x/time v0.5.0 google.golang.org/grpc v1.71.1 google.golang.org/protobuf v1.36.6 ) @@ -50,8 +45,6 @@ require ( github.com/bits-and-blooms/bitset v1.2.0 // indirect github.com/cenkalti/backoff/v4 v4.1.0 // indirect github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 // indirect - github.com/coreos/go-semver v0.3.0 // indirect - github.com/coreos/go-systemd/v22 v22.3.2 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/docker/cli v20.10.7+incompatible // indirect github.com/docker/docker v20.10.7+incompatible // indirect @@ -90,11 +83,7 @@ require ( github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect github.com/xeipuuv/gojsonschema v1.2.0 // indirect - go.etcd.io/etcd/client/pkg/v3 v3.5.13 // indirect - go.uber.org/atomic v1.7.0 // indirect - go.uber.org/multierr v1.6.0 // indirect golang.org/x/sys v0.29.0 // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f // indirect gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect gopkg.in/ini.v1 v1.51.0 // indirect diff --git a/go.sum b/go.sum index f775229a..98d2f133 100644 --- a/go.sum +++ b/go.sum @@ -86,12 +86,9 @@ github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVn github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6/go.mod h1:GL3xCUCBDV3CZiTSEKksMWbLE66hEyuu9qyDOOqM47Y= github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk= github.com/coreos/etcd v3.3.13+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE= -github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM= github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk= github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= -github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI= -github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= github.com/coreos/pkg v0.0.0-20180928190104-399ea9e2e55f/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA= github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= @@ -141,7 +138,6 @@ github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs= github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= -github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/gofrs/uuid v4.0.0+incompatible h1:1SD/1F5pU8p29ybwgQSwpQk+mwdRrXCYuPhW6m+TnJw= github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ= @@ -467,12 +463,6 @@ github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9de github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/zenazn/goji v0.9.0/go.mod h1:7S9M489iMyHBNxwZnk9/EHS098H4/F6TATF2mIxtB1Q= go.etcd.io/bbolt v1.3.2/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU= -go.etcd.io/etcd/api/v3 v3.5.13 h1:8WXU2/NBge6AUF1K1gOexB6e07NgsN1hXK0rSTtgSp4= -go.etcd.io/etcd/api/v3 v3.5.13/go.mod h1:gBqlqkcMMZMVTMm4NDZloEVJzxQOQIls8splbqBDa0c= -go.etcd.io/etcd/client/pkg/v3 v3.5.13 h1:RVZSAnWWWiI5IrYAXjQorajncORbS0zI48LQlE2kQWg= -go.etcd.io/etcd/client/pkg/v3 v3.5.13/go.mod h1:XxHT4u1qU12E2+po+UVPrEeL94Um6zL58ppuJWXSAB8= -go.etcd.io/etcd/client/v3 v3.5.13 h1:o0fHTNJLeO0MyVbc7I3fsCf6nrOqn5d+diSarKnB2js= -go.etcd.io/etcd/client/v3 v3.5.13/go.mod h1:cqiAeY8b5DEEcpxvgWKsbLIWNM/8Wy2xJSDMtioMcoI= go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU= go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8= go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw= @@ -495,19 +485,13 @@ go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= go.uber.org/atomic v1.5.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= -go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw= -go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4= go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU= -go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4= -go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU= go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA= go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.13.0/go.mod h1:zwrFLgMcdUuIBviXEYEH1YKNaOBnKXsx2IPda5bBwHM= -go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U= -go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo= golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20181029021203-45a5f77698d3/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= @@ -536,8 +520,6 @@ golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u0 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM= golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU= golang.org/x/exp v0.0.0-20200331195152-e8c3332aa8e5/go.mod h1:4M0jN8W1tt0AVLNr8HDosyJCDCDuyL9N9+3m7wDWgKw= -golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM= -golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc= golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= @@ -685,8 +667,6 @@ golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= -golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= -golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= @@ -815,8 +795,6 @@ google.golang.org/genproto v0.0.0-20201210142538-e3217bee35cc/go.mod h1:FWY/as6D google.golang.org/genproto v0.0.0-20201214200347-8c77b98c765d/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no= google.golang.org/genproto v0.0.0-20210108203827-ffc7fda8c3d7/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no= google.golang.org/genproto v0.0.0-20210226172003-ab064af71705/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no= -google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422 h1:GVIKPyP/kLIyVOgOnTwFOrvQaQUzOzGMCxgFUOEmm24= -google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422/go.mod h1:b6h1vNKhxaSoEI+5jc3PJUCustfli/mRab7295pY7rw= google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f h1:OxYkA3wjPsZyBylwymxSHa7ViiW1Sml4ToBrncvFehI= google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f/go.mod h1:+2Yz8+CLJbIfL9z73EW45avw8Lmge3xVElCP9zEKi50= google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= @@ -863,12 +841,10 @@ gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bl gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= -gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gotest.tools/v3 v3.0.2 h1:kG1BFyqVHuQoVQiR1bWGnfz/fmHvvuiSPIV7rvl360E= diff --git a/pkg/code/data/nonce/memory/store.go b/pkg/code/data/nonce/memory/store.go index e657b658..6d6c090c 100644 --- a/pkg/code/data/nonce/memory/store.go +++ b/pkg/code/data/nonce/memory/store.go @@ -133,10 +133,10 @@ func (s *store) filter(items []*nonce.Record, cursor query.Cursor, limit uint64, return res } -func (s *store) filterAvailable(items []*nonce.Record) []*nonce.Record { +func (s *store) filterAvailableToClaim(items []*nonce.Record) []*nonce.Record { var res []*nonce.Record for _, item := range items { - if item.IsAvailable() { + if item.IsAvailableToClaim() { res = append(res, item) } } @@ -238,7 +238,7 @@ func (s *store) GetRandomAvailableByPurpose(ctx context.Context, env nonce.Envir items := s.findByStateAndPurpose(env, instance, nonce.StateAvailable, purpose) items = append(items, s.findByStateAndPurpose(env, instance, nonce.StateClaimed, purpose)...) - items = s.filterAvailable(items) + items = s.filterAvailableToClaim(items) if len(items) == 0 { return nil, nonce.ErrNonceNotFound } @@ -254,7 +254,7 @@ func (s *store) BatchClaimAvailableByPurpose(ctx context.Context, env nonce.Envi items := s.findByStateAndPurpose(env, instance, nonce.StateAvailable, purpose) items = append(items, s.findByStateAndPurpose(env, instance, nonce.StateClaimed, purpose)...) - items = s.filterAvailable(items) + items = s.filterAvailableToClaim(items) if len(items) == 0 { return nil, nonce.ErrNonceNotFound } diff --git a/pkg/code/data/nonce/nonce.go b/pkg/code/data/nonce/nonce.go index ba43b167..9b4484bb 100644 --- a/pkg/code/data/nonce/nonce.go +++ b/pkg/code/data/nonce/nonce.go @@ -149,18 +149,23 @@ func (r *Record) Validate() error { return nil } -func (r *Record) IsAvailable() bool { +func (r *Record) IsAvailableToClaim() bool { if r.State == StateAvailable { return true } if r.State != StateClaimed { return false } - if r.ClaimExpiresAt == nil { - return false - } + return r.ClaimExpiresAt.Before(time.Now()) +} - return time.Now().After(*r.ClaimExpiresAt) +func (r *Record) CanReserveWithSignature() bool { + if r.State == StateAvailable { + return true + } + // Allow a small buffer against expiration timestamp to account for DB + // call latency + return r.State == StateClaimed && r.ClaimExpiresAt.After(time.Now().Add(time.Second)) } func (r *Record) GetPublicKey() (ed25519.PublicKey, error) { diff --git a/pkg/code/data/nonce/tests/tests.go b/pkg/code/data/nonce/tests/tests.go index 15b7afc1..50c6224e 100644 --- a/pkg/code/data/nonce/tests/tests.go +++ b/pkg/code/data/nonce/tests/tests.go @@ -365,7 +365,8 @@ func testGetRandomAvailableByPurpose(t *testing.T, s nonce.Store) { assert.Equal(t, purpose, actual.Purpose) assert.Equal(t, nonce.EnvironmentSolana, actual.Environment) assert.Equal(t, nonce.EnvironmentInstanceSolanaMainnet, actual.EnvironmentInstance) - assert.True(t, actual.IsAvailable()) + assert.True(t, actual.IsAvailableToClaim()) + assert.True(t, actual.CanReserveWithSignature()) switch actual.State { case nonce.StateAvailable: availableState++ diff --git a/pkg/code/transaction/nonce.go b/pkg/code/transaction/nonce.go index fdc5fe56..f6c59a02 100644 --- a/pkg/code/transaction/nonce.go +++ b/pkg/code/transaction/nonce.go @@ -15,9 +15,7 @@ import ( "github.com/code-payments/code-server/pkg/solana" ) -var ( - ErrNoAvailableNonces = errors.New("no available nonces") -) +// todo: Deprecate this in favour of LocalNoncePool var ( // Temporary global lock, so we avoid any chance of double locking a nonce, @@ -83,7 +81,7 @@ func SelectAvailableNonce(ctx context.Context, data code_data.Provider, env nonc return err } - if !record.IsAvailable() { + if !record.IsAvailableToClaim() { // Unlock and try again lock.Unlock() return errors.New("selected nonce that became unavailable") @@ -154,7 +152,7 @@ func (n *SelectedNonce) MarkReservedWithSignature(ctx context.Context, sig strin return n.data.SaveNonce(ctx, n.record) } - if !n.record.IsAvailable() { + if !n.record.CanReserveWithSignature() { return errors.New("nonce must be available to reserve") } @@ -177,7 +175,7 @@ func (n *SelectedNonce) ReleaseIfNotReserved() error { return errors.New("nonce is unlocked") } - if n.record.IsAvailable() { + if n.record.IsAvailableToClaim() { return nil } diff --git a/pkg/code/transaction/nonce_pool.go b/pkg/code/transaction/nonce_pool.go new file mode 100644 index 00000000..6252336e --- /dev/null +++ b/pkg/code/transaction/nonce_pool.go @@ -0,0 +1,522 @@ +package transaction + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + "time" + + "github.com/google/uuid" + "github.com/newrelic/go-agent/v3/newrelic" + "github.com/sirupsen/logrus" + + code_data "github.com/code-payments/code-server/pkg/code/data" + "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/pointer" +) + +var ( + ErrNoAvailableNonces = errors.New("no available nonces") + ErrNoncePoolClosed = errors.New("nonce pool is closed") +) + +// NoncePoolOption configures a nonce pool. +type NoncePoolOption func(*noncePoolOpts) + +type noncePoolOpts struct { + desiredPoolSize int + + nodeID string + minExpiration time.Duration + maxExpiration time.Duration + + refreshInterval time.Duration + refreshPoolInterval time.Duration + + shutdownGracePeriod time.Duration +} + +func defaultOptions() noncePoolOpts { + return noncePoolOpts{ + desiredPoolSize: 100, + + nodeID: uuid.New().String(), + minExpiration: 3 * time.Minute, + maxExpiration: 5 * time.Minute, + + refreshInterval: 5 * time.Second, + refreshPoolInterval: 10 * time.Second, + + shutdownGracePeriod: time.Minute, + } +} + +// WithNoncePoolSize configures the desired size of the pool. +// +// The pool will use this size to determine how much to load +// when the pool starts up, or when the pool gets low. It can +// be viewed as 'target memory' in a GC, with the actual pool +// size behaving like a saw-tooth graph. +// +// The pool does not have any mechanism to shrink the pool to +// this size beyond the natural consumption of nonces. +func WithNoncePoolSize(size int) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.desiredPoolSize = size + } +} + +// WithNoncePoolNodeID configures the node id to use when claiming nonces. +func WithNoncePoolNodeID(id string) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.nodeID = id + } +} + +// WithNoncePoolMinExpiration configures the lower bound for the +// expiration window of claimed nonces. +func WithNoncePoolMinExpiration(d time.Duration) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.minExpiration = d + } +} + +// WithNoncePoolMaxExpiration configures the upper bound for the +// expiration window of claimed nonces. +func WithNoncePoolMaxExpiration(d time.Duration) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.maxExpiration = d + } +} + +// WithNoncePoolRefreshInterval specifies how often the pool should be +// scanning it's free list for refresh candidates. Candidates are claimed +// nonces whose expiration is <= 2/3 of the min expiration. +func WithNoncePoolRefreshInterval(interval time.Duration) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.refreshInterval = interval + } +} + +// WithNoncePoolRefreshPoolInterval configures the pool to refresh the set of +// nonces at this duration. If the pool size is 1/2 the desired size, nonces +// will be fetched from DB. This condition is also checked (asynchronously) every +// time a nonce is pulled from the pool. +func WithNoncePoolRefreshPoolInterval(interval time.Duration) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.refreshPoolInterval = interval + } +} + +// WithNoncePoolShutdownGracePeriod configures the amount of time given for +// cleanup on call to Close(). +func WithNoncePoolShutdownGracePeriod(duration time.Duration) NoncePoolOption { + return func(npo *noncePoolOpts) { + npo.shutdownGracePeriod = duration + } +} + +func (opts *noncePoolOpts) validate() error { + if opts.desiredPoolSize < 10 { + return errors.New("pool size must greater than 10") + } + + if opts.nodeID == "" { + return errors.New("missing node id") + } + if opts.minExpiration < 10*time.Second { + return errors.New("min expiry must be >= 10s") + } + if opts.maxExpiration < 10*time.Second { + return errors.New("max expiry must be >= 10s") + } + if opts.minExpiration > opts.maxExpiration { + return errors.New("min expiry must <= max expiry") + } + + if opts.refreshInterval < time.Second || opts.refreshInterval > opts.minExpiration/2 { + return fmt.Errorf("invalid refresh interval %v, must be between (1, minExpiration/2)", opts.refreshInterval) + } + + if opts.refreshPoolInterval < time.Second { + return fmt.Errorf("invalid refresh pool interval %v, must be greater than 1s", opts.refreshPoolInterval) + } + + if opts.shutdownGracePeriod < time.Second { + return fmt.Errorf("invalid shutdown grace period %v, must be greater than 1s", opts.shutdownGracePeriod) + } + + return nil +} + +// Nonce represents a handle to a nonce that is owned by a local nonce pool. +type Nonce struct { + pool *LocalNoncePool + record *nonce.Record +} + +// MarkReservedWithSignature marks the nonce as reserved with a signature +func (n *Nonce) MarkReservedWithSignature(ctx context.Context, sig string) error { + if len(sig) == 0 { + return errors.New("signature is empty") + } + + if n.record.Signature == sig { + return nil + } + + if n.record.State == nonce.StateReserved || len(n.record.Signature) != 0 { + return errors.New("nonce already reserved with a different signature") + } + + if !n.record.CanReserveWithSignature() { + return errors.New("nonce is not in a valid state to reserve with signature") + } + + n.record.State = nonce.StateReserved + n.record.Signature = sig + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil + return n.pool.data.SaveNonce(ctx, n.record) +} + +// ReleaseIfNotReserved releases the nonce back to the pool if +// the nonce has not yet been reserved (or more specifically, is +// still owned by the pool). +func (n *Nonce) ReleaseIfNotReserved() { + if n.record.State != nonce.StateClaimed { + return + } + if *n.record.ClaimNodeID != n.pool.opts.nodeID { + return + } + if n.record.ClaimExpiresAt.Before(time.Now()) { + return + } + + n.pool.mu.Lock() + n.pool.freeList = append(n.pool.freeList, n) + n.pool.mu.Unlock() +} + +// LocalNoncePool is a pool of nonces that are cached in memory for +// quick access. The LocalNoncePool will continually monitor the pool +// to ensure sufficient size, as well refresh nonce expiration +// times. +// +// If the pool empties before it can be refilled, ErrNoAvailableNonces +// will be returned. Therefore, the pool should be sufficiently large +// such that the consumption of poolSize/2 nonces is _slower_ than the +// operation to top up the pool. +type LocalNoncePool struct { + log *logrus.Entry + + data code_data.Provider + + metricsProvider *newrelic.Application + + env nonce.Environment + envInstance string + poolType nonce.Purpose + + opts noncePoolOpts + + workerCtx context.Context + cancelWorkerCtx context.CancelFunc + + mu sync.RWMutex + freeList []*Nonce + isClosed bool + + refreshPoolCh chan struct{} +} + +func NewLocalNoncePool( + data code_data.Provider, + metricsProvider *newrelic.Application, + env nonce.Environment, + envInstance string, + poolType nonce.Purpose, + opts ...NoncePoolOption, +) (*LocalNoncePool, error) { + np := &LocalNoncePool{ + log: logrus.StandardLogger().WithFields(logrus.Fields{ + "type": "transaction/LocalNoncePool", + "environment": env.String(), + "environment_instance": envInstance, + "pool_type": poolType.String(), + }), + + data: data, + + metricsProvider: metricsProvider, + + env: env, + envInstance: envInstance, + poolType: poolType, + + opts: defaultOptions(), + + refreshPoolCh: make(chan struct{}), + } + + np.workerCtx, np.cancelWorkerCtx = context.WithCancel(context.Background()) + + for _, o := range opts { + o(&np.opts) + } + + if err := np.opts.validate(); err != nil { + np.cancelWorkerCtx() + return nil, err + } + + _, err := np.load(np.workerCtx, np.opts.desiredPoolSize) + switch err { + case nil, nonce.ErrNonceNotFound: + default: + np.cancelWorkerCtx() + return nil, err + } + + go np.refreshPool() + go np.refreshNonces() + go np.metricsPoller() + + return np, nil +} + +func (np *LocalNoncePool) GetNonce(ctx context.Context) (*Nonce, error) { + var n *Nonce + + np.mu.Lock() + if np.isClosed { + np.mu.Unlock() + return nil, ErrNoncePoolClosed + } + + size := len(np.freeList) + if size > 0 { + n = np.freeList[0] + np.freeList = np.freeList[1:] + } + np.mu.Unlock() + + if size < np.opts.desiredPoolSize/2 { + select { + case np.refreshPoolCh <- struct{}{}: + default: + } + } + + if n == nil { + return nil, ErrNoAvailableNonces + } + return n, nil +} + +func (np *LocalNoncePool) Close() error { + log := np.log.WithField("method", "Close") + + np.mu.Lock() + defer np.mu.Unlock() + + np.isClosed = true + np.cancelWorkerCtx() + + ctx, cancel := context.WithTimeout(context.Background(), np.opts.shutdownGracePeriod) + defer cancel() + + remaining := len(np.freeList) + for _, n := range np.freeList { + n.record.State = nonce.StateAvailable + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil + + if err := np.data.SaveNonce(ctx, n.record); err != nil { + log.WithError(err).WithField("nonce", n.record.Address).Warn("Failed to release nonce on shutdown") + } else { + remaining-- + } + } + + if remaining != 0 { + return fmt.Errorf("failed to free all nonces (%d left unfreed)", remaining) + } + + return nil +} + +func (np *LocalNoncePool) load(ctx context.Context, limit int) (int, error) { + np.mu.Lock() + if np.isClosed { + np.mu.Unlock() + return 0, nil + } + np.mu.Unlock() + + now := time.Now() + records, err := np.data.BatchClaimAvailableNoncesByPurpose( + ctx, + np.env, + np.envInstance, + np.poolType, + limit, + np.opts.nodeID, + now.Add(np.opts.minExpiration), + now.Add(np.opts.maxExpiration), + ) + if err != nil { + return 0, err + } + + if len(records) == 0 { + return 0, ErrNoAvailableNonces + } + + np.mu.Lock() + for i := range records { + np.freeList = append(np.freeList, &Nonce{pool: np, record: records[i]}) + } + np.mu.Unlock() + + return len(records), nil +} + +func (np *LocalNoncePool) refreshPool() { + log := np.log.WithField("method", "refreshPool") + + for { + select { + case <-np.workerCtx.Done(): + return + case <-np.refreshPoolCh: + case <-time.After(np.opts.refreshPoolInterval): + } + + np.mu.Lock() + size := len(np.freeList) + np.mu.Unlock() + + if size >= np.opts.desiredPoolSize { + continue + } + + limit := np.opts.desiredPoolSize - size + log := log.WithField("limit", limit) + log.Debug("Refreshing nonce pool") + loaded, err := np.load(np.workerCtx, limit) + if err != nil { + log.WithError(err).Warn("Failed to refresh nonce pool") + } else if loaded < limit { + log.WithField("count", loaded).Warn("Unable to refresh nonce pool with desired number of nonces") + } else { + log.WithField("count", loaded).Debug("Nonce pool refreshed") + } + } +} + +func (np *LocalNoncePool) refreshNonces() { + for { + select { + case <-np.workerCtx.Done(): + return + case <-time.After(np.opts.refreshInterval): + } + + np.refreshNoncesNow() + } +} + +func (np *LocalNoncePool) refreshNoncesNow() { + log := np.log.WithField("method", "refreshNoncesNow") + + now := time.Now() + refreshList := make([]*Nonce, 0) + + np.mu.Lock() + if np.isClosed { + np.mu.Unlock() + return + } + + for i := 0; i < len(np.freeList); { + n := np.freeList[i] + if n.record.ClaimExpiresAt.Sub(now) > 2*np.opts.minExpiration/3 { + i++ + continue + } + + refreshList = append(refreshList, n) + np.freeList = slices.Delete(np.freeList, i, i+1) + } + np.mu.Unlock() + + if len(refreshList) == 0 { + return + } + + log.WithField("count", len(refreshList)).Debug("Refreshing nonces") + + for _, n := range refreshList { + log := log.WithField("nonce", n.record.Address) + log.Debug("Refreshing nonce") + + if time.Since(*n.record.ClaimExpiresAt) >= 0 { + log.Warn("Nonce claim is expired, abandoning") + continue + } + if time.Since(*n.record.ClaimExpiresAt) > -time.Second { + log.Warn("Nonce claim is too close to expiry, abandoning") + continue + } + + n.record.ClaimExpiresAt = pointer.Time(n.record.ClaimExpiresAt.Add(np.opts.minExpiration)) + err := np.data.SaveNonce(np.workerCtx, n.record) + if err != nil { + log.WithError(err).Warn("Failed to refresh nonce, abandoning") + } else { + np.mu.Lock() + np.freeList = append(np.freeList, n) + np.mu.Unlock() + } + } +} + +func (np *LocalNoncePool) metricsPoller() { + for { + select { + case <-np.workerCtx.Done(): + return + case <-time.After(time.Second): + np.recordPoolSizeMetricEvent() + } + } +} + +func (np *LocalNoncePool) recordPoolSizeMetricEvent() { + np.mu.Lock() + if np.isClosed { + np.mu.Unlock() + return + } + size := len(np.freeList) + np.mu.Unlock() + + kvs := np.getBaseMetricKvs() + kvs["current_nonce_pool_size"] = size + kvs["desired_nonce_pool_size"] = np.opts.desiredPoolSize + + np.metricsProvider.RecordCustomEvent("LocalNoncePoolSizePollingCheck", kvs) +} + +func (np *LocalNoncePool) getBaseMetricKvs() map[string]interface{} { + return map[string]interface{}{ + "node_id": np.opts.nodeID, + "nonce_env": np.env.String(), + "nonce_env_instance": np.envInstance, + "nonce_pool_type": np.poolType.String(), + } +} diff --git a/pkg/code/transaction/nonce_pool_test.go b/pkg/code/transaction/nonce_pool_test.go new file mode 100644 index 00000000..20dd60de --- /dev/null +++ b/pkg/code/transaction/nonce_pool_test.go @@ -0,0 +1,270 @@ +package transaction + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + + code_data "github.com/code-payments/code-server/pkg/code/data" + "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/testutil" +) + +func TestLocalNoncePool(t *testing.T) { + for _, tc := range []struct { + name string + tf func(*localNoncePoolTest) + }{ + {"HappyPath", testLocalNoncePoolHappyPath}, + {"Reload", testLocalNoncePoolReload}, + {"Refresh", testLocalNoncePoolRefresh}, + {"ReserveWithSignatureEdgeCases", testLocalNoncePoolReserveWithSignatureEdgeCases}, + } { + t.Run( + tc.name, + func(t *testing.T) { + nt := newLocalNoncePoolTest(t) + defer nt.pool.Close() + tc.tf(nt) + }, + ) + } +} + +func testLocalNoncePoolHappyPath(nt *localNoncePoolTest) { + start := time.Now() + + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeOnDemandTransaction) + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaDevnet, nonce.PurposeClientTransaction) + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentCvm, testutil.NewRandomAccount(nt.t).PublicKey().ToBase58(), nonce.PurposeClientTransaction) + + ctx := context.Background() + + // None of the refresh periods should have kicked in, so our pool is empty. + n, err := nt.pool.GetNonce(ctx) + require.ErrorIs(nt.t, err, ErrNoAvailableNonces) + require.Nil(nt.t, n) + + // Manually trigger a load + _, err = nt.pool.load(ctx, nt.pool.opts.desiredPoolSize) + require.NoError(nt.t, err) + + // Nonces should now be claimed and available in the pool + observed := map[string]*Nonce{} + for i := 0; i < nt.pool.opts.desiredPoolSize; i++ { + n, err = nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + require.NotContains(nt.t, observed, n.record.Address) + observed[n.record.Address] = n + + actual, err := nt.data.GetNonce(ctx, n.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.Environment, nonce.EnvironmentSolana) + require.Equal(nt.t, actual.EnvironmentInstance, nonce.EnvironmentInstanceSolanaMainnet) + require.Equal(nt.t, actual.Purpose, nonce.PurposeClientTransaction) + require.Equal(nt.t, actual.State, nonce.StateClaimed) + require.Equal(nt.t, *actual.ClaimNodeID, nt.pool.opts.nodeID) + require.True(nt.t, actual.ClaimExpiresAt.After(start.Add(nt.pool.opts.minExpiration))) + require.True(nt.t, actual.ClaimExpiresAt.Before(start.Add(nt.pool.opts.maxExpiration).Add(100*time.Millisecond))) + } + + // Underlying local pool is empty + n, err = nt.pool.GetNonce(ctx) + require.ErrorIs(nt.t, err, ErrNoAvailableNonces) + require.Nil(nt.t, n) + + // Reserve some nonces with a signature + reserved := map[string]*Nonce{} + for k, v := range observed { + reserved[k] = v + signature := fmt.Sprintf("signature%d", len(reserved)) + require.NoError(nt.t, v.MarkReservedWithSignature(ctx, signature)) + + actual, err := nt.data.GetNonce(ctx, v.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateReserved) + require.Equal(nt.t, actual.Signature, signature) + require.Nil(nt.t, actual.ClaimNodeID) + require.Nil(nt.t, actual.ClaimExpiresAt) + + if len(reserved) > len(observed)/2 { + break + } + } + + // Releasing back to the pool should allow us to re-use the nonces that + // weren't reserved + for _, v := range observed { + v.ReleaseIfNotReserved() + } + clear(observed) + + for i := 0; i < nt.pool.opts.desiredPoolSize-len(reserved); i++ { + n, err = nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + require.NotContains(nt.t, observed, n.record.Address) + require.NotContains(nt.t, reserved, n.record.Address) + observed[n.record.Address] = n + } + + // Underlying local pool is empty + n, err = nt.pool.GetNonce(ctx) + require.ErrorIs(nt.t, err, ErrNoAvailableNonces) + require.Nil(nt.t, n) + + // Releasing back to the pool will allow us to make nonces available to + // claim + for _, v := range observed { + v.ReleaseIfNotReserved() + } + + require.NoError(nt.t, nt.pool.Close()) + n, err = nt.pool.GetNonce(ctx) + require.ErrorIs(nt.t, err, ErrNoncePoolClosed) + require.Nil(nt.t, n) + + // Reserved nonces should stay reserved + for _, v := range reserved { + actual, err := nt.data.GetNonce(ctx, v.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateReserved) + } + // Claimed nonces should become available to claim by another pool + for _, v := range observed { + actual, err := nt.data.GetNonce(ctx, v.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateAvailable) + } +} + +func testLocalNoncePoolReload(nt *localNoncePoolTest) { + ctx := context.Background() + + nt.initializeNonces(10*nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + + _, err := nt.pool.load(ctx, nt.pool.opts.desiredPoolSize) + require.NoError(nt.t, err) + + observed := map[string]*Nonce{} + for i := 0; i < 10*nt.pool.opts.desiredPoolSize; i++ { + n, err := nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + require.NotContains(nt.t, observed, n.record.Address) + observed[n.record.Address] = n + + actual, err := nt.data.GetNonce(ctx, n.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateClaimed) + require.Equal(nt.t, *actual.ClaimNodeID, nt.pool.opts.nodeID) + require.True(nt.t, actual.ClaimExpiresAt.After(time.Now())) + + // Allow nonce pool to reload + if i%nt.pool.opts.desiredPoolSize == 0 { + time.Sleep(10 * time.Millisecond) + } + } +} + +func testLocalNoncePoolRefresh(nt *localNoncePoolTest) { + ctx := context.Background() + + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + + _, err := nt.pool.load(ctx, nt.pool.opts.desiredPoolSize) + require.NoError(nt.t, err) + + // Wait out the max expiration, so every nonce's claim needs to be refreshed + // at least once + // + // todo: Improve test running time + time.Sleep(2 * nt.pool.opts.maxExpiration) + + // All nonces should still be claimed and available for use + observed := map[string]*Nonce{} + for i := 0; i < nt.pool.opts.desiredPoolSize; i++ { + n, err := nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + require.NotContains(nt.t, observed, n.record.Address) + observed[n.record.Address] = n + + actual, err := nt.data.GetNonce(ctx, n.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateClaimed) + require.Equal(nt.t, *actual.ClaimNodeID, nt.pool.opts.nodeID) + require.True(nt.t, actual.ClaimExpiresAt.After(time.Now())) + } +} + +func testLocalNoncePoolReserveWithSignatureEdgeCases(nt *localNoncePoolTest) { + ctx := context.Background() + + nt.initializeNonces(nt.pool.opts.desiredPoolSize, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) + + _, err := nt.pool.load(ctx, nt.pool.opts.desiredPoolSize) + require.NoError(nt.t, err) + + n, err := nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + + require.Error(nt.t, n.MarkReservedWithSignature(ctx, "")) + require.NoError(nt.t, n.MarkReservedWithSignature(ctx, "signature1")) + require.NoError(nt.t, n.MarkReservedWithSignature(ctx, "signature1")) + require.Error(nt.t, n.MarkReservedWithSignature(ctx, "signature2")) + + actual, err := nt.data.GetNonce(ctx, n.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateReserved) + require.Equal(nt.t, actual.Signature, "signature1") +} + +type localNoncePoolTest struct { + t *testing.T + pool *LocalNoncePool + data code_data.DatabaseData +} + +func newLocalNoncePoolTest(t *testing.T) *localNoncePoolTest { + data := code_data.NewTestDataProvider() + + pool, err := NewLocalNoncePool( + data, + nil, + nonce.EnvironmentSolana, + nonce.EnvironmentInstanceSolanaMainnet, + nonce.PurposeClientTransaction, + WithNoncePoolRefreshInterval(time.Second), + WithNoncePoolRefreshPoolInterval(2*time.Second), + WithNoncePoolMinExpiration(10*time.Second), + WithNoncePoolMaxExpiration(15*time.Second), + ) + require.NoError(t, err) + + return &localNoncePoolTest{ + t: t, + data: data, + pool: pool, + } +} + +func (np *localNoncePoolTest) initializeNonces(amount int, env nonce.Environment, envInstance string, purpose nonce.Purpose) { + for i := 0; i < amount; i++ { + err := np.data.SaveNonce(context.Background(), &nonce.Record{ + Address: fmt.Sprintf("addr-%s-%s-%s-%d", env.String(), envInstance, purpose.String(), i), + Authority: "my authority!", + Environment: env, + EnvironmentInstance: envInstance, + Purpose: purpose, + State: nonce.StateAvailable, + }) + require.NoError(np.t, err) + } +} From a7ced3bf4af6828db64779a075198f05c2b241c8 Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 14:10:10 -0400 Subject: [PATCH 3/8] Remove deprecated random nonce selection logic --- pkg/code/data/internal.go | 4 - pkg/code/data/nonce/memory/store.go | 16 -- pkg/code/data/nonce/postgres/model.go | 47 ------ pkg/code/data/nonce/postgres/store.go | 8 - pkg/code/data/nonce/store.go | 8 - pkg/code/data/nonce/tests/tests.go | 100 ------------ pkg/code/transaction/nonce.go | 221 ------------------------- pkg/code/transaction/nonce_test.go | 227 -------------------------- 8 files changed, 631 deletions(-) delete mode 100644 pkg/code/transaction/nonce.go delete mode 100644 pkg/code/transaction/nonce_test.go diff --git a/pkg/code/data/internal.go b/pkg/code/data/internal.go index e8e10677..6b825854 100644 --- a/pkg/code/data/internal.go +++ b/pkg/code/data/internal.go @@ -135,7 +135,6 @@ type DatabaseData interface { GetNonceCountByState(ctx context.Context, env nonce.Environment, instance string, state nonce.State) (uint64, error) GetNonceCountByStateAndPurpose(ctx context.Context, env nonce.Environment, instance string, state nonce.State, purpose nonce.Purpose) (uint64, error) GetAllNonceByState(ctx context.Context, env nonce.Environment, instance string, state nonce.State, opts ...query.Option) ([]*nonce.Record, error) - GetRandomAvailableNonceByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) BatchClaimAvailableNoncesByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) SaveNonce(ctx context.Context, record *nonce.Record) error @@ -530,9 +529,6 @@ func (dp *DatabaseProvider) GetAllNonceByState(ctx context.Context, env nonce.En return dp.nonces.GetAllByState(ctx, env, instance, state, req.Cursor, req.Limit, req.SortBy) } -func (dp *DatabaseProvider) GetRandomAvailableNonceByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) { - return dp.nonces.GetRandomAvailableByPurpose(ctx, env, instance, purpose) -} func (dp *DatabaseProvider) BatchClaimAvailableNoncesByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { return dp.nonces.BatchClaimAvailableByPurpose(ctx, env, instance, purpose, limit, nodeID, minExpireAt, maxExpireAt) } diff --git a/pkg/code/data/nonce/memory/store.go b/pkg/code/data/nonce/memory/store.go index 6d6c090c..0a65bd3f 100644 --- a/pkg/code/data/nonce/memory/store.go +++ b/pkg/code/data/nonce/memory/store.go @@ -232,22 +232,6 @@ func (s *store) GetAllByState(ctx context.Context, env nonce.Environment, instan return nil, nonce.ErrNonceNotFound } -func (s *store) GetRandomAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) { - s.mu.Lock() - defer s.mu.Unlock() - - items := s.findByStateAndPurpose(env, instance, nonce.StateAvailable, purpose) - items = append(items, s.findByStateAndPurpose(env, instance, nonce.StateClaimed, purpose)...) - items = s.filterAvailableToClaim(items) - if len(items) == 0 { - return nil, nonce.ErrNonceNotFound - } - - index := rand.Intn(len(items)) - cloned := items[index].Clone() - return &cloned, nil -} - func (s *store) BatchClaimAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { s.mu.Lock() defer s.mu.Unlock() diff --git a/pkg/code/data/nonce/postgres/model.go b/pkg/code/data/nonce/postgres/model.go index a97b5be8..4fb2fdde 100644 --- a/pkg/code/data/nonce/postgres/model.go +++ b/pkg/code/data/nonce/postgres/model.go @@ -201,53 +201,6 @@ func dbGetAllByState(ctx context.Context, db *sqlx.DB, env nonce.Environment, in return res, nil } -// todo: Implementation still isn't perfect, but better than no randomness. It's -// sufficiently efficient, as long as our nonce pool is larger than the max offset. -// todo: We may need to tune the offset based on pool size and environment, but it -// should be sufficiently good enough for now. -func dbGetRandomAvailableByPurpose(ctx context.Context, db *sqlx.DB, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonceModel, error) { - res := &nonceModel{} - - // Signature null check is required because some legacy records didn't have this - // set and causes this call to fail. This is a result of the field not being - // defined at the time of record creation. - // - // todo: Fix said nonce records - query := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version - FROM ` + nonceTableName + ` - WHERE environment = $1 AND environment_instance = $2 AND ((state = $3) OR (state = $4 AND claim_expires_at < $5)) AND purpose = $6 AND signature IS NOT NULL - OFFSET FLOOR(RANDOM() * 100) - LIMIT 1 - ` - fallbackQuery := `SELECT - id, address, authority, blockhash, environment, environment_instance, purpose, state, signature, claim_node_id, claim_expires_at, version - FROM ` + nonceTableName + ` - WHERE environment = $1 AND environment_instance = $2 AND ((state = $3) OR (state = $4 AND claim_expires_at < $5)) AND purpose = $6 AND signature IS NOT NULL - LIMIT 1 - ` - - nowMs := time.Now().UnixMilli() - err := db.GetContext(ctx, res, query, env, instance, nonce.StateAvailable, nonce.StateClaimed, nowMs, purpose) - if err != nil { - err = pgutil.CheckNoRows(err, nonce.ErrNonceNotFound) - - // No nonces found. Because our query isn't perfect, fall back to a - // strategy that will guarantee to select something if an available - // nonce exists. - if err == nonce.ErrNonceNotFound { - err := db.GetContext(ctx, res, fallbackQuery, env, instance, nonce.StateAvailable, nonce.StateClaimed, nowMs, purpose) - if err != nil { - return nil, pgutil.CheckNoRows(err, nonce.ErrNonceNotFound) - } - return res, nil - } - - return nil, err - } - return res, nil -} - func dbBatchClaimAvailableByPurpose( ctx context.Context, db *sqlx.DB, diff --git a/pkg/code/data/nonce/postgres/store.go b/pkg/code/data/nonce/postgres/store.go index f523c05b..b889834b 100644 --- a/pkg/code/data/nonce/postgres/store.go +++ b/pkg/code/data/nonce/postgres/store.go @@ -72,14 +72,6 @@ func (s *store) GetAllByState(ctx context.Context, env nonce.Environment, instan return nonces, nil } -func (s *store) GetRandomAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose) (*nonce.Record, error) { - model, err := dbGetRandomAvailableByPurpose(ctx, s.db, env, instance, purpose) - if err != nil { - return nil, err - } - return fromNonceModel(model), nil -} - func (s *store) BatchClaimAvailableByPurpose(ctx context.Context, env nonce.Environment, instance string, purpose nonce.Purpose, limit int, nodeID string, minExpireAt, maxExpireAt time.Time) ([]*nonce.Record, error) { models, err := dbBatchClaimAvailableByPurpose(ctx, s.db, env, instance, purpose, limit, nodeID, minExpireAt, maxExpireAt) if err != nil { diff --git a/pkg/code/data/nonce/store.go b/pkg/code/data/nonce/store.go index 35cb23d2..7288b2c8 100644 --- a/pkg/code/data/nonce/store.go +++ b/pkg/code/data/nonce/store.go @@ -33,14 +33,6 @@ type Store interface { // Returns ErrNotFound if no records are found. GetAllByState(ctx context.Context, env Environment, instance string, state State, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*Record, error) - // GetRandomAvailableByPurpose gets a random available nonce for a purpose within - // an environment instance. - // - // Returns ErrNotFound if no records are found. - // - // Deprecated in favour of BatchClaimAvailableByPurpose - GetRandomAvailableByPurpose(ctx context.Context, env Environment, instance string, purpose Purpose) (*Record, error) - // BatchClaimAvailableByPurpose batch claims up to the specified limit. // // The returned nonces will be marked as claimed by the current node, with diff --git a/pkg/code/data/nonce/tests/tests.go b/pkg/code/data/nonce/tests/tests.go index 50c6224e..07f97425 100644 --- a/pkg/code/data/nonce/tests/tests.go +++ b/pkg/code/data/nonce/tests/tests.go @@ -23,7 +23,6 @@ func RunTests(t *testing.T, s nonce.Store, teardown func()) { testUpdate, testGetAllByState, testGetCount, - testGetRandomAvailableByPurpose, testBatchClaimAvailableByPurpose, testBatchClaimAvailableByPurposeExpirationRandomness, } { @@ -285,105 +284,6 @@ func testGetCount(t *testing.T, s nonce.Store) { }) } -func testGetRandomAvailableByPurpose(t *testing.T, s nonce.Store) { - t.Run("testGetRandomAvailableByPurpose", func(t *testing.T) { - ctx := context.Background() - - _, err := s.GetRandomAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - assert.Equal(t, nonce.ErrNonceNotFound, err) - - for _, purpose := range []nonce.Purpose{ - nonce.PurposeClientTransaction, - nonce.PurposeInternalServerProcess, - } { - for _, state := range []nonce.State{ - nonce.StateUnknown, - nonce.StateAvailable, - nonce.StateReserved, - nonce.StateReleased, - nonce.StateClaimed, - } { - for i := 0; i < 50; i++ { - record := &nonce.Record{ - Address: fmt.Sprintf("nonce_%s_%s_%d", purpose, state, i), - Authority: "authority", - Blockhash: "bh", - Environment: nonce.EnvironmentSolana, - EnvironmentInstance: nonce.EnvironmentInstanceSolanaMainnet, - Purpose: purpose, - State: state, - Signature: "", - } - if state == nonce.StateClaimed { - record.ClaimNodeID = pointer.String("node_id") - - if i < 25 { - record.ClaimExpiresAt = pointer.Time(time.Now().Add(-time.Hour)) - } else { - record.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) - } - } - require.NoError(t, s.Save(ctx, record)) - } - } - } - - for i := 0; i < 100; i++ { - record := &nonce.Record{ - Address: fmt.Sprintf("nonce_devnet_%d", i), - Authority: "authority", - Blockhash: "bh", - Environment: nonce.EnvironmentSolana, - EnvironmentInstance: nonce.EnvironmentInstanceSolanaDevnet, - Purpose: nonce.PurposeInternalServerProcess, - State: nonce.StateAvailable, - Signature: "", - } - require.NoError(t, s.Save(ctx, record)) - - record = &nonce.Record{ - Address: fmt.Sprintf("nonce_cvm_%d", i), - Authority: "authority", - Blockhash: "bh", - Environment: nonce.EnvironmentCvm, - EnvironmentInstance: "pubkey", - Purpose: nonce.PurposeClientTransaction, - State: nonce.StateClaimed, - Signature: "", - ClaimNodeID: pointer.String("node_id"), - ClaimExpiresAt: pointer.Time(time.Now().Add(-time.Hour)), - } - require.NoError(t, s.Save(ctx, record)) - } - - for _, purpose := range []nonce.Purpose{nonce.PurposeClientTransaction, nonce.PurposeInternalServerProcess} { - availableState, claimedState := 0, 0 - selectedByAddress := make(map[string]struct{}) - for i := 0; i < 100; i++ { - actual, err := s.GetRandomAvailableByPurpose(ctx, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, purpose) - require.NoError(t, err) - assert.Equal(t, purpose, actual.Purpose) - assert.Equal(t, nonce.EnvironmentSolana, actual.Environment) - assert.Equal(t, nonce.EnvironmentInstanceSolanaMainnet, actual.EnvironmentInstance) - assert.True(t, actual.IsAvailableToClaim()) - assert.True(t, actual.CanReserveWithSignature()) - switch actual.State { - case nonce.StateAvailable: - availableState++ - case nonce.StateClaimed: - claimedState++ - assert.True(t, time.Now().After(*actual.ClaimExpiresAt)) - default: - } - selectedByAddress[actual.Address] = struct{}{} - } - assert.True(t, len(selectedByAddress) > 10) - assert.NotZero(t, availableState) - assert.NotZero(t, claimedState) - } - }) -} - func testBatchClaimAvailableByPurpose(t *testing.T, s nonce.Store) { t.Run("testBatchClaimAvailableByPurpose", func(t *testing.T) { ctx := context.Background() diff --git a/pkg/code/transaction/nonce.go b/pkg/code/transaction/nonce.go deleted file mode 100644 index f6c59a02..00000000 --- a/pkg/code/transaction/nonce.go +++ /dev/null @@ -1,221 +0,0 @@ -package transaction - -import ( - "context" - "errors" - "sync" - "time" - - "github.com/mr-tron/base58" - - "github.com/code-payments/code-server/pkg/code/common" - code_data "github.com/code-payments/code-server/pkg/code/data" - "github.com/code-payments/code-server/pkg/code/data/nonce" - "github.com/code-payments/code-server/pkg/retry" - "github.com/code-payments/code-server/pkg/solana" -) - -// todo: Deprecate this in favour of LocalNoncePool - -var ( - // Temporary global lock, so we avoid any chance of double locking a nonce, - // since we can't check the status of a sync.Mutx. - globalNonceLock sync.Mutex - - // Can't used striped lock because we need to hold mutliple nonces at once, so - // deadlock would be possible. This is fine for now, given our nonce pool has - // a fixed and relatively small size. - nonceLocksMu sync.Mutex - nonceLocks map[string]*sync.Mutex -) - -func init() { - nonceLocks = make(map[string]*sync.Mutex) -} - -// SelectedNonce is a nonce that is available and selected for use in a transaction. -// Implementations should unlock the lock after using the nonce. If used, its state -// must be updated as reserved. -type SelectedNonce struct { - localLock sync.Mutex - distributedLock *sync.Mutex // todo: Use a distributed lock - isUnlocked bool - - data code_data.Provider - - record *nonce.Record - Account *common.Account - Blockhash solana.Blockhash -} - -// SelectAvailableNonce selects an available from the nonce pool within an environment -// for the specified use case. The returned nonce is marked as reserved without a signature, -// so it cannot be selected again. It's the responsibility of the external caller to make -// it available again if it doesn't get assigned a fulfillment. -func SelectAvailableNonce(ctx context.Context, data code_data.Provider, env nonce.Environment, instance string, useCase nonce.Purpose) (*SelectedNonce, error) { - var lock *sync.Mutex - var account *common.Account - var bh solana.Blockhash - var record *nonce.Record - - _, err := retry.Retry(func() error { - globalNonceLock.Lock() - defer globalNonceLock.Unlock() - - randomRecord, err := data.GetRandomAvailableNonceByPurpose(ctx, env, instance, useCase) - if err == nonce.ErrNonceNotFound { - return ErrNoAvailableNonces - } else if err != nil { - return err - } - - record = randomRecord - - lock = getNonceLock(record.Address) - lock.Lock() - - // Refetch because the state could have changed by the time we got the lock - record, err = data.GetNonce(ctx, record.Address) - if err != nil { - lock.Unlock() - return err - } - - if !record.IsAvailableToClaim() { - // Unlock and try again - lock.Unlock() - return errors.New("selected nonce that became unavailable") - } - - account, err = common.NewAccountFromPublicKeyString(record.Address) - if err != nil { - lock.Unlock() - return err - } - - untypedBlockhash, err := base58.Decode(record.Blockhash) - if err != nil { - lock.Unlock() - return err - } - copy(bh[:], untypedBlockhash) - - // Reserve the nonce for use with a fulfillment - record.State = nonce.StateReserved - record.ClaimNodeID = nil - record.ClaimExpiresAt = nil - err = data.SaveNonce(ctx, record) - if err != nil { - lock.Unlock() - return err - } - - return nil - }, retry.NonRetriableErrors(context.Canceled, ErrNoAvailableNonces), retry.Limit(5)) - if err != nil { - return nil, err - } - - return &SelectedNonce{ - distributedLock: lock, - data: data, - record: record, - Account: account, - Blockhash: bh, - }, nil -} - -// MarkReservedWithSignature marks the nonce as reserved with a signature -func (n *SelectedNonce) MarkReservedWithSignature(ctx context.Context, sig string) error { - if len(sig) == 0 { - return errors.New("signature is empty") - } - - n.localLock.Lock() - defer n.localLock.Unlock() - - if n.isUnlocked { - return errors.New("nonce is unlocked") - } - - if n.record.Signature == sig { - return nil - } - - if len(n.record.Signature) != 0 { - return errors.New("nonce already has a different signature") - } - - // Nonce is reserved without a signature, so update its signature - if n.record.State == nonce.StateReserved { - n.record.Signature = sig - return n.data.SaveNonce(ctx, n.record) - } - - if !n.record.CanReserveWithSignature() { - return errors.New("nonce must be available to reserve") - } - - n.record.State = nonce.StateReserved - n.record.Signature = sig - n.record.ClaimNodeID = nil - n.record.ClaimExpiresAt = nil - return n.data.SaveNonce(ctx, n.record) -} - -// ReleaseIfNotReserved makes a nonce available if it hasn't been reserved with -// a signature. It's recommended to call this in tandem with Unlock when the -// caller knows it's safe to go from the reserved to available state (ie. don't -// use this in uprade flows!). -func (n *SelectedNonce) ReleaseIfNotReserved() error { - n.localLock.Lock() - defer n.localLock.Unlock() - - if n.isUnlocked { - return errors.New("nonce is unlocked") - } - - if n.record.IsAvailableToClaim() { - return nil - } - - // A nonce is not fully reserved if it's state is reserved, but there is no - // assigned signature. - if n.record.State == nonce.StateReserved && len(n.record.Signature) == 0 { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - n.record.State = nonce.StateAvailable - n.record.ClaimNodeID = nil - n.record.ClaimExpiresAt = nil - return n.data.SaveNonce(ctx, n.record) - } - - return nil -} - -func (n *SelectedNonce) Unlock() { - n.localLock.Lock() - defer n.localLock.Unlock() - - if n.isUnlocked { - return - } - - n.isUnlocked = true - - n.distributedLock.Unlock() -} - -func getNonceLock(address string) *sync.Mutex { - nonceLocksMu.Lock() - defer nonceLocksMu.Unlock() - - lock, ok := nonceLocks[address] - if !ok { - var mu sync.Mutex - lock = &mu - nonceLocks[address] = &mu - } - return lock -} diff --git a/pkg/code/transaction/nonce_test.go b/pkg/code/transaction/nonce_test.go deleted file mode 100644 index 8b487ec9..00000000 --- a/pkg/code/transaction/nonce_test.go +++ /dev/null @@ -1,227 +0,0 @@ -package transaction - -import ( - "context" - "crypto/rand" - "testing" - "time" - - "github.com/mr-tron/base58" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/code-payments/code-server/pkg/code/common" - code_data "github.com/code-payments/code-server/pkg/code/data" - "github.com/code-payments/code-server/pkg/code/data/nonce" - "github.com/code-payments/code-server/pkg/code/data/vault" - "github.com/code-payments/code-server/pkg/pointer" - "github.com/code-payments/code-server/pkg/solana" - "github.com/code-payments/code-server/pkg/testutil" -) - -func TestNonce_SelectAvailableNonce(t *testing.T) { - env := setupNonceTestEnv(t) - - allNonces := generateAvailableNonces(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction, 10) - noncesByAddress := make(map[string]*nonce.Record) - for _, nonceRecord := range allNonces { - noncesByAddress[nonceRecord.Address] = nonceRecord - } - - _, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeInternalServerProcess) - assert.Equal(t, ErrNoAvailableNonces, err) - - _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentCvm, testutil.NewRandomAccount(t).PublicKey().ToBase58(), nonce.PurposeClientTransaction) - assert.Equal(t, ErrNoAvailableNonces, err) - - selectedNonces := make(map[string]struct{}) - for i := 0; i < len(noncesByAddress); i++ { - selectedNonce, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.NoError(t, err) - - _, ok := selectedNonces[selectedNonce.Account.PublicKey().ToBase58()] - assert.False(t, ok) - - nonceRecord, ok := noncesByAddress[selectedNonce.Account.PublicKey().ToBase58()] - require.True(t, ok) - assert.Equal(t, nonceRecord.Address, selectedNonce.record.Address) - assert.Equal(t, nonceRecord.Address, selectedNonce.Account.PublicKey().ToBase58()) - assert.Equal(t, nonceRecord.Blockhash, base58.Encode(selectedNonce.Blockhash[:])) - - selectedNonces[selectedNonce.record.Address] = struct{}{} - - updatedRecord, err := env.data.GetNonce(env.ctx, nonceRecord.Address) - require.NoError(t, err) - assert.Equal(t, nonce.StateReserved, updatedRecord.State) - assert.Empty(t, updatedRecord.Signature) - } - assert.Len(t, selectedNonces, len(noncesByAddress)) - - _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - assert.Equal(t, ErrNoAvailableNonces, err) -} - -func TestNonce_SelectAvailableNonceClaimed(t *testing.T) { - env := setupNonceTestEnv(t) - - expiredNonces := map[string]*nonce.Record{} - for i := 0; i < 10; i++ { - n := generateClaimedNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, true) - expiredNonces[n.Address] = n - - generateClaimedNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, false) - } - - _, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeInternalServerProcess) - assert.Equal(t, ErrNoAvailableNonces, err) - - _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentCvm, testutil.NewRandomAccount(t).PublicKey().ToBase58(), nonce.PurposeClientTransaction) - assert.Equal(t, ErrNoAvailableNonces, err) - - for i := 0; i < 10; i++ { - selectedNonce, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.NoError(t, err) - - nonceRecord, ok := expiredNonces[selectedNonce.Account.PublicKey().ToBase58()] - require.True(t, ok) - require.True(t, nonceRecord.ClaimExpiresAt.Before(time.Now())) - assert.Equal(t, nonceRecord.Address, selectedNonce.record.Address) - assert.Equal(t, nonceRecord.Address, selectedNonce.Account.PublicKey().ToBase58()) - assert.Equal(t, nonceRecord.Blockhash, base58.Encode(selectedNonce.Blockhash[:])) - delete(expiredNonces, selectedNonce.Account.PublicKey().ToBase58()) - - updatedRecord, err := env.data.GetNonce(env.ctx, selectedNonce.Account.PublicKey().ToBase58()) - require.NoError(t, err) - assert.Equal(t, nonce.StateReserved, updatedRecord.State) - assert.Empty(t, updatedRecord.Signature) - } - - _, err = SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.Equal(t, ErrNoAvailableNonces, err) -} - -func TestNonce_MarkReservedWithSignature(t *testing.T) { - env := setupNonceTestEnv(t) - - generateAvailableNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - - selectedNonce, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.NoError(t, err) - - assert.Error(t, selectedNonce.MarkReservedWithSignature(env.ctx, "")) - require.NoError(t, selectedNonce.MarkReservedWithSignature(env.ctx, "signature1")) - require.NoError(t, selectedNonce.MarkReservedWithSignature(env.ctx, "signature1")) - assert.Error(t, selectedNonce.MarkReservedWithSignature(env.ctx, "signature2")) - - updatedRecord, err := env.data.GetNonce(env.ctx, selectedNonce.Account.PublicKey().ToBase58()) - require.NoError(t, err) - assert.Equal(t, nonce.StateReserved, updatedRecord.State) - assert.Equal(t, "signature1", updatedRecord.Signature) -} - -func TestNonce_ReleaseIfNotReserved(t *testing.T) { - env := setupNonceTestEnv(t) - - generateAvailableNonce(t, env, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - - selectedNonce, err := SelectAvailableNonce(env.ctx, env.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeClientTransaction) - require.NoError(t, err) - - require.NoError(t, selectedNonce.ReleaseIfNotReserved()) - - updatedRecord, err := env.data.GetNonce(env.ctx, selectedNonce.Account.PublicKey().ToBase58()) - require.NoError(t, err) - assert.Equal(t, nonce.StateAvailable, updatedRecord.State) - assert.Empty(t, updatedRecord.Signature) - - require.NoError(t, selectedNonce.MarkReservedWithSignature(env.ctx, "signature")) - require.NoError(t, selectedNonce.ReleaseIfNotReserved()) - - updatedRecord, err = env.data.GetNonce(env.ctx, selectedNonce.Account.PublicKey().ToBase58()) - require.NoError(t, err) - assert.Equal(t, nonce.StateReserved, updatedRecord.State) - assert.Equal(t, "signature", updatedRecord.Signature) -} - -type nonceTestEnv struct { - ctx context.Context - data code_data.Provider -} - -func setupNonceTestEnv(t *testing.T) nonceTestEnv { - data := code_data.NewTestDataProvider() - - testutil.SetupRandomSubsidizer(t, data) - - return nonceTestEnv{ - ctx: context.Background(), - data: data, - } -} - -func generateAvailableNonce(t *testing.T, env nonceTestEnv, nonceEnv nonce.Environment, instance string, useCase nonce.Purpose) *nonce.Record { - nonceAccount := testutil.NewRandomAccount(t) - - var bh solana.Blockhash - rand.Read(bh[:]) - - nonceKey := &vault.Record{ - PublicKey: nonceAccount.PublicKey().ToBase58(), - PrivateKey: nonceAccount.PrivateKey().ToBase58(), - State: vault.StateReserved, - CreatedAt: time.Now(), - } - nonceRecord := &nonce.Record{ - Address: nonceAccount.PublicKey().ToBase58(), - Authority: common.GetSubsidizer().PublicKey().ToBase58(), - Blockhash: base58.Encode(bh[:]), - Environment: nonceEnv, - EnvironmentInstance: instance, - Purpose: nonce.PurposeClientTransaction, - State: nonce.StateAvailable, - } - require.NoError(t, env.data.SaveKey(env.ctx, nonceKey)) - require.NoError(t, env.data.SaveNonce(env.ctx, nonceRecord)) - return nonceRecord -} - -func generateClaimedNonce(t *testing.T, env nonceTestEnv, nonceEnv nonce.Environment, instance string, expired bool) *nonce.Record { - nonceAccount := testutil.NewRandomAccount(t) - - var bh solana.Blockhash - rand.Read(bh[:]) - - nonceKey := &vault.Record{ - PublicKey: nonceAccount.PublicKey().ToBase58(), - PrivateKey: nonceAccount.PrivateKey().ToBase58(), - State: vault.StateReserved, - CreatedAt: time.Now(), - } - nonceRecord := &nonce.Record{ - Address: nonceAccount.PublicKey().ToBase58(), - Authority: common.GetSubsidizer().PublicKey().ToBase58(), - Blockhash: base58.Encode(bh[:]), - Environment: nonceEnv, - EnvironmentInstance: instance, - Purpose: nonce.PurposeClientTransaction, - State: nonce.StateClaimed, - ClaimNodeID: pointer.String("my_node_id"), - } - if expired { - nonceRecord.ClaimExpiresAt = pointer.Time(time.Now().Add(-time.Hour)) - } else { - nonceRecord.ClaimExpiresAt = pointer.Time(time.Now().Add(time.Hour)) - } - - require.NoError(t, env.data.SaveKey(env.ctx, nonceKey)) - require.NoError(t, env.data.SaveNonce(env.ctx, nonceRecord)) - return nonceRecord -} - -func generateAvailableNonces(t *testing.T, env nonceTestEnv, nonceEnv nonce.Environment, instance string, useCase nonce.Purpose, count int) []*nonce.Record { - var nonces []*nonce.Record - for i := 0; i < count; i++ { - nonces = append(nonces, generateAvailableNonce(t, env, nonceEnv, instance, useCase)) - } - return nonces -} From e486ce0ce27bfbaab2de72ac970364c6f214dad3 Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 14:41:59 -0400 Subject: [PATCH 4/8] Setup local nonce pool in gRPC transaction service --- pkg/code/server/transaction/airdrop.go | 43 +++++++++---------------- pkg/code/server/transaction/intent.go | 8 ++--- pkg/code/server/transaction/server.go | 28 ++++++++++++++-- pkg/code/transaction/nonce_pool.go | 40 ++++++++++++++++++++--- pkg/code/transaction/nonce_pool_test.go | 19 +++++++++-- 5 files changed, 95 insertions(+), 43 deletions(-) diff --git a/pkg/code/server/transaction/airdrop.go b/pkg/code/server/transaction/airdrop.go index 5be30f4d..78e8bd67 100644 --- a/pkg/code/server/transaction/airdrop.go +++ b/pkg/code/server/transaction/airdrop.go @@ -25,9 +25,7 @@ import ( "github.com/code-payments/code-server/pkg/code/data/currency" "github.com/code-payments/code-server/pkg/code/data/fulfillment" "github.com/code-payments/code-server/pkg/code/data/intent" - "github.com/code-payments/code-server/pkg/code/data/nonce" exchange_rate_util "github.com/code-payments/code-server/pkg/code/exchangerate" - "github.com/code-payments/code-server/pkg/code/transaction" currency_lib "github.com/code-payments/code-server/pkg/currency" "github.com/code-payments/code-server/pkg/grpc/client" "github.com/code-payments/code-server/pkg/pointer" @@ -297,14 +295,13 @@ func (s *transactionServer) airdrop(ctx context.Context, intentId string, owner // Instead of constructing and validating everything manually, we could // have a proper client call SubmitIntent in a worker. - selectedNonce, err := transaction.SelectAvailableNonce(ctx, s.data, nonce.EnvironmentCvm, common.CodeVmAccount.PublicKey().ToBase58(), nonce.PurposeClientTransaction) + selectedNonce, err := s.noncePool.GetNonce(ctx) if err != nil { log.WithError(err).Warn("failure selecting available nonce") return nil, err } defer func() { selectedNonce.ReleaseIfNotReserved() - selectedNonce.Unlock() }() vixnHash := cvm.GetCompactTransferMessage(&cvm.GetCompactTransferMessageArgs{ @@ -424,34 +421,24 @@ func (s *transactionServer) airdrop(ctx context.Context, intentId string, owner return intentRecord, nil } -func (s *transactionServer) mustLoadAirdropper(ctx context.Context) { - log := s.log.WithFields(logrus.Fields{ - "method": "mustLoadAirdropper", - "key": s.conf.airdropperOwnerPublicKey.Get(ctx), - }) - - err := func() error { - vaultRecord, err := s.data.GetKey(ctx, s.conf.airdropperOwnerPublicKey.Get(ctx)) - if err != nil { - return err - } - - ownerAccount, err := common.NewAccountFromPrivateKeyString(vaultRecord.PrivateKey) - if err != nil { - return err - } +func (s *transactionServer) loadAirdropper(ctx context.Context) error { + vaultRecord, err := s.data.GetKey(ctx, s.conf.airdropperOwnerPublicKey.Get(ctx)) + if err != nil { + return err + } - timelockAccounts, err := ownerAccount.GetTimelockAccounts(common.CodeVmAccount, common.CoreMintAccount) - if err != nil { - return err - } + ownerAccount, err := common.NewAccountFromPrivateKeyString(vaultRecord.PrivateKey) + if err != nil { + return err + } - s.airdropper = timelockAccounts - return nil - }() + timelockAccounts, err := ownerAccount.GetTimelockAccounts(common.CodeVmAccount, common.CoreMintAccount) if err != nil { - log.WithError(err).Fatal("failure loading account") + return err } + + s.airdropper = timelockAccounts + return nil } func GetOldAirdropIntentId(airdropType AirdropType, reference string) string { diff --git a/pkg/code/server/transaction/intent.go b/pkg/code/server/transaction/intent.go index c20de1c7..cca70d30 100644 --- a/pkg/code/server/transaction/intent.go +++ b/pkg/code/server/transaction/intent.go @@ -25,7 +25,6 @@ import ( "github.com/code-payments/code-server/pkg/code/data/action" "github.com/code-payments/code-server/pkg/code/data/fulfillment" "github.com/code-payments/code-server/pkg/code/data/intent" - "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/code/data/timelock" "github.com/code-payments/code-server/pkg/code/transaction" "github.com/code-payments/code-server/pkg/grpc/client" @@ -300,7 +299,7 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm var actionHandlers []CreateActionHandler var actionRecords []*action.Record var fulfillments []fulfillmentWithSigningMetadata - var reservedNonces []*transaction.SelectedNonce + var reservedNonces []*transaction.Nonce var serverParameters []*transactionpb.ServerParameter for i, protoAction := range submitActionsReq.Actions { log := log.WithField("action_id", i) @@ -368,11 +367,11 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm // Select any available nonce reserved for use for a client transaction, // if it's required - var selectedNonce *transaction.SelectedNonce + var selectedNonce *transaction.Nonce var nonceAccount *common.Account var nonceBlockchash solana.Blockhash if actionHandler.RequiresNonce(j) { - selectedNonce, err = transaction.SelectAvailableNonce(ctx, s.data, nonce.EnvironmentCvm, common.CodeVmAccount.PublicKey().ToBase58(), nonce.PurposeClientTransaction) + selectedNonce, err = s.noncePool.GetNonce(ctx) if err != nil { log.WithError(err).Warn("failure selecting available nonce") return handleSubmitIntentError(streamer, err) @@ -383,7 +382,6 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm // caused a failed RPC call, and we want to avoid malicious or erroneous // clients from consuming our nonce pool! selectedNonce.ReleaseIfNotReserved() - selectedNonce.Unlock() }() nonceAccount = selectedNonce.Account nonceBlockchash = selectedNonce.Blockhash diff --git a/pkg/code/server/transaction/server.go b/pkg/code/server/transaction/server.go index 8b5d1588..91332af3 100644 --- a/pkg/code/server/transaction/server.go +++ b/pkg/code/server/transaction/server.go @@ -4,6 +4,7 @@ import ( "context" "sync" + "github.com/pkg/errors" "github.com/sirupsen/logrus" transactionpb "github.com/code-payments/code-protobuf-api/generated/go/transaction/v2" @@ -13,6 +14,8 @@ import ( auth_util "github.com/code-payments/code-server/pkg/code/auth" "github.com/code-payments/code-server/pkg/code/common" code_data "github.com/code-payments/code-server/pkg/code/data" + "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/code/transaction" "github.com/code-payments/code-server/pkg/jupiter" sync_util "github.com/code-payments/code-server/pkg/sync" ) @@ -30,6 +33,8 @@ type transactionServer struct { antispamGuard *antispam.Guard amlGuard *aml.Guard + noncePool *transaction.LocalNoncePool + airdropperLock sync.Mutex airdropper *common.TimelockAccounts @@ -53,14 +58,26 @@ func NewTransactionServer( airdropIntegration AirdropIntegration, antispamGuard *antispam.Guard, amlGuard *aml.Guard, + noncePool *transaction.LocalNoncePool, configProvider ConfigProvider, -) transactionpb.TransactionServer { +) (transactionpb.TransactionServer, error) { ctx := context.Background() conf := configProvider() stripedLockParallelization := uint(conf.stripedLockParallelization.Get(ctx)) + noncePoolEnv, noncePoolEnvInstance, noncePoolType := noncePool.GetConfiguration() + if noncePoolEnv != nonce.EnvironmentCvm { + return nil, errors.Errorf("nonce pool environment must be %s", nonce.EnvironmentCvm) + } + if noncePoolEnvInstance != common.CodeVmAccount.PublicKey().ToBase58() { + return nil, errors.Errorf("nonce pool environment instance must be %s", common.CodeVmAccount.PublicKey().ToBase58()) + } + if noncePoolType != nonce.PurposeClientTransaction { + return nil, errors.Errorf("nonce pool type must be %s", nonce.PurposeClientTransaction) + } + s := &transactionServer{ log: logrus.StandardLogger().WithField("type", "transaction/v2/server"), conf: conf, @@ -74,6 +91,8 @@ func NewTransactionServer( antispamGuard: antispamGuard, amlGuard: amlGuard, + noncePool: noncePool, + intentLocks: sync_util.NewStripedLock(stripedLockParallelization), ownerLocks: sync_util.NewStripedLock(stripedLockParallelization), giftCardLocks: sync_util.NewStripedLock(stripedLockParallelization), @@ -81,8 +100,11 @@ func NewTransactionServer( airdropper := s.conf.airdropperOwnerPublicKey.Get(ctx) if len(airdropper) > 0 && airdropper != defaultAirdropperOwnerPublicKey { - s.mustLoadAirdropper(ctx) + err := s.loadAirdropper(ctx) + if err != nil { + return nil, err + } } - return s + return s, nil } diff --git a/pkg/code/transaction/nonce_pool.go b/pkg/code/transaction/nonce_pool.go index 6252336e..16ff61a1 100644 --- a/pkg/code/transaction/nonce_pool.go +++ b/pkg/code/transaction/nonce_pool.go @@ -2,19 +2,22 @@ package transaction import ( "context" - "errors" "fmt" "slices" "sync" "time" "github.com/google/uuid" + "github.com/mr-tron/base58/base58" "github.com/newrelic/go-agent/v3/newrelic" + "github.com/pkg/errors" "github.com/sirupsen/logrus" + "github.com/code-payments/code-server/pkg/code/common" code_data "github.com/code-payments/code-server/pkg/code/data" "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/pointer" + "github.com/code-payments/code-server/pkg/solana" ) var ( @@ -153,6 +156,9 @@ func (opts *noncePoolOpts) validate() error { // Nonce represents a handle to a nonce that is owned by a local nonce pool. type Nonce struct { + Account *common.Account + Blockhash solana.Blockhash + pool *LocalNoncePool record *nonce.Record } @@ -317,6 +323,10 @@ func (np *LocalNoncePool) GetNonce(ctx context.Context) (*Nonce, error) { return n, nil } +func (np *LocalNoncePool) GetConfiguration() (nonce.Environment, string, nonce.Purpose) { + return np.env, np.envInstance, np.poolType +} + func (np *LocalNoncePool) Close() error { log := np.log.WithField("method", "Close") @@ -376,13 +386,33 @@ func (np *LocalNoncePool) load(ctx context.Context, limit int) (int, error) { return 0, ErrNoAvailableNonces } - np.mu.Lock() - for i := range records { - np.freeList = append(np.freeList, &Nonce{pool: np, record: records[i]}) + var newNonces []*Nonce + for _, record := range records { + account, err := common.NewAccountFromPublicKeyString(record.Address) + if err != nil { + return 0, errors.Wrap(err, "invalid address") + } + + decodedBh, err := base58.Decode(record.Blockhash) + if err != nil { + return 0, errors.Wrap(err, "invalid blochash") + } + var bh solana.Blockhash + copy(bh[:], decodedBh) + + newNonces = append(newNonces, &Nonce{ + Account: account, + Blockhash: bh, + pool: np, + record: record, + }) } + + np.mu.Lock() + np.freeList = append(np.freeList, newNonces...) np.mu.Unlock() - return len(records), nil + return len(newNonces), nil } func (np *LocalNoncePool) refreshPool() { diff --git a/pkg/code/transaction/nonce_pool_test.go b/pkg/code/transaction/nonce_pool_test.go index 20dd60de..578fa83c 100644 --- a/pkg/code/transaction/nonce_pool_test.go +++ b/pkg/code/transaction/nonce_pool_test.go @@ -2,14 +2,18 @@ package transaction import ( "context" + "crypto/rand" "fmt" "testing" "time" + "github.com/mr-tron/base58" "github.com/stretchr/testify/require" + "github.com/code-payments/code-server/pkg/code/common" code_data "github.com/code-payments/code-server/pkg/code/data" "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/solana" "github.com/code-payments/code-server/pkg/testutil" ) @@ -60,6 +64,8 @@ func testLocalNoncePoolHappyPath(nt *localNoncePoolTest) { require.NoError(nt.t, err) require.NotNil(nt.t, n) require.NotContains(nt.t, observed, n.record.Address) + require.Equal(nt.t, n.Account.PublicKey().ToBase58(), n.record.Address) + require.Equal(nt.t, base58.Encode(n.Blockhash[:]), n.record.Blockhash) observed[n.record.Address] = n actual, err := nt.data.GetNonce(ctx, n.record.Address) @@ -157,6 +163,8 @@ func testLocalNoncePoolReload(nt *localNoncePoolTest) { require.NoError(nt.t, err) require.NotNil(nt.t, n) require.NotContains(nt.t, observed, n.record.Address) + require.Equal(nt.t, n.Account.PublicKey().ToBase58(), n.record.Address) + require.Equal(nt.t, base58.Encode(n.Blockhash[:]), n.record.Blockhash) observed[n.record.Address] = n actual, err := nt.data.GetNonce(ctx, n.record.Address) @@ -193,6 +201,8 @@ func testLocalNoncePoolRefresh(nt *localNoncePoolTest) { require.NoError(nt.t, err) require.NotNil(nt.t, n) require.NotContains(nt.t, observed, n.record.Address) + require.Equal(nt.t, n.Account.PublicKey().ToBase58(), n.record.Address) + require.Equal(nt.t, base58.Encode(n.Blockhash[:]), n.record.Blockhash) observed[n.record.Address] = n actual, err := nt.data.GetNonce(ctx, n.record.Address) @@ -235,6 +245,8 @@ type localNoncePoolTest struct { func newLocalNoncePoolTest(t *testing.T) *localNoncePoolTest { data := code_data.NewTestDataProvider() + testutil.SetupRandomSubsidizer(t, data) + pool, err := NewLocalNoncePool( data, nil, @@ -257,9 +269,12 @@ func newLocalNoncePoolTest(t *testing.T) *localNoncePoolTest { func (np *localNoncePoolTest) initializeNonces(amount int, env nonce.Environment, envInstance string, purpose nonce.Purpose) { for i := 0; i < amount; i++ { + var bh solana.Blockhash + rand.Read(bh[:]) err := np.data.SaveNonce(context.Background(), &nonce.Record{ - Address: fmt.Sprintf("addr-%s-%s-%s-%d", env.String(), envInstance, purpose.String(), i), - Authority: "my authority!", + Address: testutil.NewRandomAccount(np.t).PublicKey().ToBase58(), + Blockhash: base58.Encode(bh[:]), + Authority: common.GetSubsidizer().PublicKey().ToBase58(), Environment: env, EnvironmentInstance: envInstance, Purpose: purpose, From f10fca54804410871f0d32a12c3d004190862437 Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 14:53:19 -0400 Subject: [PATCH 5/8] Setup local nonce pool in sequencer service --- .../async/sequencer/fulfillment_handler.go | 10 +++++----- pkg/code/async/sequencer/service.go | 19 +++++++++++++++++-- pkg/code/async/sequencer/testutil.go | 6 +++--- pkg/code/async/sequencer/worker.go | 5 +---- pkg/code/async/sequencer/worker_test.go | 14 +++++++++++++- 5 files changed, 39 insertions(+), 15 deletions(-) diff --git a/pkg/code/async/sequencer/fulfillment_handler.go b/pkg/code/async/sequencer/fulfillment_handler.go index ff21f9f9..d599328d 100644 --- a/pkg/code/async/sequencer/fulfillment_handler.go +++ b/pkg/code/async/sequencer/fulfillment_handler.go @@ -46,7 +46,7 @@ type FulfillmentHandler interface { // to the blockchain. This is an optimization for the nonce pool. Implementations // should not modify the provided fulfillment record or selected nonce, but rather // use relevant fields to make the corresponding transaction. - MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) + MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) // OnSuccess is a callback function executed on a finalized transaction. OnSuccess(ctx context.Context, fulfillmentRecord *fulfillment.Record, txnRecord *transaction.Record) error @@ -123,7 +123,7 @@ func (h *InitializeLockedTimelockAccountFulfillmentHandler) SupportsOnDemandTran return true } -func (h *InitializeLockedTimelockAccountFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) { +func (h *InitializeLockedTimelockAccountFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) { if fulfillmentRecord.FulfillmentType != fulfillment.InitializeLockedTimelockAccount { return nil, errors.New("invalid fulfillment type") } @@ -268,7 +268,7 @@ func (h *NoPrivacyTransferWithAuthorityFulfillmentHandler) SupportsOnDemandTrans return true } -func (h *NoPrivacyTransferWithAuthorityFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) { +func (h *NoPrivacyTransferWithAuthorityFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) { virtualSignatureBytes, err := base58.Decode(*fulfillmentRecord.VirtualSignature) if err != nil { return nil, err @@ -444,7 +444,7 @@ func (h *NoPrivacyWithdrawFulfillmentHandler) SupportsOnDemandTransactions() boo return true } -func (h *NoPrivacyWithdrawFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) { +func (h *NoPrivacyWithdrawFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) { virtualSignatureBytes, err := base58.Decode(*fulfillmentRecord.VirtualSignature) if err != nil { return nil, err @@ -634,7 +634,7 @@ func (h *CloseEmptyTimelockAccountFulfillmentHandler) SupportsOnDemandTransactio return true } -func (h *CloseEmptyTimelockAccountFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) { +func (h *CloseEmptyTimelockAccountFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) { if fulfillmentRecord.FulfillmentType != fulfillment.CloseEmptyTimelockAccount { return nil, errors.New("invalid fulfillment type") } diff --git a/pkg/code/async/sequencer/service.go b/pkg/code/async/sequencer/service.go index 30a9c7c8..62a4f88f 100644 --- a/pkg/code/async/sequencer/service.go +++ b/pkg/code/async/sequencer/service.go @@ -14,6 +14,8 @@ import ( "github.com/code-payments/code-server/pkg/code/data/action" "github.com/code-payments/code-server/pkg/code/data/fulfillment" "github.com/code-payments/code-server/pkg/code/data/intent" + "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/code/transaction" ) var ( @@ -28,22 +30,35 @@ type service struct { data code_data.Provider scheduler Scheduler vmIndexerClient indexerpb.IndexerClient + noncePool *transaction.LocalNoncePool fulfillmentHandlersByType map[fulfillment.Type]FulfillmentHandler actionHandlersByType map[action.Type]ActionHandler intentHandlersByType map[intent.Type]IntentHandler } -func New(data code_data.Provider, scheduler Scheduler, vmIndexerClient indexerpb.IndexerClient, configProvider ConfigProvider) async.Service { +func New(data code_data.Provider, scheduler Scheduler, vmIndexerClient indexerpb.IndexerClient, noncePool *transaction.LocalNoncePool, configProvider ConfigProvider) (async.Service, error) { + noncePoolEnv, noncePoolEnvInstance, noncePoolType := noncePool.GetConfiguration() + if noncePoolEnv != nonce.EnvironmentSolana { + return nil, errors.Errorf("nonce pool environment must be %s", nonce.EnvironmentSolana) + } + if noncePoolEnvInstance != nonce.EnvironmentInstanceSolanaMainnet { + return nil, errors.Errorf("nonce pool environment instance must be %s", nonce.EnvironmentInstanceSolanaMainnet) + } + if noncePoolType != nonce.PurposeOnDemandTransaction { + return nil, errors.Errorf("nonce pool type must be %s", nonce.PurposeOnDemandTransaction) + } + return &service{ log: logrus.StandardLogger().WithField("service", "sequencer"), conf: configProvider(), data: data, scheduler: scheduler, vmIndexerClient: vmIndexerClient, + noncePool: noncePool, // todo: validate configuration fulfillmentHandlersByType: getFulfillmentHandlers(data, vmIndexerClient), actionHandlersByType: getActionHandlers(data), intentHandlersByType: getIntentHandlers(data), - } + }, nil } func (p *service) Start(ctx context.Context, interval time.Duration) error { diff --git a/pkg/code/async/sequencer/testutil.go b/pkg/code/async/sequencer/testutil.go index eb2081f7..8897fd48 100644 --- a/pkg/code/async/sequencer/testutil.go +++ b/pkg/code/async/sequencer/testutil.go @@ -4,12 +4,12 @@ import ( "context" "errors" - "github.com/code-payments/code-server/pkg/solana" - "github.com/code-payments/code-server/pkg/solana/memo" "github.com/code-payments/code-server/pkg/code/common" "github.com/code-payments/code-server/pkg/code/data/fulfillment" "github.com/code-payments/code-server/pkg/code/data/transaction" transaction_util "github.com/code-payments/code-server/pkg/code/transaction" + "github.com/code-payments/code-server/pkg/solana" + "github.com/code-payments/code-server/pkg/solana/memo" ) type mockScheduler struct { @@ -42,7 +42,7 @@ func (h *mockFulfillmentHandler) SupportsOnDemandTransactions() bool { return h.supportsOnDemandTxnCreation } -func (h *mockFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.SelectedNonce) (*solana.Transaction, error) { +func (h *mockFulfillmentHandler) MakeOnDemandTransaction(ctx context.Context, fulfillmentRecord *fulfillment.Record, selectedNonce *transaction_util.Nonce) (*solana.Transaction, error) { if !h.supportsOnDemandTxnCreation { return nil, errors.New("not supported") } diff --git a/pkg/code/async/sequencer/worker.go b/pkg/code/async/sequencer/worker.go index 081339d9..798310f1 100644 --- a/pkg/code/async/sequencer/worker.go +++ b/pkg/code/async/sequencer/worker.go @@ -12,9 +12,7 @@ import ( "github.com/code-payments/code-server/pkg/code/common" "github.com/code-payments/code-server/pkg/code/data/fulfillment" - "github.com/code-payments/code-server/pkg/code/data/nonce" "github.com/code-payments/code-server/pkg/code/data/transaction" - transaction_util "github.com/code-payments/code-server/pkg/code/transaction" "github.com/code-payments/code-server/pkg/database/query" "github.com/code-payments/code-server/pkg/metrics" "github.com/code-payments/code-server/pkg/pointer" @@ -225,13 +223,12 @@ func (p *service) handlePending(ctx context.Context, record *fulfillment.Record) return errors.New("unexpected scheduled fulfillment without transaction data") } - selectedNonce, err := transaction_util.SelectAvailableNonce(ctx, p.data, nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeOnDemandTransaction) + selectedNonce, err := p.noncePool.GetNonce(ctx) if err != nil { return err } defer func() { selectedNonce.ReleaseIfNotReserved() - selectedNonce.Unlock() }() err = p.data.ExecuteInTx(ctx, sql.LevelDefault, func(ctx context.Context) error { diff --git a/pkg/code/async/sequencer/worker_test.go b/pkg/code/async/sequencer/worker_test.go index f52b5201..ec16bc0b 100644 --- a/pkg/code/async/sequencer/worker_test.go +++ b/pkg/code/async/sequencer/worker_test.go @@ -105,6 +105,7 @@ func TestFulfillmentWorker_StatePending_OnDemandTransactionCreation_HappyPath(t env.fulfillmentHandler.supportsOnDemandTxnCreation = true nonceRecord := env.generateAvailableNonce(t) + time.Sleep(time.Second) // todo: Better way of waiting for nonce to be claimed fulfillmentRecord := env.createAnyFulfillmentInState(t, fulfillment.StatePending) fulfillmentRecord.Signature = nil @@ -226,12 +227,23 @@ func setupWorkerEnv(t *testing.T) *workerTestEnv { db := code_data.NewTestDataProvider() scheduler := &mockScheduler{} + noncePool, err := transaction_util.NewLocalNoncePool( + db, + nil, + nonce.EnvironmentSolana, + nonce.EnvironmentInstanceSolanaMainnet, + nonce.PurposeOnDemandTransaction, + transaction_util.WithNoncePoolRefreshPoolInterval(time.Second), + ) + require.NoError(t, err) fulfillmentHandler := &mockFulfillmentHandler{} actionHandler := &mockActionHandler{} intentHandler := &mockIntentHandler{} // todo: setup a test vm indexer - worker := New(db, scheduler, nil, withManualTestOverrides(&testOverrides{})).(*service) + workerInterface, err := New(db, scheduler, nil, noncePool, withManualTestOverrides(&testOverrides{})) + require.NoError(t, err) + worker := workerInterface.(*service) for key := range worker.fulfillmentHandlersByType { worker.fulfillmentHandlersByType[key] = fulfillmentHandler } From 1f9a7b41d3bf5a88f3c3a8faed95fee1a6e3389b Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 15:00:58 -0400 Subject: [PATCH 6/8] Add validation utility to LocalNoncePool for configuration enforcement --- pkg/code/async/sequencer/service.go | 11 ++--------- pkg/code/server/transaction/server.go | 12 ++---------- pkg/code/transaction/nonce_pool.go | 17 +++++++++++++++-- 3 files changed, 19 insertions(+), 21 deletions(-) diff --git a/pkg/code/async/sequencer/service.go b/pkg/code/async/sequencer/service.go index 62a4f88f..acc1c601 100644 --- a/pkg/code/async/sequencer/service.go +++ b/pkg/code/async/sequencer/service.go @@ -37,15 +37,8 @@ type service struct { } func New(data code_data.Provider, scheduler Scheduler, vmIndexerClient indexerpb.IndexerClient, noncePool *transaction.LocalNoncePool, configProvider ConfigProvider) (async.Service, error) { - noncePoolEnv, noncePoolEnvInstance, noncePoolType := noncePool.GetConfiguration() - if noncePoolEnv != nonce.EnvironmentSolana { - return nil, errors.Errorf("nonce pool environment must be %s", nonce.EnvironmentSolana) - } - if noncePoolEnvInstance != nonce.EnvironmentInstanceSolanaMainnet { - return nil, errors.Errorf("nonce pool environment instance must be %s", nonce.EnvironmentInstanceSolanaMainnet) - } - if noncePoolType != nonce.PurposeOnDemandTransaction { - return nil, errors.Errorf("nonce pool type must be %s", nonce.PurposeOnDemandTransaction) + if err := noncePool.Validate(nonce.EnvironmentSolana, nonce.EnvironmentInstanceSolanaMainnet, nonce.PurposeOnDemandTransaction); err != nil { + return nil, err } return &service{ diff --git a/pkg/code/server/transaction/server.go b/pkg/code/server/transaction/server.go index 91332af3..d6b2ed8f 100644 --- a/pkg/code/server/transaction/server.go +++ b/pkg/code/server/transaction/server.go @@ -4,7 +4,6 @@ import ( "context" "sync" - "github.com/pkg/errors" "github.com/sirupsen/logrus" transactionpb "github.com/code-payments/code-protobuf-api/generated/go/transaction/v2" @@ -67,15 +66,8 @@ func NewTransactionServer( stripedLockParallelization := uint(conf.stripedLockParallelization.Get(ctx)) - noncePoolEnv, noncePoolEnvInstance, noncePoolType := noncePool.GetConfiguration() - if noncePoolEnv != nonce.EnvironmentCvm { - return nil, errors.Errorf("nonce pool environment must be %s", nonce.EnvironmentCvm) - } - if noncePoolEnvInstance != common.CodeVmAccount.PublicKey().ToBase58() { - return nil, errors.Errorf("nonce pool environment instance must be %s", common.CodeVmAccount.PublicKey().ToBase58()) - } - if noncePoolType != nonce.PurposeClientTransaction { - return nil, errors.Errorf("nonce pool type must be %s", nonce.PurposeClientTransaction) + if err := noncePool.Validate(nonce.EnvironmentCvm, common.CodeVmAccount.PublicKey().ToBase58(), nonce.PurposeClientTransaction); err != nil { + return nil, err } s := &transactionServer{ diff --git a/pkg/code/transaction/nonce_pool.go b/pkg/code/transaction/nonce_pool.go index 16ff61a1..9a619922 100644 --- a/pkg/code/transaction/nonce_pool.go +++ b/pkg/code/transaction/nonce_pool.go @@ -323,8 +323,21 @@ func (np *LocalNoncePool) GetNonce(ctx context.Context) (*Nonce, error) { return n, nil } -func (np *LocalNoncePool) GetConfiguration() (nonce.Environment, string, nonce.Purpose) { - return np.env, np.envInstance, np.poolType +func (np *LocalNoncePool) Validate( + env nonce.Environment, + envInstance string, + poolType nonce.Purpose, +) error { + if np.env != env { + return errors.Errorf("nonce pool environment must be %s", env) + } + if np.envInstance != envInstance { + return errors.Errorf("nonce pool environment instance must be %s", envInstance) + } + if np.poolType != poolType { + return errors.Errorf("nonce pool type must be %s", poolType) + } + return nil } func (np *LocalNoncePool) Close() error { From 3b206f8b6a8006037dfb5e5156977e0952414ad8 Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Wed, 28 May 2025 15:33:36 -0400 Subject: [PATCH 7/8] Add nonce and nonce pool method tracing --- pkg/code/async/sequencer/worker.go | 2 +- pkg/code/server/transaction/airdrop.go | 2 +- pkg/code/server/transaction/intent.go | 2 +- pkg/code/transaction/nonce_pool.go | 105 +++++++++++++++--------- pkg/code/transaction/nonce_pool_test.go | 4 +- 5 files changed, 71 insertions(+), 44 deletions(-) diff --git a/pkg/code/async/sequencer/worker.go b/pkg/code/async/sequencer/worker.go index 798310f1..c2f1ebdf 100644 --- a/pkg/code/async/sequencer/worker.go +++ b/pkg/code/async/sequencer/worker.go @@ -228,7 +228,7 @@ func (p *service) handlePending(ctx context.Context, record *fulfillment.Record) return err } defer func() { - selectedNonce.ReleaseIfNotReserved() + selectedNonce.ReleaseIfNotReserved(ctx) }() err = p.data.ExecuteInTx(ctx, sql.LevelDefault, func(ctx context.Context) error { diff --git a/pkg/code/server/transaction/airdrop.go b/pkg/code/server/transaction/airdrop.go index 78e8bd67..702f4ffe 100644 --- a/pkg/code/server/transaction/airdrop.go +++ b/pkg/code/server/transaction/airdrop.go @@ -301,7 +301,7 @@ func (s *transactionServer) airdrop(ctx context.Context, intentId string, owner return nil, err } defer func() { - selectedNonce.ReleaseIfNotReserved() + selectedNonce.ReleaseIfNotReserved(ctx) }() vixnHash := cvm.GetCompactTransferMessage(&cvm.GetCompactTransferMessageArgs{ diff --git a/pkg/code/server/transaction/intent.go b/pkg/code/server/transaction/intent.go index cca70d30..d41e2163 100644 --- a/pkg/code/server/transaction/intent.go +++ b/pkg/code/server/transaction/intent.go @@ -381,7 +381,7 @@ func (s *transactionServer) SubmitIntent(streamer transactionpb.Transaction_Subm // it's safe to put it back in the available pool. The client will have // caused a failed RPC call, and we want to avoid malicious or erroneous // clients from consuming our nonce pool! - selectedNonce.ReleaseIfNotReserved() + selectedNonce.ReleaseIfNotReserved(ctx) }() nonceAccount = selectedNonce.Account nonceBlockchash = selectedNonce.Blockhash diff --git a/pkg/code/transaction/nonce_pool.go b/pkg/code/transaction/nonce_pool.go index 9a619922..47ca486d 100644 --- a/pkg/code/transaction/nonce_pool.go +++ b/pkg/code/transaction/nonce_pool.go @@ -16,10 +16,16 @@ import ( "github.com/code-payments/code-server/pkg/code/common" code_data "github.com/code-payments/code-server/pkg/code/data" "github.com/code-payments/code-server/pkg/code/data/nonce" + "github.com/code-payments/code-server/pkg/metrics" "github.com/code-payments/code-server/pkg/pointer" "github.com/code-payments/code-server/pkg/solana" ) +const ( + nonceMetricsStructName = "transaction.Nonce" + localNoncePoolMetricsStructName = "transaction.LocalNoncePool" +) + var ( ErrNoAvailableNonces = errors.New("no available nonces") ErrNoncePoolClosed = errors.New("nonce pool is closed") @@ -165,33 +171,43 @@ type Nonce struct { // MarkReservedWithSignature marks the nonce as reserved with a signature func (n *Nonce) MarkReservedWithSignature(ctx context.Context, sig string) error { - if len(sig) == 0 { - return errors.New("signature is empty") - } + tracer := metrics.TraceMethodCall(ctx, nonceMetricsStructName, "MarkReservedWithSignature") + defer tracer.End() - if n.record.Signature == sig { - return nil - } + err := func() error { + if len(sig) == 0 { + return errors.New("signature is empty") + } - if n.record.State == nonce.StateReserved || len(n.record.Signature) != 0 { - return errors.New("nonce already reserved with a different signature") - } + if n.record.Signature == sig { + return nil + } - if !n.record.CanReserveWithSignature() { - return errors.New("nonce is not in a valid state to reserve with signature") - } + if n.record.State == nonce.StateReserved || len(n.record.Signature) != 0 { + return errors.New("nonce already reserved with a different signature") + } - n.record.State = nonce.StateReserved - n.record.Signature = sig - n.record.ClaimNodeID = nil - n.record.ClaimExpiresAt = nil - return n.pool.data.SaveNonce(ctx, n.record) + if !n.record.CanReserveWithSignature() { + return errors.New("nonce is not in a valid state to reserve with signature") + } + + n.record.State = nonce.StateReserved + n.record.Signature = sig + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil + return n.pool.data.SaveNonce(ctx, n.record) + }() + tracer.OnError(err) + return err } // ReleaseIfNotReserved releases the nonce back to the pool if // the nonce has not yet been reserved (or more specifically, is // still owned by the pool). -func (n *Nonce) ReleaseIfNotReserved() { +func (n *Nonce) ReleaseIfNotReserved(ctx context.Context) { + tracer := metrics.TraceMethodCall(ctx, nonceMetricsStructName, "ReleaseIfNotReserved") + defer tracer.End() + if n.record.State != nonce.StateClaimed { return } @@ -295,32 +311,39 @@ func NewLocalNoncePool( } func (np *LocalNoncePool) GetNonce(ctx context.Context) (*Nonce, error) { - var n *Nonce + tracer := metrics.TraceMethodCall(ctx, localNoncePoolMetricsStructName, "GetNonce") + defer tracer.End() - np.mu.Lock() - if np.isClosed { - np.mu.Unlock() - return nil, ErrNoncePoolClosed - } + n, err := func() (*Nonce, error) { + var n *Nonce - size := len(np.freeList) - if size > 0 { - n = np.freeList[0] - np.freeList = np.freeList[1:] - } - np.mu.Unlock() + np.mu.Lock() + if np.isClosed { + np.mu.Unlock() + return nil, ErrNoncePoolClosed + } - if size < np.opts.desiredPoolSize/2 { - select { - case np.refreshPoolCh <- struct{}{}: - default: + size := len(np.freeList) + if size > 0 { + n = np.freeList[0] + np.freeList = np.freeList[1:] } - } + np.mu.Unlock() - if n == nil { - return nil, ErrNoAvailableNonces - } - return n, nil + if size < np.opts.desiredPoolSize/2 { + select { + case np.refreshPoolCh <- struct{}{}: + default: + } + } + + if n == nil { + return nil, ErrNoAvailableNonces + } + return n, nil + }() + tracer.OnError(err) + return n, err } func (np *LocalNoncePool) Validate( @@ -346,6 +369,10 @@ func (np *LocalNoncePool) Close() error { np.mu.Lock() defer np.mu.Unlock() + if np.isClosed { + return nil + } + np.isClosed = true np.cancelWorkerCtx() diff --git a/pkg/code/transaction/nonce_pool_test.go b/pkg/code/transaction/nonce_pool_test.go index 578fa83c..846c5c72 100644 --- a/pkg/code/transaction/nonce_pool_test.go +++ b/pkg/code/transaction/nonce_pool_test.go @@ -106,7 +106,7 @@ func testLocalNoncePoolHappyPath(nt *localNoncePoolTest) { // Releasing back to the pool should allow us to re-use the nonces that // weren't reserved for _, v := range observed { - v.ReleaseIfNotReserved() + v.ReleaseIfNotReserved(ctx) } clear(observed) @@ -127,7 +127,7 @@ func testLocalNoncePoolHappyPath(nt *localNoncePoolTest) { // Releasing back to the pool will allow us to make nonces available to // claim for _, v := range observed { - v.ReleaseIfNotReserved() + v.ReleaseIfNotReserved(ctx) } require.NoError(nt.t, nt.pool.Close()) From ced51e7a6ba844eaedea0db254038d0c3ca1f3ea Mon Sep 17 00:00:00 2001 From: Jeff Yanta Date: Thu, 29 May 2025 09:51:24 -0400 Subject: [PATCH 8/8] Add more edge case tests --- pkg/code/transaction/nonce_pool.go | 6 ++---- pkg/code/transaction/nonce_pool_test.go | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/pkg/code/transaction/nonce_pool.go b/pkg/code/transaction/nonce_pool.go index 47ca486d..fb7e5968 100644 --- a/pkg/code/transaction/nonce_pool.go +++ b/pkg/code/transaction/nonce_pool.go @@ -284,17 +284,15 @@ func NewLocalNoncePool( refreshPoolCh: make(chan struct{}), } - np.workerCtx, np.cancelWorkerCtx = context.WithCancel(context.Background()) - for _, o := range opts { o(&np.opts) } - if err := np.opts.validate(); err != nil { - np.cancelWorkerCtx() return nil, err } + np.workerCtx, np.cancelWorkerCtx = context.WithCancel(context.Background()) + _, err := np.load(np.workerCtx, np.opts.desiredPoolSize) switch err { case nil, nonce.ErrNonceNotFound: diff --git a/pkg/code/transaction/nonce_pool_test.go b/pkg/code/transaction/nonce_pool_test.go index 846c5c72..f2e40301 100644 --- a/pkg/code/transaction/nonce_pool_test.go +++ b/pkg/code/transaction/nonce_pool_test.go @@ -234,6 +234,22 @@ func testLocalNoncePoolReserveWithSignatureEdgeCases(nt *localNoncePoolTest) { require.NoError(nt.t, err) require.Equal(nt.t, actual.State, nonce.StateReserved) require.Equal(nt.t, actual.Signature, "signature1") + + n, err = nt.pool.GetNonce(ctx) + require.NoError(nt.t, err) + require.NotNil(nt.t, n) + + n.record.State = nonce.StateUnknown + n.record.ClaimNodeID = nil + n.record.ClaimExpiresAt = nil + require.NoError(nt.t, nt.data.SaveNonce(ctx, n.record)) + + require.Error(nt.t, n.MarkReservedWithSignature(ctx, "signature3")) + + actual, err = nt.data.GetNonce(ctx, n.record.Address) + require.NoError(nt.t, err) + require.Equal(nt.t, actual.State, nonce.StateUnknown) + require.Empty(nt.t, actual.Signature) } type localNoncePoolTest struct {