feat: compress kept trace decision message

collect/collect_test.go

@@ -2254,7 +2254,10 @@ func TestSendDropDecisions(t *testing.T) {
 	coll.dropDecisionBuffer <- TraceDecision{TraceID: "trace1"}
 	close(coll.dropDecisionBuffer)
 	droppedMessage := <-messages
-	assert.Equal(t, "trace1", droppedMessage)
+
+	decompressedData, err := decompressDropDecisions([]byte(droppedMessage))
+	assert.NoError(t, err)
+	assert.Equal(t, "trace1", decompressedData)
 
 	<-closed
 
@@ -2278,7 +2281,10 @@ func TestSendDropDecisions(t *testing.T) {
 	}
 	close(coll.dropDecisionBuffer)
 	droppedMessage = <-messages
-	assert.Equal(t, "trace0,trace1,trace2,trace3,trace4", droppedMessage)
+
+	decompressedData, err = decompressDropDecisions([]byte(droppedMessage))
+	assert.NoError(t, err)
+	assert.Equal(t, "trace0,trace1,trace2,trace3,trace4", decompressedData)
 
 	<-closed
 }

collect/trace_decision.go

@@ -1,10 +1,13 @@
 package collect
 
 import (
-	"encoding/json"
+	"bytes"
+	"encoding/gob"
 	"fmt"
 	"strings"
+	"sync"
 
+	"github.com/golang/snappy"
 	"github.com/honeycombio/refinery/collect/cache"
 )
 
@@ -40,47 +43,46 @@ func newDroppedDecisionMessage(tds []TraceDecision) (string, error) {
 		}
 	}
 
-	if len(traceIDs) == 0 {
-		return "", fmt.Errorf("no valid trace IDs provided")
-	}
-
-	return strings.Join(traceIDs, ","), nil
-}
-func newKeptDecisionMessage(tds []TraceDecision) (string, error) {
-	if len(tds) == 0 {
-		return "", fmt.Errorf("no kept trace decisions provided")
-	}
-
-	data, err := json.Marshal(tds)
+	compressed, err := compress(strings.Join(traceIDs, ","))
 	if err != nil {
 		return "", err
 	}
-
-	return string(data), nil
+	return string(compressed), nil
 }
 
 func newDroppedTraceDecision(msg string) ([]TraceDecision, error) {
-	if msg == "" {
-		return nil, fmt.Errorf("empty drop message")
+	data, err := decompressDropDecisions([]byte(msg))
+	if err != nil {
+		return nil, err
 	}
-	var decisions []TraceDecision
-	for _, traceID := range strings.Split(msg, ",") {
+
+	traceIDs := strings.Split(data, ",")
+	decisions := make([]TraceDecision, 0, len(traceIDs))
+	for _, traceID := range traceIDs {
 		decisions = append(decisions, TraceDecision{
 			TraceID: traceID,
 		})
 	}
-
 	return decisions, nil
 }
 
+func newKeptDecisionMessage(tds []TraceDecision) (string, error) {
+	if len(tds) == 0 {
+		return "", fmt.Errorf("no kept trace decisions provided")
+	}
+	compressed, err := compress(tds)
+	if err != nil {
+		return "", err
+	}
+	return string(compressed), nil
+}
+
 func newKeptTraceDecision(msg string) ([]TraceDecision, error) {
-	keptDecisions := make([]TraceDecision, 0)
-	err := json.Unmarshal([]byte(msg), &keptDecisions)
+	compressed, err := decompressKeptDecisions([]byte(msg))
 	if err != nil {
 		return nil, err
 	}
-
-	return keptDecisions, nil
+	return compressed, nil
 }
 
 var _ cache.KeptTrace = &TraceDecision{}
@@ -92,16 +94,16 @@ type TraceDecision struct {
 	// keptDecision
 	Kept            bool
 	Rate            uint
-	SamplerKey      string `json:",omitempty"`
-	SamplerSelector string `json:",omitempty"`
+	SamplerKey      string
+	SamplerSelector string
 	SendReason      string
 	HasRoot         bool
 	Reason          string
-	Count           uint32 `json:",omitempty"` // number of spans in the trace
-	EventCount      uint32 `json:",omitempty"` // number of span events in the trace
-	LinkCount       uint32 `json:",omitempty"` // number of span links in the trace
+	Count           uint32
+	EventCount      uint32
+	LinkCount       uint32
 
-	keptReasonIdx uint `json:",omitempty"`
+	keptReasonIdx uint
 }
 
 func (td *TraceDecision) DescendantCount() uint32 {
@@ -135,3 +137,72 @@ func (td *TraceDecision) KeptReason() uint {
 func (td *TraceDecision) SetKeptReason(reasonIdx uint) {
 	td.keptReasonIdx = reasonIdx
 }
+
+var bufferPool = sync.Pool{
+	New: func() any { return new(bytes.Buffer) },
+}
+
+var snappyWriterPool = sync.Pool{
+	New: func() any { return snappy.NewBufferedWriter(nil) },
+}
+
+func compress(data any) ([]byte, error) {
+	// Get a buffer from the pool and reset it
+	buf := bufferPool.Get().(*bytes.Buffer)
+	buf.Reset()
+	defer bufferPool.Put(buf)
+
+	// Get a snappy writer from the pool, set it to write to the buffer, and reset it
+	compr := snappyWriterPool.Get().(*snappy.Writer)
+	compr.Reset(buf)
+	defer snappyWriterPool.Put(compr)
+
+	enc := gob.NewEncoder(compr)
+	if err := enc.Encode(data); err != nil {
+		return nil, err
+	}
+
+	// Flush snappy writer
+	if err := compr.Close(); err != nil {
+		return nil, err
+	}
+
+	// Copy the buffer’s bytes to avoid reuse issues when returning
+	return bytes.Clone(buf.Bytes()), nil
+}
+
+func decompressKeptDecisions(data []byte) ([]TraceDecision, error) {
+	// Get a buffer from the pool and set it up with data
+	buf := bufferPool.Get().(*bytes.Buffer)
+	defer bufferPool.Put(buf)
+	buf.Reset()
+	buf.Write(data)
+
+	// Snappy reader to decompress data in buffer
+	reader := snappy.NewReader(buf)
+	dec := gob.NewDecoder(reader)
+
+	var tds []TraceDecision
+	if err := dec.Decode(&tds); err != nil {
+		return nil, err
+	}
+	return tds, nil
+}
+
+func decompressDropDecisions(data []byte) (string, error) {
+	// Get a buffer from the pool and set it up with data
+	buf := bufferPool.Get().(*bytes.Buffer)
+	defer bufferPool.Put(buf)
+	buf.Reset()
+	buf.Write(data)
+
+	// Snappy reader to decompress data in buffer
+	reader := snappy.NewReader(buf)
+	dec := gob.NewDecoder(reader)
+
+	var traceIDs string
+	if err := dec.Decode(&traceIDs); err != nil {
+		return "", err
+	}
+	return traceIDs, nil
+}