capnp: Refactor Arenas

arena_test.go

@@ -1,10 +1,59 @@
 package capnp
 
 import (
-	"bytes"
 	"testing"
+
+	"capnproto.org/go/capnp/v3/exp/bufferpool"
+	"github.com/stretchr/testify/require"
 )
 
+type arenaAllocTest struct {
+	name string
+
+	// Arrange
+	init func() (Arena, map[SegmentID]*Segment)
+	size Size
+
+	// Assert
+	id   SegmentID
+	data []byte
+}
+
+func (test *arenaAllocTest) run(t *testing.T) {
+	arena, _ := test.init()
+	seg, addr, err := arena.Allocate(test.size, nil, nil)
+
+	require.NoError(t, err, "Allocate error")
+	require.Equal(t, test.id, seg.id)
+
+	// Allocate() contract is that segment data starting at addr should
+	// have anough room for test.size bytes.
+	require.Less(t, int(addr), len(seg.data))
+
+	data := seg.data[addr:]
+	require.LessOrEqual(t, test.size, Size(cap(seg.data)))
+
+	data = data[:test.size]
+	require.Equal(t, test.data, data)
+}
+
+func incrementingData(n int) []byte {
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = byte(i % 256)
+	}
+	return b
+}
+
+func segmentData(a Arena, id SegmentID) []byte {
+	seg := a.Segment(id)
+	if seg == nil {
+		return nil
+	}
+
+	return seg.Data()
+}
+
 func TestSingleSegment(t *testing.T) {
 	t.Parallel()
 	t.Helper()
@@ -13,40 +62,22 @@ func TestSingleSegment(t *testing.T) {
 		t.Parallel()
 
 		arena := SingleSegment(nil)
-		if n := arena.NumSegments(); n != 1 {
-			t.Errorf("SingleSegment(nil).NumSegments() = %d; want 1", n)
-		}
-		data, err := arena.Data(0)
-		if len(data) != 0 {
-			t.Errorf("SingleSegment(nil).Data(0) = %#v; want nil", data)
-		}
-		if err != nil {
-			t.Errorf("SingleSegment(nil).Data(0) error: %v", err)
-		}
-		_, err = arena.Data(1)
-		if err == nil {
-			t.Error("SingleSegment(nil).Data(1) succeeded; want error")
-		}
+		require.Equal(t, int64(1), arena.NumSegments())
+		data0 := segmentData(arena, 0)
+		require.Empty(t, data0)
+		data1 := segmentData(arena, 1)
+		require.Empty(t, data1)
 	})
 
 	t.Run("ExistingData", func(t *testing.T) {
 		t.Parallel()
 
 		arena := SingleSegment(incrementingData(8))
-		if n := arena.NumSegments(); n != 1 {
-			t.Errorf("SingleSegment(incrementingData(8)).NumSegments() = %d; want 1", n)
-		}
-		data, err := arena.Data(0)
-		if want := incrementingData(8); !bytes.Equal(data, want) {
-			t.Errorf("SingleSegment(incrementingData(8)).Data(0) = %#v; want %#v", data, want)
-		}
-		if err != nil {
-			t.Errorf("SingleSegment(incrementingData(8)).Data(0) error: %v", err)
-		}
-		_, err = arena.Data(1)
-		if err == nil {
-			t.Error("SingleSegment(incrementingData(8)).Data(1) succeeded; want error")
-		}
+		require.Equal(t, int64(1), arena.NumSegments())
+		data0 := segmentData(arena, 0)
+		require.Equal(t, incrementingData(8), data0)
+		data1 := segmentData(arena, 1)
+		require.Empty(t, data1)
 	})
 }
 
@@ -61,7 +92,7 @@ func TestSingleSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: []byte{},
+			data: []byte{7: 0},
 		},
 		{
 			name: "unloaded",
@@ -71,7 +102,7 @@ func TestSingleSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(16),
+			data: incrementingData(24)[16 : 16+8],
 		},
 		{
 			name: "loaded",
@@ -85,7 +116,7 @@ func TestSingleSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(16),
+			data: incrementingData(24)[16 : 16+8],
 		},
 		{
 			name: "loaded changes length",
@@ -98,7 +129,7 @@ func TestSingleSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(24),
+			data: incrementingData(32)[16 : 16+8],
 		},
 		{
 			name: "message-filled segment",
@@ -111,11 +142,12 @@ func TestSingleSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(24),
+			data: incrementingData(24)[16 : 16+8],
 		},
 	}
 	for i := range tests {
-		tests[i].run(t, i)
+		tc := tests[i]
+		t.Run(tc.name, tc.run)
 	}
 }
 
