Merge branch 'normj/response-streaming' into normj/aspnetcore-responsestreaming

Author: normj

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStream.cs

@@ -14,6 +14,7 @@
  */
 
 using System;
+using System.Buffers;
 using System.IO;
 using System.Text;
 using System.Threading;
@@ -76,34 +77,37 @@ internal ResponseStream(byte[] prelude)
         internal async Task SetHttpOutputStreamAsync(Stream httpOutputStream, CancellationToken cancellationToken = default)
         {
             _httpOutputStream = httpOutputStream;
-            _httpStreamReady.Release();
-
-            await WritePreludeAsync(cancellationToken);
-        }
 
-        private async Task WritePreludeAsync(CancellationToken cancellationToken = default)
-        {
+            // Write the prelude BEFORE releasing _httpStreamReady. This prevents a race
+            // where a handler WriteAsync that is already waiting on the semaphore could
+            // sneak in and write body data before the prelude, causing intermittent
+            // "Failed to parse prelude JSON" errors from API Gateway.
             if (_prelude?.Length > 0)
             {
                 _logger.LogDebug($"Writing prelude of {_prelude.Length} bytes to HTTP stream.");
-                await _httpStreamReady.WaitAsync(_httpStreamWaitTimeout, cancellationToken);
-                try
+
+                lock (_lock)
                 {
-                    lock (_lock)
-                    {
-                        ThrowIfCompletedOrError();
-                    }
+                    ThrowIfCompletedOrError();
+                }
 
-                    await _httpOutputStream.WriteAsync(_prelude, 0, _prelude.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(PreludeDelimiter, 0, PreludeDelimiter.Length, cancellationToken);
+                var combinedLength = _prelude.Length + PreludeDelimiter.Length;
+                var combined = ArrayPool<byte>.Shared.Rent(combinedLength);
+                try
+                {
+                    Buffer.BlockCopy(_prelude, 0, combined, 0, _prelude.Length);
+                    Buffer.BlockCopy(PreludeDelimiter, 0, combined, _prelude.Length, PreludeDelimiter.Length);
 
+                    await _httpOutputStream.WriteAsync(combined, 0, combinedLength, cancellationToken);
                     await _httpOutputStream.FlushAsync(cancellationToken);
                 }
                 finally
                 {
-                    _httpStreamReady.Release();
+                    ArrayPool<byte>.Shared.Return(combined);
                 }
             }
+
+            _httpStreamReady.Release();
         }
 
         /// <summary>

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/ResponseStreamTests.cs

@@ -284,6 +284,107 @@ public async Task SetHttpOutputStreamAsync_NullPrelude_WritesNoPreludeBytes()
             Assert.Empty(output.ToArray());
         }
 
+        // ---- Prelude + delimiter single-chunk tests (via ChunkedStreamWriter) ----
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_ProducesSingleChunk()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":200}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+
+            var wireBytes = Encoding.ASCII.GetString(rawOutput.ToArray());
+
+            // The prelude (18 bytes) + delimiter (8 bytes) = 26 bytes = 0x1A
+            // Should be exactly one chunk: "1A\r\n{prelude}{8 null bytes}\r\n"
+            var expectedDataLength = preludeJson.Length + 8; // 26
+            var expectedHex = expectedDataLength.ToString("X");
+            Assert.StartsWith($"{expectedHex}\r\n", wireBytes);
+
+            // Verify there is only one chunk header (only one hex size prefix)
+            var chunkCount = 0;
+            var remaining = wireBytes;
+            while (remaining.Length > 0)
+            {
+                var crlfIndex = remaining.IndexOf("\r\n", StringComparison.Ordinal);
+                if (crlfIndex < 0) break;
+                var sizeStr = remaining.Substring(0, crlfIndex);
+                if (int.TryParse(sizeStr, System.Globalization.NumberStyles.HexNumber, null, out var chunkSize) && chunkSize >= 0)
+                {
+                    chunkCount++;
+                    // Skip past: hex\r\n{data}\r\n
+                    remaining = remaining.Substring(crlfIndex + 2 + chunkSize + 2);
+                }
+                else
+                {
+                    break;
+                }
+            }
+            Assert.Equal(1, chunkCount);
+        }
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_DelimiterImmediatelyFollowsPrelude()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":201}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+
+            // Parse the chunk to get the raw data payload
+            var wireBytes = rawOutput.ToArray();
+            var wireStr = Encoding.ASCII.GetString(wireBytes);
+            var firstCrlf = wireStr.IndexOf("\r\n", StringComparison.Ordinal);
+            var dataStart = firstCrlf + 2;
+            var dataLength = preludeJson.Length + 8;
+            var chunkData = new byte[dataLength];
+            Array.Copy(wireBytes, dataStart, chunkData, 0, dataLength);
+
+            // First part should be the prelude JSON
+            Assert.Equal(preludeJson, chunkData[..preludeJson.Length]);
+            // Immediately followed by 8 null bytes (delimiter)
+            Assert.Equal(new byte[8], chunkData[preludeJson.Length..]);
+        }
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_HandlerDataInSeparateChunk()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":200}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+            await rs.WriteAsync(Encoding.UTF8.GetBytes("body data"), 0, 9);
+
+            var wireStr = Encoding.ASCII.GetString(rawOutput.ToArray());
+
+            // Should have exactly 2 chunks: one for prelude+delimiter, one for body
+            var chunkCount = 0;
+            var remaining = wireStr;
+            while (remaining.Length > 0)
+            {
+                var crlfIndex = remaining.IndexOf("\r\n", StringComparison.Ordinal);
+                if (crlfIndex < 0) break;
+                var sizeStr = remaining.Substring(0, crlfIndex);
+                if (int.TryParse(sizeStr, System.Globalization.NumberStyles.HexNumber, null, out var chunkSize) && chunkSize >= 0)
+                {
+                    chunkCount++;
+                    remaining = remaining.Substring(crlfIndex + 2 + chunkSize + 2);
+                }
+                else
+                {
+                    break;
+                }
+            }
+            Assert.Equal(2, chunkCount);
+        }
+
         // ---- MarkCompleted idempotency ----
 
         [Fact]