From aff6ab115c5fe1d40bea21e4cd2d53d8c2634751 Mon Sep 17 00:00:00 2001
From: Matan Markind <matan@starkware.co>
Date: Mon, 24 Jun 2024 15:52:35 +0300
Subject: [PATCH] feat(consensus): add tracing instrumentation and logging to
 consensus

---
 crates/papyrus_node/src/main.rs                    |  2 +-
 crates/sequencing/papyrus_consensus/src/lib.rs     | 10 +++++-----
 .../src/single_height_consensus.rs                 | 14 ++++++++++++++
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/crates/papyrus_node/src/main.rs b/crates/papyrus_node/src/main.rs
index c66106f293..f082e34d9a 100644
--- a/crates/papyrus_node/src/main.rs
+++ b/crates/papyrus_node/src/main.rs
@@ -108,7 +108,7 @@ fn run_consensus(
     consensus_channels: BroadcastSubscriberChannels<ConsensusMessage>,
 ) -> anyhow::Result<JoinHandle<Result<(), ConsensusError>>> {
     let Ok(validator_id) = env::var("CONSENSUS_VALIDATOR_ID") else {
-        info!("CONSENSUS_VALIDATOR_ID is not set. Not run consensus.");
+        info!("CONSENSUS_VALIDATOR_ID is not set. Not running consensus.");
         return Ok(tokio::spawn(pending()));
     };
     info!("Running consensus as validator {validator_id}");
diff --git a/crates/sequencing/papyrus_consensus/src/lib.rs b/crates/sequencing/papyrus_consensus/src/lib.rs
index 763bb85773..0ec98cba3f 100644
--- a/crates/sequencing/papyrus_consensus/src/lib.rs
+++ b/crates/sequencing/papyrus_consensus/src/lib.rs
@@ -19,7 +19,7 @@ use papyrus_network::network_manager::SubscriberReceiver;
 use papyrus_protobuf::consensus::{ConsensusMessage, Proposal};
 use single_height_consensus::SingleHeightConsensus;
 use starknet_api::block::{BlockHash, BlockNumber};
-use tracing::info;
+use tracing::{debug, info, instrument};
 use types::{ConsensusBlock, ConsensusContext, ConsensusError, ProposalInit, ValidatorId};
 
 // TODO(matan): Remove dead code allowance at the end of milestone 1.
@@ -37,6 +37,7 @@ pub mod types;
 use futures::StreamExt;
 
 // TODO(dvir): add test for this.
+#[instrument(skip(context, start_height, network_receiver), level = "info")]
 #[allow(missing_docs)]
 pub async fn run_consensus<BlockT: ConsensusBlock>(
     context: Arc<dyn ConsensusContext<Block = BlockT>>,
@@ -50,12 +51,11 @@ where
 {
     let mut current_height = start_height;
     loop {
-        info!("Starting consensus for height {current_height}");
+        debug!("Starting consensus for height {current_height}");
         let mut shc =
             SingleHeightConsensus::new(current_height, context.clone(), validator_id).await;
 
         let block = if let Some(block) = shc.start().await? {
-            info!("Proposer flow height {current_height}");
             block
         } else {
             info!("Validator flow height {current_height}");
@@ -73,8 +73,8 @@ where
         };
 
         info!(
-            "Finished consensus for height: {start_height}. Agreed on block with id: {}",
-            block.id()
+            "Finished consensus for height: {current_height}. Agreed on block with id: {:x}",
+            block.id().0
         );
         current_height = current_height.unchecked_next();
     }
diff --git a/crates/sequencing/papyrus_consensus/src/single_height_consensus.rs b/crates/sequencing/papyrus_consensus/src/single_height_consensus.rs
index 51c09ce661..1b03ca4132 100644
--- a/crates/sequencing/papyrus_consensus/src/single_height_consensus.rs
+++ b/crates/sequencing/papyrus_consensus/src/single_height_consensus.rs
@@ -6,6 +6,7 @@ use std::sync::Arc;
 
 use futures::channel::{mpsc, oneshot};
 use starknet_api::block::{BlockHash, BlockNumber};
+use tracing::{debug, info, instrument};
 
 use crate::types::{ConsensusBlock, ConsensusContext, ConsensusError, ProposalInit, ValidatorId};
 
@@ -36,11 +37,15 @@ where
         Self { height, context, validators, id }
     }
 
+    #[instrument(skip(self), fields(height=self.height.0), level = "debug")]
     pub(crate) async fn start(&mut self) -> Result<Option<BlockT>, ConsensusError> {
+        info!("Starting consensus with validators {:?}", self.validators);
+
         let proposer_id = self.context.proposer(&self.validators, self.height);
         if proposer_id != self.id {
             return Ok(None);
         }
+        debug!("Proposer flow");
 
         let (content_receiver, block_receiver) = self.context.build_proposal(self.height).await;
         let (fin_sender, fin_receiver) = oneshot::channel();
@@ -62,12 +67,21 @@ where
 
     /// Receive a proposal from a peer node. Returns only once the proposal has been fully received
     /// and processed.
+    #[instrument(
+        skip(self, init, content_receiver, fin_receiver),
+        fields(height = %self.height),
+        level = "debug",
+    )]
     pub(crate) async fn handle_proposal(
         &mut self,
         init: ProposalInit,
         content_receiver: mpsc::Receiver<<BlockT as ConsensusBlock>::ProposalChunk>,
         fin_receiver: oneshot::Receiver<BlockHash>,
     ) -> Result<Option<BlockT>, ConsensusError> {
+        debug!(
+            "Received proposal: proposal_height={}, proposer={:?}",
+            init.height.0, init.proposer
+        );
         let proposer_id = self.context.proposer(&self.validators, self.height);
         if init.height != self.height {
             let msg = format!("invalid height: expected {:?}, got {:?}", self.height, init.height);