Repository: honeycombio/refinery
Pull request #1594 (2025-06-18 23:48:21)

test: enable DryRun in collect loop benchmark test
## Which problem is this PR solving?

`EMADynamic` sampler can't maintain at SampleRate 1 due to probability. To ensure the consistency of the test result, enabling dyrun will allow all traces to be sent regardless of its sampling decision. 

## Short description of the changes

- enable `DryRun` in collect loop benchmark test




Patch:
diff --git a/route/batched_event.go b/route/batched_event.go
new file mode 100644
index 0000000..40d1f34
--- /dev/null
+++ b/route/batched_event.go
@@ -0,0 +1,120 @@
+package route
+
+import (
+	"bytes"
+	"time"
+
+	"github.com/honeycombio/refinery/types"
+	"github.com/tinylib/msgp/msgp"
+)
+
+type batchedEvent struct {
+	Timestamp        string        `json:"time"`
+	MsgPackTimestamp *time.Time    `msgpack:"time,omitempty"`
+	SampleRate       int64         `json:"samplerate" msgpack:"samplerate"`
+	Data             types.Payload `json:"data" msgpack:"data"`
+}
+
+func (b *batchedEvent) getEventTime() time.Time {
+	if b.MsgPackTimestamp != nil {
+		return b.MsgPackTimestamp.UTC()
+	}
+
+	return getEventTime(b.Timestamp)
+}
+
+func (b *batchedEvent) getSampleRate() uint {
+	if b.SampleRate == 0 {
+		return defaultSampleRate
+	}
+	return uint(b.SampleRate)
+}
+
+// UnmarshalMsg implements msgp.Unmarshaler
+// Based on generated code from the msgp tool, but modified for clarity and to
+// avoid the unnecessary unsafe string conversion.
+// For instructions on how to generate this sort of code, see
+// https://pkg.go.dev/github.com/tinylib/msgp#section-readme
+func (b *batchedEvent) UnmarshalMsg(bts []byte) (o []byte, err error) {
+	var field []byte
+	var fieldsRemaining uint32
+	fieldsRemaining, bts, err = msgp.ReadMapHeaderBytes(bts)
+	if err != nil {
+		err = msgp.WrapError(err)
+		return
+	}
+	for fieldsRemaining > 0 {
+		fieldsRemaining--
+		field, bts, err = msgp.ReadMapKeyZC(bts)
+		if err != nil {
+			err = msgp.WrapError(err)
+			return
+		}
+		switch {
+		case bytes.Equal(field, []byte("time")):
+			if msgp.IsNil(bts) {
+				bts, err = msgp.ReadNilBytes(bts)
+				if err != nil {
+					return
+				}
+				b.MsgPackTimestamp = nil
+			} else {
+				if b.MsgPackTimestamp == nil {
+					b.MsgPackTimestamp = new(time.Time)
+				}
+				*b.MsgPackTimestamp, bts, err = msgp.ReadTimeBytes(bts)
+				if err != nil {
+					err = msgp.WrapError(err, "MsgPackTimestamp")
+					return
+				}
+			}
+		case bytes.Equal(field, []byte("samplerate")):
+			b.SampleRate, bts, err = msgp.ReadInt64Bytes(bts)
+			if err != nil {
+				err = msgp.WrapError(err, "SampleRate")
+				return
+			}
+		case bytes.Equal(field, []byte("data")):
+			bts, err = b.Data.UnmarshalMsg(bts)
+			if err != nil {
+				err = msgp.WrapError(err, "Data")
+				return
+			}
+		default:
+			bts, err = msgp.Skip(bts)
+			if err != nil {
+				err = msgp.WrapError(err)
+				return
+			}
+		}
+	}
+	o = bts
+	return
+}
+
+// Create a type for []batchedEvent so we can give it an unmarshaler.
+type batchedEvents []batchedEvent
+
+// UnmarshalMsg implements msgp.Unmarshaler
+func (b *batchedEvents) UnmarshalMsg(bts []byte) (o []byte, err error) {
+	var totalValues uint32
+	totalValues, bts, err = msgp.ReadArrayHeaderBytes(bts)
+	if err != nil {
+		err = msgp.WrapError(err)
+		return
+	}
+	if cap(*b) >= int(totalValues) {
+		*b = (*b)[:totalValues]
+	} else {
+		*b = make(batchedEvents, totalValues)
+	}
+	for i := range *b {
+		bts, err = (*b)[i].UnmarshalMsg(bts)
+		if err != nil {
+			err = msgp.WrapError(err, i)
+			return
+		}
+	}
+	o = bts
+	return
+}
diff --git a/route/route.go b/route/route.go
index d8482d4..b6d60b2 100644
--- a/route/route.go
+++ b/route/route.go
@@ -22,6 +22,7 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/klauspost/compress/zstd"
 	"github.com/pelletier/go-toml/v2"
