diff --git a/src/DotNetty.Handlers/Timeout/WriteTimeoutHandler.cs b/src/DotNetty.Handlers/Timeout/WriteTimeoutHandler.cs index 76f616595..afc0e612e 100644 --- a/src/DotNetty.Handlers/Timeout/WriteTimeoutHandler.cs +++ b/src/DotNetty.Handlers/Timeout/WriteTimeoutHandler.cs @@ -85,14 +85,17 @@ public WriteTimeoutHandler(TimeSpan timeout) public override ValueTask WriteAsync(IChannelHandlerContext context, object message) { - ValueTask task = context.WriteAsync(message); + ValueTask future = context.WriteAsync(message); if (this.timeout.Ticks > 0) { + //allocating task cause we need to attach continuation + Task task = future.AsTask(); this.ScheduleTimeout(context, task); + return new ValueTask(task); } - return task; + return future; } public override void HandlerRemoved(IChannelHandlerContext context) @@ -106,20 +109,19 @@ public override void HandlerRemoved(IChannelHandlerContext context) } } - async void ScheduleTimeout(IChannelHandlerContext context, ValueTask future) + void ScheduleTimeout(IChannelHandlerContext context, Task future) { // Schedule a timeout. var task = new WriteTimeoutTask(context, future, this); - task.ScheduledTask = context.Executor.Schedule(task, this.timeout); + task.ScheduledTask = context.Executor.Schedule(task, timeout); if (!task.ScheduledTask.Completion.IsCompleted) { this.AddWriteTimeoutTask(task); // Cancel the scheduled timeout if the flush promise is complete. - await future; - WriteTimeoutTask.OperationCompleteAction(task); + future.ContinueWith(WriteTimeoutTask.OperationCompleteAction, task, TaskContinuationOptions.ExecuteSynchronously); } } @@ -147,22 +149,22 @@ protected virtual void WriteTimedOut(IChannelHandlerContext context) } } - sealed class WriteTimeoutTask : AbstractPromise, IRunnable + sealed class WriteTimeoutTask : IRunnable { readonly WriteTimeoutHandler handler; readonly IChannelHandlerContext context; - readonly ValueTask future; + readonly Task future; - public static readonly Action OperationCompleteAction = HandleOperationComplete; + public static readonly Action OperationCompleteAction = HandleOperationComplete; - public WriteTimeoutTask(IChannelHandlerContext context, ValueTask future, WriteTimeoutHandler handler) + public WriteTimeoutTask(IChannelHandlerContext context, Task future, WriteTimeoutHandler handler) { this.context = context; this.future = future; this.handler = handler; } - static void HandleOperationComplete(object state) + static void HandleOperationComplete(Task future, object state) { var writeTimeoutTask = (WriteTimeoutTask) state;