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

flow/
adapter.rs
