Skip to content

Commit

Permalink
Update simulation to set max compute unit limit and enable sig verifi…
Browse files Browse the repository at this point in the history 
…cation (#919)

* Updated simulation to set max CU limit and enabled sig verification

* Added signature to simulation tx and fixed tests

* Cleaned up code

* Updated simulation to explicitly set the configured commitment to avoid using the default

* Added check to ensure estimated compute unit limit does not exceed max after adding buffer

* Fixed linting
  • Loading branch information
@amit-momin
amit-momin authored Nov 18, 2024
1 parent b9bf6a3 commit e2db20a
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 21 deletions.
44 changes: 32 additions & 12 deletions pkg/solana/txm/txm.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"math"
"math/big"
"strings"
"sync"
Expand All @@ -20,6 +19,7 @@ import (
commontypes "github.com/smartcontractkit/chainlink-common/pkg/types"
"github.com/smartcontractkit/chainlink-common/pkg/utils"
bigmath "github.com/smartcontractkit/chainlink-common/pkg/utils/big_math"
"github.com/smartcontractkit/chainlink-common/pkg/utils/mathutil"

"github.com/smartcontractkit/chainlink-solana/pkg/solana/client"
"github.com/smartcontractkit/chainlink-solana/pkg/solana/config"
Expand All @@ -33,6 +33,7 @@ const (
MaxSigsToConfirm = 256 // max number of signatures in GetSignatureStatus call
EstimateComputeUnitLimitBuffer = 10 // percent buffer added on top of estimated compute unit limits to account for any variance
TxReapInterval = 10 * time.Second // interval of time between reaping transactions that have met the retention threshold
MaxComputeUnitLimit = 1_400_000 // max compute unit limit a transaction can have
)

var _ services.Service = (*Txm)(nil)
Expand Down Expand Up @@ -616,7 +617,7 @@ func (txm *Txm) Enqueue(ctx context.Context, accountID string, tx *solanaGo.Tran
select {
case txm.chSend <- msg:
default:
txm.lggr.Errorw("failed to enqeue tx", "queueFull", len(txm.chSend) == MaxQueueLen, "tx", msg)
txm.lggr.Errorw("failed to enqueue tx", "queueFull", len(txm.chSend) == MaxQueueLen, "tx", msg)
return fmt.Errorf("failed to enqueue transaction for %s", accountID)
}
return nil
Expand Down Expand Up @@ -646,16 +647,37 @@ func (txm *Txm) GetTransactionStatus(ctx context.Context, transactionID string)
// EstimateComputeUnitLimit estimates the compute unit limit needed for a transaction.
// It simulates the provided transaction to determine the used compute and applies a buffer to it.
func (txm *Txm) EstimateComputeUnitLimit(ctx context.Context, tx *solanaGo.Transaction) (uint32, error) {
res, err := txm.simulateTx(ctx, tx)
txCopy := *tx

// Set max compute unit limit when simulating a transaction to avoid getting an error for exceeding the default 200k compute unit limit
if computeUnitLimitErr := fees.SetComputeUnitLimit(&txCopy, fees.ComputeUnitLimit(MaxComputeUnitLimit)); computeUnitLimitErr != nil {
txm.lggr.Errorw("failed to set compute unit limit when simulating tx", "error", computeUnitLimitErr)
return 0, computeUnitLimitErr
}

// Sign and set signature in tx copy for simulation
txMsg, marshalErr := txCopy.Message.MarshalBinary()
if marshalErr != nil {
return 0, fmt.Errorf("failed to marshal tx message: %w", marshalErr)
}
sigBytes, signErr := txm.ks.Sign(ctx, txCopy.Message.AccountKeys[0].String(), txMsg)
if signErr != nil {
return 0, fmt.Errorf("failed to sign transaction: %w", signErr)
}
var sig [64]byte
copy(sig[:], sigBytes)
txCopy.Signatures = append(txCopy.Signatures, sig)

res, err := txm.simulateTx(ctx, &txCopy)
if err != nil {
return 0, err
}

// Return error if response err is non-nil to avoid broadcasting a tx destined to fail
if res.Err != nil {
sig := solanaGo.Signature{}
if len(tx.Signatures) > 0 {
sig = tx.Signatures[0]
if len(txCopy.Signatures) > 0 {
sig = txCopy.Signatures[0]
}
txm.processSimulationError("", sig, res)
return 0, fmt.Errorf("simulated tx returned error: %v", res.Err)
Expand All @@ -672,13 +694,10 @@ func (txm *Txm) EstimateComputeUnitLimit(ctx context.Context, tx *solanaGo.Trans
// Add buffer to the used compute estimate
unitsConsumed = bigmath.AddPercentage(new(big.Int).SetUint64(unitsConsumed), EstimateComputeUnitLimitBuffer).Uint64()

if unitsConsumed > math.MaxUint32 {
txm.lggr.Debug("compute units used with buffer greater than uint32 max", "unitsConsumed", unitsConsumed)
// Do not return error to allow falling back to default compute unit limit
return 0, nil
}
// Ensure unitsConsumed does not exceed the max compute unit limit for a transaction after adding buffer
unitsConsumed = mathutil.Min(unitsConsumed, MaxComputeUnitLimit)

return uint32(unitsConsumed), nil
return uint32(unitsConsumed), nil //nolint // unitsConsumed can only be a maximum of 1.4M
}

// simulateTx simulates transactions using the SimulateTx client method
Expand All @@ -690,7 +709,8 @@ func (txm *Txm) simulateTx(ctx context.Context, tx *solanaGo.Transaction) (res *
return
}

res, err = client.SimulateTx(ctx, tx, nil) // use default options (does not verify signatures)
// Simulate with signature verification enabled since it can have an impact on the compute units used
res, err = client.SimulateTx(ctx, tx, &rpc.SimulateTransactionOpts{SigVerify: true, Commitment: txm.cfg.Commitment()})
if err != nil {
// This error can occur if endpoint goes down or if invalid signature
txm.lggr.Errorw("failed to simulate tx", "error", err)
Expand Down
32 changes: 25 additions & 7 deletions pkg/solana/txm/txm_internal_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -910,18 +910,20 @@ func TestTxm_compute_unit_limit_estimation(t *testing.T) {
t.Run("simulation_succeeds", func(t *testing.T) {
// Test tx is not discarded due to confirm timeout and tracked to finalization
tx, signed := getTx(t, 1, mkey)
// add signature and compute unit limit to tx for simulation (excludes compute unit price)
simulateTx := addSigAndLimitToTx(t, mkey, solana.PublicKey{}, *tx, MaxComputeUnitLimit)
sig := randomSignature(t)
var wg sync.WaitGroup
wg.Add(3)

computeUnitConsumed := uint64(1_000_000)
computeUnitLimit := fees.ComputeUnitLimit(uint32(bigmath.AddPercentage(new(big.Int).SetUint64(computeUnitConsumed), EstimateComputeUnitLimitBuffer).Uint64()))
mc.On("SendTx", mock.Anything, signed(0, true, computeUnitLimit)).Return(sig, nil)
// First simulated before broadcast without signature or compute unit limit set
mc.On("SimulateTx", mock.Anything, tx, mock.Anything).Run(func(mock.Arguments) {
// First simulation before broadcast with signature and max compute unit limit set
mc.On("SimulateTx", mock.Anything, simulateTx, mock.Anything).Run(func(mock.Arguments) {
wg.Done()
}).Return(&rpc.SimulateTransactionResult{UnitsConsumed: &computeUnitConsumed}, nil).Once()
// Second simulated after broadcast with signature and compute unit limit set
// Second simulation after broadcast with signature and compute unit limit set
mc.On("SimulateTx", mock.Anything, signed(0, true, computeUnitLimit), mock.Anything).Run(func(mock.Arguments) {
wg.Done()
}).Return(&rpc.SimulateTransactionResult{UnitsConsumed: &computeUnitConsumed}, nil).Once()
Expand Down Expand Up @@ -982,11 +984,13 @@ func TestTxm_compute_unit_limit_estimation(t *testing.T) {

t.Run("simulation_returns_error", func(t *testing.T) {
// Test tx is not discarded due to confirm timeout and tracked to finalization
tx, signed := getTx(t, 1, mkey)
tx, _ := getTx(t, 1, mkey)
// add signature and compute unit limit to tx for simulation (excludes compute unit price)
simulateTx := addSigAndLimitToTx(t, mkey, solana.PublicKey{}, *tx, MaxComputeUnitLimit)
sig := randomSignature(t)

mc.On("SendTx", mock.Anything, signed(0, true, fees.ComputeUnitLimit(0))).Return(sig, nil).Panic("SendTx should never be called").Maybe()
mc.On("SimulateTx", mock.Anything, tx, mock.Anything).Return(&rpc.SimulateTransactionResult{Err: errors.New("tx err")}, nil).Once()
mc.On("SendTx", mock.Anything, mock.Anything).Return(sig, nil).Panic("SendTx should never be called").Maybe()
// First simulation before broadcast with max compute unit limit
mc.On("SimulateTx", mock.Anything, simulateTx, mock.Anything).Return(&rpc.SimulateTransactionResult{Err: errors.New("tx err")}, nil).Once()

// tx should NOT be able to queue
assert.Error(t, txm.Enqueue(ctx, t.Name(), tx, nil))
Expand Down Expand Up @@ -1072,3 +1076,17 @@ func TestTxm_Enqueue(t *testing.T) {
})
}
}

func addSigAndLimitToTx(t *testing.T, keystore SimpleKeystore, pubkey solana.PublicKey, tx solana.Transaction, limit fees.ComputeUnitLimit) *solana.Transaction {
txCopy := tx
// sign tx
txMsg, err := tx.Message.MarshalBinary()
require.NoError(t, err)
sigBytes, err := keystore.Sign(context.Background(), pubkey.String(), txMsg)
require.NoError(t, err)
var sig [64]byte
copy(sig[:], sigBytes)
txCopy.Signatures = append(txCopy.Signatures, sig)
require.NoError(t, fees.SetComputeUnitLimit(&txCopy, limit))
return &txCopy
}
43 changes: 41 additions & 2 deletions pkg/solana/txm/txm_unit_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (
solanaClient "github.com/smartcontractkit/chainlink-solana/pkg/solana/client"
clientmocks "github.com/smartcontractkit/chainlink-solana/pkg/solana/client/mocks"
"github.com/smartcontractkit/chainlink-solana/pkg/solana/config"
"github.com/smartcontractkit/chainlink-solana/pkg/solana/fees"
solanatxm "github.com/smartcontractkit/chainlink-solana/pkg/solana/txm"
keyMocks "github.com/smartcontractkit/chainlink-solana/pkg/solana/txm/mocks"

Expand All @@ -25,11 +26,11 @@ import (

func TestTxm_EstimateComputeUnitLimit(t *testing.T) {
t.Parallel()

ctx := tests.Context(t)

// setup mock keystore
mkey := keyMocks.NewSimpleKeystore(t)
mkey.On("Sign", mock.Anything, mock.Anything, mock.Anything).Return([]byte{}, nil)

// setup key
key, err := solana.NewRandomPrivateKey()
Expand Down Expand Up @@ -57,7 +58,17 @@ func TestTxm_EstimateComputeUnitLimit(t *testing.T) {
Blockhash: solana.Hash{},
},
}, nil).Once()
client.On("SimulateTx", mock.Anything, mock.Anything, mock.Anything).Return(&rpc.SimulateTransactionResult{
client.On("SimulateTx", mock.Anything, mock.IsType(&solana.Transaction{}), mock.IsType(&rpc.SimulateTransactionOpts{})).Run(func(args mock.Arguments) {
// Validate max compute unit limit is set in transaction
tx := args.Get(1).(*solana.Transaction)
limit, err := fees.ParseComputeUnitLimit(tx.Message.Instructions[len(tx.Message.Instructions)-1].Data)
require.NoError(t, err)
require.Equal(t, fees.ComputeUnitLimit(solanatxm.MaxComputeUnitLimit), limit)

// Validate signature verification is enabled
opts := args.Get(2).(*rpc.SimulateTransactionOpts)
require.True(t, opts.SigVerify)
}).Return(&rpc.SimulateTransactionResult{
Err: nil,
UnitsConsumed: &usedCompute,
}, nil).Once()
Expand Down Expand Up @@ -111,6 +122,34 @@ func TestTxm_EstimateComputeUnitLimit(t *testing.T) {
require.NoError(t, err)
require.Equal(t, uint32(0), computeUnitLimit)
})

t.Run("simulation returns max compute unit limit if adding buffer exceeds it", func(t *testing.T) {
usedCompute := uint64(1_400_000)
client.On("LatestBlockhash", mock.Anything).Return(&rpc.GetLatestBlockhashResult{
Value: &rpc.LatestBlockhashResult{
LastValidBlockHeight: 100,
Blockhash: solana.Hash{},
},
}, nil).Once()
client.On("SimulateTx", mock.Anything, mock.IsType(&solana.Transaction{}), mock.IsType(&rpc.SimulateTransactionOpts{})).Run(func(args mock.Arguments) {
// Validate max compute unit limit is set in transaction
tx := args.Get(1).(*solana.Transaction)
limit, err := fees.ParseComputeUnitLimit(tx.Message.Instructions[len(tx.Message.Instructions)-1].Data)
require.NoError(t, err)
require.Equal(t, fees.ComputeUnitLimit(solanatxm.MaxComputeUnitLimit), limit)

// Validate signature verification is enabled
opts := args.Get(2).(*rpc.SimulateTransactionOpts)
require.True(t, opts.SigVerify)
}).Return(&rpc.SimulateTransactionResult{
Err: nil,
UnitsConsumed: &usedCompute,
}, nil).Once()
tx := createTx(t, client, pubKey, pubKey, pubKeyReceiver, solana.LAMPORTS_PER_SOL)
computeUnitLimit, err := txm.EstimateComputeUnitLimit(ctx, tx)
require.NoError(t, err)
require.Equal(t, uint32(1_400_000), computeUnitLimit)
})
}

func createTx(t *testing.T, client solanaClient.ReaderWriter, signer solana.PublicKey, sender solana.PublicKey, receiver solana.PublicKey, amt uint64) *solana.Transaction {
Expand Down

0 comments on commit e2db20a

Please sign in to comment.