[RW2] Fix type and unit labels propagation in Remote Write v2 receiver to prioritize type-and-unit-labels feature (#17387)

* fix

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

* fix nits & update docs

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

* fix docs

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

---------

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

docs/feature_flags.md

@@ -266,7 +266,7 @@ These may not work well if the `<range>` is not a multiple of the collection int
 When enabled, Prometheus will start injecting additional, reserved `__type__`
 and `__unit__` labels as designed in the [PROM-39 proposal](https://github.com/prometheus/proposals/pull/39).
 
-Those labels are sourced from the metadata structured of the existing scrape and ingestion formats
+Those labels are sourced from the metadata structures of the existing scrape and ingestion formats
 like OpenMetrics Text, Prometheus Text, Prometheus Proto, Remote Write 2 and OTLP. All the user provided labels with
 `__type__` and `__unit__` will be overridden.
 
@@ -274,7 +274,7 @@ PromQL layer will handle those labels the same way __name__ is handled, e.g. dro
 on certain operations like `-` or `+` and affected by `promql-delayed-name-removal` feature.
 
 This feature enables important metadata information to be accessible directly with samples and PromQL layer.
- 
+
 It's especially useful for users who:
 
 * Want to be able to select metrics based on type or unit.
@@ -284,6 +284,12 @@ It's especially useful for users who:
 In future more [work is planned](https://github.com/prometheus/prometheus/issues/16610) that will depend on this e.g. rich PromQL UX that helps
 when wrong types are used on wrong functions, automatic renames, delta types and more.
 
+### Behavior with metadata records
+
+When this feature is enabled and the metadata WAL records exists, in an unlikely situation when type or unit are different across those, 
+the Prometheus outputs intends to prefer the `__type__` and `__unit__` labels values. For example on Remote Write 2.0, 
+if  the metadata record somehow (e.g. due to bug) says "counter", but `__type__="gauge"` the remote time series will be set to a gauge.
+
 ## Use Uncached IO
 
 `--enable-feature=use-uncached-io`

storage/remote/queue_manager.go

@@ -1954,20 +1954,27 @@ func populateV2TimeSeries(symbolTable *writev2.SymbolsTable, batch []timeSeries,
 	var nPendingSamples, nPendingExemplars, nPendingHistograms, nPendingMetadata, nUnexpectedMetadata int
 	for nPending, d := range batch {
 		pendingData[nPending].Samples = pendingData[nPending].Samples[:0]
-		if d.metadata != nil {
+		switch {
+		case enableTypeAndUnitLabels:
+			m := schema.NewMetadataFromLabels(d.seriesLabels)
+			pendingData[nPending].Metadata.Type = writev2.FromMetadataType(m.Type)
+			pendingData[nPending].Metadata.UnitRef = symbolTable.Symbolize(m.Unit)
+			pendingData[nPending].Metadata.HelpRef = 0 // Type and unit does not give us help.
+			// Use Help from d.metadata if available.
+			if d.metadata != nil {
+				pendingData[nPending].Metadata.HelpRef = symbolTable.Symbolize(d.metadata.Help)
+				nPendingMetadata++
+			}
+		case d.metadata != nil:
 			pendingData[nPending].Metadata.Type = writev2.FromMetadataType(d.metadata.Type)
 			pendingData[nPending].Metadata.HelpRef = symbolTable.Symbolize(d.metadata.Help)
 			pendingData[nPending].Metadata.UnitRef = symbolTable.Symbolize(d.metadata.Unit)
 			nPendingMetadata++
-		} else {
+		default:
-			var m schema.Metadata
-			if enableTypeAndUnitLabels {
-				m = schema.NewMetadataFromLabels(d.seriesLabels)
-			}
 			// Safeguard against sending garbage in case of not having metadata
 			// for whatever reason.
-			pendingData[nPending].Metadata.Type = writev2.FromMetadataType(m.Type)
+			pendingData[nPending].Metadata.Type = writev2.FromMetadataType(model.MetricTypeUnknown)
-			pendingData[nPending].Metadata.UnitRef = symbolTable.Symbolize(m.Unit)
+			pendingData[nPending].Metadata.UnitRef = 0
 			pendingData[nPending].Metadata.HelpRef = 0
 		}
 

storage/remote/queue_manager_test.go

@@ -39,6 +39,7 @@ import (
 	"github.com/prometheus/prometheus/config"
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/metadata"
 	"github.com/prometheus/prometheus/model/relabel"
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/prompb"
@@ -2492,6 +2493,136 @@ func TestPopulateV2TimeSeries_TypeAndUnitLabels(t *testing.T) {
 	}
 }
 
+// TestPopulateV2TimeSeries_MetadataAndTypeAndUnit verifies that type and unit labels are properly
+// extracted from labels even when d.metadata is not nil (agent mode scenario).
+// Regression test for https://github.com/prometheus/prometheus/issues/17381.
+func TestPopulateV2TimeSeries_MetadataAndTypeAndUnit(t *testing.T) {
+	symbolTable := writev2.NewSymbolTable()
+
+	testCases := []struct {
+		name              string
+		typeLabel         string
+		unitLabel         string
+		metadata          *metadata.Metadata
+		expectedType      writev2.Metadata_MetricType
+		expectedUnit      string
+		enableTypeAndUnit bool
+	}{
+		{
+			name:              "type_and_unit_no_meta",
+			typeLabel:         "gauge",
+			unitLabel:         "bytes",
+			metadata:          nil,
+			expectedType:      writev2.Metadata_METRIC_TYPE_GAUGE,
+			expectedUnit:      "bytes",
+			enableTypeAndUnit: true,
+		},
+		{
+			name:              "type_no_unit_no_meta",
+			typeLabel:         "counter",
+			unitLabel:         "",
+			metadata:          nil,
+			expectedType:      writev2.Metadata_METRIC_TYPE_COUNTER,
+			expectedUnit:      "",
+			enableTypeAndUnit: true,
+		},
+		{
+			name:              "no_type_and_unit_no_meta",
+			typeLabel:         "gauge",
+			unitLabel:         "bytes",
+			metadata:          nil,
+			expectedType:      writev2.Metadata_METRIC_TYPE_UNSPECIFIED,
+			expectedUnit:      "",
+			enableTypeAndUnit: false,
+		},
+		{
+			name:      "type_and_unit_and_meta",
+			typeLabel: "gauge",
+			unitLabel: "bytes",
+			metadata: &metadata.Metadata{
+				Type: model.MetricTypeGauge,
+				Unit: "bytes",
+				Help: "Test metric",
+			},
+			expectedType:      writev2.Metadata_METRIC_TYPE_GAUGE,
+			expectedUnit:      "bytes",
+			enableTypeAndUnit: true,
+		},
+		{
+			name:      "type-and-unit-overrides-meta",
+			typeLabel: "counter",
+			unitLabel: "requests",
+			metadata: &metadata.Metadata{
+				Type: model.MetricTypeGauge,
+				Unit: "bytes",
+				Help: "Test metric",
+			},
+			expectedType:      writev2.Metadata_METRIC_TYPE_COUNTER,
+			expectedUnit:      "requests",
+			enableTypeAndUnit: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			batch := make([]timeSeries, 1)
+			builder := labels.NewScratchBuilder(3)
+
+			// Simulate labels with __type__ and __unit__ as scraped with type-and-unit-labels feature.
+			builder.Add(labels.MetricName, "test_metric")
+			if tc.typeLabel != "" {
+				builder.Add("__type__", tc.typeLabel)
+			}
+			if tc.unitLabel != "" {
+				builder.Add("__unit__", tc.unitLabel)
+			}
+			builder.Sort()
+
+			batch[0] = timeSeries{
+				seriesLabels: builder.Labels(),
+				value:        123.45,
+				timestamp:    time.Now().UnixMilli(),
+				sType:        tSample,
+				metadata:     tc.metadata,
+			}
+
+			pendingData := make([]writev2.TimeSeries, 1)
+
+			symbolTable.Reset()
+			nSamples, nExemplars, nHistograms, nMetadata, _ := populateV2TimeSeries(
+				&symbolTable,
+				batch,
+				pendingData,
+				false,
+				false,
+				tc.enableTypeAndUnit,
+			)
+
+			require.Equal(t, 1, nSamples, "Should have 1 sample")
+			require.Equal(t, 0, nExemplars, "Should have 0 exemplars")
+			require.Equal(t, 0, nHistograms, "Should have 0 histograms")
+
+			// Verify type is correctly extracted.
+			require.Equal(t, tc.expectedType, pendingData[0].Metadata.Type,
+				"Type should match expected for %s", tc.name)
+
+			// Verify unit is correctly extracted.
+			unitRef := pendingData[0].Metadata.UnitRef
+			symbols := symbolTable.Symbols()
+			var actualUnit string
+			if unitRef > 0 && unitRef < uint32(len(symbols)) {
+				actualUnit = symbols[unitRef]
+			}
+			require.Equal(t, tc.expectedUnit, actualUnit, "Unit should match for %s", tc.name)
+
+			// Verify metadata count.
+			if tc.metadata != nil && tc.enableTypeAndUnit {
+				require.Equal(t, 1, nMetadata, "Should count metadata when d.metadata is provided")
+			}
+		})
+	}
+}
+
 func TestHighestTimestampOnAppend(t *testing.T) {
 	for _, protoMsg := range []remoteapi.WriteMessageType{remoteapi.WriteV1MessageType, remoteapi.WriteV2MessageType} {
 		t.Run(fmt.Sprint(protoMsg), func(t *testing.T) {