From ccf549a62d228e6852c2e6448c21729fd41a2172 Mon Sep 17 00:00:00 2001
From: Mars Hall <mars@heroku.com>
Date: Tue, 5 Mar 2019 11:19:56 -0800
Subject: [PATCH] Stricter pg backend locking during workspace creation

---
 backend/remote-state/pg/client.go | 35 +++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/backend/remote-state/pg/client.go b/backend/remote-state/pg/client.go
index 52e629860..3c94e88f3 100644
--- a/backend/remote-state/pg/client.go
+++ b/backend/remote-state/pg/client.go
@@ -72,15 +72,27 @@ func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
 		info.ID = lockID
 	}
 
-	// Try to acquire lock for the existing row.
-	query := `SELECT %s.id, pg_try_advisory_lock(%s.id) FROM %s.%s WHERE %s.name = $1`
+	// Local helper function so we can call it multiple places
+	//
+	lockUnlock := func(pgLockId string) error {
+		query := `SELECT pg_advisory_unlock(%s)`
+		row := c.Client.QueryRow(fmt.Sprintf(query, pgLockId))
+		var didUnlock []byte
+		err := row.Scan(&didUnlock)
+		if err != nil {
+			return &state.LockError{Info: c.info, Err: err}
+		}
+		return nil
+	}
+
+	// Try to acquire locks for the existing row `id` and the creation lock `-1`.
+	query := `SELECT %s.id, pg_try_advisory_lock(%s.id), pg_try_advisory_lock(-1) FROM %s.%s WHERE %s.name = $1`
 	row := c.Client.QueryRow(fmt.Sprintf(query, statesTableName, statesTableName, c.SchemaName, statesTableName, statesTableName), c.Name)
-	var pgLockId, didLock []byte
-	err = row.Scan(&pgLockId, &didLock)
+	var pgLockId, didLock, didLockForCreate []byte
+	err = row.Scan(&pgLockId, &didLock, &didLockForCreate)
 	switch {
 	case err == sql.ErrNoRows:
-		// When the row does not yet exist in state, take
-		// the `-1` lock to create the new row.
+		// No rows means we're creating the workspace. Take the creation lock.
 		innerRow := c.Client.QueryRow(`SELECT pg_try_advisory_lock(-1)`)
 		var innerDidLock []byte
 		err := innerRow.Scan(&innerDidLock)
@@ -94,8 +106,16 @@ func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
 	case err != nil:
 		return "", &state.LockError{Info: info, Err: err}
 	case string(didLock) == "false":
+		// Existing workspace is already locked. Release the attempted creation lock.
+		lockUnlock("-1")
 		return "", &state.LockError{Info: info, Err: fmt.Errorf("Workspace is already locked: %s", c.Name)}
+	case string(didLockForCreate) == "false":
+		// Someone has the creation lock already. Release the existing workspace because it might not be safe to touch.
+		lockUnlock(string(pgLockId))
+		return "", &state.LockError{Info: info, Err: fmt.Errorf("Cannot lock workspace; already locked for workspace creation: %s", c.Name)}
 	default:
+		// Existing workspace is now locked. Release the attempted creation lock.
+		lockUnlock("-1")
 		info.Path = string(pgLockId)
 	}
 	c.info = info
@@ -116,9 +136,6 @@ func (c *RemoteClient) Unlock(id string) error {
 		if err != nil {
 			return &state.LockError{Info: c.info, Err: err}
 		}
-		if string(didUnlock) == "false" {
-			return &state.LockError{Info: c.info, Err: fmt.Errorf("Workspace is already unlocked: %s", c.Name)}
-		}
 		c.info = nil
 	}
 	return nil