Don't dispose application singletons after running a code migration

2 files changed, 181 insertions, 29 deletions

diff --git a/Jellyfin.Server/Migrations/Stages/CodeMigration.cs b/Jellyfin.Server/Migrations/Stages/CodeMigration.cs
index 971b47608f..69e5cbc375 100644
--- a/Jellyfin.Server/Migrations/Stages/CodeMigration.cs
+++ b/Jellyfin.Server/Migrations/Stages/CodeMigration.cs
@@ -4,7 +4,6 @@ using System.Threading;
 using System.Threading.Tasks;
 using Jellyfin.Server.ServerSetupApp;
 using Microsoft.Extensions.DependencyInjection;
-using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.Logging;
 
 namespace Jellyfin.Server.Migrations.Stages;
@@ -22,28 +21,6 @@ internal class CodeMigration(Type migrationType, JellyfinMigrationAttribute meta
         return Metadata.Order.ToString("yyyyMMddHHmmsss", CultureInfo.InvariantCulture) + "_" + Metadata.Name!;
     }
 
-    private IServiceCollection MigrationServices(IServiceProvider serviceProvider, IStartupLogger logger)
-    {
-        var childServiceCollection = new ServiceCollection()
-            .AddSingleton(serviceProvider)
-            .AddSingleton(logger)
-            .AddSingleton(typeof(IStartupLogger<>), typeof(NestedStartupLogger<>))
-            .AddSingleton<StartupLogTopic>(logger.Topic!);
-
-        foreach (ServiceDescriptor service in serviceProvider.GetRequiredService<IServiceCollection>())
-        {
-            if (service.Lifetime == ServiceLifetime.Singleton && !service.ServiceType.IsGenericTypeDefinition)
-            {
-                childServiceCollection.AddSingleton(service.ServiceType, _ => serviceProvider.GetService(service.ServiceType)!);
-                continue;
-            }
-
-            childServiceCollection.Add(service);
-        }
-
-        return childServiceCollection;
-    }
-
     public async Task Perform(IServiceProvider? serviceProvider, IStartupLogger logger, CancellationToken cancellationToken)
     {
 #pragma warning disable CS0618 // Type or member is obsolete
@@ -55,10 +32,13 @@ internal class CodeMigration(Type migrationType, JellyfinMigrationAttribute meta
             }
             else
             {
-                using var migrationServices = MigrationServices(serviceProvider, logger).BuildServiceProvider();
-                ((IMigrationRoutine)ActivatorUtilities.CreateInstance(migrationServices, MigrationType)).Perform();
-#pragma warning restore CS0618 // Type or member is obsolete
+                var migrationServices = new MigrationServiceProvider(serviceProvider, logger);
+                await using (migrationServices.ConfigureAwait(false))
+                {
+                    ((IMigrationRoutine)ActivatorUtilities.CreateInstance(migrationServices, MigrationType)).Perform();
+                }
             }
+#pragma warning restore CS0618 // Type or member is obsolete
         }
         else if (typeof(IAsyncMigrationRoutine).IsAssignableFrom(MigrationType))
         {
@@ -68,8 +48,11 @@ internal class CodeMigration(Type migrationType, JellyfinMigrationAttribute meta
             }
             else
             {
-                using var migrationServices = MigrationServices(serviceProvider, logger).BuildServiceProvider();
-                await ((IAsyncMigrationRoutine)ActivatorUtilities.CreateInstance(migrationServices, MigrationType)).PerformAsync(cancellationToken).ConfigureAwait(false);
+                var migrationServices = new MigrationServiceProvider(serviceProvider, logger);
+                await using (migrationServices.ConfigureAwait(false))
+                {
+                    await ((IAsyncMigrationRoutine)ActivatorUtilities.CreateInstance(migrationServices, MigrationType)).PerformAsync(cancellationToken).ConfigureAwait(false);
+                }
             }
         }
         else
@@ -78,9 +61,88 @@ internal class CodeMigration(Type migrationType, JellyfinMigrationAttribute meta
         }
     }
 
