Remove TryLockPull, use TryLock per directory instead

We generate a list of all interesting directories, so we can target
the locks to the affected directories instead of using a (too) global lock

server/events/project_command_builder.go

@@ -732,14 +732,6 @@ func (p *DefaultProjectCommandBuilder) getCfg(ctx *command.Context, projectName
 // buildAllProjectCommandsByPlan builds contexts for a command for every project that has
 // pending plans in this ctx.
 func (p *DefaultProjectCommandBuilder) buildAllProjectCommandsByPlan(ctx *command.Context, commentCmd *CommentCommand) ([]command.ProjectContext, error) {
-	// Lock all dirs in this pull request (instead of a single dir) because we
-	// don't know how many dirs we'll need to run the command in.
-	unlockFn, err := p.WorkingDirLocker.TryLockPull(ctx.Pull.BaseRepo.FullName, ctx.Pull.Num)
-	if err != nil {
-		return nil, err
-	}
-	defer unlockFn()
-
 	pullDir, err := p.WorkingDir.GetPullDir(ctx.Pull.BaseRepo, ctx.Pull)
 	if err != nil {
 		return nil, err
@@ -759,6 +751,12 @@ func (p *DefaultProjectCommandBuilder) buildAllProjectCommandsByPlan(ctx *comman
 
 	var cmds []command.ProjectContext
 	for _, plan := range plans {
+		// Lock all the directories we need to run the command in
+		unlockFn, err := p.WorkingDirLocker.TryLock(ctx.Pull.BaseRepo.FullName, ctx.Pull.Num, plan.Workspace, plan.RepoRelDir)
+		if err != nil {
+			return nil, err
+		}
+		defer unlockFn()
 		commentCmds, err := p.buildProjectCommandCtx(ctx, commentCmd.CommandName(), commentCmd.SubName, plan.ProjectName, commentCmd.Flags, defaultRepoDir, plan.RepoRelDir, plan.Workspace, commentCmd.Verbose)
 		if err != nil {
 			return nil, errors.Wrapf(err, "building command for dir '%s'", plan.RepoRelDir)

server/events/working_dir_locker.go

@@ -15,7 +15,6 @@ package events
 
 import (
 	"fmt"
-	"strings"
 	"sync"
 )
 
@@ -32,62 +31,33 @@ type WorkingDirLocker interface {
 	// an error if the workspace is already locked. The error is expected to
 	// be printed to the pull request.
 	TryLock(repoFullName string, pullNum int, workspace string, path string) (func(), error)
-	// TryLockPull tries to acquire a lock for all the workspaces in this repo
-	// and pull.
-	// It returns a function that should be used to unlock the workspace and
-	// an error if the workspace is already locked. The error is expected to
-	// be printed to the pull request.
-	TryLockPull(repoFullName string, pullNum int) (func(), error)
 }
 
 // DefaultWorkingDirLocker implements WorkingDirLocker.
 type DefaultWorkingDirLocker struct {
 	// mutex prevents against multiple threads calling functions on this struct
 	// concurrently. It's only used for entry/exit to each function.
 	mutex sync.Mutex
-	// locks is a list of the keys that are locked. We then use prefix
-	// matching to determine if something is locked. It's naive but that's okay
-	// because there won't be many locks at one time.
-	locks []string
+	// locks is a set of the keys that are locked.
+	locks map[string]struct{}
 }
 
 // NewDefaultWorkingDirLocker is a constructor.
 func NewDefaultWorkingDirLocker() *DefaultWorkingDirLocker {
-	return &DefaultWorkingDirLocker{}
-}
-
-func (d *DefaultWorkingDirLocker) TryLockPull(repoFullName string, pullNum int) (func(), error) {
-	d.mutex.Lock()
-	defer d.mutex.Unlock()
-
-	pullKey := d.pullKey(repoFullName, pullNum)
-	for _, l := range d.locks {
-		if l == pullKey || strings.HasPrefix(l, pullKey+"/") {
-			return func() {}, fmt.Errorf("the Atlantis working dir is currently locked by another" +
-				" command that is running for this pull request.\n" +
-				"Wait until the previous command is complete and try again")
-		}
-	}
-	d.locks = append(d.locks, pullKey)
-	return func() {
-		d.UnlockPull(repoFullName, pullNum)
-	}, nil
+	return &DefaultWorkingDirLocker{locks: make(map[string]struct{})}
 }
 
 func (d *DefaultWorkingDirLocker) TryLock(repoFullName string, pullNum int, workspace string, path string) (func(), error) {
 	d.mutex.Lock()
 	defer d.mutex.Unlock()
 
-	pullKey := d.pullKey(repoFullName, pullNum)
 	workspaceKey := d.workspaceKey(repoFullName, pullNum, workspace, path)
-	for _, l := range d.locks {
-		if l == pullKey || l == workspaceKey {
-			return func() {}, fmt.Errorf("the %s workspace at path %s is currently locked by another"+
-				" command that is running for this pull request.\n"+
-				"Wait until the previous command is complete and try again", workspace, path)
-		}
+	if _, exists := d.locks[workspaceKey]; exists {
+		return func() {}, fmt.Errorf("the %s workspace at path %s is currently locked by another"+
+			" command that is running for this pull request.\n"+
+			"Wait until the previous command is complete and try again", workspace, path)
 	}
-	d.locks = append(d.locks, workspaceKey)
+	d.locks[workspaceKey] = struct{}{}
 	return func() {
 		d.unlock(repoFullName, pullNum, workspace, path)
 	}, nil
@@ -99,32 +69,9 @@ func (d *DefaultWorkingDirLocker) unlock(repoFullName string, pullNum int, works
 	defer d.mutex.Unlock()
 
 	workspaceKey := d.workspaceKey(repoFullName, pullNum, workspace, path)
-	d.removeLock(workspaceKey)
-}
-
-// Unlock unlocks all workspaces for this pull.
-func (d *DefaultWorkingDirLocker) UnlockPull(repoFullName string, pullNum int) {
-	d.mutex.Lock()
-	defer d.mutex.Unlock()
-
-	pullKey := d.pullKey(repoFullName, pullNum)
-	d.removeLock(pullKey)
-}
-
-func (d *DefaultWorkingDirLocker) removeLock(key string) {
-	var newLocks []string
-	for _, l := range d.locks {
-		if l != key {
-			newLocks = append(newLocks, l)
-		}
-	}
-	d.locks = newLocks
+	delete(d.locks, workspaceKey)
 }
 
 func (d *DefaultWorkingDirLocker) workspaceKey(repo string, pull int, workspace string, path string) string {
-	return fmt.Sprintf("%s/%s/%s", d.pullKey(repo, pull), workspace, path)
-}
-
-func (d *DefaultWorkingDirLocker) pullKey(repo string, pull int) string {
-	return fmt.Sprintf("%s/%d", repo, pull)
+	return fmt.Sprintf("%s/%d/%s/%s", repo, pull, workspace, path)
 }

server/events/working_dir_locker_test.go

@@ -170,47 +170,3 @@ func TestUnlockDifferentPulls(t *testing.T) {
 	_, err = locker.TryLock(repo, newPull, workspace, path)
 	Ok(t, err)
 }
-
-func TestLockPull(t *testing.T) {
-	locker := events.NewDefaultWorkingDirLocker()
-	unlock, err := locker.TryLockPull("owner/repo", 1)
-	Ok(t, err)
-
-	// Now a lock for the same pull or for a workspace should fail.
-	_, err = locker.TryLockPull("owner/repo", 1)
-	Assert(t, err != nil, "exp err")
-	_, err = locker.TryLock("owner/repo", 1, "workspace", path)
-	Assert(t, err != nil, "exp err")
-
-	// Lock for a different pull and workspace should succeed.
-	_, err = locker.TryLockPull("owner/repo", 2)
-	Ok(t, err)
-	_, err = locker.TryLock("owner/repo", 3, "workspace", path)
-	Ok(t, err)
-
-	// After unlocking, should be able to get a pull lock.
-	unlock()
-	unlock, err = locker.TryLockPull("owner/repo", 1)
-	Ok(t, err)
-
-	// If we unlock that too, should be able to get the workspace lock.
-	unlock()
-	_, err = locker.TryLock("owner/repo", 1, "workspace", path)
-	Ok(t, err)
-	unlock()
-}
-
-// If the workspace was locked first, we shouldn't be able to get the pull lock.
-func TestLockPull_WorkspaceFirst(t *testing.T) {
-	locker := events.NewDefaultWorkingDirLocker()
-	unlock, err := locker.TryLock("owner/repo", 1, "workspace", path)
-	Ok(t, err)
-
-	_, err = locker.TryLockPull("owner/repo", 1)
-	Assert(t, err != nil, "exp err")
-
-	// After unlocking the workspace, should be able to get the lock.
-	unlock()
-	_, err = locker.TryLockPull("owner/repo", 1)
-	Ok(t, err)
-}