feat(database): add transaction retry attempts (#10197)

Author: memleakd

system/Database/BaseConnection.php

@@ -18,6 +18,7 @@
 use CodeIgniter\Database\Exceptions\RetryableTransactionException;
 use CodeIgniter\Database\Exceptions\UniqueConstraintViolationException;
 use CodeIgniter\Events\Events;
+use CodeIgniter\Exceptions\InvalidArgumentException;
 use CodeIgniter\I18n\Time;
 use Exception;
 use ReflectionClass;
@@ -228,6 +229,11 @@ abstract class BaseConnection implements ConnectionInterface
      */
     protected ?DatabaseException $lastException = null;
 
+    /**
+     * The first database exception that caused the current transaction to fail.
+     */
+    protected ?DatabaseException $transFailureException = null;
+
     /**
      * Connection ID
      *
@@ -860,7 +866,7 @@ public function query(string $sql, $binds = null, bool $setEscapeFlags = true, s
             $query->setDuration($startTime, $startTime);
 
             // This will trigger a rollback if transactions are being used
-            $this->handleTransStatus();
+            $this->handleTransStatus($exception ?? $this->lastException);
 
             if (
                 $this->DBDebug
@@ -1082,44 +1088,67 @@ public function afterRollback(callable $callback): static
      * @template TReturn
      *
      * @param callable(self): TReturn $callback
+     * @param positive-int            $attempts
      *
      * @return false|TReturn
      */
-    public function transaction(callable $callback): mixed
+    public function transaction(callable $callback, int $attempts = 1): mixed
     {
+        if ($attempts < 1) {
+            throw new InvalidArgumentException('Transaction attempts must be a positive integer.');
+        }
+
         if (! $this->transEnabled) {
             return $callback($this);
         }
 
-        if (! $this->transBegin()) {
-            return false;
-        }
+        $attempts = $this->transDepth === 0 ? $attempts : 1;
+
+        for ($attempt = 1; $attempt <= $attempts; $attempt++) {
+            if (! $this->transBegin()) {
+                return false;
+            }
 
-        try {
-            $result = $callback($this);
-        } catch (Throwable $e) {
             try {
-                $this->transRollback();
-            } catch (Throwable $rollbackException) {
-                log_message('error', 'Database: Transaction callback threw an exception before rollback failed: ' . $e);
+                $result = $callback($this);
+            } catch (Throwable $e) {
+                try {
+                    $this->transRollback();
+                } catch (Throwable $rollbackException) {
+                    log_message('error', 'Database: Transaction callback threw an exception before rollback failed: ' . $e);
+
+                    throw $rollbackException;
+                } finally {
+                    if ($this->transDepth > 0) {
+                        $this->transStatus = false;
+                    } elseif ($this->transStrict === false) {
+                        $this->transStatus = true;
+                    }
+                }
 
-                throw $rollbackException;
-            } finally {
-                if ($this->transDepth > 0) {
-                    $this->transStatus = false;
-                } elseif ($this->transStrict === false) {
-                    $this->transStatus = true;
+                if ($this->transDepth === 0 && $e instanceof RetryableTransactionException && $attempt < $attempts) {
+                    $this->prepareTransactionRetry();
+
+                    continue;
                 }
+
+                throw $e;
             }
 
-            throw $e;
-        }
+            if (! $this->transComplete()) {
+                if ($this->transDepth === 0 && $this->transFailureException instanceof RetryableTransactionException && $attempt < $attempts) {
+                    $this->prepareTransactionRetry();
 
-        if (! $this->transComplete()) {
-            return false;
+                    continue;
+                }
+
+                return false;
+            }
+
+            return $result;
         }
 
-        return $result;
+        return false;
     }
 
     /**
@@ -1145,7 +1174,8 @@ public function transBegin(bool $testMode = false): bool
         // Reset the transaction failure flag.
         // If the $testMode flag is set to TRUE transactions will be rolled back
         // even if the queries produce a successful result.
-        $this->transFailure = $testMode;
+        $this->transFailure          = $testMode;
+        $this->transFailureException = null;
 
         if ($this->_transBegin()) {
             $this->transDepth++;
@@ -1219,13 +1249,24 @@ public function resetTransStatus(): static
      *
      * @internal This method is for internal database component use only
      */
-    public function handleTransStatus(): void
+    public function handleTransStatus(?DatabaseException $exception = null): void
     {
         if ($this->transDepth !== 0) {
             $this->transStatus = false;
+            $this->transFailureException ??= $exception;
         }
     }
 
+    /**
+     * Reset transaction state that should not leak into a retry attempt.
+     */
+    protected function prepareTransactionRetry(): void
+    {
+        $this->transStatus           = true;
+        $this->transFailureException = null;
+        $this->lastException         = null;
+    }
+
     /**
      * Run and clear callbacks registered for a successful transaction commit.
      */

system/Database/BasePreparedQuery.php

@@ -141,11 +141,11 @@ public function execute(...$data)
         if ($result === false) {
             $query->setDuration($startTime, $startTime);
 
-            // This will trigger a rollback if transactions are being used
-            $this->db->handleTransStatus();
-
             $databaseException = $this->createDatabaseException($exception);
 
+            // This will trigger a rollback if transactions are being used
+            $this->db->handleTransStatus($databaseException);
+
             if ($this->db->DBDebug) {
                 // We call this function in order to roll-back queries
                 // if transactions are enabled. If we don't call this here

system/Database/ConnectionInterface.php

@@ -144,10 +144,11 @@ public function afterRollback(callable $callback): static;
      * @template TReturn
      *
      * @param callable(self): TReturn $callback
+     * @param positive-int            $attempts
      *
      * @return false|TReturn
      */
-    public function transaction(callable $callback): mixed;
+    public function transaction(callable $callback, int $attempts = 1): mixed;
 
     /**
      * Returns an instance of the query builder for this connection.