feat: show flow's mem usage in INFORMATION_SCHEMA.FLOWS (#4890)

* feat: add flow mem size to sys table

* chore: rm dup def

* chore: remove unused variant

* chore: minor refactor

* refactor: per review

src/flow/Cargo.toml

@@ -40,6 +40,8 @@ datatypes.workspace = true
 enum-as-inner = "0.6.0"
 enum_dispatch = "0.3"
 futures = "0.3"
+get-size-derive2 = "0.1.2"
+get-size2 = "0.1.2"
 greptime-proto.workspace = true
 # This fork of hydroflow is simply for keeping our dependency in our org, and pin the version
 # otherwise it is the same with upstream repo

src/flow/src/adapter.rs

@@ -60,6 +60,7 @@ use crate::repr::{self, DiffRow, Row, BATCH_SIZE};
 
 mod flownode_impl;
 mod parse_expr;
+mod stat;
 #[cfg(test)]
 mod tests;
 mod util;
@@ -69,6 +70,7 @@ pub(crate) mod node_context;
 mod table_source;
 
 use crate::error::Error;
+use crate::utils::StateReportHandler;
 use crate::FrontendInvoker;
 
 // `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow
@@ -137,6 +139,8 @@ pub struct FlowWorkerManager {
     ///
     /// So that a series of event like `inserts -> flush` can be handled correctly
     flush_lock: RwLock<()>,
+    /// receive a oneshot sender to send state size report
+    state_report_handler: RwLock<Option<StateReportHandler>>,
 }
 
 /// Building FlownodeManager
@@ -170,9 +174,15 @@ impl FlowWorkerManager {
             tick_manager,
             node_id,
             flush_lock: RwLock::new(()),
+            state_report_handler: RwLock::new(None),
         }
     }
 
+    pub async fn with_state_report_handler(self, handler: StateReportHandler) -> Self {
+        *self.state_report_handler.write().await = Some(handler);
+        self
+    }
+
     /// Create a flownode manager with one worker
     pub fn new_with_worker<'s>(
         node_id: Option<u32>,
@@ -500,6 +510,27 @@ impl FlowWorkerManager {
 
 /// Flow Runtime related methods
 impl FlowWorkerManager {
+    /// Start state report handler, which will receive a sender from HeartbeatTask to send state size report back
+    ///
+    /// if heartbeat task is shutdown, this future will exit too
+    async fn start_state_report_handler(self: Arc<Self>) -> Option<JoinHandle<()>> {
+        let state_report_handler = self.state_report_handler.write().await.take();
+        if let Some(mut handler) = state_report_handler {
+            let zelf = self.clone();
+            let handler = common_runtime::spawn_global(async move {
+                while let Some(ret_handler) = handler.recv().await {
+                    let state_report = zelf.gen_state_report().await;
+                    ret_handler.send(state_report).unwrap_or_else(|err| {
+                        common_telemetry::error!(err; "Send state size report error");
+                    });
+                }
+            });
+            Some(handler)
+        } else {
+            None
+        }
+    }
+
     /// run in common_runtime background runtime
     pub fn run_background(
         self: Arc<Self>,
@@ -507,6 +538,7 @@ impl FlowWorkerManager {
     ) -> JoinHandle<()> {
         info!("Starting flownode manager's background task");
         common_runtime::spawn_global(async move {
+            let _state_report_handler = self.clone().start_state_report_handler().await;
             self.run(shutdown).await;
         })
     }

src/flow/src/adapter/stat.rs

@@ -0,0 +1,40 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+
+use common_meta::key::flow::flow_state::FlowStat;
+
+use crate::FlowWorkerManager;
+
+impl FlowWorkerManager {
+    pub async fn gen_state_report(&self) -> FlowStat {
+        let mut full_report = BTreeMap::new();
+        for worker in self.worker_handles.iter() {
+            let worker = worker.lock().await;
+            match worker.get_state_size().await {
+                Ok(state_size) => {
+                    full_report.extend(state_size.into_iter().map(|(k, v)| (k as u32, v)))
+                }
+                Err(err) => {
+                    common_telemetry::error!(err; "Get flow stat size error");
+                }
+            }
+        }
+
+        FlowStat {
+            state_size: full_report,
+        }
+    }
+}

src/flow/src/adapter/worker.rs

@@ -197,6 +197,21 @@ impl WorkerHandle {
             .fail()
         }
     }
+
+    pub async fn get_state_size(&self) -> Result<BTreeMap<FlowId, usize>, Error> {
+        let ret = self
+            .itc_client
+            .call_with_resp(Request::QueryStateSize)
+            .await?;
+        ret.into_query_state_size().map_err(|ret| {
+            InternalSnafu {
+                reason: format!(
+                    "Flow Node/Worker itc failed, expect Response::QueryStateSize, found {ret:?}"
+                ),
+            }
+            .build()
+        })
+    }
 }
 
 impl Drop for WorkerHandle {
@@ -361,6 +376,13 @@ impl<'s> Worker<'s> {
                 Some(Response::ContainTask { result: ret })
             }
             Request::Shutdown => return Err(()),
+            Request::QueryStateSize => {
+                let mut ret = BTreeMap::new();
+                for (flow_id, task_state) in self.task_states.iter() {
+                    ret.insert(*flow_id, task_state.state.get_state_size());
+                }
+                Some(Response::QueryStateSize { result: ret })
+            }
         };
         Ok(ret)
     }
