[azservicebus] Fixing an error where the RPCLink wasn't very robust when it came to connection/link errors (#17389)

The RPC link didn't have a good comprehensive set of checks for errors that came back, which would result in the response router continually failing. Fixed now to use the same GetRecoveryKind() logic as the rest of the code does. Also, added in some better logging for some of the edge error cases, just so we know what's happening. This has potentially been hiding some failure cases, including scenarios where the link should have stopped (so it can properly recover).

Also, some refactoring - the current go-amqp type structure is a bit hard to mock because concrete "client" types come back from other concrete types (for instance, amqp.Session has `GetReceiver() *amqp.Receiver`, which makes it hard to control the actual receiver object. I've added in an `amqpwrap` package where I've gone full-wrapper for client and session so I can do proper mocking in tests. That code is only used for rpc.go for now but I'll be migrating it over to others as we go.

Author: richardpark-msft

sdk/messaging/azservicebus/internal/amqpInterfaces.go

@@ -6,31 +6,14 @@ package internal
 import (
 	"context"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/amqpwrap"
 	"github.com/Azure/go-amqp"
 )
 
-// AMQPReceiver is implemented by *amqp.Receiver
-type AMQPReceiver interface {
-	IssueCredit(credit uint32) error
-	DrainCredit(ctx context.Context) error
-	Receive(ctx context.Context) (*amqp.Message, error)
-	Prefetched(ctx context.Context) (*amqp.Message, error)
-
-	// settlement functions
-	AcceptMessage(ctx context.Context, msg *amqp.Message) error
-	RejectMessage(ctx context.Context, msg *amqp.Message, e *amqp.Error) error
-	ReleaseMessage(ctx context.Context, msg *amqp.Message) error
-	ModifyMessage(ctx context.Context, msg *amqp.Message, deliveryFailed, undeliverableHere bool, messageAnnotations amqp.Annotations) error
-
-	LinkName() string
-	LinkSourceFilterValue(name string) interface{}
-}
-
-// AMQPReceiver is implemented by *amqp.Receiver
-type AMQPReceiverCloser interface {
-	AMQPReceiver
-	Close(ctx context.Context) error
-}
+type AMQPReceiver = amqpwrap.AMQPReceiver
+type AMQPReceiverCloser = amqpwrap.AMQPReceiverCloser
+type AMQPSender = amqpwrap.AMQPSender
+type AMQPSenderCloser = amqpwrap.AMQPSenderCloser
 
 // AMQPSession is implemented by *amqp.Session
 type AMQPSession interface {
@@ -44,18 +27,6 @@ type AMQPSessionCloser interface {
 	Close(ctx context.Context) error
 }
 
-// AMQPSender is implemented by *amqp.Sender
-type AMQPSender interface {
-	Send(ctx context.Context, msg *amqp.Message) error
-	MaxMessageSize() uint64
-}
-
-// AMQPSenderCloser is implemented by *amqp.Sender
-type AMQPSenderCloser interface {
-	AMQPSender
-	Close(ctx context.Context) error
-}
-
 // RPCLink is implemented by *rpc.Link
 type RPCLink interface {
 	Close(ctx context.Context) error

sdk/messaging/azservicebus/internal/amqpwrap/amqpwrap.go

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+// Package amqpwrap has some simple wrappers to make it easier to
+// abstract the go-amqp types.
+package amqpwrap
+
+import (
+	"context"
+
+	"github.com/Azure/go-amqp"
+)
+
+// AMQPReceiver is implemented by *amqp.Receiver
+type AMQPReceiver interface {
+	IssueCredit(credit uint32) error
+	DrainCredit(ctx context.Context) error
+	Receive(ctx context.Context) (*amqp.Message, error)
+	Prefetched(ctx context.Context) (*amqp.Message, error)
+
+	// settlement functions
+	AcceptMessage(ctx context.Context, msg *amqp.Message) error
+	RejectMessage(ctx context.Context, msg *amqp.Message, e *amqp.Error) error
+	ReleaseMessage(ctx context.Context, msg *amqp.Message) error
+	ModifyMessage(ctx context.Context, msg *amqp.Message, deliveryFailed, undeliverableHere bool, messageAnnotations amqp.Annotations) error
+
+	LinkName() string
+	LinkSourceFilterValue(name string) interface{}
+}
+
+// AMQPReceiverCloser is implemented by *amqp.Receiver
+type AMQPReceiverCloser interface {
+	AMQPReceiver
+	Close(ctx context.Context) error
+}
+
+// AMQPSender is implemented by *amqp.Sender
+type AMQPSender interface {
+	Send(ctx context.Context, msg *amqp.Message) error
+	MaxMessageSize() uint64
+}
+
+// AMQPSenderCloser is implemented by *amqp.Sender
+type AMQPSenderCloser interface {
+	AMQPSender
+	Close(ctx context.Context) error
+}
+
+// AMQPSession is a simple interface, implemented by *AMQPSessionWrapper.
+// It exists only so we can return AMQPReceiver/AMQPSender interfaces.
+type AMQPSession interface {
+	Close(ctx context.Context) error
+	NewReceiver(opts ...amqp.LinkOption) (AMQPReceiverCloser, error)
+	NewSender(opts ...amqp.LinkOption) (AMQPSenderCloser, error)
+}
+
+type AMQPClient interface {
+	Close() error
+	NewSession(opts ...amqp.SessionOption) (AMQPSession, error)
+}
+
+// AMQPClientWrapper is a simple interface, implemented by *AMQPClientWrapper
+// It exists only so we can return AMQPSession, which itself only exists so we can
+// return interfaces for AMQPSender and AMQPReceiver from AMQPSession.
+type AMQPClientWrapper struct {
+	Inner *amqp.Client
+}
+
+func (w *AMQPClientWrapper) Close() error {
+	return w.Inner.Close()
+}
+
+func (w *AMQPClientWrapper) NewSession(opts ...amqp.SessionOption) (AMQPSession, error) {
+	sess, err := w.Inner.NewSession(opts...)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &AMQPSessionWrapper{
+		inner: sess,
+	}, nil
+}
+
+type AMQPSessionWrapper struct {
+	inner *amqp.Session
+}
+
+func (w *AMQPSessionWrapper) Close(ctx context.Context) error {
+	return w.inner.Close(ctx)
+}
+
+func (w *AMQPSessionWrapper) NewReceiver(opts ...amqp.LinkOption) (AMQPReceiverCloser, error) {
+	return w.inner.NewReceiver(opts...)
+}
+
+func (w *AMQPSessionWrapper) NewSender(opts ...amqp.LinkOption) (AMQPSenderCloser, error) {
+	return w.inner.NewSender(opts...)
+}

sdk/messaging/azservicebus/internal/cbs.go

@@ -7,6 +7,7 @@ import (
 	"context"
 
 	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/internal/auth"
 	"github.com/Azure/go-amqp"
 )
@@ -23,7 +24,7 @@ const (
 // NegotiateClaim attempts to put a token to the $cbs management endpoint to negotiate auth for the given audience
 func NegotiateClaim(ctx context.Context, audience string, conn *amqp.Client, provider auth.TokenProvider) error {
 	link, err := NewRPCLink(RPCLinkArgs{
-		Client:  conn,
+		Client:  &amqpwrap.AMQPClientWrapper{Inner: conn},
 		Address: cbsAddress,
 	})
 

sdk/messaging/azservicebus/internal/namespace.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"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/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"
@@ -224,7 +225,7 @@ func (ns *Namespace) NewRPCLink(ctx context.Context, managementPath string) (RPC
 	}
 
 	return NewRPCLink(RPCLinkArgs{
-		Client:  client,
+		Client:  &amqpwrap.AMQPClientWrapper{Inner: client},
 		Address: managementPath,
 	})
 }

sdk/messaging/azservicebus/internal/rpc.go

@@ -10,7 +10,9 @@ import (
 	"sync"
 	"time"
 
+	azlog "github.com/Azure/azure-sdk-for-go/sdk/internal/log"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/uuid"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/amqpwrap"
 	"github.com/Azure/go-amqp"
 )
 
@@ -24,10 +26,9 @@ const (
 type (
 	// rpcLink is the bidirectional communication structure used for CBS negotiation
 	rpcLink struct {
-		session *amqp.Session
-
-		receiver amqpReceiver // *amqp.Receiver
-		sender   amqpSender   // *amqp.Sender
+		session  amqpwrap.AMQPSession
+		receiver amqpwrap.AMQPReceiverCloser // *amqp.Receiver
+		sender   amqpwrap.AMQPSenderCloser   // *amqp.Sender
 
 		clientAddress string
 		sessionID     *string
@@ -38,9 +39,10 @@ type (
 		responseMap             map[string]chan rpcResponse
 		broadcastErr            error // the error that caused the responseMap to be nil'd
 
+		logEvent azlog.Event
+
 		// for unit tests
-		uuidNewV4     func() (uuid.UUID, error)
-		messageAccept func(ctx context.Context, message *amqp.Message) error
+		uuidNewV4 func() (uuid.UUID, error)
 	}
 
 	// RPCResponse is the simplified response structure from an RPC like call
@@ -57,17 +59,6 @@ type (
 		message *amqp.Message
 		err     error
 	}
-
-	// Actually: *amqp.Receiver
-	amqpReceiver interface {
-		Receive(ctx context.Context) (*amqp.Message, error)
-		Close(ctx context.Context) error
-	}
-
-	amqpSender interface {
-		Send(ctx context.Context, msg *amqp.Message) error
-		Close(ctx context.Context) error
-	}
 )
 
 type rpcError struct {
@@ -84,8 +75,9 @@ func (e rpcError) RPCCode() int {
 }
 
 type RPCLinkArgs struct {
-	Client  *amqp.Client
-	Address string
+	Client   amqpwrap.AMQPClient
+	Address  string
+	LogEvent azlog.Event
 }
 
 // NewRPCLink will build a new request response link
@@ -110,6 +102,7 @@ func NewRPCLink(args RPCLinkArgs) (*rpcLink, error) {
 		uuidNewV4:               uuid.New,
 		responseMap:             map[string]chan rpcResponse{},
 		startResponseRouterOnce: &sync.Once{},
+		logEvent:                args.LogEvent,
 	}
 
 	sender, err := session.NewSender(
@@ -147,45 +140,57 @@ func NewRPCLink(args RPCLinkArgs) (*rpcLink, error) {
 
 	link.sender = sender
 	link.receiver = receiver
-	link.messageAccept = receiver.AcceptMessage
 
 	return link, nil
 }
 
+const responseRouterShutdownMessage = "Response router has shut down"
+
 // startResponseRouter is responsible for taking any messages received on the 'response'
 // link and forwarding it to the proper channel. The channel is being select'd by the
 // original `RPC` call.
 func (l *rpcLink) startResponseRouter() {
+	defer azlog.Writef(l.logEvent, responseRouterShutdownMessage)
+
 	for {
 		res, err := l.receiver.Receive(context.Background())
 
-		// You'll see this when the link is shutting down (either
-		// service-initiated via 'detach' or a user-initiated shutdown)
-		if isClosedError(err) {
-			l.broadcastError(err)
-			break
+		if err != nil {
+			// if the link or connection has a malfunction that would require it to restart then
+			// we need to bail out, broadcasting to all affected callers/consumers.
+			if GetRecoveryKind(err) != RecoveryKindNone {
+				azlog.Writef(l.logEvent, "Error in RPCLink, stopping response router: %s", err.Error())
+				l.broadcastError(err)
+				break
+			}
+
+			azlog.Writef(l.logEvent, "Non-fatal error in RPCLink, starting to receive again: %s", err.Error())
+			continue
 		}
 
 		// I don't believe this should happen. The JS version of this same code
 		// ignores errors as well since responses should always be correlated
 		// to actual send requests. So this is just here for completeness.
 		if res == nil {
+			azlog.Writef(l.logEvent, "RPCLink received no error, but also got no response")
 			continue
 		}
 
 		autogenMessageId, ok := res.Properties.CorrelationID.(string)
 
 		if !ok {
-			// TODO: it'd be good to track these in some way. We don't have a good way to
-			// forward this on at this point.
+			azlog.Writef(l.logEvent, "RPCLink message received without a CorrelationID %v", res)
 			continue
 		}
 
 		ch := l.deleteChannelFromMap(autogenMessageId)
 
-		if ch != nil {
-			ch <- rpcResponse{message: res, err: err}
+		if ch == nil {
+			azlog.Writef(l.logEvent, "RPCLink had no response channel for correlation ID %v", autogenMessageId)
+			continue
 		}
+
+		ch <- rpcResponse{message: res, err: err}
 	}
 }
 
@@ -228,7 +233,7 @@ func (l *rpcLink) RPC(ctx context.Context, msg *amqp.Message) (*RPCResponse, err
 
 	if err != nil {
 		l.deleteChannelFromMap(messageID)
-		return nil, fmt.Errorf("Failed to send message with ID %s: %w", messageID, err)
+		return nil, fmt.Errorf("failed to send message with ID %s: %w", messageID, err)
 	}
 
 	var res *amqp.Message
@@ -281,7 +286,7 @@ func (l *rpcLink) RPC(ctx context.Context, msg *amqp.Message) (*RPCResponse, err
 		Message:     res,
 	}
 
-	if err := l.messageAccept(ctx, res); err != nil {
+	if err := l.receiver.AcceptMessage(ctx, res); err != nil {
 		return response, fmt.Errorf("failed accepting message on rpc link: %w", err)
 	}
 
@@ -415,15 +420,6 @@ func addMessageID(message *amqp.Message, uuidNewV4 func() (uuid.UUID, error)) (*
 	return &copiedMessage, autoGenMessageID, nil
 }
 
-func isClosedError(err error) bool {
-	var detachError *amqp.DetachError
-
-	return errors.Is(err, amqp.ErrLinkClosed) ||
-		errors.As(err, &detachError) ||
-		errors.Is(err, amqp.ErrConnClosed) ||
-		errors.Is(err, amqp.ErrSessionClosed)
-}
-
 // asRPCError checks to see if the res is actually a failed request
 // (where failed means the status code was non-2xx). If so,
 // it returns true and updates the struct pointed to by err.