chore: increase parallelism of TestWorkspaceQuota (#6710)

kylecarbs · web-flow · commit abe1e89f806f · 2023-03-21T22:44:01.000Z
This does a lot of build operations, so having multiple provisioner
daemons is great.

We were actually approaching the ceiling here for test duration!
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,
diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go
@@ -2,9 +2,11 @@ package coderd_test
 
 import (
 	"context"
+	"sync"
 	"testing"
 
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/coderd/coderdtest"
@@ -37,12 +39,12 @@ func TestWorkspaceQuota(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 		max := 1
-		client := coderdenttest.New(t, &coderdenttest.Options{
+		client, _, api := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
 			UserWorkspaceQuota: max,
-			Options: &coderdtest.Options{
-				IncludeProvisionerDaemon: true,
-			},
 		})
+		coderdtest.NewProvisionerDaemon(t, api.AGPL)
+		coderdtest.NewProvisionerDaemon(t, api.AGPL)
+		coderdtest.NewProvisionerDaemon(t, api.AGPL)
 
 		user := coderdtest.CreateFirstUser(t, client)
 		coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
@@ -104,12 +106,18 @@ func TestWorkspaceQuota(t *testing.T) {
 		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 
 		// Spin up three workspaces fine
+		var wg sync.WaitGroup
 		for i := 0; i < 3; i++ {
-			workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
-			build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
-			verifyQuota(ctx, t, client, i+1, 3)
-			require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+				build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+				assert.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+			}()
 		}
+		wg.Wait()
+		verifyQuota(ctx, t, client, 3, 3)
 
 		// Next one must fail
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
