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

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;

    /// <summary>
    ///     Initializes a new instance of the <see cref="BenchmarkTransactionHolder" /> class.
    /// </summary>
    /// <param name="storage">The storage engine used to create transactions.</param>
    public BenchmarkTransactionHolder(StorageEngine storage)
    {
        _storage = storage ?? throw new ArgumentNullException(nameof(storage));
    }

    /// <summary>
    ///     Disposes this holder and rolls back any outstanding transaction.
    /// </summary>
    public void Dispose()
    {
        RollbackAndReset();
    }

    /// <summary>
    ///     Gets the current active transaction or starts a new one.
    /// </summary>
    /// <returns>The current active transaction.</returns>
    public ITransaction GetCurrentTransactionOrStart()
    {
        lock (_sync)
        {
            if (_currentTransaction == null || _currentTransaction.State != TransactionState.Active)
                _currentTransaction = _storage.BeginTransaction();

            return _currentTransaction;
        }
    }

    /// <summary>
    ///     Gets the current active transaction or starts a new one asynchronously.
    /// </summary>
    /// <returns>A task that returns the current active transaction.</returns>
    public Task<ITransaction> GetCurrentTransactionOrStartAsync()
    {
        return Task.FromResult(GetCurrentTransactionOrStart());
    }

    /// <summary>
    ///     Commits the current transaction when active and clears the holder.
    /// </summary>
    public void CommitAndReset()
    {
        lock (_sync)
        {
            if (_currentTransaction == null) return;

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

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

    /// <summary>
    ///     Rolls back the current transaction when active and clears the holder.
    /// </summary>
    public void RollbackAndReset()
    {
        lock (_sync)
        {
            if (_currentTransaction == null) return;

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

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