DB-level changes to support distributed nonce pools

Author: Jeff Yanta

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 {

pkg/code/async/nonce/pool.go

@@ -95,31 +95,36 @@ 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
 					-> StateReleased
 				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
+}

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)
 }

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
 }

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,22 +24,22 @@ 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
 
 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"