Allow a callback for Session[Initializing|Closing]Async (Azure#29945)

* Allow a callback for Session[Initializing|Closing]Async

Closes Azure#29917

* Remove remarks

* Defensive

* Approve API

* Change log

* Make things private and sneakishly hide a problem even the reviewers didn't spot

* Properly use the session

* Reduce to a single test

Author: danielmarbach

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added ability to register a callback for ` SessionInitializingAsync` and `SessionClosingAsync` to the `ServiceBusOptions`. _(A community contribution, courtesy of [danielmarbach](https://github.com/danielmarbach))_
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/api/Microsoft.Azure.WebJobs.Extensions.ServiceBus.netstandard2.0.cs

@@ -83,7 +83,9 @@ public ServiceBusOptions() { }
         public int MaxMessageBatchSize { get { throw null; } set { } }
         public int PrefetchCount { get { throw null; } set { } }
         public System.Func<Azure.Messaging.ServiceBus.ProcessErrorEventArgs, System.Threading.Tasks.Task> ProcessErrorAsync { get { throw null; } set { } }
+        public System.Func<Azure.Messaging.ServiceBus.ProcessSessionEventArgs, System.Threading.Tasks.Task> SessionClosingAsync { get { throw null; } set { } }
         public System.TimeSpan? SessionIdleTimeout { get { throw null; } set { } }
+        public System.Func<Azure.Messaging.ServiceBus.ProcessSessionEventArgs, System.Threading.Tasks.Task> SessionInitializingAsync { get { throw null; } set { } }
         public Azure.Messaging.ServiceBus.ServiceBusTransportType TransportType { get { throw null; } set { } }
         public System.Net.IWebProxy WebProxy { get { throw null; } set { } }
         string Microsoft.Azure.WebJobs.Hosting.IOptionsFormatter.Format() { throw null; }

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Config/ServiceBusExtensionConfigProvider.cs

@@ -81,6 +81,16 @@ public void Initialize(ExtensionConfigContext context)
                 LogExceptionReceivedEvent(e, _loggerFactory);
                 return Task.CompletedTask;
             };
+            Options.SessionInitializingAsync ??= (e) =>
+            {
+                LogSessionInitializingEvent(e, _loggerFactory);
+                return Task.CompletedTask;
+            };
+            Options.SessionClosingAsync ??= (e) =>
+            {
+                LogSessionClosingEvent(e, _loggerFactory);
+                return Task.CompletedTask;
+            };
 
             context
                 .AddConverter(new MessageToStringConverter())
@@ -132,5 +142,33 @@ private static LogLevel GetLogLevel(Exception ex)
                 return LogLevel.Information;
             }
         }
+
+        internal static void LogSessionInitializingEvent(ProcessSessionEventArgs e, ILoggerFactory loggerFactory)
+        {
+            try
+            {
+                var logger = loggerFactory?.CreateLogger<ServiceBusListener>();
+                string message = $"Session initializing (SessionId={e.SessionId}, SessionLockedUntil={e.SessionLockedUntil})";
+                logger?.LogInformation(0, message);
+            }
+            catch (Exception)
+            {
+                // best effort logging
+            }
+        }
+
+        internal static void LogSessionClosingEvent(ProcessSessionEventArgs e, ILoggerFactory loggerFactory)
+        {
+            try
+            {
+                var logger = loggerFactory?.CreateLogger<ServiceBusListener>();
+                string message = $"Session closing (SessionId={e.SessionId}, SessionLockedUntil={e.SessionLockedUntil})";
+                logger?.LogInformation(0, message);
+            }
+            catch (Exception)
+            {
+                // best effort logging
+            }
+        }
     }
 }

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Config/ServiceBusOptions.cs

@@ -130,6 +130,16 @@ public int MaxConcurrentSessions
         /// </summary>
         public Func<ProcessErrorEventArgs, Task> ProcessErrorAsync { get; set; }
 
+        /// <summary>
+        /// Optional handler that can be set to be notified when a new session is about to be processed.
+        /// </summary>
+        public Func<ProcessSessionEventArgs, Task> SessionInitializingAsync { get; set; }
+
+        /// <summary>
+        /// Optional handler that can be set to be notified when a session is about to be closed for processing.
+        /// </summary>
+        public Func<ProcessSessionEventArgs, Task> SessionClosingAsync { get; set; }
+
         /// <summary>
         /// Gets or sets the maximum number of messages that will be passed to each function call. This only applies for functions that receive
         /// a batch of messages. The default value is 1000.
