diff --git a/pkg/network/message.go b/pkg/network/message.go index 5eb360d9d..cbdaa554f 100644 --- a/pkg/network/message.go +++ b/pkg/network/message.go @@ -182,6 +182,11 @@ func (m *Message) Encode(br *io.BinWriter) error { if err := m.tryCompressPayload(); err != nil { return err } + growSize := 2 + 1 // header + empty payload + if m.compressedPayload != nil { + growSize += 8 + len(m.compressedPayload) // varint + byte-slice + } + br.Grow(growSize) br.WriteB(byte(m.Flags)) br.WriteB(byte(m.Command)) if m.compressedPayload != nil { diff --git a/pkg/network/message_test.go b/pkg/network/message_test.go index 879d4a9f7..c4976a702 100644 --- a/pkg/network/message_test.go +++ b/pkg/network/message_test.go @@ -52,6 +52,34 @@ func TestEncodeDecodeVersion(t *testing.T) { require.NotEqual(t, len(expected.compressedPayload), len(uncompressed)) } +func BenchmarkMessageBytes(b *testing.B) { + // shouldn't try to compress headers payload + ep := &payload.Extensible{ + Category: "consensus", + ValidBlockStart: rand.Uint32(), + ValidBlockEnd: rand.Uint32(), + Sender: util.Uint160{}, + Data: make([]byte, 300), + Witness: transaction.Witness{ + InvocationScript: make([]byte, 33), + VerificationScript: make([]byte, 40), + }, + } + random.Fill(ep.Data) + random.Fill(ep.Witness.InvocationScript) + random.Fill(ep.Witness.VerificationScript) + msg := NewMessage(CMDExtensible, ep) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := msg.Bytes() + if err != nil { + b.FailNow() + } + } +} + func TestEncodeDecodeHeaders(t *testing.T) { // shouldn't try to compress headers payload headers := &payload.Headers{Hdrs: make([]*block.Header, CompressionMinSize)}