diff --git a/src/client/src/client.rs b/src/client/src/client.rs index ed59f6701e7..1746854cd26 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -86,6 +86,7 @@ struct Inner { impl Default for Inner { fn default() -> Self { + #[allow(deprecated)] Self::with_manager_and_peers(ChannelManager::new(), Vec::new(), ClientOptions::default()) } } @@ -104,6 +105,7 @@ struct Peers { } impl Inner { + #[deprecated(note = "legacy single-manager path shares its pool between lanes")] fn with_manager_and_peers( channel_manager: ChannelManager, peers: Vec, @@ -112,6 +114,7 @@ impl Inner { Self::with_managers_and_peers(channel_manager.clone(), channel_manager, peers, options) } + // The explicit dual-manager path keeps query and control pools independent. fn with_managers_and_peers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, @@ -221,36 +224,61 @@ fn random_initial_delay(max_delay: Duration) -> Duration { } impl Client { + /// Creates a client whose query and control lanes intentionally share the default manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn new() -> Self { Default::default() } + /// Creates a client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_urls(urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_urls_and_options(urls, ClientOptions::default()) } /// Creates a client with URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_urls_and_options(urls: A, options: ClientOptions) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(ChannelManager::new(), urls, options) } + /// Creates a TLS client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls(urls: A, client_tls: ClientTlsOption) -> Result where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_tls_and_urls_and_options(urls, client_tls, ClientOptions::default()) } /// Creates a client with TLS URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls_and_options( urls: A, client_tls: ClientTlsOption, @@ -264,6 +292,7 @@ impl Client { let tls_config = load_client_tls_config(Some(client_tls)).context(error::CreateTlsChannelSnafu)?; let channel_manager = ChannelManager::with_config(channel_config, tls_config); + #[allow(deprecated)] Ok(Self::with_manager_and_urls_and_options( channel_manager, urls, @@ -271,15 +300,24 @@ impl Client { )) } + /// Creates a client with one manager shared intentionally by query and control lanes. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_manager_and_urls(channel_manager: ChannelManager, urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(channel_manager, urls, ClientOptions::default()) } - pub(crate) fn with_managers_and_urls( + /// Creates a client with query and control lanes backed by the supplied managers. + /// + /// The lanes are isolated only when the supplied managers are independently constructed; + /// this constructor does not enforce that they are distinct. + pub fn with_query_and_control_managers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -288,7 +326,7 @@ impl Client { U: AsRef, A: AsRef<[U]>, { - Self::with_managers_and_urls_and_options( + Self::with_query_and_control_managers_and_options( query_channel_manager, control_channel_manager, urls, @@ -296,7 +334,7 @@ impl Client { ) } - fn with_managers_and_urls_and_options( + fn with_query_and_control_managers_and_options( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -322,6 +360,11 @@ impl Client { } /// Creates a client with a channel manager, URLs, and custom options. + /// + /// The query and control lanes intentionally share this manager and its pool. + #[deprecated( + note = "shares this manager and its pool between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_manager_and_urls_and_options( channel_manager: ChannelManager, urls: A, @@ -332,7 +375,8 @@ impl Client { A: AsRef<[U]>, { let channel_manager_for_query = channel_manager.clone(); - Self::with_managers_and_urls_and_options( + // Legacy constructors intentionally share the manager and therefore its pool. + Self::with_query_and_control_managers_and_options( channel_manager_for_query, channel_manager, urls, @@ -529,6 +573,23 @@ impl Client { Ok(()) } + /// Returns the number of cached channels in the query and control pools for tests. + #[cfg(feature = "testing")] + pub fn channel_pool_sizes(&self) -> (usize, usize) { + let pool_size = |channel_manager: &ChannelManager| { + let mut size = 0; + channel_manager.retain_channel(|_, _| { + size += 1; + true + }); + size + }; + ( + pool_size(&self.inner.query_channel_manager), + pool_size(&self.inner.control_channel_manager), + ) + } + /// Returns peer addresses grouped by active and inactive state for tests. #[cfg(feature = "testing")] pub fn peer_addresses_by_state(&self) -> (Vec, Vec) { @@ -547,6 +608,7 @@ impl Client { } #[cfg(test)] +#[allow(deprecated)] mod tests { use std::collections::HashSet; use std::sync::Arc; @@ -654,6 +716,36 @@ mod tests { ] } + #[tokio::test] + async fn test_explicit_dual_manager_constructor_uses_isolated_reused_channel_pools() { + let query_channel_manager = ChannelManager::new(); + let control_channel_manager = ChannelManager::new(); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + ["127.0.0.1:3001"], + ); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(0, channel_pool_size(&control_channel_manager)); + + client.make_control_flight_client(false, false).unwrap(); + client.make_control_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(1, channel_pool_size(&control_channel_manager)); + } + + fn channel_pool_size(manager: &ChannelManager) -> usize { + let mut size = 0; + manager.retain_channel(|_, _| { + size += 1; + true + }); + size + } + #[test] fn test_inner() { let inner = Inner::default(); diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index 008ada00d23..4d723361e80 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -77,6 +77,7 @@ impl FlownodeManager for NodeClients { } impl NodeClients { + /// Creates node clients with independent query and control channel pools. pub fn new(config: ChannelConfig) -> Self { Self { query_channel_manager: ChannelManager::with_config(config.clone(), None), @@ -91,7 +92,7 @@ impl NodeClients { pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { - Client::with_managers_and_urls( + Client::with_query_and_control_managers( self.query_channel_manager.clone(), self.control_channel_manager.clone(), vec![datanode.addr.clone()], @@ -153,6 +154,7 @@ mod tests { assert_pool_has_one_address(&node_clients.control_channel_manager); let manager = ChannelManager::new(); + #[allow(deprecated)] let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]); legacy.make_flight_client(false, false).unwrap(); legacy.make_control_flight_client(false, false).unwrap(); diff --git a/src/client/src/database.rs b/src/client/src/database.rs index e97fff6db00..de3c722a86a 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use std::task::{Context, Poll}; use std::time::Duration; @@ -55,6 +55,7 @@ use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; use snafu::{IntoError, ResultExt}; +use tokio::sync::Notify; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; @@ -70,20 +71,47 @@ type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; const HINTS_METADATA_KEY: &str = "x-greptime-hints"; +/// Maximum time to wait for the optional trailing metrics message after +/// affected rows have already been delivered. +const FLIGHT_TRAILING_METRICS_TIMEOUT: Duration = Duration::from_secs(5); /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully -/// drained and [`Self::is_ready`] returns `true`. +/// drained and [`Self::is_ready`] returns `true`. Affected-row outputs may +/// briefly await a compatibility trailing metrics message. #[derive(Debug, Clone, Default)] pub struct OutputMetrics { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct OutputMetricsInner { metrics: RwLock>, + completion_error: RwLock>, ready: AtomicBool, + ready_notify: Notify, + compatibility_task: Mutex>, +} + +impl Default for OutputMetricsInner { + fn default() -> Self { + Self { + metrics: RwLock::new(None), + completion_error: RwLock::new(None), + ready: AtomicBool::new(false), + ready_notify: Notify::new(), + compatibility_task: Mutex::new(None), + } + } +} + +impl Drop for OutputMetricsInner { + fn drop(&mut self) { + if let Some(handle) = self.compatibility_task.get_mut().unwrap().take() { + handle.abort(); + } + } } impl OutputMetrics { @@ -98,10 +126,45 @@ impl OutputMetrics { /// Marks the terminal metrics as final for this output. pub fn mark_ready(&self) { - let _ = self + if self .inner .ready - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire); + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.inner.ready_notify.notify_waiters(); + } + } + + /// Waits until terminal metrics are final. + pub async fn wait_ready(&self) { + loop { + let notified = self.inner.ready_notify.notified(); + if self.is_ready() { + return; + } + notified.await; + } + } + + /// Returns an error encountered while completing the output, if any. + pub fn completion_error(&self) -> Option { + self.inner.completion_error.read().unwrap().clone() + } + + fn set_completion_error(&self, error: impl Into) { + *self.inner.completion_error.write().unwrap() = Some(error.into()); + } + + fn set_compatibility_task(&self, handle: tokio::task::AbortHandle) { + let mut task = self.inner.compatibility_task.lock().unwrap(); + if !self.is_ready() { + *task = Some(handle); + } + } + + fn take_compatibility_task(&self) -> Option { + self.inner.compatibility_task.lock().unwrap().take() } /// Returns whether terminal metrics are final. @@ -148,7 +211,8 @@ impl OutputMetrics { /// /// The contained [`OutputMetrics`] lets callers read stream terminal metrics /// after consuming `output`. For non-stream outputs, metrics are ready -/// immediately. +/// immediately. Flight affected-row outputs without inline metrics require +/// [`OutputMetrics::wait_ready`] before compatibility trailing metrics can be read. #[derive(Debug)] pub struct OutputWithMetrics { pub output: Output, @@ -194,6 +258,55 @@ fn parse_terminal_metrics(metrics_json: &str) -> Result { }) } +fn spawn_affected_rows_trailing_metrics_task( + terminal_metrics: &OutputMetrics, + mut reader: FlightMessageReader, +) where + S: Stream> + Send + Unpin + 'static, +{ + let metrics_ref = Arc::downgrade(&terminal_metrics.inner); + let remote_addr = reader.remote_addr().to_string(); + let task = common_runtime::spawn_global(async move { + let result = + tokio::time::timeout(FLIGHT_TRAILING_METRICS_TIMEOUT, reader.read_next()).await; + let Some(inner) = metrics_ref.upgrade() else { + return; + }; + let metrics = OutputMetrics { inner }; + match result { + Ok(Ok(Some(FlightMessage::Metrics(s)))) => match parse_terminal_metrics(&s) { + Ok(metrics_json) => metrics.update(Some(metrics_json)), + Err(error) => { + metrics.set_completion_error(error.to_string()); + warn!( + "Failed to decode trailing Flight metrics from {}: {}", + remote_addr, error + ); + } + }, + Ok(Ok(None)) => {} + Ok(Ok(Some(other))) => { + let error = format!("Unexpected trailing Flight message: {other:?}"); + metrics.set_completion_error(error.clone()); + warn!("{} from {}", error, remote_addr); + } + Ok(Err(error)) => { + let error = flight_stream_error(&remote_addr, error); + metrics.set_completion_error(error.to_string()); + warn!("{}", error); + } + Err(_) => { + let error = "Timed out waiting for trailing Flight metrics"; + metrics.set_completion_error(error); + warn!("{} from {}", error, remote_addr); + } + } + metrics.mark_ready(); + metrics.take_compatibility_task(); + }); + terminal_metrics.set_compatibility_task(task.abort_handle()); +} + struct StreamWithMetrics { stream: common_recordbatch::SendableRecordBatchStream, metrics: OutputMetrics, @@ -233,6 +346,9 @@ impl Stream for StreamWithMetrics { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let polled = Pin::new(&mut self.stream).poll_next(cx); + if let Poll::Ready(Some(Err(error))) = &polled { + self.metrics.set_completion_error(error.to_string()); + } if let Poll::Ready(None) = &polled { self.sync_terminal_metrics(); self.metrics.mark_ready(); @@ -280,32 +396,12 @@ where FlightMessage::AffectedRows { rows, metrics } => { let terminal_metrics = OutputMetrics::new(); if let Some(metrics) = metrics { + // Inline metrics are authoritative. They complete the output + // without touching the rest of the Flight stream. terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); - } - let next_message = reader - .read_next() - .await - .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; - match next_message { - None => terminal_metrics.mark_ready(), - Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { - terminal_metrics.update(Some(parse_terminal_metrics(&s)?)); - terminal_metrics.mark_ready(); - } - Some(FlightMessage::Metrics(_)) => { - return IllegalFlightMessagesSnafu { - reason: "'AffectedRows' Flight metadata already carries Metrics and cannot be followed by another Metrics message", - } - .fail(); - } - Some(other) => { - return IllegalFlightMessagesSnafu { - reason: format!( - "'AffectedRows' Flight message can only be followed by a Metrics message, got {other:?}" - ), - } - .fail(); - } + terminal_metrics.mark_ready(); + } else { + spawn_affected_rows_trailing_metrics_task(&terminal_metrics, reader); } Ok(OutputWithMetrics { output: Output::new_with_affected_rows(rows), @@ -749,7 +845,9 @@ impl Database { /// Executes a SQL query and returns the output with terminal metrics. /// /// For stream outputs, callers must consume the stream before reading final - /// terminal metrics from [`OutputWithMetrics::metrics`]. + /// terminal metrics from [`OutputWithMetrics::metrics`]. For affected-row + /// outputs without inline metrics, call [`OutputMetrics::wait_ready`] when + /// compatibility trailing metrics are required. pub async fn sql_with_terminal_metrics( &self, sql: S, @@ -986,6 +1084,7 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; + use tokio::sync::oneshot; use tonic::codegen::http::{HeaderMap, HeaderValue}; use tonic::metadata::MetadataMap; use tonic::{Code, Status}; @@ -1296,11 +1395,6 @@ mod tests { ); } - #[test] - fn test_parse_terminal_metrics_rejects_invalid_json() { - assert!(parse_terminal_metrics("{not-json}").is_err()); - } - #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { let output = output_from_flight_message_stream( @@ -1322,27 +1416,191 @@ mod tests { } #[tokio::test] - async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { + async fn test_affected_rows_inline_metrics_do_not_poll_trailer() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream( + let output = output_from_flight_message_stream( "test-peer".to_string(), futures_util::stream::iter(vec![ Ok(FlightMessage::AffectedRows { rows: 3, - metrics: Some(metrics_json.clone()), + metrics: Some(metrics_json), }), - Ok(FlightMessage::Metrics(metrics_json)), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(99))), ] as Vec>), ) .await - .unwrap_err(); + .unwrap(); - assert!( - err.to_string().contains("already carries Metrics"), - "unexpected error: {err:?}" + assert!(output.metrics.is_ready()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) ); } + #[tokio::test] + async fn test_affected_rows_without_inline_metrics_becomes_ready_after_trailer() { + // Hold the trailer until the not-ready state has been observed. + let (trailer_tx, trailer_rx) = oneshot::channel(); + let trailer = futures_util::stream::once(trailer_rx).map(|ready| { + ready.unwrap(); + Ok(FlightMessage::Metrics(terminal_metrics_json())) + }); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })]) + .chain(trailer), + ) + .await + .unwrap(); + + assert!(!output.metrics.is_ready()); + trailer_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), output.metrics.wait_ready()) + .await + .expect("terminal metrics must become ready once the trailer arrives"); + assert!(output.metrics.completion_error().is_none()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) + ); + } + + #[tokio::test] + async fn test_affected_rows_malformed_trailer_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("Invalid terminal metrics message")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_transport_trailer_error_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Err(Status::unavailable("trailer read failed").into()), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("trailer read failed")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_compatibility_reader_is_cancelled_after_second_poll_begins() { + struct DropProbe { + first: Option>, + polled: Option>, + dropped: Option>, + } + + impl Stream for DropProbe { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if let Some(message) = self.first.take() { + return Poll::Ready(Some(message)); + } + if let Some(polled) = self.polled.take() { + let _ = polled.send(()); + } + Poll::Pending + } + } + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } + } + + let (polled_tx, polled_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + DropProbe { + first: Some(Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })), + polled: Some(polled_tx), + dropped: Some(dropped_tx), + }, + ) + .await + .unwrap(); + polled_rx.await.unwrap(); + drop(output); + // Require cancellation before the five-second compatibility timeout. + tokio::time::timeout(Duration::from_secs(1), dropped_rx) + .await + .expect("the compatibility reader must be dropped once the output is dropped") + .unwrap(); + } + + #[tokio::test] + async fn test_schema_record_batch_yields_before_pending_message_stream() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "v", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let messages = futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ] as Vec>) + .chain(futures_util::stream::pending()); + let output = output_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + let OutputData::Stream(mut stream) = output.output.data else { + panic!("expected stream output"); + }; + + let batch = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(!output.metrics.is_ready()); + } + #[tokio::test] async fn test_invalid_terminal_metrics_after_record_batch_yields_batch_then_error() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 1753fdd5246..3a21e6daa4e 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,40 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::pin::Pin; - use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use futures_util::stream::Peekable; use futures_util::{Stream, StreamExt}; use snafu::{OptionExt, ResultExt}; use crate::Result; use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum FlightMessageKind { - Schema, - RecordBatch, - AffectedRows, - Metrics, -} - -impl From<&FlightMessage> for FlightMessageKind { - fn from(message: &FlightMessage) -> Self { - match message { - FlightMessage::Schema(_) => Self::Schema, - FlightMessage::RecordBatch(_) => Self::RecordBatch, - FlightMessage::AffectedRows { .. } => Self::AffectedRows, - FlightMessage::Metrics(_) => Self::Metrics, - } - } -} - pub(crate) struct FlightMessageReader { /// Remote Flight peer associated with this response stream. remote_addr: String, - messages: Peekable, + messages: S, } impl FlightMessageReader @@ -55,7 +33,7 @@ where pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { Self { remote_addr: remote_addr.into(), - messages: messages.peekable(), + messages, } } @@ -72,21 +50,6 @@ where pub(crate) async fn read_next(&mut self) -> Result> { self.messages.next().await.transpose() } - - pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { - match Pin::new(&mut self.messages).peek().await { - Some(Ok(message)) => Ok(Some(message.into())), - None => Ok(None), - Some(Err(_)) => match self.read_next().await { - // `peek` only borrows the error; consume it to preserve the source error. - Err(error) => Err(error), - Ok(_) => IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek".to_string(), - } - .fail(), - }, - } - } } pub(crate) fn decode_flight_data( diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 24876c73fcf..f9ebbb0d0bf 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -51,7 +51,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::region::codec::RegionCodec; use crate::{Client, metrics}; @@ -272,9 +272,7 @@ where "poll_flight_data_stream" )); - let mut stream_ended = false; - - while !stream_ended { + loop { let flight_message = match reader.read_next().await { Ok(Some(message)) => message, Ok(None) => break, @@ -287,59 +285,27 @@ where match flight_message { FlightMessage::RecordBatch(record_batch) => { - let result_to_yield = - RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - - // Metrics follow a batch so MergeScan can observe them before yielding it. - match reader.peek_next_message_kind().await { - Ok(Some(FlightMessageKind::Metrics)) => { - let metrics_message = match reader.read_next().await { - Ok(Some(FlightMessage::Metrics(metrics))) => metrics, - Ok(Some(_) | None) => { - yield IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek", - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(error) => { - yield Err(BoxedError::new(flight_stream_error( - &stream_addr, - error, - ))) - .context(ExternalSnafu); - break; - } - }; - let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); - metrics_ref.swap(metrics); - } - Ok(Some(FlightMessageKind::RecordBatch)) => {} - Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Ok(None) => stream_ended = true, - Err(error) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) - .context(ExternalSnafu); - break; - } - } - - yield Ok(result_to_yield); + // Deliver each batch immediately. In particular, do not + // wait for a possible following Metrics message; it is + // consumed on the next poll of this stream. + yield Ok(RecordBatch::from_df_record_batch( + schema_cloned.clone(), + record_batch, + )); } FlightMessage::Metrics(s) => { // Metrics may arrive before the next RecordBatch. - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); + match serde_json::from_str(&s) { + Ok(metrics) => { + metrics_ref.swap(Some(Arc::new(metrics))); + } + Err(error) => { + common_telemetry::warn!( + "Failed to decode region Flight metrics: {}", + error + ); + } + } continue; } _ => { @@ -1211,7 +1177,34 @@ mod tests { } #[tokio::test] - async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + async fn test_record_batch_is_yielded_without_waiting_for_next_message() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let messages = stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]) + .chain(stream::pending::>()); + let mut recordbatches = + recordbatches_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + + let batch = tokio::time::timeout(Duration::from_secs(1), recordbatches.next()) + .await + .expect("the first batch must not wait for lookahead") + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + + #[tokio::test] + async fn test_malformed_region_metrics_are_non_fatal() { let schema = test_schema(); let batch = RecordBatch::new( schema.clone(), @@ -1222,23 +1215,19 @@ mod tests { "test-peer".to_string(), stream::iter(vec![ Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics(test_metrics_json())), ]), ) .await .unwrap(); - let batch = recordbatches.next().await.unwrap().unwrap(); - assert_eq!(batch.num_rows(), 1); - - let metrics = recordbatches.metrics().unwrap(); - assert_eq!(metrics.elapsed_compute, 7); + assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1); assert!(recordbatches.next().await.is_none()); } #[tokio::test] - async fn test_record_batch_stream_preserves_peeked_record_batch() { + async fn test_record_batch_stream_preserves_following_record_batch() { let schema = test_schema(); let first_batch = RecordBatch::new( schema.clone(), diff --git a/src/cmd/src/cli.rs b/src/cmd/src/cli.rs index 09ec9646d8b..3739d4cc4fc 100644 --- a/src/cmd/src/cli.rs +++ b/src/cmd/src/cli.rs @@ -108,6 +108,7 @@ mod tests { use clap::Parser; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::ChannelManager; use common_telemetry::logging::LoggingOptions; use rand::Rng; @@ -174,7 +175,11 @@ mod tests { let mut instance = standalone.build(standalone_opts).await?; instance.start().await?; - let client = Client::with_urls([rpc_addr.as_str()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [rpc_addr.as_str()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); database .sql(r#"CREATE DATABASE "cli.export.create_table";"#) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 6cae195f3e6..b000f7e1f02 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -467,8 +467,6 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; - // Source Flight streams and sink unary responses share pooled connections. - channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index e932933ff84..3f1e1847eb8 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -91,7 +91,8 @@ impl HandlerMutable { pub enum FrontendClient { Distributed { meta_client: Arc, - chnl_mgr: ChannelManager, + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, query: QueryOptions, batch_opts: BatchingModeOptions, }, @@ -136,17 +137,19 @@ impl FrontendClient { batch_opts: BatchingModeOptions, ) -> Result { common_telemetry::info!("Frontend client build without auth"); + let cfg = ChannelConfig::new() + .connect_timeout(batch_opts.grpc_conn_timeout) + .timeout(Some(batch_opts.query_timeout)); + let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) + .context(InvalidClientConfigSnafu)?; + // Keep separate TLS config handles and pools for query and control traffic. + let query_channel_manager = ChannelManager::with_config(cfg.clone(), tls_config.clone()); + let control_channel_manager = ChannelManager::with_config(cfg, tls_config); + Ok(Self::Distributed { meta_client, - chnl_mgr: { - let cfg = ChannelConfig::new() - .connect_timeout(batch_opts.grpc_conn_timeout) - .timeout(Some(batch_opts.query_timeout)); - - let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) - .context(InvalidClientConfigSnafu)?; - ChannelManager::with_config(cfg, tls_config) - }, + query_channel_manager, + control_channel_manager, query, batch_opts, }) @@ -227,7 +230,8 @@ impl FrontendClient { frontends: &[Peer], ) -> Result, Error> { let Self::Distributed { - chnl_mgr, + query_channel_manager, + control_channel_manager, batch_opts, .. } = self @@ -240,10 +244,15 @@ impl FrontendClient { .iter() .map(|peer| { let addr = peer.addr.clone(); - let chnl_mgr = chnl_mgr.clone(); + let query_channel_manager = query_channel_manager.clone(); + let control_channel_manager = control_channel_manager.clone(); async move { - let client = Client::with_manager_and_urls(chnl_mgr, vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager, + control_channel_manager, + vec![addr.clone()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); match tokio::time::timeout(probe_timeout, database.sql("SELECT 1")).await { @@ -282,7 +291,8 @@ impl FrontendClient { ) -> Result { let Self::Distributed { meta_client: _, - chnl_mgr, + query_channel_manager, + control_channel_manager, query: _, batch_opts, } = self @@ -302,7 +312,11 @@ impl FrontendClient { for peer in frontends { let addr = peer.addr.clone(); - let client = Client::with_manager_and_urls(chnl_mgr.clone(), vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + vec![addr.clone()], + ); let database = Database::new(catalog, schema, client); let db = DatabaseWithPeer::new(database, peer); match db.try_select_one().await { @@ -1125,7 +1139,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, @@ -1162,7 +1180,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index fb64dc9c20e..2b36d46fbd3 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -826,6 +826,12 @@ impl BatchingTask { let res = res?; let (affected_rows, _) = res.output.extract_rows_and_cost(); + if matches!(&res.output.data, common_query::OutputData::AffectedRows(_)) { + res.metrics.wait_ready().await; + } + if let Some(error) = res.metrics.completion_error() { + warn!("Flow {flow_id} completed with terminal metrics error: {error}"); + } debug!( "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}", elapsed, diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 2499cd19d1d..cce0d79f7fc 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -620,7 +620,11 @@ backoff_multiplier = 3.0 expected: std::result::Result<&str, (StatusCode, &str)>, ) { let addr = frontend.server_handlers().addr(GRPC_SERVER).unwrap(); - let client = Client::with_urls([addr.to_string()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.to_string()], + ); let client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); let response = client.sql("SELECT 1").await; diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 8cb33a1e28a..464e0e3d24a 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -474,6 +474,7 @@ mod tests { use async_trait::async_trait; use auth::{UserProviderRef, static_user_provider_from_option}; use client::{Client, Database}; + use common_grpc::channel_manager::ChannelManager; use meta_client::client::MetaClientBuilder; use servers::grpc::GRPC_SERVER; use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream}; @@ -660,12 +661,20 @@ mod tests { let public_database = Database::new( "greptime", "public", - Client::with_urls([public_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [public_addr.to_string()], + ), ); let internal_database = Database::new( "greptime", "public", - Client::with_urls([internal_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [internal_addr.to_string()], + ), ); let internal_result = internal_database.sql("SELECT 1").await; diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index c30519da984..49964de24b7 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -336,6 +336,9 @@ pub mod mock { }), ) .unwrap(); + // The mock transport is a single in-memory duplex connection, so both lanes + // must resolve through this one pre-wired manager. + #[allow(deprecated)] Client::with_manager_and_urls(channel_manager, vec![datanode.addr.clone()]) } } diff --git a/src/meta-srv/src/utils/database.rs b/src/meta-srv/src/utils/database.rs index 976be8a85b6..796d63009d8 100644 --- a/src/meta-srv/src/utils/database.rs +++ b/src/meta-srv/src/utils/database.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use client::error::{ExternalSnafu, Result as ClientResult}; use client::{Client, Database, Output}; use common_error::ext::BoxedError; +use common_grpc::channel_manager::ChannelManager; use common_meta::peer::PeerDiscoveryRef; use common_telemetry::{debug, warn}; use snafu::{ResultExt, ensure}; @@ -43,6 +44,9 @@ impl<'a> DatabaseContext<'a> { } /// A cached frontend database operator used by metasrv. +/// +/// Its cached clients use independent query and control channel-manager pools; +/// the legacy single-manager client constructors intentionally remain shared. pub struct DatabaseOperator { peer_discovery: PeerDiscoveryRef, client: RwLock>, @@ -110,7 +114,11 @@ impl DatabaseOperator { debug!("Available frontend addresses: {:?}", urls); - Ok(Client::with_urls(urls)) + Ok(Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + urls, + )) } async fn maybe_init_client(&self) -> ClientResult { @@ -147,3 +155,62 @@ fn should_reset_client(result: &client::error::Result) -> bool { .map(|err| err.is_connection_error()) .unwrap_or(false) } + +#[cfg(test)] +mod tests { + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::cluster::{FrontendStatus, NodeInfo, NodeStatus}; + use common_meta::peer::{Peer, PeerDiscovery}; + + use super::*; + + struct TestPeerDiscovery; + + #[async_trait::async_trait] + impl PeerDiscovery for TestPeerDiscovery { + async fn active_frontends(&self) -> common_meta::error::Result> { + Ok(vec![NodeInfo { + peer: Peer::new(1, "127.0.0.1:3001".to_string()), + last_activity_ts: 0, + status: NodeStatus::Frontend(FrontendStatus::default()), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: Default::default(), + }]) + } + + async fn active_datanodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + + async fn active_flownodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + } + + #[tokio::test] + async fn test_build_client_uses_isolated_reused_channel_pools() { + let operator = DatabaseOperator::new(Arc::new(TestPeerDiscovery)); + let client = operator.build_client().await.unwrap(); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!((1, 0), client.channel_pool_sizes()); + + client.find_channel().unwrap(); + client.find_channel().unwrap(); + assert_eq!((1, 1), client.channel_pool_sizes()); + } +} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 28644b8ccc9..0bf466ac8ab 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,7 +44,6 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; -use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -81,29 +80,6 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; -// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared -// HTTP/2 connection are pending, trading bounded memory for request liveness. -async fn forward_record_batches( - mut stream: SendableRecordBatchStream, - batch_tx: tokio::sync::mpsc::UnboundedSender>, -) -> Result<()> { - while let Some(batch) = stream.next().await { - match batch.context(CreateRecordBatchSnafu) { - Ok(batch) => { - if batch_tx.send(Ok(batch)).is_err() { - break; - } - tokio::task::yield_now().await; - } - Err(error) => { - let _ = batch_tx.send(Err(error)); - break; - } - } - } - Ok(()) -} - fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -229,7 +205,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let stream = match data { + let mut stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -240,31 +216,32 @@ impl DatafusionQueryEngine { match dml.op { WriteOp::Insert(_) => { - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_insert_record_batches( - batch_rx, - &table_name, - table.schema(), - query_ctx.clone(), - ); - let ((), (rows, cost)) = try_join(producer, consumer).await?; - affected_rows += rows; - insert_cost += cost; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(&table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } } WriteOp::Delete => { - // Keep DELETE on the same producer/consumer schedule as INSERT so the source - // stream can continue draining while mutation RPCs are pending. - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_delete_record_batches( - batch_rx, - &table_name, - &table, - query_ctx.clone(), - ); - let (rows, ()) = try_join(consumer, producer).await?; - affected_rows += rows; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(&table_name, &table, column_vectors, query_ctx.clone()) + .await?; + } } _ => unreachable!("guarded by the 'ensure!' at the beginning"), } @@ -274,53 +251,6 @@ impl DatafusionQueryEngine { )) } - async fn consume_insert_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table_schema: Arc, - query_ctx: QueryContextRef, - ) -> Result<(usize, usize)> { - let mut affected_rows = 0; - let mut insert_cost = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table_schema.clone()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - // We ignore the insert op. - let output = self - .insert(table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - Ok((affected_rows, insert_cost)) - } - - async fn consume_delete_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table: &TableRef, - query_ctx: QueryContextRef, - ) -> Result { - let mut affected_rows = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - affected_rows += self - .delete(table_name, table, column_vectors, query_ctx.clone()) - .await?; - } - Ok(affected_rows) - } - #[tracing::instrument(skip_all)] async fn delete( &self, diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index e40c707fe98..3ec3cf14724 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -658,11 +658,15 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ( addr.to_string(), - Client::with_manager_and_urls(channel_manager, [addr]), + // The mock connector pool lives on this single manager, so both lanes + // intentionally share it (same semantics as the legacy constructor). + Client::with_query_and_control_managers(channel_manager.clone(), channel_manager, [addr]), ) } +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { use api::v1::flow::FlowRequest; use api::v1::region::{ diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 2486dddee05..97df1114eaf 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod test { use std::collections::HashMap; use std::net::SocketAddr; @@ -671,7 +673,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); - assert!(result.metrics.is_ready()); + result.metrics.wait_ready().await; assert!(result.region_watermark_map().is_none()); let err = client @@ -697,6 +699,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); + result.metrics.wait_ready().await; assert_eq!( result.region_watermark_map(), Some(std::collections::HashMap::from([previous_watermark])) diff --git a/tests-integration/src/grpc/network.rs b/tests-integration/src/grpc/network.rs index 427317fb05d..97e4c56494a 100644 --- a/tests-integration/src/grpc/network.rs +++ b/tests-integration/src/grpc/network.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { use std::io; use std::pin::Pin; diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 07c370f026d..5f99f0e2b60 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1000,8 +1000,14 @@ pub async fn setup_authenticated_grpc_database( .unwrap(); let grpc_server = setup_grpc_server_for_frontend_instance(instance, Some(user_provider)).await; let grpc_addr = grpc_server.bind_addr().unwrap().to_string(); - let mut database = - Database::new_with_dbname("greptime-public", Client::with_urls(vec![grpc_addr])); + let mut database = Database::new_with_dbname( + "greptime-public", + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_addr], + ), + ); database.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic { username: username.to_string(), password: password.to_string(), diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 5c718bea49f..b9cefce4871 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Retain coverage of the legacy constructors, including TLS configuration. +#![allow(deprecated)] + use std::sync::Arc; use api::v1::alter_table_expr::Kind; diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index 163e827e53a..a595c877940 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -64,7 +64,11 @@ impl MultiProtocolClient { let grpc_client = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls(vec![grpc_server_addr]), + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_server_addr], + ), ); let pg_client = create_postgres_client(pg_server_addr).await; let mysql_client = create_mysql_client(mysql_server_addr).await;