Retryable Transaction Implementation for Doctrine ORM

TransactionManager.php

<?php declare(strict_types = 1);

namespace ShipMonk\QueryGuard\Framework\DoctrineOrm;

use Closure;
use Doctrine\DBAL\Exception\RetryableException;
use Doctrine\DBAL\TransactionIsolationLevel;
use Doctrine\ORM\EntityManager;
use Doctrine\ORM\EntityManagerInterface;
use Doctrine\ORM\OptimisticLockException;
use Doctrine\ORM\UnitOfWork;
use Psr\Log\LoggerInterface;
use ShipMonk\QueryGuard\Framework\Exception\LogicException;
use Symfony\Component\Clock\Clock;
use Throwable;
use function random_int;

final class TransactionManager
{

    private const int TRANSACTION_RETRY_LIMIT = 3;

    /**
     * @var list<Closure(): void>
     */
    private array $afterCommitCallbacks = [];

    private bool $afterCommitCallbacksOpenForModification = false;

    public function __construct(
        private EntityManagerInterface $entityManager,
        private LoggerInterface $logger,
    )
    {
    }

    /**
     * @param Closure(): void $callback
     */
    public function afterCommit(Closure $callback): void
    {
        if ($this->afterCommitCallbacksOpenForModification) {
            $this->afterCommitCallbacks[] = $callback;

        } else {
            throw new LogicException('Cannot add more callbacks to transaction.');
        }
    }

    /**
     * @param Closure(): T $callback
     * @return T
     *
     * @template T
     *
     * @param-immediately-invoked-callable $callback
     */
    public function retryableTransaction(Closure $callback): mixed
    {
        if ($this->entityManager->getConnection()->isTransactionActive()) {
            throw new LogicException('TransactionManager::retryableTransaction() cannot be called when there is an active transaction.');
        }

        if ($this->entityManager->getUnitOfWork()->size() > 0) {
            throw new LogicException('TransactionManager::retryableTransaction() can only be called on a clean EntityManager.');
        }

        $this->entityManager->clear();
        $this->entityManager->getConnection()->setTransactionIsolation(TransactionIsolationLevel::REPEATABLE_READ);
        $failureCount = 0;

        while (true) {
            try {
                $this->afterCommitCallbacksOpenForModification = true;
                $result = $this->entityManager->wrapInTransaction($callback);
                $this->afterCommitCallbacksOpenForModification = false;

                foreach ($this->afterCommitCallbacks as $afterCommitCallback) {
                    $afterCommitCallback();
                }

                return $result;

            } catch (Throwable $e) {
                $failureCount++;
                $this->clearAndReopen();

                if ($failureCount === self::TRANSACTION_RETRY_LIMIT || !$this->shouldRetry($e)) {
                    throw $e;

                } else {
                    $this->logger->warning('Transaction failed with retryable exception, retrying', ['exception' => $e, 'failure_count' => $failureCount]);
                    Clock::get()->sleep(0.1 * (2 ** $failureCount) * random_int(90, 110) / 100);
                }

            } finally {
                $this->afterCommitCallbacks = [];
                $this->afterCommitCallbacksOpenForModification = false;
                $this->entityManager->clear();
            }
        }
    }

    private function shouldRetry(Throwable $e): bool
    {
        return $e instanceof RetryableException || $e instanceof OptimisticLockException;
    }

    private function clearAndReopen(): void
    {
        $this->entityManager->clear();

        (Closure::bind(
            function (): void {
                $this->persisters = [];
            },
            $this->entityManager->getUnitOfWork(),
            UnitOfWork::class,
        ))();

        (Closure::bind(
            function (): void {
                $this->closed = false;
            },
            $this->entityManager,
            EntityManager::class,
        ))();
    }

}