From 3c526842751042e0fa2c4df7aacd4c868b6a9c10 Mon Sep 17 00:00:00 2001
From: Yuri Nikolic <durica.nikolic@grafana.com>
Date: Wed, 18 Oct 2023 16:38:41 +0200
Subject: [PATCH] Distributor: improve wrapping of gRPC errors returned by the
 ingester

Signed-off-by: Yuri Nikolic <durica.nikolic@grafana.com>
---
 pkg/distributor/distributor.go      |  6 ++++
 pkg/distributor/distributor_test.go | 46 ++++++++++++++++++++++-------
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/pkg/distributor/distributor.go b/pkg/distributor/distributor.go
index 844656ebd78..46919557e0c 100644
--- a/pkg/distributor/distributor.go
+++ b/pkg/distributor/distributor.go
@@ -1277,6 +1277,12 @@ func handleIngesterPushError(err error) error {
 		// Wrap HTTP gRPC error with more explanatory message.
 		return httpgrpc.Errorf(int(resp.Code), "failed pushing to ingester: %s", resp.Body)
 	}
+	stat, ok := status.FromError(err)
+	if ok {
+		st := stat.Proto()
+		st.Message = fmt.Sprintf("failed pushing to ingester: %s", st.Message)
+		return status.ErrorProto(st)
+	}
 	return errors.Wrap(err, "failed pushing to ingester")
 }
 
diff --git a/pkg/distributor/distributor_test.go b/pkg/distributor/distributor_test.go
index b6192ca2714..8c17dd1d3c8 100644
--- a/pkg/distributor/distributor_test.go
+++ b/pkg/distributor/distributor_test.go
@@ -4785,10 +4785,10 @@ func TestHandleIngesterPushError(t *testing.T) {
 	outputErrorMsgPrefix := "failed pushing to ingester"
 	userID := "test"
 	errWithUserID := fmt.Errorf("user=%s: %s", userID, testErrorMsg)
-	unavailableErr := status.New(codes.Unavailable, testErrorMsg).Err()
 	test := map[string]struct {
 		ingesterPushError   error
 		expectedOutputError error
+		checkDetails        bool
 	}{
 		"no error gives no error": {
 			ingesterPushError:   nil,
@@ -4797,18 +4797,21 @@ func TestHandleIngesterPushError(t *testing.T) {
 		"a 4xx HTTP gRPC error gives a 4xx HTTP gRPC error": {
 			ingesterPushError:   httpgrpc.Errorf(http.StatusBadRequest, testErrorMsg),
 			expectedOutputError: httpgrpc.Errorf(http.StatusBadRequest, "%s: %s", outputErrorMsgPrefix, testErrorMsg),
+			checkDetails:        true,
 		},
 		"a 5xx HTTP gRPC error gives a 5xx HTTP gRPC error": {
 			ingesterPushError:   httpgrpc.Errorf(http.StatusServiceUnavailable, testErrorMsg),
 			expectedOutputError: httpgrpc.Errorf(http.StatusServiceUnavailable, "%s: %s", outputErrorMsgPrefix, testErrorMsg),
+			checkDetails:        true,
 		},
 		"a random ingester error without status gives the same wrapped error": {
 			ingesterPushError:   errWithUserID,
 			expectedOutputError: errors.Wrap(errWithUserID, outputErrorMsgPrefix),
 		},
-		"a gRPC unavailable error gives the same wrapped error": {
-			ingesterPushError:   unavailableErr,
-			expectedOutputError: errors.Wrap(unavailableErr, outputErrorMsgPrefix),
+		"a gRPC unavailable error gives a gRPC unavailable error with a wrapped message": {
+			ingesterPushError:   createGRPCErrorWithDetails(t, codes.Unavailable, testErrorMsg, mimirpb.INVALID),
+			expectedOutputError: createGRPCErrorWithDetails(t, codes.Unavailable, fmt.Sprintf("%s: %s", outputErrorMsgPrefix, testErrorMsg), mimirpb.INVALID),
+			checkDetails:        true,
 		},
 		"a context cancel error gives the same wrapped error": {
 			ingesterPushError:   context.Canceled,
@@ -4816,13 +4819,27 @@ func TestHandleIngesterPushError(t *testing.T) {
 		},
 	}
 
-	for _, testData := range test {
-		err := handleIngesterPushError(testData.ingesterPushError)
-		if testData.expectedOutputError == nil {
-			require.NoError(t, err)
-		} else {
-			require.ErrorContains(t, err, testData.expectedOutputError.Error())
-		}
+	for testName, testData := range test {
+		t.Run(testName, func(t *testing.T) {
+			err := handleIngesterPushError(testData.ingesterPushError)
+			if testData.expectedOutputError == nil {
+				require.NoError(t, err)
+			} else {
+				if testData.checkDetails {
+					expectedStat, ok := status.FromError(testData.expectedOutputError)
+					require.True(t, ok)
+
+					stat, ok := status.FromError(err)
+					require.True(t, ok)
+
+					require.Equal(t, expectedStat.Code(), stat.Code())
+					require.Equal(t, expectedStat.Message(), stat.Message())
+					require.Equal(t, expectedStat.Details(), stat.Details())
+				} else {
+					require.Errorf(t, err, testData.expectedOutputError.Error())
+				}
+			}
+		})
 	}
 }
 
@@ -4924,3 +4941,10 @@ func checkGRPCError(t *testing.T, expectedStatus *status.Status, expectedDetails
 		require.Equal(t, expectedDetails, errorDetails)
 	}
 }
+
+func createGRPCErrorWithDetails(t *testing.T, code codes.Code, message string, cause mimirpb.ErrorCause) error {
+	stat := status.New(code, message)
+	statWithDetails, err := stat.WithDetails(&mimirpb.WriteErrorDetails{Cause: cause})
+	require.NoError(t, err)
+	return statWithDetails.Err()
+}