[azeventhubs] Fix credit requesting errors (Azure#20015)

The current go-amqp can have issues if your active credits are greater than the session window size. We have a blocking-release bug here: Azure/go-amqp#240. I've also added validation to prevent asking for <= 0 credits as well, which isn't supported.

I've also updated the stress tests to build with Go1.20 - it doesn't change the requirements for customer code.

Author: richardpark-msft

sdk/messaging/azeventhubs/internal/eh/stress/Dockerfile

@@ -1,4 +1,4 @@
-FROM mcr.microsoft.com/oss/go/microsoft/golang:1.18 as build
+FROM mcr.microsoft.com/oss/go/microsoft/golang:1.20 as build
 # you'll need to run this build from the root of the azeventhubs module
 ENV GOOS=linux 
 ENV GOARCH=amd64 

sdk/messaging/azeventhubs/internal/eh/stress/stress.go

@@ -5,10 +5,8 @@ package main
 import (
 	"context"
 	"fmt"
-	"math/rand"
 	"os"
 	"sort"
-	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/eh/stress/tests"
 )
@@ -42,8 +40,6 @@ func main() {
 
 	for _, test := range tests {
 		if test.name == testName {
-			rand.Seed(time.Now().UnixNano())
-
 			if err := test.fn(context.Background()); err != nil {
 				fmt.Printf("ERROR: %s\n", err)
 				os.Exit(1)

sdk/messaging/azeventhubs/partition_client.go

@@ -21,7 +21,11 @@ import (
 const DefaultConsumerGroup = "$Default"
 
 const defaultPrefetchSize = uint32(300)
-const defaultMaxCreditSize = uint32(2048)
+
+// defaultLinkRxBuffer is the maximum number of transfer frames we can handle
+// on the Receiver. This matches the current default window size that go-amqp
+// uses for sessions.
+const defaultMaxCreditSize = uint32(5000)
 
 // StartPosition indicates the position to start receiving events within a partition.
 // The default position is Latest.
@@ -100,6 +104,14 @@ func (pc *PartitionClient) ReceiveEvents(ctx context.Context, count int, options
 
 	prefetchDisabled := pc.prefetch < 0
 
+	if count <= 0 {
+		return nil, internal.NewErrNonRetriable("count should be greater than 0")
+	}
+
+	if prefetchDisabled && count > int(defaultMaxCreditSize) {
+		return nil, internal.NewErrNonRetriable(fmt.Sprintf("count cannot exceed %d", defaultMaxCreditSize))
+	}
+
 	err := pc.links.Retry(ctx, EventConsumer, "ReceiveEvents", pc.partitionID, pc.retryOptions, func(ctx context.Context, lwid internal.LinkWithID[amqpwrap.AMQPReceiverCloser]) error {
 		events = nil
 
@@ -259,6 +271,11 @@ func newPartitionClient(args partitionClientArgs, options *PartitionClientOption
 		return nil, err
 	}
 
+	if options.Prefetch > int32(defaultMaxCreditSize) {
+		// don't allow them to set the prefetch above the session window size.
+		return nil, internal.NewErrNonRetriable(fmt.Sprintf("options.Prefetch cannot exceed %d", defaultMaxCreditSize))
+	}
+
 	client := &PartitionClient{
 		eventHub:         args.eventHub,
 		partitionID:      args.partitionID,

sdk/messaging/azeventhubs/partition_client_unit_test.go

@@ -5,11 +5,13 @@ package azeventhubs
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/go-amqp"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/test"
 	"github.com/stretchr/testify/require"
 )
 
@@ -41,6 +43,43 @@ func TestUnit_PartitionClient_PrefetchOff(t *testing.T) {
 	require.True(t, ns.Receiver.ManualCreditsSetFromOptions)
 }
 
+func TestUnit_PartitionClient_PrefetchOff_CreditLimits(t *testing.T) {
+	ns := &internal.FakeNSForPartClient{
+		Receiver: &internal.FakeAMQPReceiver{
+			Messages: fakeMessages(int(defaultMaxCreditSize)),
+		},
+	}
+
+	client, err := newPartitionClient(partitionClientArgs{
+		namespace: ns,
+	}, &PartitionClientOptions{
+		Prefetch: -1,
+	})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	// can't request over the max
+	events, err := client.ReceiveEvents(ctx, int(defaultMaxCreditSize+1), nil)
+	require.EqualError(t, err, "count cannot exceed 5000")
+	require.Empty(t, events)
+
+	// can't request negative/0 credits
+	events, err = client.ReceiveEvents(ctx, 0, nil)
+	require.EqualError(t, err, "count should be greater than 0")
+	require.Empty(t, events)
+
+	events, err = client.ReceiveEvents(ctx, -1, nil)
+	require.EqualError(t, err, "count should be greater than 0")
+	require.Empty(t, events)
+
+	// can request the max
+	events, err = client.ReceiveEvents(ctx, int(defaultMaxCreditSize), nil)
+	require.NoError(t, err)
+	require.NotEmpty(t, events)
+}
+
 func TestUnit_PartitionClient_PrefetchOffOnlyBackfillsCredits(t *testing.T) {
 	testData := []struct {
 		Name    string
@@ -116,3 +155,77 @@ func TestUnit_PartitionClient_PrefetchOn(t *testing.T) {
 		require.Equal(t, uint32(td.initialCredits-3), ns.Receiver.ActiveCredits, "All messages should have been received")
 	}
 }
+
+func TestUnit_PartitionClient_PrefetchLimit(t *testing.T) {
+	newPartitionClient := func(prefetch int32) (*PartitionClient, error) {
+		ns := &internal.FakeNSForPartClient{
+			Receiver: &internal.FakeAMQPReceiver{
+				Messages: fakeMessages(int(defaultMaxCreditSize) + 1),
+			},
+		}
+
+		client, err := newPartitionClient(partitionClientArgs{namespace: ns}, &PartitionClientOptions{
+			Prefetch: prefetch,
+		})
+
+		return client, err
+	}
+
+	t.Run("max allowed credits is defaultMaxCreditSize", func(t *testing.T) {
+		client, err := newPartitionClient(int32(defaultMaxCreditSize))
+		require.NoError(t, err)
+		require.NotNil(t, client)
+
+		test.RequireClose(t, client)
+	})
+
+	t.Run("can't request zero or negative credits", func(t *testing.T) {
+		client, err := newPartitionClient(int32(defaultMaxCreditSize))
+		require.NoError(t, err)
+		require.NotNil(t, client)
+
+		events, err := client.ReceiveEvents(context.Background(), 0, nil)
+		require.EqualError(t, err, "count should be greater than 0")
+		require.Empty(t, events)
+
+		events, err = client.ReceiveEvents(context.Background(), -1, nil)
+		require.EqualError(t, err, "count should be greater than 0")
+		require.Empty(t, events)
+
+		test.RequireClose(t, client)
+	})
+
+	t.Run("can receive more than defaultMaxCreditSize in prefetch mode", func(t *testing.T) {
+		client, err := newPartitionClient(0)
+		require.NoError(t, err)
+		require.NotNil(t, client)
+
+		// if you're using prefetch it's fine to ask for more than the `defaultMaxCreditSize`
+		// since it doesn't actually cause to request more credits than is allowed.
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		events, err := client.ReceiveEvents(ctx, int(defaultMaxCreditSize+1), nil)
+		require.NoError(t, err)
+		require.NotEmpty(t, events)
+
+		test.RequireClose(t, client)
+	})
+
+	t.Run("can't set a option.Prefetch > defaultMaxCreditSize", func(t *testing.T) {
+		// and you can't create a PartitionClient that uses more credits than allowed.
+		client, err := newPartitionClient(int32(defaultMaxCreditSize) + 1)
+		require.EqualError(t, err, fmt.Sprintf("options.Prefetch cannot exceed %d", defaultMaxCreditSize))
+		require.Nil(t, client)
+	})
+}
+
+func fakeMessages(count int) []*amqp.Message {
+	var messages []*amqp.Message
+
+	for i := 0; i < count; i++ {
+		messages = append(messages, &amqp.Message{})
+	}
+
+	return messages
+}