remove sync.Pool as requests are unbounded, update Prometheus deps

.gitignore

@@ -1,5 +1,6 @@
 cortex-tenant
 sink
+generate
 config.yml
 .out
 *.rpm

README.md

@@ -14,17 +14,17 @@ Prometheus remote write proxy which marks timeseries with a Cortex tenant ID bas
 
 Cortex tenants (separate namespaces where metrics are stored to and queried from) are identified by a `X-Scope-OrgID` HTTP header on both writes and queries.
 
-This makes it impossible to use a single Prometheus (or an HA pair) to write to multiple tenants.
+Problem is that Prometheus can't be configured to send this header. And even if it was possible to set it in the remote write configuration - it would be the same for all jobs. This makes it impossible to use a single Prometheus (or an HA pair) to write to multiple tenants.
 
 This proxy solves the problem using the following logic:
 
 - Receive Prometheus remote write
 - Search each timeseries for a specific label name and extract a tenant ID from its value.
   If the label wasn't not found then it can fall back to a configurable default ID.
   If none is configured then the write request will be rejected.
-- Optionally removes this label the timeseries
+- Optionally removes this label from the timeseries
 - Groups timeseries by tenant
-- Issues a number of parallel per-tenant HTTP requests to Cortex adding the tenant HTTP header (`X-Scope-OrgID` by default)
+- Issues a number of parallel per-tenant HTTP requests to Cortex with the relevant tenant HTTP header (`X-Scope-OrgID` by default)
 
 ## Usage
 

VERSION

@@ -1 +1 @@
-1.3.2
+1.3.3

go.mod

@@ -15,15 +15,14 @@ require (
 	github.com/klauspost/compress v1.11.7 // indirect
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/common v0.15.0
-	github.com/prometheus/prometheus v2.5.0+incompatible
+	github.com/prometheus/prometheus v1.8.2-0.20210120113717-82330b96ee74
 	github.com/sirupsen/logrus v1.7.0
-	github.com/stretchr/testify v1.5.1
+	github.com/stretchr/testify v1.6.1
 	github.com/valyala/fasthttp v1.19.0
-	golang.org/x/net v0.0.0-20201224014010-6772e930b67b // indirect
-	golang.org/x/sys v0.0.0-20210113181707-4bcb84eeeb78 // indirect
+	golang.org/x/net v0.0.0-20210119194325-5f4716e94777 // indirect
+	golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4 // indirect
 	golang.org/x/text v0.3.5 // indirect
-	google.golang.org/genproto v0.0.0-20210114201628-6edceaf6022f // indirect
+	google.golang.org/genproto v0.0.0-20210119180700-e258113e47cc // indirect
 	google.golang.org/grpc v1.35.0 // indirect
-	google.golang.org/protobuf v1.25.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0
 )

main.go

