From 7a626d2e848f739e4f6f94b275f1f1b2eeb56300 Mon Sep 17 00:00:00 2001
From: Kevin Glowacz <kjg@github.com>
Date: Wed, 16 Oct 2024 08:53:41 -0500
Subject: [PATCH] Allow the span kind to be set via StartSpanOptions (#23590)

---
 sdk/azcore/CHANGELOG.md                      |  2 ++
 sdk/azcore/runtime/policy_http_trace.go      | 10 +++++---
 sdk/azcore/runtime/policy_http_trace_test.go | 27 ++++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md
index f53e94895b..607d992fd8 100644
--- a/sdk/azcore/CHANGELOG.md
+++ b/sdk/azcore/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+* Added field `Kind` to `runtime.StartSpanOptions` to allow a kind to be set when starting a span.
+
 ### Breaking Changes
 
 ### Bugs Fixed
diff --git a/sdk/azcore/runtime/policy_http_trace.go b/sdk/azcore/runtime/policy_http_trace.go
index bc6989310b..f375195c4b 100644
--- a/sdk/azcore/runtime/policy_http_trace.go
+++ b/sdk/azcore/runtime/policy_http_trace.go
@@ -96,6 +96,8 @@ func (h *httpTracePolicy) Do(req *policy.Request) (resp *http.Response, err erro
 
 // StartSpanOptions contains the optional values for StartSpan.
 type StartSpanOptions struct {
+	// Kind indicates the kind of Span.
+	Kind tracing.SpanKind
 	// Attributes contains key-value pairs of attributes for the span.
 	Attributes []tracing.Attribute
 }
@@ -115,7 +117,6 @@ func StartSpan(ctx context.Context, name string, tracer tracing.Tracer, options
 	// we MUST propagate the active tracer before returning so that the trace policy can access it
 	ctx = context.WithValue(ctx, shared.CtxWithTracingTracer{}, tracer)
 
-	const newSpanKind = tracing.SpanKindInternal
 	if activeSpan := ctx.Value(ctxActiveSpan{}); activeSpan != nil {
 		// per the design guidelines, if a SDK method Foo() calls SDK method Bar(),
 		// then the span for Bar() must be suppressed. however, if Bar() makes a REST
@@ -131,12 +132,15 @@ func StartSpan(ctx context.Context, name string, tracer tracing.Tracer, options
 	if options == nil {
 		options = &StartSpanOptions{}
 	}
+	if options.Kind == 0 {
+		options.Kind = tracing.SpanKindInternal
+	}
 
 	ctx, span := tracer.Start(ctx, name, &tracing.SpanOptions{
-		Kind:       newSpanKind,
+		Kind:       options.Kind,
 		Attributes: options.Attributes,
 	})
-	ctx = context.WithValue(ctx, ctxActiveSpan{}, newSpanKind)
+	ctx = context.WithValue(ctx, ctxActiveSpan{}, options.Kind)
 	return ctx, func(err error) {
 		if err != nil {
 			errType := strings.Replace(fmt.Sprintf("%T", err), "*exported.", "*azcore.", 1)
diff --git a/sdk/azcore/runtime/policy_http_trace_test.go b/sdk/azcore/runtime/policy_http_trace_test.go
index 1e83508b7b..43581ec54d 100644
--- a/sdk/azcore/runtime/policy_http_trace_test.go
+++ b/sdk/azcore/runtime/policy_http_trace_test.go
@@ -267,3 +267,30 @@ func TestStartSpanWithAttributes(t *testing.T) {
 	require.True(t, startCalled)
 	require.True(t, endCalled)
 }
+
+func TestStartSpanWithKind(t *testing.T) {
+	// span no error
+	var startCalled bool
+	var endCalled bool
+	tr := tracing.NewTracer(func(ctx context.Context, spanName string, options *tracing.SpanOptions) (context.Context, tracing.Span) {
+		startCalled = true
+		require.EqualValues(t, "TestStartSpan", spanName)
+		require.NotNil(t, options)
+		// The span kind should be passed through
+		require.EqualValues(t, tracing.SpanKindClient, options.Kind)
+		spanImpl := tracing.SpanImpl{
+			End: func() { endCalled = true },
+		}
+		return ctx, tracing.NewSpan(spanImpl)
+	}, nil)
+	ctx, end := StartSpan(context.Background(), "TestStartSpan", tr, &StartSpanOptions{
+		Kind: tracing.SpanKindClient,
+	})
+	end(nil)
+	ctxTr := ctx.Value(shared.CtxWithTracingTracer{})
+	require.NotNil(t, ctxTr)
+	_, ok := ctxTr.(tracing.Tracer)
+	require.True(t, ok)
+	require.True(t, startCalled)
+	require.True(t, endCalled)
+}