Skip to content

Commit 29fb8cf

Browse files
committed
refactor: simplify console output handling and remove unnecessary state management
1 parent c2711e2 commit 29fb8cf

File tree

2 files changed

+258
-152
lines changed

2 files changed

+258
-152
lines changed

libraries/src/AWS.Lambda.Powertools.Common/Core/ConsoleWrapper.cs

Lines changed: 53 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ public class ConsoleWrapper : IConsoleWrapper
88
{
99
private static bool _override;
1010
private static TextWriter _testOutputStream;
11-
private static bool _outputResetPerformed = false;
1211
private static bool _inTestMode = false;
1312

1413
/// <inheritdoc />
@@ -20,7 +19,7 @@ public void WriteLine(string message)
2019
}
2120
else
2221
{
23-
EnsureConsoleOutputOnce();
22+
EnsureConsoleOutput();
2423
Console.WriteLine(message);
2524
}
2625
}
@@ -34,7 +33,7 @@ public void Debug(string message)
3433
}
3534
else
3635
{
37-
EnsureConsoleOutputOnce();
36+
EnsureConsoleOutput();
3837
System.Diagnostics.Debug.WriteLine(message);
3938
}
4039
}
@@ -50,9 +49,9 @@ public void Error(string message)
5049
{
5150
if (!_override)
5251
{
53-
var errordOutput = new StreamWriter(Console.OpenStandardError());
54-
errordOutput.AutoFlush = true;
55-
Console.SetError(errordOutput);
52+
var errorOutput = new StreamWriter(Console.OpenStandardError());
53+
errorOutput.AutoFlush = true;
54+
Console.SetError(errorOutput);
5655
}
5756
Console.Error.WriteLine(message);
5857
}
@@ -70,23 +69,59 @@ public static void SetOut(TextWriter consoleOut)
7069
Console.SetOut(consoleOut);
7170
}
7271

73-
private static void EnsureConsoleOutputOnce()
72+
private static void EnsureConsoleOutput()
7473
{
75-
if (_outputResetPerformed) return;
76-
OverrideLambdaLogger();
77-
_outputResetPerformed = true;
74+
// Check if we need to override console output for Lambda environment
75+
if (ShouldOverrideConsole())
76+
{
77+
OverrideLambdaLogger();
78+
}
79+
}
80+
81+
private static bool ShouldOverrideConsole()
82+
{
83+
// Don't override if we're in test mode
84+
if (_inTestMode) return false;
85+
86+
// Always override in Lambda environment to prevent Lambda's log wrapping
87+
var isLambda = !string.IsNullOrEmpty(Environment.GetEnvironmentVariable("AWS_LAMBDA_FUNCTION_NAME"));
88+
89+
return isLambda && (!_override || HasLambdaReInterceptedConsole());
90+
}
91+
92+
private static bool HasLambdaReInterceptedConsole()
93+
{
94+
// Lambda might re-intercept console between init and handler execution
95+
try
96+
{
97+
var currentOut = Console.Out;
98+
// Check if current output stream looks like it might be Lambda's wrapper
99+
var typeName = currentOut.GetType().FullName ?? "";
100+
return typeName.Contains("Lambda") || typeName == "System.IO.TextWriter+SyncTextWriter";
101+
}
102+
catch
103+
{
104+
return true; // Assume re-interception if we can't determine
105+
}
78106
}
79107

80108
private static void OverrideLambdaLogger()
81109
{
82-
if (_override)
110+
try
111+
{
112+
// Force override of LambdaLogger
113+
var standardOutput = new StreamWriter(Console.OpenStandardOutput())
114+
{
115+
AutoFlush = true
116+
};
117+
Console.SetOut(standardOutput);
118+
_override = true;
119+
}
120+
catch (Exception)
83121
{
84-
return;
122+
// Log the failure but don't throw - degraded functionality is better than crash
123+
_override = false;
85124
}
86-
// Force override of LambdaLogger
87-
var standardOutput = new StreamWriter(Console.OpenStandardOutput());
88-
standardOutput.AutoFlush = true;
89-
Console.SetOut(standardOutput);
90125
}
91126

92127
internal static void WriteLine(string logLevel, string message)
@@ -102,14 +137,14 @@ public static void ResetForTest()
102137
_override = false;
103138
_inTestMode = false;
104139
_testOutputStream = null;
105-
_outputResetPerformed = false;
106140
}
107141

108142
/// <summary>
109143
/// Clear the output reset flag
110144
/// </summary>
111145
public static void ClearOutputResetFlag()
112146
{
113-
_outputResetPerformed = false;
147+
// This method is kept for backward compatibility but no longer needed
148+
// since we removed the _outputResetPerformed flag
114149
}
115150
}

0 commit comments

Comments
 (0)