From 52ea126084aa83b07852ea21e76dd9c9c06913e1 Mon Sep 17 00:00:00 2001
From: zach painter <zachpainter77@gmail.com>
Date: Sat, 3 Aug 2024 01:20:50 -0700
Subject: [PATCH 1/3] update wrapper

---
 src/MediatR/Wrappers/RequestHandlerWrapper.cs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/MediatR/Wrappers/RequestHandlerWrapper.cs b/src/MediatR/Wrappers/RequestHandlerWrapper.cs
index 1550ecf8..dc213560 100644
--- a/src/MediatR/Wrappers/RequestHandlerWrapper.cs
+++ b/src/MediatR/Wrappers/RequestHandlerWrapper.cs
@@ -34,14 +34,14 @@ public class RequestHandlerWrapperImpl<TRequest, TResponse> : RequestHandlerWrap
     public override Task<TResponse> Handle(IRequest<TResponse> request, IServiceProvider serviceProvider,
         CancellationToken cancellationToken)
     {
-        Task<TResponse> Handler() => serviceProvider.GetRequiredService<IRequestHandler<TRequest, TResponse>>()
-            .Handle((TRequest) request, cancellationToken);
+        Task<TResponse> Handler(CancellationToken t = default) => serviceProvider.GetRequiredService<IRequestHandler<TRequest, TResponse>>()
+            .Handle((TRequest) request, t == default ? cancellationToken : t);
 
         return serviceProvider
             .GetServices<IPipelineBehavior<TRequest, TResponse>>()
             .Reverse()
             .Aggregate((RequestHandlerDelegate<TResponse>) Handler,
-                (next, pipeline) => () => pipeline.Handle((TRequest) request, next, cancellationToken))();
+                (next, pipeline) => (t) => pipeline.Handle((TRequest) request, next, t == default ? cancellationToken : t))();
     }
 }
 
@@ -55,10 +55,10 @@ public class RequestHandlerWrapperImpl<TRequest> : RequestHandlerWrapper
     public override Task<Unit> Handle(IRequest request, IServiceProvider serviceProvider,
         CancellationToken cancellationToken)
     {
-        async Task<Unit> Handler()
+        async Task<Unit> Handler(CancellationToken t = default)
         {
             await serviceProvider.GetRequiredService<IRequestHandler<TRequest>>()
-                .Handle((TRequest) request, cancellationToken);
+                .Handle((TRequest) request, t == default ? cancellationToken : t);
 
             return Unit.Value;
         }
@@ -67,6 +67,6 @@ await serviceProvider.GetRequiredService<IRequestHandler<TRequest>>()
             .GetServices<IPipelineBehavior<TRequest, Unit>>()
             .Reverse()
             .Aggregate((RequestHandlerDelegate<Unit>) Handler,
-                (next, pipeline) => () => pipeline.Handle((TRequest) request, next, cancellationToken))();
+                (next, pipeline) => (t) => pipeline.Handle((TRequest) request, next, t == default ? cancellationToken : t))();
     }
 }
\ No newline at end of file

From cca7861c091fdb4207d57e504fd646afe9917b62 Mon Sep 17 00:00:00 2001
From: zach painter <zachpainter77@gmail.com>
Date: Sat, 3 Aug 2024 01:22:56 -0700
Subject: [PATCH 2/3] modify pipline behavior abstraction

---
 src/MediatR/IPipelineBehavior.cs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/MediatR/IPipelineBehavior.cs b/src/MediatR/IPipelineBehavior.cs
index be6b1cb3..142ec707 100644
--- a/src/MediatR/IPipelineBehavior.cs
+++ b/src/MediatR/IPipelineBehavior.cs
@@ -9,7 +9,7 @@ namespace MediatR;
 /// </summary>
 /// <typeparam name="TResponse">Response type</typeparam>
 /// <returns>Awaitable task returning a <typeparamref name="TResponse"/></returns>
-public delegate Task<TResponse> RequestHandlerDelegate<TResponse>();
+public delegate Task<TResponse> RequestHandlerDelegate<TResponse>(CancellationToken t = default);
 
 /// <summary>
 /// Pipeline behavior to surround the inner handler.

