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 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)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 Items { get; } private readonly MessageContext outerContext; - private ICallFrame callFrame; + internal Action UseNestedContext; + internal Action OnContextDisposed; public MessageContext() { Items = new Dictionary(); } - public MessageContext(ICallFrame callFrame) - { - Items = new Dictionary(); - - 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; }