use gofrs/flock instead to handle file locking

johnstcn · johnstcn · commit e830baf6e265 · 2022-05-13T12:31:17.000+01:00
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -2,7 +2,6 @@ package cli
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -11,6 +10,7 @@ import (
 	"time"
 
 	"github.com/gen2brain/beeep"
+	"github.com/gofrs/flock"
 	"github.com/google/uuid"
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
@@ -197,35 +197,20 @@ func ssh() *cobra.Command {
 // avoid spamming the user with notifications in case of multiple instances
 // of the CLI running simultaneously.
 func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace) (stop func()) {
-	lockPath := filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String())
-	lockStat, err := os.Stat(lockPath)
-	if err == nil {
-		// Lock file already exists for this workspace. How old is it?
-		lockAge := time.Now().Sub(lockStat.ModTime())
-		if lockAge < 3*autostopPollInterval {
-			// Lock file exists and is still "fresh". Do nothing.
-			return func() {}
-		}
-	}
-	if !errors.Is(err, os.ErrNotExist) {
-		// No permission to write to temp? Not much we can do.
-		return func() {}
-	}
-	lockFile, err := os.Create(lockPath)
-	if err != nil {
-		// Someone got there already?
-		return func() {}
-	}
-
-	condition := notifyCondition(ctx, client, workspace.ID, lockFile)
+	lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
+	condition := notifyCondition(ctx, client, workspace.ID, lock)
 	return notify.Notify(condition, autostopPollInterval, autostopNotifyCountdown...)
 }
 
 // Notify the user if the workspace is due to shutdown.
-func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lockFile *os.File) notify.Condition {
+func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lock *flock.Flock) notify.Condition {
 	return func(now time.Time) (deadline time.Time, callback func()) {
-		// update lockFile (best effort)
-		_ = os.Chtimes(lockFile.Name(), now, now)
+		// Keep trying to regain the lock.
+		locked, err := lock.TryLockContext(ctx, autostopPollInterval)
+		if err != nil || !locked {
+			return time.Time{}, nil
+		}
+
 		ws, err := client.Workspace(ctx, workspaceID)
 		if err != nil {
 			return time.Time{}, nil
diff --git a/go.mod b/go.mod
@@ -64,6 +64,7 @@ require (
 	github.com/go-chi/httprate v0.5.3
 	github.com/go-chi/render v1.0.1
 	github.com/go-playground/validator/v10 v10.11.0
+	github.com/gofrs/flock v0.8.1
 	github.com/gohugoio/hugo v0.98.0
 	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/golang-migrate/migrate/v4 v4.15.2
diff --git a/go.sum b/go.sum
@@ -754,6 +754,8 @@ github.com/godbus/dbus/v5 v5.0.6/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5x
 github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
 github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gofiber/fiber/v2 v2.11.0/go.mod h1:oZTLWqYnqpMMuF922SjGbsYZsdpE1MCfh416HNdweIM=
+github.com/gofrs/flock v0.8.1 h1:+gYjHKf32LDeiEEFhQaotPbLuUXjY5ZqxKgXy7n59aw=
+github.com/gofrs/flock v0.8.1/go.mod h1:F1TvTiK9OcQqauNUHlbJvyl9Qa1QvF/gOUDKA14jxHU=
 github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/googleapis v1.2.0/go.mod h1:Njal3psf3qN6dwBtQfUmBZh2ybovJ0tlu3o/AC7HYjU=
