如何正确实现ValueTaskSource.SetCompleted？IValueTaskSource实现存疑

我们实现了一个继承IValueTaskSource的Message类，由于无可等待对象，无法使用async-await。该类通过向已有工作线程发送消息，返回可供调用者等待的ValueTask，以获取线程处理完成的通知。当前实现代码如下：

class Message : IValueTaskSource {
    public ValueTask Send()
    {
        /* how the message is sent is irrelevant */
        return new ValueTask(this, 0);
    }

    private Action<object> continuation;
    private object continuationState;

    void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short _, ValueTaskSourceOnCompletedFlags __)
    {
         lock(this) {
              if (GetStatus(_) == ValueTaskSourceStatus.Pending)
              {
                  this.continuation = continuation;
                  this.continuationState = state;
                  return;
              }
              continuation(continuationState); /* Suspect */
         }
    }

    public void SetCompleted()
    {
        lock (this)
        {
             /* set state completed omitted for brevity */
             continuation?.Invoke(continuationState); /* Suspect */             
        }
    }
}

潜在问题分析 #

当前实现存在两个关键风险：

1. 栈溢出风险：两处标记为/* Suspect */的直接同步调用continuation逻辑，如果调用链嵌套层级过深（比如大规模连续使用这类Message对象），会导致当前线程栈被不断占用，最终触发栈溢出。
2. 未遵循IValueTaskSource规范：忽略了ValueTaskSourceOnCompletedFlags参数，无法正确处理续体的执行上下文流转、调度上下文要求，后续若遇到需要特定同步上下文的场景（如UI线程）会出现异常。

修复方案 #

1. 避免同步调用续体，防止栈溢出

直接在当前线程同步执行续体是栈溢出的根源，正确的做法是将续体调度到线程池或指定上下文异步执行，避免占用当前栈空间。

修改OnCompleted中状态已完成时的逻辑，以及SetCompleted中的续体调用逻辑：

// 封装调度续体的工具方法
private void QueueContinuation(Action<object> continuation, object state)
{
    // 使用线程池调度，避免同步执行
    ThreadPool.QueueUserWorkItem(static s =>
    {
        var tuple = ((Action<object>, object))s;
        tuple.Item1(tuple.Item2);
    }, (continuation, state));
}

void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
{
    lock(this) {
         if (GetStatus(token) == ValueTaskSourceStatus.Pending)
         {
             // 后续结合flags处理续体，先暂存基础信息
             this.continuation = continuation;
             this.continuationState = state;
             return;
         }
         // 状态已完成，异步调度续体
         QueueContinuation(continuation, state);
    }
}

public void SetCompleted()
{
    Action<object> cont = null;
    object contState = null;
    lock (this)
    {
         /* 补充状态设置逻辑，标记为已完成 */
         if (GetStatus(0) != ValueTaskSourceStatus.Pending)
             return;
         // 取出续体后立即清空，防止重复调用
         cont = continuation;
         contState = continuationState;
         continuation = null;
         continuationState = null;
    }
    // 异步调度续体执行
    cont?.Invoke(contState);
}

2. 正确处理ValueTaskSourceOnCompletedFlags参数

这个参数定义了续体的执行规则，核心要处理两个flag：

• FlowExecutionContext：需要捕获当前执行上下文，在续体执行时恢复
• UseSchedulingContext：需要使用当前的SynchronizationContext或TaskScheduler调度续体

修改后的完整OnCompleted和调度逻辑：

private sealed class ContinuationData
{
    public Action<object> Continuation { get; set; }
    public object State { get; set; }
    public ExecutionContext ExecutionContext { get; set; }
    public SynchronizationContext SyncContext { get; set; }
    public TaskScheduler TaskScheduler { get; set; }
}

private ContinuationData _continuationData;

void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
{
    lock(this) {
         if (GetStatus(token) == ValueTaskSourceStatus.Pending)
         {
             var data = new ContinuationData
             {
                 Continuation = continuation,
                 State = state,
                 // 根据flag决定是否捕获执行上下文
                 ExecutionContext = (flags & ValueTaskSourceOnCompletedFlags.FlowExecutionContext) != 0 
                     ? ExecutionContext.Capture() 
                     : null,
                 // 根据flag决定是否捕获调度上下文
                 SyncContext = (flags & ValueTaskSourceOnCompletedFlags.UseSchedulingContext) != 0 
                     ? SynchronizationContext.Current 
                     : null,
                 TaskScheduler = (flags & ValueTaskSourceOnCompletedFlags.UseSchedulingContext) != 0 
                     ? TaskScheduler.Current 
                     : null
             };
             _continuationData = data;
             return;
         }
         // 状态已完成，按规则调度续体
         ScheduleContinuation(continuation, state, flags);
    }
}

