flow/
adapter.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//! for getting data from source and sending results to sink
16//! and communicating with other parts of the database
17#![warn(unused_imports)]
18
19use std::collections::BTreeMap;
20use std::sync::Arc;
21use std::time::{Duration, Instant, SystemTime};
22
23use api::v1::{RowDeleteRequest, RowDeleteRequests, RowInsertRequest, RowInsertRequests};
24use common_base::memory_limit::MemoryLimit;
25use common_config::Configurable;
26use common_error::ext::BoxedError;
27use common_meta::key::TableMetadataManagerRef;
28use common_options::memory::MemoryOptions;
29use common_runtime::JoinHandle;
30use common_stat::get_total_cpu_cores;
31use common_telemetry::logging::{LoggingOptions, TracingOptions};
32use common_telemetry::{debug, info, trace};
33use datatypes::schema::ColumnSchema;
34use datatypes::value::Value;
35use greptime_proto::v1;
36use itertools::{EitherOrBoth, Itertools};
37use meta_client::MetaClientOptions;
38use query::QueryEngine;
39use query::options::QueryOptions;
40use serde::{Deserialize, Serialize};
41use servers::grpc::GrpcOptions;
42use servers::http::HttpOptions;
43use session::context::QueryContext;
44use snafu::{OptionExt, ResultExt, ensure};
45use store_api::storage::{ConcreteDataType, RegionId};
46use table::metadata::TableId;
47use tokio::sync::broadcast::error::TryRecvError;
48use tokio::sync::{Mutex, RwLock, broadcast, watch};
49
50pub(crate) use crate::adapter::node_context::FlownodeContext;
51use crate::adapter::refill::RefillTask;
52use crate::adapter::table_source::ManagedTableSource;
53use crate::adapter::util::relation_desc_to_column_schemas_with_fallback;
54pub(crate) use crate::adapter::worker::{Worker, WorkerHandle, create_worker};
55use crate::batching_mode::BatchingModeOptions;
56use crate::compute::ErrCollector;
57use crate::df_optimizer::sql_to_flow_plan;
58use crate::error::{EvalSnafu, ExternalSnafu, InternalSnafu, InvalidQuerySnafu, UnexpectedSnafu};
59use crate::expr::Batch;
60use crate::metrics::{METRIC_FLOW_INSERT_ELAPSED, METRIC_FLOW_ROWS, METRIC_FLOW_RUN_INTERVAL_MS};
61use crate::repr::{self, BATCH_SIZE, DiffRow, RelationDesc, Row};
62use crate::{CreateFlowArgs, FlowId, TableName};
63
64pub(crate) mod flownode_impl;
65mod parse_expr;
66pub(crate) mod refill;
67mod stat;
68#[cfg(test)]
69mod tests;
70pub(crate) mod util;
71mod worker;
72
73pub(crate) mod node_context;
74pub(crate) mod table_source;
75
76use crate::FrontendInvoker;
77use crate::error::Error;
78
79// `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow
80pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder";
81
82pub const AUTO_CREATED_UPDATE_AT_TS_COL: &str = "update_at";
83
84/// Flow config that exists both in standalone&distributed mode
85#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
86#[serde(default)]
87pub struct FlowConfig {
88    pub num_workers: usize,
89    pub batching_mode: BatchingModeOptions,
90}
91
92impl Default for FlowConfig {
93    fn default() -> Self {
94        Self {
95            num_workers: (get_total_cpu_cores() / 2).max(1),
96            batching_mode: BatchingModeOptions::default(),
97        }
98    }
99}
100
101/// Options for flow node
102#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
103#[serde(default)]
104pub struct FlownodeOptions {
105    pub node_id: Option<u64>,
106    pub flow: FlowConfig,
107    pub grpc: GrpcOptions,
108    pub http: HttpOptions,
109    pub meta_client: Option<MetaClientOptions>,
110    pub logging: LoggingOptions,
111    pub tracing: TracingOptions,
112    pub query: QueryOptions,
113    pub memory: MemoryOptions,
114}
115
116impl Default for FlownodeOptions {
117    fn default() -> Self {
118        Self {
119            node_id: None,
120            flow: FlowConfig::default(),
121            grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:3004"),
122            http: HttpOptions::default(),
123            meta_client: None,
124            logging: LoggingOptions::default(),
125            tracing: TracingOptions::default(),
126            // flownode's query option is set to 1 to throttle flow's query so
127            // that it won't use too much cpu or memory
128            query: QueryOptions {
129                parallelism: 1,
130                allow_query_fallback: false,
131                memory_pool_size: MemoryLimit::default(),
132            },
133            memory: MemoryOptions::default(),
134        }
135    }
136}
137
138impl Configurable for FlownodeOptions {
139    fn validate_sanitize(&mut self) -> common_config::error::Result<()> {
140        if self.flow.num_workers == 0 {
141            self.flow.num_workers = (get_total_cpu_cores() / 2).max(1);
142        }
143        Ok(())
144    }
145}
146
147/// Arc-ed FlowNodeManager, cheaper to clone
148pub type FlowStreamingEngineRef = Arc<StreamingEngine>;
149
150/// FlowNodeManager manages the state of all tasks in the flow node, which should be run on the same thread
151///
152/// The choice of timestamp is just using current system timestamp for now
153///
154pub struct StreamingEngine {
155    /// The handler to the worker that will run the dataflow
156    /// which is `!Send` so a handle is used
157    pub worker_handles: Vec<WorkerHandle>,
158    /// The selector to select a worker to run the dataflow
159    worker_selector: Mutex<usize>,
160    /// The query engine that will be used to parse the query and convert it to a dataflow plan
161    pub query_engine: Arc<dyn QueryEngine>,
162    /// Getting table name and table schema from table info manager
163    table_info_source: ManagedTableSource,
164    frontend_invoker: RwLock<Option<FrontendInvoker>>,
165    /// contains mapping from table name to global id, and table schema
166    node_context: RwLock<FlownodeContext>,
167    /// Contains all refill tasks
168    refill_tasks: RwLock<BTreeMap<FlowId, RefillTask>>,
169    flow_err_collectors: RwLock<BTreeMap<FlowId, ErrCollector>>,
170    src_send_buf_lens: RwLock<BTreeMap<TableId, watch::Receiver<usize>>>,
171    tick_manager: FlowTickManager,
172    /// This node id is only available in distributed mode, on standalone mode this is guaranteed to be `None`
173    pub node_id: Option<u32>,
174    /// Lock for flushing, will be `read` by `handle_inserts` and `write` by `flush_flow`
175    ///
176    /// So that a series of event like `inserts -> flush` can be handled correctly
177    flush_lock: RwLock<()>,
178}
179
180/// Building FlownodeManager
181impl StreamingEngine {
182    /// set frontend invoker
183    pub async fn set_frontend_invoker(&self, frontend: FrontendInvoker) {
184        *self.frontend_invoker.write().await = Some(frontend);
185    }
186
187    /// Create **without** setting `frontend_invoker`
188    pub fn new(
189        node_id: Option<u32>,
190        query_engine: Arc<dyn QueryEngine>,
191        table_meta: TableMetadataManagerRef,
192    ) -> Self {
193        let srv_map = ManagedTableSource::new(
194            table_meta.table_info_manager().clone(),
195            table_meta.table_name_manager().clone(),
196        );
197        let node_context = FlownodeContext::new(Box::new(srv_map.clone()) as _);
198        let tick_manager = FlowTickManager::new();
199        let worker_handles = Vec::new();
200        StreamingEngine {
201            worker_handles,
202            worker_selector: Mutex::new(0),
203            query_engine,
204            table_info_source: srv_map,
205            frontend_invoker: RwLock::new(None),
206            node_context: RwLock::new(node_context),
207            refill_tasks: Default::default(),
208            flow_err_collectors: Default::default(),
209            src_send_buf_lens: Default::default(),
210            tick_manager,
211            node_id,
212            flush_lock: RwLock::new(()),
213        }
214    }
215
216    /// Create a flownode manager with one worker
217    pub fn new_with_workers<'s>(
218        node_id: Option<u32>,
219        query_engine: Arc<dyn QueryEngine>,
220        table_meta: TableMetadataManagerRef,
221        num_workers: usize,
222    ) -> (Self, Vec<Worker<'s>>) {
223        let mut zelf = Self::new(node_id, query_engine, table_meta);
224
225        let workers: Vec<_> = (0..num_workers)
226            .map(|_| {
227                let (handle, worker) = create_worker();
228                zelf.add_worker_handle(handle);
229                worker
230            })
231            .collect();
232        (zelf, workers)
233    }
234
235    /// add a worker handler to manager, meaning this corresponding worker is under it's manage
236    pub fn add_worker_handle(&mut self, handle: WorkerHandle) {
237        self.worker_handles.push(handle);
238    }
239}
240
241#[derive(Debug)]
242pub enum DiffRequest {
243    Insert(Vec<(Row, repr::Timestamp)>),
244    Delete(Vec<(Row, repr::Timestamp)>),
245}
246
247impl DiffRequest {
248    pub fn len(&self) -> usize {
249        match self {
250            Self::Insert(v) => v.len(),
251            Self::Delete(v) => v.len(),
252        }
253    }
254
255    pub fn is_empty(&self) -> bool {
256        self.len() == 0
257    }
258}
259
260pub fn batches_to_rows_req(batches: Vec<Batch>) -> Result<Vec<DiffRequest>, Error> {
261    let mut reqs = Vec::new();
262    for batch in batches {
263        let mut rows = Vec::with_capacity(batch.row_count());
264        for i in 0..batch.row_count() {
265            let row = batch.get_row(i).context(EvalSnafu)?;
266            rows.push((Row::new(row), 0));
267        }
268        reqs.push(DiffRequest::Insert(rows));
269    }
270    Ok(reqs)
271}
272
273/// This impl block contains methods to send writeback requests to frontend
274impl StreamingEngine {
275    /// Return the number of requests it made
276    pub async fn send_writeback_requests(&self) -> Result<usize, Error> {
277        let all_reqs = self.generate_writeback_request().await?;
278        if all_reqs.is_empty() || all_reqs.iter().all(|v| v.1.is_empty()) {
279            return Ok(0);
280        }
281        let mut req_cnt = 0;
282        for (table_name, reqs) in all_reqs {
283            if reqs.is_empty() {
284                continue;
285            }
286            let (catalog, schema) = (table_name[0].clone(), table_name[1].clone());
287            let ctx = Arc::new(QueryContext::with(&catalog, &schema));
288
289            let (is_ts_placeholder, proto_schema) = match self
290                .try_fetch_existing_table(&table_name)
291                .await?
292                .context(UnexpectedSnafu {
293                    reason: format!("Table not found: {}", table_name.join(".")),
294                }) {
295                Ok(r) => r,
296                Err(e) => {
297                    if self
298                        .table_info_source
299                        .get_opt_table_id_from_name(&table_name)
300                        .await?
301                        .is_none()
302                    {
303                        // deal with both flow&sink table no longer exists
304                        // but some output is still in output buf
305                        common_telemetry::warn!(e; "Table `{}` no longer exists, skip writeback", table_name.join("."));
306                        continue;
307                    } else {
308                        return Err(e);
309                    }
310                }
311            };
312            let schema_len = proto_schema.len();
313
314            let total_rows = reqs.iter().map(|r| r.len()).sum::<usize>();
315            trace!(
316                "Sending {} writeback requests to table {}, reqs total rows={}",
317                reqs.len(),
318                table_name.join("."),
319                reqs.iter().map(|r| r.len()).sum::<usize>()
320            );
321
322            METRIC_FLOW_ROWS
323                .with_label_values(&["out-streaming"])
324                .inc_by(total_rows as u64);
325
326            let now = self.tick_manager.tick();
327            for req in reqs {
328                match req {
329                    DiffRequest::Insert(insert) => {
330                        let rows_proto: Vec<v1::Row> = insert
331                            .into_iter()
332                            .map(|(mut row, _ts)| {
333                                // extend `update_at` col if needed
334                                // if schema include a millisecond timestamp here, and result row doesn't have it, add it
335                                if row.len() < proto_schema.len()
336                                    && proto_schema[row.len()].datatype
337                                        == greptime_proto::v1::ColumnDataType::TimestampMillisecond
338                                            as i32
339                                {
340                                    row.extend([Value::from(
341                                        common_time::Timestamp::new_millisecond(now),
342                                    )]);
343                                }
344                                // ts col, if auto create
345                                if is_ts_placeholder {
346                                    ensure!(
347                                        row.len() == schema_len - 1,
348                                        InternalSnafu {
349                                            reason: format!(
350                                                "Row len mismatch, expect {} got {}",
351                                                schema_len - 1,
352                                                row.len()
353                                            )
354                                        }
355                                    );
356                                    row.extend([Value::from(
357                                        common_time::Timestamp::new_millisecond(0),
358                                    )]);
359                                }
360                                if row.len() != proto_schema.len() {
361                                    UnexpectedSnafu {
362                                        reason: format!(
363                                            "Flow output row length mismatch, expect {} got {}, the columns in schema are: {:?}",
364                                            proto_schema.len(),
365                                            row.len(),
366                                            proto_schema.iter().map(|c|&c.column_name).collect_vec()
367                                        ),
368                                    }
369                                    .fail()?;
370                                }
371                                Ok(row.into())
372                            })
373                            .collect::<Result<Vec<_>, Error>>()?;
374                        let table_name = table_name.last().unwrap().clone();
375                        let req = RowInsertRequest {
376                            table_name,
377                            rows: Some(v1::Rows {
378                                schema: proto_schema.clone(),
379                                rows: rows_proto,
380                            }),
381                        };
382                        req_cnt += 1;
383                        self.frontend_invoker
384                            .read()
385                            .await
386                            .as_ref()
387                            .with_context(|| UnexpectedSnafu {
388                                reason: "Expect a frontend invoker for flownode to write back",
389                            })?
390                            .row_inserts(RowInsertRequests { inserts: vec![req] }, ctx.clone())
391                            .await
392                            .map_err(BoxedError::new)
393                            .with_context(|_| ExternalSnafu {})?;
394                    }
395                    DiffRequest::Delete(remove) => {
396                        info!("original remove rows={:?}", remove);
397                        let rows_proto: Vec<v1::Row> = remove
398                            .into_iter()
399                            .map(|(mut row, _ts)| {
400                                row.extend(Some(Value::from(
401                                    common_time::Timestamp::new_millisecond(0),
402                                )));
403                                row.into()
404                            })
405                            .collect::<Vec<_>>();
406                        let table_name = table_name.last().unwrap().clone();
407                        let req = RowDeleteRequest {
408                            table_name,
409                            rows: Some(v1::Rows {
410                                schema: proto_schema.clone(),
411                                rows: rows_proto,
412                            }),
413                        };
414
415                        req_cnt += 1;
416                        self.frontend_invoker
417                            .read()
418                            .await
419                            .as_ref()
420                            .with_context(|| UnexpectedSnafu {
421                                reason: "Expect a frontend invoker for flownode to write back",
422                            })?
423                            .row_deletes(RowDeleteRequests { deletes: vec![req] }, ctx.clone())
424                            .await
425                            .map_err(BoxedError::new)
426                            .with_context(|_| ExternalSnafu {})?;
427                    }
428                }
429            }
430        }
431        Ok(req_cnt)
432    }
433
434    /// Generate writeback request for all sink table
435    pub async fn generate_writeback_request(
436        &self,
437    ) -> Result<BTreeMap<TableName, Vec<DiffRequest>>, Error> {
438        trace!("Start to generate writeback request");
439        let mut output = BTreeMap::new();
440        let mut total_row_count = 0;
441        for (name, sink_recv) in self
442            .node_context
443            .write()
444            .await
445            .sink_receiver
446            .iter_mut()
447            .map(|(n, (_s, r))| (n, r))
448        {
449            let mut batches = Vec::new();
450            while let Ok(batch) = sink_recv.try_recv() {
451                total_row_count += batch.row_count();
452                batches.push(batch);
453            }
454            let reqs = batches_to_rows_req(batches)?;
455            output.insert(name.clone(), reqs);
456        }
457        trace!("Prepare writeback req: total row count={}", total_row_count);
458        Ok(output)
459    }
460
461    /// Fetch table schema and primary key from table info source, if table not exist return None
462    async fn fetch_table_pk_schema(
463        &self,
464        table_name: &TableName,
465    ) -> Result<Option<(Vec<String>, Option<usize>, Vec<ColumnSchema>)>, Error> {
466        if let Some(table_id) = self
467            .table_info_source
468            .get_opt_table_id_from_name(table_name)
469            .await?
470        {
471            let table_info = self
472                .table_info_source
473                .get_table_info_value(&table_id)
474                .await?
475                .unwrap();
476            let meta = table_info.table_info.meta;
477            let schema = meta.schema.column_schemas().to_vec();
478            let primary_keys = meta
479                .primary_key_indices
480                .into_iter()
481                .map(|i| schema[i].name.clone())
482                .collect_vec();
483            let time_index = meta.schema.timestamp_index();
484            Ok(Some((primary_keys, time_index, schema)))
485        } else {
486            Ok(None)
487        }
488    }
489
490    /// return (primary keys, schema and if the table have a placeholder timestamp column)
491    /// schema of the table comes from flow's output plan
492    ///
493    /// adjust to add `update_at` column and ts placeholder if needed
494    async fn adjust_auto_created_table_schema(
495        &self,
496        schema: &RelationDesc,
497    ) -> Result<(Vec<String>, Vec<ColumnSchema>, bool), Error> {
498        // TODO(discord9): consider remove buggy auto create by schema
499
500        // TODO(discord9): use default key from schema
501        let primary_keys = schema
502            .typ()
503            .keys
504            .first()
505            .map(|v| {
506                v.column_indices
507                    .iter()
508                    .map(|i| {
509                        schema
510                            .get_name(*i)
511                            .clone()
512                            .unwrap_or_else(|| format!("col_{i}"))
513                    })
514                    .collect_vec()
515            })
516            .unwrap_or_default();
517        let update_at = ColumnSchema::new(
518            AUTO_CREATED_UPDATE_AT_TS_COL,
519            ConcreteDataType::timestamp_millisecond_datatype(),
520            true,
521        );
522
523        let original_schema = relation_desc_to_column_schemas_with_fallback(schema);
524
525        let mut with_auto_added_col = original_schema.clone();
526        with_auto_added_col.push(update_at);
527
528        // if no time index, add one as placeholder
529        let no_time_index = schema.typ().time_index.is_none();
530        if no_time_index {
531            let ts_col = ColumnSchema::new(
532                AUTO_CREATED_PLACEHOLDER_TS_COL,
533                ConcreteDataType::timestamp_millisecond_datatype(),
534                true,
535            )
536            .with_time_index(true);
537            with_auto_added_col.push(ts_col);
538        }
539
540        Ok((primary_keys, with_auto_added_col, no_time_index))
541    }
542}
543
544/// Flow Runtime related methods
545impl StreamingEngine {
546    /// run in common_runtime background runtime
547    pub fn run_background(
548        self: Arc<Self>,
549        shutdown: Option<broadcast::Receiver<()>>,
550    ) -> JoinHandle<()> {
551        info!("Starting flownode manager's background task");
552        common_runtime::spawn_global(async move { self.run(shutdown).await })
553    }
554
555    /// log all flow errors
556    pub async fn log_all_errors(&self) {
557        for (f_id, f_err) in self.flow_err_collectors.read().await.iter() {
558            let all_errors = f_err.get_all().await;
559            if !all_errors.is_empty() {
560                let all_errors = all_errors
561                    .into_iter()
562                    .map(|i| format!("{:?}", i))
563                    .join("\n");
564                common_telemetry::error!("Flow {} has following errors: {}", f_id, all_errors);
565            }
566        }
567    }
568
569    /// Trigger dataflow running, and then send writeback request to the source sender
570    ///
571    /// note that this method didn't handle input mirror request, as this should be handled by grpc server
572    pub async fn run(&self, mut shutdown: Option<broadcast::Receiver<()>>) {
573        debug!("Starting to run");
574        let default_interval = Duration::from_secs(1);
575        let mut tick_interval = tokio::time::interval(default_interval);
576        // burst mode, so that if we miss a tick, we will run immediately to fully utilize the cpu
577        tick_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Burst);
578        let mut avg_spd = 0; // rows/sec
579        let mut since_last_run = tokio::time::Instant::now();
580        let run_per_trace = 10;
581        let mut run_cnt = 0;
582        loop {
583            // TODO(discord9): only run when new inputs arrive or scheduled to
584            let row_cnt = self.run_available(false).await.unwrap_or_else(|err| {
585                common_telemetry::error!(err;"Run available errors");
586                0
587            });
588
589            if let Err(err) = self.send_writeback_requests().await {
590                common_telemetry::error!(err;"Send writeback request errors");
591            };
592            self.log_all_errors().await;
593
594            // determine if need to shutdown
595            match &shutdown.as_mut().map(|s| s.try_recv()) {
596                Some(Ok(())) => {
597                    info!("Shutdown flow's main loop");
598                    break;
599                }
600                Some(Err(TryRecvError::Empty)) => (),
601                Some(Err(TryRecvError::Closed)) => {
602                    common_telemetry::error!("Shutdown channel is closed");
603                    break;
604                }
605                Some(Err(TryRecvError::Lagged(num))) => {
606                    common_telemetry::error!(
607                        "Shutdown channel is lagged by {}, meaning multiple shutdown cmd have been issued",
608                        num
609                    );
610                    break;
611                }
612                None => (),
613            }
614
615            // for now we want to batch rows until there is around `BATCH_SIZE` rows in send buf
616            // before trigger a run of flow's worker
617            let wait_for = since_last_run.elapsed();
618
619            // last runs insert speed
620            let cur_spd = row_cnt * 1000 / wait_for.as_millis().max(1) as usize;
621            // rapid increase, slow decay
622            avg_spd = if cur_spd > avg_spd {
623                cur_spd
624            } else {
625                (9 * avg_spd + cur_spd) / 10
626            };
627            let new_wait = BATCH_SIZE * 1000 / avg_spd.max(1); //in ms
628            let new_wait = Duration::from_millis(new_wait as u64).min(default_interval);
629
630            // print trace every `run_per_trace` times so that we can see if there is something wrong
631            // but also not get flooded with trace
632            if run_cnt >= run_per_trace {
633                trace!("avg_spd={} r/s, cur_spd={} r/s", avg_spd, cur_spd);
634                trace!("Wait for {} ms, row_cnt={}", new_wait.as_millis(), row_cnt);
635                run_cnt = 0;
636            } else {
637                run_cnt += 1;
638            }
639
640            METRIC_FLOW_RUN_INTERVAL_MS.set(new_wait.as_millis() as i64);
641            since_last_run = tokio::time::Instant::now();
642            tokio::select! {
643                _ = tick_interval.tick() => (),
644                _ = tokio::time::sleep(new_wait) => ()
645            }
646        }
647        // flow is now shutdown, drop frontend_invoker early so a ref cycle(in standalone mode) can be prevent:
648        // FlowWorkerManager.frontend_invoker -> FrontendInvoker.inserter
649        // -> Inserter.node_manager -> NodeManager.flownode -> Flownode.flow_streaming_engine.frontend_invoker
650        self.frontend_invoker.write().await.take();
651    }
652
653    /// Run all available subgraph in the flow node
654    /// This will try to run all dataflow in this node
655    ///
656    /// set `blocking` to true to wait until worker finish running
657    /// false to just trigger run and return immediately
658    /// return numbers of rows send to worker(Inaccuary)
659    /// TODO(discord9): add flag for subgraph that have input since last run
660    pub async fn run_available(&self, blocking: bool) -> Result<usize, Error> {
661        let mut row_cnt = 0;
662
663        let now = self.tick_manager.tick();
664        for worker in self.worker_handles.iter() {
665            // TODO(discord9): consider how to handle error in individual worker
666            worker.run_available(now, blocking).await?;
667        }
668        // check row send and rows remain in send buf
669        let flush_res = if blocking {
670            let ctx = self.node_context.read().await;
671            ctx.flush_all_sender().await
672        } else {
673            match self.node_context.try_read() {
674                Ok(ctx) => ctx.flush_all_sender().await,
675                Err(_) => return Ok(row_cnt),
676            }
677        };
678        match flush_res {
679            Ok(r) => {
680                common_telemetry::trace!("Total flushed {} rows", r);
681                row_cnt += r;
682            }
683            Err(err) => {
684                common_telemetry::error!("Flush send buf errors: {:?}", err);
685            }
686        };
687
688        Ok(row_cnt)
689    }
690
691    /// send write request to related source sender
692    pub async fn handle_write_request(
693        &self,
694        region_id: RegionId,
695        rows: Vec<DiffRow>,
696        batch_datatypes: &[ConcreteDataType],
697    ) -> Result<(), Error> {
698        let rows_len = rows.len();
699        let table_id = region_id.table_id();
700        let _timer = METRIC_FLOW_INSERT_ELAPSED
701            .with_label_values(&[table_id.to_string().as_str()])
702            .start_timer();
703        self.node_context
704            .read()
705            .await
706            .send(table_id, rows, batch_datatypes)
707            .await?;
708        trace!(
709            "Handling write request for table_id={} with {} rows",
710            table_id, rows_len
711        );
712        Ok(())
713    }
714}
715
716/// Create&Remove flow
717impl StreamingEngine {
718    /// remove a flow by it's id
719    pub async fn remove_flow_inner(&self, flow_id: FlowId) -> Result<(), Error> {
720        for handle in self.worker_handles.iter() {
721            if handle.contains_flow(flow_id).await? {
722                handle.remove_flow(flow_id).await?;
723                break;
724            }
725        }
726        self.node_context.write().await.remove_flow(flow_id);
727        Ok(())
728    }
729
730    /// Return task id if a new task is created, otherwise return None
731    ///
732    /// steps to create task:
733    /// 1. parse query into typed plan(and optional parse expire_after expr)
734    /// 2. render source/sink with output table id and used input table id
735    pub async fn create_flow_inner(&self, args: CreateFlowArgs) -> Result<Option<FlowId>, Error> {
736        let CreateFlowArgs {
737            flow_id,
738            sink_table_name,
739            source_table_ids,
740            create_if_not_exists,
741            or_replace,
742            expire_after,
743            eval_interval: _,
744            comment,
745            sql,
746            flow_options,
747            query_ctx,
748        } = args;
749
750        let mut node_ctx = self.node_context.write().await;
751        // assign global id to source and sink table
752        for source in &source_table_ids {
753            node_ctx
754                .assign_global_id_to_table(&self.table_info_source, None, Some(*source))
755                .await?;
756        }
757        node_ctx
758            .assign_global_id_to_table(&self.table_info_source, Some(sink_table_name.clone()), None)
759            .await?;
760
761        node_ctx.register_task_src_sink(flow_id, &source_table_ids, sink_table_name.clone());
762
763        node_ctx.query_context = query_ctx.map(Arc::new);
764        // construct a active dataflow state with it
765        let flow_plan = sql_to_flow_plan(&mut node_ctx, &self.query_engine, &sql).await?;
766
767        debug!("Flow {:?}'s Plan is {:?}", flow_id, flow_plan);
768
769        // check schema against actual table schema if exists
770        // if not exist create sink table immediately
771        if let Some((_, _, real_schema)) = self.fetch_table_pk_schema(&sink_table_name).await? {
772            let auto_schema = relation_desc_to_column_schemas_with_fallback(&flow_plan.schema);
773
774            // for column schema, only `data_type` need to be check for equality
775            // since one can omit flow's column name when write flow query
776            // print a user friendly error message about mismatch and how to correct them
777            for (idx, zipped) in auto_schema
778                .iter()
779                .zip_longest(real_schema.iter())
780                .enumerate()
781            {
782                match zipped {
783                    EitherOrBoth::Both(auto, real) => {
784                        if auto.data_type != real.data_type {
785                            InvalidQuerySnafu {
786                                    reason: format!(
787                                        "Column {}(name is '{}', flow inferred name is '{}')'s data type mismatch, expect {:?} got {:?}",
788                                        idx,
789                                        real.name,
790                                        auto.name,
791                                        real.data_type,
792                                        auto.data_type
793                                    ),
794                                }
795                                .fail()?;
796                        }
797                    }
798                    EitherOrBoth::Right(real) if real.data_type.is_timestamp() => {
799                        // if table is auto created, the last one or two column should be timestamp(update at and ts placeholder)
800                        continue;
801                    }
802                    _ => InvalidQuerySnafu {
803                        reason: format!(
804                            "schema length mismatched, expected {} found {}",
805                            real_schema.len(),
806                            auto_schema.len()
807                        ),
808                    }
809                    .fail()?,
810                }
811            }
812        } else {
813            // assign inferred schema to sink table
814            // create sink table
815            let did_create = self
816                .create_table_from_relation(
817                    &format!("flow-id={flow_id}"),
818                    &sink_table_name,
819                    &flow_plan.schema,
820                )
821                .await?;
822            if !did_create {
823                UnexpectedSnafu {
824                    reason: format!("Failed to create table {:?}", sink_table_name),
825                }
826                .fail()?;
827            }
828        }
829
830        node_ctx.add_flow_plan(flow_id, flow_plan.clone());
831
832        let _ = comment;
833        let _ = flow_options;
834
835        // TODO(discord9): add more than one handles
836        let sink_id = node_ctx.table_repr.get_by_name(&sink_table_name).unwrap().1;
837        let sink_sender = node_ctx.get_sink_by_global_id(&sink_id)?;
838
839        let source_ids = source_table_ids
840            .iter()
841            .map(|id| node_ctx.table_repr.get_by_table_id(id).unwrap().1)
842            .collect_vec();
843        let source_receivers = source_ids
844            .iter()
845            .map(|id| {
846                node_ctx
847                    .get_source_by_global_id(id)
848                    .map(|s| s.get_receiver())
849            })
850            .collect::<Result<Vec<_>, _>>()?;
851        let err_collector = ErrCollector::default();
852        self.flow_err_collectors
853            .write()
854            .await
855            .insert(flow_id, err_collector.clone());
856        // TODO(discord9): load balance?
857        let handle = self.get_worker_handle_for_create_flow().await;
858        let create_request = worker::Request::Create {
859            flow_id,
860            plan: flow_plan,
861            sink_id,
862            sink_sender,
863            source_ids,
864            src_recvs: source_receivers,
865            expire_after,
866            or_replace,
867            create_if_not_exists,
868            err_collector,
869        };
870
871        handle.create_flow(create_request).await?;
872        info!("Successfully create flow with id={}", flow_id);
873        Ok(Some(flow_id))
874    }
875
876    pub async fn flush_flow_inner(&self, flow_id: FlowId) -> Result<usize, Error> {
877        debug!("Starting to flush flow_id={:?}", flow_id);
878        // lock to make sure writes before flush are written to flow
879        // and immediately drop to prevent following writes to be blocked
880        drop(self.flush_lock.write().await);
881        let flushed_input_rows = self.node_context.read().await.flush_all_sender().await?;
882        let rows_send = self.run_available(true).await?;
883        let row = self.send_writeback_requests().await?;
884        debug!(
885            "Done to flush flow_id={:?} with {} input rows flushed, {} rows sent and {} output rows flushed",
886            flow_id, flushed_input_rows, rows_send, row
887        );
888        Ok(row)
889    }
890
891    pub async fn flow_exist_inner(&self, flow_id: FlowId) -> Result<bool, Error> {
892        let mut exist = false;
893        for handle in self.worker_handles.iter() {
894            if handle.contains_flow(flow_id).await? {
895                exist = true;
896                break;
897            }
898        }
899        Ok(exist)
900    }
901}
902
903/// FlowTickManager is a manager for flow tick, which trakc flow execution progress
904///
905/// TODO(discord9): better way to do it, and not expose flow tick even to other flow to avoid
906/// TSO coord mess
907#[derive(Clone, Debug)]
908pub struct FlowTickManager {
909    /// The starting instant of the flow, used with `start_timestamp` to calculate the current timestamp
910    start: Instant,
911    /// The timestamp when the flow started
912    start_timestamp: repr::Timestamp,
913}
914
915impl Default for FlowTickManager {
916    fn default() -> Self {
917        Self::new()
918    }
919}
920
921impl FlowTickManager {
922    pub fn new() -> Self {
923        FlowTickManager {
924            start: Instant::now(),
925            start_timestamp: SystemTime::now()
926                .duration_since(SystemTime::UNIX_EPOCH)
927                .unwrap()
928                .as_millis() as repr::Timestamp,
929        }
930    }
931
932    /// Return the current timestamp in milliseconds
933    ///
934    /// TODO(discord9): reconsider since `tick()` require a monotonic clock and also need to survive recover later
935    pub fn tick(&self) -> repr::Timestamp {
936        let current = Instant::now();
937        let since_the_epoch = current - self.start;
938        since_the_epoch.as_millis() as repr::Timestamp + self.start_timestamp
939    }
940}
flow/adapter.rs

flow/
adapter.rs
