Preallocate byte slices from prom pool

integration/microservices/prometheus.yaml

@@ -18,4 +18,5 @@ scrape_configs:
     - targets:
       - distributor:3100
       - ingester-0:3100
-      - ingester-1:3100
+      - ingester-1:3100
+      - ingester-2:3100

integration/microservices/tempo.yaml

@@ -22,7 +22,7 @@ memberlist:
   - ingester-2:7946
 
 overrides:
-  ingestion_burst_size: 100000
+  ingestion_burst_size_bytes: 10_000_000
   max_traces_per_user: 1000000
 
 server:

modules/distributor/distributor.go

@@ -268,11 +268,11 @@ func (d *Distributor) sendToIngestersViaBytes(ctx context.Context, userID string
 		localCtx = user.InjectOrgID(localCtx, userID)
 
 		req := tempopb.PushBytesRequest{
-			Requests: make([][]byte, len(indexes)),
+			Requests: make([]tempopb.PreallocRequest, len(indexes)),
 		}
 
 		for i, j := range indexes {
-			req.Requests[i] = rawRequests[j][0:]
+			req.Requests[i].Request = rawRequests[j][0:]
 		}
 
 		c, err := d.pool.GetClientFor(ingester.Addr)

modules/ingester/ingester.go

@@ -181,7 +181,7 @@ func (i *Ingester) PushBytes(ctx context.Context, req *tempopb.PushBytesRequest)
 	// Unmarshal and push each request
 	for _, v := range req.Requests {
 		r := tempopb.PushRequest{}
-		err := r.Unmarshal(v)
+		err := r.Unmarshal(v.Request)
 		if err != nil {
 			return nil, err
 		}
@@ -191,6 +191,10 @@ func (i *Ingester) PushBytes(ctx context.Context, req *tempopb.PushBytesRequest)
 			return nil, err
 		}
 	}
+
+	// Reuse request instead of handing over to GC
+	tempopb.ReuseRequest(req)
+
 	return &tempopb.PushResponse{}, nil
 }
 

pkg/tempopb/prealloc.go

@@ -0,0 +1,45 @@
+package tempopb
+
+import (
+	"github.com/prometheus/prometheus/pkg/pool"
+)
+
+var (
+	// buckets: [1KiB, 2KiB, 4KiB, 8KiB, 16KiB]
+	bytePool = pool.New(1_000, 16_000, 2, func(size int) interface{} { return make([]byte, 0, size) })
+)
+
+// PreallocRequest is a (repeated bytes requests) which preallocs slices on Unmarshal.
+type PreallocRequest struct {
+	Request []byte
+}
+
+// Unmarshal implements proto.Message.
+func (r *PreallocRequest) Unmarshal(dAtA []byte) error {
+	r.Request = bytePool.Get(len(dAtA)).([]byte)
+	r.Request = r.Request[:len(dAtA)]
+	copy(r.Request, dAtA)
+	return nil
+}
+
+// MarshalTo implements proto.Marshaller.
+func (r *PreallocRequest) MarshalTo(dAtA []byte) (int, error) {
+	copy(dAtA[:], r.Request[:])
+	return 0, nil // returned int is not used
+}
+
+// Size implements proto.Sizer.
+func (r *PreallocRequest) Size() (n int) {
+	if r == nil {
+		return 0
+	}
+	return len(r.Request)
+}
+
+// ReuseRequest puts the byte slice back into bytePool for reuse.
+func ReuseRequest(req *PushBytesRequest) {
+	for i := range req.Requests {
+		// We want to preserve the underlying allocated memory, [:0] helps us retains the cap() of the slice
+		bytePool.Put(req.Requests[i].Request[:0])
+	}
+}

pkg/tempopb/tempo.proto

@@ -3,6 +3,7 @@ syntax="proto3";
 package tempopb;
 
 import "trace/v1/trace.proto";
+import "github.com/gogo/protobuf/gogoproto/gogo.proto";
 
 service Pusher {
   rpc Push(PushRequest) returns (PushResponse) {};
@@ -37,5 +38,5 @@ message PushResponse {
 
 message PushBytesRequest {
   // pre-serialized PushRequests
-  repeated bytes requests = 1;
+  repeated bytes requests = 1 [(gogoproto.nullable) = false, (gogoproto.customtype) = "PreallocRequest"];
 }