Polishing

coder · mtojek · Sep 11, 2023 · Sep 5, 2023 · Sep 6, 2023 · Sep 6, 2023
commit 18c5a1d20f712ac88b3d021d3f9999d95a01c95a
diff --git a/provisioner/terraform/cleanup.go b/provisioner/terraform/cleanup.go
@@ -137,7 +137,6 @@ func latestAccessTime(fs afero.Fs, pluginPath string) (time.Time, error) {
 		}
 
 		var accessTime = info.ModTime() // fallback to modTime if accessTime is not available (afero)
-
 		if info.Sys() != nil {
 			timeSpec := times.Get(info)
 			accessTime = timeSpec.AccessTime()

diff --git a/provisioner/terraform/cleanup_test.go b/provisioner/terraform/cleanup_test.go
@@ -21,29 +21,30 @@ import (
 	"github.com/coder/coder/v2/testutil"
 )
 
-const (
-	cachePath = "/tmp/coder/provisioner-0/tf"
-)
+const cachePath = "/tmp/coder/provisioner-0/tf"
 
 // updateGoldenFiles is a flag that can be set to update golden files.
 var updateGoldenFiles = flag.Bool("update", false, "Update golden files")
 
 func TestPluginCache_Golden(t *testing.T) {
 	t.Parallel()
 
+	prepare := func() (afero.Fs, time.Time, slog.Logger) {
+		fs := afero.NewMemMapFs()
+		now := time.Date(2023, time.June, 3, 4, 5, 6, 0, time.UTC)
+		logger := slogtest.Make(t, nil).
+			Leveled(slog.LevelDebug).
+			Named("cleanup-test")
+		return fs, now, logger
+	}
+
 	t.Run("all plugins are stale", func(t *testing.T) {
 		t.Parallel()
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()
 
-		// prepare
-		fs := afero.NewMemMapFs()
-		now := time.Date(2023, time.June, 3, 4, 5, 6, 0, time.UTC)
-
-		logger := slogtest.Make(t, nil).
-			Leveled(slog.LevelDebug).
-			Named("cleanup-test")
+		fs, now, logger := prepare()
 
 		// given
 		// This plugin is older than 30 days.
@@ -71,13 +72,7 @@ func TestPluginCache_Golden(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()
 
-		// prepare
-		fs := afero.NewMemMapFs()
-		now := time.Date(2023, time.June, 3, 4, 5, 6, 0, time.UTC)
-
-		logger := slogtest.Make(t, nil).
-			Leveled(slog.LevelDebug).
-			Named("cleanup-test")
+		fs, now, logger := prepare()
 
 		// given
 		addPluginFile(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "terraform-provider-coder_v0.11.1", now.Add(-2*time.Hour))
@@ -104,19 +99,13 @@ func TestPluginCache_Golden(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()
 
-		// prepare
-		fs := afero.NewMemMapFs()
-		now := time.Date(2023, time.June, 3, 4, 5, 6, 0, time.UTC)
-
-		logger := slogtest.Make(t, nil).
-			Leveled(slog.LevelDebug).
-			Named("cleanup-test")
+		fs, now, logger := prepare()
 
 		// given
 		addPluginFile(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "terraform-provider-coder_v0.11.1", now.Add(-63*24*time.Hour))
 		addPluginFile(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "LICENSE", now.Add(-33*24*time.Hour))
 		addPluginFile(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "README.md", now.Add(-31*24*time.Hour))
-		addPluginFolder(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "new_folder", now.Add(-4*time.Hour))
+		addPluginFolder(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "new_folder", now.Add(-4*time.Hour)) // touched
 		addPluginFile(t, fs, "registry.terraform.io/coder/coder/0.11.1/darwin_arm64", "new_folder/foobar.tf", now.Add(-43*24*time.Hour))
 
 		addPluginFile(t, fs, "registry.terraform.io/kreuzwerker/docker/2.25.0/darwin_arm64", "terraform-provider-docker_v2.25.0", now.Add(-31*24*time.Hour))

diff --git a/provisionersdk/cleanup.go b/provisionersdk/cleanup.go
@@ -2,42 +2,43 @@ package provisionersdk
 
 import (
 	"context"
-	"os"
 	"path/filepath"
 	"time"
 
 	"github.com/djherbis/times"
+	"github.com/spf13/afero"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
 )
 
-// cleanStaleSessions browses the work directory searching for stale session
+// CleanStaleSessions browses the work directory searching for stale session
 // directories. Coder provisioner is supposed to remove them once after finishing the provisioning,
 // but there is a risk of keeping them in case of a failure.
-func cleanStaleSessions(ctx context.Context, workDirectory string, now time.Time, logger slog.Logger) error {
-	entries, err := os.ReadDir(workDirectory)
+func CleanStaleSessions(ctx context.Context, workDirectory string, fs afero.Fs, now time.Time, logger slog.Logger) error {
+	entries, err := afero.ReadDir(fs, workDirectory)
 	if err != nil {
 		return xerrors.Errorf("can't read %q directory", workDirectory)
 	}
 
-	for _, entry := range entries {
-		dirName := entry.Name()
+	for _, fi := range entries {
+		dirName := fi.Name()
 
-		if entry.IsDir() && isValidSessionDir(dirName) {
+		if fi.IsDir() && isValidSessionDir(dirName) {
 			sessionDirPath := filepath.Join(workDirectory, dirName)
-			fi, err := entry.Info()
-			if err != nil {
-				return xerrors.Errorf("can't read %q directory info: %w", sessionDirPath, err)
+
+			var accessTime = fi.ModTime() // fallback to modTime if accessTime is not available (afero)
+			if fi.Sys() != nil {
+				timeSpec := times.Get(fi)
+				accessTime = timeSpec.AccessTime()
 			}
 
-			timeSpec := times.Get(fi)
-			if timeSpec.AccessTime().Add(staleSessionRetention).After(now) {
+			if accessTime.Add(staleSessionRetention).After(now) {
 				continue
 			}
 
-			logger.Info(ctx, "remove stale session directory: %s", sessionDirPath)
-			err = os.RemoveAll(sessionDirPath)
+			logger.Info(ctx, "remove stale session directory", slog.F("session_path", sessionDirPath))
+			err = fs.RemoveAll(sessionDirPath)
 			if err != nil {
 				return xerrors.Errorf("can't remove %q directory: %w", sessionDirPath, err)
 			}

diff --git a/provisionersdk/cleanup_test.go b/provisionersdk/cleanup_test.go
@@ -0,0 +1,112 @@
+package provisionersdk_test
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/v2/provisionersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+const workDirectory = "/tmp/coder/provisioner-34/work"
+
+func TestStaleSessions(t *testing.T) {
+	t.Parallel()
+
+	prepare := func() (afero.Fs, time.Time, slog.Logger) {
+		fs := afero.NewMemMapFs()
+		now := time.Date(2023, time.June, 3, 4, 5, 6, 0, time.UTC)
+		logger := slogtest.Make(t, nil).
+			Leveled(slog.LevelDebug).
+			Named("cleanup-test")
+		return fs, now, logger
+	}
+
+	t.Run("all sessions are stale", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+
+		fs, now, logger := prepare()
+
+		// given
+		first := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, first, now.Add(-7*24*time.Hour))
+		second := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, second, now.Add(-8*24*time.Hour))
+		third := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, third, now.Add(-9*24*time.Hour))
+
+		// when
+		provisionersdk.CleanStaleSessions(ctx, workDirectory, fs, now, logger)
+
+		// then
+		entries, err := afero.ReadDir(fs, workDirectory)
+		require.NoError(t, err)
+		require.Empty(t, entries, "all session leftovers should be removed")
+	})
+
+	t.Run("one session is stale", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+
+		fs, now, logger := prepare()
+
+		// given
+		first := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, first, now.Add(-7*24*time.Hour))
+		second := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, second, now.Add(-6*24*time.Hour))
+
+		// when
+		provisionersdk.CleanStaleSessions(ctx, workDirectory, fs, now, logger)
+
+		// then
+		entries, err := afero.ReadDir(fs, workDirectory)
+		require.NoError(t, err)
+		require.Len(t, entries, 1, "one session should be present")
+		require.Equal(t, second, entries[0].Name(), 1)
+	})
+
+	t.Run("no stale sessions", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+
+		fs, now, logger := prepare()
+
+		// given
+		first := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, first, now.Add(-6*24*time.Hour))
+		second := provisionersdk.SessionDir(uuid.NewString())
+		addSessionFolder(t, fs, second, now.Add(-5*24*time.Hour))
+
+		// when
+		provisionersdk.CleanStaleSessions(ctx, workDirectory, fs, now, logger)
+
+		// then
+		entries, err := afero.ReadDir(fs, workDirectory)
+		require.NoError(t, err)
+		require.Len(t, entries, 2, "both sessions should be present")
+	})
+}
+
+func addSessionFolder(t *testing.T, fs afero.Fs, sessionName string, accessTime time.Time) {
+	err := fs.MkdirAll(filepath.Join(workDirectory, sessionName), 0755)
+	require.NoError(t, err, "can't create session folder")
+	fs.Chtimes(filepath.Join(workDirectory, sessionName), accessTime, accessTime)
+	require.NoError(t, err, "can't set times")
+}
diff --git a/provisionersdk/session.go b/provisionersdk/session.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/spf13/afero"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
@@ -41,12 +42,12 @@ func (p *protoServer) Session(stream proto.DRPCProvisioner_SessionStream) error
 		server: p.server,
 	}
 
-	err := cleanStaleSessions(s.Context(), p.opts.WorkDirectory, time.Now(), s.Logger)
+	err := CleanStaleSessions(s.Context(), p.opts.WorkDirectory, afero.NewOsFs(), time.Now(), s.Logger)
 	if err != nil {
 		return xerrors.Errorf("unable to clean stale sessions %q: %w", s.WorkDirectory, err)
 	}
 
-	s.WorkDirectory = filepath.Join(p.opts.WorkDirectory, sessionDir(sessID))
+	s.WorkDirectory = filepath.Join(p.opts.WorkDirectory, SessionDir(sessID))
 	err = os.MkdirAll(s.WorkDirectory, 0o700)
 	if err != nil {
 		return xerrors.Errorf("create work directory %q: %w", s.WorkDirectory, err)
@@ -327,6 +328,7 @@ func (r *request[R, C]) do() (C, error) {
 	}
 }
 
-func sessionDir(sessID string) string {
+// SessionDir returns the directory name with mandatory prefix.
+func SessionDir(sessID string) string {
 	return sessionDirPrefix + sessID
 }