+	"github.com/tinylib/msgp/msgp"
 	"github.com/vmihailenco/msgpack/v5"
 	"google.golang.org/grpc"
 	healthserver "google.golang.org/grpc/health"
@@ -439,7 +440,7 @@ func (r *Router) requestToEvent(ctx context.Context, req *http.Request, reqBod [
 	}
 
 	data := map[string]interface{}{}
-	err = unmarshal(req, bytes.NewReader(reqBod), &data)
+	err = unmarshal(req, reqBod, &data)
 	if err != nil {
 		return nil, err
 	}
@@ -476,8 +477,8 @@ func (r *Router) batch(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	batchedEvents := make([]batchedEvent, 0)
-	err = unmarshal(req, bytes.NewReader(reqBod), &batchedEvents)
+	batchedEvents := make(batchedEvents, 0)
+	err = unmarshal(req, reqBod, &batchedEvents)
 	if err != nil {
 		debugLog.WithField("error", err.Error()).WithField("request.url", req.URL).WithField("json_body", string(reqBod)).Logf("error parsing json")
 		r.handlerReturnWithError(w, ErrJSONFailed, err)
@@ -732,28 +733,6 @@ func (r *Router) getMaybeCompressedBody(req *http.Request) (io.Reader, error) {
 	return reader, nil
 }
 
-type batchedEvent struct {
-	Timestamp        string        `json:"time"`
-	MsgPackTimestamp *time.Time    `msgpack:"time,omitempty"`
-	SampleRate       int64         `json:"samplerate" msgpack:"samplerate"`
-	Data             types.Payload `json:"data" msgpack:"data"`
-}
-
-func (b *batchedEvent) getEventTime() time.Time {
-	if b.MsgPackTimestamp != nil {
-		return b.MsgPackTimestamp.UTC()
-	}
-
-	return getEventTime(b.Timestamp)
-}
-
-func (b *batchedEvent) getSampleRate() uint {
-	if b.SampleRate == 0 {
-		return defaultSampleRate
-	}
-	return uint(b.SampleRate)
-}
-
 // getEventTime tries to guess the time format in our time header!
 // Allowable options are
 // * RFC3339Nano
@@ -815,14 +794,18 @@ func makeDecoders(num int) (chan *zstd.Decoder, error) {
 	return zstdDecoders, nil
 }
 
-func unmarshal(r *http.Request, data io.Reader, v interface{}) error {
+func unmarshal(r *http.Request, data []byte, v interface{}) error {
 	switch r.Header.Get("Content-Type") {
 	case "application/x-msgpack", "application/msgpack":
-		decoder := msgpack.NewDecoder(data)
+		if unmarshaler, ok := v.(msgp.Unmarshaler); ok {
+			_, err := unmarshaler.UnmarshalMsg(data)
+			return err
+		}
+		decoder := msgpack.NewDecoder(bytes.NewReader(data))
 		decoder.UseLooseInterfaceDecoding(true)
 		return decoder.Decode(v)
 	default:
-		return jsoniter.NewDecoder(data).Decode(v)
+		return jsoniter.Unmarshal(data, v)
 	}
 }
 
diff --git a/types/msgp_payload_map.go b/types/msgp_payload_map.go
index a62e4d8..600f14f 100644
--- a/types/msgp_payload_map.go
+++ b/types/msgp_payload_map.go
@@ -41,7 +41,7 @@ func (m *MsgpPayloadMap) Size() int {
 }
 
 func (m *MsgpPayloadMap) Iterate() (msgpPayloadMapIter, error) {
-	if m == nil {
+	if m == nil || len(m.rawData) == 0 {
 		return msgpPayloadMapIter{}, nil
 	}
 	n, remaining, err := msgp.ReadMapHeaderBytes(m.rawData)
diff --git a/types/payload.go b/types/payload.go
index a8bee6f..76cd8ee 100644
--- a/types/payload.go
+++ b/types/payload.go
@@ -6,12 +6,13 @@ import (
 	"maps"
 
 	jsoniter "github.com/json-iterator/go"
+	"github.com/tinylib/msgp/msgp"
 	"github.com/vmihailenco/msgpack/v5"
 )
 
 type Payload struct {
 	// A serialized messagepack map used to source fields.
-	msgpMap *MsgpPayloadMap
+	msgpMap MsgpPayloadMap
 
 	// Deserialized fields, either from the internal msgpMap, or set externally.
 	memoizedFields map[string]any
@@ -23,12 +24,27 @@ func NewPayload(data map[string]any) Payload {
 	}
 }
 
+// UnmarshalMsgpack implements msgpack.Unmarshaler, but doesn't unmarshal. Instead it
+// keep a reference to serialized data.
 func (p *Payload) UnmarshalMsgpack(data []byte) error {
-	p.msgpMap = &MsgpPayloadMap{rawData: data}
-	p.memoizedFields = make(map[string]any)
+	p.msgpMap = MsgpPayloadMap{rawData: data}
 	return nil
 }
 
+// UnmarshalMsg implements msgp.Unmarshaler, similar to above but expects to be
+// part of a larger message.
+func (p *Payload) UnmarshalMsg(bts []byte) (o []byte, err error) {
+	// Oddly the msgp library doesn't export the internal size method it uses
+	// to skip data. So we will derived it from the returned slice.
+	remainder, err := msgp.Skip(bts)
+	if err != nil {
+		return nil, err
+	}
+	ourData := bts[:len(bts)-len(remainder)]
+	p.msgpMap = MsgpPayloadMap{rawData: ourData}
+	return remainder, err
+}
+
 func (p *Payload) UnmarshalJSON(data []byte) error {
 	var fields map[string]any
 	if err := jsoniter.Unmarshal(data, &fields); err != nil {
@@ -41,10 +57,6 @@ func (p *Payload) UnmarshalJSON(data []byte) error {
 // Extracts all of the listed fields from the internal msgp buffer in a single
 // pass, for efficient random access later.
 func (p *Payload) MemoizeFields(keys ...string) {
-	if p.msgpMap == nil {
-		return
-	}
-
 	if p.memoizedFields == nil {
 		p.memoizedFields = make(map[string]any, len(keys))
 	}
@@ -152,7 +164,7 @@ func (p *Payload) Set(key string, value any) {
 }
 
 func (p *Payload) IsEmpty() bool {
-	return len(p.memoizedFields) == 0 && p.msgpMap == nil
+	return len(p.memoizedFields) == 0 && p.msgpMap.Size() == 0
 }
 
 // All() allows easily iterating all values in the Payload, but this is very
@@ -203,10 +215,7 @@ func (p *Payload) All() iter.Seq2[string, any] {
 
 // Estimates data size, not very accurately, but it's fast.
 func (p *Payload) GetDataSize() int {
-	var total int
-	if p.msgpMap != nil {
-		total += p.msgpMap.Size()
-	}
+	total := p.msgpMap.Size()
 	for k, v := range p.memoizedFields {
 		total += len(k) + getByteSize(v)
 	}


Test patch:
diff --git a/collect/collect_test.go b/collect/collect_test.go
index 622c53f..a81cfc3 100644
--- a/collect/collect_test.go
+++ b/collect/collect_test.go
@@ -2724,6 +2724,7 @@ func BenchmarkCollectorWithSamplers(b *testing.B) {
 
 func setupBenchmarkCollector(b *testing.B, samplerConfig interface{}, sender *mockSender) *InMemCollector {
 	conf := &config.MockConfig{
+		DryRun: true,
 		GetTracesConfigVal: config.TracesConfig{
 			SendTicker:   config.Duration(5 * time.Millisecond),
 			SendDelay:    config.Duration(1 * time.Millisecond),
diff --git a/route/batched_event_test.go b/route/batched_event_test.go
new file mode 100644
index 0000000..af517dc
--- /dev/null
+++ b/route/batched_event_test.go
@@ -0,0 +1,69 @@
+package route
+
+import (
+	"maps"
+	"testing"
+	"time"
+
+	"github.com/honeycombio/refinery/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/vmihailenco/msgpack/v5"
+)
+
+func TestBatchedEventRoundTrip(t *testing.T) {
+	timestamp := time.Now()
+
+	events := batchedEvents{
+		{
+			MsgPackTimestamp: &timestamp,
+			Data: types.NewPayload(map[string]any{
+				"string_field": "test_value",
+				"int_field":    int64(42),
+				"float_field":  3.14159,
+				"bool_field":   true,
+				"nil_field":    nil,
+			}),
+		},
+		{
+			MsgPackTimestamp: &timestamp,
+			SampleRate:       200,
+			Data: types.NewPayload(map[string]any{
+				"another_string": "hello world",
+				"negative_int":   int64(-123),
+				"large_float":    1234.5678,
+				"false_bool":     false,
+				"nil_value":      nil,
+			}),
+		},
+	}
+
+	serialized, err := msgpack.Marshal(events)
+	require.NoError(t, err)
+
+	var deserialized batchedEvents
+	_, err = deserialized.UnmarshalMsg(serialized)
+	require.NoError(t, err)
+
+	require.Len(t, deserialized, 2)
+	for i := range deserialized {
+		assert.WithinDuration(t, timestamp, *deserialized[i].MsgPackTimestamp, 0)
+		assert.Equal(t, events[i].SampleRate, deserialized[i].SampleRate)
+		assert.Equal(t, maps.Collect(events[i].Data.All()), maps.Collect(deserialized[i].Data.All()))
+	}
+
+	// Test overwriting with another unmarshal, with swapped events order.
+	events[0], events[1] = events[1], events[0]
+	serialized, err = msgpack.Marshal(events)
+	require.NoError(t, err)
+
+	deserialized = deserialized[:0]
+	_, err = deserialized.UnmarshalMsg(serialized)
+	require.NoError(t, err)
+	require.Len(t, deserialized, 2)
+	for i := range deserialized {
+		assert.WithinDuration(t, timestamp, *deserialized[i].MsgPackTimestamp, 0)
+		assert.Equal(t, events[i].SampleRate, deserialized[i].SampleRate)
+		assert.Equal(t, maps.Collect(events[i].Data.All()), maps.Collect(deserialized[i].Data.All()))
+	}
+}
diff --git a/route/route_test.go b/route/route_test.go
index e014bdd..ce57542 100644
--- a/route/route_test.go
+++ b/route/route_test.go
@@ -139,7 +139,7 @@ func TestUnmarshal(t *testing.T) {
 		req.Header.Set("Content-Type", "nope")
 
 		var data map[string]interface{}
-		err := unmarshal(req, req.Body, &data)
+		err := unmarshal(req, readAll(t, req.Body), &data)
 		// Should succeed because invalid content type defaults to JSON
 		assert.NoError(t, err)
 	})
@@ -156,7 +156,7 @@ func TestUnmarshal(t *testing.T) {
 					req.Header.Set("Content-Type", contentType)
 
 					var result map[string]interface{}
-					err = unmarshal(req, req.Body, &result)
+					err = unmarshal(req, readAll(t, req.Body), &result)
 					require.NoError(t, err)
 
 					// Compare directly to test data
@@ -177,7 +177,7 @@ func TestUnmarshal(t *testing.T) {
 					req.Header.Set("Content-Type", contentType)
 
 					var result map[string]interface{}
-					err = unmarshal(req, req.Body, &result)
+					err = unmarshal(req, readAll(t, req.Body), &result)
 					require.NoError(t, err)
 
 					// Compare directly to test data
@@ -187,10 +187,10 @@ func TestUnmarshal(t *testing.T) {
 		})
 	})
 
-	// Test []batchedEvent unmarshaling (used in batch)
-	t.Run("[]batchedEvent", func(t *testing.T) {
+	// Test batchedEvents unmarshaling (used in batch)
+	t.Run("batchedEvents", func(t *testing.T) {
 		t.Run("json", func(t *testing.T) {
-			batchEvents := []batchedEvent{
+			batchEvents := batchedEvents{
 				{
 					Timestamp:  now.Format(time.RFC3339Nano),
 					SampleRate: 2,
@@ -210,8 +210,8 @@ func TestUnmarshal(t *testing.T) {
 					req := httptest.NewRequest("POST", "/test", bytes.NewReader(jsonData))
 					req.Header.Set("Content-Type", contentType)
 
-					var result []batchedEvent
-					err = unmarshal(req, req.Body, &result)
+					var result batchedEvents
+					err = unmarshal(req, readAll(t, req.Body), &result)
 					require.NoError(t, err)
 					require.Len(t, result, 2)
 
@@ -258,8 +258,8 @@ func TestUnmarshal(t *testing.T) {
 					req := httptest.NewRequest("POST", "/test", buf)
 					req.Header.Set("Content-Type", contentType)
 
-					var result []batchedEvent
-					err = unmarshal(req, req.Body, &result)
+					var result batchedEvents
+					err = unmarshal(req, readAll(t, req.Body), &result)
 					require.NoError(t, err)
 					require.Len(t, result, 2)
 
@@ -983,9 +983,9 @@ func newBatchRouter(t testing.TB) *Router {
 	}
 }
 
-func createBatchEvents() []batchedEvent {
+func createBatchEvents() batchedEvents {
 	now := time.Now().UTC()
-	batchEvents := []batchedEvent{
+	batchEvents := batchedEvents{
 		{
 			Timestamp:  now.Format(time.RFC3339Nano),
 			SampleRate: 2,
@@ -1144,3 +1144,9 @@ func BenchmarkRouterBatch(b *testing.B) {
 		}
 	}
 }
+
+func readAll(t testing.TB, r io.Reader) []byte {
+	got, err := io.ReadAll(r)
+	require.NoError(t, err)
+	return got
+}
diff --git a/types/msgp_payload_map_test.go b/types/msgp_payload_map_test.go
index 9cff043..0032fc2 100644
--- a/types/msgp_payload_map_test.go
+++ b/types/msgp_payload_map_test.go
@@ -150,9 +150,11 @@ func TestMsgpPayloadMap_ErrorCases(t *testing.T) {
 	_, err = arrayPayloadMap.Iterate()
 	assert.Error(t, err)
 
-	// Test empty data
+	// Test empty data - iterating is ok, but EOFs immediately.
 	emptyPayloadMap := NewMessagePackPayloadMap(nil)
-	_, err = emptyPayloadMap.Iterate()
+	iter, err := emptyPayloadMap.Iterate()
+	assert.NoError(t, err)
+	_, _, err = iter.NextKey()
 	assert.Error(t, err)
 
 	// Test truncated data, should fail on the second value.
@@ -182,7 +184,7 @@ func TestMsgpPayloadMap_ErrorCases(t *testing.T) {
 	require.NoError(t, err)
 
 	payloadMap := NewMessagePackPayloadMap(encoded)
-	iter, err := payloadMap.Iterate()
+	iter, err = payloadMap.Iterate()
 	require.NoError(t, err)
 
 	for {
diff --git a/types/payload_test.go b/types/payload_test.go
index d7205c7..09189fb 100644
--- a/types/payload_test.go
+++ b/types/payload_test.go
@@ -62,11 +62,21 @@ func TestPayload(t *testing.T) {
 	ph = NewPayload(data)
 	t.Run("from_map", doTest)
 
+	ph = Payload{}
 	msgpData, err := msgpack.Marshal(data)
 	require.NoError(t, err)
 	err = msgpack.Unmarshal(msgpData, &ph)
 	require.NoError(t, err)
+	t.Run("from_msgpack", doTest)
+
+	// Test payload with other stuff (another payload) following.
+	ph = Payload{}
+	extendedMsgpData := append(msgpData, msgpData...)
+	remainder, err := ph.UnmarshalMsg(extendedMsgpData)
+	require.NoError(t, err)
+	assert.Equal(t, msgpData, remainder)
 	t.Run("from_msgp", doTest)
+
 }
 
 func BenchmarkPayload(b *testing.B) {
@@ -99,13 +109,20 @@ func BenchmarkPayload(b *testing.B) {
 		}
 	})
 
-	b.Run("create_msgp", func(b *testing.B) {
+	b.Run("create_msgpack", func(b *testing.B) {
 		for b.Loop() {
 			var phMsgp Payload
 			_ = phMsgp.UnmarshalMsgpack(msgpData)
 		}
 	})
 
+	b.Run("create_msgp", func(b *testing.B) {
+		for b.Loop() {
+			var phMsgp Payload
+			_, _ = phMsgp.UnmarshalMsg(msgpData)
+		}
+	})
+
 	for _, num := range []int{1, 5, 10, 25} {
 		b.Run(fmt.Sprintf("get_map/%d", num), func(b *testing.B) {
 			modulo := len(keys) - num