From 83dbba0c02d6964850a30268693b57c45d0feaad Mon Sep 17 00:00:00 2001
From: zach painter <zachpainter77@gmail.com>
Date: Sat, 3 Aug 2024 01:25:16 -0700
Subject: [PATCH 3/3] add tests

---
 test/MediatR.Tests/SendTests.cs | 474 +++++++++++++++++++-------------
 1 file changed, 277 insertions(+), 197 deletions(-)

diff --git a/test/MediatR.Tests/SendTests.cs b/test/MediatR.Tests/SendTests.cs
index 553c689c..e1ea1cfd 100644
--- a/test/MediatR.Tests/SendTests.cs
+++ b/test/MediatR.Tests/SendTests.cs
@@ -1,108 +1,113 @@
-using System.Threading;
-
-using System;
-using System.Threading.Tasks;
-using Shouldly;
-using Xunit;
-using Microsoft.Extensions.DependencyInjection;
+using System.Threading;
+
+using System;
+using System.Threading.Tasks;
+using Shouldly;
+using Xunit;
+using Microsoft.Extensions.DependencyInjection;
 using System.Reflection;
-
-namespace MediatR.Tests;
-public class SendTests
-{
+using MediatR.Pipeline;
+
+namespace MediatR.Tests;
+public class SendTests
+{
     private readonly IServiceProvider _serviceProvider;
-    private Dependency _dependency;
+    private Dependency _dependency;
     private readonly IMediator _mediator;
 
     public SendTests()
     {
         _dependency = new Dependency();
         var services = new ServiceCollection();
-        services.AddMediatR(cfg => cfg.RegisterServicesFromAssemblies(typeof(Ping).Assembly));
+        services.AddMediatR(cfg =>
+        {
+            cfg.RegisterServicesFromAssemblies(typeof(Ping).Assembly);
+            cfg.AddOpenBehavior(typeof(TimeoutBehavior<,>), ServiceLifetime.Transient);
+        });
         services.AddSingleton(_dependency);
         _serviceProvider = services.BuildServiceProvider();
         _mediator = _serviceProvider.GetService<IMediator>()!;
-    }
-
-    public class Ping : IRequest<Pong>
-    {
-        public string? Message { get; set; }
-    }
-
-    public class VoidPing : IRequest
-    {
-    }
-
-    public class Pong
-    {
-        public string? Message { get; set; }
-    }
-
-    public class PingHandler : IRequestHandler<Ping, Pong>
-    {
-        public Task<Pong> Handle(Ping request, CancellationToken cancellationToken)
-        {
-            return Task.FromResult(new Pong { Message = request.Message + " Pong" });
-        }
-    }
-
-    public class Dependency
-    {
-        public bool Called { get; set; }
-        public bool CalledSpecific { get; set; }
-    }
-
-    public class VoidPingHandler : IRequestHandler<VoidPing>
-    {
-        private readonly Dependency _dependency;
-
-        public VoidPingHandler(Dependency dependency) => _dependency = dependency;
-
-        public Task Handle(VoidPing request, CancellationToken cancellationToken)
-        {
-            _dependency.Called = true;
-
-            return Task.CompletedTask;
-        }
-    }
-
-    public class GenericPing<T> : IRequest<T>
-        where T : Pong
-    {
-        public T? Pong { get; set; }
-    }
-
-    public class GenericPingHandler<T> : IRequestHandler<GenericPing<T>, T>
-        where T : Pong
-    {
-        private readonly Dependency _dependency;
-
-        public GenericPingHandler(Dependency dependency) => _dependency = dependency;
-
-        public Task<T> Handle(GenericPing<T> request, CancellationToken cancellationToken)
-        {
-            _dependency.Called = true;
-            request.Pong!.Message += " Pong";
-            return Task.FromResult(request.Pong!);
-        }
-    }
-
-    public class VoidGenericPing<T> : IRequest
-        where T : Pong
-    { }
-
-    public class VoidGenericPingHandler<T> : IRequestHandler<VoidGenericPing<T>>
-        where T : Pong
-    {
-        private readonly Dependency _dependency;
-        public VoidGenericPingHandler(Dependency dependency) => _dependency = dependency;
-
-        public Task Handle(VoidGenericPing<T> request, CancellationToken cancellationToken)
-        {
-            _dependency.Called = true;
-
-            return Task.CompletedTask;
-        }
+    }
+
+    public class Ping : IRequest<Pong>
+    {
+        public string? Message { get; set; }
+    }
+
+    public class VoidPing : IRequest
+    {
+    }
+
+    public class Pong
+    {
+        public string? Message { get; set; }
+    }
+
+    public class PingHandler : IRequestHandler<Ping, Pong>
+    {
+        public Task<Pong> Handle(Ping request, CancellationToken cancellationToken)
+        {
+            return Task.FromResult(new Pong { Message = request.Message + " Pong" });
+        }
+    }
+
+    public class Dependency
+    {
+        public bool Called { get; set; }
+        public bool CalledSpecific { get; set; }
+    }
+
+    public class VoidPingHandler : IRequestHandler<VoidPing>
+    {
+        private readonly Dependency _dependency;
+
+        public VoidPingHandler(Dependency dependency) => _dependency = dependency;
+
+        public Task Handle(VoidPing request, CancellationToken cancellationToken)
+        {
+            _dependency.Called = true;
+
+            return Task.CompletedTask;
+        }
+    }
+
+    public class GenericPing<T> : IRequest<T>
+        where T : Pong
+    {
+        public T? Pong { get; set; }
+    }
+
+    public class GenericPingHandler<T> : IRequestHandler<GenericPing<T>, T>
+        where T : Pong
+    {
+        private readonly Dependency _dependency;
+
+        public GenericPingHandler(Dependency dependency) => _dependency = dependency;
+
+        public Task<T> Handle(GenericPing<T> request, CancellationToken cancellationToken)
+        {
+            _dependency.Called = true;
+            request.Pong!.Message += " Pong";
+            return Task.FromResult(request.Pong!);
+        }
+    }
+
+    public class VoidGenericPing<T> : IRequest
+        where T : Pong
+    { }
+
+    public class VoidGenericPingHandler<T> : IRequestHandler<VoidGenericPing<T>>
+        where T : Pong
+    {
+        private readonly Dependency _dependency;
+        public VoidGenericPingHandler(Dependency dependency) => _dependency = dependency;
+
+        public Task Handle(VoidGenericPing<T> request, CancellationToken cancellationToken)
+        {
+            _dependency.Called = true;
+
+            return Task.CompletedTask;
+        }
     }
 
     public class PongExtension : Pong
@@ -123,22 +128,22 @@ public Task Handle(VoidGenericPing<PongExtension> request, CancellationToken can
         }
     }
 
-    public interface ITestInterface1 { }
-    public interface ITestInterface2 { }
+    public interface ITestInterface1 { }
+    public interface ITestInterface2 { }
     public interface ITestInterface3 { }
 
     public class TestClass1 : ITestInterface1 { }
     public class TestClass2 : ITestInterface2 { }
     public class TestClass3 : ITestInterface3 { }
 
-    public class MultipleGenericTypeParameterRequest<T1, T2, T3> : IRequest<int> 
-       where T1 : ITestInterface1
-       where T2 : ITestInterface2
-       where T3 : ITestInterface3
-    {
-        public int Foo { get; set; }
-    }
-
+    public class MultipleGenericTypeParameterRequest<T1, T2, T3> : IRequest<int>
+       where T1 : ITestInterface1
+       where T2 : ITestInterface2
+       where T3 : ITestInterface3
+    {
+        public int Foo { get; set; }
+    }
+
     public class MultipleGenericTypeParameterRequestHandler<T1, T2, T3> : IRequestHandler<MultipleGenericTypeParameterRequest<T1, T2, T3>, int>
         where T1 : ITestInterface1
         where T2 : ITestInterface2
@@ -148,92 +153,141 @@ public class MultipleGenericTypeParameterRequestHandler<T1, T2, T3> : IRequestHa
 
         public MultipleGenericTypeParameterRequestHandler(Dependency dependency) => _dependency = dependency;
 
-        public Task<int> Handle(MultipleGenericTypeParameterRequest<T1, T2, T3> request, CancellationToken cancellationToken)
-        {
-            _dependency.Called = true;
-            return Task.FromResult(1);
-        }
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_handler()
-    {
-        var response = await _mediator.Send(new Ping { Message = "Ping" });
-
-        response.Message.ShouldBe("Ping Pong");
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_void_handler()
-    {
-        await _mediator.Send(new VoidPing());
-
-        _dependency.Called.ShouldBeTrue();
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_handler_via_dynamic_dispatch()
-    {
-        object request = new Ping { Message = "Ping" };
-        var response = await _mediator.Send(request);
-
-        var pong = response.ShouldBeOfType<Pong>();
-        pong.Message.ShouldBe("Ping Pong");
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_void_handler_via_dynamic_dispatch()
-    {
-        object request = new VoidPing();
-        var response = await _mediator.Send(request);
-
-        response.ShouldBeOfType<Unit>();
-
-        _dependency.Called.ShouldBeTrue();
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_handler_by_specific_interface()
-    {
-        var response = await _mediator.Send(new Ping { Message = "Ping" });
-
-        response.Message.ShouldBe("Ping Pong");
-    }
-
-    [Fact]
-    public async Task Should_resolve_main_handler_by_given_interface()
-    {
-        // wrap requests in an array, so this test won't break on a 'replace with var' refactoring
-        var requests = new IRequest[] { new VoidPing() };
-        await _mediator.Send(requests[0]);
-
-        _dependency.Called.ShouldBeTrue();
-    }
-
-    [Fact]
-    public Task Should_raise_execption_on_null_request() => Should.ThrowAsync<ArgumentNullException>(async () => await _mediator.Send(default!));
-
-    [Fact]
-    public async Task Should_resolve_generic_handler()
+        public Task<int> Handle(MultipleGenericTypeParameterRequest<T1, T2, T3> request, CancellationToken cancellationToken)
+        {
+            _dependency.Called = true;
+            return Task.FromResult(1);
+        }
+    }
+
+    public class TimeoutBehavior<TRequest, TResponse> : IPipelineBehavior<TRequest, TResponse>
+        where TRequest : notnull
+    {
+        public async Task<TResponse> Handle(TRequest request, RequestHandlerDelegate<TResponse> next, CancellationToken cancellationToken)
+        {
+            using (var cts = new CancellationTokenSource(500))
+            {
+                return await next(cts.Token);
+            }
+        }
+    }
+
+    public class TimeoutRequest : IRequest
+    {
+    }
+
+    public class TimeoutRequest2 : IRequest<int>
+    {
+    }
+
+    public class TimeoutRequestHandler : IRequestHandler<TimeoutRequest>
+    {
+        private readonly Dependency _dependency;
+
+        public TimeoutRequestHandler(Dependency dependency) => _dependency = dependency;
+
+        public async Task Handle(TimeoutRequest request, CancellationToken cancellationToken)
+        {
+            await Task.Delay(2000, cancellationToken);
+
+            _dependency.Called = true;
+        }
+    }
+
+    public class TimeoutRequest2Handler : IRequestHandler<TimeoutRequest2, int>
+    {
+        private readonly Dependency _dependency;
+
+        public TimeoutRequest2Handler(Dependency dependency) => _dependency = dependency;
+
+        public async Task<int> Handle(TimeoutRequest2 request, CancellationToken cancellationToken)
+        {
+            await Task.Delay(2000, cancellationToken);
+
+            _dependency.Called = true;
+            return 1;
+        }
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_handler()
+    {
+        var response = await _mediator.Send(new Ping { Message = "Ping" });
+
+        response.Message.ShouldBe("Ping Pong");
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_void_handler()
+    {
+        await _mediator.Send(new VoidPing());
+
+        _dependency.Called.ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_handler_via_dynamic_dispatch()
+    {
+        object request = new Ping { Message = "Ping" };
+        var response = await _mediator.Send(request);
+
+        var pong = response.ShouldBeOfType<Pong>();
+        pong.Message.ShouldBe("Ping Pong");
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_void_handler_via_dynamic_dispatch()
+    {
+        object request = new VoidPing();
+        var response = await _mediator.Send(request);
+
+        response.ShouldBeOfType<Unit>();
+
+        _dependency.Called.ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_handler_by_specific_interface()
+    {
+        var response = await _mediator.Send(new Ping { Message = "Ping" });
+
+        response.Message.ShouldBe("Ping Pong");
+    }
+
+    [Fact]
+    public async Task Should_resolve_main_handler_by_given_interface()
+    {
+        // wrap requests in an array, so this test won't break on a 'replace with var' refactoring
+        var requests = new IRequest[] { new VoidPing() };
+        await _mediator.Send(requests[0]);
+
+        _dependency.Called.ShouldBeTrue();
+    }
+
+    [Fact]
+    public Task Should_raise_execption_on_null_request() => Should.ThrowAsync<ArgumentNullException>(async () => await _mediator.Send(default!));
+
+    [Fact]
+    public async Task Should_resolve_generic_handler()
     {
         var request = new GenericPing<Pong> { Pong = new Pong { Message = "Ping" } };
-        var result = await _mediator.Send(request);
-
-        var pong = result.ShouldBeOfType<Pong>();
-        pong.Message.ShouldBe("Ping Pong");
-
-        _dependency.Called.ShouldBeTrue();
-    }
-
-    [Fact]
-    public async Task Should_resolve_generic_void_handler()
-    {
-        var request = new VoidGenericPing<Pong>();
-        await _mediator.Send(request);
-
-        _dependency.Called.ShouldBeTrue();
-    }
-
+        var result = await _mediator.Send(request);
+
+        var pong = result.ShouldBeOfType<Pong>();
+        pong.Message.ShouldBe("Ping Pong");
+
+        _dependency.Called.ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task Should_resolve_generic_void_handler()
+    {
+        var request = new VoidGenericPing<Pong>();
+        await _mediator.Send(request);
+
+        _dependency.Called.ShouldBeTrue();
+    }
+
     [Fact]
     public async Task Should_resolve_multiple_type_parameter_generic_handler()
     {
@@ -241,16 +295,16 @@ public async Task Should_resolve_multiple_type_parameter_generic_handler()
         await _mediator.Send(request);
 
         _dependency.Called.ShouldBeTrue();
-    }
-
-    [Fact]
+    }
+
+    [Fact]
     public async Task Should_resolve_closed_handler_if_defined()
     {
         var dependency = new Dependency();
         var services = new ServiceCollection();
-        services.AddSingleton(dependency);       
+        services.AddSingleton(dependency);
         services.AddMediatR(cfg => cfg.RegisterServicesFromAssemblies(Assembly.GetExecutingAssembly()));
-        services.AddTransient<IRequestHandler<VoidGenericPing<PongExtension>>,TestClass1PingRequestHandler>();
+        services.AddTransient<IRequestHandler<VoidGenericPing<PongExtension>>, TestClass1PingRequestHandler>();
         var serviceProvider = services.BuildServiceProvider();
         var mediator = serviceProvider.GetService<IMediator>()!;
 
@@ -259,9 +313,9 @@ public async Task Should_resolve_closed_handler_if_defined()
 
         dependency.Called.ShouldBeFalse();
         dependency.CalledSpecific.ShouldBeTrue();
-    }
-
-    [Fact]
+    }
+
+    [Fact]
     public async Task Should_resolve_open_handler_if_not_defined()
     {
         var dependency = new Dependency();
@@ -277,5 +331,31 @@ public async Task Should_resolve_open_handler_if_not_defined()
 
         dependency.Called.ShouldBeTrue();
         dependency.CalledSpecific.ShouldBeFalse();
-    }
+    }
+
+    [Fact]
+    public async Task TimeoutBehavior_Void_Should_Cancel_Long_Running_Task_And_Throw_Exception()
+    {
+        var request = new TimeoutRequest();
+
+        var exception = await Should.ThrowAsync<TaskCanceledException>(() => _mediator.Send(request));
+
+        exception.ShouldNotBeNull();
+        exception.ShouldBeAssignableTo<TaskCanceledException>();
+        _dependency.Called.ShouldBeFalse();
+    }
+
+    [Fact]
+    public async Task TimeoutBehavior_NonVoid_Should_Cancel_Long_Running_Task_And_Throw_Exception()
+    {
+        var request = new TimeoutRequest2();
+        int result = 0;
+
+        var exception = await Should.ThrowAsync<TaskCanceledException>(async () => { result = await _mediator.Send(request); });
+
+        exception.ShouldNotBeNull();
+        exception.ShouldBeAssignableTo<TaskCanceledException>();
+        _dependency.Called.ShouldBeFalse();
+        result.ShouldBe(0);
+    }
 }
\ No newline at end of file