perf: optimize Prometheus remote write v2 decoding (#8873)

* test: add bench for prom decode

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* chore: merge v2 decode manually and add to bench

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* refactor: update v2 path

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* refactor: update v2 path

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* fix: use constant

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

---------

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

src/servers/Cargo.toml

@@ -180,6 +180,11 @@ common-version.workspace = true
 name = "bench_prom"
 harness = false
 
+[[bench]]
+name = "prom_decode"
+harness = false
+required-features = ["testing"]
+
 [[bench]]
 name = "to_http_output"
 harness = false

src/servers/benches/prom_decode.rs

@@ -12,21 +12,137 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::time::Duration;
 
-use api::prom_store::remote::WriteRequest;
-use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
+use api::greptime_proto::io::prometheus::write::v2 as write_v2;
+use api::prom_store::remote::{self as write_v1, WriteRequest};
+use criterion::{
+    BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main,
+};
 use prost::Message;
 use servers::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest};
+use servers::prom_remote_write::v2::test_util as remote_write_v2;
 use servers::prom_remote_write::validation::{PromValidationMode, validate_label_name};
 use servers::prom_store::to_grpc_row_insert_requests;
 
-fn bench_decode_prom_request(c: &mut Criterion) {
+fn load_fixture_v1_bytes() -> Vec<u8> {
     let mut d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
     d.push("benches");
     d.push("write_request.pb.data");
+    std::fs::read(d).expect("read write_request.pb.data fixture")
+}
 
-    let data = std::fs::read(d).unwrap();
+/// Convert a PRW v1 `WriteRequest` into an equivalent PRW v2 `Request`.
+///
+/// Labels are interned into the v2 symbol table so payloads with repeated label
+/// names/values exercise the main on-wire advantage of v2.
+fn write_request_v1_to_v2(v1: &WriteRequest) -> write_v2::Request {
+    let mut symbols = vec![String::new()];
+    let mut symbol_index: HashMap<String, u32> = HashMap::new();
+    symbol_index.insert(String::new(), 0);
+
+    let mut intern = |value: &str| -> u32 {
+        if let Some(&idx) = symbol_index.get(value) {
+            return idx;
+        }
+        let idx = symbols.len() as u32;
+        symbols.push(value.to_string());
+        symbol_index.insert(value.to_string(), idx);
+        idx
+    };
+
+    let timeseries = v1
+        .timeseries
+        .iter()
+        .map(|series| {
+            let mut labels: Vec<&write_v1::Label> = series.labels.iter().collect();
+            // PRW v2 requires lexicographically sorted label pairs.
+            labels.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.value.cmp(&b.value)));
+
+            let mut labels_refs = Vec::with_capacity(labels.len() * 2);
+            for label in labels {
+                labels_refs.push(intern(&label.name));
+                labels_refs.push(intern(&label.value));
+            }
+
+            write_v2::TimeSeries {
+                labels_refs,
+                samples: series
+                    .samples
+                    .iter()
+                    .map(|sample| write_v2::Sample {
+                        value: sample.value,
+                        timestamp: sample.timestamp,
+                        start_timestamp: 0,
+                    })
+                    .collect(),
+                histograms: Vec::new(),
+                exemplars: Vec::new(),
+                metadata: None,
+            }
+        })
+        .collect();
+
+    write_v2::Request {
+        symbols,
+        timeseries,
+    }
+}
+
+/// Synthetic sample workload with controllable label reuse.
+///
+/// `shared_labels` keeps common label names/values across series (v2-friendly).
+/// Each series also gets a unique `series_id` label so cardinality stays high.
+fn generate_sample_workload(
+    series_count: usize,
+    samples_per_series: usize,
+    shared_label_count: usize,
+) -> (WriteRequest, write_v2::Request) {
+    let shared_labels: Vec<write_v1::Label> = (0..shared_label_count)
+        .map(|i| write_v1::Label {
+            name: format!("label_{i}"),
+            value: format!("value_{}", i % 16),
+        })
+        .collect();
+
+    let timeseries = (0..series_count)
+        .map(|series_idx| {
+            let mut labels = Vec::with_capacity(shared_label_count + 2);
+            labels.push(write_v1::Label {
+                name: "__name__".to_string(),
+                value: format!("metric_{}", series_idx % 32),
+            });
+            labels.extend(shared_labels.iter().cloned());
+            labels.push(write_v1::Label {
+                name: "series_id".to_string(),
+                value: series_idx.to_string(),
+            });
+
+            write_v1::TimeSeries {
+                labels,
+                samples: (0..samples_per_series)
+                    .map(|sample_idx| write_v1::Sample {
+                        value: (series_idx * 1000 + sample_idx) as f64,
+                        timestamp: 1_700_000_000_000 + sample_idx as i64 * 1_000,
+                    })
+                    .collect(),
+                exemplars: Vec::new(),
+                histograms: Vec::new(),
+            }
+        })
+        .collect();
+
+    let v1 = WriteRequest {
+        timeseries,
+        metadata: Vec::new(),
+    };
+    let v2 = write_request_v1_to_v2(&v1);
+    (v1, v2)
+}
+
+fn bench_decode_prom_request(c: &mut Criterion) {
+    let data = load_fixture_v1_bytes();
 
     let mut group = c.benchmark_group("decode_prom_request");
     group.measurement_time(Duration::from_secs(3));
@@ -64,6 +180,128 @@ fn bench_decode_prom_request(c: &mut Criterion) {
     group.finish();
 }
 
+/// Compare Prometheus remote-write v1 vs v2 decoding efficiency.
+///
+/// Two layers are measured on equivalent logical payloads:
+/// 1. protobuf decode only (`WriteRequest` / v2 `Request`)
+/// 2. full Greptime path to row-insert requests
+///
+/// Workloads:
+/// - `fixture`: existing `write_request.pb.data` converted v1 -> v2
+/// - `synthetic_*`: generated series with shared labels (symbol-table friendly)
+#[allow(clippy::print_stderr)]
+fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) {
+    let fixture_v1_bytes = load_fixture_v1_bytes();
+    let fixture_v1 = WriteRequest::decode(fixture_v1_bytes.as_slice()).unwrap();
+    let fixture_v2 = write_request_v1_to_v2(&fixture_v1);
+    let fixture_v2_bytes = fixture_v2.encode_to_vec();
+
+    let synthetic = [
+        (
+            "synthetic_1k_series_1_sample",
+            generate_sample_workload(1_000, 1, 8),
+        ),
+        (
+            "synthetic_1k_series_10_samples",
+            generate_sample_workload(1_000, 10, 8),
+        ),
+        (
+            "synthetic_5k_series_1_sample",
+            generate_sample_workload(5_000, 1, 8),
+        ),
+    ];
+
+    let mut workloads: Vec<(&str, Vec<u8>, Vec<u8>)> =
+        vec![("fixture", fixture_v1_bytes, fixture_v2_bytes)];
+    for (name, (v1, v2)) in &synthetic {
+        workloads.push((*name, v1.encode_to_vec(), v2.encode_to_vec()));
+    }
+
+    eprintln!("\nprom remote-write v1 vs v2 payload sizes:");
+    for (name, v1_bytes, v2_bytes) in &workloads {
+        let ratio = v2_bytes.len() as f64 / v1_bytes.len() as f64;
+        eprintln!(
+            "  {name}: v1={} B, v2={} B, v2/v1={ratio:.3}",
+            v1_bytes.len(),
+            v2_bytes.len()
+        );
+    }
+
+    // --- protobuf decode only ---
+    {
+        let mut group = c.benchmark_group("prom_rw_protobuf_decode");
+        group.measurement_time(Duration::from_secs(3));
+
+        for (name, v1_bytes, v2_bytes) in &workloads {
+            group.throughput(Throughput::Bytes(v1_bytes.len() as u64));
+            group.bench_with_input(BenchmarkId::new("v1", name), v1_bytes, |b, bytes| {
+                b.iter(|| {
+                    black_box(WriteRequest::decode(black_box(bytes.as_slice())).unwrap());
+                });
+            });
+
+            group.throughput(Throughput::Bytes(v2_bytes.len() as u64));
+            group.bench_with_input(BenchmarkId::new("v2", name), v2_bytes, |b, bytes| {
+                b.iter(|| {
+                    black_box(write_v2::Request::decode(black_box(bytes.as_slice())).unwrap());
+                });
+            });
+        }
+
+        group.finish();
+    }
+
+    // --- full path: protobuf -> Greptime row inserts ---
+    {
+        let mut group = c.benchmark_group("prom_rw_decode_to_rows");
+        group.sample_size(50);
+        group.measurement_time(Duration::from_secs(5));
+
+        for (name, v1_bytes, v2_bytes) in &workloads {
+            group.throughput(Throughput::Bytes(v1_bytes.len() as u64));
+            group.bench_with_input(BenchmarkId::new("v1_custom", name), v1_bytes, |b, bytes| {
+                let mut prom_request = PromWriteRequest::default();
+                let mut processor = PromSeriesProcessor::default_processor();
+                b.iter_batched(
+                    || bytes.clone(),
+                    |bytes| {
+                        prom_request
+                            .decode(black_box(bytes), PromValidationMode::Strict, &mut processor)
+                            .unwrap();
+                        let rows = prom_request.as_row_insert_requests();
+                        black_box(&rows);
+                    },
+                    BatchSize::LargeInput,
+                );
+            });
+
+            // Baseline: stock prost WriteRequest + existing converter.
+            group.throughput(Throughput::Bytes(v1_bytes.len() as u64));
+            group.bench_with_input(BenchmarkId::new("v1_prost", name), v1_bytes, |b, bytes| {
+                b.iter(|| {
+                    let request = WriteRequest::decode(black_box(bytes.as_slice())).unwrap();
+                    black_box(to_grpc_row_insert_requests(&request).unwrap());
+                });
+            });
+
+            group.throughput(Throughput::Bytes(v2_bytes.len() as u64));
+            group.bench_with_input(BenchmarkId::new("v2_manual", name), v2_bytes, |b, bytes| {
+                b.iter(|| {
+                    black_box(
+                        remote_write_v2::decode_uncompressed_write_requests(
+                            black_box(bytes.as_slice()),
+                            true,
+                        )
+                        .unwrap(),
+                    );
+                });
+            });
+        }
+
+        group.finish();
+    }
+}
+
 /// Benchmark comparing UTF-8 string validation (`decode_string`) vs
 /// direct byte-level Prometheus label name validation (`decode_label_name`).
 fn bench_label_name_validation(c: &mut Criterion) {
@@ -155,6 +393,7 @@ fn bench_utf8_validation(c: &mut Criterion) {
 criterion_group!(
     benches,
     bench_decode_prom_request,
+    bench_prom_v1_vs_v2_decode,
     bench_label_name_validation,
     bench_utf8_validation
 );

src/servers/src/http/prom_store.rs

@@ -48,7 +48,7 @@ use crate::http::header::{
 use crate::pending_rows_batcher::PendingRowsBatcher;
 use crate::prom_remote_write::decode::PromSeriesProcessor;
 use crate::prom_remote_write::decode_remote_write_request;
-use crate::prom_remote_write::v2::{decode_remote_write_v2_request, into_write_requests};
+use crate::prom_remote_write::v2::decode_remote_write_v2;
 use crate::prom_remote_write::validation::PromValidationMode;
 use crate::prom_store::snappy_decompress;
 use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse};
@@ -235,23 +235,11 @@ async fn remote_write_v2(
     let (db, mut query_ctx, _timer) =
         prepare_remote_write_context(&params, query_ctx, REMOTE_WRITE_V2_VERSION);
 
-    let request = match decode_remote_write_v2_request(is_zstd, body) {
-        Ok(request) => request,
-        Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)),
-    };
-    if !experimental_enable_prometheus_native_histogram && request_has_native_histograms(&request) {
-        return Ok(remote_write_v2_error_response(
-            error::InvalidPromRemoteRequestSnafu {
-                msg: "prometheus remote write v2 native histogram ingestion is experimental; set prom_store.experimental_enable_prometheus_native_histogram = true to enable it"
-                    .to_string(),
-            }
-            .build(),
-            0,
-            0,
-            0,
-        ));
-    }
-    let req = match into_write_requests(request) {
+    let req = match decode_remote_write_v2(
+        is_zstd,
+        body,
+        experimental_enable_prometheus_native_histogram,
+    ) {
         Ok(req) => req,
         Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)),
     };
