retry insert active key on tx serialization failure

coder · johnstcn · Sep 7, 2023 · Sep 4, 2023 · Sep 4, 2023 · Sep 5, 2023
commit 381f0783d533e6a8dfd745c24c122e700ce7f0f1
diff --git a/enterprise/dbcrypt/dbcrypt.go b/enterprise/dbcrypt/dbcrypt.go
@@ -49,7 +49,8 @@ func New(ctx context.Context, db database.Store, ciphers ...Cipher) (database.St
 		Store:               db,
 	}
 	// nolint: gocritic // This is allowed.
-	if err := dbc.ensureEncrypted(dbauthz.AsSystemRestricted(ctx)); err != nil {
+	authCtx := dbauthz.AsSystemRestricted(ctx)
+	if err := dbc.ensureEncryptedWithRetry(authCtx); err != nil {
 		return nil, xerrors.Errorf("ensure encrypted database fields: %w", err)
 	}
 	return dbc, nil
@@ -334,6 +335,22 @@ func (db *dbCrypt) decryptField(field *string, digest sql.NullString) error {
 	return nil
 }
 
+func (db *dbCrypt) ensureEncryptedWithRetry(ctx context.Context) error {
+	err := db.ensureEncrypted(ctx)
+	if err == nil {
+		return nil
+	}
+	// If we get a serialization error, then we need to retry.
+	var pqerr *pq.Error
+	if !xerrors.As(err, &pqerr) {
+		return err
+	}
+	if pqerr.Code != "40001" { // serialization_failure
+		return err
+	}
+	return db.ensureEncrypted(ctx)
+}
+
 func (db *dbCrypt) ensureEncrypted(ctx context.Context) error {
 	return db.InTx(func(s database.Store) error {
 		// Attempt to read the encrypted test fields of the currently active keys.
@@ -364,21 +381,10 @@ func (db *dbCrypt) ensureEncrypted(ctx context.Context) error {
 		}
 
 		// If we get here, then we have a new key that we need to insert.
-		// If this conflicts with another transaction, we do not need to retry as
-		// the other transaction will have inserted the key for us.
-		if err := db.InsertDBCryptKey(ctx, database.InsertDBCryptKeyParams{
+		return db.InsertDBCryptKey(ctx, database.InsertDBCryptKeyParams{
 			Number:          highestNumber + 1,
 			ActiveKeyDigest: db.primaryCipherDigest,
 			Test:            testValue,
-		}); err != nil {
-			var pqErr *pq.Error
-			if xerrors.As(err, &pqErr) && pqErr.Code == "23505" {
-				// Unique constraint violation -> another transaction has inserted the key for us.
-				return nil
-			}
-			return err
-		}
-
-		return nil
+		})
 	}, &sql.TxOptions{Isolation: sql.LevelRepeatableRead})
 }
diff --git a/enterprise/dbcrypt/dbcrypt_internal_test.go b/enterprise/dbcrypt/dbcrypt_internal_test.go
@@ -8,10 +8,13 @@ import (
 	"io"
 	"testing"
 
+	"github.com/golang/mock/gomock"
+	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbmock"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 )
 
@@ -470,6 +473,46 @@ func TestNew(t *testing.T) {
 		require.Error(t, err)
 		require.ErrorContains(t, err, "has been revoked")
 	})
+
+	t.Run("Retry", func(t *testing.T) {
+		t.Parallel()
+		// Given: a cipher is loaded
+		cipher := initCipher(t)
+		ctx, cancel := context.WithCancel(context.Background())
+		testVal, err := cipher.Encrypt([]byte("coder"))
+		key := database.DBCryptKey{
+			Number:          1,
+			ActiveKeyDigest: sql.NullString{String: cipher.HexDigest(), Valid: true},
+			Test:            b64encode(testVal),
+		}
+		require.NoError(t, err)
+		t.Cleanup(cancel)
+
+		// And: a database that returns an error once when we try to serialize a key
+		ctrl := gomock.NewController(t)
+		mockDB := dbmock.NewMockStore(ctrl)
+
+		gomock.InOrder(
+			// First try: we get a serialization error.
+			expectTx(mockDB),
+			mockDB.EXPECT().GetDBCryptKeys(gomock.Any()).Times(1).Return([]database.DBCryptKey{}, nil),
+			mockDB.EXPECT().InsertDBCryptKey(gomock.Any(), gomock.Any()).Times(1).Return(&pq.Error{Code: "40001"}),
+			// Second try: we get the key we wanted to insert initially.
+			expectTx(mockDB),
+			mockDB.EXPECT().GetDBCryptKeys(gomock.Any()).Times(1).Return([]database.DBCryptKey{key}, nil),
+		)
+
+		_, err = New(ctx, mockDB, cipher)
+		require.NoError(t, err)
+	})
+}
+
+func expectTx(mdb *dbmock.MockStore) *gomock.Call {
+	return mdb.EXPECT().InTx(gomock.Any(), gomock.Any()).Times(1).DoAndReturn(
+		func(f func(store database.Store) error, _ *sql.TxOptions) error {
+			return f(mdb)
+		},
+	)
 }
 
 func requireEncryptedEquals(t *testing.T, c Cipher, value, expected string) {
@@ -511,3 +554,11 @@ func fakeBase64RandomData(t *testing.T, n int) string {
 	require.NoError(t, err)
 	return base64.StdEncoding.EncodeToString(b)
 }
+
+func withInTx(mTx *dbmock.MockStore) {
+	mTx.EXPECT().InTx(gomock.Any(), gomock.Any()).Times(1).DoAndReturn(
+		func(f func(store database.Store) error, _ *sql.TxOptions) error {
+			return f(mTx)
+		},
+	)
+}