From f0b29709e57066b166731c024f78d32ead77c0d1 Mon Sep 17 00:00:00 2001
From: Matheus Degiovani <opensource@matheusd.com>
Date: Thu, 14 Mar 2024 17:47:45 -0300
Subject: [PATCH] message: Fix reuse of first segment

This fixes the Message's Reset() call to allow reuse of the first
segment.

Prior to this fix, the first segment was discarded after the first Reset
call, effectively causing a new segment to be initialized on every Reset
call.

By reusing the first segment, the number of heap allocations is reduced
and therefore performance is increased in use cases where the message
object is reused.

The fix involved associtating the segment to the message and fixing
checks to ensure the data of the segment is re-allocated after the
reset.

A benchmark is included to show the current performance of this.
---
 message.go      | 12 +++++++++---
 message_test.go | 18 ++++++++++++++++++
 segment.go      |  2 ++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/message.go b/message.go
index 5ae7f3e8..d9c24d95 100644
--- a/message.go
+++ b/message.go
@@ -100,6 +100,11 @@ func (m *Message) Release() {
 func (m *Message) Reset(arena Arena) (first *Segment, err error) {
 	m.capTable.Reset()
 	for k := range m.segs {
+		// Optimization: keep the first segment so that the re-used
+		// Message does not have to allocate a new one.
+		if k == 0 && m.segs[k] == &m.firstSeg {
+			continue
+		}
 		delete(m.segs, k)
 	}
 
@@ -113,6 +118,7 @@ func (m *Message) Reset(arena Arena) (first *Segment, err error) {
 		DepthLimit:    m.DepthLimit,
 		capTable:      m.capTable,
 		segs:          m.segs,
+		firstSeg:      Segment{msg: m},
 	}
 
 	if arena != nil {
@@ -264,10 +270,10 @@ func (m *Message) Segment(id SegmentID) (*Segment, error) {
 // segment returns the segment with the given ID, with no bounds
 // checking.  The caller must be holding m.mu.
 func (m *Message) segment(id SegmentID) (*Segment, error) {
-	if m.segs == nil && id == 0 && m.firstSeg.msg != nil {
+	if m.segs == nil && id == 0 && m.firstSeg.msg != nil && m.firstSeg.data != nil {
 		return &m.firstSeg, nil
 	}
-	if s := m.segs[id]; s != nil {
+	if s := m.segs[id]; s != nil && s.data != nil {
 		return s, nil
 	}
 	if len(m.segs) == maxInt {
@@ -442,7 +448,7 @@ func alloc(s *Segment, sz Size) (*Segment, address, error) {
 		var err error
 		s, err = s.msg.allocSegment(sz)
 		if err != nil {
-			return nil, 0, err
+			return nil, 0, errors.New("allocSegment failed: " + err.Error())
 		}
 	}
 
diff --git a/message_test.go b/message_test.go
index b146bd06..b705e2e1 100644
--- a/message_test.go
+++ b/message_test.go
@@ -643,3 +643,21 @@ func (readOnlyArena) Allocate(sz Size, segs map[SegmentID]*Segment) (SegmentID,
 }
 
 var errReadOnlyArena = errors.New("Allocate called on read-only arena")
+
+func BenchmarkMessageGetFirstSegment(b *testing.B) {
+	var msg Message
+	var arena Arena = SingleSegment(nil)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, err := msg.Reset(arena)
+		if err != nil {
+			b.Fatal(err)
+		}
+		_, err = msg.Segment(0)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
diff --git a/segment.go b/segment.go
index 14b4bd4e..119e50cf 100644
--- a/segment.go
+++ b/segment.go
@@ -15,6 +15,8 @@ type SegmentID uint32
 // It is part of a Message, which can contain other segments that
 // reference each other.
 type Segment struct {
+	// msg associated with this segment. A Message instance m maintains the
+	// invariant that that all m.segs[].msg == m.
 	msg  *Message
 	id   SegmentID
 	data []byte