Allow JIT killing of Raven.Embedded.EmbeddedServer by tearing it down as part of stop sequence (#4824)

* Teardown `Raven.Embedded.EmbeddedServer` as part of stop sequence so that we kill the instance if needed as adhering to cancellation is not possible in host builder root container dispose.

* Removed static

* Add why we set `GracefulShutdownTimeout` to Zero during dispose

* Enable Raven.Embedded console and file logging

* Set `GracefulShutdownTimeout` to 1 hour, as default is 30 seconds.

* Workaround for failing tests that seems to re-use the instance. Is that even supported?

* WaitForExitAsync is needed to wait for process to exit after kill signal

* Apply suggestions from Daniel's code review

Co-authored-by: Daniel Marbach <[email protected]>

* Ported disposing of RavenDB during stop to primary instance

* Dispose Process instance `using var`, ensure EmbeddedServer.Instance.Dispose always gets invoked, added comment and improved code readability

* Keep logging behavior default as it, only create client logs with LogMode.Information

* Update src/ServiceControl.RavenDB/EmbeddedDatabase.cs

* Update src/ServiceControl.Persistence.RavenDB/IRavenPersistenceLifecycle.cs

Co-authored-by: Mauro Servienti <[email protected]>

* Improved comments and cancellation behavior

* Fix comment

---------

Co-authored-by: Mauro Servienti <[email protected]>
Co-authored-by: Daniel Marbach <[email protected]>
Co-authored-by: Andreas Öhlund <[email protected]>

Author: 4 people

src/ServiceControl.Audit.Persistence.RavenDB/IRavenPersistenceLifecycle.cs

@@ -6,5 +6,6 @@
     interface IRavenPersistenceLifecycle
     {
         Task Initialize(CancellationToken cancellationToken = default);
+        Task Stop(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Audit.Persistence.RavenDB/RavenEmbeddedPersistenceLifecycle.cs

@@ -72,6 +72,8 @@ public async Task Initialize(CancellationToken cancellationToken = default)
             }
         }
 
+        public Task Stop(CancellationToken cancellationToken = default) => database!.Stop(cancellationToken);
+
         public void Dispose()
         {
             documentStore?.Dispose();

src/ServiceControl.Audit.Persistence.RavenDB/RavenExternalPersistenceLifecycle.cs

@@ -64,6 +64,8 @@ public async Task Initialize(CancellationToken cancellationToken = default)
             }
         }
 
+        public Task Stop(CancellationToken cancellationToken = default) => Task.CompletedTask; // We are not stopping an external instance
+
         public void Dispose() => documentStore?.Dispose();
 
         IDocumentStore? documentStore;

src/ServiceControl.Audit.Persistence.RavenDB/RavenPersistenceLifecycleHostedService.cs

@@ -8,6 +8,6 @@ sealed class RavenPersistenceLifecycleHostedService(IRavenPersistenceLifecycle l
     {
         public Task StartAsync(CancellationToken cancellationToken) => lifecycle.Initialize(cancellationToken);
 
-        public Task StopAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+        public Task StopAsync(CancellationToken cancellationToken) => lifecycle.Stop(cancellationToken);
     }
 }

src/ServiceControl.Persistence.RavenDB/IRavenPersistenceLifecycle.cs

@@ -8,5 +8,6 @@ namespace ServiceControl.Persistence.RavenDB
     interface IRavenPersistenceLifecycle
     {
         Task Initialize(CancellationToken cancellationToken = default);
+        Task Stop(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence.RavenDB/RavenEmbeddedPersistenceLifecycle.cs

@@ -69,6 +69,14 @@ public async Task Initialize(CancellationToken cancellationToken)
             }
         }
 
+        public async Task Stop(CancellationToken cancellationToken)
+        {
+            if (database != null)
+            {
+                await database.Stop(cancellationToken);
+            }
+        }
+
         public void Dispose()
         {
             documentStore?.Dispose();

src/ServiceControl.Persistence.RavenDB/RavenExternalPersistenceLifecycle.cs

@@ -59,6 +59,8 @@ public async Task Initialize(CancellationToken cancellationToken)
             }
         }
 
