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