From a3057cf4b9dfd21385c91f06a493a850f4ca7028 Mon Sep 17 00:00:00 2001
From: Benjamin Ludewig <ludewig.b@gmail.com>
Date: Tue, 30 Mar 2021 02:29:32 +0200
Subject: [PATCH] add feature for serializing TraceId as byte array

---
 tarpc/Cargo.toml   |  1 +
 tarpc/src/trace.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/tarpc/Cargo.toml b/tarpc/Cargo.toml
index bd5d007d..dda50ea1 100644
--- a/tarpc/Cargo.toml
+++ b/tarpc/Cargo.toml
@@ -16,6 +16,7 @@ description = "An RPC framework for Rust with a focus on ease of use."
 default = []
 
 serde1 = ["tarpc-plugins/serde1", "serde", "serde/derive"]
+serde1-no-u128 = ["serde1"]
 tokio1 = ["tokio/rt-multi-thread"]
 serde-transport = ["serde1", "tokio1", "tokio-serde/json", "tokio-util/codec"]
 tcp = ["tokio/net"]
diff --git a/tarpc/src/trace.rs b/tarpc/src/trace.rs
index b3f03da4..eef4b820 100644
--- a/tarpc/src/trace.rs
+++ b/tarpc/src/trace.rs
@@ -47,7 +47,10 @@ pub struct Context {
 /// A 128-bit UUID identifying a trace. All spans caused by the same originating span share the
 /// same trace ID.
 #[derive(Debug, Default, PartialEq, Eq, Hash, Clone, Copy)]
-#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(
+    all(feature = "serde1", not(feature = "serde1-no-u128")),
+    derive(serde::Serialize, serde::Deserialize)
+)]
 pub struct TraceId(u128);
 
 /// A 64-bit identifier of a span within a trace. The identifier is unique within the span's trace.
@@ -95,3 +98,46 @@ impl fmt::Display for SpanId {
         Ok(())
     }
 }
+
+#[cfg(feature = "serde1-no-u128")]
+impl serde::Serialize for TraceId {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        use serde::ser::SerializeTupleStruct;
+        let mut tup = serializer.serialize_tuple_struct("TraceId", 1)?;
+        tup.serialize_field(&self.0.to_le_bytes())?;
+        tup.end()
+    }
+}
+
+#[cfg(feature = "serde1-no-u128")]
+impl<'de> serde::Deserialize<'de> for TraceId {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        struct Visitor;
+
+        impl<'de> serde::de::Visitor<'de> for Visitor {
+            type Value = TraceId;
+
+            fn expecting(&self, formatter: &mut Formatter) -> fmt::Result {
+                formatter.write_str("tuple struct TraceId")
+            }
+
+            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
+            where
+                A: serde::de::SeqAccess<'de>,
+            {
+                let bytes: [u8; 16] = seq.next_element()?.ok_or_else(|| {
+                    serde::de::Error::invalid_length(0, &"tuple struct TraceId with 1 element")
+                })?;
+                Ok(TraceId(u128::from_le_bytes(bytes)))
+            }
+        }
+
+        deserializer.deserialize_tuple_struct("TraceId", 1, Visitor)
+    }
+}