From afd56f2bbfa641364798e439b2f90d81b481749e Mon Sep 17 00:00:00 2001 From: Aleksandr Bezobchuk Date: Mon, 6 Mar 2023 12:40:57 -0500 Subject: [PATCH] feat: auction index (#12) Co-authored-by: David Terpay <35130517+davidterpay@users.noreply.github.com> --- go.mod | 4 +- mempool/bid_list.go | 58 ++++++++++++++ mempool/bid_list_test.go | 45 +++++++++++ mempool/mempool.go | 101 ++++++++++++++++++------ mempool/mempool_test.go | 162 +++++++++++++++++++++++++++++++++++++++ mempool/tx.go | 33 ++++---- mempool/utils_test.go | 68 ++++++++++++++++ 7 files changed, 430 insertions(+), 41 deletions(-) create mode 100644 mempool/bid_list.go create mode 100644 mempool/bid_list_test.go create mode 100644 mempool/mempool_test.go create mode 100644 mempool/utils_test.go diff --git a/go.mod b/go.mod index 53e4a7e..4703adb 100644 --- a/go.mod +++ b/go.mod @@ -5,12 +5,14 @@ go 1.20 require ( cosmossdk.io/api v0.3.1 cosmossdk.io/errors v1.0.0-beta.7 + github.com/cometbft/cometbft v0.37.0-alpha.3 github.com/cosmos/cosmos-proto v1.0.0-beta.2 github.com/cosmos/cosmos-sdk v0.47.0-rc2.0.20230228000043-54240ec9ab19 github.com/cosmos/gogoproto v1.4.4 github.com/golang/protobuf v1.5.2 github.com/gorilla/mux v1.8.0 github.com/grpc-ecosystem/grpc-gateway v1.16.0 + github.com/huandu/skiplist v1.2.0 github.com/stretchr/testify v1.8.1 github.com/spf13/cobra v1.6.1 google.golang.org/genproto v0.0.0-20230216225411-c8e22ba71e44 @@ -32,7 +34,6 @@ require ( github.com/btcsuite/btcd/btcec/v2 v2.3.2 // indirect github.com/cespare/xxhash v1.1.0 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect - github.com/cometbft/cometbft v0.37.0-alpha.3 // indirect github.com/cometbft/cometbft-db v0.7.0 // indirect github.com/confio/ics23/go v0.9.0 // indirect github.com/cosmos/btcutil v1.0.5 // indirect @@ -65,7 +66,6 @@ require ( github.com/hashicorp/golang-lru v0.5.5-0.20210104140557-80c98217689d // indirect github.com/hashicorp/hcl v1.0.0 // indirect github.com/hdevalence/ed25519consensus v0.1.0 // indirect - github.com/huandu/skiplist v1.2.0 // indirect github.com/inconshreveable/mousetrap v1.0.1 // indirect github.com/jmhodges/levigo v1.0.0 // indirect github.com/klauspost/compress v1.15.15 // indirect diff --git a/mempool/bid_list.go b/mempool/bid_list.go new file mode 100644 index 0000000..9d66f93 --- /dev/null +++ b/mempool/bid_list.go @@ -0,0 +1,58 @@ +package mempool + +import ( + sdk "github.com/cosmos/cosmos-sdk/types" + "github.com/huandu/skiplist" +) + +type ( + // AuctionBidList defines a list of WrappedBidTx objects, sorted by their bids. + + AuctionBidList struct { + list *skiplist.SkipList + } + + auctionBidListKey struct { + bid sdk.Coins + hash []byte + } +) + +func NewAuctionBidList() *AuctionBidList { + return &AuctionBidList{ + list: skiplist.New(skiplist.GreaterThanFunc(func(lhs, rhs any) int { + bidA := lhs.(auctionBidListKey) + bidB := rhs.(auctionBidListKey) + + switch { + case bidA.bid.IsAllGT(bidB.bid): + return 1 + + case bidA.bid.IsAllLT(bidB.bid): + return -1 + + default: + // in case of a tie in bid, sort by hash + return skiplist.ByteAsc.Compare(bidA.hash, bidB.hash) + } + })), + } +} + +// TopBid returns the WrappedBidTx with the highest bid. +func (abl *AuctionBidList) TopBid() *WrappedBidTx { + n := abl.list.Back() + if n == nil { + return nil + } + + return n.Value.(*WrappedBidTx) +} + +func (abl *AuctionBidList) Insert(wBidTx *WrappedBidTx) { + abl.list.Set(auctionBidListKey{bid: wBidTx.bid, hash: wBidTx.hash[:]}, wBidTx) +} + +func (abl *AuctionBidList) Remove(wBidTx *WrappedBidTx) { + abl.list.Remove(auctionBidListKey{bid: wBidTx.bid, hash: wBidTx.hash[:]}) +} diff --git a/mempool/bid_list_test.go b/mempool/bid_list_test.go new file mode 100644 index 0000000..2951275 --- /dev/null +++ b/mempool/bid_list_test.go @@ -0,0 +1,45 @@ +package mempool_test + +import ( + "math/rand" + "testing" + "time" + + sdk "github.com/cosmos/cosmos-sdk/types" + "github.com/skip-mev/pob/mempool" + "github.com/stretchr/testify/require" +) + +var emptyHash = [32]byte{} + +func TestAuctionBidList(t *testing.T) { + abl := mempool.NewAuctionBidList() + + require.Nil(t, abl.TopBid()) + + // insert a bid which should be the head and tail + bid1 := sdk.NewCoins(sdk.NewInt64Coin("foo", 100)) + abl.Insert(mempool.NewWrappedBidTx(nil, emptyHash, bid1)) + require.Equal(t, bid1, abl.TopBid().GetBid()) + + // insert 500 random bids between [100, 1000) + var currTopBid sdk.Coins + rng := rand.New(rand.NewSource(time.Now().UnixNano())) + for i := 0; i < 500; i++ { + randomBid := rng.Int63n(1000-100) + 100 + + bid := sdk.NewCoins(sdk.NewInt64Coin("foo", randomBid)) + abl.Insert(mempool.NewWrappedBidTx(nil, emptyHash, bid)) + + currTopBid = abl.TopBid().GetBid() + } + + // insert a bid which should be the new tail, thus the highest bid + bid2 := sdk.NewCoins(sdk.NewInt64Coin("foo", 1000)) + abl.Insert(mempool.NewWrappedBidTx(nil, emptyHash, bid2)) + require.Equal(t, bid2, abl.TopBid().GetBid()) + + // remove the top bid and ensure the new top bid is the previous top bid + abl.Remove(mempool.NewWrappedBidTx(nil, emptyHash, bid2)) + require.Equal(t, currTopBid, abl.TopBid().GetBid()) +} diff --git a/mempool/mempool.go b/mempool/mempool.go index f565daf..1a248ce 100644 --- a/mempool/mempool.go +++ b/mempool/mempool.go @@ -18,43 +18,42 @@ var _ sdkmempool.Mempool = (*AuctionMempool)(nil) // and indexing auction bids. type AuctionMempool struct { // globalIndex defines the index of all transactions in the mempool. It uses - // the SDK's builtin PriorityNonceMempool. Once a bid if selected for top-of-block, + // the SDK's builtin PriorityNonceMempool. Once a bid is selected for top-of-block, // all subsequent transactions in the mempool will be selected from this index. - globalIndex sdkmempool.PriorityNonceMempool + globalIndex *sdkmempool.PriorityNonceMempool + + // auctionIndex defines an index of auction bids. + auctionIndex *AuctionBidList // txIndex defines an index of all transactions in the mempool by hash. - txIndex map[string]*WrappedTx + txIndex map[string]sdk.Tx // txEncoder defines the sdk.Tx encoder that allows us to encode transactions // and construct their hashes. txEncoder sdk.TxEncoder - - // auctionIndex *heap.Heap[PriorityTx] } func NewAuctionMempool(txEncoder sdk.TxEncoder, opts ...sdkmempool.PriorityNonceMempoolOption) *AuctionMempool { return &AuctionMempool{ - globalIndex: *sdkmempool.NewPriorityMempool(opts...), - txIndex: make(map[string]*WrappedTx), - txEncoder: txEncoder, + globalIndex: sdkmempool.NewPriorityMempool(opts...), + auctionIndex: NewAuctionBidList(), + txIndex: make(map[string]sdk.Tx), + txEncoder: txEncoder, } } func (am *AuctionMempool) Insert(ctx context.Context, tx sdk.Tx) error { - bz, err := am.txEncoder(tx) + hash, hashStr, err := am.getTxHash(tx) if err != nil { - return fmt.Errorf("failed to encode tx: %w", err) + return err } - hash := sha256.Sum256(bz) - hashStr := base64.StdEncoding.EncodeToString(hash[:]) if _, ok := am.txIndex[hashStr]; ok { return fmt.Errorf("tx already exists: %s", hashStr) } - wrappedTx := &WrappedTx{ - Tx: tx, - hash: hash, + if err := am.globalIndex.Insert(ctx, tx); err != nil { + return fmt.Errorf("failed to insert tx into global index: %w", err) } msg, err := GetMsgAuctionBidFromTx(tx) @@ -63,21 +62,65 @@ func (am *AuctionMempool) Insert(ctx context.Context, tx sdk.Tx) error { } if msg != nil { - // TODO: Insert into auctionIndex and update wrappedTx to reflect the index - // pointer. + am.auctionIndex.Insert(NewWrappedBidTx(tx, hash, msg.GetBid())) } - if err := am.globalIndex.Insert(ctx, wrappedTx); err != nil { - return fmt.Errorf("failed to insert tx into global index: %w", err) - } - - am.txIndex[hashStr] = wrappedTx + am.txIndex[hashStr] = tx return nil } func (am *AuctionMempool) Remove(tx sdk.Tx) error { - panic("not implemented") + hash, hashStr, err := am.getTxHash(tx) + if err != nil { + return err + } + + // 1. Remove the tx from the global index + if err := am.globalIndex.Remove(tx); err != nil { + return fmt.Errorf("failed to remove tx from global index: %w", err) + } + + // 2. Remove from the transaction index + delete(am.txIndex, hashStr) + + msg, err := GetMsgAuctionBidFromTx(tx) + if err != nil { + return err + } + + // 3. Remove the bid from the auction index (if applicable). In addition, we + // remove all referenced transactions from the global and transaction indices. + if msg != nil { + am.auctionIndex.Remove(NewWrappedBidTx(tx, hash, msg.GetBid())) + + for _, refTxRaw := range msg.GetTransactions() { + refHash := sha256.Sum256(refTxRaw) + refHashStr := base64.StdEncoding.EncodeToString(refHash[:]) + + // check if we have the referenced transaction first + if refTx, ok := am.txIndex[refHashStr]; ok { + if err := am.globalIndex.Remove(refTx); err != nil { + return fmt.Errorf("failed to remove bid referenced tx from global index: %w", err) + } + } + + delete(am.txIndex, refHashStr) + } + } + + return nil +} + +// SelectTopAuctionBidTx returns the top auction bid tx in the mempool if one +// exists. +func (am *AuctionMempool) SelectTopAuctionBidTx() sdk.Tx { + wBidTx := am.auctionIndex.TopBid() + if wBidTx == nil { + return nil + } + + return wBidTx.Tx } func (am *AuctionMempool) Select(ctx context.Context, txs [][]byte) sdkmempool.Iterator { @@ -87,3 +130,15 @@ func (am *AuctionMempool) Select(ctx context.Context, txs [][]byte) sdkmempool.I func (am *AuctionMempool) CountTx() int { return am.globalIndex.CountTx() } + +func (am *AuctionMempool) getTxHash(tx sdk.Tx) ([32]byte, string, error) { + bz, err := am.txEncoder(tx) + if err != nil { + return [32]byte{}, "", fmt.Errorf("failed to encode tx: %w", err) + } + + hash := sha256.Sum256(bz) + hashStr := base64.StdEncoding.EncodeToString(hash[:]) + + return hash, hashStr, nil +} diff --git a/mempool/mempool_test.go b/mempool/mempool_test.go new file mode 100644 index 0000000..0842b0e --- /dev/null +++ b/mempool/mempool_test.go @@ -0,0 +1,162 @@ +package mempool_test + +import ( + "math/rand" + "testing" + "time" + + "github.com/cometbft/cometbft/libs/log" + cmtproto "github.com/cometbft/cometbft/proto/tendermint/types" + "github.com/cosmos/cosmos-sdk/client" + sdk "github.com/cosmos/cosmos-sdk/types" + "github.com/cosmos/cosmos-sdk/types/tx/signing" + banktypes "github.com/cosmos/cosmos-sdk/x/bank/types" + "github.com/skip-mev/pob/mempool" + auctiontypes "github.com/skip-mev/pob/x/auction/types" + "github.com/stretchr/testify/require" +) + +func TestAuctionMempool(t *testing.T) { + encCfg := createTestEncodingConfig() + amp := mempool.NewAuctionMempool(encCfg.TxConfig.TxEncoder()) + ctx := sdk.NewContext(nil, cmtproto.Header{}, false, log.NewNopLogger()) + rng := rand.New(rand.NewSource(time.Now().Unix())) + accounts := RandomAccounts(rng, 5) + + accNonces := map[string]uint64{} + for _, acc := range accounts { + accNonces[acc.Address.String()] = 0 + } + + // insert a bunch of normal txs + for i := 0; i < 1000; i++ { + p := rng.Int63n(500-1) + 1 + j := rng.Intn(len(accounts)) + acc := accounts[j] + txBuilder := encCfg.TxConfig.NewTxBuilder() + + msgs := []sdk.Msg{ + &banktypes.MsgSend{ + FromAddress: acc.Address.String(), + ToAddress: acc.Address.String(), + }, + } + err := txBuilder.SetMsgs(msgs...) + require.NoError(t, err) + + sigV2 := signing.SignatureV2{ + PubKey: acc.PrivKey.PubKey(), + Data: &signing.SingleSignatureData{ + SignMode: encCfg.TxConfig.SignModeHandler().DefaultMode(), + Signature: nil, + }, + Sequence: accNonces[acc.Address.String()], + } + err = txBuilder.SetSignatures(sigV2) + require.NoError(t, err) + + accNonces[acc.Address.String()]++ + + require.NoError(t, amp.Insert(ctx.WithPriority(p), txBuilder.GetTx())) + } + + require.Nil(t, amp.SelectTopAuctionBidTx()) + + // insert bid transactions + var highestBid sdk.Coins + biddingAccs := RandomAccounts(rng, 100) + + for _, acc := range biddingAccs { + p := rng.Int63n(500-1) + 1 + txBuilder := encCfg.TxConfig.NewTxBuilder() + + // keep track of highest bid + bid := sdk.NewCoins(sdk.NewInt64Coin("foo", p)) + if bid.IsAllGT(highestBid) { + highestBid = bid + } + + bidMsg, err := createMsgAuctionBid(encCfg.TxConfig, acc, bid) + require.NoError(t, err) + + msgs := []sdk.Msg{bidMsg} + err = txBuilder.SetMsgs(msgs...) + require.NoError(t, err) + + sigV2 := signing.SignatureV2{ + PubKey: acc.PrivKey.PubKey(), + Data: &signing.SingleSignatureData{ + SignMode: encCfg.TxConfig.SignModeHandler().DefaultMode(), + Signature: nil, + }, + Sequence: 0, + } + err = txBuilder.SetSignatures(sigV2) + require.NoError(t, err) + + require.NoError(t, amp.Insert(ctx.WithPriority(p), txBuilder.GetTx())) + + // Insert the referenced txs just to ensure that they are removed from the + // mempool in cases where they exist. + for _, refRawTx := range bidMsg.GetTransactions() { + refTx, err := encCfg.TxConfig.TxDecoder()(refRawTx) + require.NoError(t, err) + require.NoError(t, amp.Insert(ctx.WithPriority(0), refTx)) + } + } + + expectedCount := 1000 + 100 + 200 + require.Equal(t, expectedCount, amp.CountTx()) + + // select the top bid and misc txs + bidTx := amp.SelectTopAuctionBidTx() + require.Len(t, bidTx.GetMsgs(), 1) + require.Equal(t, highestBid, bidTx.GetMsgs()[0].(*auctiontypes.MsgAuctionBid).Bid) + + // remove bid tx, which should also removed the referenced txs + require.NoError(t, amp.Remove(bidTx)) + require.Equal(t, expectedCount-3, amp.CountTx()) +} + +func createMsgAuctionBid(txCfg client.TxConfig, bidder Account, bid sdk.Coins) (*auctiontypes.MsgAuctionBid, error) { + bidMsg := &auctiontypes.MsgAuctionBid{ + Bidder: bidder.Address.String(), + Bid: bid, + Transactions: make([][]byte, 2), + } + + for i := 0; i < 2; i++ { + txBuilder := txCfg.NewTxBuilder() + + msgs := []sdk.Msg{ + &banktypes.MsgSend{ + FromAddress: bidder.Address.String(), + ToAddress: bidder.Address.String(), + }, + } + if err := txBuilder.SetMsgs(msgs...); err != nil { + return nil, err + } + + sigV2 := signing.SignatureV2{ + PubKey: bidder.PrivKey.PubKey(), + Data: &signing.SingleSignatureData{ + SignMode: txCfg.SignModeHandler().DefaultMode(), + Signature: nil, + }, + Sequence: uint64(i + 1), + } + if err := txBuilder.SetSignatures(sigV2); err != nil { + return nil, err + } + + bz, err := txCfg.TxEncoder()(txBuilder.GetTx()) + if err != nil { + return nil, err + } + + bidMsg.Transactions[i] = bz + } + + return bidMsg, nil +} diff --git a/mempool/tx.go b/mempool/tx.go index e754f36..12341dc 100644 --- a/mempool/tx.go +++ b/mempool/tx.go @@ -7,23 +7,25 @@ import ( auctiontypes "github.com/skip-mev/pob/x/auction/types" ) -type ( - // WrappedTx defines a wrapper around an sdk.Tx with additional metadata. - WrappedTx struct { - sdk.Tx +// WrappedBidTx defines a wrapper around an sdk.Tx that contains a single +// MsgAuctionBid message with additional metadata. +type WrappedBidTx struct { + sdk.Tx - hash [32]byte + hash [32]byte + bid sdk.Coins +} + +func NewWrappedBidTx(tx sdk.Tx, hash [32]byte, bid sdk.Coins) *WrappedBidTx { + return &WrappedBidTx{ + Tx: tx, + hash: hash, + bid: bid, } +} - // WrappedBidTx defines a wrapper around an sdk.Tx that contains a single - // MsgAuctionBid message with additional metadata. - WrappedBidTx struct { - sdk.Tx - - hash [32]byte - bid sdk.Coins - } -) +func (wbtx *WrappedBidTx) GetHash() [32]byte { return wbtx.hash } +func (wbtx *WrappedBidTx) GetBid() sdk.Coins { return wbtx.bid } // GetMsgAuctionBidFromTx attempts to retrieve a MsgAuctionBid from an sdk.Tx if // one exists. If a MsgAuctionBid does exist and other messages are also present, @@ -47,8 +49,7 @@ func GetMsgAuctionBidFromTx(tx sdk.Tx) (*auctiontypes.MsgAuctionBid, error) { return auctionBidMsgs[0], nil default: - // A transaction with at at least one MsgAuctionBid message and some other - // message. + // a transaction with at at least one MsgAuctionBid message return nil, errors.New("invalid MsgAuctionBid transaction") } } diff --git a/mempool/utils_test.go b/mempool/utils_test.go new file mode 100644 index 0000000..cd2548d --- /dev/null +++ b/mempool/utils_test.go @@ -0,0 +1,68 @@ +package mempool_test + +import ( + "math/rand" + + "github.com/cosmos/cosmos-sdk/client" + "github.com/cosmos/cosmos-sdk/codec" + "github.com/cosmos/cosmos-sdk/codec/types" + cryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec" + "github.com/cosmos/cosmos-sdk/crypto/keys/ed25519" + "github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1" + cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types" + sdk "github.com/cosmos/cosmos-sdk/types" + "github.com/cosmos/cosmos-sdk/x/auth/tx" + banktypes "github.com/cosmos/cosmos-sdk/x/bank/types" +) + +type encodingConfig struct { + InterfaceRegistry types.InterfaceRegistry + Codec codec.Codec + TxConfig client.TxConfig + Amino *codec.LegacyAmino +} + +func createTestEncodingConfig() encodingConfig { + cdc := codec.NewLegacyAmino() + interfaceRegistry := types.NewInterfaceRegistry() + + banktypes.RegisterInterfaces(interfaceRegistry) + cryptocodec.RegisterInterfaces(interfaceRegistry) + + codec := codec.NewProtoCodec(interfaceRegistry) + + return encodingConfig{ + InterfaceRegistry: interfaceRegistry, + Codec: codec, + TxConfig: tx.NewTxConfig(codec, tx.DefaultSignModes), + Amino: cdc, + } +} + +type Account struct { + PrivKey cryptotypes.PrivKey + PubKey cryptotypes.PubKey + Address sdk.AccAddress + ConsKey cryptotypes.PrivKey +} + +func (acc Account) Equals(acc2 Account) bool { + return acc.Address.Equals(acc2.Address) +} + +func RandomAccounts(r *rand.Rand, n int) []Account { + accs := make([]Account, n) + + for i := 0; i < n; i++ { + pkSeed := make([]byte, 15) + r.Read(pkSeed) + + accs[i].PrivKey = secp256k1.GenPrivKeyFromSecret(pkSeed) + accs[i].PubKey = accs[i].PrivKey.PubKey() + accs[i].Address = sdk.AccAddress(accs[i].PubKey.Address()) + + accs[i].ConsKey = ed25519.GenPrivKeyFromSecret(pkSeed) + } + + return accs +}