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 catalog::CatalogManagerRef;
22use common_base::Plugins;
23use common_error::ext::BoxedError;
24use common_meta::key::TableMetadataManagerRef;
25use common_meta::key::flow::FlowMetadataManagerRef;
26use common_meta::node_manager::Flownode;
27use futures::TryStreamExt;
28use greptime_proto::v1::flow::{FlowRequest, FlowResponse, InsertRequests, flow_server};
29use itertools::Itertools;
30use query::QueryEngineFactory;
31use servers::add_service;
32use servers::grpc::builder::GrpcServerBuilder;
33use servers::grpc::{GrpcServer, GrpcServerConfig};
34use servers::http::HttpServerBuilder;
35use servers::metrics_handler::MetricsHandler;
36use servers::server::{ServerHandler, ServerHandlers};
37use snafu::ResultExt;
38use tonic::codec::CompressionEncoding;
39use tonic::{Request, Response, Status};
40
41use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef};
42use crate::batching_mode::engine::BatchingEngine;
43use crate::error::{
44    DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu,
45    StartServerSnafu, to_status_with_last_err,
46};
47use crate::heartbeat::HeartbeatTask;
48use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS};
49use crate::utils::{SizeReportSender, StateReportHandler};
50use crate::{BatchingExecutionFactory, Error, FlownodeOptions, FrontendClient, StreamingEngine};
51
52pub const FLOW_NODE_SERVER_NAME: &str = "FLOW_NODE_SERVER";
53/// wrapping flow node manager to avoid orphan rule with Arc<...>
54#[derive(Clone)]
55pub struct FlowService {
56    pub dual_engine: FlowDualEngineRef,
57}
58
59impl FlowService {
60    pub fn new(manager: FlowDualEngineRef) -> Self {
61        Self {
62            dual_engine: manager,
63        }
64    }
65}
66
67#[async_trait::async_trait]
68impl flow_server::Flow for FlowService {
69    async fn handle_create_remove(
70        &self,
71        request: Request<FlowRequest>,
72    ) -> Result<Response<FlowResponse>, Status> {
73        let _timer = METRIC_FLOW_PROCESSING_TIME
74            .with_label_values(&["ddl"])
75            .start_timer();
76
77        let request = request.into_inner();
78        self.dual_engine
79            .handle(request)
80            .await
81            .map_err(|err| {
82                common_telemetry::error!(err; "Failed to handle flow request");
83                err
84            })
85            .map(Response::new)
86            .map_err(to_status_with_last_err)
87    }
88
89    async fn handle_mirror_request(
90        &self,
91        request: Request<InsertRequests>,
92    ) -> Result<Response<FlowResponse>, Status> {
93        let _timer = METRIC_FLOW_PROCESSING_TIME
94            .with_label_values(&["insert"])
95            .start_timer();
96
97        let request = request.into_inner();
98        // TODO(discord9): fix protobuf import order shenanigans to remove this duplicated define
99        let mut row_count = 0;
100        let request = api::v1::region::InsertRequests {
101            requests: request
102                .requests
103                .into_iter()
104                .map(|insert| {
105                    insert.rows.as_ref().inspect(|x| row_count += x.rows.len());
106                    api::v1::region::InsertRequest {
107                        region_id: insert.region_id,
108                        rows: insert.rows,
109                        skip_wal: false,
110                        partition_expr_version: insert.partition_expr_version,
111                    }
112                })
113                .collect_vec(),
114        };
115
116        METRIC_FLOW_ROWS
117            .with_label_values(&["in"])
118            .inc_by(row_count as u64);
119
120        self.dual_engine
121            .handle_inserts(request)
122            .await
123            .map(Response::new)
124            .map_err(to_status_with_last_err)
125    }
126
127    async fn handle_mark_dirty_time_window(
128        &self,
129        reqs: Request<DirtyWindowRequests>,
130    ) -> Result<Response<FlowResponse>, Status> {
131        self.dual_engine
132            .handle_mark_window_dirty(reqs.into_inner())
133            .await
134            .map(Response::new)
135            .map_err(to_status_with_last_err)
136    }
137}
138
139#[derive(Clone)]
140pub struct FlownodeServer {
141    inner: Arc<FlownodeServerInner>,
142}
143
144/// FlownodeServerInner is the inner state of FlownodeServer,
145/// this struct mostly useful for construct/start and stop the
146/// flow node server
147struct FlownodeServerInner {
148    flow_service: FlowService,
149}
150
151impl FlownodeServer {
152    pub fn new(flow_service: FlowService) -> Self {
153        Self {
154            inner: Arc::new(FlownodeServerInner { flow_service }),
155        }
156    }
157}
158
159impl FlownodeServer {
160    pub fn create_flow_service(&self) -> flow_server::FlowServer<impl flow_server::Flow> {
161        flow_server::FlowServer::new(self.inner.flow_service.clone())
162            .accept_compressed(CompressionEncoding::Gzip)
163            .send_compressed(CompressionEncoding::Gzip)
164            .accept_compressed(CompressionEncoding::Zstd)
165            .send_compressed(CompressionEncoding::Zstd)
166    }
167}
168
169/// The flownode server instance.
170pub struct FlownodeInstance {
171    flownode_server: FlownodeServer,
172    services: ServerHandlers,
173    heartbeat_task: Option<HeartbeatTask>,
174    state_report_task: Option<common_runtime::JoinHandle<()>>,
175    consistent_check_task_started: bool,
176}
177
178impl FlownodeInstance {
179    pub async fn start(&mut self) -> Result<(), crate::Error> {
180        if let Some(task) = &self.heartbeat_task {
181            task.start().await?;
182        }
183
184        let engine = self.flow_engine();
185        // The state-report task owns the only report receiver, so keep it alive
186        // across an ordinary stop/start cycle of this instance.
187        if self.state_report_task.is_none() {
188            self.state_report_task = engine.clone().start_state_report_task().await;
189        }
190        if let Err(err) = engine.start_flow_consistent_check_task().await {
191            self.rollback_background_tasks().await;
192            return Err(err);
193        }
194        self.consistent_check_task_started = true;
195
196        if let Err(err) = self.services.start_all().await.context(StartServerSnafu) {
197            self.rollback_background_tasks().await;
198            return Err(err);
199        }
200
201        Ok(())
202    }
203    pub async fn shutdown(&mut self) -> Result<(), Error> {
204        let services_result = self
205            .services
206            .shutdown_all()
207            .await
208            .context(ShutdownServerSnafu);
209        let tasks_result = self.stop_background_tasks().await;
210
211        services_result?;
212        tasks_result?;
213        Ok(())
214    }
215
216    async fn stop_background_tasks(&mut self) -> Result<(), Error> {
217        let check_result = if self.consistent_check_task_started {
218            self.consistent_check_task_started = false;
219            self.flow_engine().stop_flow_consistent_check_task().await
220        } else {
221            Ok(())
222        };
223
224        if let Some(task) = &self.heartbeat_task {
225            task.shutdown();
226        }
227
228        check_result
229    }
230
231    async fn rollback_background_tasks(&mut self) {
232        if let Err(err) = self.stop_background_tasks().await {
233            common_telemetry::error!(err; "Failed to roll back flownode background tasks");
234        }
235    }
236
237    pub fn flownode_server(&self) -> &FlownodeServer {
238        &self.flownode_server
239    }
240
241    pub fn flow_engine(&self) -> FlowDualEngineRef {
242        self.flownode_server.inner.flow_service.dual_engine.clone()
243    }
244
245    pub fn setup_services(&mut self, services: ServerHandlers) {
246        self.services = services;
247    }
248}
249
250/// [`FlownodeInstance`] Builder
251pub struct FlownodeBuilder {
252    opts: FlownodeOptions,
253    plugins: Plugins,
254    table_meta: TableMetadataManagerRef,
255    catalog_manager: CatalogManagerRef,
256    flow_metadata_manager: FlowMetadataManagerRef,
257    heartbeat_task: Option<HeartbeatTask>,
258    /// receive a oneshot sender to send state size report
259    state_report_handler: Option<StateReportHandler>,
260    frontend_client: Arc<FrontendClient>,
261}
262
263impl FlownodeBuilder {
264    /// init flownode builder
265    pub fn new(
266        opts: FlownodeOptions,
267        plugins: Plugins,
268        table_meta: TableMetadataManagerRef,
269        catalog_manager: CatalogManagerRef,
270        flow_metadata_manager: FlowMetadataManagerRef,
271        frontend_client: Arc<FrontendClient>,
272    ) -> Self {
273        Self {
274            opts,
275            plugins,
276            table_meta,
277            catalog_manager,
278            flow_metadata_manager,
279            heartbeat_task: None,
280            state_report_handler: None,
281            frontend_client,
282        }
283    }
284
285    pub fn with_heartbeat_task(self, heartbeat_task: HeartbeatTask) -> Self {
286        let (sender, receiver) = SizeReportSender::new();
287        Self {
288            heartbeat_task: Some(heartbeat_task.with_query_stat_size(sender)),
289            state_report_handler: Some(receiver),
290            ..self
291        }
292    }
293
294    pub fn opts(&self) -> &FlownodeOptions {
295        &self.opts
296    }
297
298    pub fn table_meta(&self) -> &TableMetadataManagerRef {
299        &self.table_meta
300    }
301
302    pub fn catalog_manager(&self) -> &CatalogManagerRef {
303        &self.catalog_manager
304    }
305
306    pub fn flow_metadata_manager(&self) -> &FlowMetadataManagerRef {
307        &self.flow_metadata_manager
308    }
309
310    pub fn frontend_client(&self) -> &Arc<FrontendClient> {
311        &self.frontend_client
312    }
313
314    pub fn set_plugins(&mut self, plugins: Plugins) {
315        self.plugins = plugins;
316    }
317
318    pub async fn build(mut self) -> Result<FlownodeInstance, Error> {
319        // TODO(discord9): does this query engine need those?
320        let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
321            // query engine in flownode is only used for translate plan with resolved table source.
322            self.catalog_manager.clone(),
323            None,
324            None,
325            None,
326            None,
327            None,
328            false,
329            Default::default(),
330            self.opts.query.clone(),
331        )
332        .context(DatafusionSnafu {
333            context: "Failed to build query engine",
334        })?;
335        let manager = Arc::new(self.build_manager(query_engine_factory.query_engine()));
336        let batching = Arc::new(BatchingEngine::new_with_execution(
337            self.frontend_client.clone(),
338            query_engine_factory.query_engine(),
339            self.flow_metadata_manager.clone(),
340            self.table_meta.clone(),
341            self.catalog_manager.clone(),
342            self.opts.flow.batching_mode.clone(),
343            self.plugins.get::<Arc<dyn BatchingExecutionFactory>>(),
344        ));
345        let dual = Arc::new(FlowDualEngine::new(
346            manager.clone(),
347            batching,
348            self.flow_metadata_manager.clone(),
349            self.catalog_manager.clone(),
350            self.plugins.clone(),
351        ));
352        if let Some(handler) = self.state_report_handler.take() {
353            dual.set_state_report_handler(handler).await;
354        }
355
356        let server = FlownodeServer::new(FlowService::new(dual));
357
358        let heartbeat_task = self.heartbeat_task;
359
360        let instance = FlownodeInstance {
361            flownode_server: server,
362            services: ServerHandlers::default(),
363            heartbeat_task,
364            state_report_task: None,
365            consistent_check_task_started: false,
366        };
367        Ok(instance)
368    }
369
370    fn build_manager(&self, query_engine: Arc<dyn query::QueryEngine>) -> StreamingEngine {
371        StreamingEngine::new(
372            self.opts.node_id.map(|id| id as u32),
373            query_engine,
374            self.table_meta.clone(),
375            self.frontend_client.clone(),
376        )
377    }
378}
379
380/// Useful in distributed mode
381pub struct FlownodeServiceBuilder<'a> {
382    opts: &'a FlownodeOptions,
383    grpc_server: Option<GrpcServer>,
384    enable_http_service: bool,
385}
386
387impl<'a> FlownodeServiceBuilder<'a> {
388    pub fn new(opts: &'a FlownodeOptions) -> Self {
389        Self {
390            opts,
391            grpc_server: None,
392            enable_http_service: false,
393        }
394    }
395
396    pub fn enable_http_service(self) -> Self {
397        Self {
398            enable_http_service: true,
399            ..self
400        }
401    }
402
403    pub fn with_grpc_server(self, grpc_server: GrpcServer) -> Self {
404        Self {
405            grpc_server: Some(grpc_server),
406            ..self
407        }
408    }
409
410    pub fn with_default_grpc_server(mut self, flownode_server: &FlownodeServer) -> Self {
411        let grpc_server = Self::grpc_server_builder(self.opts, flownode_server).build();
412        self.grpc_server = Some(grpc_server);
413        self
414    }
415
416    pub fn build(mut self) -> Result<ServerHandlers, Error> {
417        let handlers = ServerHandlers::default();
418        if let Some(grpc_server) = self.grpc_server.take() {
419            let addr: SocketAddr = self.opts.grpc.bind_addr.parse().context(ParseAddrSnafu {
420                addr: &self.opts.grpc.bind_addr,
421            })?;
422            let handler: ServerHandler = (Box::new(grpc_server), addr);
423            handlers.insert(handler);
424        }
425
426        if self.enable_http_service {
427            let http_server = HttpServerBuilder::new(self.opts.http.clone())
428                .with_metrics_handler(MetricsHandler)
429                .build();
430            let addr: SocketAddr = self.opts.http.addr.parse().context(ParseAddrSnafu {
431                addr: &self.opts.http.addr,
432            })?;
433            let handler: ServerHandler = (Box::new(http_server), addr);
434            handlers.insert(handler);
435        }
436        Ok(handlers)
437    }
438
439    pub fn grpc_server_builder(
440        opts: &FlownodeOptions,
441        flownode_server: &FlownodeServer,
442    ) -> GrpcServerBuilder {
443        let config = GrpcServerConfig {
444            max_recv_message_size: opts.grpc.max_recv_message_size.as_bytes() as usize,
445            max_send_message_size: opts.grpc.max_send_message_size.as_bytes() as usize,
446            tls: opts.grpc.tls.clone(),
447            max_connection_age: opts.grpc.max_connection_age,
448        };
449        let service = flownode_server.create_flow_service();
450        let runtime = common_runtime::global_runtime();
451        let mut builder = GrpcServerBuilder::new(config, runtime);
452        add_service!(builder, service);
453        builder
454    }
455}
456
457/// get all flow ids in this flownode
458pub(crate) async fn get_all_flow_ids(
459    flow_metadata_manager: &FlowMetadataManagerRef,
460    catalog_manager: &CatalogManagerRef,
461    nodeid: Option<u64>,
462) -> Result<Vec<u32>, Error> {
463    let ret = if let Some(nodeid) = nodeid {
464        let flow_ids_one_node = flow_metadata_manager
465            .flownode_flow_manager()
466            .flows(nodeid)
467            .try_collect::<Vec<_>>()
468            .await
469            .context(ListFlowsSnafu { id: Some(nodeid) })?;
470        flow_ids_one_node.into_iter().map(|(id, _)| id).collect()
471    } else {
472        let all_catalogs = catalog_manager
473            .catalog_names()
474            .await
475            .map_err(BoxedError::new)
476            .context(ExternalSnafu)?;
477        let mut all_flow_ids = vec![];
478        for catalog in all_catalogs {
479            let flows = flow_metadata_manager
480                .flow_name_manager()
481                .flow_names(&catalog)
482                .await
483                .try_collect::<Vec<_>>()
484                .await
485                .map_err(BoxedError::new)
486                .context(ExternalSnafu)?;
487
488            all_flow_ids.extend(flows.into_iter().map(|(_, id)| id.flow_id()));
489        }
490        all_flow_ids
491    };
492
493    Ok(ret)
494}
495
496#[cfg(test)]
497mod tests {
498    use std::sync::Arc;
499
500    use catalog::memory::new_memory_catalog_manager;
501    use common_base::Plugins;
502    use common_meta::key::TableMetadataManager;
503    use common_meta::key::flow::FlowMetadataManager;
504    use common_meta::kv_backend::memory::MemoryKvBackend;
505    use query::options::QueryOptions;
506
507    use super::*;
508    use crate::batching_mode::BatchingModeOptions;
509
510    async fn new_test_flownode_server() -> (FlownodeServer, SizeReportSender) {
511        let (frontend_client, _handler) =
512            FrontendClient::from_empty_grpc_handler(QueryOptions::default());
513
514        new_test_flownode_server_with_frontend_client(
515            frontend_client,
516            BatchingModeOptions::default(),
517            None,
518        )
519        .await
520    }
521
522    async fn new_test_flownode_server_with_frontend_client(
523        frontend_client: FrontendClient,
524        batching_opts: BatchingModeOptions,
525        node_id: Option<u32>,
526    ) -> (FlownodeServer, SizeReportSender) {
527        let kv_backend = Arc::new(MemoryKvBackend::new());
528        let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
529        table_meta.init().await.unwrap();
530        let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend.clone()));
531        let catalog_manager = new_memory_catalog_manager().unwrap();
532        let query_engine = crate::test_utils::create_test_query_engine();
533
534        let frontend_client = Arc::new(frontend_client);
535        let streaming_engine = Arc::new(StreamingEngine::new(
536            node_id,
537            query_engine.clone(),
538            table_meta.clone(),
539            frontend_client.clone(),
540        ));
541        let batching_engine = Arc::new(BatchingEngine::new(
542            frontend_client,
543            query_engine,
544            flow_meta.clone(),
545            table_meta,
546            catalog_manager.clone(),
547            batching_opts,
548        ));
549        let dual_engine = Arc::new(FlowDualEngine::new(
550            streaming_engine,
551            batching_engine,
552            flow_meta,
553            catalog_manager,
554            Plugins::new(),
555        ));
556
557        let (report_sender, report_handler) = SizeReportSender::new();
558        dual_engine.set_state_report_handler(report_handler).await;
559
560        let server = FlownodeServer::new(FlowService::new(dual_engine));
561        (server, report_sender)
562    }
563}
flow/server.rs

flow/
server.rs
