feat: improve with review comments

Signed-off-by: etorreborre <etorreborre@yahoo.com>

Author: etorreborre

crates/amaru/src/stages/consensus/forward_chain/tcp_forward_chain_server.rs

@@ -59,7 +59,7 @@ impl TcpForwardChainServer {
     }
 
     /// Creates a new TcpForwardChainServer instance with a provided Acto runtime and TcpListener.
-    #[expect(clippy::unwrap_used)]
+    #[expect(clippy::expect_used)]
     pub fn create(
         store: Arc<dyn ChainStore<Header>>,
         tcp_listener: TcpListener,
@@ -92,7 +92,10 @@ impl TcpForwardChainServer {
                 // in particular, it isn’t possible to poll for new peers within the `schedule` method
                 match PeerServer::accept(&tcp_listener, network_magic).await {
                     Ok(peer) => {
-                        let our_tip = our_tip_clone.lock().unwrap().clone();
+                        let our_tip = our_tip_clone
+                            .lock()
+                            .expect("poisoned lock for our tip")
+                            .clone();
                         clients_clone.send(ClientMsg::Peer(
                             peer,
                             Tip(our_tip.point().pallas_point(), our_tip.block_height()),
@@ -123,11 +126,14 @@ impl ForwardEventListener for TcpForwardChainServer {
     async fn send(&self, event: ForwardEvent) -> anyhow::Result<()> {
         match event {
             ForwardEvent::Forward(header) => {
-                let mut our_tip = self
-                    .our_tip
-                    .lock()
-                    .map_err(|e| anyhow::anyhow!("Mutex poisoned: {}", e))?;
-                *our_tip = header.as_header_tip();
+                {
+                    let mut our_tip = self
+                        .our_tip
+                        .lock()
+                        .map_err(|e| anyhow::anyhow!("Mutex poisoned: {}", e))?;
+                    *our_tip = header.as_header_tip();
+                };
+
                 self.clients.send(ClientMsg::Op(ClientOp::Forward(header)));
                 Ok(())
             }

simulation/amaru-sim/src/lib.rs

@@ -11,6 +11,7 @@
 // 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.
+#![feature(generic_atomic)]
 
 pub mod echo;
 pub mod simulator;

simulation/amaru-sim/src/simulator/run.rs

@@ -51,9 +51,10 @@ use pure_stage::trace_buffer::TraceBuffer;
 use pure_stage::{Instant, Receiver, StageGraph, StageRef};
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 use tokio::runtime::Runtime;
-use tokio::sync::{Mutex, mpsc};
+use tokio::sync::mpsc;
 use tracing::{Span, info};
 
 /// Run the full simulation:
@@ -294,7 +295,7 @@ pub struct MockForwardEventListener {
     node_id: String,
     number_of_downstream_peers: u8,
     sender: mpsc::Sender<Envelope<ChainSyncMessage>>,
-    msg_id: Arc<Mutex<u64>>,
+    msg_id: Arc<AtomicU64>,
 }
 
 impl MockForwardEventListener {
@@ -306,7 +307,7 @@ impl MockForwardEventListener {
         Self {
             node_id,
             number_of_downstream_peers,
-            msg_id: Arc::new(Mutex::new(0)),
+            msg_id: Arc::new(AtomicU64::new(0)),
             sender,
         }
     }
@@ -337,16 +338,20 @@ impl ForwardEventListener for MockForwardEventListener {
             }
         }
 
+        // This allocates a range of message ids from
+        // self.msg_id to self.msg_id + number_of_downstream_peers
+        let base_msg_id = self
+            .msg_id
+            .fetch_add(self.number_of_downstream_peers as u64, Ordering::Relaxed);
+
         for i in 1..=self.number_of_downstream_peers {
-            let envelope = {
-                let mut msg_id = self.msg_id.lock().await;
-                let envelope = Envelope {
-                    src: self.node_id.clone(),
-                    dest: format!("c{}", i),
-                    body: message(&event, *msg_id),
-                };
-                *msg_id += 1;
-                envelope
+            let dest = format!("c{}", i);
+            let msg_id = base_msg_id + i as u64;
+            println!("msg id {}", msg_id);
+            let envelope = Envelope {
+                src: self.node_id.clone(),
+                dest,
+                body: message(&event, msg_id),
             };
             self.sender.send(envelope).await?;
         }