如何配置UnitOfWork/Repository以适配await切换线程的场景？

在生产环境常见场景中，调用await Task.Delay()（模拟外部服务调用）后，RequestHandler中剩余代码会切换到另一个线程执行——这是await的设计特性，但此时向仓储添加审计条目会抛出异常，因为UnitOfWork是在之前的线程中开启的。

我希望请求处理器中的所有操作都处于同一个工作单元/事务中，同时也认为使用await调用是合理的，请问正确的处理方案是什么？

查阅资料后得知，多数观点认为工作单元应是短生命周期且仅在单线程运行，这点我认同。我曾考虑用.GetAwaiter().GetResult()替代await，但这显然不是正确方案，不仅有死锁风险，而且await编译正常却运行报错的情况也让人困惑。

以下是模拟该场景的代码：

请求处理器代码 #

using System.Threading;
using System.Threading.Tasks;
using MediatR;
using MultiThreadedUnitOfWork.Models;
using MultiThreadedUnitOfWork.UnitOfWorkAndRepositories;

namespace MultiThreadedUnitOfWork.RequestsAndHandlers.Multi_Threaded
{
    public class GetUserNameByIdRequestHandlerMultiThreaded : IRequestHandler<GetUserNameByIdRequestMultiThreaded, GetUserNameByIdResponseMultiThreaded>
    {
        private readonly IRepository _repository;

        public GetUserNameByIdRequestHandlerMultiThreaded(IRepository repository)
        {
            _repository = repository;
        }

        public async Task<GetUserNameByIdResponseMultiThreaded> Handle(GetUserNameByIdRequestMultiThreaded request, CancellationToken cancellationToken)
        {
            var user = _repository.Get<User>(request.Id);

            // 模拟调用外部服务
            await Task.Delay(1000, cancellationToken);

            // 此后代码在另一个线程执行
            var auditEntry = new Audit(nameof(User));

            // 该行报错，因为工作单元在之前的线程中开启
            _repository.Add(auditEntry);

            return new GetUserNameByIdResponseMultiThreaded(user.Name);
        }
    }
}

测试场景代码 #

using System;
using System.Threading.Tasks;
using MediatR;
using MultiThreadedUnitOfWork.IoC;
using MultiThreadedUnitOfWork.Models;
using MultiThreadedUnitOfWork.RequestsAndHandlers.Multi_Threaded;
using MultiThreadedUnitOfWork.RequestsAndHandlers.Single_Threaded;
using MultiThreadedUnitOfWork.UnitOfWorkAndRepositories;
using NUnit.Framework;

namespace MultiThreadedUnitOfWork.Tests
{
    [TestFixture]
    public class Test
    {
        [OneTimeSetUp]
        public void OneTimeSetUp()
        {
            _mediator = Container.Resolve<IMediator>();
            _repository = Container.Resolve<IRepository>();

            using (var uow = UnitOfWork.Start())
            {
                var user = new User("Test User");

                _repository.Add(user);

                uow.Complete();
            }
        }

        [TearDown]
        public void TearDown()
        {
            if (UnitOfWork.Current != null)
            {
                UnitOfWork.Current.Dispose();
            }
        }

        private IMediator _mediator;
        private IRepository _repository;

        [Test]
        public async Task MultiThreadedTest()
        {
            var request = new GetUserNameByIdRequestMultiThreaded(1);

            GetUserNameByIdResponseMultiThreaded response;

            try
            {
                using (var uow = UnitOfWork.Start())
                {
                    response = await _mediator.Send(request);

                    uow.Complete();
                }
            }
            catch (Exception e)
            {
                Console.WriteLine(e);
                throw;
            }

            Assert.That(response, Is.Not.Null);
        }
    }
}

问题根源排查 #

经排查，问题出在PerThreadRequestState类。该类使用[ThreadStatic]标记存储状态，导致线程切换后无法获取到之前线程中存储的UnitOfWork实例。但由于网站部分功能依赖该类，无法直接移除，因此需要替代方案。

问题类代码 #

PerThreadRequestState

using System;
using System.Collections.Generic;

namespace MultiThreadedUnitOfWork.RequestStates
{
    public class PerThreadRequestState : IRequestState
    {
        [ThreadStatic]
        private static IDictionary<string, object> _state;
        
        private static IDictionary<string, object> State => _state ?? (_state = new Dictionary<string, object>());

        public void Store<T>(T something)
        {
            lock (State)
            {
                State[typeof(T).FullName] = something;
            }
        }
        
        public T Get<T>()
        {
            lock (State)
            {
                if (State.TryGetValue(typeof(T).FullName, out var state)) 
                    return (T)state;

                return default;
            }
        }
    }
}

UnitOfWork实现

using System.Transactions;
using MultiThreadedUnitOfWork.IoC;
using MultiThreadedUnitOfWork.RequestStates;
using NHibernate;
using IsolationLevel = System.Data.IsolationLevel;

namespace MultiThreadedUnitOfWork.UnitOfWorkAndRepositories
{
    public sealed class UnitOfWorkWithNHibernateTransaction : IUnitOfWork
    {
        internal UnitOfWorkWithNHibernateTransaction()
        {
            if (Current != null)
                throw new UnitOfWorkException("Unit of work should not be nested");

            var sessionFactory = Container.Resolve<ISessionFactory>();
            Session = sessionFactory.OpenSession();

            if (Transaction.Current == null)
                Session.BeginTransaction(IsolationLevel.ReadCommitted);

            Current = this;
        }

        private static IUnitOfWork Current
        {
            get => Container.Resolve<IRequestState>().Get<IUnitOfWork>();
            set => Container.Resolve<IRequestState>().Store(value);
        }

        public ISession Session { get; private set; }
    }
}

解决方案 #

1. 替换线程绑定的状态存储为异步上下文兼容的实现 #

将PerThreadRequestState替换为基于AsyncLocal<T>的实现，AsyncLocal会在异步流中保留上下文，线程切换后仍能获取到正确的状态：

using System;
using System.Collections.Generic;

namespace MultiThreadedUnitOfWork.RequestStates
{
    public class AsyncRequestState : IRequestState
    {
        private static readonly AsyncLocal<IDictionary<string, object>> _state = new AsyncLocal<IDictionary<string, object>>();
        
        private IDictionary<string, object> State => _state.Value ?? (_state.Value = new Dictionary<string, object>());

        public void Store<T>(T something)
        {
            lock (State)
            {
                State[typeof(T).FullName] = something;
            }
        }
        
        public T Get<T>()
        {
            lock (State)
            {
                if (State.TryGetValue(typeof(T).FullName, out var state)) 
                    return (T)state;

                return default;
            }
        }
    }
}

然后在IoC容器中将IRequestState的实现从PerThreadRequestState替换为AsyncRequestState即可。

2. 确保NHibernate Session线程安全（可选） #

NHibernate的ISession本身不是线程安全的，但由于我们通过AsyncLocal确保同一异步流中使用的是同一个Session实例，因此不会有线程安全问题——因为异步流中的代码不会同时在多个线程执行，只是会切换线程。

3. 避免同步阻塞异步代码 #

绝对不要用.GetAwaiter().GetResult()或.Result来阻塞异步代码，这不仅会导致死锁风险，还违背了异步编程的设计初衷。

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