feat: support parallel processing of messages (celestiaorg#1595)

## Description

closes celestiaorg#1534

---------

Co-authored-by: Rootul P <[email protected]>

Author: rach-id

blockchain/v0/reactor.go

@@ -30,6 +30,9 @@ const (
 	statusUpdateIntervalSeconds = 10
 	// check if we should switch to consensus reactor
 	switchToConsensusIntervalSeconds = 1
+
+	// ReactorIncomingMessageQueueSize the size of the reactor's message queue.
+	ReactorIncomingMessageQueueSize = 10
 )
 
 type consensusReactor interface {
@@ -92,7 +95,7 @@ func NewBlockchainReactor(state sm.State, blockExec *sm.BlockExecutor, store *st
 		requestsCh:   requestsCh,
 		errorsCh:     errorsCh,
 	}
-	bcR.BaseReactor = *p2p.NewBaseReactor("BlockchainReactor", bcR)
+	bcR.BaseReactor = *p2p.NewBaseReactor("BlockchainReactor", bcR, p2p.WithIncomingQueueSize(ReactorIncomingMessageQueueSize))
 	return bcR
 }
 

blockchain/v1/reactor.go

@@ -24,6 +24,9 @@ const (
 
 	// ask for best height every 10s
 	statusUpdateIntervalSeconds = 10
+
+	// ReactorIncomingMessageQueueSize the size of the reactor's message queue.
+	ReactorIncomingMessageQueueSize = 10
 )
 
 var (
@@ -100,7 +103,7 @@ func NewBlockchainReactor(state sm.State, blockExec *sm.BlockExecutor, store *st
 	}
 	fsm := NewFSM(startHeight, bcR)
 	bcR.fsm = fsm
-	bcR.BaseReactor = *p2p.NewBaseReactor("BlockchainReactor", bcR)
+	bcR.BaseReactor = *p2p.NewBaseReactor("BlockchainReactor", bcR, p2p.WithIncomingQueueSize(ReactorIncomingMessageQueueSize))
 	// bcR.swReporter = behaviour.NewSwitchReporter(bcR.BaseReactor.Switch)
 
 	return bcR

consensus/byzantine_test.go

@@ -610,3 +610,7 @@ func (br *ByzantineReactor) Receive(chID byte, p p2p.Peer, m []byte) {
 	br.reactor.Receive(chID, p, m)
 }
 func (br *ByzantineReactor) InitPeer(peer p2p.Peer) p2p.Peer { return peer }
+
+func (br *ByzantineReactor) QueueUnprocessedEnvelope(e p2p.UnprocessedEnvelope) {
+	br.reactor.QueueUnprocessedEnvelope(e)
+}

consensus/reactor.go

@@ -34,6 +34,9 @@ const (
 
 	blocksToContributeToBecomeGoodPeer = 10000
 	votesToContributeToBecomeGoodPeer  = 10000
+
+	// ReactorIncomingMessageQueueSize the size of the reactor's message queue.
+	ReactorIncomingMessageQueueSize = 1000
 )
 
 //-----------------------------------------------------------------------------
@@ -65,7 +68,7 @@ func NewReactor(consensusState *State, waitSync bool, options ...ReactorOption)
 		Metrics:     NopMetrics(),
 		traceClient: trace.NoOpTracer(),
 	}
-	conR.BaseReactor = *p2p.NewBaseReactor("Consensus", conR)
+	conR.BaseReactor = *p2p.NewBaseReactor("Consensus", conR, p2p.WithIncomingQueueSize(ReactorIncomingMessageQueueSize))
 
 	for _, option := range options {
 		option(conR)

evidence/reactor.go

@@ -24,6 +24,9 @@ const (
 	broadcastEvidenceIntervalS = 10
 	// If a message fails wait this much before sending it again
 	peerRetryMessageIntervalMS = 100
+
+	// ReactorIncomingMessageQueueSize the size of the reactor's message queue.
+	ReactorIncomingMessageQueueSize = 1
 )
 
 // Reactor handles evpool evidence broadcasting amongst peers.
@@ -38,7 +41,7 @@ func NewReactor(evpool *Pool) *Reactor {
 	evR := &Reactor{
 		evpool: evpool,
 	}
-	evR.BaseReactor = *p2p.NewBaseReactor("Evidence", evR)
+	evR.BaseReactor = *p2p.NewBaseReactor("Evidence", evR, p2p.WithIncomingQueueSize(ReactorIncomingMessageQueueSize))
 	return evR
 }
 
@@ -175,7 +178,7 @@ func (evR *Reactor) broadcastEvidenceRoutine(peer p2p.Peer) {
 
 // Returns the message to send to the peer, or nil if the evidence is invalid for the peer.
 // If message is nil, we should sleep and try again.
-func (evR Reactor) prepareEvidenceMessage(
+func (evR *Reactor) prepareEvidenceMessage(
 	peer p2p.Peer,
 	ev types.Evidence,
 ) (evis []types.Evidence) {
@@ -222,6 +225,10 @@ func (evR Reactor) prepareEvidenceMessage(
 	return []types.Evidence{ev}
 }
 
+func (evR *Reactor) OnStop() {
+	evR.BaseReactor.OnStop()
+}
+
 // PeerState describes the state of a peer.
 type PeerState interface {
 	GetHeight() int64

evidence/reactor_test.go

@@ -224,6 +224,7 @@ func TestReactorBroadcastEvidenceMemoryLeak(t *testing.T) {
 	r.AddPeer(p)
 
 	_ = sendEvidence(t, pool, val, 2)
+	r.OnStop()
 }
 
 // evidenceLogger is a TestingLogger which uses a different

mempool/v0/reactor_test.go

@@ -88,6 +88,8 @@ func TestReactorConcurrency(t *testing.T) {
 
 	const numTxs = 5
 
+	reactors[0].mempool.config.Size = 10000
+	reactors[1].mempool.config.Size = 10000
 	for i := 0; i < 1000; i++ {
 		wg.Add(2)
 
@@ -119,9 +121,6 @@ func TestReactorConcurrency(t *testing.T) {
 			err := reactors[1].mempool.Update(1, []types.Tx{}, make([]*abci.ResponseDeliverTx, 0), nil, nil)
 			assert.NoError(t, err)
 		}()
-
-		// 1. flush the mempool
-		reactors[1].mempool.Flush()
 	}
 
 	wg.Wait()

p2p/base_reactor.go

@@ -1,10 +1,19 @@
 package p2p
 
 import (
+	"context"
+	"fmt"
+	"reflect"
+
+	"github.com/gogo/protobuf/proto"
 	"github.com/tendermint/tendermint/libs/service"
 	"github.com/tendermint/tendermint/p2p/conn"
+	"github.com/tendermint/tendermint/pkg/trace/schema"
 )
 
+// ProcessorFunc is the message processor function type.
+type ProcessorFunc func(context.Context, <-chan UnprocessedEnvelope) error
+
 // Reactor is responsible for handling incoming messages on one or more
 // Channel. Switch calls GetChannels when reactor is added to it. When a new
 // peer joins our node, InitPeer and AddPeer are called. RemovePeer is called
@@ -51,6 +60,15 @@ type Reactor interface {
 	// Deprecated: Reactors looking to receive data from a peer should implement ReceiveEnvelope.
 	// Receive will be deprecated in favor of ReceiveEnvelope in v0.37.
 	Receive(chID byte, peer Peer, msgBytes []byte)
+
+	// QueueUnprocessedEnvelope is called by the switch when an unprocessed
+	// envelope is received. Unprocessed envelopes are immediately buffered in a
+	// queue to avoid blocking. Incoming messages are then passed to a
+	// processing function. The default processing function unmarshals the
+	// messages in the order the sender sent them and then calls Receive on the
+	// reactor. The queue size and the processing function can be changed via
+	// passing options to the base reactor.
+	QueueUnprocessedEnvelope(e UnprocessedEnvelope)
 }
 
 type EnvelopeReceiver interface {
@@ -68,18 +86,147 @@ type EnvelopeReceiver interface {
 type BaseReactor struct {
 	service.BaseService // Provides Start, Stop, .Quit
 	Switch              *Switch
+
+	incoming chan UnprocessedEnvelope
+
+	ctx    context.Context
+	cancel context.CancelFunc
+	// processor is called with the incoming channel and is responsible for
+	// unmarshalling the messages and calling Receive on the reactor.
+	processor ProcessorFunc
 }
 
-func NewBaseReactor(name string, impl Reactor) *BaseReactor {
-	return &BaseReactor{
+type ReactorOptions func(*BaseReactor)
+
+func NewBaseReactor(name string, impl Reactor, opts ...ReactorOptions) *BaseReactor {
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	base := &BaseReactor{
+		ctx:         ctx,
+		cancel:      cancel,
 		BaseService: *service.NewBaseService(nil, name, impl),
 		Switch:      nil,
+		incoming:    make(chan UnprocessedEnvelope, 100),
+		processor:   DefaultProcessor(impl),
+	}
+	for _, opt := range opts {
+		opt(base)
+	}
+
+	go func() {
+		err := base.processor(ctx, base.incoming)
+		if err != nil {
+			err = base.Stop()
+			if err != nil {
+				panic(err)
+			}
+		}
+	}()
+
+	return base
+}
+
+// WithProcessor sets the processor function for the reactor. The processor
+// function is called with the incoming channel and is responsible for
+// unmarshalling the messages and calling Receive on the reactor.
+func WithProcessor(processor ProcessorFunc) ReactorOptions {
+	return func(br *BaseReactor) {
+		br.processor = processor
+	}
+}
+
+// WithIncomingQueueSize sets the size of the incoming message queue for a
+// reactor.
+func WithIncomingQueueSize(size int) ReactorOptions {
+	return func(br *BaseReactor) {
+		br.incoming = make(chan UnprocessedEnvelope, size)
 	}
 }
 
 func (br *BaseReactor) SetSwitch(sw *Switch) {
 	br.Switch = sw
 }
+
+// QueueUnprocessedEnvelope is called by the switch when an unprocessed
+// envelope is received. Unprocessed envelopes are immediately buffered in a
+// queue to avoid blocking. The size of the queue can be changed by passing
+// options to the base reactor.
+func (br *BaseReactor) QueueUnprocessedEnvelope(e UnprocessedEnvelope) {
+	select {
+	// if the context is done, do nothing.
+	case <-br.ctx.Done():
+	// if not, add the item to the channel.
+	case br.incoming <- e:
+	}
+}
+
+func (br *BaseReactor) OnStop() {
+	br.cancel()
+	close(br.incoming)
+}
+
+// DefaultProcessor unmarshalls the message and calls Receive on the reactor.
+// This preserves the sender's original order for all messages.
+func DefaultProcessor(impl Reactor) func(context.Context, <-chan UnprocessedEnvelope) error {
+	implChannels := impl.GetChannels()
+
+	chIDs := make(map[byte]proto.Message, len(implChannels))
+	for _, chDesc := range implChannels {
+		chIDs[chDesc.ID] = chDesc.MessageType
+	}
+	return func(ctx context.Context, incoming <-chan UnprocessedEnvelope) error {
+		for {
+			select {
+			case <-ctx.Done():
+				return nil
+			case ue, ok := <-incoming:
+				if !ok {
+					// this means the channel was closed.
+					return nil
+				}
+				mt := chIDs[ue.ChannelID]
+
+				if mt == nil {
+					return fmt.Errorf("no message type registered for channel %d", ue.ChannelID)
+				}
+
+				msg := proto.Clone(mt)
+
+				err := proto.Unmarshal(ue.Message, msg)
+				if err != nil {
+					return fmt.Errorf("unmarshaling message: %v into type: %s resulted in error %w", msg, reflect.TypeOf(mt), err)
+				}
+
+				if w, ok := msg.(Unwrapper); ok {
+					msg, err = w.Unwrap()
+					if err != nil {
+						return fmt.Errorf("unwrapping message: %v", err)
+					}
+				}
+
+				labels := []string{
+					"peer_id", string(ue.Src.ID()),
+					"chID", fmt.Sprintf("%#x", ue.ChannelID),
+				}
+
+				ue.Src.Metrics().PeerReceiveBytesTotal.With(labels...).Add(float64(len(ue.Message)))
+				ue.Src.Metrics().MessageReceiveBytesTotal.With(append(labels, "message_type", ue.Src.ValueToMetricLabel(msg))...).Add(float64(len(ue.Message)))
+				schema.WriteReceivedBytes(ue.Src.TraceClient(), string(ue.Src.ID()), ue.ChannelID, len(ue.Message))
+
+				if nr, ok := impl.(EnvelopeReceiver); ok {
+					nr.ReceiveEnvelope(Envelope{
+						ChannelID: ue.ChannelID,
+						Src:       ue.Src,
+						Message:   msg,
+					})
+				} else {
+					impl.Receive(ue.ChannelID, ue.Src, ue.Message)
+				}
+			}
+		}
+	}
+}
+
 func (*BaseReactor) GetChannels() []*conn.ChannelDescriptor        { return nil }
 func (*BaseReactor) AddPeer(peer Peer)                             {}
 func (*BaseReactor) RemovePeer(peer Peer, reason interface{})      {}

p2p/base_reactor_test.go

@@ -0,0 +1,122 @@
+package p2p_test
+
+import (
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/stretchr/testify/require"
+	"github.com/tendermint/tendermint/libs/service"
+	"github.com/tendermint/tendermint/p2p"
+	"github.com/tendermint/tendermint/p2p/conn"
+	"github.com/tendermint/tendermint/pkg/trace"
+	"github.com/tendermint/tendermint/proto/tendermint/mempool"
+)
+
+// TestBaseReactorProcessor tests the BaseReactor's message processing by
+// queueing encoded messages and adding artificial delay to the first message.
+// Depending on the processors used, the ordering of the sender could be lost.
+func TestBaseReactorProcessor(t *testing.T) {
+	// a reactor using the default processor should be able to queue
+	// messages, and they get processed in order.
+	or := NewOrderedReactor()
+
+	msgs := []string{"msg1", "msg2", "msg3"}
+	or.fillQueue(t, msgs...)
+
+	time.Sleep(300 * time.Millisecond) // wait plenty of time for the processing to finish
+
+	or.Lock()
+	require.Equal(t, len(msgs), len(or.received))
+	require.Equal(t, msgs, or.received)
+	or.Unlock()
+}
+
+var _ p2p.Reactor = &orderedReactor{}
+
+// orderedReactor is used for testing. It saves each envelope in the order it
+// receives it.
+type orderedReactor struct {
+	p2p.BaseReactor
+
+	sync.Mutex
+	received      []string
+	receivedFirst bool
+}
+
+func NewOrderedReactor() *orderedReactor {
+	r := &orderedReactor{Mutex: sync.Mutex{}}
+	r.BaseReactor = *p2p.NewBaseReactor("Ordered Rector", r, p2p.WithIncomingQueueSize(10))
+	return r
+}
+
+func (r *orderedReactor) GetChannels() []*conn.ChannelDescriptor {
+	return []*conn.ChannelDescriptor{
+		{
+			ID:                  0x99,
+			Priority:            1,
+			RecvMessageCapacity: 10,
+			MessageType:         &mempool.Txs{},
+		},
+	}
+}
+
+// ReceiveEnvelope adds a delay to the first processed envelope to test ordering.
+func (r *orderedReactor) ReceiveEnvelope(e p2p.Envelope) {
+	r.Lock()
+	f := r.receivedFirst
+	if !f {
+		r.receivedFirst = true
+		r.Unlock()
+		time.Sleep(100 * time.Millisecond)
+	} else {
+		r.Unlock()
+	}
+	r.Lock()
+	defer r.Unlock()
+
+	envMsg := e.Message.(*mempool.Txs)
+	r.received = append(r.received, string(envMsg.Txs[0]))
+}
+
+func (r *orderedReactor) fillQueue(t *testing.T, msgs ...string) {
+	peer := &imaginaryPeer{}
+	for _, msg := range msgs {
+		s, err := proto.Marshal(&mempool.Txs{Txs: [][]byte{[]byte(msg)}})
+		require.NoError(t, err)
+		r.QueueUnprocessedEnvelope(p2p.UnprocessedEnvelope{
+			Src:       peer,
+			Message:   s,
+			ChannelID: 0x99,
+		})
+	}
+}
+
+var _ p2p.IntrospectivePeer = &imaginaryPeer{}
+
+type imaginaryPeer struct {
+	service.BaseService
+}
+
+func (ip *imaginaryPeer) TraceClient() trace.Tracer       { return trace.NoOpTracer() }
+func (ip *imaginaryPeer) HasIPChanged() bool              { return false }
+func (ip *imaginaryPeer) FlushStop()                      {}
+func (ip *imaginaryPeer) ID() p2p.ID                      { return "" }
+func (ip *imaginaryPeer) RemoteIP() net.IP                { return []byte{} }
+func (ip *imaginaryPeer) RemoteAddr() net.Addr            { return nil }
+func (ip *imaginaryPeer) IsOutbound() bool                { return true }
+func (ip *imaginaryPeer) CloseConn() error                { return nil }
+func (ip *imaginaryPeer) IsPersistent() bool              { return false }
+func (ip *imaginaryPeer) NodeInfo() p2p.NodeInfo          { return nil }
+func (ip *imaginaryPeer) Status() conn.ConnectionStatus   { return conn.ConnectionStatus{} }
+func (ip *imaginaryPeer) SocketAddr() *p2p.NetAddress     { return nil }
+func (ip *imaginaryPeer) Send(byte, []byte) bool          { return true }
+func (ip *imaginaryPeer) TrySend(byte, []byte) bool       { return true }
+func (ip *imaginaryPeer) Set(key string, value any)       {}
+func (ip *imaginaryPeer) Get(key string) any              { return nil }
+func (ip *imaginaryPeer) SetRemovalFailed()               {}
+func (ip *imaginaryPeer) GetRemovalFailed() bool          { return false }
+func (ip *imaginaryPeer) Metrics() *p2p.Metrics           { return p2p.NopMetrics() }
+func (ip *imaginaryPeer) ValueToMetricLabel(i any) string { return "" }

p2p/conn/connection.go

@@ -865,7 +865,8 @@ func (ch *Channel) recvPacketMsg(packet tmp2p.PacketMsg) ([]byte, error) {
 	}
 	ch.recving = append(ch.recving, packet.Data...)
 	if packet.EOF {
-		msgBytes := ch.recving
+		msgBytes := make([]byte, len(ch.recving))
+		copy(msgBytes, ch.recving)
 
 		// clear the slice without re-allocating.
 		// http://stackoverflow.com/questions/16971741/how-do-you-clear-a-slice-in-go

p2p/mock/reactor.go

@@ -14,7 +14,7 @@ type Reactor struct {
 
 func NewReactor() *Reactor {
 	r := &Reactor{}
-	r.BaseReactor = *p2p.NewBaseReactor("Mock-PEX", r)
+	r.BaseReactor = *p2p.NewBaseReactor("Mock-PEX", r, p2p.WithIncomingQueueSize(1))
 	r.SetLogger(log.TestingLogger())
 	return r
 }

p2p/peer.go

@@ -3,7 +3,6 @@ package p2p
 import (
 	"fmt"
 	"net"
-	"reflect"
 	"time"
 
 	"github.com/gogo/protobuf/proto"
@@ -55,6 +54,13 @@ type Peer interface {
 	GetRemovalFailed() bool
 }
 
+type IntrospectivePeer interface {
+	Peer
+	Metrics() *Metrics
+	ValueToMetricLabel(i any) string
+	TraceClient() trace.Tracer
+}
+
 type EnvelopeSender interface {
 	SendEnvelope(Envelope) bool
 	TrySendEnvelope(Envelope) bool
@@ -187,6 +193,10 @@ type peer struct {
 	removalAttemptFailed bool
 }
 
+func (p *peer) TraceClient() trace.Tracer {
+	return p.traceClient
+}
+
 type PeerOption func(*peer)
 
 func WithPeerTracer(t trace.Tracer) PeerOption {
@@ -200,7 +210,7 @@ func newPeer(
 	mConfig cmtconn.MConnConfig,
 	nodeInfo NodeInfo,
 	reactorsByCh map[byte]Reactor,
-	msgTypeByChID map[byte]proto.Message,
+	_ map[byte]proto.Message,
 	chDescs []*cmtconn.ChannelDescriptor,
 	onPeerError func(Peer, interface{}),
 	mlc *metricsLabelCache,
@@ -221,7 +231,6 @@ func newPeer(
 		pc.conn,
 		p,
 		reactorsByCh,
-		msgTypeByChID,
 		chDescs,
 		onPeerError,
 		mConfig,
@@ -246,6 +255,14 @@ func (p *peer) String() string {
 //---------------------------------------------------
 // Implements service.Service
 
+func (p *peer) Metrics() *Metrics {
+	return p.metrics
+}
+
+func (p *peer) ValueToMetricLabel(i any) string {
+	return p.mlc.ValueToMetricLabel(i)
+}
+
 // SetLogger implements BaseService.
 func (p *peer) SetLogger(l log.Logger) {
 	p.Logger = l
@@ -539,7 +556,6 @@ func createMConnection(
 	conn net.Conn,
 	p *peer,
 	reactorsByCh map[byte]Reactor,
-	msgTypeByChID map[byte]proto.Message,
 	chDescs []*cmtconn.ChannelDescriptor,
 	onPeerError func(Peer, interface{}),
 	config cmtconn.MConnConfig,
@@ -552,37 +568,12 @@ func createMConnection(
 			// which does onPeerError.
 			panic(fmt.Sprintf("Unknown channel %X", chID))
 		}
-		mt := msgTypeByChID[chID]
-		msg := proto.Clone(mt)
-		err := proto.Unmarshal(msgBytes, msg)
-		if err != nil {
-			panic(fmt.Errorf("unmarshaling message: %s into type: %s", err, reflect.TypeOf(mt)))
-		}
 
-		if w, ok := msg.(Unwrapper); ok {
-			msg, err = w.Unwrap()
-			if err != nil {
-				panic(fmt.Errorf("unwrapping message: %s", err))
-			}
-		}
-
-		labels := []string{
-			"peer_id", string(p.ID()),
-			"chID", fmt.Sprintf("%#x", chID),
-		}
-
-		p.metrics.PeerReceiveBytesTotal.With(labels...).Add(float64(len(msgBytes)))
-		p.metrics.MessageReceiveBytesTotal.With(append(labels, "message_type", p.mlc.ValueToMetricLabel(msg))...).Add(float64(len(msgBytes)))
-		schema.WriteReceivedBytes(p.traceClient, string(p.ID()), chID, len(msgBytes))
-		if nr, ok := reactor.(EnvelopeReceiver); ok {
-			nr.ReceiveEnvelope(Envelope{
-				ChannelID: chID,
-				Src:       p,
-				Message:   msg,
-			})
-		} else {
-			reactor.Receive(chID, p, msgBytes)
-		}
+		reactor.QueueUnprocessedEnvelope(UnprocessedEnvelope{
+			ChannelID: chID,
+			Src:       p,
+			Message:   msgBytes,
+		})
 	}
 
 	onError := func(r interface{}) {

p2p/types.go

@@ -16,6 +16,12 @@ type Envelope struct {
 	ChannelID byte
 }
 
+type UnprocessedEnvelope struct {
+	Src       IntrospectivePeer
+	Message   []byte
+	ChannelID byte
+}
+
 // Unwrapper is a Protobuf message that can contain a variety of inner messages
 // (e.g. via oneof fields). If a Channel's message type implements Unwrapper, the
 // p2p layer will automatically unwrap inbound messages so that reactors do not have to do this themselves.

statesync/reactor.go

@@ -24,6 +24,9 @@ const (
 	ChunkChannel = byte(0x61)
 	// recentSnapshots is the number of recent snapshots to send and receive per peer.
 	recentSnapshots = 10
+
+	// ReactorIncomingMessageQueueSize the size of the reactor's message queue.
+	ReactorIncomingMessageQueueSize = 100
 )
 
 // Reactor handles state sync, both restoring snapshots for the local node and serving snapshots
@@ -55,7 +58,7 @@ func NewReactor(
 		conn:      conn,
 		connQuery: connQuery,
 	}
-	r.BaseReactor = *p2p.NewBaseReactor("StateSync", r)
+	r.BaseReactor = *p2p.NewBaseReactor("StateSync", r, p2p.WithIncomingQueueSize(ReactorIncomingMessageQueueSize))
 
 	return r
 }