Fix SQLite structured log shell lifecycle (#7461)

* Fix SQLite structured log shell lifecycle

* Address Greptile lifecycle feedback

* Tighten structured log lifecycle guards

* Address structured log background task review

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Relational/Extensions/RelationalStructuredLogsServiceCollectionExtensions.cs

@@ -3,6 +3,7 @@ using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Contracts;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Options;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Services;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Stores;
+using Elsa.Extensions;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 
@@ -24,6 +25,7 @@ public static class RelationalStructuredLogsServiceCollectionExtensions
         services.Replace(ServiceDescriptor.Singleton<IStructuredLogStore>(sp => sp.GetRequiredService<StructuredLogWriteBuffer>()));
         services.Replace(ServiceDescriptor.Singleton<IStructuredLogWriteBuffer>(sp => sp.GetRequiredService<StructuredLogWriteBuffer>()));
         services.TryAddEnumerable(ServiceDescriptor.Singleton<IStructuredLogStorageDiagnostics, StructuredLogWriteBufferStorageDiagnostics>());
+        services.AddBackgroundTask<StructuredLogWriteBufferBackgroundTask>();
         services.AddHostedService(sp => sp.GetRequiredService<StructuredLogWriteBuffer>());
 
         return services;

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Relational/Services/StructuredLogWriteBuffer.cs

@@ -13,11 +13,13 @@ public class StructuredLogWriteBuffer(
     RelationalStructuredLogStore store,
     IOptions<RelationalStructuredLogOptions> options) : IStructuredLogStore, IStructuredLogWriteBuffer, IHostedService, IAsyncDisposable
 {
+    private readonly object _lifecycleLock = new();
     private readonly Queue<StructuredLogEvent> _queue = new();
     private readonly SemaphoreSlim _signal = new(0);
-    private readonly CancellationTokenSource _stopTokenSource = new();
+    private CancellationTokenSource _stopTokenSource = new();
     private Task? _backgroundTask;
     private long _droppedWriteCount;
+    private int _activeStartCount;
     private int _disposed;
 
     public long DroppedWriteCount => Interlocked.Read(ref _droppedWriteCount);
@@ -62,21 +64,54 @@ public class StructuredLogWriteBuffer(
 
     public Task StartAsync(CancellationToken cancellationToken)
     {
-        _backgroundTask ??= Task.Run(ProcessQueueAsync, CancellationToken.None);
+        lock (_lifecycleLock)
+        {
+            _activeStartCount++;
+
+            if (_backgroundTask is { IsCompleted: false })
+                return Task.CompletedTask;
+
+            if (_stopTokenSource.IsCancellationRequested)
+            {
+                _stopTokenSource.Dispose();
+                _stopTokenSource = new();
+            }
+
+            var stopToken = _stopTokenSource.Token;
+            _backgroundTask = Task.Run(() => ProcessQueueAsync(stopToken), CancellationToken.None);
+        }
+
         return Task.CompletedTask;
     }
 
     public async Task StopAsync(CancellationToken cancellationToken)
     {
-        await _stopTokenSource.CancelAsync();
+        Task? backgroundTask;
+        CancellationTokenSource stopTokenSource;
 
-        if (_backgroundTask != null)
+        lock (_lifecycleLock)
+        {
+            if (_activeStartCount == 0)
+                return;
+
+            _activeStartCount--;
+
+            if (_activeStartCount > 0)
+                return;
+
+            backgroundTask = _backgroundTask;
+            stopTokenSource = _stopTokenSource;
+        }
+
+        await stopTokenSource.CancelAsync();
+
+        if (backgroundTask != null)
         {
             try
             {
-                await _backgroundTask.WaitAsync(cancellationToken);
+                await backgroundTask.WaitAsync(cancellationToken);
             }
-            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested || _stopTokenSource.IsCancellationRequested)
+            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested || stopTokenSource.IsCancellationRequested)
             {
                 // Expected during shutdown; remaining queued writes are flushed below.
             }
@@ -119,16 +154,25 @@ public class StructuredLogWriteBuffer(
         if (Interlocked.Exchange(ref _disposed, 1) == 1)
             return;
 
-        await _stopTokenSource.CancelAsync();
+        Task? backgroundTask;
+        CancellationTokenSource stopTokenSource;
+
+        lock (_lifecycleLock)
+        {
+            backgroundTask = _backgroundTask;
+            stopTokenSource = _stopTokenSource;
+        }
+
+        await stopTokenSource.CancelAsync();
         using var timeoutTokenSource = new CancellationTokenSource(options.Value.WriteQueue.ShutdownFlushTimeout);
 
-        if (_backgroundTask != null)
+        if (backgroundTask != null)
         {
             try
             {
-                await _backgroundTask.WaitAsync(timeoutTokenSource.Token);
+                await backgroundTask.WaitAsync(timeoutTokenSource.Token);
             }
-            catch (OperationCanceledException) when (timeoutTokenSource.IsCancellationRequested || _stopTokenSource.IsCancellationRequested)
+            catch (OperationCanceledException) when (timeoutTokenSource.IsCancellationRequested || stopTokenSource.IsCancellationRequested)
             {
                 CountPendingWritesAsDropped();
             }
@@ -144,13 +188,11 @@ public class StructuredLogWriteBuffer(
         }
 
         _signal.Dispose();
-        _stopTokenSource.Dispose();
+        stopTokenSource.Dispose();
     }
 
-    private async Task ProcessQueueAsync()
+    private async Task ProcessQueueAsync(CancellationToken cancellationToken)
     {
-        var cancellationToken = _stopTokenSource.Token;
-
         while (!cancellationToken.IsCancellationRequested)
         {
             try

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Relational/Services/StructuredLogWriteBufferBackgroundTask.cs

@@ -0,0 +1,31 @@
+using Elsa.Common;
+
+namespace Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Services;
+
+public class StructuredLogWriteBufferBackgroundTask(StructuredLogWriteBuffer writeBuffer) : IBackgroundTask
+{
+    private volatile bool _started;
+
+    public async Task StartAsync(CancellationToken cancellationToken)
+    {
+        if (_started)
+            return;
+
+        await writeBuffer.StartAsync(cancellationToken);
+        _started = true;
+    }
+
+    public Task ExecuteAsync(CancellationToken cancellationToken)
+    {
+        return Task.CompletedTask;
+    }
+
+    public async Task StopAsync(CancellationToken cancellationToken)
+    {
+        if (!_started)
+            return;
+
+        _started = false;
+        await writeBuffer.StopAsync(cancellationToken);
+    }
+}

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Relational/ShellFeatures/StructuredLogRelationalPersistenceFeature.cs

@@ -1,4 +1,5 @@
 using CShells.Features;
+using Elsa.Diagnostics.StructuredLogs.ShellFeatures;
 using JetBrains.Annotations;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -10,7 +11,7 @@ namespace Elsa.Diagnostics.StructuredLogs.Persistence.Relational.ShellFeatures;
 [ShellFeature(
     DisplayName = "Structured Log Relational Persistence",
     Description = "Provides shared relational persistence services for diagnostics structured logs",
-    DependsOn = ["Structured Logs"])]
+    DependsOn = [typeof(StructuredLogsFeature)])]
 [UsedImplicitly]
 public class StructuredLogRelationalPersistenceFeature : IShellFeature
 {

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite/Extensions/SqliteStructuredLogsModuleExtensions.cs

@@ -1,3 +1,4 @@
+using Elsa.Common;
 using Elsa.Diagnostics.StructuredLogs.Features;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Contracts;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Extensions;
@@ -45,7 +46,9 @@ public static class SqliteStructuredLogsModuleExtensions
         services.TryAddSingleton<IRelationalStructuredLogConnectionFactory, SqliteStructuredLogConnectionFactory>();
         services.TryAddSingleton<IRelationalStructuredLogDialect, SqliteStructuredLogDialect>();
         services.TryAddSingleton<IStructuredLogSchemaMigrator, SqliteStructuredLogSchemaMigrator>();
-        services.AddHostedService<SqliteStructuredLogStartupService>();
+        services.TryAddSingleton<SqliteStructuredLogStartupService>();
+        services.AddHostedService(sp => sp.GetRequiredService<SqliteStructuredLogStartupService>());
+        services.AddScoped<IStartupTask>(sp => sp.GetRequiredService<SqliteStructuredLogStartupService>());
         services.AddRelationalStructuredLogPersistence();
 
         return services;

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite/Services/SqliteStructuredLogStartupService.cs

@@ -1,3 +1,4 @@
+using Elsa.Common;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Contracts;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Services;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.Options;
@@ -9,15 +10,36 @@ namespace Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.Services;
 public class SqliteStructuredLogStartupService(
     IStructuredLogSchemaMigrator schemaMigrator,
     StructuredLogRetentionService retentionService,
-    IOptions<SqliteStructuredLogOptions> options) : IHostedService
+    IOptions<SqliteStructuredLogOptions> options) : IHostedService, IStartupTask
 {
+    private readonly SemaphoreSlim _startupLock = new(1, 1);
+    private bool _executed;
+
     public async Task StartAsync(CancellationToken cancellationToken)
     {
+        await ExecuteAsync(cancellationToken);
+    }
+
+    public async Task ExecuteAsync(CancellationToken cancellationToken)
+    {
+        await _startupLock.WaitAsync(cancellationToken);
+        try
+        {
+            if (_executed)
+                return;
+
             if (options.Value.RunMigrationsOnStartup)
                 await schemaMigrator.MigrateAsync(cancellationToken);
 
             if (options.Value.Relational.Retention.CleanupOnStartup)
                 await retentionService.CleanupAsync(cancellationToken);
+
+            _executed = true;
+        }
+        finally
+        {
+            _startupLock.Release();
+        }
     }
 
     public Task StopAsync(CancellationToken cancellationToken) => Task.CompletedTask;

src/modules/Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite/ShellFeatures/SqliteStructuredLogPersistenceFeature.cs

@@ -1,4 +1,5 @@
 using CShells.Features;
+using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.ShellFeatures;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.Extensions;
 using JetBrains.Annotations;
 using Microsoft.Extensions.DependencyInjection;
@@ -11,7 +12,7 @@ namespace Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.ShellFeatures;
 [ShellFeature(
     DisplayName = "SQLite Structured Log Persistence",
     Description = "Provides SQLite persistence for diagnostics structured logs",
-    DependsOn = ["Structured Log Relational Persistence"])]
+    DependsOn = [typeof(StructuredLogRelationalPersistenceFeature)])]
 [UsedImplicitly]
 public class SqliteStructuredLogPersistenceFeature : IShellFeature
 {

test/integration/Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.IntegrationTests/SqliteStructuredLogMigrationTests.cs

@@ -1,3 +1,4 @@
+using Elsa.Common;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.Options;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.Services;
 using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Services;
@@ -29,6 +30,29 @@ public class SqliteStructuredLogMigrationTests
         Assert.False(await host.TableExistsAsync("StructuredLogEvents"));
     }
 
+    [Fact]
+    public async Task StartupTask_RunsMigrations_WhenHostedServicesAreNotStarted()
+    {
+        await using var host = new SqliteStructuredLogTestHost(migrate: false);
+        using var scope = host.Services.CreateScope();
+        var startup = scope.ServiceProvider.GetServices<IStartupTask>().OfType<SqliteStructuredLogStartupService>().Single();
+
+        await startup.ExecuteAsync(CancellationToken.None);
+
+        Assert.True(await host.TableExistsAsync("StructuredLogEvents"));
+    }
+
+    [Fact]
+    public async Task StartupTask_UsesSameInstance_AsHostedService()
+    {
+        await using var host = new SqliteStructuredLogTestHost(migrate: false);
+        using var scope = host.Services.CreateScope();
+        var hostedService = host.Services.GetServices<IHostedService>().OfType<SqliteStructuredLogStartupService>().Single();
+        var startupTask = scope.ServiceProvider.GetServices<IStartupTask>().OfType<SqliteStructuredLogStartupService>().Single();
+
+        Assert.Same(hostedService, startupTask);
+    }
+
     [Fact]
     public async Task HostedServices_StartMigrationBeforeWriteBuffer()
     {

test/integration/Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.IntegrationTests/SqliteStructuredLogWriteQueueTests.cs

@@ -1,4 +1,8 @@
+using Elsa.Common;
 using Elsa.Diagnostics.StructuredLogs.Models;
+using Elsa.Diagnostics.StructuredLogs.Persistence.Relational.Services;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 
 namespace Elsa.Diagnostics.StructuredLogs.Persistence.Sqlite.IntegrationTests;
 
@@ -43,4 +47,96 @@ public class SqliteStructuredLogWriteQueueTests
 
         Assert.Equal(1, await host.CountRowsAsync("StructuredLogEvents"));
     }
+
+    [Fact]
+    public async Task BackgroundTask_FlushesLoggerWrites_ForShellLifecycle()
+    {
+        await using var host = new SqliteStructuredLogTestHost(migrate: false);
+        using var scope = host.Services.CreateScope();
+        var services = scope.ServiceProvider;
+
+        foreach (var startupTask in services.GetServices<IStartupTask>())
+            await startupTask.ExecuteAsync(CancellationToken.None);
+
+        var backgroundTask = services.GetServices<IBackgroundTask>().OfType<StructuredLogWriteBufferBackgroundTask>().Single();
+        await backgroundTask.StartAsync(CancellationToken.None);
+        await backgroundTask.StopAsync(CancellationToken.None);
+        await backgroundTask.StartAsync(CancellationToken.None);
+
+        try
+        {
+            var logger = services.GetRequiredService<ILoggerFactory>().CreateLogger("Elsa.Tests.ShellLifecycle");
+            logger.LogInformation("SQLite structured log shell background task test");
+
+            await WaitUntilAsync(async () => await host.CountRowsAsync("StructuredLogEvents") >= 1);
+        }
+        finally
+        {
+            await backgroundTask.StopAsync(CancellationToken.None);
+        }
+    }
+
+    [Fact]
+    public async Task BackgroundTask_Stop_DoesNotStopHostedWriteBuffer()
+    {
+        await using var host = new SqliteStructuredLogTestHost(migrate: false);
+        using var scope = host.Services.CreateScope();
+        var services = scope.ServiceProvider;
+
+        await host.StartHostedServicesAsync();
+        var backgroundTask = services.GetServices<IBackgroundTask>().OfType<StructuredLogWriteBufferBackgroundTask>().Single();
+        await backgroundTask.StartAsync(CancellationToken.None);
+        await backgroundTask.StopAsync(CancellationToken.None);
+
+        try
+        {
+            var logger = services.GetRequiredService<ILoggerFactory>().CreateLogger("Elsa.Tests.HostAndShellLifecycle");
+            logger.LogInformation("SQLite structured log hosted write buffer test");
+
+            await WaitUntilAsync(async () => await host.CountRowsAsync("StructuredLogEvents") >= 1);
+        }
+        finally
+        {
+            await host.StopHostedServicesAsync();
+        }
+    }
+
+    [Fact]
+    public async Task BackgroundTask_StopWithoutStart_DoesNotStopHostedWriteBuffer()
+    {
+        await using var host = new SqliteStructuredLogTestHost(migrate: false);
+        using var scope = host.Services.CreateScope();
+        var services = scope.ServiceProvider;
+
+        await host.StartHostedServicesAsync();
+        var backgroundTask = services.GetServices<IBackgroundTask>().OfType<StructuredLogWriteBufferBackgroundTask>().Single();
+        await backgroundTask.StopAsync(CancellationToken.None);
+
+        try
+        {
+            var logger = services.GetRequiredService<ILoggerFactory>().CreateLogger("Elsa.Tests.PartialShellLifecycle");
+            logger.LogInformation("SQLite structured log partial shell lifecycle test");
+
+            await WaitUntilAsync(async () => await host.CountRowsAsync("StructuredLogEvents") >= 1);
+        }
+        finally
+        {
+            await host.StopHostedServicesAsync();
+        }
+    }
+
+    private static async Task WaitUntilAsync(Func<ValueTask<bool>> condition)
+    {
+        using var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+
+        while (!timeoutTokenSource.IsCancellationRequested)
+        {
+            if (await condition())
+                return;
+
+            await Task.Delay(50);
+        }
+
+        Assert.Fail("The expected structured log row was not persisted before the timeout elapsed.");
+    }
 }