Fix the library scheduler never retiring its runners

2 files changed, 332 insertions, 29 deletions

diff --git a/MediaBrowser.Controller/LibraryTaskScheduler/LimitedConcurrencyLibraryScheduler.cs b/MediaBrowser.Controller/LibraryTaskScheduler/LimitedConcurrencyLibraryScheduler.cs
index 6da398129a..3c6ffe3cbd 100644
--- a/MediaBrowser.Controller/LibraryTaskScheduler/LimitedConcurrencyLibraryScheduler.cs
+++ b/MediaBrowser.Controller/LibraryTaskScheduler/LimitedConcurrencyLibraryScheduler.cs
@@ -17,7 +17,8 @@ namespace MediaBrowser.Controller.LibraryTaskScheduler;
 /// </summary>
 public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibraryScheduler, IAsyncDisposable
 {
-    private const int CleanupGracePeriod = 60;
+    private static readonly TimeSpan _cleanupGracePeriod = TimeSpan.FromSeconds(60);
+
     private readonly IHostApplicationLifetime _hostApplicationLifetime;
     private readonly ILogger<LimitedConcurrencyLibraryScheduler> _logger;
     private readonly IServerConfigurationManager _serverConfigurationManager;
@@ -31,6 +32,8 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
     private readonly Lock _taskLock = new();
 
     private readonly Channel<TaskQueueItem> _tasks = Channel.CreateUnbounded<TaskQueueItem>();
+    private readonly CancellationTokenSource _disposeTokenSource = new();
+    private readonly TimeSpan _gracePeriod;
 
     private volatile int _workCounter;
     private Task? _cleanupTask;
@@ -46,10 +49,34 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
         IHostApplicationLifetime hostApplicationLifetime,
         ILogger<LimitedConcurrencyLibraryScheduler> logger,
         IServerConfigurationManager serverConfigurationManager)
+        : this(hostApplicationLifetime, logger, serverConfigurationManager, _cleanupGracePeriod)
+    {
+    }
+
+    internal LimitedConcurrencyLibraryScheduler(
+        IHostApplicationLifetime hostApplicationLifetime,
+        ILogger<LimitedConcurrencyLibraryScheduler> logger,
+        IServerConfigurationManager serverConfigurationManager,
+        TimeSpan gracePeriod)
     {
         _hostApplicationLifetime = hostApplicationLifetime;
         _logger = logger;
         _serverConfigurationManager = serverConfigurationManager;
+        _gracePeriod = gracePeriod;
+    }
+
+    /// <summary>
+    /// Gets the number of runners the scheduler currently keeps alive.
+    /// </summary>
+    internal int ActiveRunnerCount
+    {
+        get
+        {
+            lock (_taskLock)
+            {
+                return _taskRunners.Count;
+            }
+        }
     }
 
     private void ScheduleTaskCleanup()
@@ -68,31 +95,65 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
 
         async Task RunCleanupTask()
         {
-            _logger.LogDebug("Schedule cleanup task in {CleanupGracePerioid} sec.", CleanupGracePeriod);
-            await Task.Delay(TimeSpan.FromSeconds(CleanupGracePeriod)).ConfigureAwait(false);
-            if (_disposed)
+            while (true)
             {
-                _logger.LogDebug("Abort cleaning up, already disposed.");
-                return;
-            }
+                _logger.LogDebug("Schedule cleanup task in {CleanupGracePeriod}.", _gracePeriod);
+                try
+                {
+                    await Task.Delay(_gracePeriod, _disposeTokenSource.Token).ConfigureAwait(false);
+                }
+                catch (OperationCanceledException)
+                {
+                    _logger.LogDebug("Abort cleaning up, already disposed.");
+                    return;
+                }
 
-            lock (_taskLock)
-            {
-                if (_tasks.Reader.Count > 0 || _workCounter > 0)
+                if (_disposed)
                 {
-                    _logger.LogDebug("Delay cleanup task, operations still running.");
-                    // tasks are still there so its still in use. Reschedule cleanup task.
-                    // we cannot just exit here and rely on the other invoker because there is a considerable timeframe where it could have already ended.
-                    _cleanupTask = RunCleanupTask();
+                    _logger.LogDebug("Abort cleaning up, already disposed.");
                     return;
                 }
+
+                CancellationTokenSource[] runners;
+                lock (_taskLock)
+                {
+                    if (_tasks.Reader.Count > 0 || _workCounter > 0)
+                    {
+                        _logger.LogDebug("Delay cleanup task, operations still running.");
+                        // tasks are still there so its still in use. Wait another grace period.
+                        // we cannot just exit here and rely on the other invoker because there is a considerable timeframe where it could have already ended.
+                        continue;
+                    }
+
+                    runners = [.. _taskRunners.Keys];
+
+                    // Retire the runners before they are told to stop: an operation starting while
+                    // they wind down must spawn its own instead of counting these towards the fanout.
+                    _taskRunners.Clear();
+
+                    // Hand the next operation the ability to schedule a cleanup again. Without this
+                    // the very first cleanup would be the only one that ever runs.
+                    _cleanupTask = null;
+                }
+
+                _logger.LogDebug("Cleanup runners.");
+                await StopRunners(runners).ConfigureAwait(false);
+                return;
             }
+        }
+    }
 