@@ -4,7 +4,6 @@ import (
 	"flag"
 	"os"
 	"os/signal"
-	"sync"
 	"syscall"
 
 	"net/http"
@@ -13,28 +12,8 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-const (
-	bufSize = 1024 * 128
-)
-
-type buffer struct {
-	b []byte
-}
-
-func (b *buffer) reset() {
-	b.b = b.b[:bufSize]
-}
-
 var (
 	version = "0.0.0"
-
-	bufferPool = sync.Pool{
-		New: func() interface{} {
-			return &buffer{
-				b: make([]byte, 0, bufSize),
-			}
-		},
-	}
 )
 
 func main() {

processor.go

@@ -112,6 +112,11 @@ func (p *processor) handle(ctx *fh.RequestCtx) {
 	}
 
 	if len(wrReqIn.Timeseries) == 0 {
+		// If there's metadata - just accept the request and drop it
+		if len(wrReqIn.Metadata) > 0 {
+			return
+		}
+
 		ctx.Error("No timeseries found in the request", fh.StatusBadRequest)
 		return
 	}
@@ -151,12 +156,12 @@ func (p *processor) handle(ctx *fh.RequestCtx) {
 	return
 }
 
-func (p *processor) createWriteRequests(in *prompb.WriteRequest) (map[string]*prompb.WriteRequest, error) {
+func (p *processor) createWriteRequests(wrReqIn *prompb.WriteRequest) (map[string]*prompb.WriteRequest, error) {
 	// Create per-tenant write requests
 	m := map[string]*prompb.WriteRequest{}
 
-	for _, ts := range in.Timeseries {
-		tenant, err := p.processTimeseries(ts)
+	for _, ts := range wrReqIn.Timeseries {
+		tenant, err := p.processTimeseries(&ts)
 		if err != nil {
 			return nil, err
 		}
@@ -174,11 +179,7 @@ func (p *processor) createWriteRequests(in *prompb.WriteRequest) (map[string]*pr
 }
 
 func (p *processor) unmarshal(b []byte) (*prompb.WriteRequest, error) {
-	buf := bufferPool.Get().(*buffer)
-	buf.reset()
-	defer bufferPool.Put(buf)
-
-	decoded, err := snappy.Decode(buf.b, b)
+	decoded, err := snappy.Decode(nil, b)
 	if err != nil {
 		return nil, errors.Wrap(err, "Unable to unpack Snappy")
 	}
@@ -191,19 +192,16 @@ func (p *processor) unmarshal(b []byte) (*prompb.WriteRequest, error) {
 	return req, nil
 }
 
-func (p *processor) marshal(wr *prompb.WriteRequest, bufDst []byte) (bufOut []byte, err error) {
-	buf := bufferPool.Get().(*buffer)
-	buf.reset()
-	defer bufferPool.Put(buf)
+func (p *processor) marshal(wr *prompb.WriteRequest) (bufOut []byte, err error) {
+	b := make([]byte, wr.Size())
 
 	// Marshal to Protobuf
-	l, err := wr.MarshalTo(buf.b)
-	if err != nil {
+	if _, err = wr.MarshalTo(b); err != nil {
 		return
 	}
 
 	// Compress with Snappy
-	return snappy.Encode(bufDst, buf.b[:l]), nil
+	return snappy.Encode(nil, b), nil
 }
 
 func (p *processor) dispatch(clientIP net.Addr, reqID uuid.UUID, m map[string]*prompb.WriteRequest) (res []result) {
@@ -256,19 +254,16 @@ func (p *processor) processTimeseries(ts *prompb.TimeSeries) (tenant string, err
 }
 
 func (p *processor) send(clientIP net.Addr, reqID uuid.UUID, tenant string, wr *prompb.WriteRequest) (code int, body []byte, err error) {
-	buf := bufferPool.Get().(*buffer)
-	buf.reset()
-
 	req := fh.AcquireRequest()
 	resp := fh.AcquireResponse()
 
 	defer func() {
 		fh.ReleaseRequest(req)
 		fh.ReleaseResponse(resp)
-		bufferPool.Put(buf)
 	}()
 
-	if buf.b, err = p.marshal(wr, buf.b); err != nil {
+	buf, err := p.marshal(wr)
+	if err != nil {
 		return
 	}
 
@@ -281,7 +276,7 @@ func (p *processor) send(clientIP net.Addr, reqID uuid.UUID, tenant string, wr *
 	req.Header.Set(p.cfg.Tenant.Header, tenant)
 
 	req.SetRequestURI(p.cfg.Target)
-	req.SetBody(buf.b)
+	req.SetBody(buf)
 
 	if err = p.cli.DoTimeout(req, resp, p.cfg.Timeout); err != nil {
 		return

processor_test.go

@@ -42,9 +42,9 @@ var (
 		Timestamp: 101112,
 	}
 
-	testTS1 = &prompb.TimeSeries{
-		Labels: []*prompb.Label{
-			&prompb.Label{
+	testTS1 = prompb.TimeSeries{
+		Labels: []prompb.Label{
+			prompb.Label{
 				Name:  "__tenant__",
 				Value: "foobar",
 			},
@@ -55,9 +55,9 @@ var (
 		},
 	}
 
-	testTS2 = &prompb.TimeSeries{
-		Labels: []*prompb.Label{
-			&prompb.Label{
+	testTS2 = prompb.TimeSeries{
+		Labels: []prompb.Label{
+			prompb.Label{
 				Name:  "__tenant__",
 				Value: "foobaz",
 			},
@@ -68,18 +68,18 @@ var (
 		},
 	}
 
-	testTS3 = &prompb.TimeSeries{
-		Labels: []*prompb.Label{
-			&prompb.Label{
+	testTS3 = prompb.TimeSeries{
+		Labels: []prompb.Label{
+			prompb.Label{
 				Name:  "__tenantXXX",
 				Value: "foobaz",
 			},
 		},
 	}
 
-	testTS4 = &prompb.TimeSeries{
-		Labels: []*prompb.Label{
-			&prompb.Label{
+	testTS4 = prompb.TimeSeries{
+		Labels: []prompb.Label{
+			prompb.Label{
 				Name:  "__tenant__",
 				Value: "foobaz",
 			},
@@ -91,20 +91,20 @@ var (
 	}
 
 	testWRQ = &prompb.WriteRequest{
-		Timeseries: []*prompb.TimeSeries{
+		Timeseries: []prompb.TimeSeries{
 			testTS1,
 			testTS2,
 		},
 	}
 
 	testWRQ1 = &prompb.WriteRequest{
-		Timeseries: []*prompb.TimeSeries{
+		Timeseries: []prompb.TimeSeries{
 			testTS1,
 		},
 	}
 
 	testWRQ2 = &prompb.WriteRequest{
-		Timeseries: []*prompb.TimeSeries{
+		Timeseries: []prompb.TimeSeries{
 			testTS2,
 		},
 	}
@@ -151,7 +151,7 @@ func Test_handle(t *testing.T) {
 	err = p.run()
 	assert.Nil(t, err)
 
-	buf, err := p.marshal(testWRQ, nil)
+	buf, err := p.marshal(testWRQ)
 	assert.Nil(t, err)
 
 	s := &fh.Server{
@@ -243,19 +243,19 @@ func Test_processTimeseries(t *testing.T) {
 	p := newProcessor(*cfg)
 	assert.Nil(t, err)
 
-	ten, err := p.processTimeseries(testTS4)
+	ten, err := p.processTimeseries(&testTS4)
 	assert.Nil(t, err)
 	assert.Equal(t, "foobaz", ten)
 
-	ten, err = p.processTimeseries(testTS3)
+	ten, err = p.processTimeseries(&testTS3)
 	assert.Nil(t, err)
 	assert.Equal(t, "default", ten)
 
 	cfg.Tenant.Default = ""
 	p = newProcessor(*cfg)
 	assert.Nil(t, err)
 
-	ten, err = p.processTimeseries(testTS3)
+	ten, err = p.processTimeseries(&testTS3)
 	assert.NotNil(t, err)
 }
 
@@ -270,7 +270,7 @@ func Test_marshal(t *testing.T) {
 	assert.NotNil(t, err)
 
 	buf := make([]byte, 1024)
-	buf, err = p.marshal(testWRQ, buf)
+	buf, err = p.marshal(testWRQ)
 	assert.Nil(t, err)
 
 	wrq, err := p.unmarshal(buf)
@@ -299,11 +299,10 @@ func Test_createWriteRequests(t *testing.T) {
 
 func Benchmark_marshal(b *testing.B) {
 	p, _ := createProcessor()
-	buf := make([]byte, 1024)
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		buf, _ = p.marshal(testWRQ, buf)
+		buf, _ := p.marshal(testWRQ)
 		_, _ = p.unmarshal(buf)
 	}
 }