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

@@ -3,6 +3,7 @@ package capnp
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"io"
 	"sync"
 	"sync/atomic"
@@ -100,6 +101,9 @@ func (m *Message) Release() {
 func (m *Message) Reset(arena Arena) (first *Segment, err error) {
 	m.capTable.Reset()
 	for k := range m.segs {
+		if k == 0 && m.segs[k] == &m.firstSeg {
+			continue
+		}
 		delete(m.segs, k)
 	}
 
@@ -113,6 +117,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 +269,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 +447,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, fmt.Errorf("allocSegment failed: %v", err)
 		}
 	}
 

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)
+		}
+	}
+}

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