-            _logger.LogDebug("Cleanup runners.");
-            foreach (var item in _taskRunners.ToArray())
+    private static async Task StopRunners(CancellationTokenSource[] runners)
+    {
+        foreach (var runner in runners)
+        {
+            try
             {
-                await item.Key.CancelAsync().ConfigureAwait(false);
-                _taskRunners.Remove(item.Key);
+                await runner.CancelAsync().ConfigureAwait(false);
+            }
+            catch (ObjectDisposedException)
+            {
+                // The runner already stopped on its own and disposed its stop source.
             }
         }
     }
@@ -127,11 +188,14 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
             {
                 var stopToken = new CancellationTokenSource();
                 var combinedSource = CancellationTokenSource.CreateLinkedTokenSource(stopToken.Token, _hostApplicationLifetime.ApplicationStopping);
+
+                // Keyed on its own stop source, because cancelling that is what reaches the linked
+                // source the runner waits on. Cancellation does not travel the other way.
                 _taskRunners.Add(
-                    combinedSource,
+                    stopToken,
                     Task.Factory.StartNew(
                         ItemWorker,
-                        (combinedSource, stopToken),
+                        (stopToken, combinedSource),
                         combinedSource.Token,
                         TaskCreationOptions.PreferFairness,
                         TaskScheduler.Default));
@@ -145,7 +209,7 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
         _deadlockDetector.Value = stopToken.TaskStop;
         try
         {
-            while (!stopToken.GlobalStop.Token.IsCancellationRequested)
+            while (!stopToken.GlobalStop.IsCancellationRequested)
             {
                 var item = await _tasks.Reader.ReadAsync(stopToken.GlobalStop.Token).ConfigureAwait(false);
                 try
@@ -162,15 +226,24 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
                 }
             }
         }
-        catch (OperationCanceledException) when (stopToken.TaskStop.IsCancellationRequested)
+        catch (OperationCanceledException) when (stopToken.GlobalStop.IsCancellationRequested)
         {
             // thats how you do it, interupt the waiter thread. There is nothing to do here when it was on purpose.
         }
+        catch (ChannelClosedException)
+        {
+            // the scheduler was disposed and will not hand out any more work.
+        }
         finally
         {
             _logger.LogDebug("Cleanup Runner'.");
             _deadlockDetector.Value = default!;
-            _taskRunners.Remove(stopToken.TaskStop);
+
+            lock (_taskLock)
+            {
+                _taskRunners.Remove(stopToken.TaskStop);
+            }
+
             stopToken.GlobalStop.Dispose();
             stopToken.TaskStop.Dispose();
         }
@@ -195,7 +268,7 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
         finally
         {
             item.Progress.Report(100);
-            item.Done.SetResult();
+            item.Done.TrySetResult();
         }
     }
 
@@ -285,16 +358,33 @@ public sealed class LimitedConcurrencyLibraryScheduler : ILimitedConcurrencyLibr
 
         _disposed = true;
         _tasks.Writer.Complete();
-        foreach (var item in _taskRunners)
+
+        // Nobody is left to run these, so release whoever is waiting on them.
+        while (_tasks.Reader.TryRead(out var item))
         {
-            await item.Key.CancelAsync().ConfigureAwait(false);
+            item.Done.TrySetResult();
         }
 
