Fixing an issue where a processor wouldn't grab partitions even when it didn't own any. This became especially bad in conditions where all the partitions were claimed.

- Adding stress test (not enabled, but useful for debugging/manual testing)
- Adding a lot of more unit tests around different ownership splits.

Author: Richard Park

sdk/messaging/azeventhubs/internal/eh/stress/stress-test-resources.json

@@ -52,17 +52,25 @@
       "apiVersion": "[variables('apiVersion')]",
       "name": "[variables('authorizationName')]",
       "location": "[variables('location')]",
-      "dependsOn": ["[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"],
+      "dependsOn": [
+        "[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"
+      ],
       "properties": {
-        "rights": ["Listen", "Manage", "Send"]
+        "rights": [
+          "Listen",
+          "Manage",
+          "Send"
+        ]
       }
     },
     {
       "type": "Microsoft.EventHub/namespaces/eventhubs",
       "apiVersion": "[variables('apiVersion')]",
       "name": "[variables('eventHubNameFull')]",
       "location": "[variables('location')]",
-      "dependsOn": ["[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"],
+      "dependsOn": [
+        "[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"
+      ],
       "properties": {
         "messageRetentionInDays": 7,
         "partitionCount": 32
@@ -73,7 +81,9 @@
       "apiVersion": "[variables('apiVersion')]",
       "name": "[concat(variables('namespaceName'), '/default')]",
       "location": "[variables('location')]",
-      "dependsOn": ["[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"],
+      "dependsOn": [
+        "[resourceId('Microsoft.EventHub/namespaces', variables('namespaceName'))]"
+      ],
       "properties": {
         "defaultAction": "Deny",
         "virtualNetworkRules": [],
@@ -127,13 +137,15 @@
           "name": "[concat('default/', variables('containerName'))]",
           "type": "blobServices/containers",
           "apiVersion": "[variables('storageApiVersion')]",
-          "dependsOn": ["[variables('storageAccountName')]"]
+          "dependsOn": [
+            "[variables('storageAccountName')]"
+          ]
         }
       ]
     },
   ],
   "outputs": {
-    "EVENTHUB_NAME": {
+    "EVENTHUB_NAME_STRESS": {
       "type": "string",
       "value": "[variables('eventHubName')]"
     },

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

@@ -17,6 +17,7 @@ func main() {
 		fn   func(ctx context.Context) error
 	}{
 		{name: "batch", fn: tests.BatchStressTester},
+		{name: "balance", fn: tests.BalanceTester},
 		{name: "processor", fn: tests.ProcessorStressTester},
 	}
 

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

@@ -0,0 +1,324 @@
+package tests
+
+import (
+	"context"
+	"errors"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	azlog "github.com/Azure/azure-sdk-for-go/sdk/internal/log"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/checkpoints"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/bloberror"
+)
+
+const (
+	EventBalanceTest azlog.Event = "balance.test"
+)
+
+// BalanceTester checks that we can properly distribute partitions and
+// maintain it over time.
+func BalanceTester(ctx context.Context) error {
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+
+	numProcessors := fs.Int("processors", 32, "The # of processor instances to run")
+	waitTime := fs.Duration("time", time.Minute, "The amount of time to wait before exiting the test")
+	strategy := fs.String("strategy", string(azeventhubs.ProcessorStrategyBalanced), "The partition acquisition strategy to use (balanced, greedy)")
+
+	if err := fs.Parse(os.Args[2:]); err != nil {
+		return err
+	}
+
+	testData, err := newStressTestData("balancetester", false, map[string]string{
+		"processors": fmt.Sprintf("%d", *numProcessors),
+		"waitTime":   waitTime.String(),
+		"strategy":   *strategy,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	containerName := testData.runID
+
+	wg := sync.WaitGroup{}
+	failuresChan := make(chan error, *numProcessors)
+
+	testCtx, cancelTest := context.WithTimeout(context.Background(), *waitTime)
+	defer cancelTest()
+
+	ownersMap := newOwners()
+
+	go func() {
+	Loop:
+		for {
+			select {
+			case <-ctx.Done():
+				break Loop
+			case <-time.After(30 * time.Second):
+				log.Printf("Current: %s", ownersMap)
+			}
+		}
+	}()
+
+	for i := 0; i < *numProcessors; i++ {
+		wg.Add(1)
+
+		go func(i int) {
+			defer wg.Done()
+
+			args := processArgs{
+				Name:          fmt.Sprintf("proc%02d", i),
+				Owners:        ownersMap,
+				CS:            testData.ConnectionString,
+				HubName:       testData.HubName,
+				StorageCS:     testData.StorageConnectionString,
+				ContainerName: containerName,
+				Strategy:      azeventhubs.ProcessorStrategy(*strategy),
+			}
+
+			if err := process(testCtx, args); err != nil {
+				failuresChan <- err
+				cancelTest()
+				return
+			}
+		}(i)
+	}
+
+	wg.Wait()
+	close(failuresChan)
+	cancelTest()
+
+	// any errors?
+	for err := range failuresChan {
+		testData.TC.TrackException(err)
+		fmt.Printf("ERROR: %s\n", err)
+		return err
+	}
+
+	// grab the # of partitions
+	err = func() error {
+		pc, err := azeventhubs.NewProducerClientFromConnectionString(testData.ConnectionString, testData.HubName, nil)
+
+		if err != nil {
+			return err
+		}
+
+		defer func() {
+			_ = pc.Close(ctx)
+		}()
+
+		ehProps, err := pc.GetEventHubProperties(ctx, nil)
+
+		if err != nil {
+			return err
+		}
+
+		// look through the dictionary now - how does everything look? Did all processors get at least
+		// one partition?
+		if len(ownersMap.byProcessor) != *numProcessors {
+			return errors.New("not all partition processors got a partition")
+		}
+
+		// now let's make sure everyone only took a fair share
+		min := len(ehProps.PartitionIDs) / *numProcessors
+
+		for _, owned := range ownersMap.byProcessor {
+			if owned < min {
+				return fmt.Errorf("a processor had %d partitions, but should have least %d", owned, min)
+			}
+		}
+
+		return nil
+	}()
+
+	if err != nil {
+		err := fmt.Errorf("unbalanced owners, map\n%s: %w", ownersMap, err)
+		testData.TC.TrackException(err)
+		return err
+	}
+
+	log.Printf("BALANCED")
+	return nil
+}
+
+type owners struct {
+	mu          sync.Mutex
+	byPartition map[string]ownerValue
+	byProcessor map[string]int
+}
+
+type ownerValue struct {
+	When time.Time
+	Who  string
+}
+
+func (v ownerValue) String() string {
+	return fmt.Sprintf("%s: %s", v.Who, time.Since(v.When)/time.Second)
+}
+
+func newOwners() *owners {
+	o := &owners{
+		byPartition: map[string]ownerValue{},
+		byProcessor: map[string]int{},
+	}
+
+	return o
+}
+
+func partKey(partID string) string {
+	return fmt.Sprintf("%2s", partID)
+}
+
+func (o *owners) Set(processorName string, partitionID string) {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+
+	partKey := partKey(partitionID)
+
+	o.byPartition[partKey] = ownerValue{
+		When: time.Now(),
+		Who:  processorName,
+	}
+
+	o.byProcessor[processorName]++
+	azlog.Writef(EventBalanceTest, "[gain] processor: %s, partition: %s, total: %d", processorName, partitionID, o.byProcessor[processorName])
+}
+
+func (o *owners) Remove(processorName string, partitionID string) {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+
+	partKey := partKey(partitionID)
+
+	if o.byPartition[partKey].Who == processorName {
+		delete(o.byPartition, partKey)
+	}
+
+	o.byProcessor[processorName]--
+	azlog.Writef(EventBalanceTest, "[loss] processor: %s, partition: %s, total: %d", processorName, partitionID, o.byProcessor[processorName])
+}
+
+func stringize[T any](m map[string]T, stringer func(v T) string) string {
+	var keys []string
+
+	for k := range m {
+		keys = append(keys, k)
+	}
+
+	sort.Strings(keys)
+
+	sb := strings.Builder{}
+
+	for _, k := range keys {
+		sb.WriteString(fmt.Sprintf("  %s: %s\n", k, stringer(m[k])))
+	}
+
+	return sb.String()
+}
+
+func (o *owners) String() string {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+
+	return fmt.Sprintf("By partition:\n%s\nBy processor:\n%s\n",
+		stringize(o.byPartition, func(v ownerValue) string { return v.String() }),
+		stringize(o.byProcessor, func(v int) string { return fmt.Sprintf("%d", v) }),
+	)
+}
+
+type processArgs struct {
+	Name   string
+	Owners *owners
+
+	CS      string
+	HubName string
+
+	StorageCS     string
+	ContainerName string
+
+	Strategy azeventhubs.ProcessorStrategy
+}
+
+func process(ctx context.Context, args processArgs) error {
+	client, err := azeventhubs.NewConsumerClientFromConnectionString(args.CS, args.HubName, azeventhubs.DefaultConsumerGroup, nil)
+
+	if err != nil {
+		return err
+	}
+
+	defer func() { _ = client.Close(ctx) }()
+
+	blobClient, err := azblob.NewClientFromConnectionString(args.StorageCS, nil)
+
+	if err != nil {
+		return err
+	}
+
+	containerClient := blobClient.ServiceClient().NewContainerClient(args.ContainerName)
+
+	if _, err := containerClient.Create(ctx, nil); err != nil {
+		if !bloberror.HasCode(err, bloberror.ContainerAlreadyExists) {
+			return err
+		}
+	}
+
+	defer func() { _, _ = containerClient.Delete(ctx, nil) }()
+
+	blobStore, err := checkpoints.NewBlobStore(containerClient, nil)
+
+	if err != nil {
+		return err
+	}
+
+	processor, err := azeventhubs.NewProcessor(client, blobStore, &azeventhubs.ProcessorOptions{
+		LoadBalancingStrategy: args.Strategy,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	go func() {
+		for {
+			pc := processor.NextPartitionClient(ctx)
+
+			if pc == nil {
+				break
+			}
+
+			args.Owners.Set(args.Name, pc.PartitionID())
+
+			go keepAlive(ctx, pc, args)
+		}
+	}()
+
+	return processor.Run(ctx)
+}
+
+func keepAlive(ctx context.Context, pc *azeventhubs.ProcessorPartitionClient, args processArgs) {
+ReceiveLoop:
+	for {
+		_, err := pc.ReceiveEvents(ctx, 1, nil)
+
+		var eventHubErr *azeventhubs.Error
+
+		switch {
+		case errors.Is(err, context.Canceled):
+			// test is over, we'll see how everything shook out.
+			break ReceiveLoop
+		case errors.As(err, &eventHubErr) && eventHubErr.Code == azeventhubs.ErrorCodeOwnershipLost:
+			// we've swapped ownership with _someone_. Record it in the map, but be careful of potential ordering
+			// and only delete if we're the one in the map!
+			args.Owners.Remove(args.Name, pc.PartitionID())
+			break ReceiveLoop
+		}
+	}
+}

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

@@ -80,7 +80,7 @@ func newStressTestData(name string, verbose bool, baggage map[string]string) (*s
 
 	variables := map[string]*string{
 		"EVENTHUB_CONNECTION_STRING":                &td.ConnectionString,
-		"EVENTHUB_NAME":                             &td.HubName,
+		"EVENTHUB_NAME_STRESS":                      &td.HubName,
 		"CHECKPOINTSTORE_STORAGE_CONNECTION_STRING": &td.StorageConnectionString,
 	}