Add traceparent header to enable distributed tracing.

client/client_test.go

@@ -435,6 +435,8 @@ func testNonJSONResponseIncludedInError(t *testing.T, statusCode int, status, er
 		Token:      "token",
 		ConfigFile: "/dev/null",
 		HTTPTransport: hc(func(r *http.Request) (*http.Response, error) {
+			// Clear traceparent header which is nondeterministic.
+			r.Header.Del("traceparent")
 			return &http.Response{
 				Proto:   "HTTP/2.0",
 				Status:  status,

httpclient/api_client.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/common"
+	"github.com/databricks/databricks-sdk-go/httpclient/traceparent"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"github.com/databricks/databricks-sdk-go/logger/httplog"
 	"github.com/databricks/databricks-sdk-go/retries"
@@ -236,6 +237,9 @@ func (c *ApiClient) attempt(
 				return c.failRequest(ctx, "failed during request visitor", err)
 			}
 		}
+		// Set traceparent for distributed tracing.
+		traceparent.AddTraceparent(request)
+
 		// request.Context() holds context potentially enhanced by visitors
 		request.Header.Set("User-Agent", useragent.FromContext(request.Context()))
 		if request.Header.Get("Content-Type") == "" && requestBody.ContentType != "" {

httpclient/api_client_test.go

@@ -714,3 +714,24 @@ func TestDefaultAuthVisitor(t *testing.T) {
 	err := c.Do(context.Background(), "GET", "/a/b", WithRequestData(map[string]any{}), authOption)
 	require.NoError(t, err)
 }
+
+func TestTraceparentHeader(t *testing.T) {
+	seenTraceparents := []string{}
+	c := NewApiClient(ClientConfig{
+		Transport: hc(func(r *http.Request) (*http.Response, error) {
+			tp := r.Header.Get("Traceparent")
+			assert.NotEmpty(t, tp)
+			assert.NotContains(t, seenTraceparents, tp)
+			seenTraceparents = append(seenTraceparents, tp)
+			return &http.Response{
+				StatusCode: 200,
+				Request:    r,
+			}, nil
+		}),
+	})
+
+	for i := 0; i < 10; i++ {
+		err := c.Do(context.Background(), "GET", "/a/b")
+		assert.NoError(t, err)
+	}
+}

httpclient/fixtures/fixture.go

@@ -6,10 +6,13 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/textproto"
 	"net/url"
 	"os"
 	"reflect"
 	"strings"
+
+	"github.com/databricks/databricks-sdk-go/httpclient/traceparent"
 )
 
 // HTTPFixture defines request structure for test
@@ -38,12 +41,17 @@ func (f HTTPFixture) AssertHeaders(req *http.Request) error {
 		return nil
 	}
 	actualHeaders := map[string]string{}
+	// remove user agent & traceparent from comparison, as it'll make fixtures too difficult
+	// to maintain in the long term
+	toSkip := map[string]struct{}{
+		textproto.CanonicalMIMEHeaderKey("User-Agent"):                   {},
+		textproto.CanonicalMIMEHeaderKey(traceparent.TRACEPARENT_HEADER): {},
+	}
 	for k := range req.Header {
-		actualHeaders[k] = req.Header.Get(k)
+		if _, skip := toSkip[k]; !skip {
+			actualHeaders[k] = req.Header.Get(k)
+		}
 	}
-	// remove user agent from comparison, as it'll make fixtures too difficult
-	// to maintain in the long term
-	delete(actualHeaders, "User-Agent")
 	if !reflect.DeepEqual(f.ExpectedHeaders, actualHeaders) {
 		expectedJSON, _ := json.MarshalIndent(f.ExpectedHeaders, "", "  ")
 		actualJSON, _ := json.MarshalIndent(actualHeaders, "", "  ")

httpclient/traceparent/traceparent.go

@@ -0,0 +1,104 @@
+package traceparent
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"net/http"
+	"strings"
+)
+
+type traceparentFlag byte
+
+const (
+	traceparentFlagSampled traceparentFlag = 1 << iota
+)
+
+const TRACEPARENT_HEADER = "traceparent"
+
+type Traceparent struct {
+	version  byte
+	traceId  [16]byte
+	parentId [8]byte
+	flags    traceparentFlag
+}
+
+func (t *Traceparent) String() string {
+	b := strings.Builder{}
+	b.WriteString(hex.EncodeToString([]byte{t.version}))
+	b.WriteRune('-')
+	b.WriteString(hex.EncodeToString(t.traceId[:]))
+	b.WriteRune('-')
+	b.WriteString(hex.EncodeToString(t.parentId[:]))
+	b.WriteRune('-')
+	b.WriteString(hex.EncodeToString([]byte{byte(t.flags)}))
+	return b.String()
+}
+
+func (t *Traceparent) Equals(other *Traceparent) bool {
+	return t.version == other.version &&
+		t.flags == other.flags &&
+		string(t.traceId[:]) == string(other.traceId[:]) &&
+		string(t.parentId[:]) == string(other.parentId[:])
+}
+
+func NewTraceparent() *Traceparent {
+	traceId := [16]byte{}
+	parentId := [8]byte{}
+	rand.Read(traceId[:])
+	rand.Read(parentId[:])
+	return &Traceparent{
+		version:  0,
+		traceId:  traceId,
+		parentId: parentId,
+		flags:    traceparentFlagSampled,
+	}
+}
+
+func FromString(s string) (*Traceparent, error) {
+	parts := strings.Split(s, "-")
+	if len(parts) != 4 {
+		return nil, fmt.Errorf("invalid traceparent string: %s", s)
+	}
+	t := &Traceparent{
+		traceId:  [16]byte{},
+		parentId: [8]byte{},
+	}
+	version, err := hex.DecodeString(parts[0])
+	if err != nil {
+		return nil, err
+	}
+	if len(version) != 1 {
+		return nil, fmt.Errorf("invalid version: %s, expected 1 byte", parts[0])
+	}
+	t.version = version[0]
+	n, err := hex.Decode(t.traceId[:], []byte(parts[1]))
+	if err != nil {
+		return nil, err
+	}
+	if n != 16 {
+		return nil, fmt.Errorf("invalid traceId: %s, expected 16 bytes", parts[1])
+	}
+	n, err = hex.Decode(t.parentId[:], []byte(parts[2]))
+	if err != nil {
+		return nil, err
+	}
+	if n != 8 {
+		return nil, fmt.Errorf("invalid parentId: %s, expected 8 bytes", parts[2])
+	}
+	flags, err := hex.DecodeString(parts[3])
+	if err != nil {
+		return nil, err
+	}
+	if len(flags) != 1 {
+		return nil, fmt.Errorf("invalid flags: %s, expected 1 byte", parts[3])
+	}
+	t.flags = traceparentFlag(flags[0])
+	return t, nil
+}
+
+func AddTraceparent(r *http.Request) {
+	if r.Header.Get(TRACEPARENT_HEADER) == "" {
+		r.Header.Set(TRACEPARENT_HEADER, NewTraceparent().String())
+	}
+}

httpclient/traceparent/traceparent_test.go

@@ -0,0 +1,53 @@
+package traceparent
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNew(t *testing.T) {
+	tp := NewTraceparent()
+	assert.Equal(t, byte(0), tp.version)
+	assert.Equal(t, byte(1), byte(tp.flags))
+}
+
+func TestEqual(t *testing.T) {
+	tp1 := NewTraceparent()
+	tp2 := &Traceparent{
+		version:  tp1.version,
+		traceId:  tp1.traceId,
+		parentId: tp1.parentId,
+		flags:    tp1.flags,
+	}
+	assert.True(t, tp1.Equals(tp2))
+}
+
+func TestTwoNewTraceparentsAreNotEqual(t *testing.T) {
+	tp1 := NewTraceparent()
+	tp2 := NewTraceparent()
+	assert.False(t, tp1.Equals(tp2))
+}
+
+var testTraceId = [16]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+var testParentId = [8]byte{0, 1, 2, 3, 4, 5, 6, 7}
+
+func TestString(t *testing.T) {
+	tp := &Traceparent{
+		version:  0,
+		traceId:  testTraceId,
+		parentId: testParentId,
+		flags:    1,
+	}
+	res := tp.String()
+	assert.Equal(t, "00-000102030405060708090a0b0c0d0e0f-0001020304050607-01", res)
+}
+
+func TestFromString(t *testing.T) {
+	tp, err := FromString("00-000102030405060708090a0b0c0d0e0f-0001020304050607-01")
+	assert.NoError(t, err)
+	assert.Equal(t, byte(0), tp.version)
+	assert.Equal(t, testTraceId, tp.traceId)
+	assert.Equal(t, testParentId, tp.parentId)
+	assert.Equal(t, byte(1), byte(tp.flags))
+}