From 07d89f454b83f47c847f58127fe79cd8113991e2 Mon Sep 17 00:00:00 2001
From: Menno van Lavieren <menno.van.lavieren@uname-it.nl>
Date: Mon, 10 Jul 2017 14:14:54 +0200
Subject: [PATCH] Bug fix in nested context disposing

---
 Tapeti/Connection/TapetiConsumer.cs | 40 +++++++++++++++++++----------
 Tapeti/Default/MessageContext.cs    | 26 ++++---------------
 2 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/Tapeti/Connection/TapetiConsumer.cs b/Tapeti/Connection/TapetiConsumer.cs
index 57c31f2..d47a498 100644
--- a/Tapeti/Connection/TapetiConsumer.cs
+++ b/Tapeti/Connection/TapetiConsumer.cs
@@ -160,10 +160,10 @@ namespace Tapeti.Connection
 
     public delegate Task Handler(MessageContext context, Func<Task> next);
 
-    public class RecursiveCaller: ICallFrame
+    public class RecursiveCaller
     {
         private Handler handle;
-        private MessageContext context;
+        private MessageContext currentContext;
         private MessageContext nextContext;
         public RecursiveCaller next;
 
@@ -174,41 +174,55 @@ namespace Tapeti.Connection
 
         internal async Task Call(MessageContext context)
         {
-            if (this.context != null)
+            if (currentContext != null)
                 throw new InvalidOperationException("Cannot simultaneously call 'next' in Middleware.");
 
             try
             {
-                this.context = context;
+                currentContext = context;
 
-                if (next != null)
-                    context.SetCallFrame(this);
+                context.UseNestedContext = next == null ? (Action<MessageContext>)null : UseNestedContext;
 
                 await handle(context, callNext);
             }
             finally
             {
-                context = null;
+                currentContext = null;
             }
         }
 
-        private Task callNext()
+        private async Task callNext()
         {
             if (next == null)
-                return Task.CompletedTask;
-
-            return next.Call(nextContext ?? context);
+                return;
+            if (nextContext != null)
+            {
+                await next.Call(nextContext);
+            }else
+            {
+                try
+                {
+                    await next.Call(currentContext);
+                }
+                finally
+                {
+                    currentContext.UseNestedContext = UseNestedContext;
+                }
+            }
         }
 
-        void ICallFrame.UseNestedContext(MessageContext context)
+        void UseNestedContext(MessageContext context)
         {
             if (nextContext != null)
                 throw new InvalidOperationException("Previous nested context was not yet disposed.");
+
+            context.OnContextDisposed = OnContextDisposed;
             nextContext = context;
         }
 
-        void ICallFrame.OnContextDisposed(MessageContext context)
+        void OnContextDisposed(MessageContext context)
         {
+            context.OnContextDisposed = null;
             if (nextContext == context)
                 nextContext = null;
         }
diff --git a/Tapeti/Default/MessageContext.cs b/Tapeti/Default/MessageContext.cs
index 728655b..999b532 100644
--- a/Tapeti/Default/MessageContext.cs
+++ b/Tapeti/Default/MessageContext.cs
@@ -7,11 +7,6 @@ using System.Linq;
 
 namespace Tapeti.Default
 {
-    public interface ICallFrame {
-        void UseNestedContext(MessageContext context);
-        void OnContextDisposed(MessageContext context);
-    }
-
     public class MessageContext : IMessageContext
     {
         public IDependencyResolver DependencyResolver { get; set; }
@@ -27,20 +22,14 @@ namespace Tapeti.Default
         public IDictionary<string, object> Items { get; }
 
         private readonly MessageContext outerContext;
-        private ICallFrame callFrame;
+        internal Action<MessageContext> UseNestedContext;
+        internal Action<MessageContext> OnContextDisposed;
 
         public MessageContext()
         {
             Items = new Dictionary<string, object>();
         }
 
-        public MessageContext(ICallFrame callFrame)
-        {
-            Items = new Dictionary<string, object>();
-
-            this.callFrame = callFrame;
-        }
-
         private MessageContext(MessageContext outerContext)
         {
             DependencyResolver = outerContext.DependencyResolver;
@@ -65,22 +54,17 @@ namespace Tapeti.Default
             foreach (var value in items.Values)
                 (value as IDisposable)?.Dispose();
 
-            callFrame?.OnContextDisposed(this);
-        }
-
-        public void SetCallFrame(ICallFrame callFrame)
-        {
-            this.callFrame = callFrame;
+            OnContextDisposed?.Invoke(this);
         }
 
         public IMessageContext SetupNestedContext()
         {
-            if (callFrame == null)
+            if (UseNestedContext == null)
                 throw new NotSupportedException("This context does not support creating nested contexts");
 
             var nested = new MessageContext(this);
 
-            callFrame.UseNestedContext(nested);
+            UseNestedContext(nested);
 
             return nested;
         }