Seralizable transactions should retry automatically

Emyrk · Emyrk · commit a8e256ac499f · 2023-03-21T16:30:39.000-05:00
diff --git a/coderd/database/db.go b/coderd/database/db.go
@@ -67,8 +67,42 @@ func (q *sqlQuerier) Ping(ctx context.Context) (time.Duration, error) {
 	return time.Since(start), err
 }
 
-// InTx performs database operations inside a transaction.
 func (q *sqlQuerier) InTx(function func(Store) error, txOpts *sql.TxOptions) error {
+	_, inTx := q.db.(*sqlx.Tx)
+	isolation := sql.LevelDefault
+	if txOpts != nil {
+		isolation = txOpts.Isolation
+	}
+
+	// If we are not already in a transaction, and we are running in serializable
+	// mode, we need to run the transaction in a retry loop. The caller should be
+	// prepared to allow retries if using serializable mode.
+	// If we are in a transaction already, the parent InTx call will handle the retry.
+	// We do not want to duplicate those retries.
+	if !inTx && isolation == sql.LevelSerializable {
+		// This is an arbitrarily chosen number.
+		const retryAmount = 3
+		var err error
+		attempts := 0
+		for attempts = 0; attempts < retryAmount; attempts++ {
+			err = q.runTx(function, txOpts)
+			if err == nil {
+				// Transaction succeeded.
+				return nil
+			}
+			if err != nil && !IsSerializedError(err) {
+				// We should only retry if the error is a serialization error.
+				return err
+			}
+		}
+		// Transaction kept failing in serializable mode.
+		return xerrors.Errorf("transaction failed after %d attempts: %w", attempts, err)
+	}
+	return q.runTx(function, txOpts)
+}
+
+// InTx performs database operations inside a transaction.
+func (q *sqlQuerier) runTx(function func(Store) error, txOpts *sql.TxOptions) error {
 	if _, ok := q.db.(*sqlx.Tx); ok {
 		// If the current inner "db" is already a transaction, we just reuse it.
 		// We do not need to handle commit/rollback as the outer tx will handle
diff --git a/coderd/database/db_test.go b/coderd/database/db_test.go
@@ -8,13 +8,44 @@ import (
 	"testing"
 
 	"github.com/google/uuid"
+	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/migrations"
 	"github.com/coder/coder/coderd/database/postgres"
 )
 
+func TestSerializedRetry(t *testing.T) {
+	t.Parallel()
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	sqlDB := testSQLDB(t)
+	db := database.New(sqlDB)
+
+	called := 0
+	txOpts := &sql.TxOptions{Isolation: sql.LevelSerializable}
+	err := db.InTx(func(tx database.Store) error {
+		// Test nested error
+		return tx.InTx(func(tx database.Store) error {
+			// The easiest way to mock a serialization failure is to
+			// return a serialization failure error.
+			called++
+			return &pq.Error{
+				Code:    "40001",
+				Message: "serialization_failure",
+			}
+		}, txOpts)
+	}, txOpts)
+	require.Error(t, err, "should fail")
+	// The double "execute transaction: execute transaction" is from the nested transactions.
+	// Just want to make sure we don't try 9 times.
+	require.Equal(t, err.Error(), "transaction failed after 3 attempts: execute transaction: execute transaction: pq: serialization_failure", "error message")
+	require.Equal(t, called, 3, "should retry 3 times")
+}
+
 func TestNestedInTx(t *testing.T) {
 	t.Parallel()
 	if testing.Short() {
diff --git a/coderd/database/errors.go b/coderd/database/errors.go
@@ -6,6 +6,14 @@ import (
 	"github.com/lib/pq"
 )
 
+func IsSerializedError(err error) bool {
+	var pqErr *pq.Error
+	if errors.As(err, &pqErr) {
+		return pqErr.Code.Name() == "serialization_failure"
+	}
+	return false
+}
+
 // IsUniqueViolation checks if the error is due to a unique violation.
 // If one or more specific unique constraints are given as arguments,
 // the error must be caused by one of them. If no constraints are given,
