flow/batching_mode/
frontend_client.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Frontend client to run flow as batching task which is time-window-aware normal query triggered every tick set by user
16
17use std::collections::HashMap;
18use std::sync::{Arc, Mutex, RwLock, Weak};
19
20use api::v1::greptime_request::Request;
21use api::v1::query_request::Query;
22use api::v1::{CreateTableExpr, QueryRequest, RowInsertRequests};
23use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics};
24use common_error::ext::BoxedError;
25use common_grpc::channel_manager::{ChannelConfig, ChannelManager, load_client_tls_config};
26use common_meta::peer::{Peer, PeerDiscovery};
27use common_query::{Output, OutputData};
28use common_telemetry::warn;
29use futures::stream::{FuturesUnordered, StreamExt};
30use meta_client::client::MetaClient;
31use query::datafusion::QUERY_PARALLELISM_HINT;
32use query::metrics::terminal_recordbatch_metrics_from_plan;
33use query::options::{FlowQueryExtensions, QueryOptions};
34use rand::rng;
35use rand::seq::SliceRandom;
36use servers::query_handler::grpc::GrpcQueryHandler;
37use session::context::{Channel, QueryContextBuilder, QueryContextRef};
38use session::hints::READ_PREFERENCE_HINT;
39use snafu::{OptionExt, ResultExt};
40use tokio::sync::SetOnce;
41
42use crate::Error;
43use crate::batching_mode::BatchingModeOptions;
44use crate::error::{
45    CreateSinkTableSnafu, ExternalSnafu, InvalidClientConfigSnafu, InvalidRequestSnafu,
46    NoAvailableFrontendSnafu, UnexpectedSnafu,
47};
48
49/// Adapter trait for [`GrpcQueryHandler`] that boxes the underlying error into [`BoxedError`].
50///
51/// This is mainly used by flownode to invoke a frontend instance in standalone mode.
52#[async_trait::async_trait]
53pub trait GrpcQueryHandlerWithBoxedError: Send + Sync + 'static {
54    async fn do_query(
55        &self,
56        query: Request,
57        ctx: QueryContextRef,
58    ) -> std::result::Result<Output, BoxedError>;
59}
60
61/// auto impl
62#[async_trait::async_trait]
63impl<T: GrpcQueryHandler + Send + Sync + 'static> GrpcQueryHandlerWithBoxedError for T {
64    async fn do_query(
65        &self,
66        query: Request,
67        ctx: QueryContextRef,
68    ) -> std::result::Result<Output, BoxedError> {
69        self.do_query(query, ctx).await.map_err(BoxedError::new)
70    }
71}
72
73#[derive(Debug, Clone)]
74pub struct HandlerMutable {
75    handler: Arc<Mutex<Option<Weak<dyn GrpcQueryHandlerWithBoxedError>>>>,
76    is_initialized: Arc<SetOnce<()>>,
77}
78
79impl HandlerMutable {
80    pub async fn set_handler(&self, handler: Weak<dyn GrpcQueryHandlerWithBoxedError>) {
81        *self.handler.lock().unwrap() = Some(handler);
82        // Ignore the error, as we allow the handler to be set multiple times.
83        let _ = self.is_initialized.set(());
84    }
85}
86
87/// A simple frontend client able to execute sql using grpc protocol
88///
89/// This is for computation-heavy query which need to offload computation to frontend, lifting the load from flownode
90#[derive(Debug, Clone)]
91pub enum FrontendClient {
92    Distributed {
93        meta_client: Arc<MetaClient>,
94        query_channel_manager: ChannelManager,
95        control_channel_manager: ChannelManager,
96        query: QueryOptions,
97        batch_opts: BatchingModeOptions,
98    },
99    Standalone {
100        /// for the sake of simplicity still use grpc even in standalone mode
101        /// notice the client here should all be lazy, so that can wait after frontend is booted then make conn
102        database_client: HandlerMutable,
103        query: QueryOptions,
104    },
105}
106
107impl FrontendClient {
108    /// Create a new empty frontend client, with a `HandlerMutable` to set the grpc handler later
109    pub fn from_empty_grpc_handler(query: QueryOptions) -> (Self, HandlerMutable) {
110        let is_initialized = Arc::new(SetOnce::new());
111        let handler = HandlerMutable {
112            handler: Arc::new(Mutex::new(None)),
113            is_initialized,
114        };
115        (
116            Self::Standalone {
117                database_client: handler.clone(),
118                query,
119            },
120            handler,
121        )
122    }
123
124    /// Waits until the frontend client is initialized.
125    pub async fn wait_initialized(&self) {
126        if let FrontendClient::Standalone {
127            database_client, ..
128        } = self
129        {
130            database_client.is_initialized.wait().await;
131        }
132    }
133
134    pub fn from_meta_client(
135        meta_client: Arc<MetaClient>,
136        query: QueryOptions,
137        batch_opts: BatchingModeOptions,
138    ) -> Result<Self, Error> {
139        common_telemetry::info!("Frontend client build without auth");
140        let cfg = ChannelConfig::new()
141            .connect_timeout(batch_opts.grpc_conn_timeout)
142            .timeout(Some(batch_opts.query_timeout));
143        let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone())
144            .context(InvalidClientConfigSnafu)?;
145        // Keep separate TLS config handles and pools for query and control traffic.
146        let query_channel_manager = ChannelManager::with_config(cfg.clone(), tls_config.clone());
147        let control_channel_manager = ChannelManager::with_config(cfg, tls_config);
148
149        Ok(Self::Distributed {
150            meta_client,
151            query_channel_manager,
152            control_channel_manager,
153            query,
154            batch_opts,
155        })
156    }
157
158    pub fn from_grpc_handler(
159        grpc_handler: Weak<dyn GrpcQueryHandlerWithBoxedError>,
160        query: QueryOptions,
161    ) -> Self {
162        let is_initialized = Arc::new(SetOnce::new_with(Some(())));
163        let handler = HandlerMutable {
164            handler: Arc::new(Mutex::new(Some(grpc_handler))),
165            is_initialized: is_initialized.clone(),
166        };
167
168        Self::Standalone {
169            database_client: handler,
170            query,
171        }
172    }
173}
174
175#[derive(Debug, Clone)]
176pub struct DatabaseWithPeer {
177    pub database: Database,
178    pub peer: Peer,
179}
180
181impl DatabaseWithPeer {
182    fn new(database: Database, peer: Peer) -> Self {
183        Self { database, peer }
184    }
185
186    /// Try sending a "SELECT 1" to the database
187    async fn try_select_one(&self) -> Result<(), Error> {
188        // notice here use `sql` for `SELECT 1` return 1 row
189        let output = self
190            .database
191            .sql("SELECT 1")
192            .await
193            .with_context(|_| InvalidRequestSnafu {
194                context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer),
195            })?;
196
197        if let OutputData::Stream(stream) = output.data {
198            common_recordbatch::util::collect(stream)
199                .await
200                .map_err(BoxedError::new)
201                .context(ExternalSnafu)?;
202        }
203
204        Ok(())
205    }
206}
207
208impl FrontendClient {
209    /// scan for available frontend from metadata
210    pub(crate) async fn scan_for_frontend(&self) -> Result<Vec<Peer>, Error> {
211        let Self::Distributed { meta_client, .. } = self else {
212            return Ok(vec![]);
213        };
214
215        meta_client
216            .active_frontends()
217            .await
218            .map(|nodes| nodes.into_iter().map(|node| node.peer).collect())
219            .map_err(BoxedError::new)
220            .context(ExternalSnafu)
221    }
222
223    /// Probes all discovered frontends without auth.
224    ///
225    /// Returns non-auth failures to allow callers to retry transient connectivity
226    /// errors. Authentication failures are returned immediately because they mean
227    /// a frontend advertised an auth-protected endpoint to flownodes.
228    pub(crate) async fn check_all_frontends_without_auth(
229        &self,
230        frontends: &[Peer],
231    ) -> Result<Vec<String>, Error> {
232        let Self::Distributed {
233            query_channel_manager,
234            control_channel_manager,
235            batch_opts,
236            ..
237        } = self
238        else {
239            return Ok(vec![]);
240        };
241
242        let probe_timeout = batch_opts.grpc_conn_timeout;
243        let mut probes = frontends
244            .iter()
245            .map(|peer| {
246                let addr = peer.addr.clone();
247                let query_channel_manager = query_channel_manager.clone();
248                let control_channel_manager = control_channel_manager.clone();
249
250                async move {
251                    let client = Client::with_query_and_control_managers(
252                        query_channel_manager,
253                        control_channel_manager,
254                        vec![addr.clone()],
255                    );
256                    let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
257
258                    match tokio::time::timeout(probe_timeout, database.sql("SELECT 1")).await {
259                        Ok(Ok(_)) => Ok(None),
260                        Ok(Err(err)) if err.tonic_code() == Some(tonic::Code::Unauthenticated) => {
261                            Err(err).context(InvalidRequestSnafu {
262                                context: format!(
263                                    "Frontend {addr} rejected unauthenticated flownode probe; ensure frontend internal_grpc is advertised to metasrv"
264                                ),
265                            })
266                        }
267                        Ok(Err(err)) => Ok(Some(format!("{addr}: {err}"))),
268                        Err(_) => Ok(Some(format!(
269                            "{addr}: health check timed out after {probe_timeout:?}"
270                        ))),
271                    }
272                }
273            })
274            .collect::<FuturesUnordered<_>>();
275
276        let mut failures = Vec::new();
277        while let Some(probe_result) = probes.next().await {
278            if let Some(failure) = probe_result? {
279                failures.push(failure);
280            }
281        }
282
283        Ok(failures)
284    }
285
286    /// Get a frontend discovered by metasrv and verified with a query probe.
287    async fn get_random_active_frontend(
288        &self,
289        catalog: &str,
290        schema: &str,
291    ) -> Result<DatabaseWithPeer, Error> {
292        let Self::Distributed {
293            meta_client: _,
294            query_channel_manager,
295            control_channel_manager,
296            query: _,
297            batch_opts,
298            ..
299        } = self
300        else {
301            return UnexpectedSnafu {
302                reason: "Expect distributed mode",
303            }
304            .fail();
305        };
306
307        let mut interval = tokio::time::interval(batch_opts.grpc_conn_timeout);
308        interval.tick().await;
309        for retry in 0..batch_opts.experimental_grpc_max_retries {
310            let mut frontends = self.scan_for_frontend().await?;
311            // shuffle the frontends to avoid always pick the same one
312            frontends.shuffle(&mut rng());
313
314            for peer in frontends {
315                let addr = peer.addr.clone();
316                let client = Client::with_query_and_control_managers(
317                    query_channel_manager.clone(),
318                    control_channel_manager.clone(),
319                    vec![addr.clone()],
320                );
321                let database = Database::new(catalog, schema, client);
322                let db = DatabaseWithPeer::new(database, peer);
323                match db.try_select_one().await {
324                    Ok(_) => return Ok(db),
325                    Err(e) => {
326                        warn!(
327                            "Failed to connect to frontend {} on retry={}: \n{e:?}",
328                            addr, retry
329                        );
330                    }
331                }
332            }
333            // no available frontend
334            // sleep and retry
335            interval.tick().await;
336        }
337
338        NoAvailableFrontendSnafu {
339            timeout: batch_opts.grpc_conn_timeout,
340            context: "No available frontend found that is able to process query",
341        }
342        .fail()
343    }
344
345    pub async fn create(
346        &self,
347        create: CreateTableExpr,
348        catalog: &str,
349        schema: &str,
350    ) -> Result<u32, Error> {
351        self.handle(
352            Request::Ddl(api::v1::DdlRequest {
353                expr: Some(api::v1::ddl_request::Expr::CreateTable(create.clone())),
354            }),
355            catalog,
356            schema,
357            &mut None,
358        )
359        .await
360        .map_err(BoxedError::new)
361        .with_context(|_| CreateSinkTableSnafu {
362            create: create.clone(),
363        })
364    }
365
366    /// Execute a SQL statement on the frontend.
367    pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result<Output, Error> {
368        match self {
369            FrontendClient::Distributed { .. } => {
370                let db = self.get_random_active_frontend(catalog, schema).await?;
371                db.database
372                    .sql(sql)
373                    .await
374                    .map_err(BoxedError::new)
375                    .context(ExternalSnafu)
376            }
377            FrontendClient::Standalone {
378                database_client, ..
379            } => {
380                let ctx = QueryContextBuilder::default()
381                    .current_catalog(catalog.to_string())
382                    .current_schema(schema.to_string())
383                    .build();
384                let ctx = Arc::new(ctx);
385                {
386                    let database_client = {
387                        database_client
388                            .handler
389                            .lock()
390                            .map_err(|e| {
391                                UnexpectedSnafu {
392                                    reason: format!("Failed to lock database client: {e}"),
393                                }
394                                .build()
395                            })?
396                            .as_ref()
397                            .context(UnexpectedSnafu {
398                                reason: "Standalone's frontend instance is not set",
399                            })?
400                            .upgrade()
401                            .context(UnexpectedSnafu {
402                                reason: "Failed to upgrade database client",
403                            })?
404                    };
405                    let req = Request::Query(QueryRequest {
406                        query: Some(Query::Sql(sql.to_string())),
407                    });
408                    database_client
409                        .do_query(req, ctx)
410                        .await
411                        .map_err(BoxedError::new)
412                        .context(ExternalSnafu)
413                }
414            }
415        }
416    }
417
418    /// Execute row inserts on the frontend.
419    pub async fn row_inserts(
420        &self,
421        catalog: &str,
422        schema: &str,
423        requests: RowInsertRequests,
424        hints: &[(&str, &str)],
425    ) -> Result<u32, Error> {
426        match self {
427            FrontendClient::Distributed { .. } => {
428                let db = self.get_random_active_frontend(catalog, schema).await?;
429                db.database
430                    .row_inserts_with_hints(requests, hints)
431                    .await
432                    .with_context(|_| InvalidRequestSnafu {
433                        context: format!("Failed to handle row inserts at {:?}", db.peer),
434                    })
435            }
436            FrontendClient::Standalone {
437                database_client, ..
438            } => {
439                let extensions = HashMap::from_iter(
440                    hints
441                        .iter()
442                        .map(|(key, value)| ((*key).to_string(), (*value).to_string())),
443                );
444                let ctx = QueryContextBuilder::default()
445                    .current_catalog(catalog.to_string())
446                    .current_schema(schema.to_string())
447                    .extensions(extensions)
448                    .build();
449                let ctx = Arc::new(ctx);
450                {
451                    let database_client = {
452                        database_client
453                            .handler
454                            .lock()
455                            .unwrap()
456                            .as_ref()
457                            .context(UnexpectedSnafu {
458                                reason: "Standalone's frontend instance is not set",
459                            })?
460                            .upgrade()
461                            .context(UnexpectedSnafu {
462                                reason: "Failed to upgrade database client",
463                            })?
464                    };
465                    let resp: common_query::Output = database_client
466                        .do_query(Request::RowInserts(requests), ctx)
467                        .await
468                        .map_err(BoxedError::new)
469                        .context(ExternalSnafu)?;
470                    match resp.data {
471                        OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| {
472                            UnexpectedSnafu {
473                                reason: format!("Failed to convert rows to u32: {}", rows),
474                            }
475                            .build()
476                        })?),
477                        _ => UnexpectedSnafu {
478                            reason: "Unexpected output data",
479                        }
480                        .fail(),
481                    }
482                }
483            }
484        }
485    }
486
487    /// Execute a flow query and return terminal metrics. `snapshot_seqs` are
488    /// optional read upper bounds used only by snapshot-fenced repair chunks.
489    pub(crate) async fn query_with_terminal_metrics(
490        &self,
491        catalog: &str,
492        schema: &str,
493        request: QueryRequest,
494        extensions: &[(&str, &str)],
495        snapshot_seqs: &HashMap<u64, u64>,
496        peer_desc: &mut Option<PeerDesc>,
497    ) -> Result<OutputWithMetrics, Error> {
498        let flow_extensions = build_flow_extensions(extensions)?;
499        match self {
500            FrontendClient::Distributed {
501                query, batch_opts, ..
502            } => {
503                let query_parallelism = query.parallelism.to_string();
504                let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
505                if query.parallelism > 0 {
506                    hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
507                }
508                let db = self.get_random_active_frontend(catalog, schema).await?;
509                *peer_desc = Some(PeerDesc::Dist {
510                    peer: db.peer.clone(),
511                });
512                db.database
513                    .flight_request()
514                    .with_hints(&hints)
515                    .with_flow_extensions(extensions)
516                    .with_snapshot_seqs(snapshot_seqs)
517                    .with_timeout(batch_opts.experimental_flight_do_get_timeout)
518                    .query_with_terminal_metrics(request)
519                    .await
520                    .map_err(BoxedError::new)
521                    .context(ExternalSnafu)
522            }
523            FrontendClient::Standalone {
524                database_client,
525                query,
526                ..
527            } => {
528                *peer_desc = Some(PeerDesc::Standalone);
529                let mut extensions_map = HashMap::new();
530                if query.parallelism > 0 {
531                    extensions_map.insert(
532                        QUERY_PARALLELISM_HINT.to_string(),
533                        query.parallelism.to_string(),
534                    );
535                }
536                for (key, value) in extensions {
537                    extensions_map.insert((*key).to_string(), (*value).to_string());
538                }
539                let ctx = QueryContextBuilder::default()
540                    .current_catalog(catalog.to_string())
541                    .current_schema(schema.to_string())
542                    .extensions(extensions_map)
543                    .channel(Channel::Internal)
544                    .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs.clone())))
545                    .build();
546                let ctx = Arc::new(ctx);
547                let database_client = {
548                    database_client
549                        .handler
550                        .lock()
551                        .map_err(|e| {
552                            UnexpectedSnafu {
553                                reason: format!("Failed to lock database client: {e}"),
554                            }
555                            .build()
556                        })?
557                        .as_ref()
558                        .context(UnexpectedSnafu {
559                            reason: "Standalone's frontend instance is not set",
560                        })?
561                        .upgrade()
562                        .context(UnexpectedSnafu {
563                            reason: "Failed to upgrade database client",
564                        })?
565                };
566                database_client
567                    .do_query(Request::Query(request), ctx.clone())
568                    .await
569                    .map(|output| {
570                        wrap_standalone_output_with_terminal_metrics(output, &flow_extensions)
571                    })
572                    .map_err(BoxedError::new)
573                    .context(ExternalSnafu)
574            }
575        }
576    }
577
578    /// Handle an insert request with one attempt.
579    ///
580    /// Unlike [`Self::handle`], this does not use the batching retry policy. It
581    /// is intended for stateless streaming sinks, where retrying an insert can
582    /// duplicate rows.
583    pub(crate) async fn handle_insert_once(
584        &self,
585        req: api::v1::greptime_request::Request,
586        catalog: &str,
587        schema: &str,
588        peer_desc: &mut Option<PeerDesc>,
589    ) -> Result<u32, Error> {
590        match self {
591            FrontendClient::Distributed { .. } => {
592                let db = self.get_random_active_frontend(catalog, schema).await?;
593
594                *peer_desc = Some(PeerDesc::Dist {
595                    peer: db.peer.clone(),
596                });
597
598                db.database
599                    .handle(req.clone())
600                    .await
601                    .with_context(|_| InvalidRequestSnafu {
602                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
603                    })
604            }
605            FrontendClient::Standalone {
606                database_client,
607                query,
608                ..
609            } => {
610                let ctx = QueryContextBuilder::default()
611                    .current_catalog(catalog.to_string())
612                    .current_schema(schema.to_string())
613                    .channel(Channel::Internal)
614                    .extensions(if query.parallelism > 0 {
615                        HashMap::from([(
616                            QUERY_PARALLELISM_HINT.to_string(),
617                            query.parallelism.to_string(),
618                        )])
619                    } else {
620                        HashMap::new()
621                    })
622                    .build();
623                let ctx = Arc::new(ctx);
624                let database_client = {
625                    database_client
626                        .handler
627                        .lock()
628                        .unwrap()
629                        .as_ref()
630                        .context(UnexpectedSnafu {
631                            reason: "Standalone's frontend instance is not set",
632                        })?
633                        .upgrade()
634                        .context(UnexpectedSnafu {
635                            reason: "Failed to upgrade database client",
636                        })?
637                };
638                let resp: common_query::Output = database_client
639                    .do_query(req, ctx)
640                    .await
641                    .map_err(BoxedError::new)
642                    .context(ExternalSnafu)?;
643                match resp.data {
644                    common_query::OutputData::AffectedRows(rows) => rows.try_into().map_err(|_| {
645                        UnexpectedSnafu {
646                            reason: format!("Failed to convert rows to u32: {}", rows),
647                        }
648                        .build()
649                    }),
650                    _ => UnexpectedSnafu {
651                        reason: "Unexpected output data",
652                    }
653                    .fail(),
654                }
655            }
656        }
657    }
658
659    /// Handle a request to frontend
660    pub(crate) async fn handle(
661        &self,
662        req: api::v1::greptime_request::Request,
663        catalog: &str,
664        schema: &str,
665        peer_desc: &mut Option<PeerDesc>,
666    ) -> Result<u32, Error> {
667        match self {
668            FrontendClient::Distributed {
669                query, batch_opts, ..
670            } => {
671                let db = self.get_random_active_frontend(catalog, schema).await?;
672
673                *peer_desc = Some(PeerDesc::Dist {
674                    peer: db.peer.clone(),
675                });
676
677                let query_parallelism = query.parallelism.to_string();
678                let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
679                if query.parallelism > 0 {
680                    hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
681                }
682                db.database
683                    .handle_with_retry(
684                        req.clone(),
685                        batch_opts.experimental_grpc_max_retries,
686                        &hints,
687                    )
688                    .await
689                    .with_context(|_| InvalidRequestSnafu {
690                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
691                    })
692            }
693            FrontendClient::Standalone {
694                database_client,
695                query,
696                ..
697            } => {
698                let ctx = QueryContextBuilder::default()
699                    .current_catalog(catalog.to_string())
700                    .current_schema(schema.to_string())
701                    .extensions(if query.parallelism > 0 {
702                        HashMap::from([(
703                            QUERY_PARALLELISM_HINT.to_string(),
704                            query.parallelism.to_string(),
705                        )])
706                    } else {
707                        HashMap::new()
708                    })
709                    .build();
710                let ctx = Arc::new(ctx);
711                {
712                    let database_client = {
713                        database_client
714                            .handler
715                            .lock()
716                            .unwrap()
717                            .as_ref()
718                            .context(UnexpectedSnafu {
719                                reason: "Standalone's frontend instance is not set",
720                            })?
721                            .upgrade()
722                            .context(UnexpectedSnafu {
723                                reason: "Failed to upgrade database client",
724                            })?
725                    };
726                    let resp: common_query::Output = database_client
727                        .do_query(req, ctx)
728                        .await
729                        .map_err(BoxedError::new)
730                        .context(ExternalSnafu)?;
731                    match resp.data {
732                        OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| {
733                            UnexpectedSnafu {
734                                reason: format!("Failed to convert rows to u32: {}", rows),
735                            }
736                            .build()
737                        })?),
738                        _ => UnexpectedSnafu {
739                            reason: "Unexpected output data",
740                        }
741                        .fail(),
742                    }
743                }
744            }
745        }
746    }
747}
748
749fn build_flow_extensions(extensions: &[(&str, &str)]) -> Result<FlowQueryExtensions, Error> {
750    let flow_extensions = HashMap::from_iter(
751        extensions
752            .iter()
753            .map(|(key, value)| ((*key).to_string(), (*value).to_string())),
754    );
755    FlowQueryExtensions::parse_flow_extensions(&flow_extensions)
756        .map_err(BoxedError::new)
757        .context(ExternalSnafu)
758        .map(|extensions| extensions.unwrap_or_default())
759}
760
761fn wrap_standalone_output_with_terminal_metrics(
762    output: Output,
763    flow_extensions: &FlowQueryExtensions,
764) -> OutputWithMetrics {
765    let should_collect_region_watermark = flow_extensions.should_collect_region_watermark();
766    let terminal_metrics =
767        if should_collect_region_watermark && !matches!(&output.data, OutputData::Stream(_)) {
768            output
769                .meta
770                .plan
771                .clone()
772                .and_then(terminal_recordbatch_metrics_from_plan)
773        } else {
774            None
775        };
776    let result = OutputWithMetrics::from_output(output);
777    if let Some(metrics) = terminal_metrics {
778        result.metrics.update(Some(metrics));
779    }
780    result
781}
782
783/// Describe a peer of frontend
784#[derive(Debug, Default, Clone)]
785pub enum PeerDesc {
786    /// The query failed before a frontend peer was selected.
787    #[default]
788    Unknown,
789    /// Distributed mode's frontend peer address
790    Dist {
791        /// frontend peer address
792        peer: Peer,
793    },
794    /// Standalone mode
795    Standalone,
796}
797
798impl std::fmt::Display for PeerDesc {
799    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
800        match self {
801            PeerDesc::Unknown => write!(f, "unknown"),
802            PeerDesc::Dist { peer } => write!(f, "{}", peer.addr),
803            PeerDesc::Standalone => write!(f, "standalone"),
804        }
805    }
806}
807
808#[cfg(test)]
809mod tests {
810    use std::pin::Pin;
811    use std::sync::atomic::{AtomicUsize, Ordering};
812    use std::task::{Context, Poll};
813    use std::time::Duration;
814
815    use api::v1::query_request::Query;
816    use arrow_flight::flight_service_server::FlightServiceServer;
817    use arrow_flight::{FlightData, Ticket};
818    use common_grpc::flight::FlightEncoder;
819    use common_query::{Output, OutputData};
820    use common_recordbatch::adapter::RecordBatchMetrics;
821    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
822    use datatypes::arrow::datatypes::Schema as ArrowSchema;
823    use datatypes::prelude::{ConcreteDataType, VectorRef};
824    use datatypes::schema::{ColumnSchema, Schema};
825    use datatypes::vectors::Int32Vector;
826    use futures::StreamExt;
827    use servers::grpc::flight::{FlightCraft, FlightCraftWrapper, TonicStream};
828    use tokio::net::TcpListener;
829    use tokio::task::JoinHandle;
830    use tokio::time::timeout;
831    use tokio_stream::wrappers::TcpListenerStream;
832    use tonic::{Request as TonicRequest, Response as TonicResponse, Status};
833
834    use super::*;
835
836    #[derive(Debug)]
837    struct NoopHandler;
838
839    struct MockMetricsStream {
840        schema: datatypes::schema::SchemaRef,
841        batch: Option<RecordBatch>,
842        metrics: RecordBatchMetrics,
843        terminal_metrics_only: bool,
844    }
845
846    impl futures::Stream for MockMetricsStream {
847        type Item = common_recordbatch::error::Result<RecordBatch>;
848
849        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
850            Poll::Ready(self.batch.take().map(Ok))
851        }
852
853        fn size_hint(&self) -> (usize, Option<usize>) {
854            (
855                usize::from(self.batch.is_some()),
856                Some(usize::from(self.batch.is_some())),
857            )
858        }
859    }
860
861    impl RecordBatchStream for MockMetricsStream {
862        fn name(&self) -> &str {
863            "MockMetricsStream"
864        }
865
866        fn schema(&self) -> datatypes::schema::SchemaRef {
867            self.schema.clone()
868        }
869
870        fn output_ordering(&self) -> Option<&[OrderOption]> {
871            None
872        }
873
874        fn metrics(&self) -> Option<RecordBatchMetrics> {
875            if self.terminal_metrics_only && self.batch.is_some() {
876                return None;
877            }
878            Some(self.metrics.clone())
879        }
880    }
881
882    #[derive(Debug)]
883    struct MetricsHandler;
884
885    #[derive(Debug)]
886    struct InsertOnceHandler {
887        calls: Arc<AtomicUsize>,
888    }
889
890    #[derive(Debug)]
891    struct ExtensionAwareHandler;
892
893    #[derive(Debug)]
894    struct SnapshotBindingHandler;
895
896    #[derive(Debug)]
897    struct RejectUnauthenticatedFlight;
898
899    #[derive(Debug)]
900    struct SlowFlight;
901
902    struct DelayedEofFlight {
903        schema_sent: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
904        release: Arc<tokio::sync::Notify>,
905    }
906
907    #[derive(Debug)]
908    struct LateStreamErrorFlight;
909
910    struct WaitForConcurrentFlight {
911        barrier: Arc<tokio::sync::Barrier>,
912    }
913
914    #[async_trait::async_trait]
915    impl GrpcQueryHandlerWithBoxedError for NoopHandler {
916        async fn do_query(
917            &self,
918            _query: Request,
919            _ctx: QueryContextRef,
920        ) -> std::result::Result<Output, BoxedError> {
921            Ok(Output::new_with_affected_rows(0))
922        }
923    }
924
925    #[async_trait::async_trait]
926    impl GrpcQueryHandlerWithBoxedError for InsertOnceHandler {
927        async fn do_query(
928            &self,
929            _query: Request,
930            ctx: QueryContextRef,
931        ) -> std::result::Result<Output, BoxedError> {
932            assert_eq!(ctx.channel(), Channel::Internal);
933            self.calls.fetch_add(1, Ordering::SeqCst);
934            Ok(Output::new_with_affected_rows(1))
935        }
936    }
937
938    #[async_trait::async_trait]
939    impl GrpcQueryHandlerWithBoxedError for MetricsHandler {
940        async fn do_query(
941            &self,
942            _query: Request,
943            _ctx: QueryContextRef,
944        ) -> std::result::Result<Output, BoxedError> {
945            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
946                "v",
947                ConcreteDataType::int32_datatype(),
948                false,
949            )]));
950            let batch = RecordBatch::new(
951                schema.clone(),
952                vec![Arc::new(Int32Vector::from_slice([1, 2])) as VectorRef],
953            )
954            .unwrap();
955            Ok(Output::new_with_stream(Box::pin(MockMetricsStream {
956                schema,
957                batch: Some(batch),
958                metrics: RecordBatchMetrics {
959                    region_watermarks: vec![common_recordbatch::adapter::RegionWatermarkEntry {
960                        region_id: 42,
961                        watermark: Some(99),
962                    }],
963                    ..Default::default()
964                },
965                terminal_metrics_only: true,
966            })))
967        }
968    }
969
970    #[async_trait::async_trait]
971    impl GrpcQueryHandlerWithBoxedError for ExtensionAwareHandler {
972        async fn do_query(
973            &self,
974            _query: Request,
975            ctx: QueryContextRef,
976        ) -> std::result::Result<Output, BoxedError> {
977            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
978            assert_eq!(ctx.channel(), Channel::Internal);
979            Ok(Output::new_with_affected_rows(1))
980        }
981    }
982
983    #[async_trait::async_trait]
984    impl GrpcQueryHandlerWithBoxedError for SnapshotBindingHandler {
985        async fn do_query(
986            &self,
987            _query: Request,
988            ctx: QueryContextRef,
989        ) -> std::result::Result<Output, BoxedError> {
990            assert_eq!(ctx.extension("flow.return_region_seq"), Some("true"));
991            assert_eq!(ctx.get_snapshot(1), Some(10));
992            assert_eq!(ctx.get_snapshot(2), Some(20));
993            ctx.set_snapshot(42, 99);
994            Ok(Output::new_with_affected_rows(1))
995        }
996    }
997
998    #[async_trait::async_trait]
999    impl FlightCraft for RejectUnauthenticatedFlight {
1000        async fn do_get(
1001            &self,
1002            _request: TonicRequest<Ticket>,
1003        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1004            Err(Status::unauthenticated("auth failed"))
1005        }
1006    }
1007
1008    #[async_trait::async_trait]
1009    impl FlightCraft for SlowFlight {
1010        async fn do_get(
1011            &self,
1012            _request: TonicRequest<Ticket>,
1013        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1014            tokio::time::sleep(Duration::from_secs(60)).await;
1015            Err(Status::unavailable("slow response"))
1016        }
1017    }
1018
1019    #[async_trait::async_trait]
1020    impl FlightCraft for DelayedEofFlight {
1021        async fn do_get(
1022            &self,
1023            _request: TonicRequest<Ticket>,
1024        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1025            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
1026            let schema_sent = self.schema_sent.lock().unwrap().take();
1027            let schema_stream = futures::stream::once(async move {
1028                if let Some(schema_sent) = schema_sent {
1029                    let _ = schema_sent.send(());
1030                }
1031                Ok(schema)
1032            });
1033            let release = self.release.clone();
1034            let delayed_eof = futures::stream::unfold(release, |release| async move {
1035                release.notified().await;
1036                None::<(std::result::Result<FlightData, Status>, _)>
1037            });
1038
1039            Ok(TonicResponse::new(Box::pin(
1040                schema_stream.chain(delayed_eof),
1041            )))
1042        }
1043    }
1044
1045    #[async_trait::async_trait]
1046    impl FlightCraft for LateStreamErrorFlight {
1047        async fn do_get(
1048            &self,
1049            _request: TonicRequest<Ticket>,
1050        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1051            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
1052            let stream =
1053                futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]);
1054            Ok(TonicResponse::new(Box::pin(stream)))
1055        }
1056    }
1057
1058    #[async_trait::async_trait]
1059    impl FlightCraft for WaitForConcurrentFlight {
1060        async fn do_get(
1061            &self,
1062            _request: TonicRequest<Ticket>,
1063        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
1064            self.barrier.wait().await;
1065            Err(Status::unavailable("probe started concurrently"))
1066        }
1067    }
1068
1069    async fn start_flight_server<T: FlightCraft>(handler: T) -> (String, JoinHandle<()>) {
1070        let listener = TcpListener::bind("127.0.0.1:0")
1071            .await
1072            .expect("bind test flight server");
1073        let addr = listener.local_addr().expect("local addr").to_string();
1074        let server = tokio::spawn(async move {
1075            tonic::transport::Server::builder()
1076                .add_service(FlightServiceServer::new(FlightCraftWrapper(handler)))
1077                .serve_with_incoming(TcpListenerStream::new(listener))
1078                .await
1079                .expect("serve test flight server");
1080        });
1081
1082        (addr, server)
1083    }
1084
1085    #[tokio::test]
1086    async fn wait_initialized() {
1087        let (client, handler_mut) =
1088            FrontendClient::from_empty_grpc_handler(QueryOptions::default());
1089
1090        assert!(
1091            timeout(Duration::from_millis(50), client.wait_initialized())
1092                .await
1093                .is_err()
1094        );
1095
1096        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1097        handler_mut.set_handler(Arc::downgrade(&handler)).await;
1098
1099        timeout(Duration::from_secs(1), client.wait_initialized())
1100            .await
1101            .expect("wait_initialized should complete after handler is set");
1102
1103        timeout(Duration::from_millis(10), client.wait_initialized())
1104            .await
1105            .expect("wait_initialized should be a no-op once initialized");
1106
1107        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1108        let client =
1109            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1110        assert!(
1111            timeout(Duration::from_millis(10), client.wait_initialized())
1112                .await
1113                .is_ok()
1114        );
1115
1116        let meta_client = Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend));
1117        let client = FrontendClient::from_meta_client(
1118            meta_client,
1119            QueryOptions::default(),
1120            BatchingModeOptions::default(),
1121        )
1122        .unwrap();
1123        assert!(
1124            timeout(Duration::from_millis(10), client.wait_initialized())
1125                .await
1126                .is_ok()
1127        );
1128    }
1129
1130    #[tokio::test]
1131    async fn test_handle_insert_once_calls_standalone_handler_once() {
1132        let calls = Arc::new(AtomicUsize::new(0));
1133        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(InsertOnceHandler {
1134            calls: calls.clone(),
1135        });
1136        let client =
1137            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1138        let mut peer_desc = None;
1139
1140        let affected_rows = client
1141            .handle_insert_once(
1142                Request::RowInserts(api::v1::RowInsertRequests { inserts: vec![] }),
1143                "greptime",
1144                "public",
1145                &mut peer_desc,
1146            )
1147            .await
1148            .unwrap();
1149
1150        assert_eq!(affected_rows, 1);
1151        assert_eq!(calls.load(Ordering::SeqCst), 1);
1152        assert!(peer_desc.is_none());
1153    }
1154
1155    #[tokio::test]
1156    async fn test_query_with_terminal_metrics_tracks_watermark_in_standalone_mode() {
1157        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(MetricsHandler);
1158        let client =
1159            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1160        let mut peer_desc = None;
1161
1162        let result = client
1163            .query_with_terminal_metrics(
1164                "greptime",
1165                "public",
1166                QueryRequest {
1167                    query: Some(Query::Sql("select 1".to_string())),
1168                },
1169                &[],
1170                &HashMap::new(),
1171                &mut peer_desc,
1172            )
1173            .await
1174            .unwrap();
1175        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1176
1177        let terminal_metrics = result.metrics.clone();
1178        assert!(!result.metrics.is_ready());
1179        assert!(terminal_metrics.get().is_none());
1180
1181        let OutputData::Stream(mut stream) = result.output.data else {
1182            panic!("expected stream output");
1183        };
1184        while stream.next().await.is_some() {}
1185
1186        assert!(terminal_metrics.is_ready());
1187        assert_eq!(
1188            terminal_metrics.region_watermark_map(),
1189            Some(HashMap::from([(42_u64, 99_u64)]))
1190        );
1191    }
1192
1193    #[tokio::test]
1194    async fn test_query_with_terminal_metrics_forwards_flow_extensions_in_standalone_mode() {
1195        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(ExtensionAwareHandler);
1196        let client =
1197            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1198        let mut peer_desc = None;
1199
1200        let result = client
1201            .query_with_terminal_metrics(
1202                "greptime",
1203                "public",
1204                QueryRequest {
1205                    query: Some(Query::Sql("insert into t select 1".to_string())),
1206                },
1207                &[("flow.return_region_seq", "true")],
1208                &HashMap::new(),
1209                &mut peer_desc,
1210            )
1211            .await
1212            .unwrap();
1213        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1214
1215        assert!(result.metrics.is_ready());
1216        assert!(result.region_watermark_map().is_none());
1217    }
1218
1219    #[tokio::test]
1220    async fn test_query_with_terminal_metrics_uses_standalone_snapshot_bounds() {
1221        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(SnapshotBindingHandler);
1222        let client =
1223            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1224        let mut peer_desc = None;
1225
1226        let result = client
1227            .query_with_terminal_metrics(
1228                "greptime",
1229                "public",
1230                QueryRequest {
1231                    query: Some(Query::Sql("insert into t select * from src".to_string())),
1232                },
1233                &[("flow.return_region_seq", "true")],
1234                &HashMap::from([(1, 10), (2, 20)]),
1235                &mut peer_desc,
1236            )
1237            .await
1238            .unwrap();
1239        assert!(matches!(peer_desc, Some(PeerDesc::Standalone)));
1240
1241        assert!(result.metrics.is_ready());
1242        assert_eq!(result.region_watermark_map(), None);
1243    }
1244
1245    #[tokio::test]
1246    async fn test_query_with_terminal_metrics_rejects_invalid_flow_extensions() {
1247        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
1248        let client =
1249            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
1250        let mut peer_desc = None;
1251
1252        let err = client
1253            .query_with_terminal_metrics(
1254                "greptime",
1255                "public",
1256                QueryRequest {
1257                    query: Some(Query::Sql("select 1".to_string())),
1258                },
1259                &[("flow.return_region_seq", "not-a-bool")],
1260                &HashMap::new(),
1261                &mut peer_desc,
1262            )
1263            .await
1264            .unwrap_err();
1265
1266        assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq"));
1267    }
1268
1269    #[tokio::test]
1270    async fn test_try_select_one_waits_for_stream_eof() {
1271        let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel();
1272        let release = Arc::new(tokio::sync::Notify::new());
1273        let (addr, server) = start_flight_server(DelayedEofFlight {
1274            schema_sent: Mutex::new(Some(schema_sent)),
1275            release: release.clone(),
1276        })
1277        .await;
1278        let database = Database::new(
1279            DEFAULT_CATALOG_NAME,
1280            DEFAULT_SCHEMA_NAME,
1281            Client::with_query_and_control_managers(
1282                ChannelManager::new(),
1283                ChannelManager::new(),
1284                [addr.as_str()],
1285            ),
1286        );
1287        let db = DatabaseWithPeer::new(
1288            database,
1289            Peer {
1290                id: 1,
1291                addr: addr.clone(),
1292            },
1293        );
1294        let mut probe = tokio::spawn(async move { db.try_select_one().await });
1295
1296        timeout(Duration::from_secs(1), schema_sent_rx)
1297            .await
1298            .expect("server should send the schema")
1299            .expect("schema signal should be sent");
1300        assert!(
1301            timeout(Duration::from_millis(100), &mut probe)
1302                .await
1303                .is_err(),
1304            "SELECT 1 must wait for the delayed stream tail and EOF"
1305        );
1306
1307        release.notify_one();
1308        timeout(Duration::from_secs(1), &mut probe)
1309            .await
1310            .expect("SELECT 1 should complete after EOF")
1311            .expect("probe task should not panic")
1312            .expect("SELECT 1 should succeed after EOF");
1313        server.abort();
1314    }
1315
1316    #[tokio::test]
1317    async fn test_try_select_one_propagates_late_stream_error() {
1318        let (addr, server) = start_flight_server(LateStreamErrorFlight).await;
1319        let database = Database::new(
1320            DEFAULT_CATALOG_NAME,
1321            DEFAULT_SCHEMA_NAME,
1322            Client::with_query_and_control_managers(
1323                ChannelManager::new(),
1324                ChannelManager::new(),
1325                [addr.as_str()],
1326            ),
1327        );
1328        let db = DatabaseWithPeer::new(
1329            database,
1330            Peer {
1331                id: 1,
1332                addr: addr.clone(),
1333            },
1334        );
1335
1336        let err = db.try_select_one().await.unwrap_err();
1337        server.abort();
1338
1339        assert!(format!("{err:?}").contains("late stream error"));
1340    }
1341
1342    #[tokio::test]
1343    async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() {
1344        let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await;
1345        let client = FrontendClient::from_meta_client(
1346            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1347            QueryOptions::default(),
1348            BatchingModeOptions::default(),
1349        )
1350        .unwrap();
1351
1352        let err = client
1353            .check_all_frontends_without_auth(&[Peer {
1354                id: 1,
1355                addr: addr.clone(),
1356            }])
1357            .await
1358            .unwrap_err();
1359        server.abort();
1360
1361        let Error::InvalidRequest {
1362            context, source, ..
1363        } = err
1364        else {
1365            panic!("expected InvalidRequest, got {err:?}");
1366        };
1367        assert!(context.contains(&addr));
1368        assert!(context.contains("rejected unauthenticated flownode probe"));
1369        assert_eq!(source.tonic_code(), Some(tonic::Code::Unauthenticated));
1370    }
1371
1372    #[tokio::test]
1373    async fn test_check_all_frontends_without_auth_uses_grpc_connection_timeout() {
1374        let (addr, server) = start_flight_server(SlowFlight).await;
1375        let client = FrontendClient::from_meta_client(
1376            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1377            QueryOptions::default(),
1378            BatchingModeOptions {
1379                grpc_conn_timeout: Duration::from_millis(50),
1380                ..Default::default()
1381            },
1382        )
1383        .unwrap();
1384
1385        let failures = client
1386            .check_all_frontends_without_auth(&[Peer {
1387                id: 1,
1388                addr: addr.clone(),
1389            }])
1390            .await
1391            .unwrap();
1392        server.abort();
1393
1394        assert_eq!(failures.len(), 1);
1395        assert!(failures[0].contains(&addr));
1396        assert!(failures[0].contains("health check timed out"));
1397    }
1398
1399    #[tokio::test]
1400    async fn test_check_all_frontends_without_auth_checks_frontends_concurrently() {
1401        let barrier = Arc::new(tokio::sync::Barrier::new(2));
1402        let (addr1, server1) = start_flight_server(WaitForConcurrentFlight {
1403            barrier: barrier.clone(),
1404        })
1405        .await;
1406        let (addr2, server2) = start_flight_server(WaitForConcurrentFlight { barrier }).await;
1407        let client = FrontendClient::from_meta_client(
1408            Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend)),
1409            QueryOptions::default(),
1410            BatchingModeOptions {
1411                grpc_conn_timeout: Duration::from_millis(500),
1412                ..Default::default()
1413            },
1414        )
1415        .unwrap();
1416
1417        let failures = timeout(
1418            Duration::from_secs(2),
1419            client.check_all_frontends_without_auth(&[
1420                Peer {
1421                    id: 1,
1422                    addr: addr1.clone(),
1423                },
1424                Peer {
1425                    id: 2,
1426                    addr: addr2.clone(),
1427                },
1428            ]),
1429        )
1430        .await
1431        .expect("concurrent probes should complete before per-peer timeouts")
1432        .unwrap();
1433        server1.abort();
1434        server2.abort();
1435
1436        assert_eq!(failures.len(), 2);
1437        assert!(failures.iter().any(|failure| failure.contains(&addr1)));
1438        assert!(failures.iter().any(|failure| failure.contains(&addr2)));
1439        assert!(
1440            failures
1441                .iter()
1442                .all(|failure| !failure.contains("health check timed out")),
1443            "sequential probes would time out before both requests reach the barrier: {failures:?}"
1444        );
1445    }
1446}
flow/batching_mode/frontend_client.rs

flow/batching_mode/
frontend_client.rs
