CBDD/tests/CBDD.Tests.Benchmark/BenchmarkTransactionHolder.cs

using ZB.MOM.WW.CBDD.Core;
using ZB.MOM.WW.CBDD.Core.Storage;
using ZB.MOM.WW.CBDD.Core.Transactions;

namespace ZB.MOM.WW.CBDD.Tests.Benchmark;

internal sealed class BenchmarkTransactionHolder : ITransactionHolder, IDisposable
{
    private readonly StorageEngine _storage;
    private readonly object _sync = new();
    private ITransaction? _currentTransaction;

    public BenchmarkTransactionHolder(StorageEngine storage)
    {
        _storage = storage ?? throw new ArgumentNullException(nameof(storage));
    }

    public ITransaction GetCurrentTransactionOrStart()
    {
        lock (_sync)
        {
            if (_currentTransaction == null || _currentTransaction.State != TransactionState.Active)
            {
                _currentTransaction = _storage.BeginTransaction();
            }

            return _currentTransaction;
        }
    }

    public Task<ITransaction> GetCurrentTransactionOrStartAsync()
    {
        return Task.FromResult(GetCurrentTransactionOrStart());
    }

    public void CommitAndReset()
    {
        lock (_sync)
        {
            if (_currentTransaction == null)
            {
                return;
            }

            if (_currentTransaction.State == TransactionState.Active ||
                _currentTransaction.State == TransactionState.Preparing)
            {
                _currentTransaction.Commit();
            }

            _currentTransaction.Dispose();
            _currentTransaction = null;
        }
    }

    public void RollbackAndReset()
    {
        lock (_sync)
        {
            if (_currentTransaction == null)
            {
                return;
            }

            if (_currentTransaction.State == TransactionState.Active ||
                _currentTransaction.State == TransactionState.Preparing)
            {
                _currentTransaction.Rollback();
            }

            _currentTransaction.Dispose();
            _currentTransaction = null;
        }
    }

    public void Dispose()
    {
        RollbackAndReset();
    }
}