fix: avoid race between replicas on start (#12344)

deansheather · johnstcn · web-flow · commit bedd2c592284 · 2024-02-28T16:14:11.000Z
DERP mesh key setup would do a SELECT and then an INSERT on failure, without a lock. During some testing with multiple replicas, I managed to cause a replica to crash due to them initializing simultaneously.

Fixes:

Encountered an error running "coder server"
create coder API: insert mesh key: pq: duplicate key value violates unique constraint "site_configs_key_key"

Co-authored-by: Cian Johnston &lt;cian@coder.com&gt;
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1874,6 +1874,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 		check.Args(u.ID).Asserts(rbac.ResourceSystem, rbac.ActionRead)
 	}))
 	s.Run("GetDERPMeshKey", s.Subtest(func(db database.Store, check *expects) {
+		db.InsertDERPMeshKey(context.Background(), "testing")
 		check.Args().Asserts(rbac.ResourceSystem, rbac.ActionRead)
 	}))
 	s.Run("InsertDERPMeshKey", s.Subtest(func(db database.Store, check *expects) {
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1761,6 +1761,9 @@ func (q *FakeQuerier) GetDERPMeshKey(_ context.Context) (string, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
+	if q.derpMeshKey == "" {
+		return "", sql.ErrNoRows
+	}
 	return q.derpMeshKey, nil
 }
 
diff --git a/coderd/database/lock.go b/coderd/database/lock.go
@@ -9,6 +9,7 @@ const (
 	// Keep the unused iota here so we don't need + 1 every time
 	lockIDUnused = iota
 	LockIDDeploymentSetup
+	LockIDEnterpriseDeploymentSetup
 )
 
 // GenLockID generates a unique and consistent lock ID from a given string.
diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go
@@ -15,6 +15,7 @@ import (
 	"tailscale.com/types/key"
 
 	"github.com/coder/coder/v2/cli/clibase"
+	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/enterprise/audit"
 	"github.com/coder/coder/v2/enterprise/audit/backends"
@@ -37,21 +38,41 @@ func (r *RootCmd) Server(_ func()) *clibase.Cmd {
 		}
 
 		options.DERPServer = derp.NewServer(key.NewNode(), tailnet.Logger(options.Logger.Named("derp")))
-		meshKey, err := options.Database.GetDERPMeshKey(ctx)
-		if err != nil {
+
+		var meshKey string
+		err := options.Database.InTx(func(tx database.Store) error {
+			// This will block until the lock is acquired, and will be
+			// automatically released when the transaction ends.
+			err := tx.AcquireLock(ctx, database.LockIDEnterpriseDeploymentSetup)
+			if err != nil {
+				return xerrors.Errorf("acquire lock: %w", err)
+			}
+
+			meshKey, err = tx.GetDERPMeshKey(ctx)
+			if err == nil {
+				return nil
+			}
 			if !errors.Is(err, sql.ErrNoRows) {
-				return nil, nil, xerrors.Errorf("get mesh key: %w", err)
+				return xerrors.Errorf("get DERP mesh key: %w", err)
 			}
 			meshKey, err = cryptorand.String(32)
 			if err != nil {
-				return nil, nil, xerrors.Errorf("generate mesh key: %w", err)
+				return xerrors.Errorf("generate DERP mesh key: %w", err)
 			}
-			err = options.Database.InsertDERPMeshKey(ctx, meshKey)
+			err = tx.InsertDERPMeshKey(ctx, meshKey)
 			if err != nil {
-				return nil, nil, xerrors.Errorf("insert mesh key: %w", err)
+				return xerrors.Errorf("insert DERP mesh key: %w", err)
 			}
+			return nil
+		}, nil)
+		if err != nil {
+			return nil, nil, err
+		}
+		if meshKey == "" {
+			return nil, nil, xerrors.New("mesh key is empty")
 		}
 		options.DERPServer.SetMeshKey(meshKey)
+
 		options.Auditor = audit.NewAuditor(
 			options.Database,
 			audit.DefaultFilter,

Original file line number	Diff line number	Diff line change
`@@ -1761,6 +1761,9 @@ func (q *FakeQuerier) GetDERPMeshKey(_ context.Context) (string, error) {`
`1761`	`1761`	`q.mutex.RLock()`
`1762`	`1762`	`defer q.mutex.RUnlock()`
`1763`	`1763`
	`1764`	`+ if q.derpMeshKey == "" {`
	`1765`	`+ return "", sql.ErrNoRows`
	`1766`	`+ }`
`1764`	`1767`	`return q.derpMeshKey, nil`
`1765`	`1768`	`}`
`1766`	`1769`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ const (`
`9`	`9`	`// Keep the unused iota here so we don't need + 1 every time`
`10`	`10`	`lockIDUnused = iota`
`11`	`11`	`LockIDDeploymentSetup`
	`12`	`+ LockIDEnterpriseDeploymentSetup`
`12`	`13`	`)`
`13`	`14`
`14`	`15`	`// GenLockID generates a unique and consistent lock ID from a given string.`