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

flow/
server.rs
