flow/
server.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//! Implementation of grpc service for flow node
16
17use std::net::SocketAddr;
18use std::sync::Arc;
19
20use api::v1::flow::DirtyWindowRequests;
21use api::v1::{RowDeleteRequests, RowInsertRequests};
22use cache::{PARTITION_INFO_CACHE_NAME, TABLE_FLOWNODE_SET_CACHE_NAME, TABLE_ROUTE_CACHE_NAME};
23use catalog::CatalogManagerRef;
24use common_base::Plugins;
25use common_datasource::object_store::LocalFileAccess;
26use common_error::ext::BoxedError;
27use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef};
28use common_meta::key::TableMetadataManagerRef;
29use common_meta::key::flow::FlowMetadataManagerRef;
30use common_meta::kv_backend::KvBackendRef;
31use common_meta::node_manager::{Flownode, NodeManagerRef};
32use common_meta::procedure_executor::ProcedureExecutorRef;
33use common_query::Output;
34use common_runtime::JoinHandle;
35use common_telemetry::tracing::info;
36use futures::TryStreamExt;
37use greptime_proto::v1::flow::{FlowRequest, FlowResponse, InsertRequests, flow_server};
38use itertools::Itertools;
39use operator::delete::Deleter;
40use operator::insert::Inserter;
41use operator::statement::StatementExecutor;
42use partition::cache::PartitionInfoCacheRef;
43use partition::manager::PartitionRuleManager;
44use query::{QueryEngine, QueryEngineFactory};
45use servers::add_service;
46use servers::grpc::builder::GrpcServerBuilder;
47use servers::grpc::{GrpcServer, GrpcServerConfig};
48use servers::http::HttpServerBuilder;
49use servers::metrics_handler::MetricsHandler;
50use servers::server::{ServerHandler, ServerHandlers};
51use session::context::QueryContextRef;
52use snafu::{OptionExt, ResultExt};
53use tokio::sync::{Mutex, broadcast, oneshot};
54use tonic::codec::CompressionEncoding;
55use tonic::{Request, Response, Status};
56
57use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef};
58use crate::adapter::{FlowStreamingEngineRef, create_worker};
59use crate::batching_mode::engine::BatchingEngine;
60use crate::error::{
61    CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu,
62    ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err,
63};
64use crate::heartbeat::HeartbeatTask;
65use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS};
66use crate::transform::register_function_to_query_engine;
67use crate::utils::{SizeReportSender, StateReportHandler};
68use crate::{Error, FlownodeOptions, FrontendClient, StreamingEngine};
69
70pub const FLOW_NODE_SERVER_NAME: &str = "FLOW_NODE_SERVER";
71/// wrapping flow node manager to avoid orphan rule with Arc<...>
72#[derive(Clone)]
73pub struct FlowService {
74    pub dual_engine: FlowDualEngineRef,
75}
76
77impl FlowService {
78    pub fn new(manager: FlowDualEngineRef) -> Self {
79        Self {
80            dual_engine: manager,
81        }
82    }
83}
84
85#[async_trait::async_trait]
86impl flow_server::Flow for FlowService {
87    async fn handle_create_remove(
88        &self,
89        request: Request<FlowRequest>,
90    ) -> Result<Response<FlowResponse>, Status> {
91        let _timer = METRIC_FLOW_PROCESSING_TIME
92            .with_label_values(&["ddl"])
93            .start_timer();
94
95        let request = request.into_inner();
96        self.dual_engine
97            .handle(request)
98            .await
99            .map_err(|err| {
100                common_telemetry::error!(err; "Failed to handle flow request");
101                err
102            })
103            .map(Response::new)
104            .map_err(to_status_with_last_err)
105    }
106
107    async fn handle_mirror_request(
108        &self,
109        request: Request<InsertRequests>,
110    ) -> Result<Response<FlowResponse>, Status> {
111        let _timer = METRIC_FLOW_PROCESSING_TIME
112            .with_label_values(&["insert"])
113            .start_timer();
114
115        let request = request.into_inner();
116        // TODO(discord9): fix protobuf import order shenanigans to remove this duplicated define
117        let mut row_count = 0;
118        let request = api::v1::region::InsertRequests {
119            requests: request
120                .requests
121                .into_iter()
122                .map(|insert| {
123                    insert.rows.as_ref().inspect(|x| row_count += x.rows.len());
124                    api::v1::region::InsertRequest {
125                        region_id: insert.region_id,
126                        rows: insert.rows,
127                        skip_wal: false,
128                        partition_expr_version: insert.partition_expr_version,
129                    }
130                })
131                .collect_vec(),
132        };
133
134        METRIC_FLOW_ROWS
135            .with_label_values(&["in"])
136            .inc_by(row_count as u64);
137
138        self.dual_engine
139            .handle_inserts(request)
140            .await
141            .map(Response::new)
142            .map_err(to_status_with_last_err)
143    }
144
145    async fn handle_mark_dirty_time_window(
146        &self,
147        reqs: Request<DirtyWindowRequests>,
148    ) -> Result<Response<FlowResponse>, Status> {
149        self.dual_engine
150            .handle_mark_window_dirty(reqs.into_inner())
151            .await
152            .map(Response::new)
153            .map_err(to_status_with_last_err)
154    }
155}
156
157#[derive(Clone)]
158pub struct FlownodeServer {
159    inner: Arc<FlownodeServerInner>,
160}
161
162/// FlownodeServerInner is the inner state of FlownodeServer,
163/// this struct mostly useful for construct/start and stop the
164/// flow node server
165struct FlownodeServerInner {
166    /// worker shutdown signal, not to be confused with server_shutdown_tx
167    worker_shutdown_tx: Mutex<broadcast::Sender<()>>,
168    /// server shutdown signal for shutdown grpc server
169    server_shutdown_tx: Mutex<broadcast::Sender<()>>,
170    /// streaming task handler
171    streaming_task_handler: Mutex<Option<JoinHandle<()>>>,
172    /// state report task handler
173    state_report_task_handler: Mutex<Option<JoinHandle<()>>>,
174    flow_service: FlowService,
175}
176
177impl FlownodeServer {
178    pub fn new(flow_service: FlowService) -> Self {
179        let (tx, _rx) = broadcast::channel::<()>(1);
180        let (server_tx, _server_rx) = broadcast::channel::<()>(1);
181        Self {
182            inner: Arc::new(FlownodeServerInner {
183                flow_service,
184                worker_shutdown_tx: Mutex::new(tx),
185                server_shutdown_tx: Mutex::new(server_tx),
186                streaming_task_handler: Mutex::new(None),
187                state_report_task_handler: Mutex::new(None),
188            }),
189        }
190    }
191
192    /// Start the background task for streaming computation.
193    ///
194    /// Should be called only after heartbeat is establish, hence can get cluster info
195    async fn start_workers(&self) -> Result<(), Error> {
196        let manager_ref = self.inner.flow_service.dual_engine.clone();
197        let mut state_report_task_handler = self.inner.state_report_task_handler.lock().await;
198        let started_state_report_task = state_report_task_handler.is_none();
199        if state_report_task_handler.is_none() {
200            *state_report_task_handler = manager_ref.clone().start_state_report_task().await;
201        }
202        drop(state_report_task_handler);
203        let handle = manager_ref
204            .streaming_engine()
205            .run_background(Some(self.inner.worker_shutdown_tx.lock().await.subscribe()));
206        self.inner
207            .streaming_task_handler
208            .lock()
209            .await
210            .replace(handle);
211
212        if let Err(err) = self
213            .inner
214            .flow_service
215            .dual_engine
216            .start_flow_consistent_check_task()
217            .await
218        {
219            self.rollback_started_workers(started_state_report_task)
220                .await;
221            return Err(err);
222        }
223
224        Ok(())
225    }
226
227    async fn rollback_started_workers(&self, abort_state_report_task: bool) {
228        let tx = self.inner.worker_shutdown_tx.lock().await;
229        if tx.send(()).is_err() {
230            info!("Receiver dropped, the flow node server has already shutdown");
231        }
232        drop(tx);
233
234        if let Some(handle) = self.inner.streaming_task_handler.lock().await.take() {
235            handle.abort();
236        }
237
238        if abort_state_report_task
239            && let Some(handle) = self.inner.state_report_task_handler.lock().await.take()
240        {
241            handle.abort();
242        }
243    }
244
245    /// Stop the background task for streaming computation.
246    async fn stop_workers(&self) -> Result<(), Error> {
247        let tx = self.inner.worker_shutdown_tx.lock().await;
248        if tx.send(()).is_err() {
249            info!("Receiver dropped, the flow node server has already shutdown");
250        }
251        // Keep state_report_task_handler alive across worker restarts.
252        // Dropping it here would permanently lose the report channel receiver.
253        self.inner
254            .flow_service
255            .dual_engine
256            .stop_flow_consistent_check_task()
257            .await?;
258        Ok(())
259    }
260}
261
262impl FlownodeServer {
263    pub fn create_flow_service(&self) -> flow_server::FlowServer<impl flow_server::Flow> {
264        flow_server::FlowServer::new(self.inner.flow_service.clone())
265            .accept_compressed(CompressionEncoding::Gzip)
266            .send_compressed(CompressionEncoding::Gzip)
267            .accept_compressed(CompressionEncoding::Zstd)
268            .send_compressed(CompressionEncoding::Zstd)
269    }
270}
271
272/// The flownode server instance.
273pub struct FlownodeInstance {
274    flownode_server: FlownodeServer,
275    services: ServerHandlers,
276    heartbeat_task: Option<HeartbeatTask>,
277}
278
279impl FlownodeInstance {
280    pub async fn start(&mut self) -> Result<(), crate::Error> {
281        if let Some(task) = &self.heartbeat_task {
282            task.start().await?;
283        }
284
285        self.flownode_server.start_workers().await?;
286
287        self.services.start_all().await.context(StartServerSnafu)?;
288
289        Ok(())
290    }
291    pub async fn shutdown(&mut self) -> Result<(), Error> {
292        self.services
293            .shutdown_all()
294            .await
295            .context(ShutdownServerSnafu)?;
296
297        self.flownode_server.stop_workers().await?;
298
299        if let Some(task) = &self.heartbeat_task {
300            task.shutdown();
301        }
302
303        Ok(())
304    }
305
306    pub fn flownode_server(&self) -> &FlownodeServer {
307        &self.flownode_server
308    }
309
310    pub fn flow_engine(&self) -> FlowDualEngineRef {
311        self.flownode_server.inner.flow_service.dual_engine.clone()
312    }
313
314    pub fn setup_services(&mut self, services: ServerHandlers) {
315        self.services = services;
316    }
317}
318
319/// [`FlownodeInstance`] Builder
320pub struct FlownodeBuilder {
321    opts: FlownodeOptions,
322    plugins: Plugins,
323    table_meta: TableMetadataManagerRef,
324    catalog_manager: CatalogManagerRef,
325    flow_metadata_manager: FlowMetadataManagerRef,
326    heartbeat_task: Option<HeartbeatTask>,
327    /// receive a oneshot sender to send state size report
328    state_report_handler: Option<StateReportHandler>,
329    frontend_client: Arc<FrontendClient>,
330}
331
332impl FlownodeBuilder {
333    /// init flownode builder
334    pub fn new(
335        opts: FlownodeOptions,
336        plugins: Plugins,
337        table_meta: TableMetadataManagerRef,
338        catalog_manager: CatalogManagerRef,
339        flow_metadata_manager: FlowMetadataManagerRef,
340        frontend_client: Arc<FrontendClient>,
341    ) -> Self {
342        Self {
343            opts,
344            plugins,
345            table_meta,
346            catalog_manager,
347            flow_metadata_manager,
348            heartbeat_task: None,
349            state_report_handler: None,
350            frontend_client,
351        }
352    }
353
354    pub fn with_heartbeat_task(self, heartbeat_task: HeartbeatTask) -> Self {
355        let (sender, receiver) = SizeReportSender::new();
356        Self {
357            heartbeat_task: Some(heartbeat_task.with_query_stat_size(sender)),
358            state_report_handler: Some(receiver),
359            ..self
360        }
361    }
362
363    pub fn opts(&self) -> &FlownodeOptions {
364        &self.opts
365    }
366
367    pub fn table_meta(&self) -> &TableMetadataManagerRef {
368        &self.table_meta
369    }
370
371    pub fn catalog_manager(&self) -> &CatalogManagerRef {
372        &self.catalog_manager
373    }
374
375    pub fn flow_metadata_manager(&self) -> &FlowMetadataManagerRef {
376        &self.flow_metadata_manager
377    }
378
379    pub fn frontend_client(&self) -> &Arc<FrontendClient> {
380        &self.frontend_client
381    }
382
383    pub fn set_plugins(&mut self, plugins: Plugins) {
384        self.plugins = plugins;
385    }
386
387    pub async fn build(mut self) -> Result<FlownodeInstance, Error> {
388        // TODO(discord9): does this query engine need those?
389        let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
390            // query engine in flownode is only used for translate plan with resolved table source.
391            self.catalog_manager.clone(),
392            None,
393            None,
394            None,
395            None,
396            None,
397            false,
398            Default::default(),
399            self.opts.query.clone(),
400        )
401        .context(DatafusionSnafu {
402            context: "Failed to build query engine",
403        })?;
404        let manager = Arc::new(
405            self.build_manager(query_engine_factory.query_engine())
406                .await?,
407        );
408        let batching = Arc::new(BatchingEngine::new(
409            self.frontend_client.clone(),
410            query_engine_factory.query_engine(),
411            self.flow_metadata_manager.clone(),
412            self.table_meta.clone(),
413            self.catalog_manager.clone(),
414            self.opts.flow.batching_mode.clone(),
415        ));
416        let dual = Arc::new(FlowDualEngine::new(
417            manager.clone(),
418            batching,
419            self.flow_metadata_manager.clone(),
420            self.catalog_manager.clone(),
421            self.plugins.clone(),
422        ));
423        if let Some(handler) = self.state_report_handler.take() {
424            dual.set_state_report_handler(handler).await;
425        }
426
427        let server = FlownodeServer::new(FlowService::new(dual));
428
429        let heartbeat_task = self.heartbeat_task;
430
431        let instance = FlownodeInstance {
432            flownode_server: server,
433            services: ServerHandlers::default(),
434            heartbeat_task,
435        };
436        Ok(instance)
437    }
438
439    /// build [`FlowWorkerManager`], note this doesn't take ownership of `self`,
440    /// nor does it actually start running the worker.
441    async fn build_manager(
442        &mut self,
443        query_engine: Arc<dyn QueryEngine>,
444    ) -> Result<StreamingEngine, Error> {
445        let table_meta = self.table_meta.clone();
446
447        register_function_to_query_engine(&query_engine);
448
449        let num_workers = self.opts.flow.num_workers;
450
451        let node_id = self.opts.node_id.map(|id| id as u32);
452
453        let mut man = StreamingEngine::new(node_id, query_engine, table_meta);
454        for worker_id in 0..num_workers {
455            let (tx, rx) = oneshot::channel();
456
457            let _handle = std::thread::Builder::new()
458                .name(format!("flow-worker-{}", worker_id))
459                .spawn(move || {
460                    let (handle, mut worker) = create_worker();
461                    let _ = tx.send(handle);
462                    info!("Flow Worker started in new thread");
463                    worker.run();
464                });
465            let worker_handle = rx.await.map_err(|e| {
466                UnexpectedSnafu {
467                    reason: format!("Failed to receive worker handle: {}", e),
468                }
469                .build()
470            })?;
471            man.add_worker_handle(worker_handle);
472        }
473        info!("Flow Node Manager started");
474        Ok(man)
475    }
476}
477
478/// Useful in distributed mode
479pub struct FlownodeServiceBuilder<'a> {
480    opts: &'a FlownodeOptions,
481    grpc_server: Option<GrpcServer>,
482    enable_http_service: bool,
483}
484
485impl<'a> FlownodeServiceBuilder<'a> {
486    pub fn new(opts: &'a FlownodeOptions) -> Self {
487        Self {
488            opts,
489            grpc_server: None,
490            enable_http_service: false,
491        }
492    }
493
494    pub fn enable_http_service(self) -> Self {
495        Self {
496            enable_http_service: true,
497            ..self
498        }
499    }
500
501    pub fn with_grpc_server(self, grpc_server: GrpcServer) -> Self {
502        Self {
503            grpc_server: Some(grpc_server),
504            ..self
505        }
506    }
507
508    pub fn with_default_grpc_server(mut self, flownode_server: &FlownodeServer) -> Self {
509        let grpc_server = Self::grpc_server_builder(self.opts, flownode_server).build();
510        self.grpc_server = Some(grpc_server);
511        self
512    }
513
514    pub fn build(mut self) -> Result<ServerHandlers, Error> {
515        let handlers = ServerHandlers::default();
516        if let Some(grpc_server) = self.grpc_server.take() {
517            let addr: SocketAddr = self.opts.grpc.bind_addr.parse().context(ParseAddrSnafu {
518                addr: &self.opts.grpc.bind_addr,
519            })?;
520            let handler: ServerHandler = (Box::new(grpc_server), addr);
521            handlers.insert(handler);
522        }
523
524        if self.enable_http_service {
525            let http_server = HttpServerBuilder::new(self.opts.http.clone())
526                .with_metrics_handler(MetricsHandler)
527                .build();
528            let addr: SocketAddr = self.opts.http.addr.parse().context(ParseAddrSnafu {
529                addr: &self.opts.http.addr,
530            })?;
531            let handler: ServerHandler = (Box::new(http_server), addr);
532            handlers.insert(handler);
533        }
534        Ok(handlers)
535    }
536
537    pub fn grpc_server_builder(
538        opts: &FlownodeOptions,
539        flownode_server: &FlownodeServer,
540    ) -> GrpcServerBuilder {
541        let config = GrpcServerConfig {
542            max_recv_message_size: opts.grpc.max_recv_message_size.as_bytes() as usize,
543            max_send_message_size: opts.grpc.max_send_message_size.as_bytes() as usize,
544            tls: opts.grpc.tls.clone(),
545            max_connection_age: opts.grpc.max_connection_age,
546        };
547        let service = flownode_server.create_flow_service();
548        let runtime = common_runtime::global_runtime();
549        let mut builder = GrpcServerBuilder::new(config, runtime);
550        add_service!(builder, service);
551        builder
552    }
553}
554
555/// Basically a tiny frontend that communicates with datanode, different from [`FrontendClient`] which
556/// connect to a real frontend instead, this is used for flow's streaming engine. And is for simple query.
557///
558/// For heavy query use [`FrontendClient`] which offload computation to frontend, lifting the load from flownode
559#[derive(Clone)]
560pub struct FrontendInvoker {
561    inserter: Arc<Inserter>,
562    deleter: Arc<Deleter>,
563    statement_executor: Arc<StatementExecutor>,
564}
565
566impl FrontendInvoker {
567    pub fn new(
568        inserter: Arc<Inserter>,
569        deleter: Arc<Deleter>,
570        statement_executor: Arc<StatementExecutor>,
571    ) -> Self {
572        Self {
573            inserter,
574            deleter,
575            statement_executor,
576        }
577    }
578
579    pub async fn build_from(
580        flow_streaming_engine: FlowStreamingEngineRef,
581        catalog_manager: CatalogManagerRef,
582        kv_backend: KvBackendRef,
583        layered_cache_registry: LayeredCacheRegistryRef,
584        procedure_executor: ProcedureExecutorRef,
585        node_manager: NodeManagerRef,
586        origin_frontend_addr: String,
587    ) -> Result<FrontendInvoker, Error> {
588        let table_route_cache: TableRouteCacheRef =
589            layered_cache_registry.get().context(CacheRequiredSnafu {
590                name: TABLE_ROUTE_CACHE_NAME,
591            })?;
592        let partition_info_cache: PartitionInfoCacheRef =
593            layered_cache_registry.get().context(CacheRequiredSnafu {
594                name: PARTITION_INFO_CACHE_NAME,
595            })?;
596
597        let partition_manager = Arc::new(PartitionRuleManager::new(
598            kv_backend.clone(),
599            table_route_cache.clone(),
600            partition_info_cache.clone(),
601        ));
602
603        let table_flownode_cache: TableFlownodeSetCacheRef =
604            layered_cache_registry.get().context(CacheRequiredSnafu {
605                name: TABLE_FLOWNODE_SET_CACHE_NAME,
606            })?;
607
608        // TODO(auto_create_table): flow sink tables are created through a controlled
609        // `CREATE FLOW` path, not client writes, so they are intentionally exempt from
610        // the frontend's global auto-create switch. Revisit if flow should honor it.
611        let inserter = Arc::new(Inserter::new(
612            catalog_manager.clone(),
613            partition_manager.clone(),
614            node_manager.clone(),
615            table_flownode_cache,
616            true,
617        ));
618
619        let deleter = Arc::new(Deleter::new(
620            catalog_manager.clone(),
621            partition_manager.clone(),
622            node_manager.clone(),
623        ));
624
625        let query_engine = flow_streaming_engine.query_engine.clone();
626
627        let statement_executor = Arc::new(StatementExecutor::new(
628            catalog_manager.clone(),
629            query_engine.clone(),
630            procedure_executor.clone(),
631            kv_backend.clone(),
632            layered_cache_registry.clone(),
633            inserter.clone(),
634            partition_manager,
635            None,
636            origin_frontend_addr,
637            LocalFileAccess::Disabled,
638        ));
639
640        let invoker = FrontendInvoker::new(inserter, deleter, statement_executor);
641        Ok(invoker)
642    }
643}
644
645impl FrontendInvoker {
646    pub async fn row_inserts(
647        &self,
648        requests: RowInsertRequests,
649        ctx: QueryContextRef,
650    ) -> common_frontend::error::Result<Output> {
651        let _timer = METRIC_FLOW_PROCESSING_TIME
652            .with_label_values(&["output_insert"])
653            .start_timer();
654
655        self.inserter
656            .handle_row_inserts(requests, ctx, &self.statement_executor, false, false)
657            .await
658            .map_err(BoxedError::new)
659            .context(common_frontend::error::ExternalSnafu)
660    }
661
662    pub async fn row_deletes(
663        &self,
664        requests: RowDeleteRequests,
665        ctx: QueryContextRef,
666    ) -> common_frontend::error::Result<Output> {
667        let _timer = METRIC_FLOW_PROCESSING_TIME
668            .with_label_values(&["output_delete"])
669            .start_timer();
670
671        self.deleter
672            .handle_row_deletes(requests, ctx)
673            .await
674            .map_err(BoxedError::new)
675            .context(common_frontend::error::ExternalSnafu)
676    }
677
678    pub fn statement_executor(&self) -> Arc<StatementExecutor> {
679        self.statement_executor.clone()
680    }
681}
682
683/// get all flow ids in this flownode
684pub(crate) async fn get_all_flow_ids(
685    flow_metadata_manager: &FlowMetadataManagerRef,
686    catalog_manager: &CatalogManagerRef,
687    nodeid: Option<u64>,
688) -> Result<Vec<u32>, Error> {
689    let ret = if let Some(nodeid) = nodeid {
690        let flow_ids_one_node = flow_metadata_manager
691            .flownode_flow_manager()
692            .flows(nodeid)
693            .try_collect::<Vec<_>>()
694            .await
695            .context(ListFlowsSnafu { id: Some(nodeid) })?;
696        flow_ids_one_node.into_iter().map(|(id, _)| id).collect()
697    } else {
698        let all_catalogs = catalog_manager
699            .catalog_names()
700            .await
701            .map_err(BoxedError::new)
702            .context(ExternalSnafu)?;
703        let mut all_flow_ids = vec![];
704        for catalog in all_catalogs {
705            let flows = flow_metadata_manager
706                .flow_name_manager()
707                .flow_names(&catalog)
708                .await
709                .try_collect::<Vec<_>>()
710                .await
711                .map_err(BoxedError::new)
712                .context(ExternalSnafu)?;
713
714            all_flow_ids.extend(flows.into_iter().map(|(_, id)| id.flow_id()));
715        }
716        all_flow_ids
717    };
718
719    Ok(ret)
720}
721
722#[cfg(test)]
723mod tests {
724    use std::sync::Arc;
725    use std::time::Duration;
726
727    use api::v1::HealthCheckRequest;
728    use api::v1::health_check_client::HealthCheckClient;
729    use api::v1::meta::Role;
730    use catalog::memory::new_memory_catalog_manager;
731    use common_base::Plugins;
732    use common_meta::key::TableMetadataManager;
733    use common_meta::key::flow::FlowMetadataManager;
734    use common_meta::kv_backend::memory::MemoryKvBackend;
735    use meta_client::client::MetaClient;
736    use query::options::QueryOptions;
737    use servers::grpc::GRPC_SERVER;
738
739    use super::*;
740    use crate::adapter::flownode_impl::FlowDualEngine;
741    use crate::batching_mode::BatchingModeOptions;
742    use crate::batching_mode::engine::BatchingEngine;
743    use crate::utils::SizeReportSender;
744
745    async fn new_test_flownode_server() -> (FlownodeServer, SizeReportSender) {
746        let (frontend_client, _handler) =
747            FrontendClient::from_empty_grpc_handler(QueryOptions::default());
748
749        new_test_flownode_server_with_frontend_client(
750            frontend_client,
751            BatchingModeOptions::default(),
752            None,
753        )
754        .await
755    }
756
757    async fn new_test_flownode_server_with_frontend_client(
758        frontend_client: FrontendClient,
759        batching_opts: BatchingModeOptions,
760        node_id: Option<u32>,
761    ) -> (FlownodeServer, SizeReportSender) {
762        let kv_backend = Arc::new(MemoryKvBackend::new());
763        let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
764        table_meta.init().await.unwrap();
765        let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend.clone()));
766        let catalog_manager = new_memory_catalog_manager().unwrap();
767        let query_engine = crate::test_utils::create_test_query_engine();
768
769        let streaming_engine = Arc::new(StreamingEngine::new(
770            node_id,
771            query_engine.clone(),
772            table_meta.clone(),
773        ));
774        let batching_engine = Arc::new(BatchingEngine::new(
775            Arc::new(frontend_client),
776            query_engine,
777            flow_meta.clone(),
778            table_meta,
779            catalog_manager.clone(),
780            batching_opts,
781        ));
782        let dual_engine = Arc::new(FlowDualEngine::new(
783            streaming_engine,
784            batching_engine,
785            flow_meta,
786            catalog_manager,
787            Plugins::new(),
788        ));
789
790        let (report_sender, report_handler) = SizeReportSender::new();
791        dual_engine.set_state_report_handler(report_handler).await;
792
793        let server = FlownodeServer::new(FlowService::new(dual_engine));
794        (server, report_sender)
795    }
796
797    #[tokio::test]
798    async fn test_state_report_handler_survives_worker_restart() {
799        let (server, report_sender) = new_test_flownode_server().await;
800
801        server.start_workers().await.unwrap();
802        report_sender.query(Duration::from_secs(3)).await.unwrap();
803
804        server.stop_workers().await.unwrap();
805        report_sender.query(Duration::from_secs(3)).await.unwrap();
806
807        server.start_workers().await.unwrap();
808        report_sender.query(Duration::from_secs(3)).await.unwrap();
809
810        server.stop_workers().await.unwrap();
811    }
812
813    #[tokio::test]
814    async fn test_start_workers_rolls_back_on_check_task_start_failure() {
815        let batching_opts = BatchingModeOptions {
816            experimental_frontend_scan_timeout: Duration::from_millis(1),
817            ..Default::default()
818        };
819        let frontend_client = FrontendClient::from_meta_client(
820            Arc::new(MetaClient::new(0, Role::Frontend)),
821            QueryOptions::default(),
822            batching_opts.clone(),
823        )
824        .unwrap();
825        let (server, _report_sender) =
826            new_test_flownode_server_with_frontend_client(frontend_client, batching_opts, Some(1))
827                .await;
828
829        server.start_workers().await.unwrap_err();
830
831        assert!(server.inner.streaming_task_handler.lock().await.is_none());
832        assert!(
833            server
834                .inner
835                .state_report_task_handler
836                .lock()
837                .await
838                .is_none()
839        );
840    }
841
842    #[tokio::test]
843    async fn test_service_builder_registers_reachable_health_check() {
844        // Arrange: compose the production gRPC service with an ephemeral local listener.
845        let (flownode_server, _report_sender) = new_test_flownode_server().await;
846        let mut opts = FlownodeOptions::default();
847        opts.grpc.bind_addr = "127.0.0.1:0".to_string();
848        let mut services = FlownodeServiceBuilder::new(&opts)
849            .with_default_grpc_server(&flownode_server)
850            .build()
851            .unwrap();
852        services.start_all().await.unwrap();
853        let addr = services.addr(GRPC_SERVER).unwrap();
854
855        // Act: call the shared health handler through the registered production server.
856        let mut client = HealthCheckClient::connect(format!("http://{addr}"))
857            .await
858            .unwrap();
859        let result = client.health_check(HealthCheckRequest {}).await;
860
861        services.shutdown_all().await.unwrap();
862
863        // Assert: the service composition exposes a healthy endpoint.
864        assert!(result.is_ok());
865    }
866}
flow/server.rs

flow/
server.rs
