Fixes for stateful swap system

jeffyanta · jeffyanta · commit 24d0d6183745 · 2025-11-26T12:53:08.000-05:00
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	filippo.io/edwards25519 v1.1.0
 	github.com/aws/aws-sdk-go-v2 v0.17.0
 	github.com/bits-and-blooms/bloom/v3 v3.1.0
-	github.com/code-payments/code-protobuf-api v1.19.1-0.20251124021538-11fc87a9e233
+	github.com/code-payments/code-protobuf-api v1.19.1-0.20251126151344-3fd1dad96a6f
 	github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba
 	github.com/emirpasic/gods v1.12.0
 	github.com/envoyproxy/protoc-gen-validate v1.2.1
diff --git a/go.sum b/go.sum
@@ -80,8 +80,8 @@ github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnht
 github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk=
 github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I=
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
-github.com/code-payments/code-protobuf-api v1.19.1-0.20251124021538-11fc87a9e233 h1:JTrB4gOeGjaguLD00bBWdIiO7fdQmvXkyWyOV9ZNSnk=
-github.com/code-payments/code-protobuf-api v1.19.1-0.20251124021538-11fc87a9e233/go.mod h1:fl4xu32MeNpGZR3wFhwEeKO0qVDuxYBNkqnvuADt6cA=
+github.com/code-payments/code-protobuf-api v1.19.1-0.20251126151344-3fd1dad96a6f h1:iqZ/lwIKAyfbyV1LKfoDfevdOcDu5QxkyeHd8rrH5KQ=
+github.com/code-payments/code-protobuf-api v1.19.1-0.20251126151344-3fd1dad96a6f/go.mod h1:fl4xu32MeNpGZR3wFhwEeKO0qVDuxYBNkqnvuADt6cA=
 github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba h1:Bkp+gmeb6Y2PWXfkSCTMBGWkb2P1BujRDSjWeI+0j5I=
 github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba/go.mod h1:jSiifpiBpyBQ8q2R0MGEbkSgWC6sbdRTyDBntmW+j1E=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw=
diff --git a/pkg/code/async/geyser/external_deposit.go b/pkg/code/async/geyser/external_deposit.go
@@ -99,7 +99,7 @@ func initiateExternalDepositIntoVm(ctx context.Context, data code_data.Provider,
 		return errors.Wrap(err, "error getting timelock accounts")
 	}
 
-	err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, data, vmIndexerClient, vmConfig.Vm, userAuthority, true)
+	err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, data, vmIndexerClient, mint, userAuthority, true)
 	if err != nil {
 		return errors.Wrap(err, "error ensuring vta is initialized")
 	}
diff --git a/pkg/code/async/nonce/pool.go b/pkg/code/async/nonce/pool.go
@@ -199,7 +199,7 @@ func (p *service) handleReleased(ctx context.Context, record *nonce.Record) erro
 	// Nonces that exist but we don't yet know their stored blockhash.
 
 	// Fetch the Solana transaction where the nonce would be consumed
-	var txn *transaction.Record
+	var txn *solana.ConfirmedTransaction
 	var err error
 	switch record.Environment {
 	case nonce.EnvironmentSolana:
@@ -220,7 +220,7 @@ func (p *service) handleReleased(ctx context.Context, record *nonce.Record) erro
 	}
 
 	// Sanity check the Solana transaction is in a finalized or failed state
