feat: pass tracing context in proto

GreptimeTeam · Nov 14, 2023 · a30131d · a30131d
1 parent 7eb2e78
commit a30131d
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 14 deletions.
diff --git a/proto/greptime/v1/common.proto b/proto/greptime/v1/common.proto
@@ -29,10 +29,9 @@ message RequestHeader {
   AuthHeader authorization = 3;
   // The `dbname` for the request
   string dbname = 4;
-  // TraceID of request
-  uint64 trace_id = 5;
-  // SpanID of request
-  uint64 span_id = 6;
+  // Encoded trace_id & span_id, follow the w3c Trace Context
+  // https://www.w3.org/TR/trace-context/#header-name
+  map<string, string> tracing_context = 5;
 }
 
 message ResponseHeader { Status status = 1; }

diff --git a/proto/greptime/v1/meta/common.proto b/proto/greptime/v1/meta/common.proto
@@ -25,6 +25,9 @@ message RequestHeader {
   // member_id is the ID of the sender server.
   uint64 member_id = 3;
   Role role = 4;
+  // Encoded trace_id & span_id, follow the w3c Trace Context
+  // https://www.w3.org/TR/trace-context/#header-name
+  map<string, string> tracing_context = 5;
 }
 
 enum Role {
@@ -49,9 +52,7 @@ message Peer {
   string addr = 2;
 }
 
-message TableId {
-  uint32 id = 1;
-}
+message TableId { uint32 id = 1; }
 
 message TableName {
   string catalog_name = 1;

diff --git a/proto/greptime/v1/region/server.proto b/proto/greptime/v1/region/server.proto
@@ -27,10 +27,8 @@ import "greptime/v1/ddl.proto";
 service Region { rpc Handle(RegionRequest) returns (RegionResponse); }
 
 message RegionRequestHeader {
-  // TraceID of request
-  uint64 trace_id = 1;
-  // SpanID of request
-  uint64 span_id = 2;
+  // Encoded trace_id & span_id, follow the w3c Trace Context https://www.w3.org/TR/trace-context/#header-name
+  map<string, string> tracing_context = 5;
   // DB Name of request, tracking only
   string dbname = 3;
 }

diff --git a/src/v1/meta.rs b/src/v1/meta.rs
@@ -59,12 +59,17 @@ impl Eq for Peer {}
 
 impl RequestHeader {
     #[inline]
-    pub fn new((cluster_id, member_id): (u64, u64), role: Role) -> Self {
+    pub fn new(
+        (cluster_id, member_id): (u64, u64),
+        role: Role,
+        tracing_context: HashMap<String, String>,
+    ) -> Self {
         Self {
             protocol_version: PROTOCOL_VERSION,
             cluster_id,
             member_id,
             role: role.into(),
+            tracing_context,
         }
     }
 }
@@ -138,15 +143,25 @@ macro_rules! gen_set_header {
     ($req: ty) => {
         impl $req {
             #[inline]
-            pub fn set_header(&mut self, (cluster_id, member_id): (u64, u64), role: Role) {
+            pub fn set_header(
+                &mut self,
+                (cluster_id, member_id): (u64, u64),
+                role: Role,
+                tracing_context: HashMap<String, String>,
+            ) {
                 match self.header.as_mut() {
                     Some(header) => {
                         header.cluster_id = cluster_id;
                         header.member_id = member_id;
                         header.role = role.into();
+                        header.tracing_context = tracing_context;
                     }
                     None => {
-                        self.header = Some(RequestHeader::new((cluster_id, member_id), role));
+                        self.header = Some(RequestHeader::new(
+                            (cluster_id, member_id),
+                            role,
+                            tracing_context,
+                        ));
                     }
                 }
             }