Confine cleanup and use held session locks

internal/artifacts/local.go

@@ -1,6 +1,7 @@
 package artifacts
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -14,7 +15,7 @@ import (
 	"gitea.maximumdirect.net/eric/narratio/internal/pathsafe"
 )
 
-// ErrLockConflict is returned when a session lock already exists.
+// ErrLockConflict is returned when a session lock is currently held.
 var ErrLockConflict = errors.New("session workdir is already locked")
 
 // LockHandle tracks a held lock for a session work directory.
@@ -178,32 +179,74 @@ func (s *LocalStore) AcquireSessionLockFor(campaign, sessionID string) (*LockHan
 	return s.acquireSessionLockForPaths(paths)
 }
 
-func (s *LocalStore) acquireSessionLockForPaths(paths SessionPaths) (*LockHandle, error) {
-	f, err := os.OpenFile(paths.LockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, fileops.WorkspaceFileMode)
+// AcquireSessionLockForContext waits for a session lock until it becomes
+// available or ctx is cancelled.
+func (s *LocalStore) AcquireSessionLockForContext(ctx context.Context, campaign, sessionID string) (*LockHandle, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	paths, err := s.EnsureLayoutFor(campaign, sessionID)
 	if err != nil {
-		if errors.Is(err, os.ErrExist) {
-			return nil, fmt.Errorf("%w: %s", ErrLockConflict, paths.LockPath)
+		return nil, err
+	}
+	for {
+		lock, err := s.acquireSessionLockForPaths(paths)
+		if !errors.Is(err, ErrLockConflict) {
+			return lock, err
 		}
+		timer := time.NewTimer(100 * time.Millisecond)
+		select {
+		case <-ctx.Done():
+			if !timer.Stop() {
+				<-timer.C
+			}
+			return nil, fmt.Errorf("wait for session lock %q: %w", paths.LockPath, ctx.Err())
+		case <-timer.C:
+		}
+	}
+}
+
+func (s *LocalStore) acquireSessionLockForPaths(paths SessionPaths) (*LockHandle, error) {
+	f, err := fileops.OpenFileConfined(paths.LockPath, os.O_CREATE|os.O_RDWR, fileops.WorkspaceFileMode)
+	if err != nil {
 		return nil, fmt.Errorf("acquire lock %q: %w", paths.LockPath, err)
 	}
+	if err := acquireHeldFileLockFn(f); err != nil {
+		closeErr := f.Close()
+		if errors.Is(err, errHeldLockConflict) {
+			return nil, fmt.Errorf("%w: %s", ErrLockConflict, paths.LockPath)
+		}
+		if closeErr != nil {
+			return nil, errors.Join(fmt.Errorf("acquire lock %q: hold: %w", paths.LockPath, err), closeErr)
+		}
+		return nil, fmt.Errorf("acquire lock %q: hold: %w", paths.LockPath, err)
+	}
+	failed := true
+	defer func() {
+		if failed {
+			_ = releaseHeldFileLockFn(f)
+			_ = f.Close()
+		}
+	}()
 	if err := f.Chmod(fileops.WorkspaceFileMode); err != nil {
-		_ = f.Close()
-		_ = os.Remove(paths.LockPath)
 		return nil, fmt.Errorf("acquire lock %q: set permissions: %w", paths.LockPath, err)
 	}
+	if err := f.Truncate(0); err != nil {
+		return nil, fmt.Errorf("acquire lock %q: clear metadata: %w", paths.LockPath, err)
+	}
+	if _, err := f.Seek(0, io.SeekStart); err != nil {
+		return nil, fmt.Errorf("acquire lock %q: seek metadata: %w", paths.LockPath, err)
+	}
 
 	metadata := "pid=" + strconv.Itoa(os.Getpid()) + "\nacquired_at=" + time.Now().UTC().Format(time.RFC3339Nano) + "\n"
 	if _, err := io.WriteString(f, metadata); err != nil {
-		_ = f.Close()
-		_ = os.Remove(paths.LockPath)
 		return nil, fmt.Errorf("acquire lock %q: write metadata: %w", paths.LockPath, err)
 	}
 	if err := f.Sync(); err != nil {
-		_ = f.Close()
-		_ = os.Remove(paths.LockPath)
 		return nil, fmt.Errorf("acquire lock %q: sync: %w", paths.LockPath, err)
 	}
 
+	failed = false
 	return &LockHandle{path: paths.LockPath, file: f}, nil
 }
 
@@ -213,23 +256,21 @@ func (s *LocalStore) ReleaseSessionLock(lock *LockHandle) error {
 		return nil
 	}
 
-	var closeErr error
-	if lock.file != nil {
-		closeErr = lock.file.Close()
-		lock.file = nil
+	if lock.file == nil {
+		return nil
 	}
-
-	removeErr := os.Remove(lock.path)
-	if errors.Is(removeErr, os.ErrNotExist) {
-		removeErr = nil
+	releaseErr := releaseHeldFileLockFn(lock.file)
+	closeErr := closeHeldLockFile(lock.file)
+	lock.file = nil
+	if releaseErr != nil && closeErr != nil {
+		return fmt.Errorf("release lock %q: %w", lock.path, errors.Join(releaseErr, closeErr))
+	}
+	if releaseErr != nil {
+		return fmt.Errorf("release lock %q: unlock: %w", lock.path, releaseErr)
 	}
-
 	if closeErr != nil {
 		return fmt.Errorf("release lock %q: close: %w", lock.path, closeErr)
 	}
-	if removeErr != nil {
-		return fmt.Errorf("release lock %q: remove: %w", lock.path, removeErr)
-	}
 	return nil
 }
 

internal/artifacts/local_test.go

@@ -1,11 +1,14 @@
 package artifacts
 
 import (
+	"context"
 	"errors"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestEnsureLayoutCreatesExpectedDirectories(t *testing.T) {
@@ -93,8 +96,148 @@ func TestLockAcquireRelease(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Exists() error = %v", err)
 	}
-	if exists {
-		t.Fatalf("expected lock file %q to be removed", lock.path)
+	if !exists {
+		t.Fatalf("expected lock metadata file %q to remain", lock.path)
+	}
+	secondLock, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+	if err != nil {
+		t.Fatalf("AcquireSessionLockFor() after release error = %v", err)
+	}
+	if err := store.ReleaseSessionLock(secondLock); err != nil {
+		t.Fatalf("ReleaseSessionLock(second lock) error = %v", err)
+	}
+}
+
+func TestLockWaitHonorsCancellation(t *testing.T) {
+	store := NewLocalStore(t.TempDir())
+	lock, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+	if err != nil {
+		t.Fatalf("AcquireSessionLockFor() error = %v", err)
+	}
+	defer func() { _ = store.ReleaseSessionLock(lock) }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err = store.AcquireSessionLockForContext(ctx, "sample-campaign", "session-1")
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("AcquireSessionLockForContext() error = %v, want context cancellation", err)
+	}
+}
+
+func TestLockWaitsUntilRelease(t *testing.T) {
+	store := NewLocalStore(t.TempDir())
+	first, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+	if err != nil {
+		t.Fatalf("AcquireSessionLockFor() error = %v", err)
+	}
+	acquired := make(chan *LockHandle, 1)
+	errs := make(chan error, 1)
+	go func() {
+		lock, err := store.AcquireSessionLockForContext(context.Background(), "sample-campaign", "session-1")
+		if err != nil {
+			errs <- err
+			return
+		}
+		acquired <- lock
+	}()
+
+	select {
+	case err := <-errs:
+		t.Fatalf("AcquireSessionLockForContext() error = %v", err)
+	case lock := <-acquired:
+		_ = store.ReleaseSessionLock(lock)
+		t.Fatal("AcquireSessionLockForContext() acquired before the held lock was released")
+	case <-time.After(150 * time.Millisecond):
+	}
+	if err := store.ReleaseSessionLock(first); err != nil {
+		t.Fatalf("ReleaseSessionLock(first) error = %v", err)
+	}
+	select {
+	case err := <-errs:
+		t.Fatalf("AcquireSessionLockForContext() error = %v", err)
+	case lock := <-acquired:
+		if err := store.ReleaseSessionLock(lock); err != nil {
+			t.Fatalf("ReleaseSessionLock(waiting lock) error = %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("AcquireSessionLockForContext() did not acquire after release")
+	}
+}
+
+func TestLockRecoversAfterLockHolderDies(t *testing.T) {
+	if os.Getenv("NARRATIO_LOCK_PROCESS_HELPER") == "1" {
+		store := NewLocalStore(os.Getenv("NARRATIO_LOCK_PROCESS_ROOT"))
+		lock, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+		if err != nil {
+			os.Exit(2)
+		}
+		if err := os.WriteFile(os.Getenv("NARRATIO_LOCK_PROCESS_READY"), []byte(lock.path), 0o600); err != nil {
+			os.Exit(3)
+		}
+		select {}
+	}
+
+	root := t.TempDir()
+	ready := filepath.Join(t.TempDir(), "ready")
+	cmd := exec.Command(os.Args[0], "-test.run=^TestLockRecoversAfterLockHolderDies$")
+	cmd.Env = append(os.Environ(),
+		"NARRATIO_LOCK_PROCESS_HELPER=1",
+		"NARRATIO_LOCK_PROCESS_ROOT="+root,
+		"NARRATIO_LOCK_PROCESS_READY="+ready,
+	)
+	if err := cmd.Start(); err != nil {
+		t.Fatalf("start lock-holder process: %v", err)
+	}
+	defer func() {
+		if cmd.Process != nil {
+			_ = cmd.Process.Kill()
+		}
+		_ = cmd.Wait()
+	}()
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		if _, err := os.Stat(ready); err == nil {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatal("lock-holder process did not signal readiness")
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	store := NewLocalStore(root)
+	if _, err := store.AcquireSessionLockFor("sample-campaign", "session-1"); !errors.Is(err, ErrLockConflict) {
+		t.Fatalf("AcquireSessionLockFor() error = %v, want lock conflict", err)
+	}
+	if err := cmd.Process.Kill(); err != nil {
+		t.Fatalf("kill lock-holder process: %v", err)
+	}
+	if err := cmd.Wait(); err == nil {
+		t.Fatal("lock-holder process exited without being killed")
+	}
+	cmd.Process = nil
+
+	lock, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+	if err != nil {
+		t.Fatalf("AcquireSessionLockFor() after holder death error = %v", err)
+	}
+	if err := store.ReleaseSessionLock(lock); err != nil {
+		t.Fatalf("ReleaseSessionLock() error = %v", err)
+	}
+}
+
+func TestReleaseSessionLockReportsUnlockFailure(t *testing.T) {
+	store := NewLocalStore(t.TempDir())
+	lock, err := store.AcquireSessionLockFor("sample-campaign", "session-1")
+	if err != nil {
+		t.Fatalf("AcquireSessionLockFor() error = %v", err)
+	}
+
+	original := releaseHeldFileLockFn
+	releaseHeldFileLockFn = func(*os.File) error { return errors.New("unlock failed") }
+	t.Cleanup(func() { releaseHeldFileLockFn = original })
+	if err := store.ReleaseSessionLock(lock); err == nil || !strings.Contains(err.Error(), "unlock failed") {
+		t.Fatalf("ReleaseSessionLock() error = %v, want unlock failure", err)
 	}
 }
 

internal/artifacts/lock_common.go

@@ -0,0 +1,18 @@
+package artifacts
+
+import (
+	"errors"
+	"os"
+)
+
+var errHeldLockConflict = errors.New("held lock conflict")
+
+var acquireHeldFileLockFn = acquireHeldFileLock
+var releaseHeldFileLockFn = releaseHeldFileLock
+
+func closeHeldLockFile(file *os.File) error {
+	if file == nil {
+		return nil
+	}
+	return file.Close()
+}

internal/artifacts/lock_unix.go

@@ -0,0 +1,22 @@
+//go:build linux || darwin
+
+package artifacts
+
+import (
+	"errors"
+	"os"
+
+	"golang.org/x/sys/unix"
+)
+
+func acquireHeldFileLock(file *os.File) error {
+	err := unix.Flock(int(file.Fd()), unix.LOCK_EX|unix.LOCK_NB)
+	if errors.Is(err, unix.EWOULDBLOCK) || errors.Is(err, unix.EAGAIN) {
+		return errHeldLockConflict
+	}
+	return err
+}
+
+func releaseHeldFileLock(file *os.File) error {
+	return unix.Flock(int(file.Fd()), unix.LOCK_UN)
+}

internal/artifacts/lock_unsupported.go

@@ -0,0 +1,15 @@
+//go:build !linux && !darwin && !windows
+
+package artifacts
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+)
+
+func acquireHeldFileLock(_ *os.File) error {
+	return fmt.Errorf("held file locks are unsupported on %s", runtime.GOOS)
+}
+
+func releaseHeldFileLock(_ *os.File) error { return nil }

internal/artifacts/lock_windows.go

@@ -0,0 +1,29 @@
+//go:build windows
+
+package artifacts
+
+import (
+	"errors"
+	"os"
+
+	"golang.org/x/sys/windows"
+)
+
+func acquireHeldFileLock(file *os.File) error {
+	err := windows.LockFileEx(
+		windows.Handle(file.Fd()),
+		windows.LOCKFILE_EXCLUSIVE_LOCK|windows.LOCKFILE_FAIL_IMMEDIATELY,
+		0,
+		1,
+		0,
+		&windows.Overlapped{},
+	)
+	if errors.Is(err, windows.ERROR_LOCK_VIOLATION) {
+		return errHeldLockConflict
+	}
+	return err
+}
+
+func releaseHeldFileLock(file *os.File) error {
+	return windows.UnlockFileEx(windows.Handle(file.Fd()), 0, 1, 0, &windows.Overlapped{})
+}

internal/artifacts/store.go

@@ -1,6 +1,7 @@
 package artifacts
 
 import (
+	"context"
 	"os"
 
 	"gitea.maximumdirect.net/eric/narratio/internal/artifactmodel"
@@ -30,5 +31,6 @@ type Store interface {
 	WriteFileAtomic(path string, data []byte, perm os.FileMode) error
 	Checksum(path string) (string, error)
 	AcquireSessionLockFor(campaign, sessionID string) (*LockHandle, error)
+	AcquireSessionLockForContext(ctx context.Context, campaign, sessionID string) (*LockHandle, error)
 	ReleaseSessionLock(lock *LockHandle) error
 }