Rework to support class library programming model

Author: normj

Libraries/src/Amazon.Lambda.Core/LambdaResponseStreamFactory.cs

@@ -0,0 +1,184 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+#if NET8_0_OR_GREATER
+using System;
+using System.IO;
+using System.Runtime.Versioning;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Amazon.Lambda.Core
+{
+    /// <summary>
+    /// Factory to create Lambda response streams for writing streaming responses in AWS Lambda functions. The created streams are write-only and non-seekable.
+    /// </summary>
+    [RequiresPreviewFeatures(LambdaResponseStreamFactory.ParameterizedPreviewMessage)]
+    public class LambdaResponseStreamFactory
+    {
+        internal const string ParameterizedPreviewMessage =
+            "Response streaming is in preview till a new version of .NET Lambda runtime client that supports response streaming " +
+            "has been deployed to the .NET Lambda managed runtime. Till deployment has been made the feature can be used by deploying as an " +
+            "executable including the latest version of Amazon.Lambda.RuntimeSupport and setting the \"EnablePreviewFeatures\" in the Lambda " +
+            "project file to \"true\"";
+
+        private static Func<byte[], ILambdaResponseStream> _streamFactory;
+
+        internal static void SetLambdaResponseStream(Func<byte[], ILambdaResponseStream> streamFactory)
+        {
+            _streamFactory = streamFactory ?? throw new ArgumentNullException(nameof(streamFactory));
+        }
+
+        /// <summary>
+        /// Creates a <see cref="Stream"/> that can be used to write streaming responses back to callers of the Lambda function. Once
+        /// A Lambda function creates a response stream all output must be returned by writing to the stream; the Lambda function's handler
+        /// return value will be ignored. The stream is write-only and non-seekable.
+        /// </summary>
+        /// <returns></returns>
+        public static Stream CreateStream()
+        {
+            var runtimeResponseStream = _streamFactory(Array.Empty<byte>());
+            return new LambdaResponseStream(runtimeResponseStream);
+        }
+    }
+
+    /// <summary>
+    /// Interface for writing streaming responses in AWS Lambda functions.
+    /// Obtained by calling <see cref="LambdaResponseStreamFactory.CreateStream"/> within a handler.
+    /// </summary>
+    internal interface ILambdaResponseStream : IDisposable
+    {
+        /// <summary>
+        /// Asynchronously writes a portion of a byte array to the response stream.
+        /// </summary>
+        /// <param name="buffer">The byte array containing data to write.</param>
+        /// <param name="offset">The zero-based byte offset in buffer at which to begin copying bytes.</param>
+        /// <param name="count">The number of bytes to write.</param>
+        /// <param name="cancellationToken">Optional cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation.</returns>
+        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
+        Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default);
+
+
+        /// <summary>
+        /// Gets the total number of bytes written to the stream so far.
+        /// </summary>
+        long BytesWritten { get; }
+
+
+        /// <summary>
+        /// Gets whether an error has been reported.
+        /// </summary>
+        bool HasError { get; }
+    }
+
+    /// <summary>
+    /// A write-only, non-seekable <see cref="Stream"/> subclass that streams response data
+    /// to the Lambda Runtime API. Returned by <see cref="LambdaResponseStreamFactory.CreateStream"/>.
+    /// Integrates with standard .NET stream consumers such as <see cref="System.IO.StreamWriter"/>.
+    /// </summary>
+    [RequiresPreviewFeatures(LambdaResponseStreamFactory.ParameterizedPreviewMessage)]
+    public class LambdaResponseStream : Stream
+    {
+        private readonly ILambdaResponseStream _responseStream;
+
+        internal LambdaResponseStream(ILambdaResponseStream responseStream)
+        {
+            _responseStream = responseStream;
+        }
+
+        /// <summary>
+        /// The number of bytes written to the Lambda response stream so far.
+        /// </summary>
+        public long BytesWritten => _responseStream.BytesWritten;
+
+        /// <summary>
+        /// Asynchronously writes a byte array to the response stream.
+        /// </summary>
+        /// <param name="buffer">The byte array to write.</param>
+        /// <param name="cancellationToken">Optional cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation.</returns>
+        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
+        public async Task WriteAsync(byte[] buffer, CancellationToken cancellationToken = default)
+        {
+            if (buffer == null)
+                throw new ArgumentNullException(nameof(buffer));
+
+            await WriteAsync(buffer, 0, buffer.Length, cancellationToken);
+        }
+
+        /// <summary>
+        /// Asynchronously writes a portion of a byte array to the response stream.
+        /// </summary>
+        /// <param name="buffer">The byte array containing data to write.</param>
+        /// <param name="offset">The zero-based byte offset in buffer at which to begin copying bytes.</param>
+        /// <param name="count">The number of bytes to write.</param>
+        /// <param name="cancellationToken">Optional cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation.</returns>
+        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
+        public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
+        {
+            await _responseStream.WriteAsync(buffer, offset, count, cancellationToken);
+        }
+
+        #region Noop Overrides
+
+        /// <summary>Gets a value indicating whether the stream supports reading. Always <c>false</c>.</summary>
+        public override bool CanRead => false;
+
+        /// <summary>Gets a value indicating whether the stream supports seeking. Always <c>false</c>.</summary>
+        public override bool CanSeek => false;
+
+        /// <summary>Gets a value indicating whether the stream supports writing. Always <c>true</c>.</summary>
+        public override bool CanWrite => true;
+
+        /// <summary>
+        /// Gets the total number of bytes written to the stream so far.
+        /// Equivalent to <see cref="BytesWritten"/>.
+        /// </summary>
+        public override long Length => BytesWritten;
+
+        /// <summary>
+        /// Getting or setting the position is not supported.
+        /// </summary>
+        /// <exception cref="NotSupportedException">Always thrown.</exception>
+        public override long Position
+        {
+            get => throw new NotSupportedException("LambdaResponseStream does not support seeking.");
+            set => throw new NotSupportedException("LambdaResponseStream does not support seeking.");
+        }
+
+        /// <summary>Not supported.</summary>
+        /// <exception cref="NotImplementedException">Always thrown.</exception>
+        public override long Seek(long offset, SeekOrigin origin)
+            => throw new NotImplementedException("LambdaResponseStream does not support seeking.");
+
+        /// <summary>Not supported.</summary>
+        /// <exception cref="NotImplementedException">Always thrown.</exception>
+        public override int Read(byte[] buffer, int offset, int count)
+            => throw new NotImplementedException("LambdaResponseStream does not support reading.");
+
+        /// <summary>Not supported.</summary>
+        /// <exception cref="NotImplementedException">Always thrown.</exception>
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            => throw new NotImplementedException("LambdaResponseStream does not support reading.");
+
+        /// <summary>
+        /// Writes a sequence of bytes to the stream. Delegates to the async path synchronously.
+        /// Prefer <see cref="WriteAsync(byte[], int, int, CancellationToken)"/> to avoid blocking.
+        /// </summary>
+        public override void Write(byte[] buffer, int offset, int count)
+            => WriteAsync(buffer, offset, count, CancellationToken.None).GetAwaiter().GetResult();
+
+        /// <summary>
+        /// Flush is a no-op; data is sent to the Runtime API immediately on each write.
+        /// </summary>
+        public override void Flush() { }
+
+        /// <summary>Not supported.</summary>
+        /// <exception cref="NotSupportedException">Always thrown.</exception>
+        public override void SetLength(long value)
+            => throw new NotSupportedException("LambdaResponseStream does not support SetLength.");
+        #endregion
+    }
+}
+#endif

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs

