Task 4 (after redesign)

Author: normj

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/ResponseStream.cs

@@ -14,62 +14,70 @@
  */
 
 using System;
-using System.Collections.Generic;
-using System.Linq;
+using System.IO;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Amazon.Lambda.RuntimeSupport
 {
     /// <summary>
-    /// Internal implementation of IResponseStream.
-    /// Buffers written data as chunks for HTTP chunked transfer encoding.
+    /// Internal implementation of IResponseStream with true streaming.
+    /// Writes data directly to the HTTP output stream as chunked transfer encoding.
     /// </summary>
     internal class ResponseStream : IResponseStream
     {
+        private static readonly byte[] CrlfBytes = Encoding.ASCII.GetBytes("\r\n");
+
         private readonly long _maxResponseSize;
-        private readonly List<byte[]> _chunks;
         private long _bytesWritten;
         private bool _isCompleted;
         private bool _hasError;
         private Exception _reportedError;
         private readonly object _lock = new object();
 
+        // The live HTTP output stream, set by StreamingHttpContent when SerializeToStreamAsync is called.
+        private Stream _httpOutputStream;
+        private readonly SemaphoreSlim _httpStreamReady = new SemaphoreSlim(0, 1);
+        private readonly SemaphoreSlim _completionSignal = new SemaphoreSlim(0, 1);
+
         public long BytesWritten => _bytesWritten;
         public bool IsCompleted => _isCompleted;
         public bool HasError => _hasError;
+        internal Exception ReportedError => _reportedError;
 
-        internal IReadOnlyList<byte[]> Chunks
+        public ResponseStream(long maxResponseSize)
         {
-            get
-            {
-                lock (_lock)
-                {
-                    return _chunks.ToList();
-                }
-            }
+            _maxResponseSize = maxResponseSize;
         }
 
-        internal Exception ReportedError => _reportedError;
+        /// <summary>
+        /// Called by StreamingHttpContent.SerializeToStreamAsync to provide the HTTP output stream.
+        /// </summary>
+        internal void SetHttpOutputStream(Stream httpOutputStream)
+        {
+            _httpOutputStream = httpOutputStream;
+            _httpStreamReady.Release();
+        }
 
-        public ResponseStream(long maxResponseSize)
+        /// <summary>
+        /// Called by StreamingHttpContent.SerializeToStreamAsync to wait until the handler
+        /// finishes writing (MarkCompleted or ReportErrorAsync).
+        /// </summary>
+        internal async Task WaitForCompletionAsync()
         {
-            _maxResponseSize = maxResponseSize;
-            _chunks = new List<byte[]>();
-            _bytesWritten = 0;
-            _isCompleted = false;
-            _hasError = false;
+            await _completionSignal.WaitAsync();
         }
 
-        public Task WriteAsync(byte[] buffer, CancellationToken cancellationToken = default)
+        public async Task WriteAsync(byte[] buffer, CancellationToken cancellationToken = default)
         {
             if (buffer == null)
                 throw new ArgumentNullException(nameof(buffer));
 
-            return WriteAsync(buffer, 0, buffer.Length, cancellationToken);
+            await WriteAsync(buffer, 0, buffer.Length, cancellationToken);
         }
 
-        public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
+        public async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
         {
             if (buffer == null)
                 throw new ArgumentNullException(nameof(buffer));
@@ -78,45 +86,45 @@ public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken c
             if (count < 0 || offset + count > buffer.Length)
                 throw new ArgumentOutOfRangeException(nameof(count));
 
-            lock (_lock)
+            // Wait for the HTTP stream to be ready (first write only blocks)
+            await _httpStreamReady.WaitAsync(cancellationToken);
+            try
             {
-                ThrowIfCompletedOrError();
-
-                if (_bytesWritten + count > _maxResponseSize)
+                lock (_lock)
                 {
-                    throw new InvalidOperationException(
-                        $"Writing {count} bytes would exceed the maximum response size of {_maxResponseSize} bytes (20 MiB). " +
-                        $"Current size: {_bytesWritten} bytes.");
-                }
-
-                var chunk = new byte[count];
-                Array.Copy(buffer, offset, chunk, 0, count);
-                _chunks.Add(chunk);
-                _bytesWritten += count;
-            }
+                    ThrowIfCompletedOrError();
 
-            return Task.CompletedTask;
-        }
-
-        public Task WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
-        {
-            lock (_lock)
-            {
-                ThrowIfCompletedOrError();
+                    if (_bytesWritten + count > _maxResponseSize)
+                    {
+                        throw new InvalidOperationException(
+                            $"Writing {count} bytes would exceed the maximum response size of {_maxResponseSize} bytes (20 MiB). " +
+                            $"Current size: {_bytesWritten} bytes.");
+                    }
 
-                if (_bytesWritten + buffer.Length > _maxResponseSize)
-                {
-                    throw new InvalidOperationException(
-                        $"Writing {buffer.Length} bytes would exceed the maximum response size of {_maxResponseSize} bytes (20 MiB). " +
-                        $"Current size: {_bytesWritten} bytes.");
+                    _bytesWritten += count;
                 }
 
-                var chunk = buffer.ToArray();
-                _chunks.Add(chunk);
-                _bytesWritten += buffer.Length;
+                // Write chunk directly to the HTTP stream: size(hex) + CRLF + data + CRLF
+                var chunkSizeHex = count.ToString("X");
+                var chunkSizeBytes = Encoding.ASCII.GetBytes(chunkSizeHex);
+                await _httpOutputStream.WriteAsync(chunkSizeBytes, 0, chunkSizeBytes.Length, cancellationToken);
+                await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
+                await _httpOutputStream.WriteAsync(buffer, offset, count, cancellationToken);
+                await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
+                await _httpOutputStream.FlushAsync(cancellationToken);
+            }
+            finally
+            {
+                // Re-release so subsequent writes don't block
+                _httpStreamReady.Release();
             }
+        }
 
-            return Task.CompletedTask;
+        public async Task WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            // Convert to array and delegate — small overhead but keeps the API simple
+            var array = buffer.ToArray();
+            await WriteAsync(array, 0, array.Length, cancellationToken);
         }
 
         public Task ReportErrorAsync(Exception exception, CancellationToken cancellationToken = default)
@@ -135,6 +143,8 @@ public Task ReportErrorAsync(Exception exception, CancellationToken cancellation
                 _reportedError = exception;
             }
 
+            // Signal completion so StreamingHttpContent can write error trailers and finish
+            _completionSignal.Release();
             return Task.CompletedTask;
         }
 
