fix: shutdown

src/Directory.Build.props

@@ -1,8 +1,8 @@
 <Project>
 
     <PropertyGroup>
-        <Version>1.0.0-alpha.24</Version>
-        <PackageVersion>1.0.0-alpha.24</PackageVersion>
+        <Version>1.0.0-alpha.25</Version>
+        <PackageVersion>1.0.0-alpha.25</PackageVersion>
         <Authors>Zapto</Authors>
         <RepositoryUrl>https://github.com/zapto-dev/Mediator</RepositoryUrl>
         <Copyright>Copyright © 2025 Zapto</Copyright>

src/Mediator.DependencyInjection/DefaultBackgroundPublisher.cs

@@ -18,49 +18,62 @@ public DefaultBackgroundPublisher(IServiceScopeFactory scopeFactory)
     /// <inheritdoc />
     public void Publish(object notification)
     {
-        _ = Task.Run(async () =>
+        ObserveTask(Task.Run(async () =>
         {
             using var scope = _scopeFactory.CreateScope();
             var mediator = scope.ServiceProvider.GetRequiredService<IPublisher>();
 
             await mediator.Publish(notification, CancellationToken.None);
-        }, CancellationToken.None);
+        }, CancellationToken.None));
     }
 
     /// <inheritdoc />
     public void Publish(MediatorNamespace ns, object notification)
     {
-        _ = Task.Run(async () =>
+        ObserveTask(Task.Run(async () =>
         {
             using var scope = _scopeFactory.CreateScope();
             var mediator = scope.ServiceProvider.GetRequiredService<IPublisher>();
 
             await mediator.Publish(ns, notification, CancellationToken.None);
-        }, CancellationToken.None);
+        }, CancellationToken.None));
     }
 
     /// <inheritdoc />
     public void Publish<TNotification>(TNotification notification) where TNotification : INotification
     {
-        _ = Task.Run(async () =>
+        ObserveTask(Task.Run(async () =>
         {
             using var scope = _scopeFactory.CreateScope();
             var mediator = scope.ServiceProvider.GetRequiredService<IPublisher>();
 
             await mediator.Publish(notification, CancellationToken.None);
-        }, CancellationToken.None);
+        }, CancellationToken.None));
     }
 
     /// <inheritdoc />
     public void Publish<TNotification>(MediatorNamespace ns, TNotification notification)
         where TNotification : INotification
     {
-        _ = Task.Run(async () =>
+        ObserveTask(Task.Run(async () =>
         {
             using var scope = _scopeFactory.CreateScope();
             var mediator = scope.ServiceProvider.GetRequiredService<IPublisher>();
 
             await mediator.Publish(ns, notification, CancellationToken.None);
-        }, CancellationToken.None);
+        }, CancellationToken.None));
+    }
+
+    /// <summary>
+    /// Observes faulted fire-and-forget tasks so they do not trigger
+    /// <see cref="TaskScheduler.UnobservedTaskException"/> and crash the process.
+    /// </summary>
+    private static void ObserveTask(Task task)
+    {
+        task.ContinueWith(
+            static t => GC.KeepAlive(t.Exception),
+            CancellationToken.None,
+            TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously,
+            TaskScheduler.Default);
     }
 }

src/Mediator.DependencyInjection/Generic/Handlers/GenericNotificationHandler.cs

@@ -24,6 +24,13 @@ internal interface IHandlerRegistration
 {
     INotificationCache Owner { get; }
 
+    /// <summary>
+    /// Set to <c>true</c> under <see cref="INotificationCache.Lock"/> when the
+    /// handler is disposed. Checked (without lock) as a best-effort guard in
+    /// <see cref="InvokeAsync"/> to skip already-disposed handlers.
+    /// </summary>
+    bool IsDisposed { get; set; }
+
     ValueTask InvokeAsync(IServiceProvider provider, object notification, CancellationToken cancellationToken);
 }
 
@@ -53,7 +60,8 @@ public GenericNotificationCache(IEnumerable<GenericNotificationRegistration> reg
         }
     }
 
-    public List<Type>? HandlerTypes { get; set; }
+    private List<Type>? _handlerTypes;
+    public List<Type>? HandlerTypes { get => Volatile.Read(ref _handlerTypes); set => Volatile.Write(ref _handlerTypes, value); }
 
     public List<GenericNotificationRegistration> MatchingRegistrations { get; }
 