@@ -391,6 +413,7 @@ pub enum Request {
         flow_id: FlowId,
     },
     Shutdown,
+    QueryStateSize,
 }
 
 #[derive(Debug, EnumAsInner)]
@@ -406,6 +429,10 @@ enum Response {
         result: bool,
     },
     RunAvail,
+    QueryStateSize {
+        /// each flow tasks' state size
+        result: BTreeMap<FlowId, usize>,
+    },
 }
 
 fn create_inter_thread_call() -> (InterThreadCallClient, InterThreadCallServer) {
@@ -423,10 +450,12 @@ struct InterThreadCallClient {
 }
 
 impl InterThreadCallClient {
+    /// call without response
     fn call_no_resp(&self, req: Request) -> Result<(), Error> {
         self.arg_sender.send((req, None)).map_err(from_send_error)
     }
 
+    /// call with response
     async fn call_with_resp(&self, req: Request) -> Result<Response, Error> {
         let (tx, rx) = oneshot::channel();
         self.arg_sender
@@ -527,6 +556,7 @@ mod test {
         );
         tx.send(Batch::empty()).unwrap();
         handle.run_available(0, true).await.unwrap();
+        assert_eq!(handle.get_state_size().await.unwrap().len(), 1);
         assert_eq!(sink_rx.recv().await.unwrap(), Batch::empty());
         drop(handle);
         worker_thread_handle.join().unwrap();

src/flow/src/compute/state.rs

@@ -16,6 +16,7 @@ use std::cell::RefCell;
 use std::collections::{BTreeMap, VecDeque};
 use std::rc::Rc;
 
+use get_size2::GetSize;
 use hydroflow::scheduled::graph::Hydroflow;
 use hydroflow::scheduled::SubgraphId;
 
@@ -109,6 +110,10 @@ impl DataflowState {
     pub fn expire_after(&self) -> Option<Timestamp> {
         self.expire_after
     }
+
+    pub fn get_state_size(&self) -> usize {
+        self.arrange_used.iter().map(|x| x.read().get_size()).sum()
+    }
 }
 
 #[derive(Debug, Clone)]

src/flow/src/heartbeat.rs

@@ -24,6 +24,7 @@ use common_meta::heartbeat::handler::{
 };
 use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MailboxRef, OutgoingMessage};
 use common_meta::heartbeat::utils::outgoing_message_to_mailbox_message;
+use common_meta::key::flow::flow_state::FlowStat;
 use common_telemetry::{debug, error, info, warn};
 use greptime_proto::v1::meta::NodeInfo;
 use meta_client::client::{HeartbeatSender, HeartbeatStream, MetaClient};
@@ -34,8 +35,27 @@ use tokio::sync::mpsc;
 use tokio::time::Duration;
 
 use crate::error::ExternalSnafu;
+use crate::utils::SizeReportSender;
 use crate::{Error, FlownodeOptions};
 
+async fn query_flow_state(
+    query_stat_size: &Option<SizeReportSender>,
+    timeout: Duration,
+) -> Option<FlowStat> {
+    if let Some(report_requester) = query_stat_size.as_ref() {
+        let ret = report_requester.query(timeout).await;
+        match ret {
+            Ok(latest) => Some(latest),
+            Err(err) => {
+                error!(err; "Failed to get query stat size");
+                None
+            }
+        }
+    } else {
+        None
+    }
+}
+
 /// The flownode heartbeat task which sending `[HeartbeatRequest]` to Metasrv periodically in background.
 #[derive(Clone)]
 pub struct HeartbeatTask {
@@ -47,9 +67,14 @@ pub struct HeartbeatTask {
     resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
     start_time_ms: u64,
     running: Arc<AtomicBool>,
+    query_stat_size: Option<SizeReportSender>,
 }
 
 impl HeartbeatTask {
+    pub fn with_query_stat_size(mut self, query_stat_size: SizeReportSender) -> Self {
+        self.query_stat_size = Some(query_stat_size);
+        self
+    }
     pub fn new(
         opts: &FlownodeOptions,
         meta_client: Arc<MetaClient>,
@@ -65,6 +90,7 @@ impl HeartbeatTask {
             resp_handler_executor,
             start_time_ms: common_time::util::current_time_millis() as u64,
             running: Arc::new(AtomicBool::new(false)),
+            query_stat_size: None,
         }
     }
 
@@ -112,6 +138,7 @@ impl HeartbeatTask {
         message: Option<OutgoingMessage>,
         peer: Option<Peer>,
         start_time_ms: u64,
+        latest_report: &Option<FlowStat>,
     ) -> Option<HeartbeatRequest> {
         let mailbox_message = match message.map(outgoing_message_to_mailbox_message) {
             Some(Ok(message)) => Some(message),
@@ -121,11 +148,22 @@ impl HeartbeatTask {
             }
             None => None,
         };
+        let flow_stat = latest_report
+            .as_ref()
+            .map(|report| {
+                report
+                    .state_size
+                    .iter()
+                    .map(|(k, v)| (*k, *v as u64))
+                    .collect()
+            })
+            .map(|f| api::v1::meta::FlowStat { flow_stat_size: f });
 
         Some(HeartbeatRequest {
             mailbox_message,
             peer,
             info: Self::build_node_info(start_time_ms),
+            flow_stat,
             ..Default::default()
         })
     }
@@ -151,24 +189,27 @@ impl HeartbeatTask {
             addr: self.peer_addr.clone(),
         });
 
+        let query_stat_size = self.query_stat_size.clone();
+
         common_runtime::spawn_hb(async move {
             // note that using interval will cause it to first immediately send
             // a heartbeat
             let mut interval = tokio::time::interval(report_interval);
             interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+            let mut latest_report = None;
 
             loop {
                 let req = tokio::select! {
                     message = outgoing_rx.recv() => {
                         if let Some(message) = message {
-                            Self::create_heartbeat_request(Some(message), self_peer.clone(), start_time_ms)
+                            Self::create_heartbeat_request(Some(message), self_peer.clone(), start_time_ms, &latest_report)
                         } else {
                             // Receives None that means Sender was dropped, we need to break the current loop
                             break
                         }
                     }
                     _ = interval.tick() => {
-                        Self::create_heartbeat_request(None, self_peer.clone(), start_time_ms)
+                        Self::create_heartbeat_request(None, self_peer.clone(), start_time_ms, &latest_report)
                     }
                 };
 
@@ -180,6 +221,10 @@ impl HeartbeatTask {
                         debug!("Send a heartbeat request to metasrv, content: {:?}", req);
                     }
                 }
+                // after sending heartbeat, try to get the latest report
+                // TODO(discord9): consider a better place to update the size report
+                // set the timeout to half of the report interval so that it wouldn't delay heartbeat if something went horribly wrong
+                latest_report = query_flow_state(&query_stat_size, report_interval / 2).await;
             }
         });
     }

src/flow/src/repr.rs

@@ -22,12 +22,14 @@ use api::v1::Row as ProtoRow;
 use datatypes::data_type::ConcreteDataType;
 use datatypes::types::cast;
 use datatypes::value::Value;
+use get_size2::GetSize;
 use itertools::Itertools;
 pub(crate) use relation::{ColumnType, Key, RelationDesc, RelationType};
 use serde::{Deserialize, Serialize};
 use snafu::ResultExt;
 
 use crate::expr::error::{CastValueSnafu, EvalError, InvalidArgumentSnafu};
+use crate::utils::get_value_heap_size;
 
 /// System-wide Record count difference type. Useful for capture data change
 ///
@@ -105,6 +107,12 @@ pub struct Row {
     pub inner: Vec<Value>,
 }
 
+impl GetSize for Row {
+    fn get_heap_size(&self) -> usize {
+        self.inner.iter().map(get_value_heap_size).sum()
+    }
+}
+
 impl Row {
     /// Create an empty row
     pub fn empty() -> Self {

src/flow/src/server.rs

@@ -55,6 +55,7 @@ use crate::error::{
 };
 use crate::heartbeat::HeartbeatTask;
 use crate::transform::register_function_to_query_engine;
+use crate::utils::{SizeReportSender, StateReportHandler};
 use crate::{Error, FlowWorkerManager, FlownodeOptions};
 
 pub const FLOW_NODE_SERVER_NAME: &str = "FLOW_NODE_SERVER";
@@ -236,6 +237,8 @@ pub struct FlownodeBuilder {
     catalog_manager: CatalogManagerRef,
     flow_metadata_manager: FlowMetadataManagerRef,
     heartbeat_task: Option<HeartbeatTask>,
+    /// receive a oneshot sender to send state size report
+    state_report_handler: Option<StateReportHandler>,
 }
 
 impl FlownodeBuilder {
@@ -254,17 +257,20 @@ impl FlownodeBuilder {
             catalog_manager,
             flow_metadata_manager,
             heartbeat_task: None,
+            state_report_handler: None,
         }
     }
 
     pub fn with_heartbeat_task(self, heartbeat_task: HeartbeatTask) -> Self {
+        let (sender, receiver) = SizeReportSender::new();
         Self {
-            heartbeat_task: Some(heartbeat_task),
+            heartbeat_task: Some(heartbeat_task.with_query_stat_size(sender)),
+            state_report_handler: Some(receiver),
             ..self
         }
     }
 
-    pub async fn build(self) -> Result<FlownodeInstance, Error> {
+    pub async fn build(mut self) -> Result<FlownodeInstance, Error> {
         // TODO(discord9): does this query engine need those?
         let query_engine_factory = QueryEngineFactory::new_with_plugins(
             // query engine in flownode is only used for translate plan with resolved table source.
@@ -383,7 +389,7 @@ impl FlownodeBuilder {
     /// build [`FlowWorkerManager`], note this doesn't take ownership of `self`,
     /// nor does it actually start running the worker.
     async fn build_manager(
-        &self,
+        &mut self,
         query_engine: Arc<dyn QueryEngine>,
     ) -> Result<FlowWorkerManager, Error> {
         let table_meta = self.table_meta.clone();
@@ -402,12 +408,15 @@ impl FlownodeBuilder {
                 info!("Flow Worker started in new thread");
                 worker.run();
             });
-        let man = rx.await.map_err(|_e| {
+        let mut man = rx.await.map_err(|_e| {
             UnexpectedSnafu {
                 reason: "sender is dropped, failed to create flow node manager",
             }
             .build()
         })?;
+        if let Some(handler) = self.state_report_handler.take() {
+            man = man.with_state_report_handler(handler).await;
+        }
         info!("Flow Node Manager started");
         Ok(man)
     }

src/flow/src/utils.rs

@@ -18,16 +18,73 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::ops::Bound;
 use std::sync::Arc;
 
+use common_meta::key::flow::flow_state::FlowStat;
 use common_telemetry::trace;
+use datatypes::value::Value;
+use get_size2::GetSize;
 use smallvec::{smallvec, SmallVec};
-use tokio::sync::RwLock;
+use tokio::sync::{mpsc, oneshot, RwLock};
+use tokio::time::Instant;
 
+use crate::error::InternalSnafu;
 use crate::expr::{EvalError, ScalarExpr};
 use crate::repr::{value_to_internal_ts, DiffRow, Duration, KeyValDiffRow, Row, Timestamp};
 
 /// A batch of updates, arranged by key
 pub type Batch = BTreeMap<Row, SmallVec<[DiffRow; 2]>>;
 
+/// Get a estimate of heap size of a value
+pub fn get_value_heap_size(v: &Value) -> usize {
+    match v {
+        Value::Binary(bin) => bin.len(),
+        Value::String(s) => s.len(),
+        Value::List(list) => list.items().iter().map(get_value_heap_size).sum(),
+        _ => 0,
+    }
+}
+
+#[derive(Clone)]
+pub struct SizeReportSender {
+    inner: mpsc::Sender<oneshot::Sender<FlowStat>>,
+}
+
+impl SizeReportSender {
+    pub fn new() -> (Self, StateReportHandler) {
+        let (tx, rx) = mpsc::channel(1);
+        let zelf = Self { inner: tx };
+        (zelf, rx)
+    }
+
+    /// Query the size report, will timeout after one second if no response
+    pub async fn query(&self, timeout: std::time::Duration) -> crate::Result<FlowStat> {
+        let (tx, rx) = oneshot::channel();
+        self.inner.send(tx).await.map_err(|_| {
+            InternalSnafu {
+                reason: "failed to send size report request due to receiver dropped",
+            }
+            .build()
+        })?;
+        let timeout = tokio::time::timeout(timeout, rx);
+        timeout
+            .await
+            .map_err(|_elapsed| {
+                InternalSnafu {
+                    reason: "failed to receive size report after one second timeout",
+                }
+                .build()
+            })?
+            .map_err(|_| {
+                InternalSnafu {
+                    reason: "failed to receive size report due to sender dropped",
+                }
+                .build()
+            })
+    }
+}
+
+/// Handle the size report request, and send the report back
+pub type StateReportHandler = mpsc::Receiver<oneshot::Sender<FlowStat>>;
+
 /// A spine of batches, arranged by timestamp
 /// TODO(discord9): consider internally index by key, value, and timestamp for faster lookup
 pub type Spine = BTreeMap<Timestamp, Batch>;
@@ -49,6 +106,24 @@ pub struct KeyExpiryManager {
     event_timestamp_from_row: Option<ScalarExpr>,
 }
 
+impl GetSize for KeyExpiryManager {
+    fn get_heap_size(&self) -> usize {
+        let row_size = if let Some(row_size) = &self
+            .event_ts_to_key
+            .first_key_value()
+            .map(|(_, v)| v.first().get_heap_size())
+        {
+            *row_size
+        } else {
+            0
+        };
+        self.event_ts_to_key
+            .values()
+            .map(|v| v.len() * row_size + std::mem::size_of::<i64>())
+            .sum::<usize>()
+    }
+}
+
 impl KeyExpiryManager {
     pub fn new(
         key_expiration_duration: Option<Duration>,
@@ -154,7 +229,7 @@ impl KeyExpiryManager {
 ///
 /// Note the two way arrow between reduce operator and arrange, it's because reduce operator need to query existing state
 /// and also need to update existing state.
-#[derive(Debug, Clone, Default, Eq, PartialEq, Ord, PartialOrd)]
+#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
 pub struct Arrangement {
     /// A name or identifier for the arrangement which can be used for debugging or logging purposes.
     /// This field is not critical to the functionality but aids in monitoring and management of arrangements.
@@ -196,6 +271,61 @@ pub struct Arrangement {
 
     /// The time that the last compaction happened, also known as the current time.
     last_compaction_time: Option<Timestamp>,
+
+    /// Estimated size of the arrangement in heap size.
+    estimated_size: usize,
+    last_size_update: Instant,
+    size_update_interval: tokio::time::Duration,
+}
+
+impl Arrangement {
+    fn compute_size(&self) -> usize {
+        self.spine
+            .values()
+            .map(|v| {
+                let per_entry_size = v
+                    .first_key_value()
+                    .map(|(k, v)| {
+                        k.get_heap_size()
+                            + v.len() * v.first().map(|r| r.get_heap_size()).unwrap_or(0)
+                    })
+                    .unwrap_or(0);
+                std::mem::size_of::<i64>() + v.len() * per_entry_size
+            })
+            .sum::<usize>()
+            + self.expire_state.get_heap_size()
+            + self.name.get_heap_size()
+    }
+
+    fn update_and_fetch_size(&mut self) -> usize {
+        if self.last_size_update.elapsed() > self.size_update_interval {
+            self.estimated_size = self.compute_size();
+            self.last_size_update = Instant::now();
+        }
+        self.estimated_size
+    }
+}
+
+impl GetSize for Arrangement {
+    fn get_heap_size(&self) -> usize {
+        self.estimated_size
+    }
+}
+
+impl Default for Arrangement {
+    fn default() -> Self {
+        Self {
+            spine: Default::default(),
+            full_arrangement: false,
+            is_written: false,
+            expire_state: None,
+            last_compaction_time: None,
+            name: Vec::new(),
+            estimated_size: 0,
+            last_size_update: Instant::now(),
+            size_update_interval: tokio::time::Duration::from_secs(3),
+        }
+    }
 }
 
 impl Arrangement {
@@ -207,6 +337,9 @@ impl Arrangement {
             expire_state: None,
             last_compaction_time: None,
             name,
+            estimated_size: 0,
+            last_size_update: Instant::now(),
+            size_update_interval: tokio::time::Duration::from_secs(3),
         }
     }
 
@@ -269,6 +402,7 @@ impl Arrangement {
             // without changing the order of updates within same tick
             key_updates.sort_by_key(|(_val, ts, _diff)| *ts);
         }
+        self.update_and_fetch_size();
         Ok(max_expired_by)
     }
 
@@ -390,6 +524,7 @@ impl Arrangement {
 
         // insert the compacted batch into spine with key being `now`
         self.spine.insert(now, compacting_batch);
+        self.update_and_fetch_size();
         Ok(max_expired_by)
     }