feat(profiling): implement automatic stage profiling and reset functionality

Author: lxsaah

Cargo.lock

@@ -255,6 +255,8 @@ test-embedded:
 	cargo check --package aimdb-embassy-adapter --target thumbv7em-none-eabihf --no-default-features --features "embassy-runtime"
 	@printf "$(YELLOW)  → Checking aimdb-embassy-adapter with network support on thumbv7em-none-eabihf target$(NC)\n"
 	cargo check --package aimdb-embassy-adapter --target thumbv7em-none-eabihf --no-default-features --features "embassy-runtime,embassy-net-support"
+	@printf "$(YELLOW)  → Checking aimdb-embassy-adapter with profiling on thumbv7em-none-eabihf target$(NC)\n"
+	cargo check --package aimdb-embassy-adapter --target thumbv7em-none-eabihf --no-default-features --features "embassy-runtime,profiling"
 	@printf "$(YELLOW)  → Checking aimdb-mqtt-connector (Embassy) on thumbv7em-none-eabihf target$(NC)\n"
 	cargo check --package aimdb-mqtt-connector --target thumbv7em-none-eabihf --no-default-features --features "embassy-runtime"
 	@printf "$(YELLOW)  → Checking aimdb-mqtt-connector (Embassy + defmt) on thumbv7em-none-eabihf target$(NC)\n"

Makefile

@@ -124,6 +124,14 @@ impl AimxClient {
         Ok(records)
     }
 
+    /// Reset stage profiling counters for every record on the server.
+    ///
+    /// Requires the server to be built with the `profiling` feature and the
+    /// connection to have write permission.
+    pub async fn reset_stage_profiling(&mut self) -> ClientResult<serde_json::Value> {
+        self.send_request("profiling.reset", None).await
+    }
+
     /// Get current value of a record
     pub async fn get_record(&mut self, name: &str) -> ClientResult<serde_json::Value> {
         let params = json!({ "record": name });

aimdb-client/src/connection.rs

@@ -38,7 +38,10 @@ metrics = ["dep:metrics", "std"] # Requires std for aggregation
 
 # Automatic stage profiling (.source()/.tap()/.link() timing).
 # Independent of `metrics`; works in no_std (only needs heap + a runtime clock).
-profiling = ["alloc"]
+# `portable-atomic/critical-section` provides the 64-bit-atomic fallback on targets
+# without native 64-bit atomics (e.g. thumbv7em); it's a no-op where native atomics
+# exist, and embedded binaries already supply a `critical-section` impl.
+profiling = ["alloc", "portable-atomic/critical-section"]
 
 # Testing features
 test-utils = ["std"]

aimdb-core/Cargo.toml

@@ -1265,6 +1265,14 @@ impl<R: aimdb_executor::Spawn + 'static> AimDb<R> {
         self.inner.list_records()
     }
 
+    /// Resets stage profiling counters for every record (feature `profiling`).
+    #[cfg(feature = "profiling")]
+    pub fn reset_stage_profiling(&self) {
+        for record in &self.inner.storages {
+            record.reset_profiling();
+        }
+    }
+
     /// Try to get record's latest value as JSON by name (std only)
     ///
     /// Convenience wrapper around `AimDbInner::try_latest_as_json()`.

aimdb-core/src/builder.rs

@@ -84,7 +84,7 @@ pub use typed_record::{AnyRecord, AnyRecordExt, TypedRecord};
 
 // Stage profiling exports (feature-gated)
 #[cfg(feature = "profiling")]
-pub use profiling::{RecordProfilingMetrics, StageMetrics};
+pub use profiling::{RecordProfilingMetrics, StageMetrics, StageProfilingInfo};
 
 // Connector Infrastructure exports
 pub use connector::TopicResolverFn;

aimdb-core/src/lib.rs