+    /// <summary>
+    /// Provides the services a migration routine is constructed with.
+    /// </summary>
+    /// <remarks>
+    /// This overlays the migration scoped logging services onto a scope of the application container. Copying the
+    /// application service descriptors into a child container instead would make that child container the owner of
+    /// every singleton it forwards, so disposing it after the migration would also dispose the applications own
+    /// instance of services like the <c>ProviderManager</c> and leave the server broken until the next restart.
+    /// </remarks>
+    private sealed class MigrationServiceProvider : IServiceProvider, IServiceProviderIsService, IAsyncDisposable
+    {
+        private readonly AsyncServiceScope _scope;
+        private readonly IStartupLogger _logger;
+        private readonly IServiceProviderIsService? _isService;
+
+        public MigrationServiceProvider(IServiceProvider serviceProvider, IStartupLogger logger)
+        {
+            _scope = serviceProvider.CreateAsyncScope();
+            _logger = logger;
+            _isService = _scope.ServiceProvider.GetService<IServiceProviderIsService>();
+        }
+
+        public object? GetService(Type serviceType)
+        {
+            if (serviceType == typeof(IServiceProvider))
+            {
+                return this;
+            }
+
+            if (serviceType == typeof(IServiceProviderIsService))
+            {
+                return _isService is null ? null : this;
+            }
+
+            if (serviceType == typeof(IStartupLogger))
+            {
+                return _logger;
+            }
+
+            if (serviceType == typeof(StartupLogTopic))
+            {
+                return _logger.Topic;
+            }
+
+            if (IsCategoryLogger(serviceType))
+            {
+                var category = serviceType.GenericTypeArguments[0];
+                var baseLogger = _scope.ServiceProvider.GetRequiredService(typeof(ILogger<>).MakeGenericType(category));
+                return Activator.CreateInstance(typeof(NestedStartupLogger<>).MakeGenericType(category), baseLogger, _logger.Topic);
+            }
+
+            return _scope.ServiceProvider.GetService(serviceType);
+        }
+
+        public bool IsService(Type serviceType)
+        {
+            if (serviceType == typeof(IServiceProvider)
+                || serviceType == typeof(IServiceProviderIsService)
+                || serviceType == typeof(IStartupLogger)
+                || serviceType == typeof(StartupLogTopic)
+                || IsCategoryLogger(serviceType))
+            {
+                return true;
+            }
+
+            return _isService?.IsService(serviceType) ?? false;
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            return _scope.DisposeAsync();
+        }
+
+        private static bool IsCategoryLogger(Type serviceType)
+        {
+            return serviceType.IsConstructedGenericType && serviceType.GetGenericTypeDefinition() == typeof(IStartupLogger<>);
+        }
+    }
+
     private class NestedStartupLogger<TCategory> : StartupLogger<TCategory>
     {
-        public NestedStartupLogger(ILogger logger, StartupLogTopic topic) : base(logger, topic)
+        public NestedStartupLogger(ILogger logger, StartupLogTopic? topic) : base(logger, topic)
         {
         }
     }
diff --git a/tests/Jellyfin.Server.Tests/Migrations/CodeMigrationTests.cs b/tests/Jellyfin.Server.Tests/Migrations/CodeMigrationTests.cs
new file mode 100644
index 0000000000..68dd4486be
--- /dev/null
+++ b/tests/Jellyfin.Server.Tests/Migrations/CodeMigrationTests.cs
@@ -0,0 +1,90 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Jellyfin.Server.Migrations;
+using Jellyfin.Server.Migrations.Stages;
+using Jellyfin.Server.ServerSetupApp;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging.Abstractions;
+using Xunit;
+
+namespace Jellyfin.Server.Tests.Migrations;
+
+public class CodeMigrationTests
+{
+    [Fact]
+    public async Task Perform_LeavesApplicationSingletonsAlive()
+    {
+        var services = new ServiceCollection()
+            .AddLogging()
+            .RegisterStartupLogger()
+            .AddSingleton<ApplicationSingleton>()
+            .AddTransient<MigrationTransient>();
+        services.AddSingleton(services);
+
+        await using var serviceProvider = services.BuildServiceProvider();
+        var applicationSingleton = serviceProvider.GetRequiredService<ApplicationSingleton>();
+        var logger = new StartupLogger(NullLogger.Instance).BeginGroup($"Test migration");
+
+        var migration = new CodeMigration(
+            typeof(TestMigration),
+            new JellyfinMigrationAttribute("2026-09-05T10:00:00", nameof(TestMigration)),
+            null);
+        await migration.Perform(serviceProvider, logger, CancellationToken.None);
+
+        var performed = TestMigration.Performed;
+        Assert.NotNull(performed);
+        // The migration has to run against the applications own services, and they have to outlive it.
+        Assert.Same(applicationSingleton, performed.Singleton);
+        Assert.False(applicationSingleton.IsDisposed);
+        Assert.Same(applicationSingleton, serviceProvider.GetRequiredService<ApplicationSingleton>());
+        // Services created for the migration itself are still owned by the migration.
+        Assert.True(performed.Transient.IsDisposed);
+        // The startup logger has to stay attached to the topic of the running migration.
+        Assert.Same(logger.Topic, performed.Logger.Topic);
+    }
+
+    private sealed class ApplicationSingleton : IDisposable
+    {
+        public bool IsDisposed { get; private set; }
+
+        public void Dispose()
+        {
+            IsDisposed = true;
+        }
+    }
+
+    private sealed class MigrationTransient : IDisposable
+    {
+        public bool IsDisposed { get; private set; }
+
+        public void Dispose()
+        {
+            IsDisposed = true;
+        }
+    }
+
+    private sealed class TestMigration : IAsyncMigrationRoutine
+    {
+        public TestMigration(ApplicationSingleton singleton, MigrationTransient transient, IStartupLogger<TestMigration> logger)
+        {
+            Singleton = singleton;
+            Transient = transient;
+            Logger = logger;
+        }
+
+        public static TestMigration? Performed { get; private set; }
+
+        public ApplicationSingleton Singleton { get; }
+
+        public MigrationTransient Transient { get; }
+
+        public IStartupLogger<TestMigration> Logger { get; }
+
+        public Task PerformAsync(CancellationToken cancellationToken)
+        {
+            Performed = this;
+            return Task.CompletedTask;
+        }
+    }
+}