Repository: Azure/azure-service-bus-go
Pull request #248 (2021-08-26 23:33:58)

Use link caching for $management operations (renew locks)
@serbrech found a bug where we appear to be having issues when creating a large # of links (in his testing it's somewhere around 200-500). 

Our RPCClient (which is what they were using) created a link for _each_ call made, which added up quickly when you were doing parallel lock renewals.

This change makes it so we cache (when possible) a link per entity/address and re-use it for any other other calls that occur. This change goes hand in hand with the change made in [Azure/azure-amqp-common-go/pull/53](https://github.com/Azure/azure-amqp-common-go/pull/53)

richardpark-msft: /azp run Azure.azure-service-bus-go

azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
richardpark-msft: /azp run Azure.azure-service-bus-go


azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
richardpark-msft: /azp run Azure.azure-service-bus-go


azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
richardpark-msft: /azp run Azure.azure-service-bus-go
azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
richardpark-msft: /azp run Azure.azure-service-bus-go


azure-pipelines[bot]: <samp>
Azure Pipelines successfully started running 2 pipeline(s).<br>

</samp>
richardpark-msft: (last live test failure was due to throttling)

Patch:
diff --git a/changelog.md b/changelog.md
index 052b6fe..f900e8e 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,5 +1,11 @@
 # Change Log
 
+## `v0.11.0`
+
+- Use a singleton instance of the management link, avoiding creating a link per management
+  link operations like dispositions or lock renewals. 
+  [PR#248](https://github.com/Azure/azure-service-bus-go/pull/248)
+
 ## `v0.10.16`
 
 - fixed an issue where links weren't being closed when retrying
diff --git a/entity.go b/entity.go
index 2a254bd..1de469b 100644
--- a/entity.go
+++ b/entity.go
@@ -22,7 +22,6 @@ type (
 	}
 
 	receivingEntity struct {
-		renewMessageLockMutex sync.Mutex
 		*entity
 	}
 
@@ -192,9 +191,6 @@ func (re *receivingEntity) RenewLocks(ctx context.Context, messages ...*Message)
 	ctx, span := re.startSpanFromContext(ctx, "sb.receivingEntity.RenewLocks")
 	defer span.End()
 
-	re.renewMessageLockMutex.Lock()
-	defer re.renewMessageLockMutex.Unlock()
-
 	client, err := re.entity.GetRPCClient(ctx)
 	if err != nil {
 		tab.For(ctx).Error(err)
diff --git a/go.mod b/go.mod
index d0a39db..99130f0 100644
--- a/go.mod
+++ b/go.mod
@@ -3,9 +3,9 @@ module github.com/Azure/azure-service-bus-go
 go 1.12
 
 require (
-	github.com/Azure/azure-amqp-common-go/v3 v3.1.0
+	github.com/Azure/azure-amqp-common-go/v3 v3.1.2
 	github.com/Azure/azure-sdk-for-go v51.1.0+incompatible
-	github.com/Azure/go-amqp v0.13.11
+	github.com/Azure/go-amqp v0.13.13
 	github.com/Azure/go-autorest/autorest v0.11.18
 	github.com/Azure/go-autorest/autorest/adal v0.9.13
 	github.com/Azure/go-autorest/autorest/date v0.3.0
diff --git a/go.sum b/go.sum
index 978cafc..4f0b222 100644
--- a/go.sum
+++ b/go.sum
@@ -1,28 +1,23 @@
-github.com/Azure/azure-amqp-common-go/v3 v3.1.0 h1:1N4YSkWYWffOpQHromYdOucBSQXhNRKzqtgICy6To8Q=
-github.com/Azure/azure-amqp-common-go/v3 v3.1.0/go.mod h1:PBIGdzcO1teYoufTKMcGibdKaYZv4avS+O6LNIp8bq0=
+github.com/Azure/azure-amqp-common-go/v3 v3.1.2 h1:Xu9nboejX9a+7uOGiwEUlpD8/oN4A6f7ZmHTMmJpPkw=
+github.com/Azure/azure-amqp-common-go/v3 v3.1.2/go.mod h1:zN7QL/vfCsq3XQxQaTkg4ScO786CA2rQnZ1LXX7QryE=
 github.com/Azure/azure-sdk-for-go v51.1.0+incompatible h1:7uk6GWtUqKg6weLv2dbKnzwb0ml1Qn70AdtRccZ543w=
 github.com/Azure/azure-sdk-for-go v51.1.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
-github.com/Azure/go-amqp v0.13.0/go.mod h1:qj+o8xPCz9tMSbQ83Vp8boHahuRDl5mkNHyt1xlxUTs=
-github.com/Azure/go-amqp v0.13.11 h1:E28zKoWuzO4+D80iUD88BUorI5PqvIZ/S/77md3hIvA=
-github.com/Azure/go-amqp v0.13.11/go.mod h1:D5ZrjQqB1dyp1A+G73xeL/kNn7D5qHJIIsNNps7YNmk=
+github.com/Azure/go-amqp v0.13.13 h1:OBPwCO50EzniOyZR0M4VbGJYDxceIy3SFOnKVMJktdY=
+github.com/Azure/go-amqp v0.13.13/go.mod h1:D5ZrjQqB1dyp1A+G73xeL/kNn7D5qHJIIsNNps7YNmk=
 github.com/Azure/go-autorest v14.2.0+incompatible h1:V5VMDjClD3GiElqLWO7mz2MxNAK/vTfRHdAubSIPRgs=
 github.com/Azure/go-autorest v14.2.0+incompatible/go.mod h1:r+4oMnoxhatjLLJ6zxSWATqVooLgysK6ZNox3g/xq24=
-github.com/Azure/go-autorest/autorest v0.11.3/go.mod h1:JFgpikqFJ/MleTTxwepExTKnFUKKszPS8UavbQYUMuw=
 github.com/Azure/go-autorest/autorest v0.11.18 h1:90Y4srNYrwOtAgVo3ndrQkTYn6kf1Eg/AjTFJ8Is2aM=
 github.com/Azure/go-autorest/autorest v0.11.18/go.mod h1:dSiJPy22c3u0OtOKDNttNgqpNFY/GeWa7GH/Pz56QRA=
-github.com/Azure/go-autorest/autorest/adal v0.9.0/go.mod h1:/c022QCutn2P7uY+/oQWWNcK9YU+MH96NgK+jErpbcg=
 github.com/Azure/go-autorest/autorest/adal v0.9.13 h1:Mp5hbtOePIzM8pJVRa3YLrWWmZtoxRXqUEzCfJt3+/Q=
 github.com/Azure/go-autorest/autorest/adal v0.9.13/go.mod h1:W/MM4U6nLxnIskrw4UwWzlHfGjwUS50aOsc/I3yuU8M=
 github.com/Azure/go-autorest/autorest/date v0.3.0 h1:7gUk1U5M/CQbp9WoqinNzJar+8KY+LPI6wiWrP/myHw=
 github.com/Azure/go-autorest/autorest/date v0.3.0/go.mod h1:BI0uouVdmngYNUzGWeSYnokU+TrmwEsOqdt8Y6sso74=
-github.com/Azure/go-autorest/autorest/mocks v0.4.0/go.mod h1:LTp+uSrOhSkaKrUy935gNZuuIPPVsHlr9DSOxSayd+k=
 github.com/Azure/go-autorest/autorest/mocks v0.4.1 h1:K0laFcLE6VLTOwNgSxaGbUcLPuGXlNkbVvq4cW4nIHk=
 github.com/Azure/go-autorest/autorest/mocks v0.4.1/go.mod h1:LTp+uSrOhSkaKrUy935gNZuuIPPVsHlr9DSOxSayd+k=
 github.com/Azure/go-autorest/autorest/to v0.4.0 h1:oXVqrxakqqV1UZdSazDOPOLvOIz+XA683u8EctwboHk=
 github.com/Azure/go-autorest/autorest/to v0.4.0/go.mod h1:fE8iZBn7LQR7zH/9XU2NcPR4o9jEImooCeWJcYV/zLE=
 github.com/Azure/go-autorest/autorest/validation v0.3.1 h1:AgyqjAd94fwNAoTjl/WQXg4VvFeRFpO+UhNyRXqF1ac=
 github.com/Azure/go-autorest/autorest/validation v0.3.1/go.mod h1:yhLgjC0Wda5DYXl6JAsWyUe4KVNffhoDhG0zVzUMo3E=
-github.com/Azure/go-autorest/logger v0.2.0/go.mod h1:T9E3cAhj2VqvPOtCYAvby9aBXkZmbF5NWuPV8+WeEW8=
 github.com/Azure/go-autorest/logger v0.2.1 h1:IG7i4p/mDa2Ce4TRyAO8IHnVhAVF3RFU+ZtXWSmf4Tg=
 github.com/Azure/go-autorest/logger v0.2.1/go.mod h1:T9E3cAhj2VqvPOtCYAvby9aBXkZmbF5NWuPV8+WeEW8=
 github.com/Azure/go-autorest/tracing v0.6.0 h1:TYi4+3m5t6K48TGI9AUdb+IzbnSxvnvUMfuitfgcfuo=
@@ -32,7 +27,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/devigned/tab v0.1.1 h1:3mD6Kb1mUOYeLpJvTVSDwSg5ZsfSxfvxGRTxRsJsITA=
 github.com/devigned/tab v0.1.1/go.mod h1:XG9mPq0dFghrYvoBF3xdRrJzSTX1b7IQrvaL9mzjeJY=
-github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible h1:TcekIExNqud5crz4xD2pavyTgWiPvpYe4Xau31I0PRk=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible/go.mod h1:pbq4aXjuKjdthFRnoDwaVPLA+WlJuPGy+QneDUgJi2k=
 github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw=
@@ -95,8 +89,6 @@ github.com/ugorji/go v1.1.7/go.mod h1:kZn38zHttfInRq0xu/PH0az30d+z6vm202qpg1oXVM
 github.com/ugorji/go/codec v1.1.7 h1:2SvQaVZ1ouYrrKKwoSk2pzd4A9evlKJb9oTL+OaLUSs=
 github.com/ugorji/go/codec v1.1.7/go.mod h1:Ax+UKWsSmolVDwsd+7N3ZtXu+yMGCf907BLYF3GoBXY=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.0.0-20200728195943-123391ffb6de/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201002170205-7f63de1d35b0 h1:hb9wdF1z5waM+dSIICn1l0DkLVDT3hqhhQsDNUmHPRE=
 golang.org/x/crypto v0.0.0-20201002170205-7f63de1d35b0/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
diff --git a/namespace.go b/namespace.go
index f49efd0..d83c8de 100644
--- a/namespace.go
+++ b/namespace.go
@@ -75,6 +75,11 @@ type (
 		initRefresh sync.Once
 		// populated with the result from auto-refresh, to be called elsewhere
 		cancelRefresh func() <-chan struct{}
+
+		// for testing
+
+		// alias for 'amqp.Dial'
+		amqpDial func(addr string, opts ...amqp.ConnOption) (*amqp.Client, error)
 	}
 
 	// NamespaceOption provides structure for configuring a new Service Bus namespace
@@ -189,6 +194,7 @@ func NamespaceWithTokenProvider(provider auth.TokenProvider) NamespaceOption {
 func NewNamespace(opts ...NamespaceOption) (*Namespace, error) {
 	ns := &Namespace{
 		Environment: azure.PublicCloud,
+		amqpDial:    amqp.Dial,
 	}
 
 	for _, opt := range opts {
@@ -235,7 +241,7 @@ func (ns *Namespace) newClient(ctx context.Context) (*amqp.Client, error) {
 		return amqp.New(nConn, append(defaultConnOptions, amqp.ConnServerHostname(ns.getHostname()))...)
 	}
 
-	return amqp.Dial(ns.getAMQPHostURI(), defaultConnOptions...)
+	return ns.amqpDial(ns.getAMQPHostURI(), defaultConnOptions...)
 }
 
 // negotiateClaim performs initial authentication and starts periodic refresh of credentials.
diff --git a/rpc.go b/rpc.go
index a1290e0..5af4214 100644
--- a/rpc.go
+++ b/rpc.go
@@ -39,12 +39,26 @@ import (
 
 type (
 	rpcClient struct {
-		ec                 entityConnector
-		client             *amqp.Client
-		clientMu           sync.RWMutex
+		ec     entityConnector
+		client *amqp.Client
+
+		clientMu  sync.RWMutex
+		linkCache map[string]*rpc.Link // stores 'address' => rpc.Link.
+
 		sessionID          *string
 		isSessionFilterSet bool
 		cancelAuthRefresh  func() <-chan struct{}
+
+		// replaceable for testing
+
+		// alias of 'rpc.NewLink'
+		newRPCLink func(conn *amqp.Client, address string, opts ...rpc.LinkOption) (*rpc.Link, error)
+
+		// alias of function 'newAMQPClient'
+		newAMQPClient func(ctx context.Context, ec entityConnector) (*amqp.Client, func() <-chan struct{}, error)
+
+		// alias of 'amqp.Client.Close()'
+		closeAMQPClient func() error
 	}
 
 	rpcClientOption func(*rpcClient) error
@@ -52,36 +66,46 @@ type (
 
 func newRPCClient(ctx context.Context, ec entityConnector, opts ...rpcClientOption) (*rpcClient, error) {
 	r := &rpcClient{
-		ec: ec,
+		ec:            ec,
+		linkCache:     map[string]*rpc.Link{},
+		newRPCLink:    rpc.NewLink,
+		newAMQPClient: newAMQPClient,
 	}
 
+	r.closeAMQPClient = func() error { return r.client.Close() }
+
 	for _, opt := range opts {
 		if err := opt(r); err != nil {
 			tab.For(ctx).Error(err)
 			return nil, err
 		}
 	}
-	if err := r.newClient(ctx); err != nil {
+
+	var err error
+	r.client, r.cancelAuthRefresh, err = r.newAMQPClient(ctx, r.ec)
+
+	if err != nil {
 		tab.For(ctx).Error(err)
 		return nil, err
 	}
 	return r, nil
 }
 
-// newClient will replace the existing client and start auth auto-refresh.
-// any pre-existing client MUST be closed before calling this method.
-// NOTE: this does *not* take the write lock, callers must hold it as required!
-func (r *rpcClient) newClient(ctx context.Context) error {
-	var err error
-	r.client, err = r.ec.Namespace().newClient(ctx)
+// newAMQPClient creates a client and starts auth auto-refresh.
+func newAMQPClient(ctx context.Context, ec entityConnector) (*amqp.Client, func() <-chan struct{}, error) {
+	client, err := ec.Namespace().newClient(ctx)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
-	r.cancelAuthRefresh, err = r.ec.Namespace().negotiateClaim(ctx, r.client, r.ec.ManagementPath())
+
+	cancelAuthRefresh, err := ec.Namespace().negotiateClaim(ctx, client, ec.ManagementPath())
+
 	if err != nil {
-		return err
+		client.Close()
+		return nil, nil, err
 	}
-	return nil
+
+	return client, cancelAuthRefresh, nil
 }
 
 // Recover will attempt to close the current session and link, then rebuild them
@@ -92,7 +116,11 @@ func (r *rpcClient) Recover(ctx context.Context) error {
 	r.clientMu.Lock()
 	defer r.clientMu.Unlock()
 	_ = r.close()
-	if err := r.newClient(ctx); err != nil {
+
+	var err error
+	r.client, r.cancelAuthRefresh, err = r.newAMQPClient(ctx, r.ec)
+
+	if err != nil {
 		tab.For(ctx).Error(err)
 		return err
 	}
@@ -111,7 +139,60 @@ func (r *rpcClient) close() error {
 	if r.cancelAuthRefresh != nil {
 		<-r.cancelAuthRefresh()
 	}
-	return r.client.Close()
+
+	// we don't want to interrupt the cleanup of these links.
+	r.resetLinkCache(context.Background())
+
+	return r.closeAMQPClient()
+}
+
+// resetLinkCache removes any cached links with the assumption that we are
+// recovering and do not care about errors that occur with links.
+func (r *rpcClient) resetLinkCache(ctx context.Context) {
+	var links []*rpc.Link
+
+	for _, link := range r.linkCache {
+		links = append(links, link)
+	}
+
+	// make a new links map - Recover() also goes through here (through `close`)
+	// so we should be prepared to handle more cached links.
+	r.linkCache = map[string]*rpc.Link{}
+
+	for _, link := range links {
+		link.Close(ctx)
+	}
+}
+
+func (r *rpcClient) getCachedLink(ctx context.Context, address string) (*rpc.Link, error) {
+	r.clientMu.RLock()
+	link, ok := r.linkCache[address]
+	r.clientMu.RUnlock()
+
+	if ok {
+		return link, nil
+	}
+
+	// possibly need to create the link
+	r.clientMu.Lock()
+	defer r.clientMu.Unlock()
+
+	// might have been added in between us checking and
+	// us getting the lock.
+	link, ok = r.linkCache[address]
+
+	if ok {
+		return link, nil
+	}
+
+	link, err := r.newRPCLink(r.client, address)
+
+	if err != nil {
+		return nil, err
+	}
+
+	r.linkCache[address] = link
+	return link, nil
 }
 
 // creates a new link and sends the RPC request, recovering and retrying on certain AMQP errors
@@ -127,19 +208,35 @@ func (r *rpcClient) doRPCWithRetry(ctx context.Context, address string, msg *amq
 		var link *rpc.Link
 		var rsp *rpc.Response
 		var err error
-		link, err = rpc.NewLink(client, address, opts...)
+		isCachedLink := false
+
+		if len(opts) == 0 {
+			// Can use a cached client since there's nothing unique about this particular
+			// link for this call.
+			// It looks like creating a management link with a session filter is the
+			// only time we don't do this.
+			isCachedLink = true
+			link, err = r.getCachedLink(ctx, address)
+		} else {
+			link, err = rpc.NewLink(client, address)
+		}
+
 		if err == nil {
 			defer func() {
-				if link == nil {
+				if isCachedLink || link == nil {
+					// cached links are closed by `rpcClient.close()`
+					// (which is also called during recovery)
 					return
 				}
+
 				link.Close(ctx)
 			}()
 			rsp, err = link.RetryableRPC(ctx, times, delay, msg)
 			if err == nil {
-				return rsp, err
+				return rsp, nil
 			}
 		}
+
 		if sendCount >= amqpRetryDefaultTimes || !isAMQPTransientError(ctx, err) {
 			return nil, err
 		}
@@ -149,10 +246,12 @@ func (r *rpcClient) doRPCWithRetry(ctx context.Context, address string, msg *amq
 		_, retryErr := common.Retry(amqpRetryDefaultTimes, amqpRetryDefaultDelay, func() (interface{}, error) {
 			ctx, sp := r.startProducerSpanFromContext(ctx, "sb.rpcClient.doRPCWithRetry.tryRecover")
 			defer sp.End()
+
 			if err := r.Recover(ctx); err == nil {
 				tab.For(ctx).Debug("recovered RPC connection")
 				return nil, nil
 			}
+
 			select {
 			case <-ctx.Done():
 				return nil, ctx.Err()


Test patch:
diff --git a/internal/stress/mgmtlink/renewLocksStressTest.go b/internal/stress/mgmtlink/renewLocksStressTest.go
new file mode 100644
index 0000000..9bda36a
--- /dev/null
+++ b/internal/stress/mgmtlink/renewLocksStressTest.go
@@ -0,0 +1,114 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	servicebus "github.com/Azure/azure-service-bus-go"
+	"github.com/joho/godotenv"
+)
+
+const queuePrefetch = 1000
+const renewalsPerMessage = 2
+
+func main() {
+	ctx := context.Background()
+
+	godotenv.Load(".env")
+	cs := os.Getenv("SERVICEBUS_CONNECTION_STRING")
+	queueName := "samples"
+
+	ns, err := servicebus.NewNamespace(servicebus.NamespaceWithConnectionString(cs))
+
+	if err != nil {
+		log.Fatalf("Failed to create namespace client: %s", err.Error())
+	}
+
+	sender, err := ns.NewSender(ctx, queueName)
+
+	if err != nil {
+		log.Fatalf("Failed to create sender: %s", err.Error())
+	}
+
+	ch := make(chan bool, 100)
+
+	for i := 0; i < queuePrefetch; i++ {
+		ch <- true
+		go func(i int) {
+			defer func() { <-ch }()
+			if err = sender.Send(ctx, &servicebus.Message{Data: []byte(fmt.Sprintf("hello world %d", i))}); err != nil {
+				log.Fatalf("Failed to send message: %s", err.Error())
+			}
+		}(i)
+	}
+
+	queue, err := ns.NewQueue(queueName, servicebus.QueueWithPrefetchCount(queuePrefetch))
+
+	if err != nil {
+		log.Fatalf("Failed to create receiver: %s", err.Error())
+	}
+
+	renewals := int32(0)
+	outstandingRenewals := int32(0)
+	failedRenewals := int32(0)
+	lastRenewalWasFailure := int32(0)
+
+	wg := &sync.WaitGroup{}
+
+	go func() {
+		ticker := time.NewTicker(time.Second * 5)
+
+		for {
+			select {
+			case <-ticker.C:
+				log.Printf("Messages: [total: %d, outstanding: %d, failed: %d]", atomic.LoadInt32(&renewals), atomic.LoadInt32(&outstandingRenewals), atomic.LoadInt32(&failedRenewals))
+			}
+		}
+	}()
+
+	err = queue.Receive(ctx, servicebus.HandlerFunc(func(c context.Context, m *servicebus.Message) error {
+		wg.Add(1)
+
+		go func() {
+			worked := false
+
+			for i := 0; i < renewalsPerMessage; i++ {
+				atomic.AddInt32(&outstandingRenewals, 1)
+				if err := queue.RenewLocks(ctx, m); err != nil {
+					worked = false
+					atomic.AddInt32(&failedRenewals, 1)
+					log.Printf("ERROR renewing: %+v", err)
+				} else {
+					worked = true
+				}
+
+				atomic.AddInt32(&outstandingRenewals, -1)
+				atomic.AddInt32(&renewals, 1)
+			}
+
+			m.Complete(ctx)
+
+			if !worked {
+				log.Printf("Last renewal was a failure here.")
+				atomic.AddInt32(&lastRenewalWasFailure, 1)
+			}
+
+			wg.Done()
+		}()
+
+		return nil
+	}))
+
+	log.Printf("Last Renewal was failure: %d", lastRenewalWasFailure)
+
+	if err != nil {
+		log.Fatalf("Failed to receive messages")
+	}
+
+	wg.Wait()
+}
diff --git a/rpc_test.go b/rpc_test.go
new file mode 100644
index 0000000..62cec52
--- /dev/null
+++ b/rpc_test.go
@@ -0,0 +1,118 @@
+package servicebus
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Azure/azure-amqp-common-go/v3/rpc"
+	"github.com/Azure/go-amqp"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRPCLinkCaching(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	t.Run("WithClose", func(t *testing.T) {
+		fake := createFakeRPCClient()
+
+		// validate our internal cache properly tracked the link
+		// now getting the link again should get us the cached link.
+		addCachedLink(ctx, t, fake)
+
+		// close the client itself
+		require.NoError(t, fake.rpcClient.Close())
+
+		require.True(t, fake.amqpClientClosed, "AMQP client was closed")
+		require.Empty(t, fake.rpcClient.linkCache, "link cache is cleared when PRC client is closed")
+	})
+
+	t.Run("WithRecover", func(t *testing.T) {
+		fake := createFakeRPCClient()
+
+		// validate our internal cache properly tracked the link
+		// now getting the link again should get us the cached link.
+		addCachedLink(ctx, t, fake)
+
+		// Recover the client instead.
+		require.NoError(t, fake.rpcClient.Recover(ctx))
+
+		require.True(t, fake.amqpClientClosed, "AMQP client was closed")
+		require.Empty(t, fake.rpcClient.linkCache, "link cache is cleared when PRC client recovers")
+
+		// sanity check: just make sure nothing breaks if we close afterwards
+		require.NoError(t, fake.rpcClient.Close())
+	})
+}
+
+func addCachedLink(ctx context.Context, t *testing.T, fake *fakeRPCClient) {
+	createdLink, err := fake.rpcClient.getCachedLink(ctx, "an address")
+
+	require.NoError(t, err)
+	require.EqualValues(t, 1, len(fake.createdLinks))
+	require.NotNil(t, createdLink)
+
+	linkFromInternalCache, exists := fake.rpcClient.linkCache["an address"]
+	require.True(t, exists)
+	require.Same(t, linkFromInternalCache, createdLink)
+
+	sameOldLink, err := fake.rpcClient.getCachedLink(ctx, "an address")
+	require.NoError(t, err)
+
+	require.Same(t, sameOldLink, createdLink, "Same link instance should always be returned")
+}
+
+type fakeRPCClient struct {
+	*rpcClient
+	newAMQPClientCreated bool
+	amqpClientClosed     bool
+	createdLinks         []*rpc.Link
+}
+
+func createFakeRPCClient() *fakeRPCClient {
+	fake := &fakeRPCClient{}
+
+	fake.rpcClient = &rpcClient{
+		ec:        &fakeEntityConnector{},
+		linkCache: map[string]*rpc.Link{},
+
+		newAMQPClient: func(ctx context.Context, ec entityConnector) (*amqp.Client, func() <-chan struct{}, error) {
+			fake.newAMQPClientCreated = true
+			ch := make(chan struct{})
+			close(ch) // NOTE, it's just a little lie (ie - we'll just pretend cancellation of creds refresh worked INSTANTLY)
+
+			return &amqp.Client{}, func() <-chan struct{} {
+				return ch
+			}, nil
+		},
+		newRPCLink: func(conn *amqp.Client, address string, opts ...rpc.LinkOption) (*rpc.Link, error) {
+			fake.createdLinks = append(fake.createdLinks, &rpc.Link{})
+			return fake.createdLinks[len(fake.createdLinks)-1], nil
+		},
+		closeAMQPClient: func() error {
+			fake.amqpClientClosed = true
+			return nil
+		},
+	}
+
+	return fake
+}
+
+type fakeEntityConnector struct{}
+
+func (fc *fakeEntityConnector) ManagementPath() string {
+	return "ThisIsTheManagementPath"
+}
+
+func (fc *fakeEntityConnector) Namespace() *Namespace {
+	return &Namespace{
+		amqpDial: func(addr string, opts ...amqp.ConnOption) (*amqp.Client, error) {
+			return &amqp.Client{}, nil
+		},
+	}
+}
+
+func (fc *fakeEntityConnector) getEntity() *entity {
+	return nil
+}