p2p: add a per-message type send and receive metric (backport #9622) (#9641)

* p2p: add a per-message type send and receive metric (#9622)

* p2p: ressurrect the p2p envelope and use to calculate message metric

Add new SendEnvelope, TrySendEnvelope, BroadcastEnvelope, and ReceiveEnvelope methods in the p2p package to work with the new envelope type.

Care was taken to ensure this was performed in a non-breaking manner.

Co-authored-by: William Banfield <4561443+williambanfield@users.noreply.github.com>
Co-authored-by: William Banfield <wbanfield@gmail.com>

statesync/messages.go

@@ -16,49 +16,6 @@ const (
 	chunkMsgSize = int(16e6)
 )
 
-// mustEncodeMsg encodes a Protobuf message, panicing on error.
-func mustEncodeMsg(pb proto.Message) []byte {
-	msg := ssproto.Message{}
-	switch pb := pb.(type) {
-	case *ssproto.ChunkRequest:
-		msg.Sum = &ssproto.Message_ChunkRequest{ChunkRequest: pb}
-	case *ssproto.ChunkResponse:
-		msg.Sum = &ssproto.Message_ChunkResponse{ChunkResponse: pb}
-	case *ssproto.SnapshotsRequest:
-		msg.Sum = &ssproto.Message_SnapshotsRequest{SnapshotsRequest: pb}
-	case *ssproto.SnapshotsResponse:
-		msg.Sum = &ssproto.Message_SnapshotsResponse{SnapshotsResponse: pb}
-	default:
-		panic(fmt.Errorf("unknown message type %T", pb))
-	}
-	bz, err := msg.Marshal()
-	if err != nil {
-		panic(fmt.Errorf("unable to marshal %T: %w", pb, err))
-	}
-	return bz
-}
-
-// decodeMsg decodes a Protobuf message.
-func decodeMsg(bz []byte) (proto.Message, error) {
-	pb := &ssproto.Message{}
-	err := proto.Unmarshal(bz, pb)
-	if err != nil {
-		return nil, err
-	}
-	switch msg := pb.Sum.(type) {
-	case *ssproto.Message_ChunkRequest:
-		return msg.ChunkRequest, nil
-	case *ssproto.Message_ChunkResponse:
-		return msg.ChunkResponse, nil
-	case *ssproto.Message_SnapshotsRequest:
-		return msg.SnapshotsRequest, nil
-	case *ssproto.Message_SnapshotsResponse:
-		return msg.SnapshotsResponse, nil
-	default:
-		return nil, fmt.Errorf("unknown message type %T", msg)
-	}
-}
-
 // validateMsg validates a message.
 func validateMsg(pb proto.Message) error {
 	if pb == nil {

statesync/messages_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/gogo/protobuf/proto"
 	"github.com/stretchr/testify/require"
 
+	"github.com/tendermint/tendermint/p2p"
 	ssproto "github.com/tendermint/tendermint/proto/tendermint/statesync"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 )
@@ -99,8 +100,9 @@ func TestStateSyncVectors(t *testing.T) {
 
 	for _, tc := range testCases {
 		tc := tc
-
-		bz := mustEncodeMsg(tc.msg)
+		w := tc.msg.(p2p.Wrapper).Wrap()
+		bz, err := proto.Marshal(w)
+		require.NoError(t, err)
 
 		require.Equal(t, tc.expBytes, hex.EncodeToString(bz), tc.testName)
 	}

statesync/reactor.go

@@ -5,6 +5,8 @@ import (
 	"sort"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
+
 	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/config"
 	tmsync "github.com/tendermint/tendermint/libs/sync"
@@ -66,12 +68,14 @@ func (r *Reactor) GetChannels() []*p2p.ChannelDescriptor {
 			Priority:            5,
 			SendQueueCapacity:   10,
 			RecvMessageCapacity: snapshotMsgSize,
+			MessageType:         &ssproto.Message{},
 		},
 		{
 			ID:                  ChunkChannel,
 			Priority:            3,
 			SendQueueCapacity:   10,
 			RecvMessageCapacity: chunkMsgSize,
+			MessageType:         &ssproto.Message{},
 		},
 	}
 }
@@ -100,27 +104,21 @@ func (r *Reactor) RemovePeer(peer p2p.Peer, reason interface{}) {
 }
 
 // Receive implements p2p.Reactor.
-func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
+func (r *Reactor) ReceiveEnvelope(e p2p.Envelope) {
 	if !r.IsRunning() {
 		return
 	}
 
-	msg, err := decodeMsg(msgBytes)
+	err := validateMsg(e.Message)
 	if err != nil {
-		r.Logger.Error("Error decoding message", "src", src, "chId", chID, "err", err)
-		r.Switch.StopPeerForError(src, err)
-		return
-	}
-	err = validateMsg(msg)
-	if err != nil {
-		r.Logger.Error("Invalid message", "peer", src, "msg", msg, "err", err)
-		r.Switch.StopPeerForError(src, err)
+		r.Logger.Error("Invalid message", "peer", e.Src, "msg", e.Message, "err", err)
+		r.Switch.StopPeerForError(e.Src, err)
 		return
 	}
 
-	switch chID {
+	switch e.ChannelID {
 	case SnapshotChannel:
-		switch msg := msg.(type) {
+		switch msg := e.Message.(type) {
 		case *ssproto.SnapshotsRequest:
 			snapshots, err := r.recentSnapshots(recentSnapshots)
 			if err != nil {
@@ -129,14 +127,17 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 			}
 			for _, snapshot := range snapshots {
 				r.Logger.Debug("Advertising snapshot", "height", snapshot.Height,
-					"format", snapshot.Format, "peer", src.ID())
-				src.Send(chID, mustEncodeMsg(&ssproto.SnapshotsResponse{
-					Height:   snapshot.Height,
-					Format:   snapshot.Format,
-					Chunks:   snapshot.Chunks,
-					Hash:     snapshot.Hash,
-					Metadata: snapshot.Metadata,
-				}))
+					"format", snapshot.Format, "peer", e.Src.ID())
+				p2p.SendEnvelopeShim(e.Src, p2p.Envelope{ //nolint: staticcheck
+					ChannelID: e.ChannelID,
+					Message: &ssproto.SnapshotsResponse{
+						Height:   snapshot.Height,
+						Format:   snapshot.Format,
+						Chunks:   snapshot.Chunks,
+						Hash:     snapshot.Hash,
+						Metadata: snapshot.Metadata,
+					},
+				}, r.Logger)
 			}
 
 		case *ssproto.SnapshotsResponse:
@@ -146,8 +147,8 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 				r.Logger.Debug("Received unexpected snapshot, no state sync in progress")
 				return
 			}
-			r.Logger.Debug("Received snapshot", "height", msg.Height, "format", msg.Format, "peer", src.ID())
-			_, err := r.syncer.AddSnapshot(src, &snapshot{
+			r.Logger.Debug("Received snapshot", "height", msg.Height, "format", msg.Format, "peer", e.Src.ID())
+			_, err := r.syncer.AddSnapshot(e.Src, &snapshot{
 				Height:   msg.Height,
 				Format:   msg.Format,
 				Chunks:   msg.Chunks,
@@ -157,7 +158,7 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 			// TODO: We may want to consider punishing the peer for certain errors
 			if err != nil {
 				r.Logger.Error("Failed to add snapshot", "height", msg.Height, "format", msg.Format,
-					"peer", src.ID(), "err", err)
+					"peer", e.Src.ID(), "err", err)
 				return
 			}
 
@@ -166,10 +167,10 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 		}
 
 	case ChunkChannel:
-		switch msg := msg.(type) {
+		switch msg := e.Message.(type) {
 		case *ssproto.ChunkRequest:
 			r.Logger.Debug("Received chunk request", "height", msg.Height, "format", msg.Format,
-				"chunk", msg.Index, "peer", src.ID())
+				"chunk", msg.Index, "peer", e.Src.ID())
 			resp, err := r.conn.LoadSnapshotChunkSync(abci.RequestLoadSnapshotChunk{
 				Height: msg.Height,
 				Format: msg.Format,
@@ -181,30 +182,33 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 				return
 			}
 			r.Logger.Debug("Sending chunk", "height", msg.Height, "format", msg.Format,
-				"chunk", msg.Index, "peer", src.ID())
-			src.Send(ChunkChannel, mustEncodeMsg(&ssproto.ChunkResponse{
-				Height:  msg.Height,
-				Format:  msg.Format,
-				Index:   msg.Index,
-				Chunk:   resp.Chunk,
-				Missing: resp.Chunk == nil,
-			}))
+				"chunk", msg.Index, "peer", e.Src.ID())
+			p2p.SendEnvelopeShim(e.Src, p2p.Envelope{ //nolint: staticcheck
+				ChannelID: ChunkChannel,
+				Message: &ssproto.ChunkResponse{
+					Height:  msg.Height,
+					Format:  msg.Format,
+					Index:   msg.Index,
+					Chunk:   resp.Chunk,
+					Missing: resp.Chunk == nil,
+				},
+			}, r.Logger)
 
 		case *ssproto.ChunkResponse:
 			r.mtx.RLock()
 			defer r.mtx.RUnlock()
 			if r.syncer == nil {
-				r.Logger.Debug("Received unexpected chunk, no state sync in progress", "peer", src.ID())
+				r.Logger.Debug("Received unexpected chunk, no state sync in progress", "peer", e.Src.ID())
 				return
 			}
 			r.Logger.Debug("Received chunk, adding to sync", "height", msg.Height, "format", msg.Format,
-				"chunk", msg.Index, "peer", src.ID())
+				"chunk", msg.Index, "peer", e.Src.ID())
 			_, err := r.syncer.AddChunk(&chunk{
 				Height: msg.Height,
 				Format: msg.Format,
 				Index:  msg.Index,
 				Chunk:  msg.Chunk,
-				Sender: src.ID(),
+				Sender: e.Src.ID(),
 			})
 			if err != nil {
 				r.Logger.Error("Failed to add chunk", "height", msg.Height, "format", msg.Format,
@@ -217,10 +221,28 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 		}
 
 	default:
-		r.Logger.Error("Received message on invalid channel %x", chID)
+		r.Logger.Error("Received message on invalid channel %x", e.ChannelID)
 	}
 }
 
+func (r *Reactor) Receive(chID byte, peer p2p.Peer, msgBytes []byte) {
+	var msg *ssproto.Message
+	err := proto.Unmarshal(msgBytes, msg)
+	if err != nil {
+		panic(err)
+	}
+	um, err := msg.Unwrap()
+	if err != nil {
+		panic(err)
+	}
+
+	r.ReceiveEnvelope(p2p.Envelope{
+		ChannelID: chID,
+		Src:       peer,
+		Message:   um,
+	})
+}
+
 // recentSnapshots fetches the n most recent snapshots from the app
 func (r *Reactor) recentSnapshots(n uint32) ([]*snapshot, error) {
 	resp, err := r.conn.ListSnapshotsSync(abci.RequestListSnapshots{})
@@ -269,7 +291,11 @@ func (r *Reactor) Sync(stateProvider StateProvider, discoveryTime time.Duration)
 	hook := func() {
 		r.Logger.Debug("Requesting snapshots from known peers")
 		// Request snapshots from all currently connected peers
-		r.Switch.Broadcast(SnapshotChannel, mustEncodeMsg(&ssproto.SnapshotsRequest{}))
+
+		r.Switch.BroadcastEnvelope(p2p.Envelope{
+			ChannelID: SnapshotChannel,
+			Message:   &ssproto.SnapshotsRequest{},
+		})
 	}
 
 	hook()

statesync/reactor_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -53,10 +54,18 @@ func TestReactor_Receive_ChunkRequest(t *testing.T) {
 			peer.On("ID").Return(p2p.ID("id"))
 			var response *ssproto.ChunkResponse
 			if tc.expectResponse != nil {
-				peer.On("Send", ChunkChannel, mock.Anything).Run(func(args mock.Arguments) {
-					msg, err := decodeMsg(args[1].([]byte))
+				peer.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+					e, ok := i.(p2p.Envelope)
+					return ok && e.ChannelID == ChunkChannel
+				})).Run(func(args mock.Arguments) {
+					e := args[0].(p2p.Envelope)
+
+					// Marshal to simulate a wire roundtrip.
+					bz, err := proto.Marshal(e.Message)
 					require.NoError(t, err)
-					response = msg.(*ssproto.ChunkResponse)
+					err = proto.Unmarshal(bz, e.Message)
+					require.NoError(t, err)
+					response = e.Message.(*ssproto.ChunkResponse)
 				}).Return(true)
 			}
 
@@ -71,7 +80,11 @@ func TestReactor_Receive_ChunkRequest(t *testing.T) {
 				}
 			})
 
-			r.Receive(ChunkChannel, peer, mustEncodeMsg(tc.request))
+			r.ReceiveEnvelope(p2p.Envelope{
+				ChannelID: ChunkChannel,
+				Src:       peer,
+				Message:   tc.request,
+			})
 			time.Sleep(100 * time.Millisecond)
 			assert.Equal(t, tc.expectResponse, response)
 
@@ -131,10 +144,18 @@ func TestReactor_Receive_SnapshotsRequest(t *testing.T) {
 			peer := &p2pmocks.Peer{}
 			if len(tc.expectResponses) > 0 {
 				peer.On("ID").Return(p2p.ID("id"))
-				peer.On("Send", SnapshotChannel, mock.Anything).Run(func(args mock.Arguments) {
-					msg, err := decodeMsg(args[1].([]byte))
+				peer.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+					e, ok := i.(p2p.Envelope)
+					return ok && e.ChannelID == SnapshotChannel
+				})).Run(func(args mock.Arguments) {
+					e := args[0].(p2p.Envelope)
+
+					// Marshal to simulate a wire roundtrip.
+					bz, err := proto.Marshal(e.Message)
 					require.NoError(t, err)
-					responses = append(responses, msg.(*ssproto.SnapshotsResponse))
+					err = proto.Unmarshal(bz, e.Message)
+					require.NoError(t, err)
+					responses = append(responses, e.Message.(*ssproto.SnapshotsResponse))
 				}).Return(true)
 			}
 
@@ -149,7 +170,11 @@ func TestReactor_Receive_SnapshotsRequest(t *testing.T) {
 				}
 			})
 
-			r.Receive(SnapshotChannel, peer, mustEncodeMsg(&ssproto.SnapshotsRequest{}))
+			r.ReceiveEnvelope(p2p.Envelope{
+				ChannelID: SnapshotChannel,
+				Src:       peer,
+				Message:   &ssproto.SnapshotsRequest{},
+			})
 			time.Sleep(100 * time.Millisecond)
 			assert.Equal(t, tc.expectResponses, responses)
 

statesync/syncer.go

@@ -126,7 +126,11 @@ func (s *syncer) AddSnapshot(peer p2p.Peer, snapshot *snapshot) (bool, error) {
 // to discover snapshots, later we may want to do retries and stuff.
 func (s *syncer) AddPeer(peer p2p.Peer) {
 	s.logger.Debug("Requesting snapshots from peer", "peer", peer.ID())
-	peer.Send(SnapshotChannel, mustEncodeMsg(&ssproto.SnapshotsRequest{}))
+	e := p2p.Envelope{
+		ChannelID: SnapshotChannel,
+		Message:   &ssproto.SnapshotsRequest{},
+	}
+	p2p.SendEnvelopeShim(peer, e, s.logger) //nolint: staticcheck
 }
 
 // RemovePeer removes a peer from the pool.
@@ -467,11 +471,14 @@ func (s *syncer) requestChunk(snapshot *snapshot, chunk uint32) {
 	}
 	s.logger.Debug("Requesting snapshot chunk", "height", snapshot.Height,
 		"format", snapshot.Format, "chunk", chunk, "peer", peer.ID())
-	peer.Send(ChunkChannel, mustEncodeMsg(&ssproto.ChunkRequest{
-		Height: snapshot.Height,
-		Format: snapshot.Format,
-		Index:  chunk,
-	}))
+	p2p.SendEnvelopeShim(peer, p2p.Envelope{ //nolint: staticcheck
+		ChannelID: ChunkChannel,
+		Message: &ssproto.ChunkRequest{
+			Height: snapshot.Height,
+			Format: snapshot.Format,
+			Index:  chunk,
+		},
+	}, s.logger)
 }
 
 // verifyApp verifies the sync, checking the app hash, last block height and app version

statesync/syncer_test.go

@@ -98,13 +98,27 @@ func TestSyncer_SyncAny(t *testing.T) {
 	// Adding a couple of peers should trigger snapshot discovery messages
 	peerA := &p2pmocks.Peer{}
 	peerA.On("ID").Return(p2p.ID("a"))
-	peerA.On("Send", SnapshotChannel, mustEncodeMsg(&ssproto.SnapshotsRequest{})).Return(true)
+	peerA.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+		e, ok := i.(p2p.Envelope)
+		if !ok {
+			return false
+		}
+		req, ok := e.Message.(*ssproto.SnapshotsRequest)
+		return ok && e.ChannelID == SnapshotChannel && req != nil
+	})).Return(true)
 	syncer.AddPeer(peerA)
 	peerA.AssertExpectations(t)
 
 	peerB := &p2pmocks.Peer{}
 	peerB.On("ID").Return(p2p.ID("b"))
-	peerB.On("Send", SnapshotChannel, mustEncodeMsg(&ssproto.SnapshotsRequest{})).Return(true)
+	peerB.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+		e, ok := i.(p2p.Envelope)
+		if !ok {
+			return false
+		}
+		req, ok := e.Message.(*ssproto.SnapshotsRequest)
+		return ok && e.ChannelID == SnapshotChannel && req != nil
+	})).Return(true)
 	syncer.AddPeer(peerB)
 	peerB.AssertExpectations(t)
 
