From 1d8ba7bcb350a85f1f86b388b1847a1fadb9d9e3 Mon Sep 17 00:00:00 2001
From: Joel Hendrix <jhendrix@microsoft.com>
Date: Mon, 26 Jun 2023 15:42:15 -0700
Subject: [PATCH] Add support for shallow cloning azcore.Client instances
 (#21065)

* Add support for shallow cloning azcore.Client instances

This allows for multiple clients to share the same underlying pipeline
while having the correct client name string in traces.

* improved field names
---
 sdk/azcore/CHANGELOG.md |  1 +
 sdk/azcore/core.go      | 25 ++++++++++++++++++++-
 sdk/azcore/core_test.go | 48 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md
index c6c833fe5b29..0a2c2982a7bb 100644
--- a/sdk/azcore/CHANGELOG.md
+++ b/sdk/azcore/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Features Added
 
 - `messaging/CloudEvent` allows you to serialize/deserialize CloudEvents, as described in the CloudEvents 1.0 specification: [link](https://github.com/cloudevents/spec)
+* Added method `WithClientName()` to type `azcore.Client` to support shallow cloning of a client with a new name used for tracing.
 
 ### Breaking Changes
 
diff --git a/sdk/azcore/core.go b/sdk/azcore/core.go
index a8f8a34752f8..e10b1ea093b3 100644
--- a/sdk/azcore/core.go
+++ b/sdk/azcore/core.go
@@ -73,6 +73,11 @@ type ClientOptions = policy.ClientOptions
 type Client struct {
 	pl runtime.Pipeline
 	tr tracing.Tracer
+
+	// cached on the client to support shallow copying with new values
+	tp        tracing.Provider
+	modVer    string
+	namespace string
 }
 
 // NewClient creates a new Client instance with the provided values.
@@ -103,7 +108,14 @@ func NewClient(clientName, moduleVersion string, plOpts runtime.PipelineOptions,
 	if tr.Enabled() && plOpts.TracingNamespace != "" {
 		tr.SetAttributes(tracing.Attribute{Key: "az.namespace", Value: plOpts.TracingNamespace})
 	}
-	return &Client{pl: pl, tr: tr}, nil
+
+	return &Client{
+		pl:        pl,
+		tr:        tr,
+		tp:        options.TracingProvider,
+		modVer:    moduleVersion,
+		namespace: plOpts.TracingNamespace,
+	}, nil
 }
 
 // Pipeline returns the pipeline for this client.
@@ -115,3 +127,14 @@ func (c *Client) Pipeline() runtime.Pipeline {
 func (c *Client) Tracer() tracing.Tracer {
 	return c.tr
 }
+
+// WithClientName returns a shallow copy of the Client with its tracing client name changed to clientName.
+// Note that the values for module name and version will be preserved from the source Client.
+//   - clientName - the fully qualified name of the client ("package.Client"); this is used by the tracing provider when creating spans
+func (c *Client) WithClientName(clientName string) *Client {
+	tr := c.tp.NewTracer(clientName, c.modVer)
+	if tr.Enabled() && c.namespace != "" {
+		tr.SetAttributes(tracing.Attribute{Key: "az.namespace", Value: c.namespace})
+	}
+	return &Client{pl: c.pl, tr: tr, tp: c.tp, modVer: c.modVer, namespace: c.namespace}
+}
diff --git a/sdk/azcore/core_test.go b/sdk/azcore/core_test.go
index 78c4af58bc55..bd00c7c0bd50 100644
--- a/sdk/azcore/core_test.go
+++ b/sdk/azcore/core_test.go
@@ -172,3 +172,51 @@ func TestNewClientTracingEnabled(t *testing.T) {
 	require.NoError(t, err)
 	require.EqualValues(t, "az.namespace:Widget.Factory", attrString)
 }
+
+func TestClientWithClientName(t *testing.T) {
+	srv, close := mock.NewServer()
+	defer close()
+
+	var clientName string
+	var modVersion string
+	var attrString string
+	client, err := NewClient("module/package.Client", "v1.0.0", runtime.PipelineOptions{TracingNamespace: "Widget.Factory"}, &policy.ClientOptions{
+		TracingProvider: tracing.NewProvider(func(name, version string) tracing.Tracer {
+			clientName = name
+			modVersion = version
+			return tracing.NewTracer(func(ctx context.Context, spanName string, options *tracing.SpanOptions) (context.Context, tracing.Span) {
+				require.NotNil(t, options)
+				for _, attr := range options.Attributes {
+					if attr.Key == "az.namespace" {
+						v, ok := attr.Value.(string)
+						require.True(t, ok)
+						attrString = attr.Key + ":" + v
+					}
+				}
+				return ctx, tracing.Span{}
+			}, nil)
+		}, nil),
+		Transport: srv,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, client)
+	require.NotZero(t, client.Pipeline())
+	require.NotZero(t, client.Tracer())
+	require.EqualValues(t, "package.Client", clientName)
+	require.EqualValues(t, "v1.0.0", modVersion)
+
+	const requestEndpoint = "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/fakeResourceGroupo/providers/Microsoft.Storage/storageAccounts/fakeAccountName"
+	req, err := exported.NewRequest(context.WithValue(context.Background(), shared.CtxWithTracingTracer{}, client.Tracer()), http.MethodGet, srv.URL()+requestEndpoint)
+	require.NoError(t, err)
+	srv.SetResponse()
+	_, err = client.Pipeline().Do(req)
+	require.NoError(t, err)
+	require.EqualValues(t, "az.namespace:Widget.Factory", attrString)
+
+	newClient := client.WithClientName("other.Client")
+	require.EqualValues(t, "other.Client", clientName)
+	require.EqualValues(t, "v1.0.0", modVersion)
+	_, err = newClient.Pipeline().Do(req)
+	require.NoError(t, err)
+	require.EqualValues(t, "az.namespace:Widget.Factory", attrString)
+}