From 6ccec483de25d9d72772652d499faa5639cb6917 Mon Sep 17 00:00:00 2001
From: martincostello <martin@martincostello.com>
Date: Fri, 27 Oct 2023 16:34:17 +0100
Subject: [PATCH] Potential fix for AOT compilation

A potential fix for #1732 using `RuntimeFeature.IsDynamicCodeSupported` to guard against the allocation regression that avoiding the infinite generic recursion causes through boxing.
---
 .../Utils/Pipeline/CompositeComponent.cs      | 59 ++++++++++++++++++-
 1 file changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs b/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
index abaa8692f4f..31027146389 100644
--- a/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
+++ b/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
@@ -1,4 +1,5 @@
-using Polly.Telemetry;
+using System.Runtime.CompilerServices;
+using Polly.Telemetry;
 
 namespace Polly.Utils.Pipeline;
 
@@ -128,11 +129,26 @@ private sealed class DelegatingComponent : PipelineComponent
 
         public PipelineComponent? Next { get; set; }
 
+        public override ValueTask DisposeAsync() => default;
+
         internal override ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState>(
             Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>> callback,
             ResilienceContext context,
             TState state)
         {
+#if NET6_0_OR_GREATER
+            return RuntimeFeature.IsDynamicCodeSupported ? ExecuteCoreImpl(callback, context, state) : ExecuteCoreAot(callback, context, state);
+#else
+            return ExecuteCoreImpl(callback, context, state);
+#endif
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private ValueTask<Outcome<TResult>> ExecuteCoreImpl<TResult, TState>(
+            Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>> callback,
+            ResilienceContext context,
+            TState state)
+        {
             return _component.ExecuteCore(
                 static (context, state) =>
                 {
@@ -147,6 +163,45 @@ internal override ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState>(
                 (Next, callback, state));
         }
 
-        public override ValueTask DisposeAsync() => default;
+#if NET6_0_OR_GREATER
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private ValueTask<Outcome<TResult>> ExecuteCoreAot<TResult, TState>(
+            Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>> callback,
+            ResilienceContext context,
+            TState state)
+        {
+            // Custom state object is used to cast the callback and state to prevent infinite
+            // generic type recursion warning IL3054 when referenced in a native AoT application.
+            // See https://github.com/App-vNext/Polly/issues/1732 for further context.
+            return _component.ExecuteCore(
+                static (context, wrapper) =>
+                {
+                    var callback = (Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>>)wrapper.Callback;
+                    var state = (TState)wrapper.State;
+                    if (context.CancellationToken.IsCancellationRequested)
+                    {
+                        return Outcome.FromExceptionAsValueTask<TResult>(new OperationCanceledException(context.CancellationToken).TrySetStackTrace());
+                    }
+
+                    return wrapper.Next.ExecuteCore(callback, context, state);
+                },
+                context,
+                new StateWrapper(Next!, callback, state!));
+        }
+
+        private struct StateWrapper
+        {
+            public StateWrapper(PipelineComponent next, object callback, object state)
+            {
+                Next = next;
+                Callback = callback;
+                State = state;
+            }
+
+            public PipelineComponent Next;
+            public object Callback;
+            public object State;
+        }
+#endif
     }
 }