[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,7 @@ 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 {
@@ -170,13 +179,16 @@ func newClientImpl(creds clientCreds, args clientImplArgs) (*Client, error) {
 	nsOptions = append(nsOptions, args.NSOptions...)
 
 	client.namespace, err = internal.NewNamespace(nsOptions...)
+	client.tracer = newTracer(tracingProvider, getFullyQualifiedNamespace(creds))
+
 	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 +208,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 +233,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 +256,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 +284,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 +354,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 +390,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,48 @@ 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(), client.tracer, tracing.NewSpanConfig("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(), client.tracer, tracing.NewSpanConfig("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, sc *tracing.SpanConfig) error
 
 	// RecoverIfNeeded will check if an error requires recovery, and will recover
 	// the link or, possibly, the connection.
@@ -66,6 +67,12 @@ type AMQPLinks interface {
 
 	// Prefix is the current logging prefix, usable for logging and continuity.
 	Prefix() string
+
+	// Tracer returns the tracer for the AMQPLinks instance.
+	Tracer() tracing.Tracer
+
+	// SetTracer sets the tracer for the AMQPLinks instance.
+	SetTracer(tracing.Tracer)
 }
 
 // AMQPLinksImpl manages the set of AMQP links (and detritus) typically needed to work
@@ -85,6 +92,8 @@ type AMQPLinksImpl struct {
 	//   PR: https://github.com/Azure/azure-sdk-for-go/pull/16847
 	id LinkID
 
+	tracer tracing.Tracer
+
 	entityPath     string
 	managementPath string
 	audience       string
@@ -122,6 +131,7 @@ type AMQPLinksImpl struct {
 type CreateLinkFunc func(ctx context.Context, session amqpwrap.AMQPSession) (amqpwrap.AMQPSenderCloser, amqpwrap.AMQPReceiverCloser, error)
 
 type NewAMQPLinksArgs struct {
+	Tracer              tracing.Tracer
 	NS                  NamespaceForAMQPLinks
 	EntityPath          string
 	CreateLinkFunc      CreateLinkFunc
@@ -132,6 +142,7 @@ type NewAMQPLinksArgs struct {
 // management link for a specific entity path.
 func NewAMQPLinks(args NewAMQPLinksArgs) AMQPLinks {
 	l := &AMQPLinksImpl{
+		tracer:              args.Tracer,
 		entityPath:          args.EntityPath,
 		managementPath:      fmt.Sprintf("%s/$management", args.EntityPath),
 		audience:            args.NS.GetEntityAudience(args.EntityPath),
@@ -145,6 +156,14 @@ func NewAMQPLinks(args NewAMQPLinksArgs) AMQPLinks {
 	return l
 }
 
+func (links *AMQPLinksImpl) Tracer() tracing.Tracer {
+	return links.tracer
+}
+
+func (links *AMQPLinksImpl) SetTracer(tracer tracing.Tracer) {
+	links.tracer = tracer
+}
+
 // ManagementPath is the management path for the associated entity.
 func (links *AMQPLinksImpl) ManagementPath() string {
 	return links.managementPath
@@ -315,7 +334,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, sc *tracing.SpanConfig) error {
 	var lastID LinkID
 
 	didQuickRetry := false
@@ -324,7 +343,7 @@ func (links *AMQPLinksImpl) Retry(ctx context.Context, eventName log.Event, oper
 		return links.getRecoveryKindFunc(err) == RecoveryKindFatal
 	}
 
-	return utils.Retry(ctx, eventName, links.Prefix()+"("+operation+")", func(ctx context.Context, args *utils.RetryFnArgs) error {
+	return utils.Retry(ctx, links.tracer, eventName, links.Prefix()+"("+operation+")", func(ctx context.Context, args *utils.RetryFnArgs) error {
 		if err := links.RecoverIfNeeded(ctx, lastID, args.LastErr); err != nil {
 			return err
 		}
@@ -366,7 +385,7 @@ func (links *AMQPLinksImpl) Retry(ctx context.Context, eventName log.Event, oper
 		}
 
 		return nil
-	}, isFatalErrorFunc, o)
+	}, isFatalErrorFunc, o, sc)
 }
 
 // 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"
 )
@@ -42,6 +43,8 @@ type FakeAMQPSession struct {
 type FakeAMQPLinks struct {
 	AMQPLinks
 
+	Tr tracing.Tracer
+
 	Closed              int
 	CloseIfNeededCalled int
 
@@ -198,14 +201,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, sc *tracing.SpanConfig) error {
+	var err error
+	ctx, endSpan := tracing.StartSpan(ctx, l.Tr, sc)
+	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) {
@@ -216,6 +224,14 @@ func (l *FakeAMQPLinks) Prefix() string {
 	return "prefix"
 }
 
+func (l *FakeAMQPLinks) Tracer() tracing.Tracer {
+	return l.Tr
+}
+
+func (l *FakeAMQPLinks) SetTracer(t tracing.Tracer) {
+	l.Tr = t
+}
+
 func (l *FakeAMQPLinks) Close(ctx context.Context, permanently bool) error {
 	if permanently {
 		l.permanently = true

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"