Skip to content

Commit

Permalink
Support attaching Uber Trace ID to every request (#937)
Browse files Browse the repository at this point in the history
  • Loading branch information
RobiNino committed Apr 30, 2024
1 parent 5bf715f commit 9d4dfaf
Showing 1 changed file with 22 additions and 0 deletions.
22 changes: 22 additions & 0 deletions http/httpclient/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 9d4dfaf

Please sign in to comment.