Spiking

Particular · Jan 27, 2022 · f5376bd · f5376bd
1 parent 4f4b3b0
commit f5376bd
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/src/NServiceBus.Core/Extensibility/ContextBag.cs b/src/NServiceBus.Core/Extensibility/ContextBag.cs
@@ -1,6 +1,7 @@
 namespace NServiceBus.Extensibility
 {
     using System.Collections.Generic;
+    using Pipeline;
 
     /// <summary>
     /// A string object bag of context objects.
@@ -143,8 +144,15 @@ internal void Merge(ContextBag context)
             }
         }
 
+        internal IBehavior[] Behaviors
+        {
+            get => behaviors ?? parentBag?.Behaviors;
+            set => behaviors = value;
+        }
+
         ContextBag parentBag;
 
         Dictionary<string, object> stash = new Dictionary<string, object>();
+        IBehavior[] behaviors;
     }
 }
diff --git a/src/NServiceBus.Core/Pipeline/Pipeline.cs b/src/NServiceBus.Core/Pipeline/Pipeline.cs
@@ -42,6 +42,7 @@ public Pipeline(IServiceProvider builder, PipelineModifications pipelineModifica
 
         public Task Invoke(TContext context)
         {
+            context.Extensions.Behaviors = behaviors;
             return pipeline(context);
         }
 

diff --git a/src/NServiceBus.Core/Pipeline/PipelineExecutionExtensions.cs b/src/NServiceBus.Core/Pipeline/PipelineExecutionExtensions.cs
@@ -5,6 +5,7 @@
     using System.Linq;
     using System.Linq.Expressions;
     using System.Reflection;
+    using System.Runtime.CompilerServices;
     using System.Threading.Tasks;
     using FastExpressionCompiler;
     using Pipeline;
@@ -48,28 +49,36 @@ public static Delegate CreatePipelineExecutionExpression(this IBehavior[] behavi
                 var genericArguments = behaviorInterfaceType.GetGenericArguments();
                 var inContextType = genericArguments[0];
 
+                MethodInfo method = typeof(PipelineExecutionExtensions).GetMethod("Next", BindingFlags.Static | BindingFlags.Public);
+                methodInfo = method.MakeGenericMethod(genericArguments);
+
                 var inContextParameter = Expression.Parameter(inContextType, $"context{i}");
 
                 if (i == behaviorCount)
                 {
                     var outContextType = genericArguments[1];
                     var doneDelegate = CreateDoneDelegate(outContextType, i);
-                    lambdaExpression = CreateBehaviorCallDelegate(currentBehavior, methodInfo, inContextParameter, doneDelegate, expressions);
+                    lambdaExpression = CreateBehaviorCallDelegate(methodInfo, inContextParameter, doneDelegate, i, expressions);
                     continue;
                 }
 
-                lambdaExpression = CreateBehaviorCallDelegate(currentBehavior, methodInfo, inContextParameter, lambdaExpression, expressions);
+                lambdaExpression = CreateBehaviorCallDelegate(methodInfo, inContextParameter, lambdaExpression, i, expressions);
             }
 
             return lambdaExpression;
         }
 
+        public static Task Next<TInContext, TOutContext>(TInContext ctx, Func<TOutContext, Task> next, int index)
+            where TInContext : IBehaviorContext
+            where TOutContext : IBehaviorContext =>
+            ((IBehavior<TInContext, TOutContext>)ctx.Extensions.Behaviors[index]).Invoke(ctx, next);
+
         /// <code>
         /// context{i} => behavior.Invoke(context{i}, context{i+1} => previous)
         /// </code>>
-        static Delegate CreateBehaviorCallDelegate(IBehavior currentBehavior, MethodInfo methodInfo, ParameterExpression outerContextParam, Delegate previous, List<Expression> expressions = null)
+        static Delegate CreateBehaviorCallDelegate(MethodInfo methodInfo, ParameterExpression outerContextParam, Delegate previous, int i, List<Expression> expressions = null)
         {
-            var body = Expression.Call(Expression.Constant(currentBehavior), methodInfo, outerContextParam, Expression.Constant(previous));
+            var body = Expression.Call(null, methodInfo, outerContextParam, Expression.Constant(previous), Expression.Constant(i));
             var lambdaExpression = Expression.Lambda(body, outerContextParam);
             expressions?.Add(lambdaExpression);
             return lambdaExpression.CompileFast();