@@ -144,6 +154,8 @@ internal void MarkCompleted()
             {
                 _isCompleted = true;
             }
+            // Signal completion so StreamingHttpContent can write the final chunk and finish
+            _completionSignal.Release();
         }
 
         private void ThrowIfCompletedOrError()
@@ -156,7 +168,8 @@ private void ThrowIfCompletedOrError()
 
         public void Dispose()
         {
-            // Nothing to dispose - all data is in managed memory
+            // Ensure completion is signaled if not already
+            try { _completionSignal.Release(); } catch (SemaphoreFullException) { }
         }
     }
 }

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/ResponseStreamContext.cs

@@ -13,6 +13,9 @@
  * permissions and limitations under the License.
  */
 
+using System.Threading;
+using System.Threading.Tasks;
+
 namespace Amazon.Lambda.RuntimeSupport
 {
     /// <summary>
@@ -39,5 +42,21 @@ internal class ResponseStreamContext
         /// The ResponseStream instance if created.
         /// </summary>
         public ResponseStream Stream { get; set; }
+
+        /// <summary>
+        /// The RuntimeApiClient used to start the streaming HTTP POST.
+        /// </summary>
+        public RuntimeApiClient RuntimeApiClient { get; set; }
+
+        /// <summary>
+        /// Cancellation token for the current invocation.
+        /// </summary>
+        public CancellationToken CancellationToken { get; set; }
+
+        /// <summary>
+        /// The Task representing the in-flight HTTP POST to the Runtime API.
+        /// Started when CreateStream() is called, completes when the stream is finalized.
+        /// </summary>
+        public Task SendTask { get; set; }
     }
 }

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/StreamingHttpContent.cs

@@ -39,18 +39,24 @@ public StreamingHttpContent(ResponseStream responseStream)
 
         protected override async Task SerializeToStreamAsync(Stream stream, TransportContext context)
         {
-            foreach (var chunk in _responseStream.Chunks)
-            {
-                await WriteChunkAsync(stream, chunk);
-            }
+            // Hand the HTTP output stream to ResponseStream so WriteAsync calls
+            // can write chunks directly to it.
+            _responseStream.SetHttpOutputStream(stream);
 
-            await WriteFinalChunkAsync(stream);
+            // Wait for the handler to finish writing (MarkCompleted or ReportErrorAsync)
+            await _responseStream.WaitForCompletionAsync();
 
+            // Write final chunk
+            await stream.WriteAsync(FinalChunkBytes, 0, FinalChunkBytes.Length);
+
+            // Write error trailers if present
             if (_responseStream.HasError)
             {
                 await WriteErrorTrailersAsync(stream, _responseStream.ReportedError);
             }
 
+            // Write final CRLF to end the chunked message
+            await stream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length);
             await stream.FlushAsync();
         }
 
@@ -60,22 +66,6 @@ protected override bool TryComputeLength(out long length)
             return false;
         }
 
-        private async Task WriteChunkAsync(Stream stream, byte[] data)
-        {
-            var chunkSizeHex = data.Length.ToString("X");
-            var chunkSizeBytes = Encoding.ASCII.GetBytes(chunkSizeHex);
-
-            await stream.WriteAsync(chunkSizeBytes, 0, chunkSizeBytes.Length);
-            await stream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length);
-            await stream.WriteAsync(data, 0, data.Length);
-            await stream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length);
-        }
-
-        private async Task WriteFinalChunkAsync(Stream stream)
-        {
-            await stream.WriteAsync(FinalChunkBytes, 0, FinalChunkBytes.Length);
-        }
-
         private async Task WriteErrorTrailersAsync(Stream stream, Exception exception)
         {
             var exceptionInfo = ExceptionInfo.GetExceptionInfo(exception);
@@ -88,8 +78,6 @@ private async Task WriteErrorTrailersAsync(Stream stream, Exception exception)
             var errorBodyHeader = $"{StreamingConstants.ErrorBodyTrailer}: {errorBodyJson}\r\n";
             var errorBodyBytes = Encoding.UTF8.GetBytes(errorBodyHeader);
             await stream.WriteAsync(errorBodyBytes, 0, errorBodyBytes.Length);
-
-            await stream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length);
         }
     }
 }