rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 23:12:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(servers): bound ready Flight batch admission and cover lifecycle

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-15 16:18:21 +08:00
parent 931050a42d
commit 6573fc97e8
1 changed files with 485 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/servers/src/grpc/flight/stream.rs
+485 -105
View File
@@ -264,10 +264,8 @@ impl FlightRecordBatchStream {
let mut metrics = StreamMetrics::new(should_send_partial_metrics);
let mut last_metrics_str = None;
let mut sent_first_batch = false;
let recordbatch_schema = recordbatches.schema();
let schema = recordbatches.schema().arrow_schema().clone();
let schema = recordbatch_schema.arrow_schema().clone();
let start = Instant::now();
if let Err(e) = tx.send(Ok(FlightMessage::Schema(schema))).await {
warn!(e; "stop sending Flight data");
@@ -275,92 +273,48 @@ impl FlightRecordBatchStream {
}
metrics.send_schema_duration += start.elapsed();
loop {
let start = Instant::now();
let batch_or_err = if should_send_partial_metrics && can_send_metrics_before_batch {
match time::timeout(
FLIGHT_METRICS_HEARTBEAT_INTERVAL,
recordbatches.next().in_current_span(),
)
.await
{
Ok(result) => result,
Err(_) => {
if let Some(metrics_str) = recordbatches
.metrics()
.and_then(|m| serde_json::to_string(&m).ok())
&& !Self::send_metrics_if_changed(
&mut tx,
&mut metrics,
&mut last_metrics_str,
metrics_str,
)
.await
{
return;
}
metrics.fetch_content_duration += start.elapsed();
continue;
}
}
} else {
recordbatches.next().in_current_span().await
};
metrics.fetch_content_duration += start.elapsed();
let Some(batch_or_err) = batch_or_err else {
break;
};
match batch_or_err {
Ok(recordbatch) => {
metrics.total_rows += recordbatch.num_rows();
metrics.record_batch_count += 1;
metrics.total_bytes += recordbatch.df_record_batch().get_array_memory_size();
let mut batches = vec![recordbatch];
let mut eof = false;
let mut stream_error = None;
if !should_send_partial_metrics && sent_first_batch {
let mut rows = batches[0].num_rows();
let mut bytes = batches[0].df_record_batch().get_array_memory_size();
while batches.len() < MAX_BATCHES && rows < MAX_ROWS && bytes < MAX_BYTES {
let start = Instant::now();
let next =
poll_fn(|cx| Poll::Ready(recordbatches.as_mut().poll_next(cx)))
.await;
metrics.fetch_content_duration += start.elapsed();
match next {
Poll::Ready(Some(Ok(recordbatch))) => {
rows += recordbatch.num_rows();
bytes += recordbatch.df_record_batch().get_array_memory_size();
metrics.total_rows += recordbatch.num_rows();
metrics.record_batch_count += 1;
metrics.total_bytes +=
recordbatch.df_record_batch().get_array_memory_size();
batches.push(recordbatch);
}
Poll::Ready(Some(e)) => {
stream_error = Some(e);
break;
}
Poll::Ready(None) => {
eof = true;
break;
}
Poll::Pending => break,
if should_send_partial_metrics {
loop {
let start = Instant::now();
let batch_or_err = if can_send_metrics_before_batch {
match time::timeout(
FLIGHT_METRICS_HEARTBEAT_INTERVAL,
recordbatches.next().in_current_span(),
)
.await
{
Ok(result) => result,
Err(_) => {
if let Some(metrics_str) = recordbatches
.metrics()
.and_then(|m| serde_json::to_string(&m).ok())
&& !Self::send_metrics_if_changed(
&mut tx,
&mut metrics,
&mut last_metrics_str,
metrics_str,
)
.await
{
return;
}
metrics.fetch_content_duration += start.elapsed();
continue;
}
}
sent_first_batch = true;
let batches = if batches.len() >= 2 {
match merge_record_batches(recordbatch_schema.clone(), &batches) {
Ok(merged) => vec![merged],
Err(_) => batches,
}
} else {
batches
};
for recordbatch in batches {
} else {
recordbatches.next().in_current_span().await
};
metrics.fetch_content_duration += start.elapsed();
let Some(batch_or_err) = batch_or_err else {
break;
};
match batch_or_err {
Ok(recordbatch) => {
metrics.total_rows += recordbatch.num_rows();
metrics.record_batch_count += 1;
metrics.total_bytes +=
recordbatch.df_record_batch().get_array_memory_size();
let start = Instant::now();
if let Err(e) = tx
.send(Ok(FlightMessage::RecordBatch(
@@ -372,9 +326,18 @@ impl FlightRecordBatchStream {
return;
}
metrics.send_record_batch_duration += start.elapsed();
if let Some(metrics_str) = recordbatches
.metrics()
.and_then(|m| serde_json::to_string(&m).ok())
&& {
last_metrics_str = Some(metrics_str.clone());
!Self::send_metrics(&mut tx, &mut metrics, metrics_str).await
}
{
return;
}
}
if let Some(Err(e)) = stream_error {
Err(e) => {
if e.status_code().should_log_error() {
error!("{e:?}");
}
@@ -384,36 +347,141 @@ impl FlightRecordBatchStream {
}
return;
}
if eof {
break;
}
}
} else {
let mut sent_first_batch = false;
let mut pending = None;
loop {
let (batch_or_err, fetched) = match pending.take() {
Some(recordbatch) => (Ok(recordbatch), false),
None => {
let start = Instant::now();
let batch_or_err = recordbatches.next().in_current_span().await;
metrics.fetch_content_duration += start.elapsed();
let Some(batch_or_err) = batch_or_err else {
break;
};
(batch_or_err, true)
}
if should_send_partial_metrics
&& let Some(metrics_str) = recordbatches
.metrics()
.and_then(|m| serde_json::to_string(&m).ok())
&& {
last_metrics_str = Some(metrics_str.clone());
!Self::send_metrics(&mut tx, &mut metrics, metrics_str).await
};
let recordbatch = match batch_or_err {
Ok(recordbatch) => recordbatch,
Err(e) => {
if e.status_code().should_log_error() {
error!("{e:?}");
}
let e = Err(e).context(error::CollectRecordbatchSnafu);
if let Err(e) = tx.send(e.map_err(|x| x.into())).await {
warn!(e; "stop sending Flight data");
}
{
return;
}
};
let batch_rows = recordbatch.num_rows();
let batch_bytes = recordbatch.df_record_batch().get_array_memory_size();
if fetched {
metrics.total_rows += batch_rows;
metrics.record_batch_count += 1;
metrics.total_bytes += batch_bytes;
}
Err(e) => {
if !sent_first_batch || batch_rows >= MAX_ROWS || batch_bytes >= MAX_BYTES {
let start = Instant::now();
if let Err(e) = tx
.send(Ok(FlightMessage::RecordBatch(
recordbatch.into_df_record_batch(),
)))
.await
{
warn!(e; "stop sending Flight data");
return;
}
metrics.send_record_batch_duration += start.elapsed();
sent_first_batch = true;
continue;
}
let mut batches = vec![recordbatch];
let mut rows = batch_rows;
let mut bytes = batch_bytes;
let mut eof = false;
let mut stream_error = None;
loop {
if batches.len() == MAX_BATCHES || rows >= MAX_ROWS || bytes >= MAX_BYTES {
break;
}
let start = Instant::now();
let next =
poll_fn(|cx| Poll::Ready(recordbatches.as_mut().poll_next(cx))).await;
metrics.fetch_content_duration += start.elapsed();
match next {
Poll::Ready(Some(Ok(recordbatch))) => {
let batch_rows = recordbatch.num_rows();
let batch_bytes = recordbatch.df_record_batch().get_array_memory_size();
if batch_rows > MAX_ROWS - rows || batch_bytes > MAX_BYTES - bytes {
metrics.total_rows += batch_rows;
metrics.record_batch_count += 1;
metrics.total_bytes += batch_bytes;
pending = Some(recordbatch);
break;
}
metrics.total_rows += batch_rows;
metrics.record_batch_count += 1;
metrics.total_bytes += batch_bytes;
rows += batch_rows;
bytes += batch_bytes;
batches.push(recordbatch);
}
Poll::Ready(Some(Err(e))) => {
stream_error = Some(e);
break;
}
Poll::Ready(None) => {
eof = true;
break;
}
Poll::Pending => break,
}
}
let batches = if batches.len() >= 2 {
match merge_record_batches(recordbatch_schema.clone(), &batches) {
Ok(merged) => vec![merged],
Err(_) => batches,
}
} else {
batches
};
for recordbatch in batches {
let start = Instant::now();
if let Err(e) = tx
.send(Ok(FlightMessage::RecordBatch(
recordbatch.into_df_record_batch(),
)))
.await
{
warn!(e; "stop sending Flight data");
return;
}
metrics.send_record_batch_duration += start.elapsed();
}
if let Some(e) = stream_error {
if e.status_code().should_log_error() {
error!("{e:?}");
}
let e = Err(e).context(error::CollectRecordbatchSnafu);
if let Err(e) = tx.send(e.map_err(|x| x.into())).await {
warn!(e; "stop sending Flight data");
}
return;
}
if eof {
break;
}
}
}
// make last package to pass metrics
// Make the last package pass metrics exactly once at EOF.
if let Some(metrics_str) = recordbatches
.metrics()
.and_then(|m| serde_json::to_string(&m).ok())
@@ -506,6 +574,7 @@ mod test {
enum ScriptedItem {
Batch(common_recordbatch::error::Result<RecordBatch>),
Pending,
PersistentPending,
}
struct ScriptedBatchStream {
@@ -514,6 +583,11 @@ mod test {
poll_count: Arc<std::sync::atomic::AtomicUsize>,
}
struct DropFlagStream {
schema: SchemaRef,
dropped: Arc<std::sync::atomic::AtomicBool>,
}
fn query_context_with_matching_capability() -> Arc<QueryContext> {
let query_ctx = QueryContext::arc();
let remote_query_id = query_ctx
@@ -588,7 +662,7 @@ mod test {
}
fn metrics(&self) -> Option<RecordBatchMetrics> {
None
Some(RecordBatchMetrics::default())
}
}
@@ -604,11 +678,44 @@ mod test {
cx.waker().wake_by_ref();
Poll::Pending
}
Some(ScriptedItem::PersistentPending) => {
self.items.push_front(ScriptedItem::PersistentPending);
Poll::Pending
}
None => Poll::Ready(None),
}
}
}
impl RecordBatchStream for DropFlagStream {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn output_ordering(&self) -> Option<&[OrderOption]> {
None
}
fn metrics(&self) -> Option<RecordBatchMetrics> {
None
}
}
impl Stream for DropFlagStream {
type Item = common_recordbatch::error::Result<RecordBatch>;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Pending
}
}
impl Drop for DropFlagStream {
fn drop(&mut self) {
self.dropped
.store(true, std::sync::atomic::Ordering::Relaxed);
}
}
fn int_batch(schema: SchemaRef, values: impl IntoIterator<Item = i32>) -> RecordBatch {
RecordBatch::new(
schema,
@@ -617,8 +724,9 @@ mod test {
.unwrap()
}
async fn flight_messages(
async fn flight_messages_with_context(
recordbatches: SendableRecordBatchStream,
query_ctx: Arc<QueryContext>,
) -> Vec<TonicResult<FlightMessage>> {
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
@@ -626,7 +734,7 @@ mod test {
)),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
query_ctx,
);
let decoder = &mut FlightDecoder::default();
let mut messages = Vec::new();
@@ -643,6 +751,78 @@ mod test {
messages
}
async fn flight_messages(
recordbatches: SendableRecordBatchStream,
) -> Vec<TonicResult<FlightMessage>> {
flight_messages_with_context(recordbatches, QueryContext::arc()).await
}
#[tokio::test]
async fn test_drop_cancels_and_releases_upstream_stream() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false));
let stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(Box::pin(
DropFlagStream {
schema,
dropped: dropped.clone(),
},
))),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
);
drop(stream);
tokio::time::timeout(Duration::from_secs(1), async {
while !dropped.load(std::sync::atomic::Ordering::Relaxed) {
tokio::task::yield_now().await;
}
})
.await
.expect("dropping Flight stream must release upstream");
}
#[tokio::test]
async fn test_first_batch_is_sent_before_any_additional_poll() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let recordbatches = Box::pin(ScriptedBatchStream {
schema: schema.clone(),
items: VecDeque::from([
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::PersistentPending,
]),
poll_count: poll_count.clone(),
});
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
);
let decoder = &mut FlightDecoder::default();
let schema_data = stream.next().await.unwrap().unwrap();
assert!(matches!(
decoder.try_decode(&schema_data).unwrap().unwrap(),
FlightMessage::Schema(_)
));
let first = stream.next().await.unwrap().unwrap();
assert!(matches!(
decoder.try_decode(&first).unwrap().unwrap(),
FlightMessage::RecordBatch(_)
));
}
#[tokio::test]
async fn test_ready_only_coalesces_later_batches_without_polling_after_first() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
@@ -674,6 +854,111 @@ mod test {
assert_eq!(merged.num_rows(), 2);
}
#[tokio::test]
async fn test_ready_only_exact_row_cap_flushes() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let recordbatches = Box::pin(ScriptedBatchStream {
schema: schema.clone(),
items: VecDeque::from([
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..1023))),
ScriptedItem::PersistentPending,
]),
poll_count: poll_count.clone(),
});
let (tx, mut rx) = mpsc::channel::<TonicResult<FlightMessage>>(1);
let handle = tokio::spawn(FlightRecordBatchStream::flight_data_stream(
recordbatches,
tx,
false,
false,
));
assert!(matches!(
rx.next().await.unwrap().unwrap(),
FlightMessage::Schema(_)
));
tokio::time::timeout(Duration::from_secs(1), async {
while poll_count.load(std::sync::atomic::Ordering::Relaxed) != 3 {
tokio::task::yield_now().await;
}
})
.await
.expect("exact-cap group must be formed");
assert!(matches!(
rx.next().await.unwrap().unwrap(),
FlightMessage::RecordBatch(_)
));
let FlightMessage::RecordBatch(batch) = rx.next().await.unwrap().unwrap() else {
panic!("expected the exact-cap group");
};
assert_eq!(batch.num_rows(), 1024);
handle.abort();
}
#[tokio::test]
async fn test_pending_batch_is_sent_before_upstream_is_polled_again() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let recordbatches = Box::pin(ScriptedBatchStream {
schema: schema.clone(),
items: VecDeque::from([
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..1024))),
ScriptedItem::PersistentPending,
]),
poll_count: poll_count.clone(),
});
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
);
let decoder = &mut FlightDecoder::default();
let schema_data = stream.next().await.unwrap().unwrap();
assert!(matches!(
decoder.try_decode(&schema_data).unwrap().unwrap(),
FlightMessage::Schema(_)
));
let first = stream.next().await.unwrap().unwrap();
assert!(matches!(
decoder.try_decode(&first).unwrap().unwrap(),
FlightMessage::RecordBatch(_)
));
tokio::time::timeout(Duration::from_secs(1), async {
while poll_count.load(std::sync::atomic::Ordering::Relaxed) != 3 {
tokio::task::yield_now().await;
}
})
.await
.expect("second group must fetch the held batch");
let held = stream.next().await.unwrap().unwrap();
assert!(matches!(
decoder.try_decode(&held).unwrap().unwrap(),
FlightMessage::RecordBatch(_)
));
let pending = stream.next().await.unwrap().unwrap();
let FlightMessage::RecordBatch(batch) = decoder.try_decode(&pending).unwrap().unwrap()
else {
panic!("expected the held oversized batch");
};
assert_eq!(batch.num_rows(), 1024);
}
#[tokio::test]
async fn test_ready_only_flushes_before_pending_and_before_error() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
@@ -705,6 +990,101 @@ mod test {
assert_eq!(messages.len(), 5);
}
#[tokio::test]
async fn test_verbose_ready_batches_preserve_batch_metrics_order() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let query_ctx = QueryContext::arc();
query_ctx.set_explain_verbose(true);
let messages = flight_messages_with_context(
Box::pin(ScriptedBatchStream {
schema: schema.clone(),
items: VecDeque::from([
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))),
]),
poll_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}),
query_ctx,
)
.await;
assert!(matches!(messages[1], Ok(FlightMessage::RecordBatch(_))));
assert!(matches!(messages[2], Ok(FlightMessage::Metrics(_))));
assert!(matches!(messages[3], Ok(FlightMessage::RecordBatch(_))));
assert!(matches!(messages[4], Ok(FlightMessage::Metrics(_))));
}
#[tokio::test]
async fn test_ready_only_preserves_oversized_and_empty_batch_caps() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let oversized = int_batch(schema.clone(), 0..1024);
let mut items = vec![
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))),
ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))),
ScriptedItem::Batch(Ok(oversized)),
];
items.extend(
(0..16).map(|_| ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone())))),
);
let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let messages = flight_messages(Box::pin(ScriptedBatchStream {
schema,
items: items.into(),
poll_count: poll_count.clone(),
}))
.await;
let batches = messages
.iter()
.filter_map(|message| match message.as_ref().unwrap() {
FlightMessage::RecordBatch(batch) => Some(batch.num_rows()),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(batches, vec![1, 1, 1024, 0]);
assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 20);
}
#[tokio::test]
async fn test_ready_only_preserves_byte_oversized_batch() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::string_datatype(),
false,
)]));
let oversized = RecordBatch::new(
schema.clone(),
vec![Arc::new(datatypes::vectors::StringVector::from_slice(
&vec!["x".repeat(300); 1023],
)) as VectorRef],
)
.unwrap();
let messages = flight_messages(Box::pin(ScriptedBatchStream {
schema: schema.clone(),
items: VecDeque::from([
ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone()))),
ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone()))),
ScriptedItem::Batch(Ok(oversized)),
]),
poll_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}))
.await;
let batches = messages
.iter()
.filter_map(|message| match message.as_ref().unwrap() {
FlightMessage::RecordBatch(batch) => Some(batch.num_rows()),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(batches, vec![0, 0, 1023]);
}
#[tokio::test]
async fn test_ready_only_coalesces_dictionary_batches_with_different_mappings() {
let arrow_schema = Arc::new(datatypes::arrow::datatypes::Schema::new(vec![