Snap for 11975806 from 30aef144b20f17f8ac0aaeb0a694340f259fdedc to 24Q3-release Change-Id: Ida9f40aee56f25635124200a4f17ffc2a27d5256

diff --git a/rust/daemon/src/wireless/uwb.rs b/rust/daemon/src/wireless/uwb.rs
index 4fceb11..426dbb7 100644
--- a/rust/daemon/src/wireless/uwb.rs
+++ b/rust/daemon/src/wireless/uwb.rs

@@ -25,7 +25,7 @@
 use crate::uwb::ranging_estimator::{SharedState, UwbRangingEstimator};
 use crate::wireless::packet::handle_response;
 
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
 use std::sync::{Arc, Mutex};
 use std::thread;
 
@@ -52,8 +52,8 @@
     pica_id: Handle,
     uci_stream_writer: UnboundedSender<Vec<u8>>,
     state: AtomicBool,
-    tx_count: i32,
-    rx_count: i32, // TODO(b/330788870): Increment rx_count after handle_response
+    tx_count: AtomicI32,
+    rx_count: Arc<AtomicI32>,
 }
 
 impl Drop for Uwb {
@@ -68,21 +68,21 @@
         self.uci_stream_writer
             .unbounded_send(packet.clone().into())
             .expect("UciStream Receiver Disconnected");
+        let _ = self.tx_count.fetch_add(1, Ordering::SeqCst);
     }
 
     fn reset(&self) {
         self.state.store(true, Ordering::SeqCst);
-        // TODO: Reset packet counts
-        // self.tx_count = 0;
-        // self.rx_count = 0;
+        self.tx_count.store(0, Ordering::SeqCst);
+        self.rx_count.store(0, Ordering::SeqCst);
     }
 
     fn get(&self) -> ProtoChip {
         let mut chip_proto = ProtoChip::new();
         let uwb_proto = ProtoRadio {
             state: self.state.load(Ordering::SeqCst).into(),
-            tx_count: self.tx_count,
-            rx_count: self.rx_count,
+            tx_count: self.tx_count.load(Ordering::SeqCst),
+            rx_count: self.rx_count.load(Ordering::SeqCst),
             ..Default::default()
         };
         chip_proto.mut_uwb().clone_from(&uwb_proto);
@@ -130,12 +130,14 @@
         .add_device(Box::pin(uci_stream_receiver), Box::pin(uci_sink_sender.sink_err_into()))
         .unwrap();
     PICA_HANDLE_TO_STATE.insert(pica_id, chip_id);
+
+    let rx_count = Arc::new(AtomicI32::new(0));
     let uwb = Uwb {
         pica_id,
         uci_stream_writer: uci_stream_sender,
         state: AtomicBool::new(true),
-        tx_count: 0,
-        rx_count: 0,
+        tx_count: AtomicI32::new(0),
+        rx_count: rx_count.clone(),
     };
 
     // Spawn a future for obtaining packet from pica and invoking handle_response_rust
@@ -143,6 +145,7 @@
         let mut uci_sink_receiver = uci_sink_receiver;
         while let Some(packet) = uci_sink_receiver.next().await {
             handle_response(chip_id, &Bytes::from(packet));
+            rx_count.fetch_add(1, Ordering::SeqCst);
         }
     });
     Box::new(uwb)