private void ScheduleContinuation(Action<object> continuation, object state, ValueTaskSourceOnCompletedFlags flags)
{
    // 优先使用同步上下文调度（如UI线程）
    if ((flags & ValueTaskSourceOnCompletedFlags.UseSchedulingContext) != 0)
    {
        var syncContext = SynchronizationContext.Current;
        if (syncContext != null && syncContext != SynchronizationContext.Default)
        {
            syncContext.Post(static s =>
            {
                var tuple = ((Action<object>, object))s;
                tuple.Item1(tuple.Item2);
            }, (continuation, state));
            return;
        }

        var scheduler = TaskScheduler.Current;
        if (scheduler != TaskScheduler.Default)
        {
            Task.Factory.StartNew(static s =>
            {
                var tuple = ((Action<object>, object))s;
                tuple.Item1(tuple.Item2);
            }, (continuation, state), CancellationToken.None, TaskCreationOptions.None, scheduler);
            return;
        }
    }

    // 处理执行上下文流转
    if ((flags & ValueTaskSourceOnCompletedFlags.FlowExecutionContext) != 0)
    {
        var context = ExecutionContext.Capture();
        ExecutionContext.Run(context, static s =>
        {
            var tuple = ((Action<object>, object))s;
            tuple.Item1(tuple.Item2);
        }, (continuation, state));
        return;
    }

    // 默认线程池调度
    ThreadPool.QueueUserWorkItem(static s =>
    {
        var tuple = ((Action<object>, object))s;
        tuple.Item1(tuple.Item2);
    }, (continuation, state));
}

public void SetCompleted()
{
    ContinuationData data = null;
    lock (this)
    {
         /* 补充状态设置逻辑，标记为已完成 */
         if (GetStatus(0) != ValueTaskSourceStatus.Pending)
             return;
         data = _continuationData;
         _continuationData = null;
    }

    if (data == null)
        return;

    // 根据捕获的上下文调度续体
    if (data.SyncContext != null && data.SyncContext != SynchronizationContext.Default)
    {
        data.SyncContext.Post(static s =>
        {
            var d = (ContinuationData)s;
            if (d.ExecutionContext != null)
            {
                ExecutionContext.Run(d.ExecutionContext, static innerS =>
                {
                    var innerD = (ContinuationData)innerS;
                    innerD.Continuation(innerD.State);
                }, d);
            }
            else
            {
                d.Continuation(d.State);
            }
        }, data);
    }
    else if (data.TaskScheduler != null && data.TaskScheduler != TaskScheduler.Default)
    {
        Task.Factory.StartNew(static s =>
        {
            var d = (ContinuationData)s;
            if (d.ExecutionContext != null)
            {
                ExecutionContext.Run(d.ExecutionContext, static innerS =>
                {
                    var innerD = (ContinuationData)innerS;
                    innerD.Continuation(innerD.State);
                }, d);
            }
            else
            {
                d.Continuation(d.State);
            }
        }, data, CancellationToken.None, TaskCreationOptions.None, data.TaskScheduler);
    }
    else if (data.ExecutionContext != null)
    {
        ExecutionContext.Run(data.ExecutionContext, static s =>
        {
            var d = (ContinuationData)s;
            d.Continuation(d.State);
        }, data);
    }
    else
    {
        ThreadPool.QueueUserWorkItem(static s =>
        {
            var d = (ContinuationData)s;
            d.Continuation(d.State);
        }, data);
    }
}

3. 完善状态管理逻辑

补充GetStatus方法的状态判断，确保状态流转正确，避免重复执行续体：

private enum MessageState
{
    Pending,
    Completed
}

private MessageState _state = MessageState.Pending;

public ValueTaskSourceStatus GetStatus(short token)
{
    lock (this)
    {
        return _state switch
        {
            MessageState.Pending => ValueTaskSourceStatus.Pending,
            MessageState.Completed => ValueTaskSourceStatus.Succeeded
        };
    }
}

public void SetCompleted()
{
    ContinuationData data = null;
    lock (this)
    {
         if (_state != MessageState.Pending)
             return;
         _state = MessageState.Completed;
         data = _continuationData;
         _continuationData = null;
    }

    // 后续调度续体的逻辑同上
    // ...
}

总结 #

通过异步调度续体替代同步调用，解决栈溢出风险；同时正确处理ValueTaskSourceOnCompletedFlags参数，符合IValueTaskSource的规范要求，能适配更多场景，避免后续大规模使用时出现异常。

内容的提问来源于stack exchange，提问作者Joshua