From cb454f7c4844b2d99b84e7c8aef555532fc85530 Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Thu, 21 Nov 2024 16:54:00 +0100
Subject: [PATCH] Make request IDs non-zero to make use of NPO and reduce
 struct size

---
 libs/telemetry/src/capturing/ng/collector.rs | 60 ++++++++------------
 libs/telemetry/src/capturing/ng/layer.rs     |  6 +-
 2 files changed, 28 insertions(+), 38 deletions(-)

diff --git a/libs/telemetry/src/capturing/ng/collector.rs b/libs/telemetry/src/capturing/ng/collector.rs
index b06ed80bcf7..1551a1ca757 100644
--- a/libs/telemetry/src/capturing/ng/collector.rs
+++ b/libs/telemetry/src/capturing/ng/collector.rs
@@ -20,6 +20,16 @@ use crate::models::{LogLevel, SpanKind, TraceSpan};
 #[repr(transparent)]
 struct SerializableNonZeroU64(NonZeroU64);
 
+impl SerializableNonZeroU64 {
+    pub fn into_u64(self) -> u64 {
+        self.0.get()
+    }
+
+    pub fn from_u64(value: u64) -> Option<Self> {
+        NonZeroU64::new(value).map(Self)
+    }
+}
+
 impl Serialize for SerializableNonZeroU64 {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -45,34 +55,6 @@ impl<'de> Deserialize<'de> for SerializableNonZeroU64 {
     }
 }
 
-#[derive(Debug, Display, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[display(fmt = "{}", _0)]
-#[repr(transparent)]
-struct SerializableU64(u64);
-
-impl Serialize for SerializableU64 {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: serde::Serializer,
-    {
-        // Serialize as string to preserve full u64 precision in JavaScript. Otherwise values
-        // larger than 2^53 - 1 will be parsed as floats on the client side, making it possible for
-        // IDs to collide.
-        self.to_string().serialize(serializer)
-    }
-}
-
-impl<'de> Deserialize<'de> for SerializableU64 {
-    fn deserialize<D>(deserializer: D) -> Result<SerializableU64, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        let value = String::deserialize(deserializer)?;
-        let value = value.parse().map_err(serde::de::Error::custom)?;
-        Ok(SerializableU64(value))
-    }
-}
-
 /// A unique identifier for a span. It maps directly to [`tracing::span::Id`] assigned by
 /// [`tracing_subscriber::registry::Registry`].
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
@@ -105,16 +87,26 @@ impl From<tracing::span::Id> for SpanId {
 /// still don't necessarily have to be unique for the whole lifetime of the process).
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
 #[repr(transparent)]
-pub struct RequestId(SerializableU64);
+pub struct RequestId(SerializableNonZeroU64);
 
 impl RequestId {
     pub fn next() -> Self {
         static NEXT_ID: AtomicU64 = AtomicU64::new(1);
-        Self(SerializableU64(NEXT_ID.fetch_add(1, Ordering::Relaxed)))
+
+        let mut id = 0;
+        while id == 0 {
+            id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
+        }
+
+        Self(SerializableNonZeroU64(NonZeroU64::new(id).unwrap()))
     }
 
     pub fn into_u64(self) -> u64 {
-        self.0 .0
+        self.0.into_u64()
+    }
+
+    pub fn from_u64(value: u64) -> Option<Self> {
+        SerializableNonZeroU64::from_u64(value).map(Self)
     }
 }
 
@@ -124,12 +116,6 @@ impl Default for RequestId {
     }
 }
 
-impl From<u64> for RequestId {
-    fn from(id: u64) -> Self {
-        Self(SerializableU64(id))
-    }
-}
-
 #[derive(Debug, Clone)]
 #[cfg_attr(test, derive(Serialize))]
 pub struct CollectedSpan {
diff --git a/libs/telemetry/src/capturing/ng/layer.rs b/libs/telemetry/src/capturing/ng/layer.rs
index 66ead397f0f..a3ae930f1a1 100644
--- a/libs/telemetry/src/capturing/ng/layer.rs
+++ b/libs/telemetry/src/capturing/ng/layer.rs
@@ -188,7 +188,11 @@ impl<'a> field::Visit for SpanAttributeVisitor<'a> {
 
     fn record_u64(&mut self, field: &field::Field, value: u64) {
         match field.name() {
-            REQUEST_ID_FIELD => self.span_builder.set_request_id(value.into()),
+            REQUEST_ID_FIELD => {
+                if let Some(request_id) = RequestId::from_u64(value) {
+                    self.span_builder.set_request_id(request_id);
+                }
+            }
             _ => self.span_builder.insert_attribute(field.name(), value.into()),
         }
     }