src/Mediator.DependencyInjection/Generic/Handlers/GenericRequestHandler.cs

@@ -43,7 +43,8 @@ public GenericRequestCache(IEnumerable<GenericRequestRegistration> registrations
         }
     }
 
-    public Type? RequestHandlerType { get; set; }
+    private Type? _requestHandlerType;
+    public Type? RequestHandlerType { get => Volatile.Read(ref _requestHandlerType); set => Volatile.Write(ref _requestHandlerType, value); }
 
     public List<GenericRequestRegistration> MatchingRegistrations { get; }
 }
@@ -67,7 +68,8 @@ public GenericRequestCache(IEnumerable<GenericRequestRegistration> registrations
         }
     }
 
-    public Type? RequestHandlerType { get; set; }
+    private Type? _requestHandlerType;
+    public Type? RequestHandlerType { get => Volatile.Read(ref _requestHandlerType); set => Volatile.Write(ref _requestHandlerType, value); }
 
     public List<GenericRequestRegistration> MatchingRegistrations { get; }
 }

src/Mediator.DependencyInjection/Generic/Handlers/GenericStreamRequestHandler.cs

@@ -41,7 +41,8 @@ public GenericStreamRequestCache(IEnumerable<GenericStreamRequestRegistration> r
         }
     }
 
-    public Type? RequestHandlerType { get; set; }
+    private Type? _requestHandlerType;
+    public Type? RequestHandlerType { get => Volatile.Read(ref _requestHandlerType); set => Volatile.Write(ref _requestHandlerType, value); }
 
     public List<GenericStreamRequestRegistration> MatchingRegistrations { get; }
 }

src/Mediator.DependencyInjection/Generic/NotificationAttributeHandler.cs

@@ -19,7 +19,13 @@ static NotificationAttributeHandler()
 
             if (attribute == null) continue;
 
-            var notificationType = method.GetParameters()[0].ParameterType;
+            var parameters = method.GetParameters();
+            if (parameters.Length == 0)
+            {
+                throw new InvalidOperationException($"Method '{method.Name}' on '{typeof(T).FullName}' is decorated with [NotificationHandler] but has no parameters. The first parameter must be the notification type.");
+            }
+
+            var notificationType = parameters[0].ParameterType;
 
             if (!typeof(INotification).IsAssignableFrom(notificationType))
             {
@@ -45,8 +51,14 @@ public static IDisposable RegisterHandlers(IServiceProvider serviceProvider, obj
             var registration = new HandlerRegistration(cache, invoker, handler, middleware);
 
             cache.Lock.Wait();
-            cache.Registrations.Add(registration);
-            cache.Lock.Release();
+            try
+            {
+                cache.Registrations.Add(registration);
+            }
+            finally
+            {
+                cache.Lock.Release();
+            }
 
             registrations.Add(registration);
         }
@@ -70,8 +82,15 @@ public void Dispose()
                 var owner = registration.Owner;
 
                 owner.Lock.Wait();
-                owner.Registrations.Remove(registration);
-                owner.Lock.Release();
+                try
+                {
+                    registration.IsDisposed = true;
+                    owner.Registrations.Remove(registration);
+                }
+                finally
+                {
+                    owner.Lock.Release();
+                }
             }
         }
     }
@@ -97,7 +116,14 @@ public HandlerRegistration(
 
         public Func<T, IServiceProvider, object, CancellationToken, ValueTask> Invoke { get; }
 
-        public bool IsDisposed { get; set; }
+        private bool _isDisposed;
+
+        public bool IsDisposed
+        {
+            get => Volatile.Read(ref _isDisposed);
+            // Written under Owner.Lock, but Volatile.Write ensures visibility for lock-free reads in InvokeAsync
+            set => Volatile.Write(ref _isDisposed, value);
+        }
 
         public ValueTask InvokeAsync(IServiceProvider provider, object notification, CancellationToken cancellationToken)
         {

src/Mediator.DependencyInjection/ServiceProviderMediator.cs

@@ -79,7 +79,7 @@
             {
                 invokeAsync = cb =>
                 {
-                    _ = Task.Run(() => invoker(cb));
+                    _ = invoker(cb);
                     return Task.CompletedTask;
                 };
             }
@@ -89,7 +89,7 @@
             }
         }
 
         return (IDisposable) typeof(NotificationAttributeHandler<>).MakeGenericType(handler.GetType())
             .GetMethod(nameof(NotificationAttributeHandler<object>.RegisterHandlers), BindingFlags.Static | BindingFlags.Public)!
             .Invoke(null, new[] { _provider, handler, invokeAsync });
     }

