From 35c5e6633fa2236f980a38ce77e21c817421e062 Mon Sep 17 00:00:00 2001
From: Ganesh Vernekar <ganeshvern@gmail.com>
Date: Tue, 3 Dec 2024 17:11:56 -0500
Subject: [PATCH] block-builder: Do not hard fail when native histogram is
 disabled and tried to ingest one (#10111)

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
---
 pkg/blockbuilder/tsdb.go      |  6 ++++
 pkg/blockbuilder/tsdb_test.go | 54 +++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/pkg/blockbuilder/tsdb.go b/pkg/blockbuilder/tsdb.go
index d8c1376230d..ee2d610fe78 100644
--- a/pkg/blockbuilder/tsdb.go
+++ b/pkg/blockbuilder/tsdb.go
@@ -122,6 +122,8 @@ func (b *TSDBBuilder) Process(ctx context.Context, rec *kgo.Record, lastBlockMax
 
 		allSamplesProcessed = true
 		discardedSamples    = 0
+
+		nativeHistogramsIngestionEnabled = b.limits.NativeHistogramsIngestionEnabled(userID)
 	)
 	for _, ts := range req.Timeseries {
 		mimirpb.FromLabelAdaptersOverwriteLabels(&labelsBuilder, ts.Labels, &nonCopiedLabels)
@@ -163,6 +165,10 @@ func (b *TSDBBuilder) Process(ctx context.Context, rec *kgo.Record, lastBlockMax
 			}
 		}
 
+		if !nativeHistogramsIngestionEnabled {
+			continue
+		}
+
 		for _, h := range ts.Histograms {
 			if h.Timestamp >= blockMax {
 				// We will process this sample in the next cycle.
diff --git a/pkg/blockbuilder/tsdb_test.go b/pkg/blockbuilder/tsdb_test.go
index e85e801f088..9b1ad07cfc7 100644
--- a/pkg/blockbuilder/tsdb_test.go
+++ b/pkg/blockbuilder/tsdb_test.go
@@ -454,6 +454,60 @@ func TestTSDBBuilderLimits(t *testing.T) {
 	require.Equal(t, uint64(100), db.Head().NumSeries())
 }
 
+// TestTSDBBuilderNativeHistogramEnabledError tests that when native histograms are disabled for a tenant,
+// the TSDB builder does not error out when trying to ingest native histogram for that tenant.
+func TestTSDBBuilderNativeHistogramEnabledError(t *testing.T) {
+	var (
+		user1 = "user1"
+		user2 = "user2"
+	)
+
+	limits := map[string]*validation.Limits{
+		user1: {
+			NativeHistogramsIngestionEnabled: true,
+		},
+		user2: {
+			NativeHistogramsIngestionEnabled: false,
+		},
+	}
+	overrides, err := validation.NewOverrides(defaultLimitsTestConfig(), validation.NewMockTenantLimits(limits))
+	require.NoError(t, err)
+
+	metrics := newTSDBBBuilderMetrics(prometheus.NewPedanticRegistry())
+	builder := NewTSDBBuilder(log.NewNopLogger(), t.TempDir(), mimir_tsdb.BlocksStorageConfig{}, overrides, metrics, 0)
+	t.Cleanup(func() {
+		require.NoError(t, builder.Close())
+	})
+
+	var (
+		processingRange = time.Hour.Milliseconds()
+		lastEnd         = 2 * processingRange
+		currEnd         = 3 * processingRange
+		ts              = lastEnd + (processingRange / 2)
+	)
+	for seriesID := 1; seriesID <= 100; seriesID++ {
+		for userID := range limits {
+			rec := &kgo.Record{
+				Key:   []byte(userID),
+				Value: createWriteRequest(t, strconv.Itoa(seriesID), nil, histogramSample(ts)),
+			}
+			allProcessed, err := builder.Process(context.Background(), rec, lastEnd, currEnd, false)
+			require.NoError(t, err)
+			require.Equal(t, true, allProcessed)
+		}
+	}
+
+	// user1 had native histograms enabled. We should see it in the TSDB.
+	db, err := builder.getOrCreateTSDB(tsdbTenant{tenantID: user1})
+	require.NoError(t, err)
+	require.Equal(t, uint64(100), db.Head().NumSeries())
+
+	// user2 had native histograms disabled. Nothing should be in the TSDB.
+	db, err = builder.getOrCreateTSDB(tsdbTenant{tenantID: user2})
+	require.NoError(t, err)
+	require.Equal(t, uint64(0), db.Head().NumSeries())
+}
+
 func defaultLimitsTestConfig() validation.Limits {
 	limits := validation.Limits{}
 	flagext.DefaultValues(&limits)