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

flow/
server.rs