-	if txn.ConfirmationState != transaction.ConfirmationFinalized && txn.ConfirmationState != transaction.ConfirmationFailed {
+	if txn == nil {
 		return nil
 	}
 
diff --git a/pkg/code/async/nonce/util.go b/pkg/code/async/nonce/util.go
@@ -66,24 +66,19 @@ func (p *service) getLatestBlockhash(ctx context.Context) (*solana.Blockhash, er
 	return &bh, nil
 }
 
-func (p *service) getTransaction(ctx context.Context, signature string) (*transaction.Record, error) {
+func (p *service) getTransaction(ctx context.Context, signature string) (*solana.ConfirmedTransaction, error) {
 	return p.getTransactionFromBlockchain(ctx, signature)
 }
 
-func (p *service) getTransactionFromBlockchain(ctx context.Context, signature string) (*transaction.Record, error) {
-	stx, err := p.data.GetBlockchainTransaction(ctx, signature, solana.CommitmentFinalized)
+func (p *service) getTransactionFromBlockchain(ctx context.Context, signature string) (*solana.ConfirmedTransaction, error) {
+	txn, err := p.data.GetBlockchainTransaction(ctx, signature, solana.CommitmentFinalized)
 	if err == solana.ErrSignatureNotFound {
 		return nil, transaction.ErrNotFound
 	}
 	if err != nil {
 		return nil, err
 	}
-
-	tx, err := transaction.FromConfirmedTransaction(stx)
-	if err != nil {
-		return nil, err
-	}
-	return tx, nil
+	return txn, nil
 }
 
 func (p *service) getRentAmount(ctx context.Context) (uint64, error) {
diff --git a/pkg/code/async/swap/config.go b/pkg/code/async/swap/config.go
@@ -3,8 +3,6 @@ package async_swap
 import (
 	"github.com/code-payments/code-server/pkg/config"
 	"github.com/code-payments/code-server/pkg/config/env"
-	"github.com/code-payments/code-server/pkg/config/memory"
-	"github.com/code-payments/code-server/pkg/config/wrapper"
 )
 
 const (
@@ -15,8 +13,7 @@ const (
 )
 
 type conf struct {
-	batchSize                     config.Uint64
-	enableCachedTransactionLookup config.Bool
+	batchSize config.Uint64
 }
 
 // ConfigProvider defines how config values are pulled
@@ -26,8 +23,7 @@ type ConfigProvider func() *conf
 func WithEnvConfigs() ConfigProvider {
 	return func() *conf {
 		return &conf{
-			batchSize:                     env.NewUint64Config(BatchSizeConfigEnvName, defaultFulfillmentBatchSize),
-			enableCachedTransactionLookup: wrapper.NewBoolConfig(memory.NewConfig(false), false),
+			batchSize: env.NewUint64Config(BatchSizeConfigEnvName, defaultFulfillmentBatchSize),
 		}
 	}
 }
diff --git a/pkg/code/async/swap/util.go b/pkg/code/async/swap/util.go
@@ -179,35 +179,6 @@ func (p *service) markNonceReleasedDueToSubmittedTransaction(ctx context.Context
 	return p.data.SaveNonce(ctx, nonceRecord)
 }
 
-func (p *service) getTransaction(ctx context.Context, record *swap.Record) (*transaction.Record, error) {
-	if record.TransactionSignature == nil || len(*record.TransactionSignature) == 0 {
-		return nil, transaction.ErrNotFound
-	}
-
-	if p.conf.enableCachedTransactionLookup.Get(ctx) {
-		return p.data.GetTransaction(ctx, *record.TransactionSignature)
-	}
-
-	return p.getTransactionFromBlockchain(ctx, record)
-}
-
-func (p *service) getTransactionFromBlockchain(ctx context.Context, record *swap.Record) (*transaction.Record, error) {
-	stx, err := p.data.GetBlockchainTransaction(ctx, *record.TransactionSignature, solana.CommitmentFinalized)
-	if err == solana.ErrSignatureNotFound {
-		return nil, transaction.ErrNotFound
-	}
-	if err != nil {
-		return nil, err
-	}
-
-	tx, err := transaction.FromConfirmedTransaction(stx)
-	if err != nil {
-		return nil, err
-	}
-
-	return tx, nil
-}
-
 func getDeltaQuarksFromTokenBalances(tokenAccount *common.Account, tokenBalances *solana.TransactionTokenBalances) (int64, error) {
 	var preQuarkBalance, postQuarkBalance int64
 	var err error
diff --git a/pkg/code/async/swap/worker.go b/pkg/code/async/swap/worker.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/code-payments/code-server/pkg/code/data/intent"
 	"github.com/code-payments/code-server/pkg/code/data/swap"
-	"github.com/code-payments/code-server/pkg/code/data/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/retry"
@@ -106,17 +105,15 @@ func (p *service) handleStateSubmitting(ctx context.Context, record *swap.Record
 		return err
 	}
 
-	finalizedTxn, err := p.getTransaction(ctx, record)
+	finalizedTxn, err := p.data.GetBlockchainTransaction(ctx, *record.TransactionSignature, solana.CommitmentFinalized)
 	if err != nil && err != solana.ErrSignatureNotFound {
 		return errors.Wrap(err, "error getting finalized transaction")
 	}
 
 	if finalizedTxn != nil {
-		if finalizedTxn.HasErrors || finalizedTxn.ConfirmationState == transaction.ConfirmationFailed {
+		if finalizedTxn.Err != nil || finalizedTxn.Meta.Err != nil {
 			return p.markSwapFailed(ctx, record)
-		}
-
-		if finalizedTxn.ConfirmationState == transaction.ConfirmationFinalized {
+		} else {
 			err = p.updateBalances(ctx, record)
 			if err != nil {
 				return errors.Wrap(err, "error updating balances")
diff --git a/pkg/code/common/vm.go b/pkg/code/common/vm.go
@@ -110,8 +110,8 @@ func GetVmConfigForMint(ctx context.Context, data code_data.Provider, mint *Acco
 	}
 }
 
-func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_data.Provider, vmIndexerClient indexerpb.IndexerClient, vm, owner *Account, waitForInitialization bool) error {
-	vmConfig, err := GetVmConfigForMint(ctx, data, vm)
+func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_data.Provider, vmIndexerClient indexerpb.IndexerClient, mint, owner *Account, waitForInitialization bool) error {
+	vmConfig, err := GetVmConfigForMint(ctx, data, mint)
 	if err != nil {
 		return err
 	}
@@ -147,12 +147,12 @@ func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_da
 	}
 
 	for range 60 {
-		time.Sleep(time.Second)
-
-		_, _, err := GetVirtualTimelockAccountLocationInMemory(ctx, vmIndexerClient, vm, owner)
+		_, _, err := GetVirtualTimelockAccountLocationInMemory(ctx, vmIndexerClient, vmConfig.Vm, owner)
 		if err == nil {
 			return nil
 		}
+
+		time.Sleep(time.Second)
 	}
 
 	return errors.New("timed out waiting for initialization")
diff --git a/pkg/code/server/transaction/errors.go b/pkg/code/server/transaction/errors.go
@@ -326,6 +326,8 @@ func handleStartSwapError(streamer transactionpb.Transaction_StartSwapServer, er
 		return status.Error(codes.DeadlineExceeded, err.Error())
 	case context.Canceled:
 		return status.Error(codes.Canceled, err.Error())
+	case transaction.ErrNoAvailableNonces, transaction.ErrNoncePoolNotFound:
+		return status.Error(codes.Unavailable, "")
 	}
 	return status.Error(codes.Internal, "rpc server failure")
 }
diff --git a/pkg/code/server/transaction/swap.go b/pkg/code/server/transaction/swap.go
@@ -190,6 +190,7 @@ func (s *transactionServer) StartSwap(streamer transactionpb.Transaction_StartSw
 		nonce.EnvironmentSolana,
 		nonce.EnvironmentInstanceSolanaMainnet,
 		nonce.PurposeClientSwap,
+		s.noncePools...,
 	)
 	if err != nil {
 		log.WithError(err).Warn("failure selecting nonce pool")
@@ -531,7 +532,7 @@ func (s *transactionServer) Swap(streamer transactionpb.Transaction_SwapServer)
 	}
 
 	if swapRecord.State != swap.StateFunded {
-		return handleSwapError(streamer, NewSwapValidationError("swap is not in a funded state"))
+		return handleSwapError(streamer, NewSwapValidationErrorf("swap state is %s", swapRecord.State))
 	}
 
 	if owner.PublicKey().ToBase58() == swapAuthority.PublicKey().ToBase58() {
@@ -558,7 +559,7 @@ func (s *transactionServer) Swap(streamer transactionpb.Transaction_SwapServer)
 	// Section: On-demand account creation
 	//
 
-	err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, s.data, s.vmIndexerClient, destinationVmConfig.Vm, owner, true)
+	err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, s.data, s.vmIndexerClient, toMint, owner, true)
 	if err != nil {
 		log.WithError(err).Warn("timed out waiting for destination timelock account initialization")
 		return handleSwapError(streamer, err)
@@ -630,7 +631,6 @@ func (s *transactionServer) Swap(streamer transactionpb.Transaction_SwapServer)
 
 	txn.SetBlockhash(solana.Blockhash(blockhash))
 
-	marshalledTxn := txn.Marshal()
 	marshalledTxnMessage := txn.Message.Marshal()
 
 	//
@@ -723,6 +723,8 @@ func (s *transactionServer) Swap(streamer transactionpb.Transaction_SwapServer)
 		return handleSwapError(streamer, err)
 	}
 
+	marshalledTxn := txn.Marshal()
+
 	txnSignature := base58.Encode(txn.Signature())
 
 	//

Original file line number	Diff line number	Diff line change
`@@ -99,7 +99,7 @@ func initiateExternalDepositIntoVm(ctx context.Context, data code_data.Provider,`
`99`	`99`	`return errors.Wrap(err, "error getting timelock accounts")`
`100`	`100`	`}`
`101`	`101`
`102`		`- err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, data, vmIndexerClient, vmConfig.Vm, userAuthority, true)`
	`102`	`+ err = common.EnsureVirtualTimelockAccountIsInitialized(ctx, data, vmIndexerClient, mint, userAuthority, true)`
`103`	`103`	`if err != nil {`
`104`	`104`	`return errors.Wrap(err, "error ensuring vta is initialized")`
`105`	`105`	`}`
Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,7 @@ func (p service) handleReleased(ctx context.Context, record nonce.Record) erro`
`199`	`199`	`// Nonces that exist but we don't yet know their stored blockhash.`
`200`	`200`
`201`	`201`	`// Fetch the Solana transaction where the nonce would be consumed`
`202`		`- var txn *transaction.Record`
	`202`	`+ var txn *solana.ConfirmedTransaction`
`203`	`203`	`var err error`
`204`	`204`	`switch record.Environment {`
`205`	`205`	`case nonce.EnvironmentSolana:`
`@@ -220,7 +220,7 @@ func (p service) handleReleased(ctx context.Context, record nonce.Record) erro`
`220`	`220`	`}`
`221`	`221`
`222`	`222`	`// Sanity check the Solana transaction is in a finalized or failed state`
`223`		`- if txn.ConfirmationState != transaction.ConfirmationFinalized && txn.ConfirmationState != transaction.ConfirmationFailed {`
	`223`	`+ if txn == nil {`
`224`	`224`	`return nil`
`225`	`225`	`}`
`226`	`226`
Original file line number	Diff line number	Diff line change
`@@ -66,24 +66,19 @@ func (p service) getLatestBlockhash(ctx context.Context) (solana.Blockhash, er`
`66`	`66`	`return &bh, nil`
`67`	`67`	`}`
`68`	`68`
`69`		`-func (p service) getTransaction(ctx context.Context, signature string) (transaction.Record, error) {`
	`69`	`+func (p service) getTransaction(ctx context.Context, signature string) (solana.ConfirmedTransaction, error) {`
`70`	`70`	`return p.getTransactionFromBlockchain(ctx, signature)`
`71`	`71`	`}`
`72`	`72`
`73`		`-func (p service) getTransactionFromBlockchain(ctx context.Context, signature string) (transaction.Record, error) {`
`74`		`- stx, err := p.data.GetBlockchainTransaction(ctx, signature, solana.CommitmentFinalized)`
	`73`	`+func (p service) getTransactionFromBlockchain(ctx context.Context, signature string) (solana.ConfirmedTransaction, error) {`
	`74`	`+ txn, err := p.data.GetBlockchainTransaction(ctx, signature, solana.CommitmentFinalized)`
`75`	`75`	`if err == solana.ErrSignatureNotFound {`
`76`	`76`	`return nil, transaction.ErrNotFound`
`77`	`77`	`}`
`78`	`78`	`if err != nil {`
`79`	`79`	`return nil, err`
`80`	`80`	`}`
`81`		`-`
`82`		`- tx, err := transaction.FromConfirmedTransaction(stx)`
`83`		`- if err != nil {`
`84`		`- return nil, err`
`85`		`- }`
`86`		`- return tx, nil`
	`81`	`+ return txn, nil`
`87`	`82`	`}`
`88`	`83`
`89`	`84`	`func (p *service) getRentAmount(ctx context.Context) (uint64, error) {`
Original file line number	Diff line number	Diff line change
`@@ -3,8 +3,6 @@ package async_swap`
`3`	`3`	`import (`
`4`	`4`	`"github.com/code-payments/code-server/pkg/config"`
`5`	`5`	`"github.com/code-payments/code-server/pkg/config/env"`
`6`		`- "github.com/code-payments/code-server/pkg/config/memory"`
`7`		`- "github.com/code-payments/code-server/pkg/config/wrapper"`
`8`	`6`	`)`
`9`	`7`
`10`	`8`	`const (`
`@@ -15,8 +13,7 @@ const (`
`15`	`13`	`)`
`16`	`14`
`17`	`15`	`type conf struct {`
`18`		`- batchSize config.Uint64`
`19`		`- enableCachedTransactionLookup config.Bool`
	`16`	`+ batchSize config.Uint64`
`20`	`17`	`}`
`21`	`18`
`22`	`19`	`// ConfigProvider defines how config values are pulled`
`@@ -26,8 +23,7 @@ type ConfigProvider func() *conf`
`26`	`23`	`func WithEnvConfigs() ConfigProvider {`
`27`	`24`	`return func() *conf {`
`28`	`25`	`return &conf{`
`29`		`- batchSize: env.NewUint64Config(BatchSizeConfigEnvName, defaultFulfillmentBatchSize),`
`30`		`- enableCachedTransactionLookup: wrapper.NewBoolConfig(memory.NewConfig(false), false),`
	`26`	`+ batchSize: env.NewUint64Config(BatchSizeConfigEnvName, defaultFulfillmentBatchSize),`
`31`	`27`	`}`
`32`	`28`	`}`
`33`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -110,8 +110,8 @@ func GetVmConfigForMint(ctx context.Context, data code_data.Provider, mint *Acco`
`110`	`110`	`}`
`111`	`111`	`}`
`112`	`112`
`113`		`-func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_data.Provider, vmIndexerClient indexerpb.IndexerClient, vm, owner *Account, waitForInitialization bool) error {`
`114`		`- vmConfig, err := GetVmConfigForMint(ctx, data, vm)`
	`113`	`+func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_data.Provider, vmIndexerClient indexerpb.IndexerClient, mint, owner *Account, waitForInitialization bool) error {`
	`114`	`+ vmConfig, err := GetVmConfigForMint(ctx, data, mint)`
`115`	`115`	`if err != nil {`
`116`	`116`	`return err`
`117`	`117`	`}`
`@@ -147,12 +147,12 @@ func EnsureVirtualTimelockAccountIsInitialized(ctx context.Context, data code_da`
`147`	`147`	`}`
`148`	`148`
`149`	`149`	`for range 60 {`
`150`		`- time.Sleep(time.Second)`
`151`		`-`
`152`		`- _, _, err := GetVirtualTimelockAccountLocationInMemory(ctx, vmIndexerClient, vm, owner)`
	`150`	`+ _, _, err := GetVirtualTimelockAccountLocationInMemory(ctx, vmIndexerClient, vmConfig.Vm, owner)`
`153`	`151`	`if err == nil {`
`154`	`152`	`return nil`
`155`	`153`	`}`
	`154`	`+`
	`155`	`+ time.Sleep(time.Second)`
`156`	`156`	`}`
`157`	`157`
`158`	`158`	`return errors.New("timed out waiting for initialization")`
Original file line number	Diff line number	Diff line change
`@@ -326,6 +326,8 @@ func handleStartSwapError(streamer transactionpb.Transaction_StartSwapServer, er`
`326`	`326`	`return status.Error(codes.DeadlineExceeded, err.Error())`
`327`	`327`	`case context.Canceled:`
`328`	`328`	`return status.Error(codes.Canceled, err.Error())`
	`329`	`+ case transaction.ErrNoAvailableNonces, transaction.ErrNoncePoolNotFound:`
	`330`	`+ return status.Error(codes.Unavailable, "")`
`329`	`331`	`}`
`330`	`332`	`return status.Error(codes.Internal, "rpc server failure")`
`331`	`333`	`}`