[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 {
@@ -166,13 +175,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},
@@ -192,6 +204,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},
@@ -216,6 +229,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,
@@ -238,6 +252,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},
@@ -265,6 +280,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},
@@ -334,6 +350,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,
@@ -369,3 +386,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/stretchr/testify/require"
 	"nhooyr.io/websocket"
 )
@@ -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/go.mod

@@ -5,7 +5,7 @@ go 1.18
 retract v1.1.2 // Breaks customers in situations where close is slow/infinite.
 
 require (
-	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.13.0
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.16.0
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0
 	github.com/Azure/go-amqp v1.1.0
@@ -34,9 +34,9 @@ require (
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	golang.org/x/crypto v0.25.0 // indirect
-	golang.org/x/net v0.27.0 // indirect
-	golang.org/x/sys v0.22.0 // indirect
-	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/crypto v0.27.0 // indirect
+	golang.org/x/net v0.29.0 // indirect
+	golang.org/x/sys v0.25.0 // indirect
+	golang.org/x/text v0.18.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

sdk/messaging/azservicebus/go.sum

@@ -1,5 +1,5 @@
-github.com/Azure/azure-sdk-for-go/sdk/azcore v1.13.0 h1:GJHeeA2N7xrG3q30L2UXDyuWRzDM900/65j70wcM4Ww=
-github.com/Azure/azure-sdk-for-go/sdk/azcore v1.13.0/go.mod h1:l38EPgmsp71HHLq9j7De57JcKOWPyhrsW1Awm1JS6K0=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.16.0 h1:JZg6HRh6W6U4OLl6lk7BZ7BLisIzM9dG1R50zUk9C/M=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.16.0/go.mod h1:YL1xnZ6QejvQHWJrX/AvhFl4WW4rqHVoKspWNVwFk0M=
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0 h1:tfLQ34V6F7tVSwoTf/4lH5sE0o6eCJuNDTmH09nDpbc=
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0/go.mod h1:9kIvujWAA58nmPmWB1m23fyWic1kYZMxD9CxaWn4Qpg=
 github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 h1:ywEEhmNahHBihViHepv3xPBn1663uRv2t2q/ESv9seY=
@@ -35,14 +35,14 @@ github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30=
-golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M=
+golang.org/x/crypto v0.27.0 h1:GXm2NjJrPaiv/h1tb2UH8QfgC/hOf/+z0p6PT8o1w7A=
+golang.org/x/crypto v0.27.0/go.mod h1:1Xngt8kV6Dvbssa53Ziq6Eqn0HqbZi5Z6R0ZpwQzt70=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
-golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys=
-golang.org/x/net v0.27.0/go.mod h1:dDi0PyhWNoiUOrAS8uXv/vnScO4wnHQO4mj9fn/RytE=
+golang.org/x/net v0.29.0 h1:5ORfpBpCs4HzDYoodCDBbwHzdR5UrLBZ3sOnUJmFoHo=
+golang.org/x/net v0.29.0/go.mod h1:gLkgy8jTGERgjzMic6DS9+SP0ajcu6Xu3Orq/SpETg0=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -51,13 +51,13 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
-golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
+golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
-golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
+golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224=
+golang.org/x/text v0.18.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=

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)