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/common/meta/src/key.rs

@@ -137,6 +137,7 @@ use self::schema_name::{SchemaManager, SchemaNameKey, SchemaNameValue};
 use self::table_route::{TableRouteManager, TableRouteValue};
 use self::tombstone::TombstoneManager;
 use crate::error::{self, Result, SerdeJsonSnafu};
+use crate::key::flow::flow_state::FlowStateValue;
 use crate::key::node_address::NodeAddressValue;
 use crate::key::table_route::TableRouteKey;
 use crate::key::txn_helper::TxnOpGetResponseSet;
@@ -1262,7 +1263,8 @@ impl_metadata_value! {
     FlowRouteValue,
     TableFlowValue,
     NodeAddressValue,
-    SchemaNameValue
+    SchemaNameValue,
+    FlowStateValue
 }
 
 impl_optional_metadata_value! {

src/common/meta/src/key/flow.rs

@@ -15,6 +15,7 @@
 pub mod flow_info;
 pub(crate) mod flow_name;
 pub(crate) mod flow_route;
+pub mod flow_state;
 pub(crate) mod flownode_flow;
 pub(crate) mod table_flow;
 
@@ -35,6 +36,7 @@ use crate::ensure_values;
 use crate::error::{self, Result};
 use crate::key::flow::flow_info::FlowInfoManager;
 use crate::key::flow::flow_name::FlowNameManager;
+use crate::key::flow::flow_state::FlowStateManager;
 use crate::key::flow::flownode_flow::FlownodeFlowManager;
 pub use crate::key::flow::table_flow::{TableFlowManager, TableFlowManagerRef};
 use crate::key::txn_helper::TxnOpGetResponseSet;
@@ -102,6 +104,8 @@ pub struct FlowMetadataManager {
     flownode_flow_manager: FlownodeFlowManager,
     table_flow_manager: TableFlowManager,
     flow_name_manager: FlowNameManager,
+    /// only metasrv have access to itself's memory backend, so for other case it should be None
+    flow_state_manager: Option<FlowStateManager>,
     kv_backend: KvBackendRef,
 }
 
@@ -114,6 +118,7 @@ impl FlowMetadataManager {
             flow_name_manager: FlowNameManager::new(kv_backend.clone()),
             flownode_flow_manager: FlownodeFlowManager::new(kv_backend.clone()),
             table_flow_manager: TableFlowManager::new(kv_backend.clone()),
+            flow_state_manager: None,
             kv_backend,
         }
     }
@@ -123,6 +128,10 @@ impl FlowMetadataManager {
         &self.flow_name_manager
     }
 
