From 9f7a37eefda7cc96156001fa83d7b740db43811a Mon Sep 17 00:00:00 2001
From: Luca Cominardi <luca.cominardi@gmail.com>
Date: Wed, 6 Dec 2023 13:08:49 +0100
Subject: [PATCH] Fix stats feature (#614)

---
 .github/workflows/ci.yml                         | 4 ++--
 io/zenoh-transport/src/common/batch.rs           | 8 ++++++++
 io/zenoh-transport/src/multicast/link.rs         | 4 ++--
 io/zenoh-transport/src/unicast/universal/link.rs | 2 +-
 zenoh/src/net/routing/pubsub.rs                  | 2 +-
 zenoh/src/net/routing/queries.rs                 | 4 ++--
 6 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 60fb69a576..42f32fbef0 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -54,11 +54,11 @@ jobs:
         env:
           CARGO_REGISTRIES_CRATES_IO_PROTOCOL: sparse
 
-      - name: Clippy unstable
+      - name: Clippy unstable and stats
         uses: actions-rs/cargo@v1
         with:
           command: clippy
-          args: --all-targets --features unstable -- -D warnings
+          args: --all-targets --features unstable --features stats -- -D warnings
         env:
           CARGO_REGISTRIES_CRATES_IO_PROTOCOL: sparse
 
diff --git a/io/zenoh-transport/src/common/batch.rs b/io/zenoh-transport/src/common/batch.rs
index cd029a9435..5305d0a50c 100644
--- a/io/zenoh-transport/src/common/batch.rs
+++ b/io/zenoh-transport/src/common/batch.rs
@@ -182,6 +182,10 @@ impl WBatch {
     pub fn clear(&mut self) {
         self.buffer.clear();
         self.codec.clear();
+        #[cfg(feature = "stats")]
+        {
+            self.stats.clear();
+        }
         if let Some(h) = self.header.get() {
             let mut writer = self.buffer.writer();
             let _ = writer.write_u8(h.get());
@@ -314,6 +318,10 @@ impl RBatch {
         }
     }
 
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
     #[inline(always)]
     pub const fn is_empty(&self) -> bool {
         self.buffer.is_empty()
diff --git a/io/zenoh-transport/src/multicast/link.rs b/io/zenoh-transport/src/multicast/link.rs
index 937216dd08..fbb917c281 100644
--- a/io/zenoh-transport/src/multicast/link.rs
+++ b/io/zenoh-transport/src/multicast/link.rs
@@ -477,7 +477,7 @@ async fn tx_task(
                 #[cfg(feature = "stats")]
                 {
                     stats.inc_tx_t_msgs(batch.stats.t_msgs);
-                    stats.inc_tx_bytes(bytes.len());
+                    stats.inc_tx_bytes(batch.len() as usize);
                 }
                 // Reinsert the batch into the queue
                 pipeline.refill(batch, priority);
@@ -595,7 +595,7 @@ async fn rx_task(
         match action {
             Action::Read((batch, locator)) => {
                 #[cfg(feature = "stats")]
-                transport.stats.inc_rx_bytes(zslice.len());
+                transport.stats.inc_rx_bytes(batch.len());
 
                 // Deserialize all the messages from the current ZBuf
                 transport.read_messages(
diff --git a/io/zenoh-transport/src/unicast/universal/link.rs b/io/zenoh-transport/src/unicast/universal/link.rs
index 3a7eafbc52..74db7f751e 100644
--- a/io/zenoh-transport/src/unicast/universal/link.rs
+++ b/io/zenoh-transport/src/unicast/universal/link.rs
@@ -184,7 +184,7 @@ async fn tx_task(
                     #[cfg(feature = "stats")]
                     {
                         stats.inc_tx_t_msgs(batch.stats.t_msgs);
-                        stats.inc_tx_bytes(bytes.len());
+                        stats.inc_tx_bytes(batch.len() as usize);
                     }
 
                     // Reinsert the batch into the queue
diff --git a/zenoh/src/net/routing/pubsub.rs b/zenoh/src/net/routing/pubsub.rs
index 75a49f1137..e9dc80b024 100644
--- a/zenoh/src/net/routing/pubsub.rs
+++ b/zenoh/src/net/routing/pubsub.rs
@@ -1734,7 +1734,7 @@ macro_rules! inc_stats {
     ) => {
         paste::paste! {
             if let Some(stats) = $face.stats.as_ref() {
-                use zenoh_buffers::SplitBuffer;
+                use zenoh_buffers::buffer::Buffer;
                 match &$body {
                     PushBody::Put(p) => {
                         stats.[<$txrx _z_put_msgs>].[<inc_ $space>](1);
diff --git a/zenoh/src/net/routing/queries.rs b/zenoh/src/net/routing/queries.rs
index a9d5ffd698..06b81a998b 100644
--- a/zenoh/src/net/routing/queries.rs
+++ b/zenoh/src/net/routing/queries.rs
@@ -2004,7 +2004,7 @@ macro_rules! inc_req_stats {
     ) => {
         paste::paste! {
             if let Some(stats) = $face.stats.as_ref() {
-                use zenoh_buffers::SplitBuffer;
+                use zenoh_buffers::buffer::Buffer;
                 match &$body {
                     RequestBody::Put(p) => {
                         stats.[<$txrx _z_put_msgs>].[<inc_ $space>](1);
@@ -2036,7 +2036,7 @@ macro_rules! inc_res_stats {
     ) => {
         paste::paste! {
             if let Some(stats) = $face.stats.as_ref() {
-                use zenoh_buffers::SplitBuffer;
+                use zenoh_buffers::buffer::Buffer;
                 match &$body {
                     ResponseBody::Put(p) => {
                         stats.[<$txrx _z_put_msgs>].[<inc_ $space>](1);