Skip to main content

flow/adapter/
flownode_impl.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//! impl `FlowNode` trait for FlowNodeManager so standalone can call them
16use std::collections::{HashMap, HashSet};
17use std::sync::Arc;
18use std::sync::atomic::AtomicBool;
19
20use api::v1::flow::{
21    CreateRequest, DirtyWindowRequests, DropRequest, FlowRequest, FlowResponse, FlushFlow,
22    flow_request,
23};
24use api::v1::region::InsertRequests;
25use catalog::CatalogManager;
26use common_base::Plugins;
27use common_error::ext::BoxedError;
28use common_meta::ddl::create_flow::FlowType;
29use common_meta::error::Result as MetaResult;
30use common_meta::key::flow::FlowMetadataManager;
31use common_meta::key::flow::flow_state::FlowStat;
32use common_runtime::JoinHandle;
33use common_telemetry::{error, info, trace, warn};
34use datatypes::value::Value;
35use futures::TryStreamExt;
36use itertools::Itertools;
37use operator::utils::try_to_session_query_context;
38use session::context::QueryContextBuilder;
39use snafu::{IntoError, OptionExt, ResultExt, ensure};
40use store_api::storage::{RegionId, TableId};
41use tokio::sync::{Mutex, RwLock};
42
43use crate::adapter::{CreateFlowArgs, StreamingEngine};
44use crate::batching_mode::engine::BatchingEngine;
45use crate::engine::{FlowEngine, FlowStatProvider};
46use crate::error::{
47    CreateFlowSnafu, ExternalSnafu, FlowNotFoundSnafu, FlowNotRecoveredSnafu,
48    IllegalCheckTaskStateSnafu, InsertIntoFlowSnafu, InternalSnafu, JoinTaskSnafu, ListFlowsSnafu,
49    NoAvailableFrontendSnafu, SyncCheckTaskSnafu, UnexpectedSnafu, UnsupportedSnafu,
50};
51use crate::metrics::{METRIC_FLOW_ROWS, METRIC_FLOW_TASK_COUNT};
52use crate::repr::{self, DiffRow};
53use crate::utils::StateReportHandler;
54use crate::{Error, FlowId};
55
56/// Ref to [`FlowDualEngine`]
57pub type FlowDualEngineRef = Arc<FlowDualEngine>;
58
59/// Manage both streaming and batching mode engine
60///
61/// including create/drop/flush flow
62/// and redirect insert requests to the appropriate engine
63pub struct FlowDualEngine {
64    streaming_engine: Arc<StreamingEngine>,
65    batching_engine: Arc<BatchingEngine>,
66    /// receive a oneshot sender to send state report
67    state_report_handler: RwLock<Option<StateReportHandler>>,
68    /// helper struct for faster query flow by table id or vice versa
69    src_table2flow: RwLock<SrcTableToFlow>,
70    flow_metadata_manager: Arc<FlowMetadataManager>,
71    catalog_manager: Arc<dyn CatalogManager>,
72    check_task: tokio::sync::Mutex<Option<ConsistentCheckTask>>,
73    plugins: Plugins,
74    done_recovering: AtomicBool,
75}
76
77impl FlowDualEngine {
78    pub fn new(
79        streaming_engine: Arc<StreamingEngine>,
80        batching_engine: Arc<BatchingEngine>,
81        flow_metadata_manager: Arc<FlowMetadataManager>,
82        catalog_manager: Arc<dyn CatalogManager>,
83        plugins: Plugins,
84    ) -> Self {
85        Self {
86            streaming_engine,
87            batching_engine,
88            state_report_handler: Default::default(),
89            src_table2flow: RwLock::new(SrcTableToFlow::default()),
90            flow_metadata_manager,
91            catalog_manager,
92            check_task: Mutex::new(None),
93            plugins,
94            done_recovering: AtomicBool::new(false),
95        }
96    }
97
98    /// Set `done_recovering` to true
99    /// indicate that we are ready to handle requests
100    pub fn set_done_recovering(&self) {
101        info!("FlowDualEngine done recovering");
102        self.done_recovering
103            .store(true, std::sync::atomic::Ordering::Release);
104    }
105
106    /// Check if `done_recovering` is true
107    pub fn is_recover_done(&self) -> bool {
108        self.done_recovering
109            .load(std::sync::atomic::Ordering::Acquire)
110    }
111
112    /// wait for recovering to be done, this will only happen when flownode just started
113    async fn wait_for_all_flow_recover(&self, waiting_req_cnt: usize) -> Result<(), Error> {
114        if self.is_recover_done() {
115            return Ok(());
116        }
117
118        warn!(
119            "FlowDualEngine is not done recovering, {} insert request waiting for recovery",
120            waiting_req_cnt
121        );
122        // wait 3 seconds, check every 1 second
123        // TODO(discord9): make this configurable
124        let mut retry = 0;
125        let max_retry = 3;
126        while retry < max_retry && !self.is_recover_done() {
127            warn!(
128                "FlowDualEngine is not done recovering, retry {} in 1s",
129                retry
130            );
131            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
132            retry += 1;
133        }
134        if retry == max_retry {
135            return FlowNotRecoveredSnafu.fail();
136        } else {
137            info!("FlowDualEngine is done recovering");
138        }
139        // TODO(discord9): also put to centralized logging for flow once it implemented
140        Ok(())
141    }
142
143    pub fn plugins(&self) -> &Plugins {
144        &self.plugins
145    }
146
147    /// Determine if the engine is in distributed mode
148    pub fn is_distributed(&self) -> bool {
149        self.streaming_engine.node_id.is_some()
150    }
151
152    pub fn streaming_engine(&self) -> Arc<StreamingEngine> {
153        self.streaming_engine.clone()
154    }
155
156    pub fn batching_engine(&self) -> Arc<BatchingEngine> {
157        self.batching_engine.clone()
158    }
159
160    pub async fn set_state_report_handler(&self, handler: StateReportHandler) {
161        *self.state_report_handler.write().await = Some(handler);
162    }
163
164    pub async fn gen_state_report(&self) -> FlowStat {
165        let streaming = self.streaming_engine.flow_stat().await;
166        let batching = self.batching_engine.flow_stat().await;
167
168        let mut state_size = streaming.state_size;
169        state_size.extend(batching.state_size);
170
171        let mut last_exec_time_map = streaming.last_exec_time_map;
172        last_exec_time_map.extend(batching.last_exec_time_map);
173
174        FlowStat {
175            state_size,
176            last_exec_time_map,
177        }
178    }
179
180    /// Start state report task, which receives a sender from heartbeat task and sends report back.
181    ///
182    /// if heartbeat task is shutdown, this future exits too.
183    pub async fn start_state_report_task(self: Arc<Self>) -> Option<JoinHandle<()>> {
184        let state_report_handler = self.state_report_handler.write().await.take();
185        if let Some(mut handler) = state_report_handler {
186            let zelf = self.clone();
187            let handler = common_runtime::spawn_global(async move {
188                while let Some(ret_handler) = handler.recv().await {
189                    let state_report = zelf.gen_state_report().await;
190                    ret_handler.send(state_report).unwrap_or_else(|err| {
191                        common_telemetry::error!(err; "Send state report error");
192                    });
193                }
194            });
195            Some(handler)
196        } else {
197            None
198        }
199    }
200
201    /// In distributed mode, scan periodically(1s) until all advertised frontends
202    /// accept unauthenticated queries, or timeout. In standalone mode, return
203    /// immediately.
204    async fn wait_for_available_frontend(&self, timeout: std::time::Duration) -> Result<(), Error> {
205        if !self.is_distributed() {
206            return Ok(());
207        }
208        let frontend_client = self.batching_engine().frontend_client.clone();
209        let sleep_duration = std::time::Duration::from_millis(1_000);
210        let now = std::time::Instant::now();
211        loop {
212            let frontend_list = frontend_client.scan_for_frontend().await?;
213            if !frontend_list.is_empty() {
214                let fe_list = frontend_list
215                    .iter()
216                    .map(|peer| &peer.addr)
217                    .collect::<Vec<_>>();
218                let probe_failures = frontend_client
219                    .check_all_frontends_without_auth(&frontend_list)
220                    .await?;
221                if probe_failures.is_empty() {
222                    info!(
223                        "Available frontend found and unauthenticated probe succeeded: {:?}",
224                        fe_list
225                    );
226                    return Ok(());
227                }
228                warn!(
229                    "Unauthenticated frontend probe failed, will retry. frontends={:?}, failures={:?}",
230                    fe_list, probe_failures
231                );
232            }
233            let elapsed = now.elapsed();
234            tokio::time::sleep(sleep_duration).await;
235            info!("Waiting for available frontend, elapsed={:?}", elapsed);
236            if elapsed >= timeout {
237                return NoAvailableFrontendSnafu {
238                    timeout,
239                    context: "No frontend accepted unauthenticated flownode probe",
240                }
241                .fail();
242            }
243        }
244    }
245
246    /// Try to sync with check task, this is only used in drop flow&flush flow, so a flow id is required
247    ///
248    /// the need to sync is to make sure flush flow actually get called
249    async fn try_sync_with_check_task(
250        &self,
251        flow_id: FlowId,
252        allow_drop: bool,
253    ) -> Result<(), Error> {
254        // this function rarely get called so adding some log is helpful
255        info!("Try to sync with check task for flow {}", flow_id);
256        let mut retry = 0;
257        let max_retry = 10;
258        // keep trying to trigger consistent check
259        while retry < max_retry {
260            if let Some(task) = self.check_task.lock().await.as_ref() {
261                task.trigger(false, allow_drop).await?;
262                break;
263            }
264            retry += 1;
265            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
266        }
267
268        if retry == max_retry {
269            error!(
270                "Can't sync with check task for flow {} with allow_drop={}",
271                flow_id, allow_drop
272            );
273            return SyncCheckTaskSnafu {
274                flow_id,
275                allow_drop,
276            }
277            .fail();
278        }
279        info!("Successfully sync with check task for flow {}", flow_id);
280
281        Ok(())
282    }
283
284    /// Spawn a task to consistently check if all flow tasks in metasrv is created on flownode,
285    /// so on startup, this will create all missing flow tasks, and constantly check at a interval
286    async fn check_flow_consistent(
287        &self,
288        allow_create: bool,
289        allow_drop: bool,
290    ) -> Result<(), Error> {
291        // use nodeid to determine if this is standalone/distributed mode, and retrieve all flows in this node(in distributed mode)/or all flows(in standalone mode)
292        let nodeid = self.streaming_engine.node_id;
293        let should_exists: Vec<_> = if let Some(nodeid) = nodeid {
294            // nodeid is available, so we only need to check flows on this node
295            // which also means we are in distributed mode
296            let to_be_recover = self
297                .flow_metadata_manager
298                .flownode_flow_manager()
299                .flows(nodeid.into())
300                .try_collect::<Vec<_>>()
301                .await
302                .context(ListFlowsSnafu {
303                    id: Some(nodeid.into()),
304                })?;
305            to_be_recover.into_iter().map(|(id, _)| id).collect()
306        } else {
307            // nodeid is not available, so we need to check all flows
308            // which also means we are in standalone mode
309            let all_catalogs = self
310                .catalog_manager
311                .catalog_names()
312                .await
313                .map_err(BoxedError::new)
314                .context(ExternalSnafu)?;
315            let mut all_flow_ids = vec![];
316            for catalog in all_catalogs {
317                let flows = self
318                    .flow_metadata_manager
319                    .flow_name_manager()
320                    .flow_names(&catalog)
321                    .await
322                    .try_collect::<Vec<_>>()
323                    .await
324                    .map_err(BoxedError::new)
325                    .context(ExternalSnafu)?;
326
327                all_flow_ids.extend(flows.into_iter().map(|(_, id)| id.flow_id()));
328            }
329            all_flow_ids
330        };
331        let should_exists = should_exists
332            .into_iter()
333            .map(|i| i as FlowId)
334            .collect::<HashSet<_>>();
335        let actual_exists = self.list_flows().await?.into_iter().collect::<HashSet<_>>();
336        let to_be_created = should_exists
337            .iter()
338            .filter(|id| !actual_exists.contains(id))
339            .collect::<Vec<_>>();
340        let to_be_dropped = actual_exists
341            .iter()
342            .filter(|id| !should_exists.contains(id))
343            .collect::<Vec<_>>();
344
345        if !to_be_created.is_empty() {
346            if allow_create {
347                info!(
348                    "Recovering {} flows: {:?}",
349                    to_be_created.len(),
350                    to_be_created
351                );
352                let mut errors = vec![];
353                for flow_id in to_be_created.clone() {
354                    let flow_id = *flow_id;
355                    let info = self
356                        .flow_metadata_manager
357                        .flow_info_manager()
358                        .get(flow_id as u32)
359                        .await
360                        .map_err(BoxedError::new)
361                        .context(ExternalSnafu)?
362                        .context(FlowNotFoundSnafu { id: flow_id })?;
363
364                    let sink_table_name = [
365                        info.sink_table_name().catalog_name.clone(),
366                        info.sink_table_name().schema_name.clone(),
367                        info.sink_table_name().table_name.clone(),
368                    ];
369                    let args = CreateFlowArgs {
370                        flow_id,
371                        sink_table_name,
372                        source_table_ids: info.source_table_ids().to_vec(),
373                        // because recover should only happen on restart the `create_if_not_exists` and `or_replace` can be arbitrary value(since flow doesn't exist)
374                        // but for the sake of consistency and to make sure recover of flow actually happen, we set both to true
375                        // (which is also fine since checks for not allow both to be true is on metasrv and we already pass that)
376                        create_if_not_exists: true,
377                        or_replace: true,
378                        expire_after: info.expire_after(),
379                        eval_interval: info.eval_interval(),
380                        comment: Some(info.comment().clone()),
381                        sql: info.raw_sql().clone(),
382                        flow_options: info.options().clone(),
383                        query_ctx: info
384                            .query_context()
385                            .clone()
386                            .map(|ctx| {
387                                try_to_session_query_context(ctx)
388                                    .map_err(BoxedError::new)
389                                    .context(ExternalSnafu)
390                            })
391                            .transpose()?
392                            // or use default QueryContext with catalog_name from info
393                            // to keep compatibility with old version
394                            .or_else(|| {
395                                Some(
396                                    QueryContextBuilder::default()
397                                        .current_catalog(info.catalog_name().clone())
398                                        .build(),
399                                )
400                            }),
401                    };
402                    if let Err(err) = self
403                        .create_flow(args)
404                        .await
405                        .map_err(BoxedError::new)
406                        .with_context(|_| CreateFlowSnafu {
407                            sql: info.raw_sql().clone(),
408                        })
409                    {
410                        errors.push((flow_id, err));
411                    }
412                }
413                if errors.is_empty() {
414                    info!("Recover flows successfully, flows: {:?}", to_be_created);
415                }
416
417                for (flow_id, err) in errors {
418                    warn!("Failed to recreate flow {}, err={:#?}", flow_id, err);
419                }
420            } else {
421                warn!(
422                    "Flows do not exist in flownode for node {:?}, flow_ids={:?}",
423                    nodeid, to_be_created
424                );
425            }
426        }
427        if !to_be_dropped.is_empty() {
428            if allow_drop {
429                info!("Dropping flows: {:?}", to_be_dropped);
430                let mut errors = vec![];
431                for flow_id in to_be_dropped {
432                    let flow_id = *flow_id;
433                    if let Err(err) = self.remove_flow(flow_id).await {
434                        errors.push((flow_id, err));
435                    }
436                }
437                for (flow_id, err) in errors {
438                    warn!("Failed to drop flow {}, err={:#?}", flow_id, err);
439                }
440            } else {
441                warn!(
442                    "Flows do not exist in metadata for node {:?}, flow_ids={:?}",
443                    nodeid, to_be_dropped
444                );
445            }
446        }
447        Ok(())
448    }
449
450    // TODO(discord9): consider sync this with heartbeat(might become necessary in the future)
451    pub async fn start_flow_consistent_check_task(self: &Arc<Self>) -> Result<(), Error> {
452        let mut check_task = self.check_task.lock().await;
453        ensure!(
454            check_task.is_none(),
455            IllegalCheckTaskStateSnafu {
456                reason: "Flow consistent check task already exists",
457            }
458        );
459        let task = ConsistentCheckTask::start_check_task(self).await?;
460        *check_task = Some(task);
461        Ok(())
462    }
463
464    pub async fn stop_flow_consistent_check_task(&self) -> Result<(), Error> {
465        info!("Stopping flow consistent check task");
466        let mut check_task = self.check_task.lock().await;
467
468        ensure!(
469            check_task.is_some(),
470            IllegalCheckTaskStateSnafu {
471                reason: "Flow consistent check task does not exist",
472            }
473        );
474
475        check_task.take().unwrap().stop().await?;
476        info!("Stopped flow consistent check task");
477        Ok(())
478    }
479
480    /// Reconciles in-memory flow tasks from persisted metadata.
481    pub async fn reconcile_flows_from_metadata(&self) -> Result<(), Error> {
482        self.check_flow_consistent(true, true).await
483    }
484
485    /// TODO(discord9): also add a `exists` api using flow metadata manager's `exists` method
486    async fn flow_exist_in_metadata(&self, flow_id: FlowId) -> Result<bool, Error> {
487        self.flow_metadata_manager
488            .flow_info_manager()
489            .get(flow_id as u32)
490            .await
491            .map_err(BoxedError::new)
492            .context(ExternalSnafu)
493            .map(|info| info.is_some())
494    }
495}
496
497struct ConsistentCheckTask {
498    handle: JoinHandle<()>,
499    shutdown_tx: tokio::sync::mpsc::Sender<()>,
500    trigger_tx: tokio::sync::mpsc::Sender<(bool, bool, tokio::sync::oneshot::Sender<()>)>,
501}
502
503impl ConsistentCheckTask {
504    async fn start_check_task(engine: &Arc<FlowDualEngine>) -> Result<Self, Error> {
505        let engine = engine.clone();
506        let min_refresh_duration = engine
507            .batching_engine()
508            .batch_opts
509            .experimental_min_refresh_duration;
510        let frontend_scan_timeout = engine
511            .batching_engine()
512            .batch_opts
513            .experimental_frontend_scan_timeout;
514        engine
515            .wait_for_available_frontend(frontend_scan_timeout)
516            .await?;
517        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
518        let (trigger_tx, mut trigger_rx) =
519            tokio::sync::mpsc::channel::<(bool, bool, tokio::sync::oneshot::Sender<()>)>(10);
520        let handle = common_runtime::spawn_global(async move {
521            // Recover flows after the startup frontend probe succeeds.
522            let mut recover_retry = 0;
523            while let Err(err) = engine.check_flow_consistent(true, false).await {
524                recover_retry += 1;
525                error!(
526                    "Failed to recover flows:\n {err:?}, retry {} in {}s",
527                    recover_retry,
528                    min_refresh_duration.as_secs()
529                );
530                tokio::time::sleep(min_refresh_duration).await;
531            }
532
533            engine.set_done_recovering();
534
535            // then do check flows, with configurable allow_create and allow_drop
536            let (mut allow_create, mut allow_drop) = (false, false);
537            let mut ret_signal: Option<tokio::sync::oneshot::Sender<()>> = None;
538            loop {
539                if let Err(err) = engine.check_flow_consistent(allow_create, allow_drop).await {
540                    error!(err; "Failed to check flow consistent");
541                }
542                if let Some(done) = ret_signal.take() {
543                    let _ = done.send(());
544                }
545                tokio::select! {
546                    _ = rx.recv() => break,
547                    incoming = trigger_rx.recv() => if let Some(incoming) = incoming {
548                        (allow_create, allow_drop) = (incoming.0, incoming.1);
549                        ret_signal = Some(incoming.2);
550                    },
551                    _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
552                        (allow_create, allow_drop) = (false, false);
553                    },
554                }
555            }
556        });
557        Ok(ConsistentCheckTask {
558            handle,
559            shutdown_tx: tx,
560            trigger_tx,
561        })
562    }
563
564    async fn trigger(&self, allow_create: bool, allow_drop: bool) -> Result<(), Error> {
565        let (tx, rx) = tokio::sync::oneshot::channel();
566        self.trigger_tx
567            .send((allow_create, allow_drop, tx))
568            .await
569            .map_err(|_| {
570                IllegalCheckTaskStateSnafu {
571                    reason: "Failed to send trigger signal",
572                }
573                .build()
574            })?;
575        rx.await.map_err(|_| {
576            IllegalCheckTaskStateSnafu {
577                reason: "Failed to receive trigger signal",
578            }
579            .build()
580        })?;
581        Ok(())
582    }
583
584    async fn stop(self) -> Result<(), Error> {
585        self.shutdown_tx.send(()).await.map_err(|_| {
586            IllegalCheckTaskStateSnafu {
587                reason: "Failed to send shutdown signal",
588            }
589            .build()
590        })?;
591        // abort so no need to wait
592        self.handle.abort();
593        Ok(())
594    }
595}
596
597#[derive(Default)]
598struct SrcTableToFlow {
599    /// mapping of table ids to flow ids for streaming mode
600    stream: HashMap<TableId, HashSet<FlowId>>,
601    /// mapping of table ids to flow ids for batching mode
602    batch: HashMap<TableId, HashSet<FlowId>>,
603    /// mapping of flow ids to (flow type, source table ids)
604    flow_infos: HashMap<FlowId, (FlowType, Vec<TableId>)>,
605}
606
607impl SrcTableToFlow {
608    fn in_stream(&self, table_id: TableId) -> bool {
609        self.stream.contains_key(&table_id)
610    }
611    fn in_batch(&self, table_id: TableId) -> bool {
612        self.batch.contains_key(&table_id)
613    }
614    fn add_flow(&mut self, flow_id: FlowId, flow_type: FlowType, src_table_ids: Vec<TableId>) {
615        let mapping = match flow_type {
616            FlowType::Streaming => &mut self.stream,
617            FlowType::Batching => &mut self.batch,
618        };
619
620        for src_table in src_table_ids.clone() {
621            mapping
622                .entry(src_table)
623                .and_modify(|flows| {
624                    flows.insert(flow_id);
625                })
626                .or_insert_with(|| {
627                    let mut set = HashSet::new();
628                    set.insert(flow_id);
629                    set
630                });
631        }
632        self.flow_infos.insert(flow_id, (flow_type, src_table_ids));
633    }
634
635    fn remove_flow(&mut self, flow_id: FlowId) {
636        let mapping = match self.get_flow_type(flow_id) {
637            Some(FlowType::Streaming) => &mut self.stream,
638            Some(FlowType::Batching) => &mut self.batch,
639            None => return,
640        };
641        if let Some((_, src_table_ids)) = self.flow_infos.remove(&flow_id) {
642            for src_table in src_table_ids {
643                if let Some(flows) = mapping.get_mut(&src_table) {
644                    flows.remove(&flow_id);
645                }
646            }
647        }
648    }
649
650    fn get_flow_type(&self, flow_id: FlowId) -> Option<FlowType> {
651        self.flow_infos
652            .get(&flow_id)
653            .map(|(flow_type, _)| flow_type)
654            .cloned()
655    }
656}
657
658impl FlowEngine for FlowDualEngine {
659    async fn create_flow(&self, args: CreateFlowArgs) -> Result<Option<FlowId>, Error> {
660        let flow_type = args
661            .flow_options
662            .get(FlowType::FLOW_TYPE_KEY)
663            .map(|s| s.as_str());
664
665        let flow_type = match flow_type {
666            Some(FlowType::BATCHING) => FlowType::Batching,
667            Some(FlowType::STREAMING) => FlowType::Streaming,
668            None => FlowType::Batching,
669            Some(flow_type) => {
670                return InternalSnafu {
671                    reason: format!("Invalid flow type: {}", flow_type),
672                }
673                .fail();
674            }
675        };
676
677        let flow_id = args.flow_id;
678        let src_table_ids = args.source_table_ids.clone();
679
680        let res = match flow_type {
681            FlowType::Batching => self.batching_engine.create_flow(args).await,
682            FlowType::Streaming => self.streaming_engine.create_flow(args).await,
683        }?;
684
685        self.src_table2flow
686            .write()
687            .await
688            .add_flow(flow_id, flow_type, src_table_ids);
689
690        Ok(res)
691    }
692
693    async fn remove_flow(&self, flow_id: FlowId) -> Result<(), Error> {
694        let flow_type = self.src_table2flow.read().await.get_flow_type(flow_id);
695
696        match flow_type {
697            Some(FlowType::Batching) => self.batching_engine.remove_flow(flow_id).await,
698            Some(FlowType::Streaming) => self.streaming_engine.remove_flow(flow_id).await,
699            None => {
700                // this can happen if flownode just restart, and is stilling creating the flow
701                // since now that this flow should dropped, we need to trigger the consistent check and allow drop
702                // this rely on drop flow ddl delete metadata first, see src/common/meta/src/ddl/drop_flow.rs
703                warn!(
704                    "Flow {} is not exist in the underlying engine, but exist in metadata",
705                    flow_id
706                );
707                self.try_sync_with_check_task(flow_id, true).await?;
708
709                Ok(())
710            }
711        }?;
712        // remove mapping
713        self.src_table2flow.write().await.remove_flow(flow_id);
714        Ok(())
715    }
716
717    async fn flush_flow(&self, flow_id: FlowId) -> Result<usize, Error> {
718        // sync with check task
719        self.try_sync_with_check_task(flow_id, false).await?;
720        let flow_type = self.src_table2flow.read().await.get_flow_type(flow_id);
721        match flow_type {
722            Some(FlowType::Batching) => self.batching_engine.flush_flow(flow_id).await,
723            Some(FlowType::Streaming) => self.streaming_engine.flush_flow(flow_id).await,
724            None => {
725                warn!(
726                    "Currently flow={flow_id} doesn't exist in flownode, ignore flush_flow request"
727                );
728                Ok(0)
729            }
730        }
731    }
732
733    async fn flow_exist(&self, flow_id: FlowId) -> Result<bool, Error> {
734        let flow_type = self.src_table2flow.read().await.get_flow_type(flow_id);
735        // not using `flow_type.is_some()` to make sure the flow is actually exist in the underlying engine
736        match flow_type {
737            Some(FlowType::Batching) => self.batching_engine.flow_exist(flow_id).await,
738            Some(FlowType::Streaming) => self.streaming_engine.flow_exist(flow_id).await,
739            None => Ok(false),
740        }
741    }
742
743    async fn list_flows(&self) -> Result<impl IntoIterator<Item = FlowId>, Error> {
744        let stream_flows = self.streaming_engine.list_flows().await?;
745        let batch_flows = self.batching_engine.list_flows().await?;
746
747        Ok(stream_flows.into_iter().chain(batch_flows))
748    }
749
750    async fn handle_flow_inserts(
751        &self,
752        request: api::v1::region::InsertRequests,
753    ) -> Result<(), Error> {
754        self.wait_for_all_flow_recover(request.requests.len())
755            .await?;
756        // TODO(discord9): make as little clone as possible
757        let mut to_stream_engine = Vec::with_capacity(request.requests.len());
758        let mut to_batch_engine = request.requests;
759
760        let mut batching_row_cnt = 0;
761        let mut streaming_row_cnt = 0;
762
763        {
764            // not locking this, or recover flows will be starved when also handling flow inserts
765            let src_table2flow = self.src_table2flow.read().await;
766            to_batch_engine.retain(|req| {
767                let region_id = RegionId::from(req.region_id);
768                let table_id = region_id.table_id();
769                let is_in_stream = src_table2flow.in_stream(table_id);
770                let is_in_batch = src_table2flow.in_batch(table_id);
771                if is_in_stream {
772                    streaming_row_cnt += req.rows.as_ref().map(|rs| rs.rows.len()).unwrap_or(0);
773                    to_stream_engine.push(req.clone());
774                }
775                if is_in_batch {
776                    batching_row_cnt += req.rows.as_ref().map(|rs| rs.rows.len()).unwrap_or(0);
777                    return true;
778                }
779                if !is_in_batch && !is_in_stream {
780                    // TODO(discord9): also put to centralized logging for flow once it implemented
781                    warn!("Table {} is not any flow's source table", table_id)
782                }
783                false
784            });
785            // drop(src_table2flow);
786            // can't use drop due to https://github.com/rust-lang/rust/pull/128846
787        }
788
789        METRIC_FLOW_ROWS
790            .with_label_values(&["in-streaming"])
791            .inc_by(streaming_row_cnt as u64);
792
793        METRIC_FLOW_ROWS
794            .with_label_values(&["in-batching"])
795            .inc_by(batching_row_cnt as u64);
796
797        let streaming_engine = self.streaming_engine.clone();
798        let stream_handler: JoinHandle<Result<(), Error>> =
799            common_runtime::spawn_global(async move {
800                streaming_engine
801                    .handle_flow_inserts(api::v1::region::InsertRequests {
802                        requests: to_stream_engine,
803                    })
804                    .await?;
805                Ok(())
806            });
807        self.batching_engine
808            .handle_flow_inserts(api::v1::region::InsertRequests {
809                requests: to_batch_engine,
810            })
811            .await?;
812        stream_handler.await.context(JoinTaskSnafu)??;
813
814        Ok(())
815    }
816
817    async fn handle_mark_window_dirty(
818        &self,
819        req: api::v1::flow::DirtyWindowRequests,
820    ) -> Result<(), Error> {
821        self.batching_engine.handle_mark_window_dirty(req).await
822    }
823}
824
825#[async_trait::async_trait]
826impl common_meta::node_manager::Flownode for FlowDualEngine {
827    async fn handle(&self, request: FlowRequest) -> MetaResult<FlowResponse> {
828        let query_ctx = request
829            .header
830            .and_then(|h| h.query_context)
831            .map(|ctx| ctx.into());
832        match request.body {
833            Some(flow_request::Body::Create(CreateRequest {
834                flow_id: Some(task_id),
835                source_table_ids,
836                sink_table_name: Some(sink_table_name),
837                create_if_not_exists,
838                expire_after,
839                eval_interval,
840                comment,
841                sql,
842                flow_options,
843                or_replace,
844            })) => {
845                let source_table_ids = source_table_ids.into_iter().map(|id| id.id).collect_vec();
846                let sink_table_name = [
847                    sink_table_name.catalog_name,
848                    sink_table_name.schema_name,
849                    sink_table_name.table_name,
850                ];
851                let expire_after = expire_after.map(|e| e.value);
852                let args = CreateFlowArgs {
853                    flow_id: task_id.id as u64,
854                    sink_table_name,
855                    source_table_ids,
856                    create_if_not_exists,
857                    or_replace,
858                    expire_after,
859                    eval_interval: eval_interval.map(|e| e.seconds),
860                    comment: Some(comment),
861                    sql: sql.clone(),
862                    flow_options,
863                    query_ctx,
864                };
865                let ret = self
866                    .create_flow(args)
867                    .await
868                    .map_err(BoxedError::new)
869                    .with_context(|_| CreateFlowSnafu { sql: sql.clone() })
870                    .map_err(to_meta_err(snafu::location!()))?;
871                METRIC_FLOW_TASK_COUNT.inc();
872                Ok(FlowResponse {
873                    affected_flows: ret
874                        .map(|id| greptime_proto::v1::FlowId { id: id as u32 })
875                        .into_iter()
876                        .collect_vec(),
877                    ..Default::default()
878                })
879            }
880            Some(flow_request::Body::Drop(DropRequest {
881                flow_id: Some(flow_id),
882            })) => {
883                self.remove_flow(flow_id.id as u64)
884                    .await
885                    .map_err(to_meta_err(snafu::location!()))?;
886                METRIC_FLOW_TASK_COUNT.dec();
887                Ok(Default::default())
888            }
889            Some(flow_request::Body::Flush(FlushFlow {
890                flow_id: Some(flow_id),
891            })) => {
892                let row = self
893                    .flush_flow(flow_id.id as u64)
894                    .await
895                    .map_err(to_meta_err(snafu::location!()))?;
896                Ok(FlowResponse {
897                    affected_flows: vec![flow_id],
898                    affected_rows: row as u64,
899                    ..Default::default()
900                })
901            }
902            other => common_meta::error::InvalidFlowRequestBodySnafu { body: other }.fail(),
903        }
904    }
905
906    async fn handle_inserts(&self, request: InsertRequests) -> MetaResult<FlowResponse> {
907        FlowEngine::handle_flow_inserts(self, request)
908            .await
909            .map(|_| Default::default())
910            .map_err(to_meta_err(snafu::location!()))
911    }
912
913    async fn handle_mark_window_dirty(&self, req: DirtyWindowRequests) -> MetaResult<FlowResponse> {
914        self.batching_engine()
915            .handle_mark_dirty_time_window(req)
916            .await
917            .map(|_| FlowResponse::default())
918            .map_err(to_meta_err(snafu::location!()))
919    }
920}
921
922/// return a function to convert `crate::error::Error` to `common_meta::error::Error`
923fn to_meta_err(
924    location: snafu::Location,
925) -> impl FnOnce(crate::error::Error) -> common_meta::error::Error {
926    move |err: crate::error::Error| -> common_meta::error::Error {
927        match err {
928            crate::error::Error::FlowNotFound { id, .. } => {
929                common_meta::error::Error::FlowNotFound {
930                    flow_name: format!("flow_id={id}"),
931                    location,
932                }
933            }
934            _ => common_meta::error::Error::External {
935                location,
936                source: BoxedError::new(err),
937            },
938        }
939    }
940}
941
942impl FlowEngine for StreamingEngine {
943    async fn create_flow(&self, args: CreateFlowArgs) -> Result<Option<FlowId>, Error> {
944        self.create_flow_inner(args).await
945    }
946
947    async fn remove_flow(&self, flow_id: FlowId) -> Result<(), Error> {
948        self.remove_flow_inner(flow_id).await
949    }
950
951    async fn flush_flow(&self, flow_id: FlowId) -> Result<usize, Error> {
952        self.flush_flow_inner(flow_id).await
953    }
954
955    async fn flow_exist(&self, flow_id: FlowId) -> Result<bool, Error> {
956        self.flow_exist_inner(flow_id).await
957    }
958
959    async fn list_flows(&self) -> Result<impl IntoIterator<Item = FlowId>, Error> {
960        Ok(self
961            .flow_err_collectors
962            .read()
963            .await
964            .keys()
965            .cloned()
966            .collect::<Vec<_>>())
967    }
968
969    async fn handle_flow_inserts(
970        &self,
971        request: api::v1::region::InsertRequests,
972    ) -> Result<(), Error> {
973        self.handle_inserts_inner(request).await
974    }
975
976    async fn handle_mark_window_dirty(
977        &self,
978        _req: api::v1::flow::DirtyWindowRequests,
979    ) -> Result<(), Error> {
980        UnsupportedSnafu {
981            reason: "handle_mark_window_dirty in streaming engine",
982        }
983        .fail()
984    }
985}
986
987/// Simple helper enum for fetching value from row with default value
988#[derive(Debug, Clone)]
989enum FetchFromRow {
990    Idx(usize),
991    Default(Value),
992}
993
994impl FetchFromRow {
995    /// Panic if idx is out of bound
996    fn fetch(&self, row: &repr::Row) -> Value {
997        match self {
998            FetchFromRow::Idx(idx) => row.get(*idx).unwrap().clone(),
999            FetchFromRow::Default(v) => v.clone(),
1000        }
1001    }
1002}
1003
1004impl StreamingEngine {
1005    async fn handle_inserts_inner(
1006        &self,
1007        request: InsertRequests,
1008    ) -> std::result::Result<(), Error> {
1009        // using try_read to ensure two things:
1010        // 1. flush wouldn't happen until inserts before it is inserted
1011        // 2. inserts happening concurrently with flush wouldn't be block by flush
1012        let _flush_lock = self.flush_lock.try_read();
1013        for write_request in request.requests {
1014            let region_id = write_request.region_id;
1015            let table_id = RegionId::from(region_id).table_id();
1016
1017            let (insert_schema, rows_proto) = write_request
1018                .rows
1019                .map(|r| (r.schema, r.rows))
1020                .unwrap_or_default();
1021
1022            // TODO(discord9): reconsider time assignment mechanism
1023            let now = self.tick_manager.tick();
1024
1025            let (table_types, fetch_order) = {
1026                let ctx = self.node_context.read().await;
1027
1028                // TODO(discord9): also check schema version so that altered table can be reported
1029                let table_schema = ctx.table_source.table_from_id(&table_id).await?;
1030                let default_vals = table_schema
1031                    .default_values
1032                    .iter()
1033                    .zip(table_schema.relation_desc.typ().column_types.iter())
1034                    .map(|(v, ty)| {
1035                        v.as_ref().and_then(|v| {
1036                            match v.create_default(ty.scalar_type(), ty.nullable()) {
1037                                Ok(v) => Some(v),
1038                                Err(err) => {
1039                                    common_telemetry::error!(err; "Failed to create default value");
1040                                    None
1041                                }
1042                            }
1043                        })
1044                    })
1045                    .collect_vec();
1046
1047                let table_types = table_schema
1048                    .relation_desc
1049                    .typ()
1050                    .column_types
1051                    .clone()
1052                    .into_iter()
1053                    .map(|t| t.scalar_type)
1054                    .collect_vec();
1055                let table_col_names = table_schema.relation_desc.names;
1056                let table_col_names = table_col_names
1057                    .iter().enumerate()
1058                    .map(|(idx,name)| match name {
1059                        Some(name) => Ok(name.clone()),
1060                        None => InternalSnafu {
1061                            reason: format!("Expect column {idx} of table id={table_id} to have name in table schema, found None"),
1062                        }
1063                        .fail(),
1064                    })
1065                    .collect::<Result<Vec<_>, _>>()?;
1066                let name_to_col = HashMap::<_, _>::from_iter(
1067                    insert_schema
1068                        .iter()
1069                        .enumerate()
1070                        .map(|(i, name)| (&name.column_name, i)),
1071                );
1072
1073                let fetch_order: Vec<FetchFromRow> = table_col_names
1074                    .iter()
1075                    .zip(default_vals)
1076                    .map(|(col_name, col_default_val)| {
1077                        name_to_col
1078                            .get(col_name)
1079                            .copied()
1080                            .map(FetchFromRow::Idx)
1081                            .or_else(|| col_default_val.clone().map(FetchFromRow::Default))
1082                            .with_context(|| UnexpectedSnafu {
1083                                reason: format!(
1084                                    "Column not found: {}, default_value: {:?}",
1085                                    col_name, col_default_val
1086                                ),
1087                            })
1088                    })
1089                    .try_collect()?;
1090
1091                trace!("Reordering columns: {:?}", fetch_order);
1092                (table_types, fetch_order)
1093            };
1094
1095            // TODO(discord9): use column instead of row
1096            let rows: Vec<DiffRow> = rows_proto
1097                .into_iter()
1098                .map(|r| {
1099                    let r = repr::Row::from(r);
1100                    let reordered = fetch_order.iter().map(|i| i.fetch(&r)).collect_vec();
1101                    repr::Row::new(reordered)
1102                })
1103                .map(|r| (r, now, 1))
1104                .collect_vec();
1105            if let Err(err) = self
1106                .handle_write_request(region_id.into(), rows, &table_types)
1107                .await
1108            {
1109                let err = BoxedError::new(err);
1110                let flow_ids = self
1111                    .node_context
1112                    .read()
1113                    .await
1114                    .get_flow_ids(table_id)
1115                    .into_iter()
1116                    .flatten()
1117                    .cloned()
1118                    .collect_vec();
1119                let err = InsertIntoFlowSnafu {
1120                    region_id,
1121                    flow_ids,
1122                }
1123                .into_error(err);
1124                common_telemetry::error!(err; "Failed to handle write request");
1125                return Err(err);
1126            }
1127        }
1128        Ok(())
1129    }
1130}