@@ -147,9 +161,9 @@ func TestSyncer_SyncAny(t *testing.T) {
 	chunkRequests := make(map[uint32]int)
 	chunkRequestsMtx := tmsync.Mutex{}
 	onChunkRequest := func(args mock.Arguments) {
-		pb, err := decodeMsg(args[1].([]byte))
-		require.NoError(t, err)
-		msg := pb.(*ssproto.ChunkRequest)
+		e, ok := args[0].(p2p.Envelope)
+		require.True(t, ok)
+		msg := e.Message.(*ssproto.ChunkRequest)
 		require.EqualValues(t, 1, msg.Height)
 		require.EqualValues(t, 1, msg.Format)
 		require.LessOrEqual(t, msg.Index, uint32(len(chunks)))
@@ -162,8 +176,14 @@ func TestSyncer_SyncAny(t *testing.T) {
 		chunkRequests[msg.Index]++
 		chunkRequestsMtx.Unlock()
 	}
-	peerA.On("Send", ChunkChannel, mock.Anything).Maybe().Run(onChunkRequest).Return(true)
-	peerB.On("Send", ChunkChannel, mock.Anything).Maybe().Run(onChunkRequest).Return(true)
+	peerA.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+		e, ok := i.(p2p.Envelope)
+		return ok && e.ChannelID == ChunkChannel
+	})).Maybe().Run(onChunkRequest).Return(true)
+	peerB.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+		e, ok := i.(p2p.Envelope)
+		return ok && e.ChannelID == ChunkChannel
+	})).Maybe().Run(onChunkRequest).Return(true)
 
 	// The first time we're applying chunk 2 we tell it to retry the snapshot and discard chunk 1,
 	// which should cause it to keep the existing chunk 0 and 2, and restart restoration from