feat: auction index (#12)

Co-authored-by: David Terpay <35130517+davidterpay@users.noreply.github.com>

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

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[:]})
+}

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())
+}

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...),
+		globalIndex:  sdkmempool.NewPriorityMempool(opts...),
-		txIndex:     make(map[string]*WrappedTx),
+		auctionIndex: NewAuctionBidList(),
-		txEncoder:   txEncoder,
+		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{
+	if err := am.globalIndex.Insert(ctx, tx); err != nil {
-		Tx:   tx,
+		return fmt.Errorf("failed to insert tx into global index: %w", err)
-		hash: hash,
 	}
 
 	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
+		am.auctionIndex.Insert(NewWrappedBidTx(tx, hash, msg.GetBid()))
-		// pointer.
 	}
 
-	if err := am.globalIndex.Insert(ctx, wrappedTx); err != nil {
+	am.txIndex[hashStr] = tx
-		return fmt.Errorf("failed to insert tx into global index: %w", err)
-	}
-
-	am.txIndex[hashStr] = wrappedTx
 
 	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
+}

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
+}

mempool/tx.go

@@ -7,23 +7,25 @@ import (
 	auctiontypes "github.com/skip-mev/pob/x/auction/types"
 )
 
-type (
+// WrappedBidTx defines a wrapper around an sdk.Tx that contains a single
-	// WrappedTx defines a wrapper around an sdk.Tx with additional metadata.
+// MsgAuctionBid message with additional metadata.
-	WrappedTx struct {
+type WrappedBidTx struct {
-		sdk.Tx
+	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
+func (wbtx *WrappedBidTx) GetHash() [32]byte { return wbtx.hash }
-	// MsgAuctionBid message with additional metadata.
+func (wbtx *WrappedBidTx) GetBid() sdk.Coins { return wbtx.bid }
-	WrappedBidTx struct {
-		sdk.Tx
-
-		hash [32]byte
-		bid  sdk.Coins
-	}
-)
 
 // 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
+		// a transaction with at at least one MsgAuctionBid message
-		// message.
 		return nil, errors.New("invalid MsgAuctionBid transaction")
 	}
 }

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
+}