Use DefferedWebSocketResult to defer Task execution for NativeAOT context

Author: AlexMacocian

Net.Sdk.Web.Extensions/Net.Sdk.Web.Extensions.csproj

@@ -6,7 +6,7 @@
     <Nullable>enable</Nullable>
     <OutputType>Library</OutputType>
     <IsPackable>true</IsPackable>
-    <Version>0.8.9</Version>
+    <Version>0.8.10</Version>
     <Authors>Alexandru Macocian</Authors>
     <PackageLicenseFile>LICENSE</PackageLicenseFile>
     <PackageRequireLicenseAcceptance>true</PackageRequireLicenseAcceptance>

Net.Sdk.Web.Extensions/WebApplicationExtensions.cs

@@ -34,16 +34,22 @@ public static IEndpointConventionBuilder UseWebSocketRoute<TWebSocketRoute>(this
         where TWebSocketRoute : WebSocketRouteBase
     {
         app.ThrowIfNull();
-        return app.MapGet(route, HandleWebSocketRoute<TWebSocketRoute>);
+        return app.MapGet(route, HandleDeferredWebSocketRoute<TWebSocketRoute>);
     }
 
-    private static async Task<IResult> HandleWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
+    private static DeferredWebSocketResult HandleDeferredWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
+        where TWebSocketRoute : WebSocketRouteBase
+    {
+        return new DeferredWebSocketResult(route, HandleWebSocketRoute(context, route, logger));
+    }
+
+    private static async Task HandleWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
         where TWebSocketRoute : WebSocketRouteBase
     {
         if (!context.WebSockets.IsWebSocketRequest)
         {
             logger.LogError("WebSocket request expected");
-            return Results.BadRequest("WebSocket request expected");
+            return;
         }
 
         var routeFilters = GetRouteFilters<TWebSocketRoute>(context).ToList();
@@ -64,22 +70,18 @@ private static async Task<IResult> HandleWebSocketRoute<TWebSocketRoute>(HttpCon
         {
             var processingTask = new Func<Task>(() => ProcessWebSocketRequest(route, context));
             await BeginProcessingPipeline(actionExecutingContext, actionExecutedContext, processingTask);
-            return Results.Empty;
         }
         catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
         {
             logger.LogInformation("Websocket closed prematurely. Marking as closed");
-            return Results.Empty;
         }
         catch (OperationCanceledException)
         {
             logger.LogInformation("Websocket closed prematurely. Marking as closed");
-            return Results.Empty;
         }
         catch (Exception ex)
         {
             logger.LogError(ex, "Encountered exception while handling websocket. Closing");
-            return Results.Empty;
         }
         finally
         {

Net.Sdk.Web.Extensions/Websockets/DeferredWebSocketResult.cs

@@ -0,0 +1,23 @@
+using System.Core.Extensions;
+
+namespace Net.Sdk.Web.Websockets;
+
+public sealed class DeferredWebSocketResult(
+    WebSocketRouteBase route,
+    Task inner) : IResult
+{
+    private readonly WebSocketRouteBase route = route.ThrowIfNull();
+    private readonly Task inner = inner.ThrowIfNull();
+
+    public async Task ExecuteAsync(HttpContext httpContext)
+    {
+        try
+        {
+            await this.inner;
+        }
+        finally
+        {
+            await this.route.SocketClosed();
+        }
+    }
+}