@@ -20,6 +20,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Amazon.Lambda.RuntimeSupport.Bootstrap;
+using Amazon.Lambda.RuntimeSupport.Client.ResponseStreaming;
 using Amazon.Lambda.RuntimeSupport.Helpers;
 
 namespace Amazon.Lambda.RuntimeSupport
@@ -225,6 +226,19 @@ internal LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler,
                 return;
             }
 #if NET8_0_OR_GREATER
+
+            try
+            {
+                // Initalize in Amazon.Lambda.Core the factory for creating the response stream and related logic for supporting response streaming.
+                ResponseStreamLambdaCoreInitializerIsolated.InitializeCore();
+            }
+            catch (TypeLoadException)
+            {
+                _logger.LogDebug("Failed to configure Amazon.Lambda.Core with factory to create response stream. This happens when the version of Amazon.Lambda.Core referenced by the Lambda function is out of date.");
+            }
+
+
+
             // Check if Initialization type is SnapStart, and invoke the snapshot restore logic.
             if (_configuration.IsInitTypeSnapstart)
             {
@@ -363,7 +377,7 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                 var runtimeApiClient = Client as RuntimeApiClient;
                 if (runtimeApiClient != null)
                 {
-                    LambdaResponseStreamFactory.InitializeInvocation(
+                    ResponseStreamFactory.InitializeInvocation(
                         invocation.LambdaContext.AwsRequestId,
                         isMultiConcurrency,
                         runtimeApiClient,
@@ -385,7 +399,7 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                     {
                         WriteUnhandledExceptionToLog(exception);
 
-                        var responseStream = LambdaResponseStreamFactory.GetStreamIfCreated(isMultiConcurrency);
+                        var responseStream = ResponseStreamFactory.GetStreamIfCreated(isMultiConcurrency);
                         if (responseStream != null)
                         {
                             responseStream.ReportError(exception);
@@ -400,20 +414,20 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                         _logger.LogInformation("Finished invoking handler");
                     }
 
-                    var streamIfCreated = LambdaResponseStreamFactory.GetStreamIfCreated(isMultiConcurrency);
+                    var streamIfCreated = ResponseStreamFactory.GetStreamIfCreated(isMultiConcurrency);
                     if (streamIfCreated != null)
                     {
                         streamIfCreated.MarkCompleted();
 
                         // If streaming was started, await the HTTP send task to ensure it completes
-                        var sendTask = LambdaResponseStreamFactory.GetSendTask(isMultiConcurrency);
+                        var sendTask = ResponseStreamFactory.GetSendTask(isMultiConcurrency);
                         if (sendTask != null)
                         {
                             // Wait for the streaming response to finish sending before allowing the next invocation to be processed. This ensures that responses are sent in the order the invocations were received.
                             await sendTask;
                         }
 
-                        streamIfCreated.ManualDispose();
+                        streamIfCreated.Dispose();
                     }
                     else if (invokeSucceeded)
                     {
@@ -454,7 +468,7 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                 {
                     if (runtimeApiClient != null)
                     {
-                        LambdaResponseStreamFactory.CleanupInvocation(isMultiConcurrency);
+                        ResponseStreamFactory.CleanupInvocation(isMultiConcurrency);
                     }
                     invocation.Dispose();
                 }

…aResponseStream.ILambdaResponseStream.cs …trap/ResponseStreaming/ResponseStream.csLibraries/src/Amazon.Lambda.RuntimeSupport/Client/LambdaResponseStream.ILambdaResponseStream.cs renamed to Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStream.cs Libraries/src/Amazon.Lambda.RuntimeSupport/Client/LambdaResponseStream.ILambdaResponseStream.cs renamed to Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStream.cs

@@ -20,13 +20,12 @@
 using System.Threading.Tasks;
 using Amazon.Lambda.RuntimeSupport.Helpers;
 
-namespace Amazon.Lambda.RuntimeSupport
+namespace Amazon.Lambda.RuntimeSupport.Client.ResponseStreaming
 {
     /// <summary>
-    /// A write-only, non-seekable <see cref="Stream"/> subclass that streams response data
-    /// to the Lambda Runtime API. Returned by <see cref="LambdaResponseStreamFactory.CreateStream"/>.
+    /// Represents the writable stream used by Lambda handlers to write response data for streaming invocations.
     /// </summary>
-    public partial class LambdaResponseStream : Stream, ILambdaResponseStream
+    internal class ResponseStream
     {
         private static readonly byte[] CrlfBytes = Encoding.ASCII.GetBytes("\r\n");
 
@@ -38,6 +37,7 @@ public partial class LambdaResponseStream : Stream, ILambdaResponseStream
 
         // The live HTTP output stream, set by StreamingHttpContent when SerializeToStreamAsync is called.
         private Stream _httpOutputStream;
+        private bool _disposedValue;
         private readonly SemaphoreSlim _httpStreamReady = new SemaphoreSlim(0, 1);
         private readonly SemaphoreSlim _completionSignal = new SemaphoreSlim(0, 1);
 
@@ -56,12 +56,7 @@ public partial class LambdaResponseStream : Stream, ILambdaResponseStream
 
         internal Exception ReportedError => _reportedError;
 
-        internal LambdaResponseStream()
-            : this(Array.Empty<byte>())
-        {
-        }
-
-        internal LambdaResponseStream(byte[] prelude)
+        internal ResponseStream(byte[] prelude)
         {
             _prelude = prelude;
         }
@@ -125,18 +120,10 @@ internal async Task WaitForCompletionAsync(CancellationToken cancellationToken =
             await _completionSignal.WaitAsync(cancellationToken);
         }
 
-        /// <summary>
-        /// Asynchronously writes a byte array to the response stream.
-        /// </summary>
-        /// <param name="buffer">The byte array to write.</param>
-        /// <param name="cancellationToken">Optional cancellation token.</param>
-        /// <returns>A task representing the asynchronous operation.</returns>
-        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
-        public async Task WriteAsync(byte[] buffer, CancellationToken cancellationToken = default)
+        internal async Task WriteAsync(byte[] buffer, CancellationToken cancellationToken = default)
         {
             if (buffer == null)
                 throw new ArgumentNullException(nameof(buffer));
-
             await WriteAsync(buffer, 0, buffer.Length, cancellationToken);
         }
 
@@ -149,7 +136,7 @@ public async Task WriteAsync(byte[] buffer, CancellationToken cancellationToken
         /// <param name="cancellationToken">Optional cancellation token.</param>
         /// <returns>A task representing the asynchronous operation.</returns>
         /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
-        public override async 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));
@@ -232,25 +219,42 @@ internal void MarkCompleted()
             }
         }
 
-        /// <summary>
-        /// The resouces like the SemaphoreSlims are manually disposed by LambdaBootstrap after each invocation instead of relying on the
-        /// Dipose pattern because we don't want the user's Lambda function to trigger Releasing and disposing the semaphores when
-        /// invocation of the user's code ends.
-        /// </summary>
-        internal void ManualDispose()
-        {
-            try { _httpStreamReady.Release(); } catch (SemaphoreFullException) { /* Ignore if already released */ }
-            _httpStreamReady.Dispose();
-            try { _completionSignal.Release(); } catch (SemaphoreFullException) { /* Ignore if already released */ }
-            _completionSignal.Dispose();
-        }
-
         private void ThrowIfCompletedOrError()
         {
             if (_isCompleted)
                 throw new InvalidOperationException("Cannot write to a completed stream.");
             if (_hasError)
                 throw new InvalidOperationException("Cannot write to a stream after an error has been reported.");
         }
+
+        /// <summary>
+        /// Disposes the stream. After calling Dispose, no further writes or error reports should be made.
+        /// </summary>
+        /// <param name="disposing"></param>
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!_disposedValue)
+            {
+                if (disposing)
+                {
+                    try { _httpStreamReady.Release(); } catch (SemaphoreFullException) { /* Ignore if already released */ }
+                    _httpStreamReady.Dispose();
+                    try { _completionSignal.Release(); } catch (SemaphoreFullException) { /* Ignore if already released */ }
+                    _completionSignal.Dispose();
+                }
+
+                _disposedValue = true;
+            }
+        }
+
+        /// <summary>
+        /// Dispose of the stream. After calling Dispose, no further writes or error reports should be made.
+        /// </summary>
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
     }
 }