using System;

using System.Threading;

using System.Threading.Tasks;

​

namespace ConsoleTest

{

    internal class Program

    {

        private static async Task Main()

        {

            var uowDirect = new UnitOfWork(false);

​

            await using (uowDirect)

            {

                // Do stuff like calling repositories (with or without await).

                await Task.Yield();

​

                await uowDirect.CompleteAsync();

            } // <-- DisposeAsync called here

​

            // Disaster here: We are still in the ambient TransactionScope (which is disposed already, but still)

            // created by the now disposed unit of work.

​

            Console.WriteLine("Fake transaction for UnitOfWork direct: " + uowDirect.FakeTransaction.Value); //prints Created, which means the transaction still exists and will cause failures

​

​

            UnitOfWork? uowECRun = new UnitOfWork(true);

​

            await using (uowECRun)

            {

                // Do stuff like calling repositories (with or without await).

                await Task.Yield();

​

                Console.WriteLine("Fake transaction for UnitOfWork EC.Run(): " + uowECRun.FakeTransaction.Value);

​

                await uowECRun.CompleteAsync();

            } // <-- DisposeAsync called here

​

            // Disaster averted because the constructor's EC changes are contained in a copy of the EC

​

            Console.WriteLine("Fake transaction for UnitOfWork EC.Run(): " + uowECRun.FakeTransaction.Value); //prints empty string, meaning the transaction does not exist and everything is golden

        }

    }

​

    public class UnitOfWork : IAsyncDisposable

    {

        private ExecutionContext executionContext = null!;

​

        public readonly AsyncLocal<string> FakeTransaction = new();

​

        private void UpdateExecutionContext() => this.executionContext = ExecutionContext.Capture()!;

​

        public UnitOfWork(bool useCapturedExecutionContext)

        {

            if (useCapturedExecutionContext)

                ExecutionContext.Run(ExecutionContext.Capture()!, init, null);

            else

                init(null);

​

            void init(object? state)

            {

                this.FakeTransaction.Value = "Created"; // Create TransactionScope

                // Setup other things.

                this.UpdateExecutionContext();

            }

        }

​

        public async Task CompleteAsync()

        {

            ExecutionContext.Restore(this.executionContext);

​

            try

            {

                if (this.FakeTransaction.Value != "Created")

                    throw new InvalidOperationException($"Expected 'Created', got: '{this.FakeTransaction.Value}'");

​

                // Flush changes to the database or whatever. Not important here.

                await Task.Yield();

​

                this.FakeTransaction.Value = "Did some work";

            }

            finally

            {

                this.UpdateExecutionContext();

            }

        }

​

        public async ValueTask DisposeAsync()

        {

            ExecutionContext.Restore(this.executionContext);

​

            try

            {

                await Task.Yield(); //MyOwnCleanupAsync(); // Note: This will still make use of the current transaction scope.

​

                if (this.FakeTransaction.Value != "Did some work")

                    throw new InvalidOperationException($"Expected 'Did some work', got: '{this.FakeTransaction.Value}'");

​

                this.FakeTransaction.Value = "Committed";

            }

            finally

            {

                this.UpdateExecutionContext();

            }

        }

    }

}

​