Add more tracing around shuffle

ipa-core/src/protocol/ipa_prf/shuffle/base.rs

@@ -38,6 +38,7 @@ where
     }
 }
 
+#[tracing::instrument(name = "h1_shuffle", skip_all, fields(%batch_size))]
 async fn run_h1<C, I, S, Zl, Zr>(
     ctx: &C,
     batch_size: usize,
@@ -79,6 +80,7 @@ where
     Ok(res)
 }
 
+#[tracing::instrument(name = "h2_shuffle", skip_all, fields(%batch_size))]
 async fn run_h2<C, I, S, Zl, Zr>(
     ctx: &C,
     batch_size: usize,
@@ -151,6 +153,7 @@ where
     Ok(res)
 }
 
+#[tracing::instrument(name = "h3_shuffle", skip_all, fields(%batch_size))]
 async fn run_h3<C, S, Zl, Zr>(
     ctx: &C,
     batch_size: usize,
@@ -289,6 +292,7 @@ where
 
 // ---------------------------- helper communication ------------------------------------ //
 
+#[tracing::instrument(name = "send_to", skip_all, fields(step, direction, sz = items.len()))]
 async fn send_to_peer<C, S>(
     items: &[S],
     ctx: &C,
@@ -311,6 +315,7 @@ where
     Ok(())
 }
 
+#[tracing::instrument(name = "receive_from", skip_all, fields(step, direction, sz = batch_size))]
 async fn receive_from_peer_into<C, S>(
     buf: &mut Vec<S>,
     batch_size: usize,