src/Mediator.Hosting/MediatorExtensions.cs

@@ -1,5 +1,6 @@
 using System;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
 using Zapto.Mediator.Options;
 using Zapto.Mediator.Services;
 
@@ -16,7 +17,7 @@ public static IMediatorBuilder AddHostingBackgroundScheduler(this IMediatorBuild
 
         mediatorBuilder.Services.AddSingleton<BackgroundQueueService>();
         mediatorBuilder.Services.AddHostedService<BackgroundQueueHostedService>();
-        mediatorBuilder.Services.AddSingleton<IBackgroundPublisher, BackgroundPublisher>();
+        mediatorBuilder.Services.Replace(ServiceDescriptor.Singleton<IBackgroundPublisher, BackgroundPublisher>());
 
         if (configure is not null)
         {

src/Mediator.Hosting/Services/BackgroundQueueService.cs

@@ -2,10 +2,8 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
-using System.Security.Cryptography;
 using System.Threading;
 using System.Threading.Tasks;
-using MediatR;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -36,18 +34,22 @@ public void QueueBackgroundWorkItem(Func<Task> workItem, object notification)
             throw new OperationCanceledException("Cannot schedule work item since the application is stopping");
         }
 
-        _workItems.Enqueue(workItem);
-
         lock (_workers)
         {
+            _workItems.Enqueue(workItem);
+
             if (_workers.Count < _options.Value.MaxDegreeOfParallelism)
             {
                 var worker = new Worker
                 {
                     Notification = notification
                 };
 
-                worker.Task = Task.Factory.StartNew(() => ProcessBackgroundWorkItem(worker), TaskCreationOptions.LongRunning);
+                worker.Task = Task.Factory.StartNew(
+                    () => ProcessBackgroundWorkItem(worker),
+                    CancellationToken.None,
+                    TaskCreationOptions.None,
+                    TaskScheduler.Default).Unwrap();
                 _workers.Add(worker);
             }
         }
@@ -93,9 +95,13 @@ public async Task WaitForBackgroundTasksAsync(CancellationToken cancellationToke
                 resultingTask = await Task.WhenAny(task, tcs.Task);
 #endif
             }
+            catch (OperationCanceledException)
+            {
+                throw;
+            }
             catch
             {
-                // ignore
+                // worker exceptions are already logged at the source; ignore here
             }
 
 #if !NET
@@ -127,6 +133,7 @@ private async Task ProcessBackgroundWorkItem(Worker worker)
             // Process next work item
             if (!_workItems.TryDequeue(out var workItem))
             {
+                await Task.Yield();
                 continue;
             }
 

src/Mediator.SourceGenerator/SenderGenerator.cs

@@ -14,7 +14,7 @@ namespace Zapto.Mediator.Generator;
 [Generator]
 public class SenderGenerator : IIncrementalGenerator
 {
-    private bool _generateAssemblyInfo;
+    private readonly bool _generateAssemblyInfo;
 
     private static readonly string[] Interfaces =
     [

src/Mediator/Exceptions/HandlerNotFoundException.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Runtime.Serialization;
 
 namespace Zapto.Mediator;
 
@@ -9,10 +8,6 @@ public HandlerNotFoundException()
 	{
 	}
 
-	protected HandlerNotFoundException(SerializationInfo info, StreamingContext context) : base(info, context)
-	{
-	}
-
 	public HandlerNotFoundException(string message) : base(message)
 	{
 	}

src/Mediator/Exceptions/NamespaceHandlerNotFoundException.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Runtime.Serialization;
 
 namespace Zapto.Mediator;
 
@@ -9,10 +8,6 @@ public NamespaceHandlerNotFoundException()
 	{
 	}
 
-	protected NamespaceHandlerNotFoundException(SerializationInfo info, StreamingContext context) : base(info, context)
-	{
-	}
-
 	public NamespaceHandlerNotFoundException(string message) : base(message)
 	{
 	}