From b4d10d665c0c58b81d340ea659a769ecc27cf62a Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Tue, 31 Oct 2023 13:44:23 -0700
Subject: [PATCH] Add send and receive trace spans
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

```
2023-10-31T20:35:01.600364Z TRACE oprf_ipa_query{sz=10000}:send{i=914 total=1658 to=H2 gate="protocol/binary_validator/row1/attributed_breakdown_key/bit7"}: ipa::helpers::gateway::send: new
2023-10-31T20:35:01.600382Z TRACE oprf_ipa_query{sz=10000}:send{i=914 total=1658 to=H2 gate="protocol/binary_validator/row1/attributed_breakdown_key/bit7"}: ipa::helpers::gateway::send: close time.busy=2.42µs time.idle=15.7µs
```

it is a bit verbose, but logging enter/exit allows to see if send was completed (because it may be blocked if it is out of bounds) and the timing.
---
 src/helpers/gateway/receive.rs |  1 +
 src/helpers/gateway/send.rs    |  1 +
 src/helpers/mod.rs             | 12 +++++++++++-
 src/protocol/mod.rs            |  6 ++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/helpers/gateway/receive.rs b/src/helpers/gateway/receive.rs
index 282ff68e2..b039c8653 100644
--- a/src/helpers/gateway/receive.rs
+++ b/src/helpers/gateway/receive.rs
@@ -43,6 +43,7 @@ impl<T: Transport, M: Message> ReceivingEnd<T, M> {
     /// ## Panics
     /// This will panic if message size does not fit into 8 bytes and it somehow got serialized
     /// and sent to this helper.
+    #[tracing::instrument(level = "trace", "receive", skip_all, fields(i = %record_id, from = ?self.channel_id.role, gate = ?self.channel_id.gate.as_ref()))]
     pub async fn receive(&self, record_id: RecordId) -> Result<M, Error> {
         self.unordered_rx
             .recv::<M, _>(record_id)
diff --git a/src/helpers/gateway/send.rs b/src/helpers/gateway/send.rs
index 4eb876af0..0d32de973 100644
--- a/src/helpers/gateway/send.rs
+++ b/src/helpers/gateway/send.rs
@@ -98,6 +98,7 @@ impl<M: Message> SendingEnd<M> {
     /// call.
     ///
     /// [`set_total_records`]: crate::protocol::context::Context::set_total_records
+    #[tracing::instrument(level = "trace", "send", skip_all, fields(i = %record_id, total = %self.inner.total_records, to = ?self.channel_id.role, gate = ?self.channel_id.gate.as_ref()))]
     pub async fn send(&self, record_id: RecordId, msg: M) -> Result<(), Error> {
         let r = self.inner.send(record_id, msg).await;
         metrics::increment_counter!(RECORDS_SENT,
diff --git a/src/helpers/mod.rs b/src/helpers/mod.rs
index 373070736..359bf305f 100644
--- a/src/helpers/mod.rs
+++ b/src/helpers/mod.rs
@@ -1,5 +1,5 @@
 use std::{
-    fmt::{Debug, Formatter},
+    fmt::{Debug, Display, Formatter},
     num::NonZeroUsize,
 };
 
@@ -441,6 +441,16 @@ pub enum TotalRecords {
     Indeterminate,
 }
 
+impl Display for TotalRecords {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            TotalRecords::Unspecified => write!(f, "unspecified"),
+            TotalRecords::Specified(v) => write!(f, "{v}"),
+            TotalRecords::Indeterminate => write!(f, "∞"),
+        }
+    }
+}
+
 impl TotalRecords {
     #[must_use]
     pub fn is_specified(&self) -> bool {
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index cc0ae5ab8..8c5e72c7d 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -86,6 +86,12 @@ impl TryFrom<&str> for QueryId {
 #[cfg_attr(feature = "enable-serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct RecordId(u32);
 
+impl Display for RecordId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 impl From<u32> for RecordId {
     fn from(v: u32) -> Self {
         RecordId(v)