@@ -0,0 +1,67 @@
+//! Serializable snapshot of stage profiling metrics, for remote introspection.
+
+extern crate alloc;
+use alloc::{string::String, vec::Vec};
+
+use serde::{Deserialize, Serialize};
+
+use crate::profiling::{RecordProfilingMetrics, StageEntry};
+
+/// A point-in-time snapshot of one execution stage's timing metrics.
+///
+/// Carried in `RecordMetadata::stage_profiling` and exposed by the
+/// `get_stage_profiling` MCP tool. All times are wall-clock nanoseconds.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct StageProfilingInfo {
+    /// Stage kind: `"source"`, `"tap"`, `"link"`, or `"transform"`.
+    pub stage_type: String,
+    /// Registration index within the stage kind (0-based).
+    pub index: usize,
+    /// Name assigned via `.with_name("...")`, if any.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub name: Option<String>,
+    /// Number of recorded invocations.
+    pub call_count: u64,
+    /// Cumulative wall-clock time, nanoseconds.
+    pub total_time_ns: u64,
+    /// Mean wall-clock time per invocation, nanoseconds.
+    pub avg_time_ns: u64,
+    /// Fastest recorded invocation, nanoseconds.
+    pub min_time_ns: u64,
+    /// Slowest recorded invocation, nanoseconds.
+    pub max_time_ns: u64,
+}
+
+impl StageProfilingInfo {
+    fn from_entry(stage_type: &str, index: usize, entry: &StageEntry) -> Self {
+        let m = &entry.metrics;
+        Self {
+            stage_type: String::from(stage_type),
+            index,
+            name: entry.name.clone(),
+            call_count: m.call_count(),
+            total_time_ns: m.total_time_ns(),
+            avg_time_ns: m.avg_time_ns(),
+            min_time_ns: m.min_time_ns(),
+            max_time_ns: m.max_time_ns(),
+        }
+    }
+}
+
+impl RecordProfilingMetrics {
+    /// Returns a serializable snapshot of every registered stage's metrics,
+    /// ordered sources → taps → links → transforms.
+    pub fn snapshot(&self) -> Vec<StageProfilingInfo> {
+        let mut out = Vec::new();
+        for (kind, stages) in [
+            ("source", self.sources()),
+            ("tap", self.taps()),
+            ("link", self.links()),
+        ] {
+            for (i, entry) in stages.iter().enumerate() {
+                out.push(StageProfilingInfo::from_entry(kind, i, entry));
+            }
+        }
+        out
+    }
+}

aimdb-core/src/profiling/info.rs

@@ -15,9 +15,11 @@
 //! answers *which* stage is slow, not *why*. Use `tracing` / `tokio-console` for
 //! CPU-vs-I/O analysis.
 
+mod info;
 mod record_profiling;
 mod stage_metrics;
 
+pub use info::StageProfilingInfo;
 pub use record_profiling::{RecordProfilingMetrics, StageEntry};
 pub use stage_metrics::StageMetrics;
 

aimdb-core/src/profiling/mod.rs

@@ -579,6 +579,8 @@ where
         "graph.nodes" => handle_graph_nodes(db, request.id).await,
         "graph.edges" => handle_graph_edges(db, request.id).await,
         "graph.topo_order" => handle_graph_topo_order(db, request.id).await,
+        #[cfg(feature = "profiling")]
+        "profiling.reset" => handle_profiling_reset(db, config, request.id).await,
         _ => {
             #[cfg(feature = "tracing")]
             tracing::warn!("Unknown method: {}", request.method);
@@ -625,6 +627,37 @@ where
     Response::success(request_id, json!(records))
 }
 
+/// Handles profiling.reset method
+///
+/// Clears stage profiling counters for every record. Requires write permission.
+#[cfg(all(feature = "std", feature = "profiling"))]
+async fn handle_profiling_reset<R>(
+    db: &Arc<AimDb<R>>,
+    config: &AimxConfig,
+    request_id: u64,
+) -> Response
+where
+    R: crate::RuntimeAdapter + crate::Spawn + 'static,
+{
+    if matches!(
+        config.security_policy,
+        crate::remote::SecurityPolicy::ReadOnly
+    ) {
+        return Response::error(
+            request_id,
+            "permission_denied",
+            "profiling.reset requires write permission (ReadOnly security policy)".to_string(),
+        );
+    }
+
+    db.reset_stage_profiling();
+
+    #[cfg(feature = "tracing")]
+    tracing::info!("Stage profiling counters reset");
+
+    Response::success(request_id, json!({ "reset": true }))
+}
+
 /// Handles record.get method
 ///
 /// Returns the current value of a record as JSON.

aimdb-core/src/remote/handler.rs

@@ -77,6 +77,13 @@ pub struct RecordMetadata {
     #[cfg(feature = "metrics")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub occupancy: Option<(usize, usize)>,
+
+    // ===== Stage profiling (feature-gated) =====
+    /// Per-stage timing metrics (`.source()`/`.tap()`/`.link()`), if the
+    /// `profiling` feature is enabled and any stage has been registered.
+    #[cfg(feature = "profiling")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stage_profiling: Option<std::vec::Vec<crate::profiling::StageProfilingInfo>>,
 }
 
 impl RecordMetadata {
@@ -132,6 +139,8 @@ impl RecordMetadata {
             dropped_count: None,
             #[cfg(feature = "metrics")]
             occupancy: None,
+            #[cfg(feature = "profiling")]
+            stage_profiling: None,
         }
     }
 
@@ -162,6 +171,18 @@ impl RecordMetadata {
         }
         self
     }
+
+    /// Attaches a stage profiling snapshot (profiling feature only).
+    #[cfg(feature = "profiling")]
+    pub fn with_stage_profiling(
+        mut self,
+        stages: std::vec::Vec<crate::profiling::StageProfilingInfo>,
+    ) -> Self {
+        if !stages.is_empty() {
+            self.stage_profiling = Some(stages);
+        }
+        self
+    }
 }
 
 #[cfg(test)]