@@ -200,9 +210,28 @@ string IOptionsFormatter.Format()
 
         internal async Task ExceptionReceivedHandler(ProcessErrorEventArgs args)
         {
-            if (ProcessErrorAsync != null)
+            var processErrorAsync = ProcessErrorAsync;
+            if (processErrorAsync != null)
+            {
+                await processErrorAsync(args).ConfigureAwait(false);
+            }
+        }
+
+        internal async Task SessionInitializingHandler(ProcessSessionEventArgs args)
+        {
+            var sessionInitializingAsync = SessionInitializingAsync;
+            if (sessionInitializingAsync != null)
+            {
+                await sessionInitializingAsync(args).ConfigureAwait(false);
+            }
+        }
+
+        internal async Task SessionClosingHandler(ProcessSessionEventArgs args)
+        {
+            var sessionClosingAsync = SessionClosingAsync;
+            if (sessionClosingAsync != null)
             {
-                await ProcessErrorAsync(args).ConfigureAwait(false);
+                await sessionClosingAsync(args).ConfigureAwait(false);
             }
         }
 

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Primitives/MessagingProvider.cs

@@ -208,6 +208,8 @@ protected internal virtual ServiceBusSessionProcessor CreateSessionProcessor(Ser
                 processor = client.CreateSessionProcessor(entityPath, options);
             }
             processor.ProcessErrorAsync += Options.ExceptionReceivedHandler;
+            processor.SessionInitializingAsync += Options.SessionInitializingHandler;
+            processor.SessionClosingAsync += Options.SessionClosingHandler;
             return processor;
         }
     }

sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/tests/ServiceBusSessionsEndToEndTests.cs

@@ -574,6 +574,30 @@ public async Task TestSingle_ReceiveFromFunction()
             }
         }
 
+        [Test]
+        public async Task TestSingle_CustomSessionHandlers()
+        {
+            await WriteQueueMessage("{'Name': 'Test1', 'Value': 'Value'}", "session1");
+            var host = BuildHost<TestCustomSessionHandlers>(SetCustomSessionHandlers);
+            using (host)
+            {
+                bool result1 = _waitHandle1.WaitOne(SBTimeoutMills);
+                bool result2 = _waitHandle2.WaitOne(SBTimeoutMills);
+
+                Assert.True(result1);
+                Assert.True(result2);
+                await host.StopAsync();
+            }
+        }
+
+        private static Action<IHostBuilder> SetCustomSessionHandlers =>
+            builder => builder.ConfigureWebJobs(b =>
+                b.AddServiceBus(sbOptions =>
+                {
+                    sbOptions.SessionInitializingAsync = TestCustomSessionHandlers.SessionInitializingHandler;
+                    sbOptions.SessionClosingAsync = TestCustomSessionHandlers.SessionClosingHandler;
+                }));
+
         [Test]
         public async Task TestBatch_ReceiveFromFunction()
         {
@@ -1106,6 +1130,30 @@ public static async Task RunAsync(
             }
         }
 
+        public class TestCustomSessionHandlers
+        {
+            public static async Task RunAsync(
+                [ServiceBusTrigger(FirstQueueNameKey, IsSessionsEnabled = true)]
+                ServiceBusReceivedMessage message,
+                ServiceBusSessionMessageActions sessionActions)
+            {
+                await sessionActions.CompleteMessageAsync(message);
+                sessionActions.ReleaseSession();
+            }
+
+            public static Task SessionInitializingHandler(ProcessSessionEventArgs arg)
+            {
+                _waitHandle1.Set();
+                return Task.CompletedTask;
+            }
+
+            public static Task SessionClosingHandler(ProcessSessionEventArgs arg)
+            {
+                _waitHandle2.Set();
+                return Task.CompletedTask;
+            }
+        }
+
         public class TestReceiveFromFunction_Batch
         {
             public static ServiceBusReceiveActions ReceiveActions { get; private set; }