-        if (_cleanupTask is not null)
+        CancellationTokenSource[] runners;
+        Task? cleanupTask;
+        lock (_taskLock)
         {
-            await _cleanupTask.ConfigureAwait(false);
-            _cleanupTask?.Dispose();
+            runners = [.. _taskRunners.Keys];
+            _taskRunners.Clear();
+            cleanupTask = _cleanupTask;
         }
+
+        await StopRunners(runners).ConfigureAwait(false);
+
+        // Cuts the grace period short instead of holding up shutdown for the rest of it.
+        await _disposeTokenSource.CancelAsync().ConfigureAwait(false);
+
+        if (cleanupTask is not null)
+        {
+            await cleanupTask.ConfigureAwait(false);
+        }
+
+        _disposeTokenSource.Dispose();
     }
 
     private class TaskQueueItem
diff --git a/tests/Jellyfin.Controller.Tests/LibraryTaskScheduler/LimitedConcurrencyLibrarySchedulerTests.cs b/tests/Jellyfin.Controller.Tests/LibraryTaskScheduler/LimitedConcurrencyLibrarySchedulerTests.cs
new file mode 100644
index 0000000000..21a719a78f
--- /dev/null
+++ b/tests/Jellyfin.Controller.Tests/LibraryTaskScheduler/LimitedConcurrencyLibrarySchedulerTests.cs
@@ -0,0 +1,213 @@
+using System;
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using MediaBrowser.Controller.Configuration;
+using MediaBrowser.Controller.LibraryTaskScheduler;
+using MediaBrowser.Model.Configuration;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging.Abstractions;
+using Moq;
+using Xunit;
+
+namespace Jellyfin.Controller.Tests.LibraryTaskScheduler
+{
+    public class LimitedConcurrencyLibrarySchedulerTests
+    {
+        private static readonly TimeSpan _shortGracePeriod = TimeSpan.FromMilliseconds(50);
+
+        // Generous, because these only ever wait for something that should already have happened.
+        private static readonly TimeSpan _timeout = TimeSpan.FromSeconds(10);
+
+        [Fact]
+        public async Task Enqueue_ProcessesEveryItem()
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping);
+            await using (scheduler)
+            {
+                var data = Enumerable.Range(0, 100).ToArray();
+                var processed = new ConcurrentBag<int>();
+
+                await scheduler.Enqueue(
+                    data,
+                    (item, _) =>
+                    {
+                        processed.Add(item);
+                        return Task.CompletedTask;
+                    },
+                    new Progress<double>(),
+                    CancellationToken.None);
+
+                Assert.Equal(data, processed.Order());
+            }
+        }
+
+        [Fact]
+        public async Task Enqueue_WithFailingWorker_StillCompletes()
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping);
+            await using (scheduler)
+            {
+                await scheduler.Enqueue(
+                    Enumerable.Range(0, 20).ToArray(),
+                    (item, _) => item % 2 == 0 ? throw new InvalidOperationException("boom") : Task.CompletedTask,
+                    new Progress<double>(),
+                    CancellationToken.None);
+            }
+        }
+
+        /// <summary>
+        /// The runners wait on a source linked to <see cref="IHostApplicationLifetime.ApplicationStopping"/>,
+        /// so a shutdown has to reach them. It does not travel from the linked source back to the one
+        /// the cleanup cancels, which is what made them immortal.
+        /// </summary>
+        [Fact]
+        public async Task ApplicationStopping_RetiresRunners()
+        {
+            using var appStopping = new CancellationTokenSource();
+
+            // Long enough that the cleanup cannot be what retires them.
+            var scheduler = CreateScheduler(appStopping, gracePeriod: TimeSpan.FromMinutes(5));
+            await using (scheduler)
+            {
+                await RunOneOperation(scheduler);
+                Assert.True(scheduler.ActiveRunnerCount > 0);
+
+                await appStopping.CancelAsync();
+
+                await WaitForAsync(() => scheduler.ActiveRunnerCount == 0);
+            }
+        }
+
+        /// <summary>
+        /// The cleanup used to be a one shot: it never released the scheduling slot it took, so
+        /// every runner spawned after the first pass stayed around for the lifetime of the server.
+        /// </summary>
+        [Fact]
+        public async Task Enqueue_RetiresIdleRunnersAfterEveryOperation()
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping);
+            await using (scheduler)
+            {
+                for (var round = 0; round < 3; round++)
+                {
+                    await RunOneOperation(scheduler);
+                    Assert.True(scheduler.ActiveRunnerCount > 0, $"no runner spawned in round {round}");
+
+                    await WaitForAsync(() => scheduler.ActiveRunnerCount == 0);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Disposing used to sit out the rest of the cleanup grace period, holding up shutdown for
+        /// up to a minute.
+        /// </summary>
+        [Fact]
+        public async Task DisposeAsync_DoesNotWaitOutTheGracePeriod()
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping, gracePeriod: TimeSpan.FromMinutes(5));
+
+            await RunOneOperation(scheduler);
+
+            var stopwatch = Stopwatch.StartNew();
+            await scheduler.DisposeAsync();
+
+            Assert.True(stopwatch.Elapsed < _timeout, $"disposing took {stopwatch.Elapsed}");
+        }
+
+        [Fact]
+        public async Task Enqueue_AfterDispose_DoesNothing()
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping);
+            await scheduler.DisposeAsync();
+
+            var processed = 0;
+            await scheduler.Enqueue(
+                Enumerable.Range(0, 10).ToArray(),
+                (_, _) =>
+                {
+                    Interlocked.Increment(ref processed);
+                    return Task.CompletedTask;
+                },
+                new Progress<double>(),
+                CancellationToken.None);
+
+            Assert.Equal(0, processed);
+        }
+
+        [Theory]
+        [InlineData(1)]
+        [InlineData(4)]
+        public async Task Enqueue_FromWithinAWorker_DoesNotDeadlock(int fanout)
+        {
+            using var appStopping = new CancellationTokenSource();
+            var scheduler = CreateScheduler(appStopping, fanout: fanout);
+            await using (scheduler)
+            {
+                var inner = 0;
+
+                var outer = scheduler.Enqueue(
+                    Enumerable.Range(0, 8).ToArray(),
+                    (_, _) => scheduler.Enqueue(
+                        Enumerable.Range(0, 4).ToArray(),
+                        (_, _) =>
+                        {
+                            Interlocked.Increment(ref inner);
+                            return Task.CompletedTask;
+                        },
+                        new Progress<double>(),
+                        CancellationToken.None),
+                    new Progress<double>(),
+                    CancellationToken.None);
+
+                await outer.WaitAsync(_timeout, TestContext.Current.CancellationToken);
+
+                Assert.Equal(32, inner);
+            }
+        }
+
+        private static LimitedConcurrencyLibraryScheduler CreateScheduler(
+            CancellationTokenSource appStopping,
+            int fanout = 4,
+            TimeSpan? gracePeriod = null)
+        {
+            var lifetime = new Mock<IHostApplicationLifetime>();
+            lifetime.SetupGet(x => x.ApplicationStopping).Returns(() => appStopping.Token);
+
+            var configurationManager = new Mock<IServerConfigurationManager>();
+            configurationManager.SetupGet(x => x.Configuration)
+                .Returns(new ServerConfiguration { LibraryScanFanoutConcurrency = fanout });
+
+            return new LimitedConcurrencyLibraryScheduler(
+                lifetime.Object,
+                NullLogger<LimitedConcurrencyLibraryScheduler>.Instance,
+                configurationManager.Object,
+                gracePeriod ?? _shortGracePeriod);
+        }
+
+        private static Task RunOneOperation(LimitedConcurrencyLibraryScheduler scheduler)
+            => scheduler.Enqueue(
+                Enumerable.Range(0, 8).ToArray(),
+                (_, _) => Task.CompletedTask,
+                new Progress<double>(),
+                CancellationToken.None);
+
+        private static async Task WaitForAsync(Func<bool> condition)
+        {
+            var stopwatch = Stopwatch.StartNew();
+            while (!condition())
+            {
+                Assert.True(stopwatch.Elapsed < _timeout, "timed out waiting for the scheduler to settle");
+                await Task.Delay(20, TestContext.Current.CancellationToken);
+            }
+        }
+    }
+}