@@ -302,15 +290,6 @@ async fn remote_write_v2(
     Ok((StatusCode::NO_CONTENT, headers).into_response())
 }
 
-fn request_has_native_histograms(
-    request: &api::greptime_proto::io::prometheus::write::v2::Request,
-) -> bool {
-    request
-        .timeseries
-        .iter()
-        .any(|series| !series.histograms.is_empty())
-}
-
 fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option<axum::response::Response> {
     params
         .get_vm_proto_version

src/servers/src/prom_remote_write/README.md

@@ -11,28 +11,33 @@ Remote write v2 enters through `remote_write_v2` in
 ```mermaid
 flowchart TD
     A["HTTP /v1/prometheus/write"] --> B["remote_write_v2"]
-    B --> C["decode_remote_write_v2_request"]
-    C --> D["into_write_requests"]
+    B --> C["decode_remote_write_v2"]
+    C --> D["borrow symbols and time-series bytes"]
+    D --> E["decode leaf messages and build rows"]
 
-    D --> E["samples ContextReq"]
-    D --> F["histograms ContextReq"]
+    E --> F["samples ContextReq"]
+    E --> G["histograms ContextReq"]
 
-    E --> G["write_prometheus_rows_with_progress"]
-    G --> H["metric engine / pending batcher when enabled"]
+    F --> H["write_prometheus_rows_with_progress"]
+    H --> I["metric engine / pending batcher when enabled"]
 
-    F --> I["histogram ContextOpt"]
-    I --> J["write_prometheus_rows_with_progress"]
-    J --> K["same metric-engine flag as samples, no batcher"]
-    K --> L["table: <metric>"]
+    G --> J["histogram ContextOpt"]
+    J --> K["write_prometheus_rows_with_progress"]
+    K --> L["same metric-engine flag as samples, no batcher"]
+    L --> M["table: <metric>"]
 
-    L --> M["field: configured native-histogram Struct"]
-    M --> N["struct children: counts, spans, buckets, sum, schema"]
-    H --> P["written headers and counters"]
-    K --> P
+    M --> N["field: configured native-histogram Struct"]
+    N --> O["struct children: counts, spans, buckets, sum, schema"]
+    I --> P["written headers and counters"]
+    L --> P
 ```
 
 The conversion step splits one v2 request into two `ContextReq`s:
 
+The `decode` codec metric covers decompression and the borrowed request
+envelope. The `convert` metric covers per-series scans, leaf prost decoding,
+validation, and row construction.
+
 - samples keep the existing sample table name and can use the metric-engine
   physical table path and pending rows batcher;
 - native histograms keep the existing metric table name.

src/servers/src/row_writer.rs

@@ -92,6 +92,26 @@ impl TableData {
         Ok(index)
     }
 
+    /// Ensures a default column schema without allocating its name when it already exists.
+    pub(crate) fn ensure_column_by_name(
+        &mut self,
+        name: &str,
+        datatype: ColumnDataType,
+        semantic_type: SemanticType,
+    ) -> Result<usize> {
+        if let Some(index) = self.column_indexes.get(name).copied() {
+            check_schema(datatype, semantic_type, &self.schema[index])?;
+            return Ok(index);
+        }
+
+        self.ensure_column(ColumnSchema {
+            column_name: name.to_string(),
+            datatype: datatype as i32,
+            semantic_type: semantic_type as i32,
+            ..Default::default()
+        })
+    }
+
     #[allow(dead_code)]
     pub fn columns(&self) -> &Vec<ColumnSchema> {
         &self.schema
@@ -211,11 +231,14 @@ impl MultiTableData {
 }
 
 /// Write data as tags into the table data.
-pub fn write_tags(
+pub fn write_tags<K>(
     table_data: &mut TableData,
-    tags: impl Iterator<Item = (String, String)>,
+    tags: impl Iterator<Item = (K, String)>,
     one_row: &mut Vec<Value>,
-) -> Result<()> {
+) -> Result<()>
+where
+    K: AsRef<str> + Into<String>,
+{
     let ktv_iter = tags.map(|(k, v)| (k, ColumnDataType::String, Some(ValueData::StringValue(v))));
     write_by_semantic_type(table_data, SemanticType::Tag, ktv_iter, one_row)
 }
@@ -326,12 +349,15 @@ pub(crate) fn write_by_schema(
     Ok(())
 }
 
-fn write_by_semantic_type(
+fn write_by_semantic_type<K>(
     table_data: &mut TableData,
     semantic_type: SemanticType,
-    ktv_iter: impl Iterator<Item = (String, ColumnDataType, Option<ValueData>)>,
+    ktv_iter: impl Iterator<Item = (K, ColumnDataType, Option<ValueData>)>,
     one_row: &mut Vec<Value>,
-) -> Result<()> {
+) -> Result<()>
+where
+    K: AsRef<str> + Into<String>,
+{
     let TableData {
         schema,
         column_indexes,
@@ -339,12 +365,13 @@ fn write_by_semantic_type(
     } = table_data;
 
     for (name, datatype, value) in ktv_iter {
-        let index = column_indexes.get(&name);
+        let index = column_indexes.get(name.as_ref()).copied();
         if let Some(index) = index {
-            check_schema(datatype, semantic_type, &schema[*index])?;
-            one_row[*index].value_data = value;
+            check_schema(datatype, semantic_type, &schema[index])?;
+            one_row[index].value_data = value;
         } else {
             let index = schema.len();
+            let name = name.into();
             schema.push(ColumnSchema {
                 column_name: name.clone(),
                 datatype: datatype as i32,

src/servers/tests/prom_remote_write_v2_test.rs

@@ -77,7 +77,7 @@ fn test_decode_remote_write_v2_native_histogram_dump() {
     assert_eq!(histogram.timestamp, 1782358160412);
 
     let (sample_inserts, histogram_inserts, sample_count, histogram_count) =
-        remote_write_v2::write_requests(decoded).unwrap();
+        remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY), true).unwrap();
     assert!(sample_inserts.is_empty());
     assert_eq!(sample_count, 0);
     assert_eq!(histogram_count, 1);