*: port our changes to newest main

cmd/thanos/receive.go

@@ -472,7 +472,7 @@ func runReceive(
 		capNProtoWriter := receive.NewCapNProtoWriter(logger, dbs, &receive.CapNProtoWriterOptions{
 			TooFarInFutureTimeWindow: int64(time.Duration(*conf.tsdbTooFarInFutureTimeWindow)),
 		})
-		handler := receive.NewCapNProtoHandler(logger, capNProtoWriter)
+		handler := receive.NewCapNProtoHandler(reg, logger, capNProtoWriter)
 		listener, err := net.Listen("tcp", conf.replicationAddr)
 		if err != nil {
 			return err

pkg/receive/capnp_server.go

@@ -6,12 +6,15 @@ package receive
 import (
 	"context"
 	"net"
+	"unsafe"
 
 	"capnproto.org/go/capnp/v3"
 	"capnproto.org/go/capnp/v3/rpc"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 
 	"github.com/thanos-io/thanos/pkg/receive/writecapnp"
 	"github.com/thanos-io/thanos/pkg/runutil"
@@ -55,32 +58,51 @@ func (c *CapNProtoServer) Shutdown() {
 }
 
 type CapNProtoHandler struct {
-	writer *CapNProtoWriter
-	logger log.Logger
+	writer       *CapNProtoWriter
+	logger       log.Logger
+	handledTotal prometheus.Counter
 }
 
-func NewCapNProtoHandler(logger log.Logger, writer *CapNProtoWriter) *CapNProtoHandler {
-	return &CapNProtoHandler{logger: logger, writer: writer}
+func NewCapNProtoHandler(reg prometheus.Registerer, logger log.Logger, writer *CapNProtoWriter) *CapNProtoHandler {
+	handledTotal := promauto.With(reg).NewCounter(prometheus.CounterOpts{
+		Name: "thanos_receive_capnproto_handled_total",
+		Help: "Total number of handled CapNProto requests.",
+	})
+
+	return &CapNProtoHandler{handledTotal: handledTotal, logger: logger, writer: writer}
 }
 
-func (c CapNProtoHandler) Write(ctx context.Context, call writecapnp.Writer_write) error {
+func (c *CapNProtoHandler) Write(ctx context.Context, call writecapnp.Writer_write) error {
+	c.handledTotal.Inc()
+
 	call.Go()
 	wr, err := call.Args().Wr()
 	if err != nil {
 		return err
 	}
-	t, err := wr.Tenant()
-	if err != nil {
-		return err
-	}
-	req, err := writecapnp.NewRequest(wr)
+
+	data, err := wr.Data()
 	if err != nil {
 		return err
 	}
-	defer req.Close()
 
 	var errs writeErrors
-	errs.Add(c.writer.Write(ctx, t, req))
+	for i := 0; i < data.Len(); i++ {
+		d := data.At(i)
+		req, err := writecapnp.NewRequest(d)
+		if err != nil {
+			return err
+		}
+		defer req.Close()
+
+		tenant, err := d.TenantBytes()
+		if err != nil {
+			return err
+		}
+
+		errs.Add(c.writer.Write(ctx, unsafe.String(&tenant[0], len(tenant)), req))
+	}
+
 	if err := errs.ErrOrNil(); err != nil {
 		level.Debug(c.logger).Log("msg", "failed to handle request", "err", err)
 		result, allocErr := call.AllocResults()

pkg/receive/capnproto_server_bench_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/test/bufconn"
 
@@ -53,7 +54,7 @@ func BenchmarkCapNProtoServer_SingleConcurrentClient(b *testing.B) {
 			&CapNProtoWriterOptions{},
 		)
 		listener = bufconn.Listen(1024)
-		handler  = NewCapNProtoHandler(log.NewNopLogger(), writer)
+		handler  = NewCapNProtoHandler(prometheus.NewRegistry(), log.NewNopLogger(), writer)
 		srv      = NewCapNProtoServer(listener, handler, log.NewNopLogger())
 	)
 	go func() {

pkg/receive/capnproto_server_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/test/bufconn"
 
@@ -25,7 +26,7 @@ func TestCapNProtoServer_SingleConcurrentClient(t *testing.T) {
 			&CapNProtoWriterOptions{},
 		)
 		listener = bufconn.Listen(1024)
-		handler  = NewCapNProtoHandler(log.NewNopLogger(), writer)
+		handler  = NewCapNProtoHandler(prometheus.NewRegistry(), log.NewNopLogger(), writer)
 		srv      = NewCapNProtoServer(listener, handler, log.NewNopLogger())
 	)
 	go func() {
@@ -53,7 +54,7 @@ func TestCapNProtoServer_MultipleConcurrentClients(t *testing.T) {
 			&CapNProtoWriterOptions{},
 		)
 		listener = bufconn.Listen(1024)
-		handler  = NewCapNProtoHandler(log.NewNopLogger(), writer)
+		handler  = NewCapNProtoHandler(prometheus.NewRegistry(), log.NewNopLogger(), writer)
 		srv      = NewCapNProtoServer(listener, handler, log.NewNopLogger())
 	)
 	go func() {