+    pub fn flow_state_manager(&self) -> Option<&FlowStateManager> {
+        self.flow_state_manager.as_ref()
+    }
+
     /// Returns the [`FlowInfoManager`].
     pub fn flow_info_manager(&self) -> &FlowInfoManager {
         &self.flow_info_manager

src/common/meta/src/key/flow/flow_state.rs

@@ -0,0 +1,162 @@
+// 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 std::sync::Arc;
+
+use serde::{Deserialize, Serialize};
+
+use crate::error::{self, Result};
+use crate::key::flow::FlowScoped;
+use crate::key::{FlowId, MetadataKey, MetadataValue};
+use crate::kv_backend::KvBackendRef;
+use crate::rpc::store::PutRequest;
+
+/// The entire FlowId to Flow Size's Map is stored directly in the value part of the key.
+const FLOW_STATE_KEY: &str = "state";
+
+/// The key of flow state.
+#[derive(Debug, Clone, Copy, PartialEq)]
+struct FlowStateKeyInner;
+
+impl FlowStateKeyInner {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl<'a> MetadataKey<'a, FlowStateKeyInner> for FlowStateKeyInner {
+    fn to_bytes(&self) -> Vec<u8> {
+        FLOW_STATE_KEY.as_bytes().to_vec()
+    }
+
+    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateKeyInner> {
+        let key = std::str::from_utf8(bytes).map_err(|e| {
+            error::InvalidMetadataSnafu {
+                err_msg: format!(
+                    "FlowInfoKeyInner '{}' is not a valid UTF8 string: {e}",
+                    String::from_utf8_lossy(bytes)
+                ),
+            }
+            .build()
+        })?;
+        if key != FLOW_STATE_KEY {
+            return Err(error::InvalidMetadataSnafu {
+                err_msg: format!("Invalid FlowStateKeyInner '{key}'"),
+            }
+            .build());
+        }
+        Ok(FlowStateKeyInner::new())
+    }
+}
+
+/// The key stores the state size of the flow.
+///
+/// The layout: `__flow/state`.
+pub struct FlowStateKey(FlowScoped<FlowStateKeyInner>);
+
+impl FlowStateKey {
+    /// Returns the [FlowStateKey].
+    pub fn new() -> FlowStateKey {
+        let inner = FlowStateKeyInner::new();
+        FlowStateKey(FlowScoped::new(inner))
+    }
+}
+
+impl Default for FlowStateKey {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<'a> MetadataKey<'a, FlowStateKey> for FlowStateKey {
+    fn to_bytes(&self) -> Vec<u8> {
+        self.0.to_bytes()
+    }
+
+    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateKey> {
+        Ok(FlowStateKey(FlowScoped::<FlowStateKeyInner>::from_bytes(
+            bytes,
+        )?))
+    }
+}
+
+/// The value of flow state size
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct FlowStateValue {
+    /// For each key, the bytes of the state in memory
+    pub state_size: BTreeMap<FlowId, usize>,
+}
+
+impl FlowStateValue {
+    pub fn new(state_size: BTreeMap<FlowId, usize>) -> Self {
+        Self { state_size }
+    }
+}
+
+pub type FlowStateManagerRef = Arc<FlowStateManager>;
+
+/// The manager of [FlowStateKey]. Since state size changes frequently, we store it in memory.
+///
+/// This is only used in distributed mode. When meta-srv use heartbeat to update the flow stat report
+/// and frontned use get to get the latest flow stat report.
+pub struct FlowStateManager {
+    in_memory: KvBackendRef,
+}
+
+impl FlowStateManager {
+    pub fn new(in_memory: KvBackendRef) -> Self {
+        Self { in_memory }
+    }
+
+    pub async fn get(&self) -> Result<Option<FlowStateValue>> {
+        let key = FlowStateKey::new().to_bytes();
+        self.in_memory
+            .get(&key)
+            .await?
+            .map(|x| FlowStateValue::try_from_raw_value(&x.value))
+            .transpose()
+    }
+
+    pub async fn put(&self, value: FlowStateValue) -> Result<()> {
+        let key = FlowStateKey::new().to_bytes();
+        let value = value.try_as_raw_value()?;
+        let req = PutRequest::new().with_key(key).with_value(value);
+        self.in_memory.put(req).await?;
+        Ok(())
+    }
+}
+
+/// Flow's state report, send regularly through heartbeat message
+#[derive(Debug, Clone)]
+pub struct FlowStat {
+    /// For each key, the bytes of the state in memory
+    pub state_size: BTreeMap<u32, usize>,
+}
+
+impl From<FlowStateValue> for FlowStat {
+    fn from(value: FlowStateValue) -> Self {
+        Self {
+            state_size: value.state_size,
+        }
+    }
+}
+
+impl From<FlowStat> for FlowStateValue {
+    fn from(value: FlowStat) -> Self {
+        Self {
+            state_size: value.state_size,
+        }
+    }
+}

src/common/meta/src/kv_backend.rs

@@ -36,7 +36,7 @@ pub mod postgres;
 pub mod test;
 pub mod txn;
 
-pub type KvBackendRef = Arc<dyn KvBackend<Error = Error> + Send + Sync>;
+pub type KvBackendRef<E = Error> = Arc<dyn KvBackend<Error = E> + Send + Sync>;
 
 #[async_trait]
 pub trait KvBackend: TxnService
@@ -161,6 +161,9 @@ where
     Self::Error: ErrorExt,
 {
     fn reset(&self);
+
+    /// Upcast as `KvBackendRef`. Since https://github.com/rust-lang/rust/issues/65991 is not yet stable.
+    fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef<Self::Error>;
 }
 
-pub type ResettableKvBackendRef = Arc<dyn ResettableKvBackend<Error = Error> + Send + Sync>;
+pub type ResettableKvBackendRef<E = Error> = Arc<dyn ResettableKvBackend<Error = E> + Send + Sync>;

src/common/meta/src/kv_backend/memory.rs

@@ -16,13 +16,13 @@ use std::any::Any;
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 use std::marker::PhantomData;
-use std::sync::RwLock;
+use std::sync::{Arc, RwLock};
 
 use async_trait::async_trait;
 use common_error::ext::ErrorExt;
 use serde::Serializer;
 
-use super::ResettableKvBackend;
+use super::{KvBackendRef, ResettableKvBackend};
 use crate::kv_backend::txn::{Txn, TxnOp, TxnOpResponse, TxnRequest, TxnResponse};
 use crate::kv_backend::{KvBackend, TxnService};
 use crate::metrics::METRIC_META_TXN_REQUEST;
@@ -311,6 +311,10 @@ impl<T: ErrorExt + Send + Sync + 'static> ResettableKvBackend for MemoryKvBacken
     fn reset(&self) {
         self.clear();
     }
+
+    fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef<T> {
+        self
+    }
 }
 
 #[cfg(test)]