[RW2] Return 400 error code for wrongly-formatted histograms (#17210)

* return 400 error code

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

* fix

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

* add more cases

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

* format code

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

* nit_fixing

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

---------

Signed-off-by: pipiland2612 <nguyen.t.dang.minh@gmail.com>

model/histogram/float_histogram.go

@@ -803,16 +803,16 @@ func (h *FloatHistogram) Validate() error {
 			return fmt.Errorf("custom buckets: %w", err)
 		}
 		if h.ZeroCount != 0 {
-			return errors.New("custom buckets: must have zero count of 0")
+			return ErrHistogramCustomBucketsZeroCount
 		}
 		if h.ZeroThreshold != 0 {
-			return errors.New("custom buckets: must have zero threshold of 0")
+			return ErrHistogramCustomBucketsZeroThresh
 		}
 		if len(h.NegativeSpans) > 0 {
-			return errors.New("custom buckets: must not have negative spans")
+			return ErrHistogramCustomBucketsNegSpans
 		}
 		if len(h.NegativeBuckets) > 0 {
-			return errors.New("custom buckets: must not have negative buckets")
+			return ErrHistogramCustomBucketsNegBuckets
 		}
 	} else {
 		if err := checkHistogramSpans(h.PositiveSpans, len(h.PositiveBuckets)); err != nil {
@@ -826,7 +826,7 @@ func (h *FloatHistogram) Validate() error {
 			return fmt.Errorf("negative side: %w", err)
 		}
 		if h.CustomValues != nil {
-			return errors.New("histogram with exponential schema must not have custom bounds")
+			return ErrHistogramExpSchemaCustomBounds
 		}
 	}
 	err := checkHistogramBuckets(h.PositiveBuckets, &pCount, false)

model/histogram/generic.go

@@ -37,6 +37,11 @@ var (
 	ErrHistogramCustomBucketsInfinite   = errors.New("histogram custom bounds must be finite")
 	ErrHistogramsIncompatibleSchema     = errors.New("cannot apply this operation on histograms with a mix of exponential and custom bucket schemas")
 	ErrHistogramsIncompatibleBounds     = errors.New("cannot apply this operation on custom buckets histograms with different custom bounds")
+	ErrHistogramCustomBucketsZeroCount  = errors.New("custom buckets: must have zero count of 0")
+	ErrHistogramCustomBucketsZeroThresh = errors.New("custom buckets: must have zero threshold of 0")
+	ErrHistogramCustomBucketsNegSpans   = errors.New("custom buckets: must not have negative spans")
+	ErrHistogramCustomBucketsNegBuckets = errors.New("custom buckets: must not have negative buckets")
+	ErrHistogramExpSchemaCustomBounds   = errors.New("histogram with exponential schema must not have custom bounds")
 )
 
 func IsCustomBucketsSchema(s int32) bool {

model/histogram/histogram.go

@@ -14,7 +14,6 @@
 package histogram
 
 import (
-	"errors"
 	"fmt"
 	"math"
 	"slices"
@@ -430,16 +429,16 @@ func (h *Histogram) Validate() error {
 			return fmt.Errorf("custom buckets: %w", err)
 		}
 		if h.ZeroCount != 0 {
-			return errors.New("custom buckets: must have zero count of 0")
+			return ErrHistogramCustomBucketsZeroCount
 		}
 		if h.ZeroThreshold != 0 {
-			return errors.New("custom buckets: must have zero threshold of 0")
+			return ErrHistogramCustomBucketsZeroThresh
 		}
 		if len(h.NegativeSpans) > 0 {
-			return errors.New("custom buckets: must not have negative spans")
+			return ErrHistogramCustomBucketsNegSpans
 		}
 		if len(h.NegativeBuckets) > 0 {
-			return errors.New("custom buckets: must not have negative buckets")
+			return ErrHistogramCustomBucketsNegBuckets
 		}
 	} else {
 		if err := checkHistogramSpans(h.PositiveSpans, len(h.PositiveBuckets)); err != nil {
@@ -453,7 +452,7 @@ func (h *Histogram) Validate() error {
 			return fmt.Errorf("negative side: %w", err)
 		}
 		if h.CustomValues != nil {
-			return errors.New("histogram with exponential schema must not have custom bounds")
+			return ErrHistogramExpSchemaCustomBounds
 		}
 	}
 	err := checkHistogramBuckets(h.PositiveBuckets, &pCount, true)

storage/remote/write_handler.go

@@ -117,6 +117,24 @@ func (*writeHandler) parseProtoMsg(contentType string) (config.RemoteWriteProtoM
 	return config.RemoteWriteProtoMsgV1, nil
 }
 
+// isHistogramValidationError checks if the error is a native histogram validation error.
+func isHistogramValidationError(err error) bool {
+	// TODO: Consider adding single histogram error type instead of individual sentinel errors.
+	return errors.Is(err, histogram.ErrHistogramCountMismatch) ||
+		errors.Is(err, histogram.ErrHistogramCountNotBigEnough) ||
+		errors.Is(err, histogram.ErrHistogramNegativeBucketCount) ||
+		errors.Is(err, histogram.ErrHistogramSpanNegativeOffset) ||
+		errors.Is(err, histogram.ErrHistogramSpansBucketsMismatch) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsMismatch) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsInvalid) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsInfinite) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsZeroCount) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsZeroThresh) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsNegSpans) ||
+		errors.Is(err, histogram.ErrHistogramCustomBucketsNegBuckets) ||
+		errors.Is(err, histogram.ErrHistogramExpSchemaCustomBounds)
+}
+
 func (h *writeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	contentType := r.Header.Get("Content-Type")
 	if contentType == "" {
@@ -190,6 +208,9 @@ func (h *writeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				// Indicated an out-of-order sample is a bad request to prevent retries.
 				http.Error(w, err.Error(), http.StatusBadRequest)
 				return
+			case isHistogramValidationError(err):
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return
 			default:
 				h.logger.Error("Error while remote writing the v1 request", "err", err.Error())
 				http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -474,6 +495,11 @@ func (h *writeHandler) appendV2(app storage.Appender, req *writev2.Request, rs *
 				badRequestErrs = append(badRequestErrs, fmt.Errorf("%w for series %v", err, ls.String()))
 				continue
 			}
+			if isHistogramValidationError(err) {
+				h.logger.Error("Invalid histogram received", "err", err.Error(), "series", ls.String(), "timestamp", hp.Timestamp)
+				badRequestErrs = append(badRequestErrs, fmt.Errorf("%w for series %v", err, ls.String()))
+				continue
+			}
 			return 0, http.StatusInternalServerError, err
 		}
 

storage/remote/write_handler_test.go

@@ -806,6 +806,94 @@ func TestCommitErr_V1Message(t *testing.T) {
 	require.Equal(t, "commit error\n", string(body))
 }
 
+// Regression test for https://github.com/prometheus/prometheus/issues/17206
+func TestHistogramValidationErrorHandling(t *testing.T) {
+	testCases := []struct {
+		desc     string
+		hist     histogram.Histogram
+		expected string
+	}{
+		{
+			desc: "count mismatch",
+			hist: histogram.Histogram{
+				Schema:          2,
+				ZeroThreshold:   1e-128,
+				ZeroCount:       1,
+				Count:           10,
+				Sum:             20,
+				PositiveSpans:   []histogram.Span{{Offset: 0, Length: 1}},
+				PositiveBuckets: []int64{2},
+				NegativeSpans:   []histogram.Span{{Offset: 0, Length: 1}},
+				NegativeBuckets: []int64{3},
+				// Total: 1 (zero) + 2 (positive) + 3 (negative) = 6, but Count = 10
+			},
+			expected: "histogram's observation count should equal",
+		},
+		{
+			desc: "custom buckets zero count",
+			hist: histogram.Histogram{
+				Schema:          histogram.CustomBucketsSchema,
+				Count:           10,
+				Sum:             20,
+				ZeroCount:       1, // Invalid: custom buckets must have zero count of 0
+				PositiveSpans:   []histogram.Span{{Offset: 0, Length: 1}},
+				PositiveBuckets: []int64{10},
+				CustomValues:    []float64{1.0},
+			},
+			expected: "custom buckets: must have zero count of 0",
+		},
+	}
+
+	for _, protoMsg := range []config.RemoteWriteProtoMsg{config.RemoteWriteProtoMsgV1, config.RemoteWriteProtoMsgV2} {
+		protoName := "V1"
+		if protoMsg == config.RemoteWriteProtoMsgV2 {
+			protoName = "V2"
+		}
+
+		for _, tc := range testCases {
+			testName := fmt.Sprintf("%s %s", protoName, tc.desc)
+			t.Run(testName, func(t *testing.T) {
+				dir := t.TempDir()
+				opts := tsdb.DefaultOptions()
+				opts.EnableNativeHistograms = true
+
+				db, err := tsdb.Open(dir, nil, nil, opts, nil)
+				require.NoError(t, err)
+				t.Cleanup(func() { require.NoError(t, db.Close()) })
+
+				handler := NewWriteHandler(promslog.NewNopLogger(), nil, db.Head(), []config.RemoteWriteProtoMsg{protoMsg}, false)
+				recorder := httptest.NewRecorder()
+
+				var buf []byte
+				if protoMsg == config.RemoteWriteProtoMsgV1 {
+					ts := []prompb.TimeSeries{{
+						Labels:     []prompb.Label{{Name: "__name__", Value: "test"}},
+						Histograms: []prompb.Histogram{prompb.FromIntHistogram(1, &tc.hist)},
+					}}
+					buf, _, _, err = buildWriteRequest(nil, ts, nil, nil, nil, nil, "snappy")
+				} else {
+					st := writev2.NewSymbolTable()
+					ts := []writev2.TimeSeries{{
+						LabelsRefs: st.SymbolizeLabels(labels.FromStrings("__name__", "test"), nil),
+						Histograms: []writev2.Histogram{writev2.FromIntHistogram(1, &tc.hist)},
+					}}
+					buf, _, _, err = buildV2WriteRequest(promslog.NewNopLogger(), ts, st.Symbols(), nil, nil, nil, "snappy")
+				}
+				require.NoError(t, err)
+
+				req := httptest.NewRequest(http.MethodPost, "/api/v1/write", bytes.NewReader(buf))
+				req.Header.Set("Content-Type", remoteWriteContentTypeHeaders[protoMsg])
+				req.Header.Set("Content-Encoding", "snappy")
+
+				handler.ServeHTTP(recorder, req)
+
+				require.Equal(t, http.StatusBadRequest, recorder.Code)
+				require.Contains(t, recorder.Body.String(), tc.expected)
+			})
+		}
+	}
+}
+
 func TestCommitErr_V2Message(t *testing.T) {
 	payload, _, _, err := buildV2WriteRequest(promslog.NewNopLogger(), writeV2RequestFixture.Timeseries, writeV2RequestFixture.Symbols, nil, nil, nil, "snappy")
 	require.NoError(t, err)