feat(storage,runtime): add request-local storage stats collection (#393)

* feat(storage,runtime): connect real StorageStatsCollector at storage I/O points

Replace the NoopStorageStatsCollector stub with a per-request
RealStorageStatsCollector installed via the STORAGE_STATS_COLLECTOR
task-local, so per-request storage metrics (keys/bytes read, written,
deleted, cache hit, compaction level) are recorded at the actual I/O
points instead of being inferred from command arguments.

- Move StorageStats / StorageStatsCollector trait into the client leaf
  crate to avoid a storage <-> runtime dependency cycle; runtime::message
  re-exports them for backward compatibility.
- Instrument the Storage facade methods in storage_impl.rs (the single
  convergence layer the cmd layer calls), requiring zero Cmd::execute
  signature changes.
- Add StorageStats.bytes_deleted per issue #312.

Fixes #312

* fix(storage): correct request-local stats accounting

Record logical storage statistics only after backend outcomes are known, aggregate completed response statistics in the client and runtime metrics, and sample request-local RocksDB block-cache hits. Add fault-injection coverage for failed and partial multi-shard mutations, while keeping new serialized fields and optional collector methods backward compatible.

Constraint: Limit this follow-up to PR #393 CI failures, CodeRabbit findings, request-local statistics correctness, and directly related regression tests.

Confidence: High; failure, no-op, partial-success, aggregation, reset, serialization, and trait-compatibility paths have direct tests.

Scope-risk: Storage statistics instrumentation and runtime metrics only; uninstrumented Storage facade methods and request-attributable compaction levels remain explicitly outside this incremental PR.

Red: The previous PR Head failed license header, rustfmt, and all three Clippy jobs; statistics were charged before failed or no-op mutations and lacked storage regression coverage.

Green: WSL cargo test -p client; cargo test -p runtime; cargo test -p storage --features test-fault-injection --test storage_stats_test -- --test-threads=1 --nocapture; cargo fmt --all -- --check; CI-equivalent workspace Clippy; git diff --cached --check.

Tested: client 6/6, runtime 135/135, storage statistics 6/6, exact Rust license-header comparison, rustfmt, workspace all-features Clippy, and staged diff validation.

Not-tested: GitHub's macOS and Windows Clippy runners and SkyWalking Eyes action will run after push.

Co-authored-by: OmX <omx@oh-my-codex.dev>

---------

Co-authored-by: OmX <omx@oh-my-codex.dev>

Author: AlexStocks

Cargo.lock

@@ -11,4 +11,8 @@ rust-version.workspace = true
 async-trait = "0.1"
 resp = { path = "../resp" }
 parking_lot = { workspace = true }
-tokio = { version = "1.50", features = ["sync"] }
+serde = { workspace = true, features = ["derive"] }
+tokio = { version = "1.50", features = ["sync", "rt"] }
+
+[dev-dependencies]
+serde_json = { workspace = true }

src/client/Cargo.toml

@@ -21,6 +21,9 @@ use async_trait::async_trait;
 use resp::{ProtocolNegotiator, RespCommand, RespData, RespResult};
 use tokio::sync::Mutex;
 
+pub mod storage_stats;
+pub use storage_stats::*;
+
 #[async_trait]
 pub trait StreamTrait: Send + Sync {
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error>;

src/client/src/lib.rs

@@ -0,0 +1,327 @@
+// Copyright (c) 2024-present, arana-db Community.  All rights reserved.
+//
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You 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.
+
+//! Request-local storage-layer instrumentation.
+//!
+//! These types were originally declared in `runtime::message` (see
+//! <https://github.com/arana-db/kiwi/issues/312>). They live in the `client`
+//! crate so that `storage`, `runtime` and `cmd` can share them without
+//! creating a dependency cycle (`client` is a leaf crate that everything else
+//! already depends on).
+//!
+//! The real collector ([`RealStorageStatsCollector`]) is installed once per
+//! storage request via the [`STORAGE_STATS_COLLECTOR`] task-local and read by
+//! the storage facade after backend outcomes are known. Key and byte counters
+//! describe successful logical storage operations and their payload sizes;
+//! RocksDB-specific observations are reported separately when available.
+
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::{Arc, Mutex as StdMutex};
+
+use serde::{Deserialize, Serialize};
+
+/// Statistics about storage operations for monitoring.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
+#[serde(default)]
+pub struct StorageStats {
+    /// Number of keys read during the operation
+    pub keys_read: u64,
+    /// Number of keys written during the operation
+    pub keys_written: u64,
+    /// Number of keys deleted during the operation
+    pub keys_deleted: u64,
+    /// Logical payload read in bytes (key bytes + returned value bytes)
+    pub bytes_read: u64,
+    /// Logical payload written in bytes (key bytes + accepted value bytes)
+    pub bytes_written: u64,
+    /// Logical key bytes affected by successful delete operations
+    pub bytes_deleted: u64,
+    /// Whether the operation hit a cache (e.g. block/moka cache)
+    pub cache_hit: bool,
+    /// RocksDB compaction level accessed, when the engine can attribute one
+    pub compaction_level: Option<u32>,
+}
+
+impl StorageStats {
+    /// Merge another completed request into this aggregate without wrapping.
+    pub fn merge(&mut self, request: &Self) {
+        self.keys_read = self.keys_read.saturating_add(request.keys_read);
+        self.keys_written = self.keys_written.saturating_add(request.keys_written);
+        self.keys_deleted = self.keys_deleted.saturating_add(request.keys_deleted);
+        self.bytes_read = self.bytes_read.saturating_add(request.bytes_read);
+        self.bytes_written = self.bytes_written.saturating_add(request.bytes_written);
+        self.bytes_deleted = self.bytes_deleted.saturating_add(request.bytes_deleted);
+        self.cache_hit |= request.cache_hit;
+        if request.compaction_level.is_some() {
+            self.compaction_level = request.compaction_level;
+        }
+    }
+}
+
+/// Request-local collector for storage-layer instrumentation.
+///
+/// Implementations record sizes observed by the storage facade after the
+/// backend outcome is known, rather than inferring success in command parsing.
+pub trait StorageStatsCollector: Send + Sync {
+    /// Record a successful logical storage read.
+    fn record_read(&self, key_bytes: u64, value_bytes: u64);
+
+    /// Record a logical storage write after the backend accepts the mutation.
+    fn record_write(&self, key_bytes: u64, value_bytes: u64);
+
+    /// Record a logical delete after the backend confirms a mutation.
+    fn record_delete(&self, key_bytes: u64);
+
+    /// Record that this request was served from a storage cache.
+    fn record_cache_hit(&self) {}
+
+    /// Record the RocksDB compaction level that served this request.
+    fn record_compaction_level(&self, _level: u32) {}
+
+    /// Return the accumulated statistics for the request.
+    fn finish(&self) -> StorageStats;
+}
+
+/// Placeholder collector used until storage-layer instrumentation is wired in.
+#[derive(Debug, Default, Clone, Copy)]
+pub struct NoopStorageStatsCollector;
+
+impl StorageStatsCollector for NoopStorageStatsCollector {
+    fn record_read(&self, _key_bytes: u64, _value_bytes: u64) {}
+
+    fn record_write(&self, _key_bytes: u64, _value_bytes: u64) {}
+
+    fn record_delete(&self, _key_bytes: u64) {}
+
+    fn finish(&self) -> StorageStats {
+        StorageStats::default()
+    }
+}
+
+/// A real per-request collector backed by lock-free atomic counters.
+///
+/// Safe to share across the (single) task that executes one storage request:
+/// every field is an atomic and `finish` snapshots the accumulated values.
+pub struct RealStorageStatsCollector {
+    keys_read: AtomicU64,
+    keys_written: AtomicU64,
+    keys_deleted: AtomicU64,
+    bytes_read: AtomicU64,
+    bytes_written: AtomicU64,
+    bytes_deleted: AtomicU64,
+    cache_hit: AtomicBool,
+    compaction_level: StdMutex<Option<u32>>,
+}
+
+impl Default for RealStorageStatsCollector {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl RealStorageStatsCollector {
+    /// Create an empty collector ready to accumulate one request's stats.
+    pub fn new() -> Self {
+        Self {
+            keys_read: AtomicU64::new(0),
+            keys_written: AtomicU64::new(0),
+            keys_deleted: AtomicU64::new(0),
+            bytes_read: AtomicU64::new(0),
+            bytes_written: AtomicU64::new(0),
+            bytes_deleted: AtomicU64::new(0),
+            cache_hit: AtomicBool::new(false),
+            compaction_level: StdMutex::new(None),
+        }
+    }
+}
+
+impl StorageStatsCollector for RealStorageStatsCollector {
+    fn record_read(&self, key_bytes: u64, value_bytes: u64) {
+        self.keys_read.fetch_add(1, Ordering::Relaxed);
+        self.bytes_read
+            .fetch_add(key_bytes.saturating_add(value_bytes), Ordering::Relaxed);
+    }
+
+    fn record_write(&self, key_bytes: u64, value_bytes: u64) {
+        self.keys_written.fetch_add(1, Ordering::Relaxed);
+        self.bytes_written
+            .fetch_add(key_bytes.saturating_add(value_bytes), Ordering::Relaxed);
+    }
+
+    fn record_delete(&self, key_bytes: u64) {
+        self.keys_deleted.fetch_add(1, Ordering::Relaxed);
+        self.bytes_deleted.fetch_add(key_bytes, Ordering::Relaxed);
+    }
+
+    fn record_cache_hit(&self) {
+        self.cache_hit.store(true, Ordering::Relaxed);
+    }
+
+    fn record_compaction_level(&self, level: u32) {
+        *self
+            .compaction_level
+            .lock()
+            .expect("compaction_level poisoned") = Some(level);
+    }
+
+    fn finish(&self) -> StorageStats {
+        StorageStats {
+            keys_read: self.keys_read.load(Ordering::Relaxed),
+            keys_written: self.keys_written.load(Ordering::Relaxed),
+            keys_deleted: self.keys_deleted.load(Ordering::Relaxed),
+            bytes_read: self.bytes_read.load(Ordering::Relaxed),
+            bytes_written: self.bytes_written.load(Ordering::Relaxed),
+            bytes_deleted: self.bytes_deleted.load(Ordering::Relaxed),
+            cache_hit: self.cache_hit.load(Ordering::Relaxed),
+            compaction_level: *self
+                .compaction_level
+                .lock()
+                .expect("compaction_level poisoned"),
+        }
+    }
+}
+
+tokio::task_local! {
+    /// Per-request storage stats collector, scoped for the duration of a single
+    /// storage command execution. The storage layer reads it via
+    /// [`try_collector`].
+    pub static STORAGE_STATS_COLLECTOR: Arc<dyn StorageStatsCollector + Send + Sync>;
+}
+
+/// Returns the active request-local collector, if one was installed for the
+/// current async task.
+///
+/// Returns `None` when no collector is scoped (e.g. unit tests that call the
+/// storage engine directly, or background storage tasks), so callers can record
+/// unconditionally without special-casing those paths.
+pub fn try_collector() -> Option<Arc<dyn StorageStatsCollector + Send + Sync>> {
+    STORAGE_STATS_COLLECTOR.try_with(|c| c.clone()).ok()
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::unwrap_used)]
+
+    use super::*;
+
+    struct LegacyCollector;
+
+    impl StorageStatsCollector for LegacyCollector {
+        fn record_read(&self, _key_bytes: u64, _value_bytes: u64) {}
+
+        fn record_write(&self, _key_bytes: u64, _value_bytes: u64) {}
+
+        fn record_delete(&self, _key_bytes: u64) {}
+
+        fn finish(&self) -> StorageStats {
+            StorageStats::default()
+        }
+    }
+
+    #[test]
+    fn real_collector_accumulates_reads_writes_deletes() {
+        let c = RealStorageStatsCollector::new();
+        c.record_read(3, 5);
+        c.record_write(3, 7);
+        c.record_read(2, 4);
+        c.record_delete(3);
+
+        let stats = c.finish();
+        assert_eq!(stats.keys_read, 2);
+        assert_eq!(stats.keys_written, 1);
+        assert_eq!(stats.keys_deleted, 1);
+        assert_eq!(stats.bytes_read, (3 + 5) + (2 + 4));
+        assert_eq!(stats.bytes_written, 3 + 7);
+        assert_eq!(stats.bytes_deleted, 3);
+        assert!(!stats.cache_hit);
+    }
+
+    #[test]
+    fn noop_yields_empty_stats() {
+        assert_eq!(NoopStorageStatsCollector.finish(), StorageStats::default());
+    }
+
+    #[test]
+    fn legacy_collectors_use_default_optional_recorders() {
+        LegacyCollector.record_cache_hit();
+        LegacyCollector.record_compaction_level(3);
+        assert_eq!(LegacyCollector.finish(), StorageStats::default());
+    }
+
+    #[test]
+    fn legacy_serialized_stats_default_new_fields() {
+        let stats: StorageStats = serde_json::from_str(
+            r#"{
+                "keys_read": 1,
+                "keys_written": 2,
+                "keys_deleted": 3,
+                "bytes_read": 4,
+                "bytes_written": 5,
+                "cache_hit": true,
+                "compaction_level": null
+            }"#,
+        )
+        .unwrap();
+
+        assert_eq!(stats.bytes_deleted, 0);
+        assert_eq!(stats.keys_read, 1);
+        assert!(stats.cache_hit);
+    }
+
+    #[test]
+    fn trait_object_records_cache_and_compaction_details() {
+        let collector: Arc<dyn StorageStatsCollector + Send + Sync> =
+            Arc::new(RealStorageStatsCollector::new());
+
+        collector.record_cache_hit();
+        collector.record_compaction_level(3);
+
+        let stats = collector.finish();
+        assert!(stats.cache_hit);
+        assert_eq!(stats.compaction_level, Some(3));
+    }
+
+    #[test]
+    fn scoped_collector_is_visible_to_try_collector() {
+        // Mirrors how `storage_server` installs the collector and how the
+        // storage facade reads it via `try_collector` — without needing a real
+        // RocksDB instance.
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+        runtime.block_on(async {
+            let collector: Arc<dyn StorageStatsCollector + Send + Sync> =
+                Arc::new(RealStorageStatsCollector::new());
+            STORAGE_STATS_COLLECTOR
+                .scope(Arc::clone(&collector), async {
+                    if let Some(c) = try_collector() {
+                        c.record_write(3, 5);
+                        c.record_read(2, 4);
+                    } else {
+                        panic!("try_collector() must return the scoped collector");
+                    }
+                })
+                .await;
+            let stats = collector.finish();
+            assert_eq!(stats.keys_written, 1);
+            assert_eq!(stats.bytes_written, 8);
+            assert_eq!(stats.keys_read, 1);
+            assert_eq!(stats.bytes_read, 6);
+        });
+    }
+}