rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-06 14:00:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: pass snapshot read bounds over flight (#8279)

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-06-11 20:05:02 +08:00
committed by GitHub
parent d76ecf67bd
commit 3608c90b2e
11 changed files with 273 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
src/client/src/database.rs
View File

@@ -36,7 +36,9 @@ use common_catalog::build_db_string;
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::ext::BoxedError;
use common_grpc::flight::do_put::DoPutResponse;
use common_grpc::flight::{FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage};
use common_grpc::flight::{
FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY,
};
use common_query::Output;
use common_recordbatch::adapter::RecordBatchMetrics;
use common_recordbatch::error::ExternalSnafu;
@@ -559,7 +561,28 @@ impl Database {
let value = serde_json::to_string(&flow_extensions.to_vec())
.expect("flow extension pairs should serialize");
let key = AsciiMetadataKey::from_static(FLOW_EXTENSIONS_METADATA_KEY);
Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value)
}
fn put_snapshot_seqs(
metadata: &mut MetadataMap,
snapshot_seqs: &std::collections::HashMap<u64, u64>,
) -> Result<()> {
if snapshot_seqs.is_empty() {
return Ok(());
}
let value =
serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize");
Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value)
}
fn put_metadata_value(
metadata: &mut MetadataMap,
key: &'static str,
value: String,
) -> Result<()> {
let key = AsciiMetadataKey::from_static(key);
let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?;
metadata.insert(key, value);
Ok(())
@@ -660,7 +683,7 @@ impl Database {
let request = Request::Query(QueryRequest {
query: Some(Query::Sql(sql.as_ref().to_string())),
});
self.do_get(request, hints, &[])
self.do_get(request, hints, &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
}
@@ -683,6 +706,7 @@ impl Database {
},
hints,
&[],
&Default::default(),
)
.await
}
@@ -695,6 +719,7 @@ impl Database {
},
&[],
&[],
&Default::default(),
)
.await
.map(OutputWithMetrics::into_output)
@@ -709,9 +734,15 @@ impl Database {
request: QueryRequest,
hints: &[(&str, &str)],
flow_extensions: &[(&str, &str)],
snapshot_seqs: &std::collections::HashMap<u64, u64>,
) -> Result<OutputWithMetrics> {
self.do_get(Request::Query(request), hints, flow_extensions)
.await
self.do_get(
Request::Query(request),
hints,
flow_extensions,
snapshot_seqs,
)
.await
}
/// Creates a new table using the provided table expression.
@@ -719,7 +750,7 @@ impl Database {
let request = Request::Ddl(DdlRequest {
expr: Some(DdlExpr::CreateTable(expr)),
});
self.do_get(request, &[], &[])
self.do_get(request, &[], &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
}
@@ -729,7 +760,7 @@ impl Database {
let request = Request::Ddl(DdlRequest {
expr: Some(DdlExpr::AlterTable(expr)),
});
self.do_get(request, &[], &[])
self.do_get(request, &[], &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
}
@@ -739,6 +770,7 @@ impl Database {
request: Request,
hints: &[(&str, &str)],
flow_extensions: &[(&str, &str)],
snapshot_seqs: &std::collections::HashMap<u64, u64>,
) -> Result<OutputWithMetrics> {
let request = self.to_rpc_request(request);
let request = Ticket {
@@ -749,6 +781,7 @@ impl Database {
let metadata = request.metadata_mut();
Self::put_hints(metadata, hints)?;
Self::put_flow_extensions(metadata, flow_extensions)?;
Self::put_snapshot_seqs(metadata, snapshot_seqs)?;
let mut client = self.client.make_flight_client(false, false)?;
@@ -928,6 +961,25 @@ mod tests {
);
}
#[test]
fn test_put_snapshot_seqs_preserves_u64_precision() {
let mut metadata = MetadataMap::new();
let snapshot_seqs = std::collections::HashMap::from([
(u64::MAX, u64::MAX - 1),
(9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64),
]);
Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap();
let value = metadata
.get(SNAPSHOT_SEQS_METADATA_KEY)
.unwrap()
.to_str()
.unwrap();
let decoded: std::collections::HashMap<u64, u64> = serde_json::from_str(value).unwrap();
assert_eq!(decoded, snapshot_seqs);
}
#[test]
fn test_flight_ctx() {
let mut ctx = FlightContext::default();

2
src/common/grpc/src/flight.rs
View File

@@ -39,6 +39,8 @@ use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result};
/// Flight metadata key used to carry flow query extensions as JSON pairs.
pub const FLOW_EXTENSIONS_METADATA_KEY: &str = "x-greptime-flow-extensions";
/// Flight metadata key used to carry query snapshot read upper bounds as JSON.
pub const SNAPSHOT_SEQS_METADATA_KEY: &str = "x-greptime-snapshot-seqs";
#[derive(Debug, Clone)]
pub enum FlightMessage {

7
src/flow/src/batching_mode/frontend_client.rs
View File

@@ -411,7 +411,12 @@ impl FrontendClient {
peer: db.peer.clone(),
});
db.database
.query_with_terminal_metrics_and_flow_extensions(request, &hints, extensions)
.query_with_terminal_metrics_and_flow_extensions(
request,
&hints,
extensions,
&Default::default(),
)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)

21
src/mito2/src/engine.rs
View File

@@ -132,6 +132,7 @@ use crate::engine::puffin_index::{IndexEntryContext, collect_index_entries_from_
use crate::error::{
IncrementalQueryStaleSnafu, InvalidRequestSnafu, JoinSnafu, MitoManifestInfoSnafu, RecvSnafu,
RegionNotFoundSnafu, Result, SerdeJsonSnafu, SerializeColumnMetadataSnafu,
SnapshotFenceStaleSnafu,
};
#[cfg(feature = "enterprise")]
use crate::extension::BoxedExtensionRangeProviderFactory;
@@ -1045,6 +1046,26 @@ impl EngineInner {
);
}
if let Some(given_seq) = request.memtable_max_sequence
&& !request.skip_sst_files
{
// Explicit snapshot fences that include SST reads are enforceable
// only while the requested upper bound is not older than the
// region's flushed frontier. If H has already been flushed into SST,
// mito cannot apply a memtable-only sequence upper bound to that
// SST scan, so fail and let Flow rebind the fenced repair instead
// of reading rows beyond H.
let min_enforceable_seq = version.flushed_sequence;
ensure!(
given_seq >= min_enforceable_seq,
SnapshotFenceStaleSnafu {
region_id,
given_seq,
min_enforceable_seq,
}
);
}
// Get cache.
let cache_manager = self.workers.cache_manager();

21
src/mito2/src/engine/scan_test.rs
View File

@@ -247,9 +247,10 @@ async fn test_scan_with_min_sst_sequence() {
}
#[tokio::test]
async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() {
async fn test_full_snapshot_upper_bound_returns_outdated_after_late_flush() {
let mut env =
TestEnv::with_prefix("test_full_snapshot_upper_bound_does_not_constrain_sst_rows").await;
TestEnv::with_prefix("test_full_snapshot_upper_bound_returns_outdated_after_late_flush")
.await;
let engine = env.create_engine(MitoConfig::default()).await;
let region_id = RegionId::new(1, 1);
@@ -277,7 +278,7 @@ async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() {
test_util::put_rows(&engine, region_id, second_rows).await;
test_util::flush_region(&engine, region_id, None).await;
let scanner = engine
let err = engine
.scanner(
region_id,
ScanRequest {
@@ -286,14 +287,14 @@ async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() {
},
)
.await
.unwrap();
.err()
.expect("expect stale snapshot fence error");
let stream = scanner.scan().await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let pretty = batches.pretty_print().unwrap();
assert!(pretty.contains("1970-01-01T00:00:03"));
assert!(pretty.contains("1970-01-01T00:00:04"));
assert_eq!(StatusCode::RequestOutdated, err.status_code());
let err_msg = err.to_string();
assert!(err_msg.contains("STALE_SNAPSHOT_FENCE"));
assert!(err_msg.contains(&region_id.to_string()));
assert!(err_msg.contains(&format!("given_seq: {snapshot_upper_bound}")));
}
#[tokio::test]

16
src/mito2/src/error.rs
View File

@@ -235,6 +235,20 @@ pub enum Error {
location: Location,
},
#[snafu(display(
"STALE_SNAPSHOT_FENCE: snapshot upper bound stale, region: {}, given_seq: {}, min_enforceable_seq: {}, retry_hint: REBIND_SNAPSHOT_FENCE",
region_id,
given_seq,
min_enforceable_seq
))]
SnapshotFenceStale {
region_id: RegionId,
given_seq: u64,
min_enforceable_seq: u64,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Old manifest missing for region {}", region_id))]
MissingOldManifest {
region_id: RegionId,
@@ -1356,7 +1370,7 @@ impl ErrorExt for Error {
| SerializePartitionExpr { .. }
| InvalidSourceAndTargetRegion { .. } => StatusCode::InvalidArguments,
IncrementalQueryStale { .. } => StatusCode::RequestOutdated,
IncrementalQueryStale { .. } | SnapshotFenceStale { .. } => StatusCode::RequestOutdated,
RegionMetadataNotFound { .. }
| Join { .. }

24
src/query/src/dummy_catalog.rs
View File

@@ -407,7 +407,15 @@ fn decide_flow_scan(query_ctx: &QueryContext, region_id: RegionId) -> Result<Flo
let memtable_max_sequence = query_ctx.get_snapshot(region_id.as_u64());
// `skip_sst_files` is only valid for memtable-only incremental deltas,
// identified by a lower checkpoint bound. A snapshot upper bound without an
// incremental lower bound is a fenced full-snapshot read and must keep SSTs
// in the scan so mito can reject stale upper bounds after H has been flushed
// into SSTs, instead of silently bypassing the stale-fence check. If a future
// incremental delta also carries an upper bound, the lower-bound stale check
// still proves whether memtable-only is safe.
let skip_sst_files = apply_incremental
&& memtable_min_sequence.is_some()
&& flow_extensions.incremental_mode == Some(FlowIncrementalMode::MemtableOnly);
Ok(FlowScanDecision {
@@ -742,16 +750,23 @@ mod tests {
assert_eq!(request.sst_min_sequence, Some(90));
assert_eq!(request.memtable_min_sequence, None);
assert!(!request.snapshot_on_scan);
assert!(!request.skip_sst_files);
}
#[test]
fn test_scan_request_from_query_context_reuses_existing_snapshot_for_incremental_scan() {
let region_id = test_region_id();
let query_ctx = QueryContextBuilder::default()
.extensions(HashMap::from([(
FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
format!(r#"{{"{}":10}}"#, region_id.as_u64()),
)]))
.extensions(HashMap::from([
(
FLOW_INCREMENTAL_MODE.to_string(),
"memtable_only".to_string(),
),
(
FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
format!(r#"{{"{}":10}}"#, region_id.as_u64()),
),
]))
.snapshot_seqs(Arc::new(RwLock::new(HashMap::from([(
region_id.as_u64(),
42_u64,
@@ -763,6 +778,7 @@ mod tests {
assert_eq!(request.memtable_min_sequence, Some(10));
assert_eq!(request.memtable_max_sequence, Some(42));
assert!(!request.snapshot_on_scan);
assert!(request.skip_sst_files);
}
#[test]

42
src/servers/src/grpc/flight.rs
View File

@@ -14,6 +14,7 @@
mod stream;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -30,6 +31,7 @@ use common_error::ext::ErrorExt;
use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse};
use common_grpc::flight::{
FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage,
SNAPSHOT_SEQS_METADATA_KEY,
};
use common_memory_manager::MemoryGuard;
use common_query::{Output, OutputData};
@@ -195,12 +197,13 @@ impl FlightCraft for GreptimeRequestHandler {
) -> TonicResult<Response<TonicStream<FlightData>>> {
let mut hints = hint_headers::extract_hints(request.metadata());
hints.extend(extract_flow_extensions(request.metadata())?);
let snapshot_seqs = extract_snapshot_seqs(request.metadata())?;
let ticket = request.into_inner().ticket;
let request =
GreptimeRequest::decode(ticket.as_ref()).context(error::InvalidFlightTicketSnafu)?;
let query_ctx =
create_query_context(Channel::Grpc, request.header.as_ref(), hints.clone())?;
create_query_context(Channel::Grpc, request.header.as_ref(), hints, snapshot_seqs)?;
// Validate flow hint syntax at the transport boundary before dispatching the request.
// This does not authorize or execute anything; `handle_request()` below still performs
// the normal frontend handling and auth checks before query execution.
@@ -218,7 +221,9 @@ impl FlightCraft for GreptimeRequestHandler {
);
let flight_compression = self.flight_compression;
async {
let output = self.handle_request(request, hints).await?;
let output = self
.handle_request_with_query_ctx(request, query_ctx.clone())
.await?;
let stream = to_flight_data_stream(
output,
TracingContext::from_current_span(),
@@ -542,21 +547,30 @@ impl Stream for PutRecordBatchRequestStream {
fn extract_flow_extensions(
metadata: &tonic::metadata::MetadataMap,
) -> TonicResult<Vec<(String, String)>> {
let Some(value) = metadata.get(FLOW_EXTENSIONS_METADATA_KEY) else {
return Ok(vec![]);
Ok(extract_json_metadata(metadata, FLOW_EXTENSIONS_METADATA_KEY)?.unwrap_or_default())
}
fn extract_snapshot_seqs(
metadata: &tonic::metadata::MetadataMap,
) -> TonicResult<HashMap<u64, u64>> {
Ok(extract_json_metadata(metadata, SNAPSHOT_SEQS_METADATA_KEY)?.unwrap_or_default())
}
fn extract_json_metadata<T: serde::de::DeserializeOwned>(
metadata: &tonic::metadata::MetadataMap,
key: &'static str,
) -> TonicResult<Option<T>> {
let Some(value) = metadata.get(key) else {
return Ok(None);
};
let value = value.to_str().map_err(|e| {
Status::invalid_argument(format!(
"Invalid {FLOW_EXTENSIONS_METADATA_KEY} metadata value: {e}"
))
})?;
let value = value
.to_str()
.map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata value: {e}")))?;
serde_json::from_str::<Vec<(String, String)>>(value).map_err(|e| {
Status::invalid_argument(format!(
"Invalid {FLOW_EXTENSIONS_METADATA_KEY} metadata JSON: {e}"
))
})
let parsed = serde_json::from_str::<T>(value)
.map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata JSON: {e}")))?;
Ok(Some(parsed))
}
fn to_flight_data_stream(

20
src/servers/src/grpc/greptime_handler.rs
View File

@@ -14,7 +14,9 @@
//! Handler for Greptime Database service. It's implemented by frontend.
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::time::Instant;
use api::helper::request_type;
@@ -74,13 +76,22 @@ impl GreptimeRequestHandler {
&self,
request: GreptimeRequest,
hints: Vec<(String, String)>,
) -> Result<Output> {
let header = request.header.as_ref();
let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?;
self.handle_request_with_query_ctx(request, query_ctx).await
}
pub(crate) async fn handle_request_with_query_ctx(
&self,
request: GreptimeRequest,
query_ctx: QueryContextRef,
) -> Result<Output> {
let query = request.request.context(InvalidQuerySnafu {
reason: "Expecting non-empty GreptimeRequest.",
})?;
let header = request.header.as_ref();
let query_ctx = create_query_context(Channel::Grpc, header, hints)?;
let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?;
query_ctx.set_current_user(user_info);
@@ -181,6 +192,7 @@ pub(crate) fn create_query_context(
channel: Channel,
header: Option<&RequestHeader>,
mut extensions: Vec<(String, String)>,
snapshot_seqs: HashMap<u64, u64>,
) -> Result<QueryContextRef> {
let (catalog, schema) = header
.map(|header| {
@@ -214,7 +226,8 @@ pub(crate) fn create_query_context(
.current_catalog(catalog)
.current_schema(schema)
.timezone(timezone)
.channel(channel);
.channel(channel)
.snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs)));
if let Some(x) = extensions
.iter()
@@ -314,8 +327,10 @@ mod tests {
"spoofed-regs".to_string(),
),
],
HashMap::from([(7, 88)]),
)
.unwrap();
assert_eq!(query_context.get_snapshot(7), Some(88));
assert_eq!(query_context.current_catalog(), "cat-a-log");
assert_eq!(query_context.current_schema(), DEFAULT_SCHEMA_NAME);
assert_eq!(
@@ -354,6 +369,7 @@ mod tests {
REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
"spoofed-query-id".to_string(),
)],
HashMap::new(),
)
.unwrap();

7
src/servers/src/grpc/prom_query_gateway.rs
View File

@@ -80,7 +80,12 @@ impl PrometheusGateway for PrometheusGatewayService {
};
let header = inner.header.as_ref();
let query_ctx = create_query_context(Channel::Promql, header, Default::default())?;
let query_ctx = create_query_context(
Channel::Promql,
header,
Default::default(),
Default::default(),
)?;
let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?;
query_ctx.set_current_user(user_info);

88
tests-integration/src/grpc/flight.rs
View File

@@ -14,11 +14,13 @@
#[cfg(test)]
mod test {
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use api::v1::auth_header::AuthScheme;
use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, SemanticType};
use api::v1::query_request::Query;
use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType};
use arrow_flight::FlightDescriptor;
use auth::user_provider_from_option;
use client::{Client, Database};
@@ -259,6 +261,90 @@ mod test {
);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_distributed_flight_snapshot_seqs_rejects_stale_sst_fence() {
common_telemetry::init_default_ut_logging();
let db = GreptimeDbClusterBuilder::new(
"test_distributed_flight_snapshot_seqs_rejects_stale_sst_fence",
)
.await
.build(false)
.await;
let runtime = common_runtime::global_runtime().clone();
let greptime_request_handler = GreptimeRequestHandler::new(
db.frontend.instance.clone(),
user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").ok(),
Some(runtime.clone()),
FlightCompression::default(),
);
let mut grpc_server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime)
.flight_handler(Arc::new(greptime_request_handler))
.build();
grpc_server
.start("127.0.0.1:0".parse::<SocketAddr>().unwrap())
.await
.unwrap();
let client = Client::with_urls(vec![grpc_server.bind_addr().unwrap().to_string()]);
let mut client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
client.set_auth(AuthScheme::Basic(Basic {
username: "greptime_user".to_string(),
password: "greptime_pwd".to_string(),
}));
create_table(&client).await;
test_put_record_batches(&client, create_record_batches(1)).await;
client.sql("admin flush_table('foo')").await.unwrap();
let regions = db.list_all_regions().await;
assert_eq!(regions.len(), 1);
let (region_id, region) = regions.into_iter().next().unwrap();
let stale_snapshot_seq = region.find_committed_sequence();
test_put_record_batches(&client, create_record_batches(10)).await;
client.sql("admin flush_table('foo')").await.unwrap();
assert!(
region.find_committed_sequence() > stale_snapshot_seq,
"second flush should make the first snapshot sequence stale"
);
let result = client
.query_with_terminal_metrics_and_flow_extensions(
QueryRequest {
query: Some(Query::Sql(
"select ts, a, `B` from foo order by ts".to_string(),
)),
},
&[],
&[("flow.return_region_seq", "true")],
&HashMap::from([(region_id.as_u64(), stale_snapshot_seq)]),
)
.await
.unwrap();
let OutputData::Stream(mut stream) = result.output.data else {
panic!("expected stream output");
};
let mut stale_fence_error = None;
while let Some(batch) = stream.next().await {
if let Err(err) = batch {
stale_fence_error = Some(format!("{err:?}"));
break;
}
}
let err_msg = stale_fence_error.expect("expected stale snapshot fence rejection");
assert!(
err_msg.contains("STALE_SNAPSHOT_FENCE")
|| err_msg.contains("RequestOutdated")
|| err_msg.contains("snapshot upper bound stale"),
"expected stale snapshot fence rejection, got: {err_msg}"
);
}
async fn test_put_record_batches(client: &Database, record_batches: Vec<RecordBatch>) {
let requests_count = record_batches.len();
let schema = record_batches[0].schema.arrow_schema().clone();