perf(header): MsgID func only needs to unmarshal Commit from ExtendedHeader (#4149)

Author: renaynay

header/headertest/serde_test.go

@@ -3,10 +3,13 @@ package headertest
 import (
 	"testing"
 
+	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/blake2b"
 
 	"github.com/celestiaorg/celestia-node/header"
+	"github.com/celestiaorg/celestia-node/share/eds/edstest"
 )
 
 func TestMarshalUnmarshalExtendedHeader(t *testing.T) {
@@ -30,6 +33,51 @@ func TestMarshalUnmarshalExtendedHeader(t *testing.T) {
 	equalExtendedHeader(t, in, out)
 }
 
+func TestMsgIDEquivalency(t *testing.T) {
+	randHeader := RandExtendedHeader(t)
+	bin, err := randHeader.MarshalBinary()
+	require.NoError(t, err)
+
+	oldMsgIDFunc := func(message *pubsub_pb.Message) string {
+		mID := func(data []byte) string {
+			hash := blake2b.Sum256(data)
+			return string(hash[:])
+		}
+
+		h, _ := header.UnmarshalExtendedHeader(message.Data)
+		if h == nil || h.RawHeader.ValidateBasic() != nil {
+			return mID(message.Data)
+		}
+
+		return h.Commit.BlockID.String()
+	}
+
+	inboundMsg := &pubsub_pb.Message{Data: bin}
+
+	expectedHash := oldMsgIDFunc(inboundMsg)
+	gotHash := header.MsgID(inboundMsg)
+
+	assert.Equal(t, expectedHash, gotHash)
+}
+
+// Before changes (with 256 EDS and 100 validators):
+// BenchmarkMsgID-8   	    5203	    224681 ns/op	  511253 B/op	    4252 allocs/op
+// After changes (with 256 EDS and 100 validators):
+// BenchmarkMsgID-8   	   23559	     48399 ns/op	  226858 B/op	    1282 allocs/op
+func BenchmarkMsgID(b *testing.B) {
+	eds := edstest.RandomAxisRoots(b, 256)
+	randHeader := RandExtendedHeaderWithRoot(b, eds)
+	bin, err := randHeader.MarshalBinary()
+	require.NoError(b, err)
+	msg := &pubsub_pb.Message{Data: bin}
+
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		_ = header.MsgID(msg)
+	}
+}
+
 func equalExtendedHeader(t *testing.T, in, out *header.ExtendedHeader) {
 	// ValidatorSet.totalVotingPower is not set (is a cached value that can be recomputed client side)
 	assert.Equal(t, in.ValidatorSet.Validators, out.ValidatorSet.Validators)

header/serde.go

@@ -64,18 +64,30 @@ func UnmarshalExtendedHeader(data []byte) (*ExtendedHeader, error) {
 	return out, nil
 }
 
-// msgID computes an id for a pubsub message
+// unmarhsalCommit exists to assist the MsgID function in generating a unique
+// message ID without the additional allocations that the full
+// UnmarshalExtendedHeader would cause.
+func unmarshalCommit(data []byte) (*core.Commit, error) {
+	in := &header_pb.ExtendedHeader{}
+	err := in.Unmarshal(data)
+	if err != nil {
+		return nil, err
+	}
+
+	return core.CommitFromProto(in.Commit)
+}
+
+// MsgID computes an id for a pubsub message
 // TODO(@Wondertan): This cause additional allocations per each recvd message in the topic
-//
-//	Find a way to avoid those.
+// TODO(@renaynay): We will still allocate now but we're minimizing surface only to Commit
 func MsgID(pmsg *pb.Message) string {
 	mID := func(data []byte) string {
 		hash := blake2b.Sum256(data)
 		return string(hash[:])
 	}
 
-	h, _ := UnmarshalExtendedHeader(pmsg.Data)
-	if h == nil || h.RawHeader.ValidateBasic() != nil {
+	commit, err := unmarshalCommit(pmsg.Data)
+	if commit == nil || err != nil {
 		// There is nothing we can do about the error, and it will be anyway caught during validation.
 		// We also *have* to return some ID for the msg, so give the hash of even faulty msg
 		return mID(pmsg.Data)
@@ -92,5 +104,5 @@ func MsgID(pmsg *pb.Message) string {
 	//
 	// To solve the nondeterminism problem above, we don't compute msg id on message body and take
 	// the actual header hash as an id.
-	return h.Commit.BlockID.String()
+	return commit.BlockID.String()
 }