azservicebus.Receiver.ReceiveMessages ignoring context timeout

[adeturner]

Bug Report

github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus v0.4.1
go version go1.18.1 linux/amd64

• What happened?

From what I can see, azservicebus.Receiver.ReceiveMessages hangs when no messages are present in the queue. The call does not cancel according to context timeout setting (either customer set or internally). This makes it tricky to write KEDA driven serverless customer code for e.g. container apps

• What did you expect or want to happen?

Preferred that ReceiveMessages times out based on the customer controlled setting, or the internal values if required (20ms/1second)

• How can we reproduce it?

Please request if you would like me to provide some code to run

• Anything we should know about your environment.

n/a

• Investigation

ctx, _ := context.WithTimeout(context.TODO(), time.Second)
var receiver *azservicebus.Receiver
messages, err = receiver.ReceiveMessages(ctx, limit, nil)

fetchMessages doesnt yield to context timeout

defaultTimeAfterFirstMessage is set to 20ms or 1 second.

I propose a change to fetchMessages something like the below

Note: due to the interfacing its unclear to me whether the same change is required in internal.AMQPReceiver, but a quick look at Receive makes it looks like it handles things ok there

func fetchMessages(ctx context.Context, receiver internal.AMQPReceiver, maxMessages int, defaultTimeAfterFirstMessage time.Duration, messages *[]*ReceivedMessage) error {
	if err := receiver.IssueCredit(uint32(maxMessages)); err != nil {
		return err
	}

	var cancel context.CancelFunc

	for {

// BEGIN CHANGE
		select {
		case <-ctx.Done():
			// context timeout
			return nil
		default:
// END CHANGE
			amqpMessage, err := receiver.Receive(ctx)

			if err != nil {
				return err
			}

			*messages = append(*messages, newReceivedMessage(amqpMessage))

			if len(*messages) == maxMessages {
				return nil
			}

//BEGIN CHANGE
			// if cancel == nil {
			if _, deadlineSet := ctx.Deadline(); !deadlineSet {
				// replace the context that we're using for everything with a new one that will cancel
				// after a period of time.
				 ctx, cancel = context.WithTimeout(ctx, defaultTimeAfterFirstMessage)
				defer cancel()
			}
// END CHANGE
		}
	}
}

[adeturner]

Apologies, I'd convinced myself that there was an issue, but I was overwriting my ctx. #facepalm