@@ -127,40 +159,22 @@ func TestMultiSegment(t *testing.T) {
 		t.Parallel()
 
 		arena := MultiSegment(nil)
-		if n := arena.NumSegments(); n != 0 {
-			t.Errorf("MultiSegment(nil).NumSegments() = %d; want 1", n)
-		}
-		_, err := arena.Data(0)
-		if err == nil {
-			t.Error("MultiSegment(nil).Data(0) succeeded; want error")
-		}
+		require.Equal(t, int64(0), arena.NumSegments())
+		data0 := segmentData(arena, 0)
+		require.Empty(t, data0)
 	})
 
 	t.Run("ExistingData", func(t *testing.T) {
 		t.Parallel()
 
 		arena := MultiSegment([][]byte{incrementingData(8), incrementingData(24)})
-		if n := arena.NumSegments(); n != 2 {
-			t.Errorf("MultiSegment(...).NumSegments() = %d; want 2", n)
-		}
-		data, err := arena.Data(0)
-		if want := incrementingData(8); !bytes.Equal(data, want) {
-			t.Errorf("MultiSegment(...).Data(0) = %#v; want %#v", data, want)
-		}
-		if err != nil {
-			t.Errorf("MultiSegment(...).Data(0) error: %v", err)
-		}
-		data, err = arena.Data(1)
-		if want := incrementingData(24); !bytes.Equal(data, want) {
-			t.Errorf("MultiSegment(...).Data(1) = %#v; want %#v", data, want)
-		}
-		if err != nil {
-			t.Errorf("MultiSegment(...).Data(1) error: %v", err)
-		}
-		_, err = arena.Data(2)
-		if err == nil {
-			t.Error("MultiSegment(...).Data(2) succeeded; want error")
-		}
+		require.Equal(t, int64(2), arena.NumSegments())
+		data0 := segmentData(arena, 0)
+		require.Equal(t, incrementingData(8), data0)
+		data1 := segmentData(arena, 1)
+		require.Equal(t, incrementingData(24), data1)
+		data2 := segmentData(arena, 2)
+		require.Empty(t, data2)
 	})
 }
 
@@ -175,7 +189,7 @@ func TestMultiSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: []byte{},
+			data: []byte{7: 0},
 		},
 		{
 			name: "space in unloaded segment",
@@ -185,7 +199,7 @@ func TestMultiSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(16),
+			data: incrementingData(24)[16 : 16+8],
 		},
 		{
 			name: "space in loaded segment",
@@ -199,7 +213,7 @@ func TestMultiSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(16),
+			data: incrementingData(24)[16 : 16+8],
 		},
 		{
 			name: "space in loaded segment changes length",
@@ -212,24 +226,27 @@ func TestMultiSegmentAllocate(t *testing.T) {
 			},
 			size: 8,
 			id:   0,
-			data: incrementingData(24),
+			data: incrementingData(24)[16 : 16+8],
 		},
 		{
-			name: "message-filled segment",
+			name: "first segment is filled",
 			init: func() (Arena, map[SegmentID]*Segment) {
 				buf := incrementingData(24)
 				segs := map[SegmentID]*Segment{
 					0: {id: 0, data: buf},
 				}
-				return MultiSegment([][]byte{buf[:16]}), segs
+				msa := MultiSegment([][]byte{buf[:16:16]})
+				msa.bp = &bufferpool.Default
+				return msa, segs
 			},
 			size: 8,
 			id:   1,
-			data: []byte{},
+			data: []byte{7: 0},
 		},
 	}
 
 	for i := range tests {
-		tests[i].run(t, i)
+		tc := tests[i]
+		t.Run(tc.name, tc.run)
 	}
 }

capability.go

@@ -59,7 +59,7 @@ func (i Interface) Message() *Message {
 	if i.seg == nil {
 		return nil
 	}
-	return i.seg.msg
+	return i.seg.Message()
 }
 
 // IsValid returns whether the interface is valid.

codec_test.go

@@ -81,12 +81,12 @@ type tooManySegsArena struct {
 
 func (t *tooManySegsArena) NumSegments() int64 { return 1<<32 + 1 }
 
-func (t *tooManySegsArena) Data(id SegmentID) ([]byte, error) {
-	return nil, errors.New("no data")
+func (t *tooManySegsArena) Segment(id SegmentID) *Segment {
+	return nil
 }
 
-func (t *tooManySegsArena) Allocate(minsz Size, segs map[SegmentID]*Segment) (SegmentID, []byte, error) {
-	return 0, nil, errors.New("cannot allocate")
+func (t *tooManySegsArena) Allocate(minsz Size, msg *Message, seg *Segment) (*Segment, address, error) {
+	return nil, 0, errors.New("cannot allocate")
 }
 
 func (t *tooManySegsArena) Release() {}

integration_test.go

@@ -1851,13 +1851,13 @@ func BenchmarkUnmarshal_Reuse(b *testing.B) {
 		data[i].data, data[i].a = buf, *a
 	}
 	msg := new(capnp.Message)
-	ta := new(testArena)
+	ta := capnp.NewReadOnlySingleSegment(nil)
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		testIdx := r.Intn(len(data))
-		*ta = testArena(data[testIdx].data[8:]) // [8:] to skip header
 		msg.Release()
+		ta.ReplaceData(data[testIdx].data[8:]) // [8:] to skip header
 		msg.Arena = ta
 		a, err := air.ReadRootBenchmarkA(msg)
 		if err != nil {
@@ -1913,25 +1913,6 @@ func BenchmarkDecode(b *testing.B) {
 	}
 }
 
-type testArena []byte
-
-func (ta testArena) NumSegments() int64 {
-	return 1
-}
-
-func (ta testArena) Data(id capnp.SegmentID) ([]byte, error) {
-	if id != 0 {
-		return nil, errors.New("test arena: requested non-zero segment")
-	}
-	return []byte(ta), nil
-}
-
-func (ta testArena) Allocate(capnp.Size, map[capnp.SegmentID]*capnp.Segment) (capnp.SegmentID, []byte, error) {
-	return 0, nil, errors.New("test arena: can't allocate")
-}
-
-func (ta testArena) Release() {}
-
 func TestPointerTraverseDefense(t *testing.T) {
 	t.Parallel()
 	const limit = 128

list.go

@@ -100,7 +100,7 @@ func (p List) Message() *Message {
 	if p.seg == nil {
 		return nil
 	}
-	return p.seg.msg
+	return p.seg.Message()
 }
 
 // IsValid returns whether the list is valid.