From 9d4dfaf050413deffe959c471921e00bf6f82641 Mon Sep 17 00:00:00 2001
From: Robi Nino <robi.nino@gmail.com>
Date: Tue, 30 Apr 2024 16:29:02 +0300
Subject: [PATCH] Support attaching Uber Trace ID to every request (#937)

---
 http/httpclient/client.go | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/http/httpclient/client.go b/http/httpclient/client.go
index 658bd7f6e..b101be52c 100644
--- a/http/httpclient/client.go
+++ b/http/httpclient/client.go
@@ -36,8 +36,14 @@ type HttpClient struct {
 const (
 	apiKeyPrefix        = "AKCp8"
 	apiKeyMinimalLength = 73
+	uberTraceIdHeader   = "uber-trace-id"
 )
 
+// If set, the Uber Trace ID header will be attached to every request.
+// This allows users to easily identify which logs on the server side are related to requests sent from this client.
+// Should be set using SetUberTraceIdToken.
+var uberTraceIdToken string
+
 func IsApiKey(key string) bool {
 	return strings.HasPrefix(key, apiKeyPrefix) && len(key) >= apiKeyMinimalLength
 }
@@ -171,6 +177,7 @@ func (jc *HttpClient) doRequest(req *http.Request, content []byte, followRedirec
 	setAuthentication(req, httpClientsDetails)
 	addUserAgentHeader(req)
 	copyHeaders(httpClientsDetails, req)
+	addUberTraceIdHeaderIfSet(req)
 
 	if !followRedirect || (followRedirect && req.Method == http.MethodPost) {
 		jc.client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
@@ -220,6 +227,21 @@ func copyHeaders(httpClientsDetails httputils.HttpClientDetails, req *http.Reque
 	}
 }
 
+// Generate an Uber Trace ID token that will be attached to every request.
+// Format of the header: {trace-id}:{span-id}:{parent-span-id}:{flags}
+// We set the trace-id and span-id to the same value, and the rest to 0.
+func SetUberTraceIdToken(traceIdToken string) {
+	uberTraceIdToken = fmt.Sprintf("%s:%s:0:0", traceIdToken, traceIdToken)
+}
+
+// If a trace ID is set, this function will attach the Uber Trace ID header to every request.
+func addUberTraceIdHeaderIfSet(req *http.Request) {
+	if uberTraceIdToken == "" {
+		return
+	}
+	req.Header.Set(uberTraceIdHeader, uberTraceIdToken)
+}
+
 func setRequestHeaders(httpClientsDetails httputils.HttpClientDetails, size int64, req *http.Request) {
 	copyHeaders(httpClientsDetails, req)
 	length := strconv.FormatInt(size, 10)