Log streaming resource cleanup (#102)

server/controllers/events/events_controller_e2e_test.go

@@ -1075,11 +1075,12 @@ func setupE2E(t *testing.T, repoDir string) (events_controllers.VCSEventsControl
 		TestingMode:   true,
 		CommandRunner: commandRunner,
 		PullCleaner: &events.PullClosedExecutor{
-			Locker:             lockingClient,
-			VCSClient:          e2eVCSClient,
-			WorkingDir:         workingDir,
-			DB:                 boltdb,
-			PullClosedTemplate: &events.PullClosedEventTemplate{},
+			Locker:                   lockingClient,
+			VCSClient:                e2eVCSClient,
+			WorkingDir:               workingDir,
+			DB:                       boltdb,
+			PullClosedTemplate:       &events.PullClosedEventTemplate{},
+			LogStreamResourceCleaner: projectCmdOutputHandler,
 		},
 		Logger:                       logger,
 		Parser:                       eventParser,

server/events/models/fixtures/fixtures.go

@@ -13,14 +13,20 @@
 
 package fixtures
 
-import "github.com/runatlantis/atlantis/server/events/models"
+import (
+	"fmt"
+
+	"github.com/runatlantis/atlantis/server/events/models"
+	"github.com/runatlantis/atlantis/server/events/yaml/valid"
+)
 
 var Pull = models.PullRequest{
 	Num:        1,
 	HeadCommit: "16ca62f65c18ff456c6ef4cacc8d4826e264bb17",
 	HeadBranch: "branch",
 	Author:     "lkysow",
 	URL:        "url",
+	BaseRepo:   GithubRepo,
 }
 
 var GithubRepo = models.Repo{
@@ -50,3 +56,11 @@ var GitlabRepo = models.Repo{
 var User = models.User{
 	Username: "lkysow",
 }
+
+var projectName = "test-project"
+
+var Project = valid.Project{
+	Name: &projectName,
+}
+
+var PullInfo = fmt.Sprintf("%s/%d/%s", GithubRepo.FullName, Pull.Num, *Project.Name)

server/events/models/models.go

@@ -447,7 +447,11 @@ func (p ProjectCommandContext) GetShowResultFileName() string {
 
 // Gets a unique identifier for the current pull request as a single string
 func (p ProjectCommandContext) PullInfo() string {
-	return fmt.Sprintf("%s/%d/%s", p.BaseRepo.FullName, p.Pull.Num, p.ProjectName)
+	return BuildPullInfo(p.BaseRepo.FullName, p.Pull.Num, p.ProjectName)
+}
+
+func BuildPullInfo(repoName string, pullNum int, projectName string) string {
+	return fmt.Sprintf("%s/%d/%s", repoName, pullNum, projectName)
 }
 
 // SplitRepoFullName splits a repo full name up into its owner and repo

server/events/pull_closed_executor.go

@@ -29,6 +29,7 @@ import (
 	"github.com/runatlantis/atlantis/server/core/locking"
 	"github.com/runatlantis/atlantis/server/events/models"
 	"github.com/runatlantis/atlantis/server/events/vcs"
+	"github.com/runatlantis/atlantis/server/handlers"
 )
 
 //go:generate pegomock generate -m --use-experimental-model-gen --package mocks -o mocks/mock_pull_cleaner.go PullCleaner
@@ -43,12 +44,13 @@ type PullCleaner interface {
 // PullClosedExecutor executes the tasks required to clean up a closed pull
 // request.
 type PullClosedExecutor struct {
-	Locker             locking.Locker
-	VCSClient          vcs.Client
-	WorkingDir         WorkingDir
-	Logger             logging.SimpleLogging
-	DB                 *db.BoltDB
-	PullClosedTemplate PullCleanupTemplate
+	Locker                   locking.Locker
+	Logger                   logging.SimpleLogging
+	DB                       *db.BoltDB
+	PullClosedTemplate       PullCleanupTemplate
+	LogStreamResourceCleaner handlers.ResourceCleaner
+	VCSClient                vcs.Client
+	WorkingDir               WorkingDir
 }
 
 type templatedProject struct {
@@ -73,6 +75,24 @@ var pullClosedTemplate = template.Must(template.New("").Parse(
 
 // CleanUpPull cleans up after a closed pull request.
 func (p *PullClosedExecutor) CleanUpPull(repo models.Repo, pull models.PullRequest) error {
+	pullStatus, err := p.DB.GetPullStatus(pull)
+	if err != nil {
+		// Log and continue to clean up other resources.
+		p.Logger.Err("retrieving pull status: %s", err)
+	}
+
+	if pullStatus != nil {
+		for _, project := range pullStatus.Projects {
+			// TODO [ORCA-943]: Set projectName to "<dir>/<workspace>" when project name is not set.
+			// Upstream atlantis only requires project name to be set if there's more than one project
+			// with same dir and workspace. If a project name has not been set, we'll use the dir and
+			// workspace to build project key.
+			// Source: https://www.runatlantis.io/docs/repo-level-atlantis-yaml.html#reference
+			projectKey := models.BuildPullInfo(pullStatus.Pull.BaseRepo.FullName, pull.Num, project.ProjectName)
+			p.LogStreamResourceCleaner.CleanUp(projectKey)
+		}
+	}
+
 	if err := p.WorkingDir.Delete(repo, pull); err != nil {
 		return errors.Wrap(err, "cleaning workspace")
 	}

server/events/pull_closed_executor_test.go

@@ -14,10 +14,15 @@
 package events_test
 
 import (
-	"errors"
+	"io/ioutil"
 	"testing"
 
+	"github.com/pkg/errors"
 	"github.com/runatlantis/atlantis/server/core/db"
+	bolt "go.etcd.io/bbolt"
+
+	"github.com/runatlantis/atlantis/server/handlers"
+	"github.com/stretchr/testify/assert"
 
 	. "github.com/petergtz/pegomock"
 	lockmocks "github.com/runatlantis/atlantis/server/core/locking/mocks"
@@ -27,18 +32,25 @@ import (
 	"github.com/runatlantis/atlantis/server/events/models"
 	"github.com/runatlantis/atlantis/server/events/models/fixtures"
 	vcsmocks "github.com/runatlantis/atlantis/server/events/vcs/mocks"
+	handlermocks "github.com/runatlantis/atlantis/server/handlers/mocks"
+	loggermocks "github.com/runatlantis/atlantis/server/logging/mocks"
 	. "github.com/runatlantis/atlantis/testing"
 )
 
 func TestCleanUpPullWorkspaceErr(t *testing.T) {
 	t.Log("when workspace.Delete returns an error, we return it")
 	RegisterMockTestingT(t)
 	w := mocks.NewMockWorkingDir()
+	tmp, cleanup := TempDir(t)
+	defer cleanup()
+	db, err := db.New(tmp)
+	Ok(t, err)
 	pce := events.PullClosedExecutor{
 		WorkingDir:         w,
 		PullClosedTemplate: &events.PullClosedEventTemplate{},
+		DB:                 db,
 	}
-	err := errors.New("err")
+	err = errors.New("err")
 	When(w.Delete(fixtures.GithubRepo, fixtures.Pull)).ThenReturn(err)
 	actualErr := pce.CleanUpPull(fixtures.GithubRepo, fixtures.Pull)
 	Equals(t, "cleaning workspace: err", actualErr.Error())
@@ -49,12 +61,17 @@ func TestCleanUpPullUnlockErr(t *testing.T) {
 	RegisterMockTestingT(t)
 	w := mocks.NewMockWorkingDir()
 	l := lockmocks.NewMockLocker()
+	tmp, cleanup := TempDir(t)
+	defer cleanup()
+	db, err := db.New(tmp)
+	Ok(t, err)
 	pce := events.PullClosedExecutor{
 		Locker:             l,
 		WorkingDir:         w,
+		DB:                 db,
 		PullClosedTemplate: &events.PullClosedEventTemplate{},
 	}
-	err := errors.New("err")
+	err = errors.New("err")
 	When(l.UnlockByPull(fixtures.GithubRepo.FullName, fixtures.Pull.Num)).ThenReturn(nil, err)
 	actualErr := pce.CleanUpPull(fixtures.GithubRepo, fixtures.Pull)
 	Equals(t, "cleaning up locks: err", actualErr.Error())
@@ -72,7 +89,6 @@ func TestCleanUpPullNoLocks(t *testing.T) {
 	Ok(t, err)
 	pce := events.PullClosedExecutor{
 		Locker:             l,
-		VCSClient:          cp,
 		WorkingDir:         w,
 		DB:                 db,
 		PullClosedTemplate: &events.PullClosedEventTemplate{},
@@ -150,18 +166,18 @@ func TestCleanUpPullComments(t *testing.T) {
 	}
 	for _, c := range cases {
 		func() {
-			w := mocks.NewMockWorkingDir()
 			cp := vcsmocks.NewMockClient()
 			l := lockmocks.NewMockLocker()
+			w := mocks.NewMockWorkingDir()
 			tmp, cleanup := TempDir(t)
 			defer cleanup()
 			db, err := db.New(tmp)
 			Ok(t, err)
 			pce := events.PullClosedExecutor{
 				Locker:             l,
+				DB:                 db,
 				VCSClient:          cp,
 				WorkingDir:         w,
-				DB:                 db,
 				PullClosedTemplate: &events.PullClosedEventTemplate{},
 			}
 			t.Log("testing: " + c.Description)
@@ -175,3 +191,99 @@ func TestCleanUpPullComments(t *testing.T) {
 		}()
 	}
 }
+
+func TestCleanUpLogStreaming(t *testing.T) {
+	RegisterMockTestingT(t)
+
+	t.Run("Should Clean Up Log Streaming Resources When PR is closed", func(t *testing.T) {
+		prjStatusUpdater := handlermocks.NewMockProjectStatusUpdater()
+		prjJobURLGenerator := handlermocks.NewMockProjectJobURLGenerator()
+
+		// Create Log streaming resources
+		prjCmdOutput := make(chan *models.ProjectCmdOutputLine)
+		prjCmdOutHandler := handlers.NewAsyncProjectCommandOutputHandler(prjCmdOutput, prjStatusUpdater, prjJobURLGenerator, logger)
+		ctx := models.ProjectCommandContext{
+			BaseRepo:    fixtures.GithubRepo,
+			Pull:        fixtures.Pull,
+			ProjectName: *fixtures.Project.Name,
+		}
+
+		go prjCmdOutHandler.Handle()
+		prjCmdOutHandler.Send(ctx, "Test Message")
+
+		// Create boltdb and add pull request.
+		var lockBucket = "bucket"
+		var configBucket = "configBucket"
+		var pullsBucketName = "pulls"
+
+		f, err := ioutil.TempFile("", "")
+		if err != nil {
+			panic(errors.Wrap(err, "failed to create temp file"))
+		}
+		path := f.Name()
+		f.Close() // nolint: errcheck
+
+		// Open the database.
+		boltDB, err := bolt.Open(path, 0600, nil)
+		if err != nil {
+			panic(errors.Wrap(err, "could not start bolt DB"))
+		}
+		if err := boltDB.Update(func(tx *bolt.Tx) error {
+			if _, err := tx.CreateBucketIfNotExists([]byte(pullsBucketName)); err != nil {
+				return errors.Wrap(err, "failed to create bucket")
+			}
+			return nil
+		}); err != nil {
+			panic(errors.Wrap(err, "could not create bucket"))
+		}
+		db, _ := db.NewWithDB(boltDB, lockBucket, configBucket)
+		result := []models.ProjectResult{
+			{
+				RepoRelDir:  fixtures.GithubRepo.FullName,
+				Workspace:   "default",
+				ProjectName: *fixtures.Project.Name,
+			},
+		}
+
+		// Create a new record for pull
+		_, err = db.UpdatePullWithResults(fixtures.Pull, result)
+		Ok(t, err)
+
+		workingDir := mocks.NewMockWorkingDir()
+		locker := lockmocks.NewMockLocker()
+		client := vcsmocks.NewMockClient()
+		logger := loggermocks.NewMockSimpleLogging()
+
+		pullClosedExecutor := events.PullClosedExecutor{
+			Locker:                   locker,
+			WorkingDir:               workingDir,
+			DB:                       db,
+			VCSClient:                client,
+			PullClosedTemplate:       &events.PullClosedEventTemplate{},
+			LogStreamResourceCleaner: prjCmdOutHandler,
+			Logger:                   logger,
+		}
+
+		locks := []models.ProjectLock{
+			{
+				Project:   models.NewProject(fixtures.GithubRepo.FullName, ""),
+				Workspace: "default",
+			},
+		}
+		When(locker.UnlockByPull(fixtures.GithubRepo.FullName, fixtures.Pull.Num)).ThenReturn(locks, nil)
+
+		// Clean up.
+		err = pullClosedExecutor.CleanUpPull(fixtures.GithubRepo, fixtures.Pull)
+		Ok(t, err)
+
+		close(prjCmdOutput)
+		_, _, comment, _ := client.VerifyWasCalledOnce().CreateComment(matchers.AnyModelsRepo(), AnyInt(), AnyString(), AnyString()).GetCapturedArguments()
+		expectedComment := "Locks and plans deleted for the projects and workspaces modified in this pull request:\n\n" + "- dir: `.` workspace: `default`"
+		Equals(t, expectedComment, comment)
+
+		// Assert log streaming resources are cleaned up.
+		dfPrjCmdOutputHandler := prjCmdOutHandler.(*handlers.AsyncProjectCommandOutputHandler)
+		assert.Empty(t, dfPrjCmdOutputHandler.GetProjectOutputBuffer(ctx.PullInfo()))
+		assert.Empty(t, dfPrjCmdOutputHandler.GetReceiverBufferForPull(ctx.PullInfo()))
+	})
+}

server/events/yaml/valid/repo_cfg.go

@@ -116,6 +116,11 @@ func (p Project) GetName() string {
 	if p.Name != nil {
 		return *p.Name
 	}
+	// TODO
+	// Upstream atlantis only requires project name to be set if there's more than one project
+	// with same dir and workspace. If a project name has not been set, we'll use the dir and
+	// workspace to build project key.
+	// Source: https://www.runatlantis.io/docs/repo-level-atlantis-yaml.html#reference
 	return ""
 }
 

server/feature/allocator.go

@@ -6,7 +6,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/runatlantis/atlantis/server/events/vcs"
 	"github.com/runatlantis/atlantis/server/logging"
-	"github.com/thomaspoignant/go-feature-flag"
+	ffclient "github.com/thomaspoignant/go-feature-flag"
 	"github.com/thomaspoignant/go-feature-flag/ffuser"
 )