Distributor: /all_user_stats now show API and Rule Ingest Rate correctly

CHANGELOG.md

@@ -26,6 +26,7 @@
 * [BUGFIX] Cassandra Storage: Fix endpoint TLS host verification. #2109
 * [BUGFIX] Experimental TSDB: fixed response status code from `422` to `500` when an error occurs while iterating chunks with the experimental blocks storage. #2402
 * [BUGFIX] Ring: Fixed a situation where upgrading from pre-1.0 cortex with a rolling strategy caused new 1.0 ingesters to lose their zone value in the ring until manually forced to re-register. #2404
+* [BUGFIX] Distributor: `/all_user_stats` now show API and Rule Ingest Rate correctly. #2457
 
 ## 1.0.0 / 2020-04-02
 

pkg/distributor/distributor.go

@@ -536,7 +536,7 @@ func (d *Distributor) Push(ctx context.Context, req *client.WriteRequest) (*clie
 		if sp := opentracing.SpanFromContext(ctx); sp != nil {
 			localCtx = opentracing.ContextWithSpan(localCtx, sp)
 		}
-		return d.send(localCtx, ingester, timeseries, metadata)
+		return d.send(localCtx, ingester, timeseries, metadata, req.Source)
 	}, func() { client.ReuseSlice(req.Timeseries) })
 	if err != nil {
 		return nil, err
@@ -566,7 +566,7 @@ func sortLabelsIfNeeded(labels []client.LabelAdapter) {
 	})
 }
 
-func (d *Distributor) send(ctx context.Context, ingester ring.IngesterDesc, timeseries []client.PreallocTimeseries, metadata []*client.MetricMetadata) error {
+func (d *Distributor) send(ctx context.Context, ingester ring.IngesterDesc, timeseries []client.PreallocTimeseries, metadata []*client.MetricMetadata, source client.WriteRequest_SourceEnum) error {
 	h, err := d.ingesterPool.GetClientFor(ingester.Addr)
 	if err != nil {
 		return err
@@ -576,6 +576,7 @@ func (d *Distributor) send(ctx context.Context, ingester ring.IngesterDesc, time
 	req := client.WriteRequest{
 		Timeseries: timeseries,
 		Metadata:   metadata,
+		Source:     source,
 	}
 	_, err = c.Push(ctx, &req)
 

pkg/ruler/compat.go

@@ -39,6 +39,8 @@ func (a *appender) AddFast(l labels.Labels, ref uint64, t int64, v float64) erro
 }
 
 func (a *appender) Commit() error {
+	// Since a.pusher is distributor, client.ReuseSlice will be called in a.pusher.Push.
+	// We shouldn't call client.ReuseSlice here.
 	_, err := a.pusher.Push(user.InjectOrgID(context.Background(), a.userID), client.ToWriteRequest(a.labels, a.samples, nil, client.RULE))
 	a.labels = nil
 	a.samples = nil

pkg/util/push/push.go

@@ -17,14 +17,16 @@ func Handler(cfg distributor.Config, push func(context.Context, *client.WriteReq
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		compressionType := util.CompressionTypeFor(r.Header.Get("X-Prometheus-Remote-Write-Version"))
 		var req client.PreallocWriteRequest
-		req.Source = client.API
 		_, err := util.ParseProtoReader(r.Context(), r.Body, int(r.ContentLength), cfg.MaxRecvMsgSize, &req, compressionType)
 		logger := util.WithContext(r.Context(), util.Logger)
 		if err != nil {
 			level.Error(logger).Log("err", err.Error())
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
+		if req.Source == 0 {
+			req.Source = client.API
+		}
 
 		if _, err := push(r.Context(), &req.WriteRequest); err != nil {
 			resp, ok := httpgrpc.HTTPResponseFromError(err)

pkg/util/push/push_test.go

@@ -0,0 +1,95 @@
+package push
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/golang/snappy"
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/cortexproject/cortex/pkg/distributor"
+	"github.com/cortexproject/cortex/pkg/ingester/client"
+)
+
+func TestHandler_remoteWrite(t *testing.T) {
+	req := createRequest(t, createPrometheusRemoteWriteProtobuf(t))
+	resp := httptest.NewRecorder()
+	handler := Handler(distributor.Config{MaxRecvMsgSize: 100000}, verifyWriteRequestHandler(t, client.API))
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, 200, resp.Code)
+}
+
+func TestHandler_cortexWriteRequest(t *testing.T) {
+	req := createRequest(t, createCortexWriteRequestProtobuf(t))
+	resp := httptest.NewRecorder()
+	handler := Handler(distributor.Config{MaxRecvMsgSize: 100000}, verifyWriteRequestHandler(t, client.RULE))
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, 200, resp.Code)
+}
+
+func verifyWriteRequestHandler(t *testing.T, expectSource client.WriteRequest_SourceEnum) func(ctx context.Context, request *client.WriteRequest) (response *client.WriteResponse, err error) {
+	t.Helper()
+	return func(ctx context.Context, request *client.WriteRequest) (response *client.WriteResponse, err error) {
+		assert.Len(t, request.Timeseries, 1)
+		assert.Equal(t, "__name__", request.Timeseries[0].Labels[0].Name)
+		assert.Equal(t, "foo", request.Timeseries[0].Labels[0].Value)
+		assert.Equal(t, expectSource, request.Source)
+		return &client.WriteResponse{}, nil
+	}
+}
+
+func createRequest(t *testing.T, protobuf []byte) *http.Request {
+	t.Helper()
+	inoutBytes := snappy.Encode(nil, protobuf)
+	req, err := http.NewRequest("POST", "http://localhost/", bytes.NewReader(inoutBytes))
+	require.NoError(t, err)
+	req.Header.Add("Content-Encoding", "snappy")
+	req.Header.Set("Content-Type", "application/x-protobuf")
+	req.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")
+	return req
+}
+
+func createPrometheusRemoteWriteProtobuf(t *testing.T) []byte {
+	t.Helper()
+	input := prompb.WriteRequest{
+		Timeseries: []prompb.TimeSeries{
+			{
+				Labels: []prompb.Label{
+					{Name: "__name__", Value: "foo"},
+				},
+				Samples: []prompb.Sample{
+					{Value: 1, Timestamp: time.Date(2020, 4, 1, 0, 0, 0, 0, time.UTC).UnixNano()},
+				},
+			},
+		},
+	}
+	inoutBytes, err := input.Marshal()
+	require.NoError(t, err)
+	return inoutBytes
+}
+func createCortexWriteRequestProtobuf(t *testing.T) []byte {
+	t.Helper()
+	ts := client.PreallocTimeseries{
+		TimeSeries: &client.TimeSeries{
+			Labels: []client.LabelAdapter{
+				{Name: "__name__", Value: "foo"},
+			},
+			Samples: []client.Sample{
+				{Value: 1, TimestampMs: time.Date(2020, 4, 1, 0, 0, 0, 0, time.UTC).UnixNano()},
+			},
+		},
+	}
+	input := client.WriteRequest{
+		Timeseries: []client.PreallocTimeseries{ts},
+		Source:     client.RULE,
+	}
+	inoutBytes, err := input.Marshal()
+	require.NoError(t, err)
+	return inoutBytes
+}