[azservicebus] Adding in an example of how you could write your own lock renewer (Azure#23168)

Just a simple example for people that want to write their own renewer.

Also, for convenience, adding in a .ps1 file to generate the .env file if you already have a resource set up.

Fixes Azure#20857

Author: richardpark-msft

sdk/messaging/azservicebus/env.ps1

@@ -0,0 +1,46 @@
+#!/usr/bin/env pwsh
+  
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+# This script will bootstrap your .env file, given a resource group that's been deployed to.
+# We expect a premium and a standard Service Bus namespace for our tests.
+#
+# You might need to install the Azure Powershell module first and login:
+#
+# Install-Module -Name Az -Repository PSGallery -Force
+# Connect-AzAccount -UseDeviceAuthentication
+# 
+
+# TODO:
+$rg = "<your resource group should go here>"
+
+$contents = ""
+
+Get-AzServiceBusNamespace -ResourceGroup $rg
+| ForEach-Object {
+    $cs = (Get-AzServiceBusKey -ResourceGroup $rg -NamespaceName $_.Name -AuthorizationRuleName RootManageSharedAccessKey).PrimaryConnectionString
+    $endpoint = $_.ServiceBusEndpoint.Replace("https://", "").Replace(":443/", "")
+
+    if ($_.SkuTier -eq "Standard") {
+        # this is in the .bicep file for this - we create a few keys with different permissions for testing.
+        $noManageCs = (Get-AzServiceBusKey -ResourceGroup $rg -NamespaceName $_.Name -AuthorizationRuleName NoManage).PrimaryConnectionString
+        $sendOnlyCs = (Get-AzServiceBusKey -ResourceGroup $rg -NamespaceName $_.Name -AuthorizationRuleName SendOnly).PrimaryConnectionString
+        $listenOnlyCs = (Get-AzServiceBusKey -ResourceGroup $rg -NamespaceName $_.Name -AuthorizationRuleName ListenOnly).PrimaryConnectionString
+
+        $contents += "`n# standard`n"
+        $contents += "SERVICEBUS_CONNECTION_STRING=$cs`n"
+        $contents += "SERVICEBUS_ENDPOINT=$endpoint`n"
+        $contents += "SERVICEBUS_CONNECTION_STRING_NO_MANAGE=$noManageCs`n"
+        $contents += "SERVICEBUS_CONNECTION_STRING_SEND_ONLY=$sendOnlyCs`n"
+        $contents += "SERVICEBUS_CONNECTION_STRING_LISTEN_ONLY=$listenOnlyCs`n"
+    }
+    else {
+        # we do a little bit of testing on premium.
+        $contents += "`n# premium`n"
+        $contents += "SERVICEBUS_CONNECTION_STRING_PREMIUM=$cs`n"
+        $contents += "SERVICEBUS_ENDPOINT_PREMIUM=$endpoint`n"
+    }    
+} 
+
+Out-File ".env" -InputObject $contents

sdk/messaging/azservicebus/example_client_autoRenewLocks_test.go

@@ -0,0 +1,83 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azservicebus_test
+
+import (
+	"context"
+	"errors"
+	"log"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+)
+
+// autoRenewMessageLock runs a loop calling [azservicebus.Receiver.RenewLock] until the message has been completed, the lock
+// has expired or the passed in context has been cancelled.
+//
+// This function returns nil if the message's lock has expired.
+//
+// - receiver - the receiver that received 'msg'
+// - msg - the received message
+// - configuredLockDuration - the lock duration, configured at the queue or subscription level.
+func autoRenewMessageLock(ctx context.Context, receiver *azservicebus.Receiver, msg *azservicebus.ReceivedMessage, configuredLockDuration time.Duration) error {
+	for {
+		log.Printf("Renewing lock for %s", msg.MessageID)
+		if err := receiver.RenewMessageLock(ctx, msg, nil); err != nil {
+			// You get this error if your lock has expired _or_ if you've settled the message.
+			// We can ignore it here because the error will also be returned by the settlement
+			// method (CompleteMessage, AbandonMessage, etc..)
+			if sbErr := (*azservicebus.Error)(nil); errors.As(err, &sbErr) && sbErr.Code == azservicebus.CodeLockLost {
+				log.Printf("Lock was lost for message with ID %s", msg.MessageID)
+				return nil
+			}
+
+			return err
+		}
+		log.Printf("Renewed lock for %s", msg.MessageID)
+
+		select {
+		case <-ctx.Done():
+			log.Printf("Cancelling lock renewal for %s", msg.MessageID)
+			return ctx.Err()
+		case <-time.After(configuredLockDuration / 2):
+		}
+	}
+}
+
+func Example_autoRenewLocks() {
+	// This is configurable, on the service. Change this value to match since it's used
+	// to figure out how often to renew a lock.
+	const queueOrSubscriptionLockDuration = 30 * time.Second
+
+	messages, err := receiver.ReceiveMessages(context.TODO(), 1, nil)
+
+	if err != nil {
+		//  TODO: Update the following line with your application specific error handling logic
+		log.Printf("ERROR: %s", err)
+		return
+	}
+
+	for _, msg := range messages {
+		go func(msg *azservicebus.ReceivedMessage) {
+			//
+			// Each goroutine will _eventually_ stop once its message has been settled.
+			//
+			// If you want to be more proactive you can create a separate context for each message
+			// and cancel that when you settle.
+			//
+			err := autoRenewMessageLock(context.TODO(), receiver, msg, queueOrSubscriptionLockDuration)
+
+			if err != nil {
+				//  TODO: Update the following line with your application specific error handling logic
+				log.Printf("ERROR: %s", err)
+				return
+			}
+		}(msg)
+	}
+
+	// TODO: Process the messages.
+	//
+	// Our lock renewal goroutines will stop once RenewLock returns that the lock has been lost. This happens
+	// if the lock has expired or if we've settled the message.
+}

sdk/messaging/azservicebus/receiver2_test.go

@@ -0,0 +1,136 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azservicebus_test
+
+import (
+	"context"
+	"log"
+	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/atom"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/test"
+	"github.com/stretchr/testify/require"
+)
+
+// NOTE: tests should migrate into here since this is the actual _test package.
+
+func TestRenewLocks(t *testing.T) {
+	client := test.NewClient(t, test.NewClientArgs[azservicebus.ClientOptions, azservicebus.Client]{
+		NewClientFromConnectionString: azservicebus.NewClientFromConnectionString,
+		NewClient:                     azservicebus.NewClient,
+	}, nil)
+
+	queueName, cleanupQueue := test.CreateExpiringQueue(t, &atom.QueueDescription{
+		LockDuration: to.Ptr("PT5S"),
+	})
+
+	defer cleanupQueue()
+
+	sender, err := client.NewSender(queueName, nil)
+	require.NoError(t, err)
+
+	defer test.RequireClose(t, sender)
+
+	receiver, err := client.NewReceiverForQueue(queueName, nil)
+	require.NoError(t, err)
+
+	defer test.RequireClose(t, receiver)
+
+	err = sender.SendMessage(context.Background(), &azservicebus.Message{
+		Body: []byte("hello world"),
+	}, nil)
+	require.NoError(t, err)
+
+	messages, err := receiver.ReceiveMessages(context.Background(), 1, nil)
+	require.NoError(t, err)
+
+	errCh := make(chan error)
+
+	go func() {
+		err = autoRenewMessageLock(context.Background(), receiver, messages[0], 5*time.Second)
+		errCh <- err
+	}()
+
+	time.Sleep(10 * time.Second) // make sure we renew a few times.
+
+	log.Printf("Completing message")
+	err = receiver.CompleteMessage(context.Background(), messages[0], nil)
+	require.NoError(t, err)
+	log.Printf("Completed message")
+
+	select {
+	case err := <-errCh:
+		require.NoError(t, err)
+	case <-time.After(5 * time.Second):
+		require.Fail(t, "goroutine took longer than 5 seconds to complete")
+	}
+}
+
+func TestRenewLocksEarlyCancel(t *testing.T) {
+	client := test.NewClient(t, test.NewClientArgs[azservicebus.ClientOptions, azservicebus.Client]{
+		NewClientFromConnectionString: azservicebus.NewClientFromConnectionString,
+		NewClient:                     azservicebus.NewClient,
+	}, nil)
+
+	queueName, cleanupQueue := test.CreateExpiringQueue(t, &atom.QueueDescription{
+		// this time we're bumping up the time. We want to make sure we're cancelling.
+		LockDuration: to.Ptr("PT1M"),
+	})
+
+	defer cleanupQueue()
+
+	sender, err := client.NewSender(queueName, nil)
+	require.NoError(t, err)
+
+	defer test.RequireClose(t, sender)
+
+	receiver, err := client.NewReceiverForQueue(queueName, nil)
+	require.NoError(t, err)
+
+	defer test.RequireClose(t, receiver)
+
+	err = sender.SendMessage(context.Background(), &azservicebus.Message{
+		Body: []byte("hello world"),
+	}, nil)
+	require.NoError(t, err)
+
+	messages, err := receiver.ReceiveMessages(context.Background(), 1, nil)
+	require.NoError(t, err)
+
+	t.Run("ContextAlreadyCancelled", func(t *testing.T) {
+		// pre-cancelled.
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		log.Printf("Auto renewing lock, but the context is already cancelled")
+		err = autoRenewMessageLock(ctx, receiver, messages[0], time.Minute)
+		log.Printf("Auto renewer returned %s", err)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
+	t.Run("ContextCancelledAfterTenSeconds", func(t *testing.T) {
+		// now cancel _after_ the first renewal has taken place (by sleeping a bit...)
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+
+		errCh := make(chan error)
+
+		log.Printf("Auto renewing lock, context will cancel after 10 seconds")
+		go func() {
+			err = autoRenewMessageLock(ctx, receiver, messages[0], time.Minute)
+			errCh <- err
+		}()
+
+		select {
+		case err := <-errCh:
+			log.Printf("Auto renewer returned %s", err)
+			require.ErrorIs(t, err, context.DeadlineExceeded)
+		case <-time.After(15 * time.Second):
+			require.Fail(t, "goroutine took longer than 15 seconds to cancel")
+		}
+	})
+}