github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/mempool/v1/reactor.go (about)

     1  package v1
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/gogo/protobuf/proto"
     9  
    10  	cfg "github.com/badrootd/celestia-core/config"
    11  	"github.com/badrootd/celestia-core/libs/clist"
    12  	"github.com/badrootd/celestia-core/libs/log"
    13  	cmtsync "github.com/badrootd/celestia-core/libs/sync"
    14  	"github.com/badrootd/celestia-core/mempool"
    15  	"github.com/badrootd/celestia-core/p2p"
    16  	"github.com/badrootd/celestia-core/pkg/trace"
    17  	"github.com/badrootd/celestia-core/pkg/trace/schema"
    18  	protomem "github.com/badrootd/celestia-core/proto/tendermint/mempool"
    19  	"github.com/badrootd/celestia-core/types"
    20  )
    21  
    22  // Reactor handles mempool tx broadcasting amongst peers.
    23  // It maintains a map from peer ID to counter, to prevent gossiping txs to the
    24  // peers you received it from.
    25  type Reactor struct {
    26  	p2p.BaseReactor
    27  	config      *cfg.MempoolConfig
    28  	mempool     *TxMempool
    29  	ids         *mempoolIDs
    30  	traceClient *trace.Client
    31  }
    32  
    33  type mempoolIDs struct {
    34  	mtx       cmtsync.RWMutex
    35  	peerMap   map[p2p.ID]uint16
    36  	nextID    uint16              // assumes that a node will never have over 65536 active peers
    37  	activeIDs map[uint16]struct{} // used to check if a given peerID key is used, the value doesn't matter
    38  }
    39  
    40  // ReserveForPeer searches for the next unused ID and assigns it to the
    41  // peer.
    42  func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) {
    43  	ids.mtx.Lock()
    44  	defer ids.mtx.Unlock()
    45  
    46  	curID := ids.nextPeerID()
    47  	ids.peerMap[peer.ID()] = curID
    48  	ids.activeIDs[curID] = struct{}{}
    49  }
    50  
    51  // nextPeerID returns the next unused peer ID to use.
    52  // This assumes that ids's mutex is already locked.
    53  func (ids *mempoolIDs) nextPeerID() uint16 {
    54  	if len(ids.activeIDs) == mempool.MaxActiveIDs {
    55  		panic(fmt.Sprintf("node has maximum %d active IDs and wanted to get one more", mempool.MaxActiveIDs))
    56  	}
    57  
    58  	_, idExists := ids.activeIDs[ids.nextID]
    59  	for idExists {
    60  		ids.nextID++
    61  		_, idExists = ids.activeIDs[ids.nextID]
    62  	}
    63  	curID := ids.nextID
    64  	ids.nextID++
    65  	return curID
    66  }
    67  
    68  // Reclaim returns the ID reserved for the peer back to unused pool.
    69  func (ids *mempoolIDs) Reclaim(peer p2p.Peer) {
    70  	ids.mtx.Lock()
    71  	defer ids.mtx.Unlock()
    72  
    73  	removedID, ok := ids.peerMap[peer.ID()]
    74  	if ok {
    75  		delete(ids.activeIDs, removedID)
    76  		delete(ids.peerMap, peer.ID())
    77  	}
    78  }
    79  
    80  // GetForPeer returns an ID reserved for the peer.
    81  func (ids *mempoolIDs) GetForPeer(peer p2p.Peer) uint16 {
    82  	ids.mtx.RLock()
    83  	defer ids.mtx.RUnlock()
    84  
    85  	return ids.peerMap[peer.ID()]
    86  }
    87  
    88  func newMempoolIDs() *mempoolIDs {
    89  	return &mempoolIDs{
    90  		peerMap:   make(map[p2p.ID]uint16),
    91  		activeIDs: map[uint16]struct{}{0: {}},
    92  		nextID:    1, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx
    93  	}
    94  }
    95  
    96  // NewReactor returns a new Reactor with the given config and mempool.
    97  func NewReactor(config *cfg.MempoolConfig, mempool *TxMempool, traceClient *trace.Client) *Reactor {
    98  	memR := &Reactor{
    99  		config:      config,
   100  		mempool:     mempool,
   101  		ids:         newMempoolIDs(),
   102  		traceClient: traceClient,
   103  	}
   104  	memR.BaseReactor = *p2p.NewBaseReactor("Mempool", memR)
   105  	return memR
   106  }
   107  
   108  // InitPeer implements Reactor by creating a state for the peer.
   109  func (memR *Reactor) InitPeer(peer p2p.Peer) p2p.Peer {
   110  	memR.ids.ReserveForPeer(peer)
   111  	return peer
   112  }
   113  
   114  // SetLogger sets the Logger on the reactor and the underlying mempool.
   115  func (memR *Reactor) SetLogger(l log.Logger) {
   116  	memR.Logger = l
   117  }
   118  
   119  // OnStart implements p2p.BaseReactor.
   120  func (memR *Reactor) OnStart() error {
   121  	if !memR.config.Broadcast {
   122  		memR.Logger.Info("Tx broadcasting is disabled")
   123  	}
   124  
   125  	// run a separate go routine to check for time based TTLs
   126  	if memR.mempool.config.TTLDuration > 0 {
   127  		go func() {
   128  			ticker := time.NewTicker(memR.mempool.config.TTLDuration)
   129  			for {
   130  				select {
   131  				case <-ticker.C:
   132  					memR.mempool.CheckToPurgeExpiredTxs()
   133  				case <-memR.Quit():
   134  					return
   135  				}
   136  			}
   137  		}()
   138  	}
   139  
   140  	return nil
   141  }
   142  
   143  // GetChannels implements Reactor by returning the list of channels for this
   144  // reactor.
   145  func (memR *Reactor) GetChannels() []*p2p.ChannelDescriptor {
   146  	largestTx := make([]byte, memR.config.MaxTxBytes)
   147  	batchMsg := protomem.Message{
   148  		Sum: &protomem.Message_Txs{
   149  			Txs: &protomem.Txs{Txs: [][]byte{largestTx}},
   150  		},
   151  	}
   152  
   153  	return []*p2p.ChannelDescriptor{
   154  		{
   155  			ID:                  mempool.MempoolChannel,
   156  			Priority:            5,
   157  			RecvMessageCapacity: batchMsg.Size(),
   158  			MessageType:         &protomem.Message{},
   159  		},
   160  	}
   161  }
   162  
   163  // AddPeer implements Reactor.
   164  // It starts a broadcast routine ensuring all txs are forwarded to the given peer.
   165  func (memR *Reactor) AddPeer(peer p2p.Peer) {
   166  	if memR.config.Broadcast {
   167  		go memR.broadcastTxRoutine(peer)
   168  	}
   169  }
   170  
   171  // RemovePeer implements Reactor.
   172  func (memR *Reactor) RemovePeer(peer p2p.Peer, reason interface{}) {
   173  	memR.ids.Reclaim(peer)
   174  	// broadcast routine checks if peer is gone and returns
   175  }
   176  
   177  // ReceiveEnvelope implements Reactor.
   178  // It adds any received transactions to the mempool.
   179  func (memR *Reactor) ReceiveEnvelope(e p2p.Envelope) {
   180  	memR.Logger.Debug("Receive", "src", e.Src, "chId", e.ChannelID, "msg", e.Message)
   181  	switch msg := e.Message.(type) {
   182  	case *protomem.Txs:
   183  		for _, tx := range msg.Txs {
   184  			schema.WriteMempoolTx(
   185  				memR.traceClient,
   186  				e.Src.ID(),
   187  				tx,
   188  				schema.TransferTypeDownload,
   189  				schema.V1VersionFieldValue,
   190  			)
   191  		}
   192  		protoTxs := msg.GetTxs()
   193  		if len(protoTxs) == 0 {
   194  			memR.Logger.Error("received tmpty txs from peer", "src", e.Src)
   195  			return
   196  		}
   197  		txInfo := mempool.TxInfo{SenderID: memR.ids.GetForPeer(e.Src)}
   198  		if e.Src != nil {
   199  			txInfo.SenderP2PID = e.Src.ID()
   200  		}
   201  
   202  		var err error
   203  		for _, tx := range protoTxs {
   204  			ntx := types.Tx(tx)
   205  			err = memR.mempool.CheckTx(ntx, nil, txInfo)
   206  			if errors.Is(err, mempool.ErrTxInCache) {
   207  				memR.Logger.Debug("Tx already exists in cache", "tx", ntx.String())
   208  			} else if err != nil {
   209  				memR.Logger.Info("Could not check tx", "tx", ntx.String(), "err", err)
   210  			}
   211  		}
   212  	default:
   213  		memR.Logger.Error("unknown message type", "src", e.Src, "chId", e.ChannelID, "msg", e.Message)
   214  		memR.Switch.StopPeerForError(e.Src, fmt.Errorf("mempool cannot handle message of type: %T", e.Message))
   215  		return
   216  	}
   217  
   218  	// broadcasting happens from go routines per peer
   219  }
   220  
   221  func (memR *Reactor) Receive(chID byte, peer p2p.Peer, msgBytes []byte) {
   222  	msg := &protomem.Message{}
   223  	err := proto.Unmarshal(msgBytes, msg)
   224  	if err != nil {
   225  		panic(err)
   226  	}
   227  	uw, err := msg.Unwrap()
   228  	if err != nil {
   229  		panic(err)
   230  	}
   231  	memR.ReceiveEnvelope(p2p.Envelope{
   232  		ChannelID: chID,
   233  		Src:       peer,
   234  		Message:   uw,
   235  	})
   236  }
   237  
   238  // PeerState describes the state of a peer.
   239  type PeerState interface {
   240  	GetHeight() int64
   241  }
   242  
   243  // Send new mempool txs to peer.
   244  func (memR *Reactor) broadcastTxRoutine(peer p2p.Peer) {
   245  	peerID := memR.ids.GetForPeer(peer)
   246  	var next *clist.CElement
   247  
   248  	for {
   249  		// In case of both next.NextWaitChan() and peer.Quit() are variable at the same time
   250  		if !memR.IsRunning() || !peer.IsRunning() {
   251  			return
   252  		}
   253  
   254  		// This happens because the CElement we were looking at got garbage
   255  		// collected (removed). That is, .NextWait() returned nil. Go ahead and
   256  		// start from the beginning.
   257  		if next == nil {
   258  			select {
   259  			case <-memR.mempool.TxsWaitChan(): // Wait until a tx is available
   260  				if next = memR.mempool.TxsFront(); next == nil {
   261  					continue
   262  				}
   263  
   264  			case <-peer.Quit():
   265  				return
   266  
   267  			case <-memR.Quit():
   268  				return
   269  			}
   270  		}
   271  
   272  		// Make sure the peer is up to date.
   273  		peerState, ok := peer.Get(types.PeerStateKey).(PeerState)
   274  		if !ok {
   275  			// Peer does not have a state yet. We set it in the consensus reactor, but
   276  			// when we add peer in Switch, the order we call reactors#AddPeer is
   277  			// different every time due to us using a map. Sometimes other reactors
   278  			// will be initialized before the consensus reactor. We should wait a few
   279  			// milliseconds and retry.
   280  			time.Sleep(mempool.PeerCatchupSleepIntervalMS * time.Millisecond)
   281  			continue
   282  		}
   283  
   284  		// Allow for a lag of 1 block.
   285  		memTx := next.Value.(*WrappedTx)
   286  		if peerState.GetHeight() < memTx.height-1 {
   287  			time.Sleep(mempool.PeerCatchupSleepIntervalMS * time.Millisecond)
   288  			continue
   289  		}
   290  
   291  		// NOTE: Transaction batching was disabled due to
   292  		// https://github.com/cometbft/cometbft/issues/5796
   293  		if !memTx.HasPeer(peerID) {
   294  			success := p2p.SendEnvelopeShim(peer, p2p.Envelope{ //nolint: staticcheck
   295  				ChannelID: mempool.MempoolChannel,
   296  				Message:   &protomem.Txs{Txs: [][]byte{memTx.tx}},
   297  			}, memR.Logger)
   298  			if !success {
   299  				time.Sleep(mempool.PeerCatchupSleepIntervalMS * time.Millisecond)
   300  				continue
   301  			} else {
   302  				// record that we have sent the peer the transaction
   303  				// to avoid doing it a second time
   304  				memTx.SetPeer(peerID)
   305  			}
   306  			schema.WriteMempoolTx(
   307  				memR.traceClient,
   308  				peer.ID(),
   309  				memTx.tx,
   310  				schema.TransferTypeUpload,
   311  				schema.V1VersionFieldValue,
   312  			)
   313  		}
   314  
   315  		select {
   316  		case <-next.NextWaitChan():
   317  			// see the start of the for loop for nil check
   318  			next = next.Next()
   319  
   320  		case <-peer.Quit():
   321  			return
   322  
   323  		case <-memR.Quit():
   324  			return
   325  		}
   326  	}
   327  }
   328  
   329  //-----------------------------------------------------------------------------
   330  // Messages
   331  
   332  // TxsMessage is a Message containing transactions.
   333  type TxsMessage struct {
   334  	Txs []types.Tx
   335  }
   336  
   337  // String returns a string representation of the TxsMessage.
   338  func (m *TxsMessage) String() string {
   339  	return fmt.Sprintf("[TxsMessage %v]", m.Txs)
   340  }