diff --git a/src/Logsmith.Generator/Emission/MethodEmitter.cs b/src/Logsmith.Generator/Emission/MethodEmitter.cs index 60e4252..bba95a9 100644 --- a/src/Logsmith.Generator/Emission/MethodEmitter.cs +++ b/src/Logsmith.Generator/Emission/MethodEmitter.cs @@ -206,6 +206,8 @@ internal static string EmitMethodBody(LogMethodInfo method) sb.AppendLine($" global::Logsmith.LogManager.Dispatch(in __entry, __utf8Message, __state, WriteProperties_{method.MethodName});"); } + // Return any ArrayPool buffer rented during overflow + sb.AppendLine(" writer.Dispose();"); sb.AppendLine(" }"); return sb.ToString(); } diff --git a/src/Logsmith/Utf8LogWriter.cs b/src/Logsmith/Utf8LogWriter.cs index ef149fc..15c6db5 100644 --- a/src/Logsmith/Utf8LogWriter.cs +++ b/src/Logsmith/Utf8LogWriter.cs @@ -1,4 +1,5 @@ using System.Buffers; +using System.Runtime.CompilerServices; using System.Text; using System.Text.Unicode; @@ -6,19 +7,31 @@ namespace Logsmith; public ref struct Utf8LogWriter { - private readonly Span _buffer; + private Span _buffer; private int _position; + private byte[]? _rented; public Utf8LogWriter(Span buffer) { _buffer = buffer; _position = 0; + _rented = null; + } + + public void Dispose() + { + var rented = _rented; + if (rented is not null) + { + _rented = null; + ArrayPool.Shared.Return(rented); + } } public void Write(ReadOnlySpan utf8Literal) { if (utf8Literal.Length > _buffer.Length - _position) - return; + Grow(utf8Literal.Length); utf8Literal.CopyTo(_buffer[_position..]); _position += utf8Literal.Length; @@ -26,18 +39,18 @@ public void Write(ReadOnlySpan utf8Literal) public void WriteFormatted(in T value) where T : IUtf8SpanFormattable { - if (value.TryFormat(_buffer[_position..], out int bytesWritten, default, null)) - { - _position += bytesWritten; - } + int bytesWritten; + while (!value.TryFormat(_buffer[_position..], out bytesWritten, default, null)) + Grow(_buffer.Length); + _position += bytesWritten; } public void WriteFormatted(in T value, ReadOnlySpan format) where T : IUtf8SpanFormattable { - if (value.TryFormat(_buffer[_position..], out int bytesWritten, format, null)) - { - _position += bytesWritten; - } + int bytesWritten; + while (!value.TryFormat(_buffer[_position..], out bytesWritten, format, null)) + Grow(_buffer.Length); + _position += bytesWritten; } public void WriteString(string? value) @@ -50,6 +63,15 @@ public void WriteString(string? value) var status = Utf8.FromUtf16(value, _buffer[_position..], out _, out int bytesWritten); if (status == OperationStatus.Done) + { + _position += bytesWritten; + return; + } + + // Buffer too small — grow to exact requirement and retry + Grow(Encoding.UTF8.GetByteCount(value)); + status = Utf8.FromUtf16(value, _buffer[_position..], out _, out bytesWritten); + if (status == OperationStatus.Done) { _position += bytesWritten; } @@ -58,4 +80,20 @@ public void WriteString(string? value) public ReadOnlySpan GetWritten() => _buffer[.._position]; public int BytesWritten => _position; + + [MethodImpl(MethodImplOptions.NoInlining)] + private void Grow(int needed) + { + int required = _position + needed; + int newSize = Math.Max(_buffer.Length * 2, required); + var newArray = ArrayPool.Shared.Rent(newSize); + _buffer[.._position].CopyTo(newArray); + + var old = _rented; + _rented = newArray; + _buffer = newArray; + + if (old is not null) + ArrayPool.Shared.Return(old); + } } diff --git a/tests/Logsmith.Tests/Utf8LogWriterTests.cs b/tests/Logsmith.Tests/Utf8LogWriterTests.cs index 94b0be0..6dac6e7 100644 --- a/tests/Logsmith.Tests/Utf8LogWriterTests.cs +++ b/tests/Logsmith.Tests/Utf8LogWriterTests.cs @@ -15,6 +15,7 @@ public void WriteFormatted_WithFormat_AppliesFormatString() var result = Encoding.UTF8.GetString(writer.GetWritten()); Assert.That(result, Is.EqualTo("3.14")); + writer.Dispose(); } [Test] @@ -27,5 +28,97 @@ public void WriteFormatted_WithEmptyFormat_SameAsDefault() var result = Encoding.UTF8.GetString(writer.GetWritten()); Assert.That(result, Is.EqualTo("42")); + writer.Dispose(); + } + + [Test] + public void WriteString_WithCurlyBraces_PreservesBraces() + { + Span buffer = stackalloc byte[512]; + var writer = new Utf8LogWriter(buffer); + + writer.Write("Vulkan validation: "u8); + writer.WriteString("{VkBuffer 0x00000001234ABCDEF} - Memory type not suitable"); + + var result = Encoding.UTF8.GetString(writer.GetWritten()); + Assert.That(result, Is.EqualTo("Vulkan validation: {VkBuffer 0x00000001234ABCDEF} - Memory type not suitable")); + writer.Dispose(); + } + + [Test] + public void WriteString_OverflowBuffer_FallsBackToArrayPool() + { + // Buffer only fits the literal prefix, not the string param + Span buffer = stackalloc byte[30]; + var writer = new Utf8LogWriter(buffer); + + writer.Write("Prefix: "u8); // 8 bytes, leaves 22 bytes + writer.WriteString("{VkBuffer 0x00000001234ABCDEF} - Memory type not suitable"); // 57 bytes — overflows + + var result = Encoding.UTF8.GetString(writer.GetWritten()); + Assert.That(result, Is.EqualTo("Prefix: {VkBuffer 0x00000001234ABCDEF} - Memory type not suitable")); + writer.Dispose(); + } + + [Test] + public void WriteString_LongMessage_PreservedViaArrayPoolFallback() + { + // Max generated buffer is 4096 bytes. A single string param can easily exceed that — + // stack traces, serialized payloads, validation messages, SQL queries, etc. + const int maxGeneratedBuffer = 4096; + Span buffer = stackalloc byte[maxGeneratedBuffer]; + var writer = new Utf8LogWriter(buffer); + + writer.Write("Payload: "u8); // 9 bytes + + // 8KB string — double the max buffer. Not unusual for serialized objects, + // full stack traces, or verbose diagnostic output. + var longMessage = new string('X', 8192); + + writer.WriteString(longMessage); + + var result = Encoding.UTF8.GetString(writer.GetWritten()); + Assert.That(result, Is.EqualTo("Payload: " + longMessage)); + Assert.That(writer.BytesWritten, Is.EqualTo(9 + 8192)); + writer.Dispose(); + } + + [Test] + public void Write_LiteralOverflow_PreservedViaArrayPoolFallback() + { + Span buffer = stackalloc byte[4]; + var writer = new Utf8LogWriter(buffer); + + writer.Write("Hello, world!"u8); // 13 bytes into 4-byte buffer + + var result = Encoding.UTF8.GetString(writer.GetWritten()); + Assert.That(result, Is.EqualTo("Hello, world!")); + writer.Dispose(); + } + + [Test] + public void WriteFormatted_Overflow_PreservedViaArrayPoolFallback() + { + Span buffer = stackalloc byte[4]; + var writer = new Utf8LogWriter(buffer); + + writer.WriteFormatted(123456789L); // needs ~9 bytes, only 4 available + + var result = Encoding.UTF8.GetString(writer.GetWritten()); + Assert.That(result, Is.EqualTo("123456789")); + writer.Dispose(); + } + + [Test] + public void Dispose_NoOverflow_NoOp() + { + Span buffer = stackalloc byte[128]; + var writer = new Utf8LogWriter(buffer); + + writer.Write("hello"u8); + + // Dispose on stackalloc-only path should be a no-op (no rented buffer to return) + writer.Dispose(); + Assert.Pass(); } }