rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-01 16:36:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: address review issue

Browse Source
This commit is contained in:
Ning Sun
2026-09-26 10:00:47 +08:00
parent 415942e6ce
commit 92a528b1e5
2 changed files with 349 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/servers/src/http/prom_store.rs
+141 -65
View File
@@ -396,39 +396,46 @@ async fn write_prometheus_rows_with_progress(
prom_store_with_metric_engine: bool,
mut batches: Vec<PromWriteBatch>,
) -> std::result::Result<PromWriteOutcome, PromWriteError> {
// 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<PromWriteBatch>,
prom_store_with_metric_engine: bool,
) -> std::result::Result<PromWriteOutcome, PromWriteError> {
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<Arc<LogicalTablePendingRowsBatcher>>,
@@ -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<PromWriteBatch>,
sample_batch_count: usize,
prom_store_with_metric_engine: bool,
) -> std::result::Result<PromWriteV2Outcome, PromWriteV2Error> {
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<B: PromWriteBatcher + ?Sized>(
prom_store_handler: PromStoreProtocolHandlerRef,
batcher: &B,
prom_store_with_metric_engine: bool,
sample_batches: Vec<PromWriteBatch>,
histogram_batches: Vec<PromWriteBatch>,
sample_batch_count: usize,
batches: Vec<PromWriteBatch>,
) -> std::result::Result<PromWriteV2Outcome, PromWriteV2Error> {
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")
};
tests-integration/tests/http.rs
+208 -4
View File
@@ -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::<servers::interceptor::PromStoreProtocolInterceptorRef<servers::error::Error>>(
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