Auto set span headers for each request

Samze · Samze · commit b48daa3bc794 · 2025-01-21T13:58:11.000-07:00
diff --git a/api/cloudcontroller/wrapper/trace_request.go b/api/cloudcontroller/wrapper/trace_request.go
@@ -12,9 +12,9 @@ type CCTraceHeaderRequest struct {
 }
 
 // NewCCTraceHeaderRequest returns a pointer to a CCTraceHeaderRequest wrapper.
-func NewCCTraceHeaderRequest(trace, span string) *CCTraceHeaderRequest {
+func NewCCTraceHeaderRequest(trace string) *CCTraceHeaderRequest {
 	return &CCTraceHeaderRequest{
-		headers: shared.NewTraceHeaders(trace, span),
+		headers: shared.NewTraceHeaders(trace),
 	}
 }
 
diff --git a/api/cloudcontroller/wrapper/trace_request_test.go b/api/cloudcontroller/wrapper/trace_request_test.go
@@ -23,16 +23,14 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		makeErr  error
 
 		traceHeader string
-		spanHeader  string
 	)
 
 	BeforeEach(func() {
 		fakeConnection = new(cloudcontrollerfakes.FakeConnection)
 
 		traceHeader = "trace-id"
-		spanHeader = "span-id"
 
-		wrapper = NewCCTraceHeaderRequest(traceHeader, spanHeader).Wrap(fakeConnection)
+		wrapper = NewCCTraceHeaderRequest(traceHeader).Wrap(fakeConnection)
 
 		body := bytes.NewReader([]byte("foo"))
 
@@ -54,7 +52,7 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		It("Adds the request headers", func() {
 			Expect(makeErr).NotTo(HaveOccurred())
 			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
-			Expect(request.Header.Get("X-B3-SpanId")).To(Equal(spanHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
 		})
 
 		It("Calls the inner connection", func() {
diff --git a/api/router/wrapper/trace_request.go b/api/router/wrapper/trace_request.go
@@ -12,9 +12,9 @@ type RoutingTraceHeaderRequest struct {
 }
 
 // NewRoutingTraceHeaderRequest returns a pointer to a RoutingTraceHeaderRequest wrapper.
-func NewRoutingTraceHeaderRequest(trace, span string) *RoutingTraceHeaderRequest {
+func NewRoutingTraceHeaderRequest(trace string) *RoutingTraceHeaderRequest {
 	return &RoutingTraceHeaderRequest{
-		headers: shared.NewTraceHeaders(trace, span),
+		headers: shared.NewTraceHeaders(trace),
 	}
 }
 
diff --git a/api/router/wrapper/trace_request_test.go b/api/router/wrapper/trace_request_test.go
@@ -23,16 +23,13 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		makeErr  error
 
 		traceHeader string
-		spanHeader  string
 	)
 
 	BeforeEach(func() {
 		fakeConnection = new(routerfakes.FakeConnection)
 
 		traceHeader = "trace-id"
-		spanHeader = "span-id"
-
-		wrapper = NewRoutingTraceHeaderRequest(traceHeader, spanHeader).Wrap(fakeConnection)
+		wrapper = NewRoutingTraceHeaderRequest(traceHeader).Wrap(fakeConnection)
 
 		body := bytes.NewReader([]byte("foo"))
 
@@ -54,7 +51,7 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		It("Adds the request headers", func() {
 			Expect(makeErr).NotTo(HaveOccurred())
 			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
-			Expect(request.Header.Get("X-B3-SpanId")).To(Equal(spanHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
 		})
 
 		It("Calls the inner connection", func() {
diff --git a/api/shared/trace_headers.go b/api/shared/trace_headers.go
@@ -2,6 +2,8 @@ package shared
 
 import (
 	"net/http"
+
+	"code.cloudfoundry.org/cli/util/trace"
 )
 
 const (
@@ -12,14 +14,12 @@ const (
 // TraceHeaders sets b3 trace headers to requests.
 type TraceHeaders struct {
 	b3trace string
-	b3span  string
 }
 
 // NewTraceHeaders returns a pointer to a TraceHeaderRequest.
-func NewTraceHeaders(trace, span string) *TraceHeaders {
+func NewTraceHeaders(trace string) *TraceHeaders {
 	return &TraceHeaders{
 		b3trace: trace,
-		b3span:  span,
 	}
 }
 
@@ -30,6 +30,8 @@ func (t *TraceHeaders) SetHeaders(request *http.Request) {
 		request.Header.Add(B3TraceIDHeader, t.b3trace)
 	}
 	if request.Header.Get(B3SpanIDHeader) == "" {
-		request.Header.Add(B3SpanIDHeader, t.b3span)
+		request.Header.Add(B3SpanIDHeader, trace.GenerateRandomTraceID(16))
 	}
+
+	// request.Header.Add(("B3", request.Header.Get(B3TraceIDHeader)+request.Header.Get(B3SpanIDHeader)))
 }
diff --git a/api/shared/trace_headers_test.go b/api/shared/trace_headers_test.go
@@ -13,7 +13,7 @@ var _ = Describe("B3 Trace Headers", func() {
 	Describe("SetHeaders", func() {
 		Context("when there are already headers set", func() {
 			It("does not add the headers", func() {
-				traceHeaders := NewTraceHeaders("new_trace_id", "new_span_id")
+				traceHeaders := NewTraceHeaders("new_trace_id")
 				request := &http.Request{
 					Header: http.Header{},
 				}
@@ -28,14 +28,14 @@ var _ = Describe("B3 Trace Headers", func() {
 
 		Context("when there are no headers set", func() {
 			It("adds the headers", func() {
-				traceHeaders := NewTraceHeaders("new_trace_id", "new_span_id")
+				traceHeaders := NewTraceHeaders("new_trace_id")
 				request := &http.Request{
 					Header: http.Header{},
 				}
 				traceHeaders.SetHeaders(request)
 
 				Expect(request.Header.Get("X-B3-TraceId")).To(Equal("new_trace_id"))
-				Expect(request.Header.Get("X-B3-SpanId")).To(Equal("new_span_id"))
+				Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
 			})
 		})
 	})
diff --git a/api/uaa/wrapper/trace_request.go b/api/uaa/wrapper/trace_request.go
@@ -14,9 +14,9 @@ type UAATraceHeaderRequest struct {
 }
 
 // NewUAATraceHeaderRequest returns a pointer to a UAATraceHeaderRequest wrapper.
-func NewUAATraceHeaderRequest(trace, span string) *UAATraceHeaderRequest {
+func NewUAATraceHeaderRequest(trace string) *UAATraceHeaderRequest {
 	return &UAATraceHeaderRequest{
-		headers: shared.NewTraceHeaders(trace, span),
+		headers: shared.NewTraceHeaders(trace),
 	}
 }
 
diff --git a/api/uaa/wrapper/trace_request_test.go b/api/uaa/wrapper/trace_request_test.go
@@ -23,16 +23,14 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		makeErr  error
 
 		traceHeader string
-		spanHeader  string
 	)
 
 	BeforeEach(func() {
 		fakeConnection = new(uaafakes.FakeConnection)
 
 		traceHeader = "trace-id"
-		spanHeader = "span-id"
 
-		wrapper = NewUAATraceHeaderRequest(traceHeader, spanHeader).Wrap(fakeConnection)
+		wrapper = NewUAATraceHeaderRequest(traceHeader).Wrap(fakeConnection)
 
 		body := bytes.NewReader([]byte("foo"))
 
@@ -54,7 +52,7 @@ var _ = Describe("CCTraceHeaderRequest", func() {
 		It("Adds the request headers", func() {
 			Expect(makeErr).NotTo(HaveOccurred())
 			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
-			Expect(request.Header.Get("X-B3-SpanId")).To(Equal(spanHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
 		})
 
 		It("Calls the inner connection", func() {
diff --git a/command/config.go b/command/config.go
@@ -16,7 +16,6 @@ type Config interface {
 	AuthorizationEndpoint() string
 	APIVersion() string
 	B3TraceID() string
-	B3SpanID() string
 	BinaryName() string
 	BinaryVersion() string
 	CFPassword() string
diff --git a/command/v7/shared/new_clients.go b/command/v7/shared/new_clients.go
@@ -47,7 +47,7 @@ func NewWrappedCloudControllerClient(config command.Config, ui command.UI, extra
 	}
 
 	ccWrappers = append(ccWrappers, extraWrappers...)
-	ccWrappers = append(ccWrappers, ccWrapper.NewCCTraceHeaderRequest(config.B3TraceID(), config.B3SpanID()))
+	ccWrappers = append(ccWrappers, ccWrapper.NewCCTraceHeaderRequest(config.B3TraceID()))
 	ccWrappers = append(ccWrappers, ccWrapper.NewRetryRequest(config.RequestRetryCount()))
 
 	return ccv3.NewClient(ccv3.Config{
@@ -86,7 +86,7 @@ func newWrappedUAAClient(config command.Config, ui command.UI) (*uaa.Client, err
 
 	uaaAuthWrapper := uaaWrapper.NewUAAAuthentication(uaaClient, config)
 	uaaClient.WrapConnection(uaaAuthWrapper)
-	uaaClient.WrapConnection(uaaWrapper.NewUAATraceHeaderRequest(config.B3TraceID(), config.B3SpanID()))
+	uaaClient.WrapConnection(uaaWrapper.NewUAATraceHeaderRequest(config.B3TraceID()))
 	uaaClient.WrapConnection(uaaWrapper.NewRetryRequest(config.RequestRetryCount()))
 
 	err = uaaClient.SetupResources(config.UAAEndpoint(), config.AuthorizationEndpoint())
@@ -123,6 +123,8 @@ func newWrappedRoutingClient(config command.Config, ui command.UI, uaaClient *ua
 	authWrapper := routingWrapper.NewUAAAuthentication(uaaClient, config)
 
 	routingWrappers = append(routingWrappers, authWrapper)
+	routingWrappers = append(routingWrappers, routingWrapper.NewRoutingTraceHeaderRequest(config.B3TraceID()))
+
 	routingConfig.Wrappers = routingWrappers
 
 	routingClient := router.NewClient(routingConfig)
diff --git a/go.mod b/go.mod
@@ -24,6 +24,7 @@ require (
 	github.com/distribution/reference v0.6.0
 	github.com/fatih/color v1.18.0
 	github.com/google/go-querystring v1.1.0
+	github.com/google/uuid v1.6.0
 	github.com/jessevdk/go-flags v1.6.1
 	github.com/lunixbochs/vtclean v1.0.0
 	github.com/mattn/go-colorable v0.1.13
diff --git a/util/configv3/env.go b/util/configv3/env.go
@@ -6,7 +6,7 @@ import (
 	"strings"
 	"time"
 
-	"code.cloudfoundry.org/cli/util/random"
+	"code.cloudfoundry.org/cli/util/trace"
 )
 
 // EnvOverride represents all the environment variables read by the CF CLI
@@ -23,7 +23,6 @@ type EnvOverride struct {
 	CFTrace          string
 	CFUsername       string
 	CFB3TraceID      string
-	CFB3SpanID       string
 	DockerPassword   string
 	CNBCredentials   string
 	Experimental     string
@@ -167,14 +166,7 @@ func (config *Config) StartupTimeout() time.Duration {
 
 func (config *Config) B3TraceID() string {
 	if config.ENV.CFB3TraceID == "" {
-		config.ENV.CFB3TraceID = random.GenerateHex(32)
+		config.ENV.CFB3TraceID = trace.GenerateUUIDTraceID()
 	}
 	return config.ENV.CFB3TraceID
 }
-
-func (config *Config) B3SpanID() string {
-	if config.ENV.CFB3SpanID == "" {
-		config.ENV.CFB3SpanID = random.GenerateHex(16)
-	}
-	return config.ENV.CFB3SpanID
-}
diff --git a/util/configv3/load_config.go b/util/configv3/load_config.go
@@ -128,7 +128,6 @@ func LoadConfig(flags ...FlagOverride) (*Config, error) {
 		CFTrace:          os.Getenv("CF_TRACE"),
 		CFUsername:       os.Getenv("CF_USERNAME"),
 		CFB3TraceID:      os.Getenv("CF_B3_TRACE_ID"),
-		CFB3SpanID:       os.Getenv("CF_B3_SPAN_ID"),
 		DockerPassword:   os.Getenv("CF_DOCKER_PASSWORD"),
 		CNBCredentials:   os.Getenv("CNB_REGISTRY_CREDS"),
 		Experimental:     os.Getenv("CF_CLI_EXPERIMENTAL"),
diff --git a/util/random/hex.go b/util/random/hex.go
diff --git a/util/trace/trace.go b/util/trace/trace.go
@@ -0,0 +1,25 @@
+package trace
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"strings"
+
+	"github.com/google/uuid"
+)
+
+// GenerateUUIDTraceID returns a UUID v4 string with the dashes removed as a 32 lower-hex encoded string.
+func GenerateUUIDTraceID() string {
+	uuidV4 := uuid.New()
+	return strings.ReplaceAll(uuidV4.String(), "-", "")
+}
+
+// GenerateRandomTraceID returns a random hex string of the given length.
+func GenerateRandomTraceID(length int) string {
+	b := make([]byte, length/2)
+	if _, err := rand.Read(b); err != nil {
+		panic(err)
+	}
+
+	return hex.EncodeToString(b)
+}
diff --git a/util/trace/trace_suite_test.go b/util/trace/trace_suite_test.go
@@ -0,0 +1,13 @@
+package trace_test
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"testing"
+)
+
+func TestTrace(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Trace Suite")
+}
diff --git a/util/trace/trace_test.go b/util/trace/trace_test.go
@@ -0,0 +1,19 @@
+package trace_test
+
+import (
+	"code.cloudfoundry.org/cli/util/trace"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("GenerateHex", func() {
+	It("returns random trace id", func() {
+		Expect(trace.GenerateUUIDTraceID()).To(HaveLen(32))
+	})
+})
+
+var _ = Describe("GenerateRandomTraceID", func() {
+	It("returns random trace id", func() {
+		Expect(trace.GenerateRandomTraceID(32)).To(HaveLen(16))
+	})
+})

Original file line number	Diff line number	Diff line change
`@@ -12,9 +12,9 @@ type CCTraceHeaderRequest struct {`
`12`	`12`	`}`
`13`	`13`
`14`	`14`	`// NewCCTraceHeaderRequest returns a pointer to a CCTraceHeaderRequest wrapper.`
`15`		`-func NewCCTraceHeaderRequest(trace, span string) *CCTraceHeaderRequest {`
	`15`	`+func NewCCTraceHeaderRequest(trace string) *CCTraceHeaderRequest {`
`16`	`16`	`return &CCTraceHeaderRequest{`
`17`		`- headers: shared.NewTraceHeaders(trace, span),`
	`17`	`+ headers: shared.NewTraceHeaders(trace),`
`18`	`18`	`}`
`19`	`19`	`}`
`20`	`20`
Original file line number	Diff line number	Diff line change
`@@ -12,9 +12,9 @@ type RoutingTraceHeaderRequest struct {`
`12`	`12`	`}`
`13`	`13`
`14`	`14`	`// NewRoutingTraceHeaderRequest returns a pointer to a RoutingTraceHeaderRequest wrapper.`
`15`		`-func NewRoutingTraceHeaderRequest(trace, span string) *RoutingTraceHeaderRequest {`
	`15`	`+func NewRoutingTraceHeaderRequest(trace string) *RoutingTraceHeaderRequest {`
`16`	`16`	`return &RoutingTraceHeaderRequest{`
`17`		`- headers: shared.NewTraceHeaders(trace, span),`
	`17`	`+ headers: shared.NewTraceHeaders(trace),`
`18`	`18`	`}`
`19`	`19`	`}`
`20`	`20`
Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,8 @@ package shared`
`2`	`2`
`3`	`3`	`import (`
`4`	`4`	`"net/http"`
	`5`	`+`
	`6`	`+ "code.cloudfoundry.org/cli/util/trace"`
`5`	`7`	`)`
`6`	`8`
`7`	`9`	`const (`
`@@ -12,14 +14,12 @@ const (`
`12`	`14`	`// TraceHeaders sets b3 trace headers to requests.`
`13`	`15`	`type TraceHeaders struct {`
`14`	`16`	`b3trace string`
`15`		`- b3span string`
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`// NewTraceHeaders returns a pointer to a TraceHeaderRequest.`
`19`		`-func NewTraceHeaders(trace, span string) *TraceHeaders {`
	`20`	`+func NewTraceHeaders(trace string) *TraceHeaders {`
`20`	`21`	`return &TraceHeaders{`
`21`	`22`	`b3trace: trace,`
`22`		`- b3span: span,`
`23`	`23`	`}`
`24`	`24`	`}`
`25`	`25`
`@@ -30,6 +30,8 @@ func (t TraceHeaders) SetHeaders(request http.Request) {`
`30`	`30`	`request.Header.Add(B3TraceIDHeader, t.b3trace)`
`31`	`31`	`}`
`32`	`32`	`if request.Header.Get(B3SpanIDHeader) == "" {`
`33`		`- request.Header.Add(B3SpanIDHeader, t.b3span)`
	`33`	`+ request.Header.Add(B3SpanIDHeader, trace.GenerateRandomTraceID(16))`
`34`	`34`	`}`
	`35`	`+`
	`36`	`+ // request.Header.Add(("B3", request.Header.Get(B3TraceIDHeader)+request.Header.Get(B3SpanIDHeader)))`
`35`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,9 +14,9 @@ type UAATraceHeaderRequest struct {`
`14`	`14`	`}`
`15`	`15`
`16`	`16`	`// NewUAATraceHeaderRequest returns a pointer to a UAATraceHeaderRequest wrapper.`
`17`		`-func NewUAATraceHeaderRequest(trace, span string) *UAATraceHeaderRequest {`
	`17`	`+func NewUAATraceHeaderRequest(trace string) *UAATraceHeaderRequest {`
`18`	`18`	`return &UAATraceHeaderRequest{`
`19`		`- headers: shared.NewTraceHeaders(trace, span),`
	`19`	`+ headers: shared.NewTraceHeaders(trace),`
`20`	`20`	`}`
`21`	`21`	`}`
`22`	`22`