+        public Task Stop(CancellationToken cancellationToken) => Task.CompletedTask;
+
         public void Dispose() => documentStore?.Dispose();
 
         IDocumentStore? documentStore;

src/ServiceControl.Persistence.RavenDB/RavenPersistenceLifecycleHostedService.cs

@@ -10,6 +10,6 @@ class RavenPersistenceLifecycleHostedService(IRavenPersistenceLifecycle persiste
     {
         public Task StartAsync(CancellationToken cancellationToken) => persistenceLifecycle.Initialize(cancellationToken);
 
-        public Task StopAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+        public Task StopAsync(CancellationToken cancellationToken) => persistenceLifecycle.Stop(cancellationToken);
     }
 }

src/ServiceControl.RavenDB/EmbeddedDatabase.cs

@@ -15,6 +15,7 @@ namespace ServiceControl.RavenDB
     using Raven.Client.Documents.Conventions;
     using Raven.Client.ServerWide.Operations;
     using Raven.Embedded;
+    using Sparrow.Logging;
 
     public sealed class EmbeddedDatabase : IDisposable
     {
@@ -52,6 +53,20 @@ public static EmbeddedDatabase Start(EmbeddedDatabaseConfiguration databaseConfi
 
             var nugetPackagesPath = Path.Combine(databaseConfiguration.DbPath, "Packages", "NuGet");
 
+            var logMode = Enum.Parse<LogMode>(databaseConfiguration.LogsMode);
+
+            if (logMode == LogMode.Information) // Most verbose
+            {
+                LoggingSource.Instance.EnableConsoleLogging();
+                LoggingSource.Instance.SetupLogMode(
+                    logMode,
+                    Path.Combine(databaseConfiguration.LogPath, "Raven.Embedded"),
+                    retentionTime: TimeSpan.FromDays(14),
+                    retentionSize: 1024 * 1024 * 10,
+                    compress: false
+                );
+            }
+
             Logger.InfoFormat("Loading RavenDB license from {0}", licenseFileNameAndServerDirectory.LicenseFileName);
             var serverOptions = new ServerOptions
             {
@@ -85,6 +100,7 @@ public static EmbeddedDatabase Start(EmbeddedDatabaseConfiguration databaseConfi
 
         void Start(ServerOptions serverOptions)
         {
+            this.serverOptions = serverOptions;
             EmbeddedServer.Instance.ServerProcessExited += OnServerProcessExited;
             EmbeddedServer.Instance.StartServer(serverOptions);
 
@@ -162,23 +178,95 @@ public async Task<IDocumentStore> Connect(CancellationToken cancellationToken)
 
         public async Task DeleteDatabase(string dbName)
         {
-            using var store = await EmbeddedServer.Instance.GetDocumentStoreAsync(new DatabaseOptions(dbName) { SkipCreatingDatabase = true });
+            using var store = await EmbeddedServer.Instance.GetDocumentStoreAsync(new DatabaseOptions(dbName)
+            {
+                SkipCreatingDatabase = true
+            });
             await store.Maintenance.Server.SendAsync(new DeleteDatabasesOperation(dbName, true));
         }
 
+        public async Task Stop(CancellationToken cancellationToken)
+        {
+            Logger.Debug("Stopping RavenDB server");
+            EmbeddedServer.Instance.ServerProcessExited -= OnServerProcessExited;
+
+            await shutdownTokenSource.CancelAsync();
+
+            // This is a workaround until the EmbeddedServer properly supports cancellation!
+            //
+            // EmbeddedServer does not have an async Stop method, the Dispose operation blocks and waits
+            // until its GracefulShutdownTimeout is reached and then does a Process.Kill. Due to this behavior this can
+            // be shorter or longer than the allowed stop duration.
+            //
+            // When Task.WhenAny is called, 2 things can happen:
+            //
+            // a. The Task.Delay gets cancelled first
+            // b. The EmbeddedServer.Dispose completes first
+            //
+            // If the Task.Delay gets cancelled first this means Dispose is still running and
+            // then we try and kill the process, if not disposed completed and we're done.
+
+            serverOptions!.GracefulShutdownTimeout = TimeSpan.FromHours(1); // During Stop/Dispose we manually control this
+
+
+            // Dispose always need to be invoked, even when already cancelled
+            var disposeTask = Task.Run(() => EmbeddedServer.Instance.Dispose(), CancellationToken.None);
+
+            // Runs "infinite" but will eventually get cancelled
+            var waitForCancellationTask = Task.Delay(Timeout.Infinite, cancellationToken);
+
+            var delayIsCancelled = waitForCancellationTask == await Task.WhenAny(disposeTask, waitForCancellationTask);
+
+            if (delayIsCancelled)
+            {
+                int processId = 0;
+                try
+                {
+                    // We always want to try and kill the process, even when already cancelled
+                    processId = await EmbeddedServer.Instance.GetServerProcessIdAsync(CancellationToken.None);
+                    Logger.WarnFormat("Killing RavenDB server PID {0} because host cancelled", processId);
+                    using var ravenChildProcess = Process.GetProcessById(processId);
+                    ravenChildProcess.Kill(entireProcessTree: true);
+                    // Kill only signals
+                    Logger.WarnFormat("Waiting for RavenDB server PID {0} to exit... ", processId);
+                    // When WaitForExitAsync returns, the process could still exist but in a frozen state to flush
+                    // memory mapped pages to storage.
+                    await ravenChildProcess.WaitForExitAsync(CancellationToken.None);
+                }
+                catch (Exception e)
+                {
+                    Logger.ErrorFormat("Failed to kill RavenDB server PID {0} shutdown\n{1}", processId, e);
+                }
+            }
+
+            serverOptions = null!;
+            Logger.Debug("Stopped RavenDB server");
+        }
+
         public void Dispose()
         {
             if (disposed)
             {
                 return;
             }
 
-            EmbeddedServer.Instance.ServerProcessExited -= OnServerProcessExited;
+            if (serverOptions != null)
+            {
+                EmbeddedServer.Instance.ServerProcessExited -= OnServerProcessExited;
+            }
 
             shutdownTokenSource.Cancel();
-            Logger.Debug("Disposing RavenDB server");
-            EmbeddedServer.Instance.Dispose();
-            Logger.Debug("Dispose RavenDB server");
+
+            if (serverOptions != null)
+            {
+                // Set GracefulShutdownTimeout to Zero and exit ASAP, under normal operation instance would already
+                // have been allowed to gracefully stop during "Stop" method.
+                serverOptions!.GracefulShutdownTimeout = TimeSpan.Zero;
+                Logger.Debug("Disposing RavenDB server");
+                EmbeddedServer.Instance.Dispose();
+                Logger.Debug("Disposed RavenDB server");
+            }
+
             shutdownTokenSource.Dispose();
             applicationStoppingRegistration.Dispose();
 
@@ -213,6 +301,7 @@ static long DataSize(EmbeddedDatabaseConfiguration configuration)
                 {
                     return -1;
                 }
+
                 return info.Length;
             }
             catch
@@ -262,8 +351,9 @@ static long DirSize(DirectoryInfo d)
         readonly EmbeddedDatabaseConfiguration configuration;
         readonly CancellationToken shutdownCancellationToken;
         readonly CancellationTokenRegistration applicationStoppingRegistration;
+        ServerOptions? serverOptions;
 
         static TimeSpan delayBetweenRestarts = TimeSpan.FromSeconds(60);
         static readonly ILog Logger = LogManager.GetLogger<EmbeddedDatabase>();
     }
-}
+}