From faa637d792aabf8a872bd7a67d09889223fbee60 Mon Sep 17 00:00:00 2001
From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com>
Date: Sat, 27 Jul 2024 23:54:03 +0100
Subject: [PATCH] Use AsyncLocal to find the correct ModuleLogger

---
 src/ModularPipelines.Build/ReleaseNotes.md           | 5 +----
 src/ModularPipelines/Engine/ModuleExecutor.cs        | 7 +++++--
 src/ModularPipelines/Logging/ModuleLogger.cs         | 7 +++----
 src/ModularPipelines/Logging/ModuleLoggerProvider.cs | 5 +++++
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/ModularPipelines.Build/ReleaseNotes.md b/src/ModularPipelines.Build/ReleaseNotes.md
index 2216e687f6..b838930f3a 100644
--- a/src/ModularPipelines.Build/ReleaseNotes.md
+++ b/src/ModularPipelines.Build/ReleaseNotes.md
@@ -1,4 +1 @@
-- Show exception (if one was thrown) within the GitHub markdown summary
-
-# Breaking
-- GitVersion.yml has been updated to v6. If using the `GitVersionInformation` model, some properties may have been added/removed/changed. 
\ No newline at end of file
+- Use `AsyncLocal` to find the correct `IModuleLogger` when using the `IModuleLoggerProvider` helper.
\ No newline at end of file
diff --git a/src/ModularPipelines/Engine/ModuleExecutor.cs b/src/ModularPipelines/Engine/ModuleExecutor.cs
index 82ec9308a0..c7e5c87a79 100644
--- a/src/ModularPipelines/Engine/ModuleExecutor.cs
+++ b/src/ModularPipelines/Engine/ModuleExecutor.cs
@@ -6,6 +6,7 @@
 using ModularPipelines.Attributes;
 using ModularPipelines.Exceptions;
 using ModularPipelines.Extensions;
+using ModularPipelines.Logging;
 using ModularPipelines.Models;
 using ModularPipelines.Modules;
 using ModularPipelines.Options;
@@ -150,7 +151,7 @@ private Task<ModuleBase> StartModule(ModuleBase module)
     {
         lock (_moduleDictionaryLock)
         {
-            return _moduleExecutionTasks.GetOrAdd(module, async @base =>
+            return _moduleExecutionTasks.GetOrAdd(module, @base => Task.Run(async () =>
             {
                 _logger.LogDebug("Starting Module {Module}", module.GetType().Name);
 
@@ -163,6 +164,8 @@ private Task<ModuleBase> StartModule(ModuleBase module)
 
                 try
                 {
+                    ModuleLogger.Values.Value = module.Context.Logger;
+                    
                     await _pipelineSetupExecutor.OnBeforeModuleStartAsync(module);
 
                     await module.StartInternal();
@@ -180,7 +183,7 @@ private Task<ModuleBase> StartModule(ModuleBase module)
                         await _moduleDisposer.DisposeAsync(module);
                     }
                 }
-            });
+            }));
         }
     }
 
diff --git a/src/ModularPipelines/Logging/ModuleLogger.cs b/src/ModularPipelines/Logging/ModuleLogger.cs
index d8c49e60b3..75b3eaf939 100644
--- a/src/ModularPipelines/Logging/ModuleLogger.cs
+++ b/src/ModularPipelines/Logging/ModuleLogger.cs
@@ -7,10 +7,12 @@ namespace ModularPipelines.Logging;
 
 internal abstract class ModuleLogger : IModuleLogger
 {
+    internal static readonly AsyncLocal<IModuleLogger> Values = new();
+
     protected static readonly object DisposeLock = new();
     protected static readonly object LogLock = new();
     protected Exception? _exception;
-
+    
     internal DateTime LastLogWritten { get; set; } = DateTime.MinValue;
 
     public abstract void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter);
@@ -34,7 +36,6 @@ internal class ModuleLogger<T> : ModuleLogger, IModuleLogger, ILogger<T>
 {
     private readonly ILogger<T> _defaultLogger;
     private readonly ISecretObfuscator _secretObfuscator;
-    private readonly ISecretProvider _secretProvider;
     private readonly IConsoleWriter _consoleWriter;
     private readonly ISmartCollapsableLoggingStringBlockProvider _collapsableLoggingStringBlockProvider;
 
@@ -46,13 +47,11 @@ internal class ModuleLogger<T> : ModuleLogger, IModuleLogger, ILogger<T>
     public ModuleLogger(ILogger<T> defaultLogger,
         IModuleLoggerContainer moduleLoggerContainer,
         ISecretObfuscator secretObfuscator,
-        ISecretProvider secretProvider,
         IConsoleWriter consoleWriter,
         ISmartCollapsableLoggingStringBlockProvider collapsableLoggingStringBlockProvider)
     {
         _defaultLogger = defaultLogger;
         _secretObfuscator = secretObfuscator;
-        _secretProvider = secretProvider;
         _consoleWriter = consoleWriter;
         _collapsableLoggingStringBlockProvider = collapsableLoggingStringBlockProvider;
         moduleLoggerContainer.AddLogger(this);
diff --git a/src/ModularPipelines/Logging/ModuleLoggerProvider.cs b/src/ModularPipelines/Logging/ModuleLoggerProvider.cs
index 82ce90cfdd..2d0166721c 100644
--- a/src/ModularPipelines/Logging/ModuleLoggerProvider.cs
+++ b/src/ModularPipelines/Logging/ModuleLoggerProvider.cs
@@ -30,6 +30,11 @@ public IModuleLogger GetLogger()
         {
             return _moduleLogger;
         }
+
+        if (ModuleLogger.Values.Value != null)
+        {
+            return _moduleLogger = ModuleLogger.Values.Value;
+        }
         
         var stackFrames = new StackTrace().GetFrames().ToList();