From 92a528b1e5cec7f7d7588732f84bc01f7678c824 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sat, 26 Sep 2026 10:00:47 +0800 Subject: [PATCH] fix: address review issue --- src/servers/src/http/prom_store.rs | 206 +++++++++++++++++++--------- tests-integration/tests/http.rs | 212 ++++++++++++++++++++++++++++- 2 files changed, 349 insertions(+), 69 deletions(-) diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 42b4da6d12c..44abb18a5be 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -396,39 +396,46 @@ async fn write_prometheus_rows_with_progress( prom_store_with_metric_engine: bool, mut batches: Vec, ) -> std::result::Result { - // The bulk encode produces millisecond batches only; a physical table - // The bulk encode produces millisecond batches only; write targets and - // existing destination tables with another time index unit must stay on - // the ordinary insert path, which converts the requests to each - // destination table's unit. - let batcher = match (prom_store_with_metric_engine, pending_rows_batcher) { - (true, Some(batcher)) if batcher.accepts_bulk_time_indexes(batches.iter()).await => { - Some(batcher) - } - _ => None, - }; - if let Some(batcher) = batcher { + if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { + // Preflight before the bulk eligibility decision: pre_write hooks + // may redirect the contexts (e.g. to a per-tenant schema), so + // eligibility must be evaluated against the prepared batches — + // and the fallback below reuses them without re-running hooks or + // admission. preflight_prometheus_rows(&prom_store_handler, &mut batches) .await .map_err(|error| PromWriteError { error, rows_written: 0, })?; - let mut rows_written = 0; - for (temp_ctx, reqs) in batches { - let rows = batcher - .submit(reqs, temp_ctx) - .await - .map_err(|error| PromWriteError { - error, - rows_written, - })?; - rows_written += rows; + // The bulk encode produces millisecond batches only; write targets + // and existing destination tables with another time index unit stay + // on the ordinary insert path, which converts the requests to each + // destination table's unit. + if batcher.accepts_bulk_time_indexes(batches.iter()).await { + let mut rows_written = 0; + for (temp_ctx, reqs) in batches { + let rows = + batcher + .submit(reqs, temp_ctx) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + rows_written += rows; + } + return Ok(PromWriteOutcome { + write_cost: 0, + rows_written, + }); } - return Ok(PromWriteOutcome { - write_cost: 0, - rows_written, - }); + return write_prepared_prometheus_rows_with_progress( + prom_store_handler, + batches, + prom_store_with_metric_engine, + ) + .await; } let row_counts = batches @@ -467,6 +474,33 @@ async fn write_prometheus_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path; the pre_write hooks and admission have already run. +async fn write_prepared_prometheus_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut rows_written = 0; + for (ctx, request) in batches { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + write_cost += output.meta.cost; + rows_written += rows; + } + Ok(PromWriteOutcome { + write_cost, + rows_written, + }) +} + async fn write_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, pending_rows_batcher: Option>, @@ -494,31 +528,44 @@ async fn write_prometheus_v2_rows_with_progress( }); } - let batcher_eligible = match (&pending_rows_batcher, prom_store_with_metric_engine) { - (Some(batcher), true) => { - batcher - .accepts_bulk_time_indexes(sample_batches.iter().chain(&histogram_batches)) - .await + let sample_batch_count = sample_batches.len(); + let mut batches = sample_batches; + batches.extend(histogram_batches); + + if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { + // Same ordering as the v1 path: preflight (which may redirect the + // contexts) before the bulk eligibility decision, and the fallback + // reuses the prepared batches without re-running hooks or admission. + preflight_prometheus_rows(&prom_store_handler, &mut batches) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written: 0, + histograms_written: 0, + })?; + // The bulk encode produces millisecond batches only; write targets + // and existing destination tables with another time index unit stay + // on the ordinary insert path, which converts the requests to each + // destination table's unit. + if batcher.accepts_bulk_time_indexes(batches.iter()).await { + return write_batched_prometheus_v2_rows_with_progress( + prom_store_handler, + batcher.as_ref(), + prom_store_with_metric_engine, + sample_batch_count, + batches, + ) + .await; } - _ => false, - }; - if batcher_eligible { - // Safety: `batcher_eligible` is only true when `pending_rows_batcher` - // is `Some`. - let batcher = pending_rows_batcher.as_deref().unwrap(); - return write_batched_prometheus_v2_rows_with_progress( + return write_prepared_prometheus_v2_rows_with_progress( prom_store_handler, - batcher, + batches, + sample_batch_count, prom_store_with_metric_engine, - sample_batches, - histogram_batches, ) .await; } - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); let row_counts = batches .iter() .map(|(_, request)| prom_write_row_count(request)) @@ -566,24 +613,49 @@ async fn write_prometheus_v2_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path with v2 partial-progress accounting; the pre_write hooks and +/// admission have already run. +async fn write_prepared_prometheus_v2_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + sample_batch_count: usize, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut samples_written = 0; + let mut histograms_written = 0; + for (index, (ctx, request)) in batches.into_iter().enumerate() { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written, + histograms_written, + })?; + write_cost += output.meta.cost; + if index < sample_batch_count { + samples_written += rows; + } else { + histograms_written += rows; + } + } + Ok(PromWriteV2Outcome { + write_cost, + samples_written, + histograms_written, + }) +} + async fn write_batched_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, batcher: &B, prom_store_with_metric_engine: bool, - sample_batches: Vec, - histogram_batches: Vec, + sample_batch_count: usize, + batches: Vec, ) -> std::result::Result { - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); - preflight_prometheus_rows(&prom_store_handler, &mut batches) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written: 0, - histograms_written: 0, - })?; - let mut samples_written = 0; let mut histograms_written = 0; let mut write_cost = 0; @@ -884,14 +956,18 @@ mod tests { events: events.clone(), }; - let Ok(outcome) = write_batched_prometheus_v2_rows_with_progress( - handler, - &batcher, - true, - vec![test_prom_write_batch("sample")], - vec![test_prom_write_batch("histogram")], - ) - .await + let mut batches = vec![ + test_prom_write_batch("sample"), + test_prom_write_batch("histogram"), + ]; + // The caller preflights before choosing the bulk path; the batched + // writer consumes the prepared batches. + preflight_prometheus_rows(&handler, &mut batches) + .await + .unwrap(); + let Ok(outcome) = + write_batched_prometheus_v2_rows_with_progress(handler, &batcher, true, 1, batches) + .await else { panic!("mixed remote write should succeed") }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 3d1ee007eb7..238497753d6 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -15,18 +15,20 @@ use std::collections::BTreeMap; use std::io::Write; use std::str::FromStr; +use std::sync::Arc; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::metadata::MetricType as RemoteWriteV2MetricType; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Sample as RemoteWriteV2Sample, - TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Request as RemoteWriteV2Request, + Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, TimeSeries, WriteRequest, }; +use api::v1::RowInsertRequests; use auth::{UserProviderRef, user_provider_from_option}; use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode}; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -81,11 +83,13 @@ use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store::{self, mock_timeseries_new_label}; use servers::request_memory_limiter::ServerMemoryLimiter; +use session::context::QueryContextRef; use standalone::options::StandaloneOptions; use table::table_name::TableName; use tests_integration::test_util::{ - MockInstanceImpl, StorageType, assert_wal_delta, build_test_prom_server, setup_test_http_app, - setup_test_http_app_with_frontend, setup_test_http_app_with_frontend_and_slow_query_threshold, + MockInstanceImpl, StorageType, TestGuard, assert_wal_delta, build_test_prom_server, + setup_test_http_app, setup_test_http_app_with_frontend, + setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, setup_test_prom_app_with_frontend_batched, }; @@ -169,6 +173,8 @@ macro_rules! http_tests { test_prometheus_remote_write_v2_native_histogram, test_prometheus_remote_write_batched, test_prometheus_remote_write_batched_mixed_time_index_units, + test_prometheus_remote_write_batched_interceptor_time_index_units, + test_prometheus_remote_write_v2_batched_interceptor_time_index_units, test_prometheus_remote_special_labels, test_prometheus_remote_schema_labels, test_prometheus_remote_write_with_pipeline, @@ -3873,6 +3879,204 @@ pub async fn test_prometheus_remote_write_batched_mixed_time_index_units(store_t guard.remove_all().await; } +/// A Prometheus write interceptor that redirects every remote write to the +/// `tenant_redirect` schema, like a per-tenant redirection: the incoming +/// context targets `public` while the destination tables live elsewhere. +struct PromSchemaRedirectInterceptor; + +impl servers::interceptor::PromStoreProtocolInterceptor for PromSchemaRedirectInterceptor { + type Error = servers::error::Error; + + fn pre_write( + &self, + _write_req: &RowInsertRequests, + ctx: QueryContextRef, + ) -> servers::error::Result<()> { + ctx.set_current_schema("tenant_redirect"); + Ok(()) + } +} + +async fn setup_redirecting_batched_prom_app( + store_type: StorageType, + name: &str, +) -> (TestClient, TestGuard) { + let plugins = Plugins::default(); + plugins.insert::>( + Arc::new(PromSchemaRedirectInterceptor), + ); + let standalone = tests_integration::standalone::GreptimeDbStandaloneBuilder::new(name) + .with_default_store_type(store_type) + .with_plugin(plugins) + .build() + .await; + let server = build_test_prom_server(standalone.fe_instance().clone(), true) + .with_greptime_config_options(standalone.opts.datanode_options().to_toml().unwrap()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + (client, standalone.guard) +} + +async fn create_redirect_schema_with_microsecond_table(client: &TestClient, table: &str) { + for sql in [ + "create database if not exists tenant_redirect", + "CREATE TABLE tenant_redirect.phy_us \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('physical_metric_table' = 'true')", + &format!( + "CREATE TABLE tenant_redirect.{table} \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('on_physical_table' = 'phy_us')" + ), + ] { + let res = client + .get(format!("/v1/sql?sql={sql}").as_str()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "setup: {sql}"); + } +} + +/// Regression test: with batching enabled and a pre_write interceptor that +/// redirects the context to a schema holding same-named non-millisecond +/// tables, the bulk eligibility must be evaluated against the redirected +/// destinations (after preflight) and fall back to the ordinary insert path +/// instead of failing in the bulk encode. +pub async fn test_prometheus_remote_write_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_batched_interceptor_units").await; + create_redirect_schema_with_microsecond_table(&client, "intercept_metric").await; + + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "intercept_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.intercept_metric", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + +/// The v2 variant of the interceptor regression test: a mixed samples-plus- +/// histograms request whose existing sample destination lives in the +/// redirected schema with a microsecond time index must fall back to the +/// ordinary insert path (both series land), not fail in the bulk encode. +pub async fn test_prometheus_remote_write_v2_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_v2_batched_interceptor_units") + .await; + create_redirect_schema_with_microsecond_table(&client, "v2_mixed_sample").await; + + // One sample series targeting the existing microsecond table, plus one + // histogram series for a new (millisecond) table. + let mut symbols = vec![String::new()]; + let mut symbol = |value: &str| { + symbols.push(value.to_string()); + (symbols.len() - 1) as u32 + }; + let name_ref = symbol("__name__"); + let sample_ref = symbol("v2_mixed_sample"); + let job_ref = symbol("job"); + let demo_ref = symbol("demo"); + let name_ref2 = symbol("__name__"); + let histo_ref = symbol("v2_mixed_histo"); + let request = RemoteWriteV2Request { + symbols, + timeseries: vec![ + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref, sample_ref, job_ref, demo_ref], + samples: vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ..Default::default() + }, + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref2, histo_ref, job_ref, demo_ref], + histograms: vec![Histogram { + count: Some(Count::CountInt(1)), + sum: 1.0, + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![1], + timestamp: 1000, + start_timestamp: 1000, + ..Default::default() + }], + ..Default::default() + }, + ], + }; + let compressed = prom_store::snappy_compress(&request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_sample", + "[[1]]", + ) + .await; + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_histo", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + /// Covers the batched (pending-rows-batcher) Prometheus remote write path, which /// bypasses `PromStoreProtocolHandler::write`. Verifies the metric table is created /// asynchronously and still carries the Prometheus semantic identity stamped on the