Merge pull request #1058 from zachpainter77/timeout-behavior-support

Timeout behavior support
jbogard · Sep 11, 2024 · 0d3bf4c · 0d3bf4c
2 parents 69b54c8 + 0d5ef88
commit 0d3bf4c
Show file tree

Hide file tree

Showing 3 changed files with 278 additions and 201 deletions.
diff --git 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.

diff --git 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))();
  }
 }