Fixed shutdown behavior

4 files changed, 17 insertions, 8 deletions

diff --git a/Jellyfin.Database/Jellyfin.Database.Implementations/IJellyfinDatabaseProvider.cs b/Jellyfin.Database/Jellyfin.Database.Implementations/IJellyfinDatabaseProvider.cs
index 72a6f819e..b27a88971 100644
--- a/Jellyfin.Database/Jellyfin.Database.Implementations/IJellyfinDatabaseProvider.cs
+++ b/Jellyfin.Database/Jellyfin.Database.Implementations/IJellyfinDatabaseProvider.cs
@@ -8,7 +8,7 @@ namespace Jellyfin.Server.Implementations;
 /// <summary>
 /// Defines the type and extension points for multi database support.
 /// </summary>
-public interface IJellyfinDatabaseProvider : IAsyncDisposable
+public interface IJellyfinDatabaseProvider
 {
     /// <summary>
     /// Gets or Sets the Database Factory when initialisaition is done.
@@ -33,4 +33,11 @@ public interface IJellyfinDatabaseProvider : IAsyncDisposable
     /// <param name="cancellationToken">The token to abort the operation.</param>
     /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
     Task RunScheduledOptimisation(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// If supported this should perform any actions that are required on stopping the jellyfin server.
+    /// </summary>
+    /// <param name="cancellationToken">The token that will be used to abort the operation.</param>
+    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
+    Task RunShutdownTask(CancellationToken cancellationToken);
 }
diff --git a/Jellyfin.Database/Jellyfin.Database.Providers.PgSql/PgSqlDatabaseProvider.cs b/Jellyfin.Database/Jellyfin.Database.Providers.PgSql/PgSqlDatabaseProvider.cs
index 1dae3401b..e6ae2acea 100644
--- a/Jellyfin.Database/Jellyfin.Database.Providers.PgSql/PgSqlDatabaseProvider.cs
+++ b/Jellyfin.Database/Jellyfin.Database.Providers.PgSql/PgSqlDatabaseProvider.cs
@@ -68,8 +68,8 @@ public sealed class PgSqlDatabaseProvider : IJellyfinDatabaseProvider
     }
 
     /// <inheritdoc/>
-    public ValueTask DisposeAsync()
+    public Task RunShutdownTask(CancellationToken cancellationToken)
     {
-        return ValueTask.CompletedTask;
+        return Task.CompletedTask;
     }
 }
diff --git a/Jellyfin.Database/Jellyfin.Database.Providers.SqLite/SqliteDatabaseProvider.cs b/Jellyfin.Database/Jellyfin.Database.Providers.SqLite/SqliteDatabaseProvider.cs
index 8ef5b6af5..907ea9156 100644
--- a/Jellyfin.Database/Jellyfin.Database.Providers.SqLite/SqliteDatabaseProvider.cs
+++ b/Jellyfin.Database/Jellyfin.Database.Providers.SqLite/SqliteDatabaseProvider.cs
@@ -64,13 +64,13 @@ public sealed class SqliteDatabaseProvider : IJellyfinDatabaseProvider
     }
 
     /// <inheritdoc/>
-    public async ValueTask DisposeAsync()
+    public async Task RunShutdownTask(CancellationToken cancellationToken)
     {
         // Run before disposing the application
-        var context = await DbContextFactory!.CreateDbContextAsync().ConfigureAwait(false);
+        var context = await DbContextFactory!.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);
         await using (context.ConfigureAwait(false))
         {
-            await context.Database.ExecuteSqlRawAsync("PRAGMA optimize").ConfigureAwait(false);
+            await context.Database.ExecuteSqlRawAsync("PRAGMA optimize", cancellationToken).ConfigureAwait(false);
         }
 
         SqliteConnection.ClearAllPools();
diff --git a/Jellyfin.Server/Program.cs b/Jellyfin.Server/Program.cs
index fd23b7e25..bb93ef138 100644
--- a/Jellyfin.Server/Program.cs
+++ b/Jellyfin.Server/Program.cs
@@ -4,6 +4,7 @@ using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using CommandLine;
 using Emby.Server.Implementations;
@@ -197,8 +198,9 @@ namespace Jellyfin.Server
                     _logger.LogInformation("Running query planner optimizations in the database... This might take a while");
 
                     var databaseProvider = appHost.ServiceProvider.GetRequiredService<IJellyfinDatabaseProvider>();
-
-                    await databaseProvider.DisposeAsync().ConfigureAwait(false);
+                    var shutdownSource = new CancellationTokenSource();
+                    shutdownSource.CancelAfter((int)TimeSpan.FromSeconds(60).TotalMicroseconds);
+                    await databaseProvider.RunShutdownTask(shutdownSource.Token).ConfigureAwait(false);
                 }
 
                 host?.Dispose();