feat(meta-service): add expirable keys stat to metrics and cluster-status API (#18179)

This commit counts the user keys and the expirable keys when building a
snapshot, and save these two number in the snapshot meta data.

These two values are:
- reported to metrics(http API `/v1/metrics`) under the keys:
  - `snapshot_primary_index_count`
  - `snapshot_expire_index_count`
- and can be retrieved in the cluster-status http
  API(`/v1/cluster/status`), in the key `snapshot_key_space_stat`, which
  is a `BTreeMap` of `String` to `u64`:
  ```
  "snapshot_key_space_stat": {
    "exp-": 1,
    "kv--": 5
  },
  ```

Also refactors MetaKvApi implementation into a standalone struct for
better separation of concerns.

This commit is Non-Breaking despite on-disk format modification.

Author: drmingdrmer

src/meta/raft-store/src/leveled_store/key_spaces_impl.rs

@@ -21,10 +21,11 @@ use std::io::Error;
 use crate::leveled_store::map_api::MapKey;
 use crate::leveled_store::map_api::MapKeyDecode;
 use crate::leveled_store::map_api::MapKeyEncode;
+use crate::leveled_store::map_api::MapKeyPrefix;
 use crate::state_machine::ExpireKey;
 
 impl MapKeyEncode for String {
-    const PREFIX: &'static str = "kv--";
+    const PREFIX: MapKeyPrefix = "kv--";
 
     fn encode<W: fmt::Write>(&self, mut w: W) -> Result<(), fmt::Error> {
         w.write_str(self.as_str())
@@ -38,7 +39,7 @@ impl MapKeyDecode for String {
 }
 
 impl MapKeyEncode for ExpireKey {
-    const PREFIX: &'static str = "exp-";
+    const PREFIX: MapKeyPrefix = "exp-";
 
     fn encode<W: fmt::Write>(&self, mut w: W) -> Result<(), fmt::Error> {
         // max u64 len is 20: 18446744073709551616

src/meta/raft-store/src/leveled_store/leveled_map/compacting_data.rs

@@ -88,6 +88,7 @@ impl<'a> CompactingData<'a> {
     ///
     /// It returns a small chunk of sys data that is always copied across levels,
     /// and a stream contains `kv` and `expire` entries.
+    /// The stream Item is 2 items tuple of key, and value with seq.
     ///
     /// The exported stream contains encoded `String` key and rotbl value [`SeqMarked`]
     pub async fn compact(
@@ -109,15 +110,18 @@ impl<'a> CompactingData<'a> {
 
         let strm = (*self.immutable_levels).expire_map().range(..).await?;
         let expire_strm = strm.map(|item: Result<(ExpireKey, Marked<String>), io::Error>| {
-            let (k, v) = item?;
-            RotblCodec::encode_key_seq_marked(&k, v).map_err(|e| with_context(e, &k))
+            let (expire_key, marked_string) = item?;
+
+            RotblCodec::encode_key_seq_marked(&expire_key, marked_string)
+                .map_err(|e| with_context(e, &expire_key))
         });
 
         // kv: prefix: `kv--/`
 
         let strm = (*self.immutable_levels).str_map().range(..).await?;
         let kv_strm = strm.map(|item: Result<(String, Marked), io::Error>| {
             let (k, v) = item?;
+
             RotblCodec::encode_key_seq_marked(&k, v).map_err(|e| with_context(e, &k))
         });
 

src/meta/raft-store/src/leveled_store/map_api.rs

@@ -29,9 +29,15 @@ pub use map_api::IOResultStream;
 use crate::marked::Marked;
 use crate::state_machine::ExpireKey;
 
+pub type MapKeyPrefix = &'static str;
+
 pub trait MapKeyEncode {
     /// PREFIX is the prefix of the key used to define key space in the on-disk storage.
-    const PREFIX: &'static str;
+    const PREFIX: MapKeyPrefix;
+
+    fn prefix(&self) -> MapKeyPrefix {
+        Self::PREFIX
+    }
 
     fn encode<W: Write>(&self, w: W) -> Result<(), fmt::Error>;
 }

src/meta/service/src/meta_service/meta_node.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::collections::BTreeSet;
 use std::net::Ipv4Addr;
 use std::sync::atomic::AtomicI32;
@@ -91,6 +92,7 @@ use crate::message::LeaveRequest;
 use crate::meta_service::errors::grpc_error_to_network_err;
 use crate::meta_service::forwarder::MetaForwarder;
 use crate::meta_service::meta_leader::MetaLeader;
+use crate::meta_service::meta_node_kv_api_impl::MetaKVApi;
 use crate::meta_service::meta_node_status::MetaNodeStatus;
 use crate::meta_service::watcher::DispatcherHandle;
 use crate::meta_service::watcher::WatchTypes;
@@ -466,6 +468,17 @@ impl MetaNode {
                 // metrics about server storage
                 server_metrics::set_raft_log_size(meta_node.get_raft_log_size().await);
                 server_metrics::set_snapshot_key_count(meta_node.get_snapshot_key_count().await);
+                {
+                    let stat = meta_node.get_snapshot_key_space_stat().await;
+
+                    server_metrics::set_snapshot_primary_index_count(
+                        stat.get("kv--").copied().unwrap_or_default(),
+                    );
+
+                    server_metrics::set_snapshot_expire_index_count(
+                        stat.get("exp-").copied().unwrap_or_default(),
+                    )
+                }
 
                 last_leader = mm.current_leader;
             }
@@ -874,6 +887,10 @@ impl MetaNode {
             .unwrap_or_default()
     }
 
+    async fn get_snapshot_key_space_stat(&self) -> BTreeMap<String, u64> {
+        self.raft_store.get_snapshot_key_space_stat().await
+    }
+
     pub async fn get_status(&self) -> Result<MetaNodeStatus, MetaError> {
         let voters = self
             .raft_store
@@ -892,6 +909,7 @@ impl MetaNode {
 
         let raft_log_status = self.get_raft_log_stat().await.into();
         let snapshot_key_count = self.get_snapshot_key_count().await;
+        let snapshot_key_space_stat = self.get_snapshot_key_space_stat().await;
 
         let metrics = self.raft.metrics().borrow().clone();
 
@@ -910,6 +928,7 @@ impl MetaNode {
             endpoint: endpoint.to_string(),
             raft_log: raft_log_status,
             snapshot_key_count,
+            snapshot_key_space_stat,
             state: format!("{:?}", metrics.state),
             is_leader: metrics.state == openraft::ServerState::Leader,
             current_term: metrics.current_term,
@@ -1175,6 +1194,11 @@ impl MetaNode {
 
         Ok(stream_sender)
     }
+
+    /// Get a kvapi::KVApi implementation.
+    pub fn kv_api(&self) -> MetaKVApi {
+        MetaKVApi::new(self)
+    }
 }
 
 pub(crate) fn event_filter_from_filter_type(filter_type: FilterType) -> EventFilter {

src/meta/service/src/meta_service/meta_node_kv_api_impl.rs

@@ -34,18 +34,24 @@ use crate::message::ForwardRequest;
 use crate::meta_service::MetaNode;
 use crate::metrics::server_metrics;
 
-/// Impl kvapi::KVApi for MetaNode.
-///
-/// Write through raft-log.
-/// Read through local state machine, which may not be consistent.
-/// E.g. Read is not guaranteed to see a write.
+/// A wrapper of MetaNode that implements kvapi::KVApi.
+pub struct MetaKVApi<'a> {
+    inner: &'a MetaNode,
+}
+
+impl<'a> MetaKVApi<'a> {
+    pub fn new(inner: &'a MetaNode) -> Self {
+        Self { inner }
+    }
+}
+
 #[async_trait]
-impl kvapi::KVApi for MetaNode {
+impl<'a> kvapi::KVApi for MetaKVApi<'a> {
     type Error = MetaAPIError;
 
     async fn upsert_kv(&self, act: UpsertKV) -> Result<UpsertKVReply, Self::Error> {
         let ent = LogEntry::new(Cmd::UpsertKV(act));
-        let rst = self.write(ent).await?;
+        let rst = self.inner.write(ent).await?;
 
         match rst {
             AppliedState::KV(x) => Ok(x),
@@ -62,6 +68,7 @@ impl kvapi::KVApi for MetaNode {
         };
 
         let res = self
+            .inner
             .handle_forwardable_request(ForwardRequest::new(1, MetaGrpcReadReq::MGetKV(req)))
             .await;
 
@@ -81,6 +88,7 @@ impl kvapi::KVApi for MetaNode {
         };
 
         let res = self
+            .inner
             .handle_forwardable_request(ForwardRequest::new(1, MetaGrpcReadReq::ListKV(req)))
             .await;
 
@@ -98,7 +106,7 @@ impl kvapi::KVApi for MetaNode {
         info!("MetaNode::transaction(): {}", txn);
 
         let ent = LogEntry::new(Cmd::Transaction(txn));
-        let rst = self.write(ent).await?;
+        let rst = self.inner.write(ent).await?;
 
         match rst {
             AppliedState::TxnReply(x) => Ok(x),
@@ -108,3 +116,32 @@ impl kvapi::KVApi for MetaNode {
         }
     }
 }
+
+/// Impl kvapi::KVApi for MetaNode.
+///
+/// Write through raft-log.
+/// Read through local state machine, which may not be consistent.
+/// E.g. Read is not guaranteed to see a write.
+#[async_trait]
+impl kvapi::KVApi for MetaNode {
+    type Error = MetaAPIError;
+
+    async fn upsert_kv(&self, act: UpsertKV) -> Result<UpsertKVReply, Self::Error> {
+        self.kv_api().upsert_kv(act).await
+    }
+
+    #[fastrace::trace]
+    async fn get_kv_stream(&self, keys: &[String]) -> Result<KVStream<Self::Error>, Self::Error> {
+        self.kv_api().get_kv_stream(keys).await
+    }
+
+    #[fastrace::trace]
+    async fn list_kv(&self, prefix: &str) -> Result<KVStream<Self::Error>, Self::Error> {
+        self.kv_api().list_kv(prefix).await
+    }
+
+    #[fastrace::trace]
+    async fn transaction(&self, txn: TxnRequest) -> Result<TxnReply, Self::Error> {
+        self.kv_api().transaction(txn).await
+    }
+}

src/meta/service/src/meta_service/meta_node_status.rs

@@ -20,7 +20,7 @@ use databend_common_meta_types::node::Node;
 use databend_common_meta_types::raft_types::LogId;
 use databend_common_meta_types::raft_types::NodeId;
 
-#[derive(serde::Serialize)]
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
 pub struct MetaNodeStatus {
     pub id: NodeId,
 
@@ -39,6 +39,11 @@ pub struct MetaNodeStatus {
     /// Total number of keys in current snapshot
     pub snapshot_key_count: u64,
 
+    /// The count of keys in each key space in the snapshot data.
+    #[serde(default)]
+    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
+    pub snapshot_key_space_stat: BTreeMap<String, u64>,
+
     /// Server state, one of "Follower", "Learner", "Candidate", "Leader".
     pub state: String,
 

src/meta/service/src/metrics/meta_metrics.rs

@@ -53,8 +53,18 @@ pub mod server_metrics {
         node_is_health: Gauge,
         leader_changes: Counter,
         applying_snapshot: Gauge,
+
+        /// Primary index is index by string key. Each primary index has an optional expire index key.
+        ///
+        /// `snapshot_key_count = snapshot_primary_index_count + snapshot_expire_index_count`
         snapshot_key_count: Gauge,
 
+        /// `snapshot_key_count = snapshot_primary_index_count + snapshot_expire_index_count`
+        snapshot_primary_index_count: Gauge,
+
+        /// `snapshot_key_count = snapshot_primary_index_count + snapshot_expire_index_count`
+        snapshot_expire_index_count: Gauge,
+
         raft_log_cache_items: Gauge,
         raft_log_cache_used_size: Gauge,
         raft_log_wal_open_chunk_size: Gauge,
@@ -83,6 +93,9 @@ pub mod server_metrics {
                 leader_changes: Counter::default(),
                 applying_snapshot: Gauge::default(),
                 snapshot_key_count: Gauge::default(),
+                snapshot_primary_index_count: Gauge::default(),
+                snapshot_expire_index_count: Gauge::default(),
+
                 raft_log_cache_items: Gauge::default(),
                 raft_log_cache_used_size: Gauge::default(),
                 raft_log_wal_open_chunk_size: Gauge::default(),
@@ -128,6 +141,16 @@ pub mod server_metrics {
                 "number of keys in the last snapshot",
                 metrics.snapshot_key_count.clone(),
             );
+            registry.register(
+                key!("snapshot_primary_index_count"),
+                "number of primary keys in the last snapshot",
+                metrics.snapshot_primary_index_count.clone(),
+            );
+            registry.register(
+                key!("snapshot_expire_index_count"),
+                "number of expire index keys in the last snapshot",
+                metrics.snapshot_expire_index_count.clone(),
+            );
 
             registry.register(
                 key!("raft_log_cache_items"),
@@ -228,6 +251,12 @@ pub mod server_metrics {
     pub fn set_snapshot_key_count(n: u64) {
         SERVER_METRICS.snapshot_key_count.set(n as i64);
     }
+    pub fn set_snapshot_primary_index_count(n: u64) {
+        SERVER_METRICS.snapshot_primary_index_count.set(n as i64);
+    }
+    pub fn set_snapshot_expire_index_count(n: u64) {
+        SERVER_METRICS.snapshot_expire_index_count.set(n as i64);
+    }
 
     pub fn set_raft_log_stat(st: RaftLogStat) {
         SERVER_METRICS

src/meta/service/src/store/store_inner.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::fs;
 use std::io;
 use std::io::ErrorKind;
@@ -192,7 +193,7 @@ impl RaftStoreInner {
             w.acquire_compactor().await
         };
 
-        let (sys_data, mut strm) = compactor
+        let (mut sys_data, mut strm) = compactor
             .compact()
             .await
             .map_err(|e| StorageError::read_snapshot(None, &e))?;
@@ -224,6 +225,16 @@ impl RaftStoreInner {
                 .await
                 .map_err(|e| StorageError::read_snapshot(None, &e))?
             {
+                // The first 4 chars are key space, such as: "kv--/" or "exp-/"
+                // Get the first 4 chars as key space.
+                let prefix = &ent.0.as_str()[..4];
+                let ks = sys_data.key_counts_mut();
+                if let Some(count) = ks.get_mut(prefix) {
+                    *count += 1;
+                } else {
+                    ks.insert(prefix.to_string(), 1);
+                }
+
                 tx.send(WriteEntry::Data(ent))
                     .await
                     .map_err(|e| StorageError::write_snapshot(Some(signature.clone()), &e))?;
@@ -283,6 +294,21 @@ impl RaftStoreInner {
         Some(db.stat().key_num)
     }
 
+    /// Returns the count of keys in each key space from the snapshot data.
+    ///
+    /// Key spaces include:
+    /// - `"exp-"`: expire index data
+    /// - `"kv--"`: key-value data
+    ///
+    /// Returns an empty map if no snapshot exists.
+    pub(crate) async fn get_snapshot_key_space_stat(&self) -> BTreeMap<String, u64> {
+        let sm = self.state_machine.read().await;
+        let Some(db) = sm.levels().persisted() else {
+            return Default::default();
+        };
+        db.sys_data().key_counts().clone()
+    }
+
     /// Install a snapshot to build a state machine from it and replace the old state machine with the new one.
     #[fastrace::trace]
     pub async fn do_install_snapshot(&self, db: DB) -> Result<(), MetaStorageError> {