refactor: gas estimation APIs (#546)

* refactor: gas estimation APIs * changelog * lint
osmosis-labs · Nov 3, 2024 · 662aac4 · 662aac4
1 parent 2a46eeb
commit 662aac4
Show file tree

Hide file tree

Showing 17 changed files with 874 additions and 350 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,6 +35,11 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 # Changelog
 
+## Unreleased
+
+- #526 - Refactor gas estimation APIs
+- #524 - Claimbot
+
 ## v26.1.0
 
 e42b32bc SQS-412 | Active Orders Query: SSE (#518)

diff --git a/domain/cosmos/tx/gas.go b/domain/cosmos/tx/gas.go
diff --git a/domain/cosmos/tx/msg_simulator.go b/domain/cosmos/tx/msg_simulator.go
@@ -0,0 +1,162 @@
+package tx
+
+import (
+	"context"
+
+	cosmosclient "github.com/cosmos/cosmos-sdk/client"
+	txclient "github.com/cosmos/cosmos-sdk/client/tx"
+	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	txtypes "github.com/cosmos/cosmos-sdk/types/tx"
+	signingtypes "github.com/cosmos/cosmos-sdk/types/tx/signing"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/osmosis-labs/osmosis/v26/app/params"
+	txfeestypes "github.com/osmosis-labs/osmosis/v26/x/txfees/types"
+	"github.com/osmosis-labs/sqs/domain/keyring"
+
+	gogogrpc "github.com/cosmos/gogoproto/grpc"
+)
+
+// MsgSimulator is an interface for calculating gas for a transaction.
+type MsgSimulator interface {
+	BuildTx(
+		ctx context.Context,
+		keyring keyring.Keyring,
+		txfeesClient txfeestypes.QueryClient,
+		encodingConfig params.EncodingConfig,
+		account *authtypes.BaseAccount,
+		chainID string,
+		msg ...sdk.Msg,
+	) (cosmosclient.TxBuilder, error)
+
+	// SimulateMsgs simulates the execution of the given messages and returns the simulation response,
+	// adjusted gas used, and any error encountered. It uses the provided gRPC client, encoding config,
+	// account details, and chain ID to create a transaction factory for the simulation.
+	SimulateMsgs(
+		encodingConfig cosmosclient.TxConfig,
+		account *authtypes.BaseAccount,
+		chainID string,
+		msgs []sdk.Msg,
+	) (*txtypes.SimulateResponse, uint64, error)
+}
+
+// NewGasCalculator creates a new GasCalculator instance.
+func NewGasCalculator(clientCtx gogogrpc.ClientConn, calculateGas CalculateGasFn) MsgSimulator {
+	return &txGasCalulator{
+		clientCtx:    clientCtx,
+		calculateGas: calculateGas,
+	}
+}
+
+// CalculateGasFn is a function type that calculates the gas for a transaction.
+type CalculateGasFn func(clientCtx gogogrpc.ClientConn, txf txclient.Factory, msgs ...sdk.Msg) (*txtypes.SimulateResponse, uint64, error)
+
+// txGasCalulator is a GasCalculator implementation that uses simulated transactions to calculate gas.
+type txGasCalulator struct {
+	clientCtx    gogogrpc.ClientConn
+	calculateGas CalculateGasFn
+}
+
+// BuildTx constructs a transaction using the provided parameters and messages.
+// Returns a TxBuilder and any error encountered.
+func (c *txGasCalulator) BuildTx(
+	ctx context.Context,
+	keyring keyring.Keyring,
+	txfeesClient txfeestypes.QueryClient,
+	encodingConfig params.EncodingConfig,
+	account *authtypes.BaseAccount,
+	chainID string,
+	msg ...sdk.Msg,
+) (cosmosclient.TxBuilder, error) {
+	key := keyring.GetKey()
+	privKey := &secp256k1.PrivKey{Key: key.Bytes()}
+
+	// Create and sign the transaction
+	txBuilder := encodingConfig.TxConfig.NewTxBuilder()
+
+	err := txBuilder.SetMsgs(msg...)
+	if err != nil {
+		return nil, err
+	}
+
+	_, gas, err := c.SimulateMsgs(
+		encodingConfig.TxConfig,
+		account,
+		chainID,
+		msg,
+	)
+	if err != nil {
+		return nil, err
+	}
+	txBuilder.SetGasLimit(gas)
+
+	feecoin, err := CalculateFeeCoin(ctx, txfeesClient, gas)
+	if err != nil {
+		return nil, err
+	}
+
+	txBuilder.SetFeeAmount(sdk.NewCoins(feecoin))
+
+	sigV2 := BuildSignatures(privKey.PubKey(), nil, account.Sequence)
+	err = txBuilder.SetSignatures(sigV2)
+	if err != nil {
+		return nil, err
+	}
+
+	signerData := BuildSignerData(chainID, account.AccountNumber, account.Sequence)
+
+	signed, err := txclient.SignWithPrivKey(
+		ctx,
+		signingtypes.SignMode_SIGN_MODE_DIRECT, signerData,
+		txBuilder, privKey, encodingConfig.TxConfig, account.Sequence)
+	if err != nil {
+		return nil, err
+	}
+
+	err = txBuilder.SetSignatures(signed)
+	if err != nil {
+		return nil, err
+	}
+
+	return txBuilder, nil
+}
+
+// SimulateMsgs implements MsgSimulator.
+func (c *txGasCalulator) SimulateMsgs(encodingConfig cosmosclient.TxConfig, account *authtypes.BaseAccount, chainID string, msgs []sdk.Msg) (*txtypes.SimulateResponse, uint64, error) {
+	txFactory := txclient.Factory{}
+	txFactory = txFactory.WithTxConfig(encodingConfig)
+	txFactory = txFactory.WithAccountNumber(account.AccountNumber)
+	txFactory = txFactory.WithSequence(account.Sequence)
+	txFactory = txFactory.WithChainID(chainID)
+	txFactory = txFactory.WithGasAdjustment(1.05)
+
+	// Estimate transaction
+	gasResult, adjustedGasUsed, err := c.calculateGas(
+		c.clientCtx,
+		txFactory,
+		msgs...,
+	)
+	if err != nil {
+		return nil, adjustedGasUsed, err
+	}
+
+	return gasResult, adjustedGasUsed, nil
+}
+
+// CalculateGas calculates the gas required for a transaction using the provided transaction factory and messages.
+func CalculateGas(
+	clientCtx gogogrpc.ClientConn,
+	txf txclient.Factory,
+	msgs ...sdk.Msg,
+) (*txtypes.SimulateResponse, uint64, error) {
+	gasResult, adjustedGasUsed, err := txclient.CalculateGas(
+		clientCtx,
+		txf,
+		msgs...,
+	)
+	if err != nil {
+		return nil, adjustedGasUsed, err
+	}
+
+	return gasResult, adjustedGasUsed, nil
+}
diff --git a/domain/cosmos/tx/msg_simulator_test.go b/domain/cosmos/tx/msg_simulator_test.go
@@ -0,0 +1,160 @@
+package tx_test
+
+import (
+	"context"
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	txtypes "github.com/cosmos/cosmos-sdk/types/tx"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/osmosis-labs/sqs/domain/cosmos/tx"
+	"github.com/osmosis-labs/sqs/domain/mocks"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSimulateMsgs(t *testing.T) {
+	tests := []struct {
+		name                     string
+		account                  *authtypes.BaseAccount
+		chainID                  string
+		msgs                     []sdk.Msg
+		setupMocks               func(calculator mocks.GetCalculateGasMock) tx.CalculateGasFn
+		expectedSimulateResponse *txtypes.SimulateResponse
+		expectedGas              uint64
+		expectedError            error
+	}{
+		{
+			name:    "Successful simulation",
+			account: &authtypes.BaseAccount{AccountNumber: 1, Sequence: 1},
+			chainID: "test-chain",
+			msgs:    []sdk.Msg{newMsg("sender", "contract", `{}`)},
+			setupMocks: func(calculator mocks.GetCalculateGasMock) tx.CalculateGasFn {
+				return calculator(&txtypes.SimulateResponse{GasInfo: &sdk.GasInfo{GasUsed: 100000}}, 50, nil)
+			},
+			expectedSimulateResponse: &txtypes.SimulateResponse{GasInfo: &sdk.GasInfo{GasUsed: 100000}},
+			expectedGas:              50,
+			expectedError:            nil,
+		},
+		{
+			name:    "Simulation error",
+			account: &authtypes.BaseAccount{AccountNumber: 2, Sequence: 2},
+			chainID: "test-chain",
+			msgs:    []sdk.Msg{},
+			setupMocks: func(calculator mocks.GetCalculateGasMock) tx.CalculateGasFn {
+				return calculator(&txtypes.SimulateResponse{}, 3, assert.AnError)
+			},
+			expectedSimulateResponse: nil,
+			expectedGas:              3,
+			expectedError:            assert.AnError,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Setup the mock
+			calculateGasFnMock := tt.setupMocks(mocks.DefaultGetCalculateGasMock)
+
+			// Create the gas calculator
+			gasCalculator := tx.NewGasCalculator(nil, calculateGasFnMock)
+
+			// Call the function
+			result, gas, err := gasCalculator.SimulateMsgs(
+				encodingConfig.TxConfig,
+				tt.account,
+				tt.chainID,
+				tt.msgs,
+			)
+
+			// Assert the results
+			assert.Equal(t, tt.expectedSimulateResponse, result)
+			assert.Equal(t, tt.expectedGas, gas)
+			if tt.expectedError != nil {
+				assert.Error(t, err)
+				assert.Equal(t, tt.expectedError, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestBuildTx(t *testing.T) {
+	testCases := []struct {
+		name          string
+		setupMocks    func(calculator mocks.GetCalculateGasMock, txFeesClient *mocks.TxFeesQueryClient, keyring *mocks.Keyring) tx.CalculateGasFn
+		account       *authtypes.BaseAccount
+		chainID       string
+		msgs          []sdk.Msg
+		expectedJSON  []byte
+		expectedError bool
+	}{
+		{
+			name: "Valid transaction",
+			setupMocks: func(calculator mocks.GetCalculateGasMock, txFeesClient *mocks.TxFeesQueryClient, keyring *mocks.Keyring) tx.CalculateGasFn {
+				keyring.WithGetKey("6cf5103c60c939a5f38e383b52239c5296c968579eec1c68a47d70fbf1d19159")
+				txFeesClient.WithBaseDenom("eth", nil)
+				txFeesClient.WithGetEipBaseFee("0.1", nil)
+
+				return calculator(&txtypes.SimulateResponse{GasInfo: &sdk.GasInfo{GasUsed: 100000}}, 50, nil)
+			},
+			account: &authtypes.BaseAccount{
+				Sequence:      13,
+				AccountNumber: 1,
+			},
+			chainID:       "test-chain",
+			msgs:          []sdk.Msg{newMsg("sender", "contract", `{"payload": "hello contract"}`)},
+			expectedJSON:  []byte(`{"body":{"messages":[{"@type":"/cosmwasm.wasm.v1.MsgExecuteContract","sender":"sender","contract":"contract","msg":{"payload":"hello contract"},"funds":[]}],"memo":"","timeout_height":"0","extension_options":[],"non_critical_extension_options":[]},"auth_info":{"signer_infos":[{"public_key":{"@type":"/cosmos.crypto.secp256k1.PubKey","key":"A+9dbfKKCHgfmiV2XUWelqidYzZhHR+KtNMvcSzWjdPQ"},"mode_info":{"single":{"mode":"SIGN_MODE_DIRECT"}},"sequence":"13"}],"fee":{"amount":[{"denom":"eth","amount":"5"}],"gas_limit":"50","payer":"","granter":""},"tip":null},"signatures":["aRlC8F2MnDA50tNNTJUk7zPvH/xc5c3Av+yaGQEiU0l0AXJxUdzOUxWHiC74D9ltvbsk0HzWbb+2uetCjdQdfA=="]}`),
+			expectedError: false,
+		},
+		{
+			name: "Error building transaction",
+			setupMocks: func(calculator mocks.GetCalculateGasMock, txFeesClient *mocks.TxFeesQueryClient, keyring *mocks.Keyring) tx.CalculateGasFn {
+				keyring.WithGetKey("6cf5103c60c939a5f38e383b52239c5296c968579eec1c68a47d70fbf1d19159")
+
+				return calculator(&txtypes.SimulateResponse{}, 50, assert.AnError)
+			},
+			account: &authtypes.BaseAccount{
+				Sequence:      8,
+				AccountNumber: 51,
+			},
+			expectedError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			txFeesClient := mocks.TxFeesQueryClient{}
+			keyring := mocks.Keyring{}
+
+			// Setup the mock
+			calculateGasFnMock := tc.setupMocks(mocks.DefaultGetCalculateGasMock, &txFeesClient, &keyring)
+
+			// Create the gas calculator
+			msgSimulator := tx.NewGasCalculator(nil, calculateGasFnMock)
+
+			txBuilder, err := msgSimulator.BuildTx(
+				context.Background(),
+				&keyring,
+				&txFeesClient,
+				encodingConfig,
+				tc.account,
+				tc.chainID,
+				tc.msgs...,
+			)
+
+			if tc.expectedError {
+				assert.Error(t, err)
+				assert.Nil(t, txBuilder)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, txBuilder)
+
+				txJSONBytes, err := encodingConfig.TxConfig.TxJSONEncoder()(txBuilder.GetTx())
+				assert.NoError(t, err)
+
+				// Add more specific assertions here based on the expected output
+				assert.Equal(t, string(tc.expectedJSON), string(txJSONBytes))
+			}
+		})
+	}
+}