Canceled distributor push requests return 499

Signed-off-by: wangguoliang <iamwgliang@gmail.com>

Author: wgliang

pkg/distributor/distributor_test.go

@@ -290,6 +290,34 @@ func TestDistributor_Push(t *testing.T) {
 	}
 }
 
+func TestDistributor_ContextCanceledRequest(t *testing.T) {
+	t.Cleanup(func() {
+		_ = func() time.Time { return time.Now() }
+	})
+
+	ds, ings, _, _ := prepare(t, prepConfig{
+		numIngesters:    3,
+		happyIngesters:  3,
+		numDistributors: 1,
+	})
+
+	// Lock all mockIngester instances, so they will be waiting
+	for i := range ings {
+		ings[i].Lock()
+		defer func(ing *mockIngester) {
+			ing.Unlock()
+		}(ings[i])
+	}
+
+	ctx := user.InjectOrgID(context.Background(), "user")
+	ctx, cancel := context.WithCancel(ctx)
+	cancel()
+	request := makeWriteRequest(123456789000, 1, 1)
+	_, err := ds[0].Push(ctx, request)
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func TestDistributor_MetricsCleanup(t *testing.T) {
 	dists, _, regs, _ := prepare(t, prepConfig{
 		numDistributors: 1,

pkg/util/push/push.go

@@ -2,6 +2,7 @@ package push
 
 import (
 	"context"
+	"errors"
 	"net/http"
 
 	"github.com/go-kit/log/level"
@@ -13,6 +14,8 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/log"
 )
 
+const statusClientCanceledRequest = 499
+
 // Func defines the type of the push. It is similar to http.HandlerFunc.
 type Func func(context.Context, *cortexpb.WriteRequest) (*cortexpb.WriteResponse, error)
 
@@ -42,6 +45,12 @@ func Handler(maxRecvMsgSize int, sourceIPs *middleware.SourceIPExtractor, push F
 		}
 
 		if _, err := push(ctx, &req.WriteRequest); err != nil {
+			if errors.Is(err, context.Canceled) {
+				http.Error(w, err.Error(), statusClientCanceledRequest)
+				level.Warn(logger).Log("msg", "push request canceled", "err", err)
+				return
+			}
+
 			resp, ok := httpgrpc.HTTPResponseFromError(err)
 			if !ok {
 				http.Error(w, err.Error(), http.StatusInternalServerError)

pkg/util/push/push_test.go

@@ -3,6 +3,7 @@ package push
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -34,6 +35,17 @@ func TestHandler_cortexWriteRequest(t *testing.T) {
 	assert.Equal(t, 200, resp.Code)
 }
 
+func TestHandler_contextCanceledRequest(t *testing.T) {
+	req := createRequest(t, createCortexWriteRequestProtobuf(t, false))
+	resp := httptest.NewRecorder()
+	sourceIPs, _ := middleware.NewSourceIPs("SomeField", "(.*)")
+	handler := Handler(100000, sourceIPs, func(_ context.Context, _ *cortexpb.WriteRequest) (*cortexpb.WriteResponse, error) {
+		return nil, fmt.Errorf("the request failed: %w", context.Canceled)
+	})
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, 499, resp.Code)
+}
+
 func TestHandler_ignoresSkipLabelNameValidationIfSet(t *testing.T) {
 	for _, req := range []*http.Request{
 		createRequest(t, createCortexWriteRequestProtobuf(t, true)),