Setup local nonce pool in gRPC transaction service

Author: Jeff Yanta

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 {

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

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,15 +91,20 @@ func NewTransactionServer(
 		antispamGuard: antispamGuard,
 		amlGuard:      amlGuard,
 
+		noncePool: noncePool,
+
 		intentLocks:   sync_util.NewStripedLock(stripedLockParallelization),
 		ownerLocks:    sync_util.NewStripedLock(stripedLockParallelization),
 		giftCardLocks: sync_util.NewStripedLock(stripedLockParallelization),
 	}
 
 	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
 }

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() {

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,