[azservicebus] Enable distributed tracing

sdk/messaging/azservicebus/client.go

@@ -16,7 +16,9 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/log"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/amqpwrap"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/conn"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/exported"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 )
 
 // Client provides methods to create Sender and Receiver
@@ -31,6 +33,7 @@ type Client struct {
 
 	linksMu      *sync.Mutex
 	links        map[uint64]amqpwrap.Closeable
+	tracer       tracing.Tracer
 	creds        clientCreds
 	namespace    internal.NamespaceForAMQPLinks
 	retryOptions RetryOptions
@@ -53,6 +56,10 @@ type ClientOptions struct {
 	// For an example, see ExampleNewClient_usingWebsockets() function in example_client_test.go.
 	NewWebSocketConn func(ctx context.Context, args NewWebSocketConnArgs) (net.Conn, error)
 
+	// TracingProvider configures the tracing provider.
+	// It defaults to a no-op tracer.
+	TracingProvider tracing.Provider
+
 	// RetryOptions controls how often operations are retried from this client and any
 	// Receivers and Senders created from this client.
 	RetryOptions RetryOptions
@@ -133,6 +140,7 @@ func newClientImpl(creds clientCreds, args clientImplArgs) (*Client, error) {
 	}
 
 	var err error
+	var tracingProvider = tracing.Provider{}
 	var nsOptions []internal.NamespaceOption
 
 	if client.creds.connectionString != "" {
@@ -146,6 +154,8 @@ func newClientImpl(creds clientCreds, args clientImplArgs) (*Client, error) {
 	}
 
 	if args.ClientOptions != nil {
+		tracingProvider = args.ClientOptions.TracingProvider
+
 		client.retryOptions = args.ClientOptions.RetryOptions
 
 		if args.ClientOptions.TLSConfig != nil {
@@ -167,16 +177,21 @@ func newClientImpl(creds clientCreds, args clientImplArgs) (*Client, error) {
 		nsOptions = append(nsOptions, internal.NamespaceWithRetryOptions(args.ClientOptions.RetryOptions))
 	}
 
+	client.tracer = newTracer(tracingProvider, getFullyQualifiedNamespace(creds))
+	nsOptions = append(nsOptions, internal.NamespaceWithTracer(client.tracer))
+
 	nsOptions = append(nsOptions, args.NSOptions...)
 
 	client.namespace, err = internal.NewNamespace(nsOptions...)
+
 	return client, err
 }
 
 // NewReceiverForQueue creates a Receiver for a queue. A receiver allows you to receive messages.
 func (client *Client) NewReceiverForQueue(queueName string, options *ReceiverOptions) (*Receiver, error) {
 	id, cleanupOnClose := client.getCleanupForCloseable()
 	receiver, err := newReceiver(newReceiverArgs{
+		tracer:              client.tracer,
 		cleanupOnClose:      cleanupOnClose,
 		ns:                  client.namespace,
 		entity:              entity{Queue: queueName},
@@ -196,6 +211,7 @@ func (client *Client) NewReceiverForQueue(queueName string, options *ReceiverOpt
 func (client *Client) NewReceiverForSubscription(topicName string, subscriptionName string, options *ReceiverOptions) (*Receiver, error) {
 	id, cleanupOnClose := client.getCleanupForCloseable()
 	receiver, err := newReceiver(newReceiverArgs{
+		tracer:              client.tracer,
 		cleanupOnClose:      cleanupOnClose,
 		ns:                  client.namespace,
 		entity:              entity{Topic: topicName, Subscription: subscriptionName},
@@ -220,6 +236,7 @@ type NewSenderOptions struct {
 func (client *Client) NewSender(queueOrTopic string, options *NewSenderOptions) (*Sender, error) {
 	id, cleanupOnClose := client.getCleanupForCloseable()
 	sender, err := newSender(newSenderArgs{
+		tracer:         client.tracer,
 		ns:             client.namespace,
 		queueOrTopic:   queueOrTopic,
 		cleanupOnClose: cleanupOnClose,
@@ -242,6 +259,7 @@ func (client *Client) AcceptSessionForQueue(ctx context.Context, queueName strin
 	sessionReceiver, err := newSessionReceiver(
 		ctx,
 		newSessionReceiverArgs{
+			tracer:         client.tracer,
 			sessionID:      &sessionID,
 			ns:             client.namespace,
 			entity:         entity{Queue: queueName},
@@ -269,6 +287,7 @@ func (client *Client) AcceptSessionForSubscription(ctx context.Context, topicNam
 	sessionReceiver, err := newSessionReceiver(
 		ctx,
 		newSessionReceiverArgs{
+			tracer:         client.tracer,
 			sessionID:      &sessionID,
 			ns:             client.namespace,
 			entity:         entity{Topic: topicName, Subscription: subscriptionName},
@@ -338,6 +357,7 @@ func (client *Client) acceptNextSessionForEntity(ctx context.Context, entity ent
 	sessionReceiver, err := newSessionReceiver(
 		ctx,
 		newSessionReceiverArgs{
+			tracer:            client.tracer,
 			sessionID:         nil,
 			ns:                client.namespace,
 			entity:            entity,
@@ -373,3 +393,17 @@ func (client *Client) getCleanupForCloseable() (uint64, func()) {
 		client.linksMu.Unlock()
 	}
 }
+
+// getFullyQualifiedNamespace returns fullyQualifiedNamespace from clientCreds if it is set.
+// Otherwise, it parses the connection string and returns the FullyQualifiedNamespace from it.
+// If both are empty, it returns an empty string.
+func getFullyQualifiedNamespace(creds clientCreds) string {
+	if creds.fullyQualifiedNamespace != "" {
+		return creds.fullyQualifiedNamespace
+	}
+	csp, err := conn.ParseConnectionString(creds.connectionString)
+	if err != nil {
+		return ""
+	}
+	return csp.FullyQualifiedNamespace
+}

sdk/messaging/azservicebus/client_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/sas"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/test"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 	"github.com/coder/websocket"
 	"github.com/stretchr/testify/require"
 )
@@ -471,6 +472,54 @@ func TestNewClientUnitTests(t *testing.T) {
 		require.EqualValues(t, ns.FQDN, "mysb.windows.servicebus.net")
 	})
 
+	t.Run("TracerIsSetUp", func(t *testing.T) {
+		hostName := "fake.servicebus.windows.net"
+		// when tracing provider is not set, use a no-op tracer.
+		client, err := NewClient(hostName, struct{ azcore.TokenCredential }{}, nil)
+		require.NoError(t, err)
+		require.Zero(t, client.tracer)
+		require.False(t, client.tracer.Enabled())
+
+		// when tracing provider is set, the tracer is set up with the provider.
+		provider := tracing.NewSpanValidator(t, tracing.SpanMatcher{
+			Name:   "TestSpan",
+			Status: tracing.SpanStatusUnset,
+			Attributes: []tracing.Attribute{
+				{Key: tracing.MessagingSystem, Value: "servicebus"},
+				{Key: tracing.ServerAddress, Value: hostName},
+			},
+		})
+		client, err = NewClient(hostName, struct{ azcore.TokenCredential }{}, &ClientOptions{
+			TracingProvider: provider,
+		})
+		require.NoError(t, err)
+		require.NotZero(t, client.tracer)
+		require.True(t, client.tracer.Enabled())
+
+		// ensure attributes are set up correctly.
+		_, endSpan := tracing.StartSpan(context.Background(), &tracing.TracerOptions{
+			Tracer:   client.tracer,
+			SpanName: "TestSpan",
+		})
+		endSpan(nil)
+
+		// attributes should be set up when using a connection string.
+		fakeConnectionString := "Endpoint=sb://fake.servicebus.windows.net/;SharedAccessKeyName=TestName;SharedAccessKey=TestKey"
+		client, err = NewClientFromConnectionString(fakeConnectionString, &ClientOptions{
+			TracingProvider: provider,
+		})
+		require.NoError(t, err)
+		require.NotZero(t, client.tracer)
+		require.True(t, client.tracer.Enabled())
+
+		// ensure attributes are set up correctly.
+		_, endSpan = tracing.StartSpan(context.Background(), &tracing.TracerOptions{
+			Tracer:   client.tracer,
+			SpanName: "TestSpan",
+		})
+		endSpan(nil)
+	})
+
 	t.Run("RetryOptionsArePropagated", func(t *testing.T) {
 		// retry options are passed and copied along several routes, just make sure it's properly propagated.
 		// NOTE: session receivers are checked in a separate test because they require actual SB access.

sdk/messaging/azservicebus/internal/amqpLinks.go

@@ -15,6 +15,7 @@ import (
 	azlog "github.com/Azure/azure-sdk-for-go/sdk/internal/log"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/amqpwrap"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/exported"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/utils"
 )
 
@@ -42,7 +43,7 @@ type AMQPLinks interface {
 	Get(ctx context.Context) (*LinksWithID, error)
 
 	// Retry will run your callback, recovering links when necessary.
-	Retry(ctx context.Context, name log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions) error
+	Retry(ctx context.Context, name log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions, to *tracing.TracerOptions) error
 
 	// RecoverIfNeeded will check if an error requires recovery, and will recover
 	// the link or, possibly, the connection.
@@ -315,7 +316,7 @@ func (l *AMQPLinksImpl) Get(ctx context.Context) (*LinksWithID, error) {
 	}, nil
 }
 
-func (links *AMQPLinksImpl) Retry(ctx context.Context, eventName log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions) error {
+func (links *AMQPLinksImpl) Retry(ctx context.Context, eventName log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions, to *tracing.TracerOptions) error {
 	var lastID LinkID
 
 	didQuickRetry := false
@@ -366,7 +367,7 @@ func (links *AMQPLinksImpl) Retry(ctx context.Context, eventName log.Event, oper
 		}
 
 		return nil
-	}, isFatalErrorFunc, o)
+	}, isFatalErrorFunc, o, to)
 }
 
 // EntityPath is the full entity path for the queue/topic/subscription.

sdk/messaging/azservicebus/internal/amqpLinks_test.go

@@ -455,7 +455,7 @@ func TestAMQPLinksCBSLinkStillOpen(t *testing.T) {
 	}, exported.RetryOptions{
 		RetryDelay:    -1,
 		MaxRetryDelay: time.Millisecond,
-	})
+	}, nil)
 
 	defer func() {
 		err := links.Close(context.Background(), true)
@@ -629,7 +629,7 @@ func TestAMQPLinksRetry(t *testing.T) {
 		// we do setDefaults() before we run.
 		RetryDelay:    time.Millisecond,
 		MaxRetryDelay: time.Millisecond,
-	})
+	}, nil)
 
 	var connErr *amqp.ConnError
 	require.ErrorAs(t, err, &connErr)

sdk/messaging/azservicebus/internal/amqp_test_utils.go

@@ -11,6 +11,7 @@ import (
 	azlog "github.com/Azure/azure-sdk-for-go/sdk/internal/log"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/amqpwrap"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/exported"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/utils"
 	"github.com/Azure/go-amqp"
 )
@@ -198,14 +199,19 @@ func (l *FakeAMQPLinks) Get(ctx context.Context) (*LinksWithID, error) {
 	}
 }
 
-func (l *FakeAMQPLinks) Retry(ctx context.Context, eventName log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions) error {
+func (l *FakeAMQPLinks) Retry(ctx context.Context, eventName log.Event, operation string, fn RetryWithLinksFn, o exported.RetryOptions, to *tracing.TracerOptions) error {
+	var err error
+	ctx, endSpan := tracing.StartSpan(ctx, to)
+	defer func() { endSpan(err) }()
+
 	lwr, err := l.Get(ctx)
 
 	if err != nil {
 		return err
 	}
 
-	return fn(ctx, lwr, &utils.RetryFnArgs{})
+	err = fn(ctx, lwr, &utils.RetryFnArgs{})
+	return err
 }
 
 func (l *FakeAMQPLinks) Writef(evt azlog.Event, format string, args ...any) {

sdk/messaging/azservicebus/internal/amqplinks_unit_test.go

@@ -80,7 +80,7 @@ func TestAMQPLinksRetriesUnit(t *testing.T) {
 				return testData.Err
 			}, exported.RetryOptions{
 				RetryDelay: time.Millisecond,
-			})
+			}, nil)
 
 			require.Equal(t, testData.Err, err)
 			require.Equal(t, testData.Attempts, attempts)
@@ -222,7 +222,7 @@ func TestAMQPCloseLinkTimeout_Receiver_CancellationDuringClose(t *testing.T) {
 	err := links.Retry(userCtx, exported.EventConn, "Test", func(ctx context.Context, tmpLWID *LinksWithID, args *utils.RetryFnArgs) error {
 		lwid = tmpLWID
 		return nil
-	}, exported.RetryOptions{})
+	}, exported.RetryOptions{}, nil)
 
 	require.NoError(t, err)
 	require.NotNil(t, lwid)
@@ -241,7 +241,7 @@ func TestAMQPCloseLinkTimeout_Receiver_CancellationDuringClose(t *testing.T) {
 	err = links.Retry(context.Background(), exported.EventConn, "Test", func(ctx context.Context, tmpLWID *LinksWithID, args *utils.RetryFnArgs) error {
 		lwid = tmpLWID
 		return nil
-	}, exported.RetryOptions{})
+	}, exported.RetryOptions{}, nil)
 
 	require.NoError(t, err)
 	require.NotNil(t, lwid)
@@ -280,7 +280,7 @@ func TestAMQPCloseLinkTimeout_Receiver_RecoverIfNeeded(t *testing.T) {
 	err := links.Retry(userCtx, exported.EventConn, "Test", func(ctx context.Context, tmpLWID *LinksWithID, args *utils.RetryFnArgs) error {
 		lwid = tmpLWID
 		return nil
-	}, exported.RetryOptions{})
+	}, exported.RetryOptions{}, nil)
 
 	require.NoError(t, err)
 	require.NotNil(t, lwid)

sdk/messaging/azservicebus/internal/constants.go

@@ -3,5 +3,7 @@
 
 package internal
 
+const ModuleName = "github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+
 // Version is the semantic version number
 const Version = "v1.7.4"

sdk/messaging/azservicebus/internal/namespace.go

@@ -21,6 +21,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/conn"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/exported"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/sbauth"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/utils"
 	"github.com/Azure/go-amqp"
 )
@@ -38,6 +39,8 @@ type (
 		//   PR: https://github.com/Azure/azure-sdk-for-go/pull/16847
 		connID uint64
 
+		tracer tracing.Tracer
+
 		FQDN          string
 		TokenProvider *sbauth.TokenProvider
 		tlsConfig     *tls.Config
@@ -168,6 +171,13 @@ func NamespaceWithNewClientFn(fn func(ctx context.Context) (amqpwrap.AMQPClient,
 	}
 }
 
+func NamespaceWithTracer(tracer tracing.Tracer) NamespaceOption {
+	return func(ns *Namespace) error {
+		ns.tracer = tracer
+		return nil
+	}
+}
+
 // NewNamespace creates a new namespace configured through NamespaceOption(s)
 func NewNamespace(opts ...NamespaceOption) (*Namespace, error) {
 	ns := &Namespace{}
@@ -428,6 +438,12 @@ func (ns *Namespace) startNegotiateClaimRenewer(ctx context.Context,
 
 			log.Writef(exported.EventAuth, "(%s) next refresh in %s", entityPath, nextClaimAt)
 
+			to := &tracing.TracerOptions{
+				Tracer:     ns.tracer,
+				SpanName:   tracing.NegotiateClaimSpanName,
+				Attributes: tracing.GetEntityPathAttributes(entityPath),
+			}
+
 			select {
 			case <-refreshCtx.Done():
 				return
@@ -442,7 +458,7 @@ func (ns *Namespace) startNegotiateClaimRenewer(ctx context.Context,
 
 						expiresOn = tmpExpiresOn
 						return nil
-					}, IsFatalSBError, ns.RetryOptions)
+					}, IsFatalSBError, ns.RetryOptions, to)
 
 					if err == nil {
 						break

sdk/messaging/azservicebus/internal/namespace_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/auth"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/sbauth"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/test"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/tracing"
 	"github.com/Azure/go-amqp"
 	"github.com/stretchr/testify/require"
 )
@@ -230,6 +231,14 @@ func TestNamespaceNegotiateClaimFails(t *testing.T) {
 
 func TestNamespaceNegotiateClaimFatalErrors(t *testing.T) {
 	ns := &Namespace{
+		tracer: tracing.NewSpanValidator(t, tracing.SpanMatcher{
+			Name:   "Namespace.NegotiateClaim",
+			Kind:   tracing.SpanKindInternal,
+			Status: tracing.SpanStatusError,
+			Attributes: []tracing.Attribute{
+				{Key: tracing.DestinationName, Value: "entity path"},
+			},
+		}).NewTracer("module", "version"),
 		TokenProvider: sbauth.NewTokenProvider(&fakeTokenCredential{expires: time.Now()}),
 	}