Allow atomic recovery updates during concurrent reads

Author: skarim

cmd/modify_test.go

@@ -2,9 +2,13 @@ package cmd
 
 import (
 	"encoding/json"
+	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -109,6 +113,72 @@ func TestModifyStateAtomicWrite(t *testing.T) {
 	assert.True(t, os.IsNotExist(err), "no .tmp file should remain after successful write")
 }
 
+func TestModifyStateReadError(t *testing.T) {
+	dir := t.TempDir()
+	path := modify.StatePath(dir)
+	require.NoError(t, os.Mkdir(path, 0700))
+	got, err := modify.LoadState(dir)
+	require.ErrorContains(t, err, "reading modify state")
+	assert.Nil(t, got)
+	var pathErr *os.PathError
+	require.ErrorAs(t, err, &pathErr)
+	assert.Equal(t, path, pathErr.Path)
+}
+
+func TestModifyStateConcurrentReadWrite(t *testing.T) {
+	dir := t.TempDir()
+	state := &modify.StateFile{
+		SchemaVersion: 1, Phase: modify.PhaseConflict,
+		OriginalBranch: "initial", ConflictBranch: "initial",
+		Snapshot: modify.Snapshot{StackMetadata: json.RawMessage("{}")},
+	}
+	require.NoError(t, modify.SaveState(dir, state))
+	stop := make(chan struct{})
+	errs := make(chan error, 2)
+	var reads atomic.Int64
+	var readers sync.WaitGroup
+	for range 2 {
+		readers.Go(func() {
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				got, err := modify.LoadState(dir)
+				if err != nil {
+					errs <- err
+					return
+				}
+				if got == nil || got.OriginalBranch == "" || got.OriginalBranch != got.ConflictBranch {
+					errs <- fmt.Errorf("reader observed an incomplete modify state")
+					return
+				}
+				reads.Add(1)
+			}
+		})
+	}
+	var writeErr error
+	for i := range 50 {
+		branch := strings.Repeat(fmt.Sprintf("%04d", i), 8192)
+		state.OriginalBranch, state.ConflictBranch = branch, branch
+		if writeErr = modify.SaveState(dir, state); writeErr != nil {
+			break
+		}
+	}
+	close(stop)
+	readers.Wait()
+	close(errs)
+	require.NoError(t, writeErr)
+	for err := range errs {
+		require.NoError(t, err)
+	}
+	assert.Positive(t, reads.Load())
+	got, err := modify.LoadState(dir)
+	require.NoError(t, err)
+	assert.Equal(t, state, got)
+}
+
 func TestCheckModifyStateGuard(t *testing.T) {
 	t.Run("no state file", func(t *testing.T) {
 		gitDir := t.TempDir()

cmd/rebase.go

@@ -543,7 +543,7 @@ func saveRebaseState(gitDir string, state *rebaseState) error {
 }
 
 func loadRebaseState(gitDir string) (*rebaseState, error) {
-	data, err := os.ReadFile(filepath.Join(gitDir, rebaseStateFile))
+	data, err := stack.ReadStateFile(filepath.Join(gitDir, rebaseStateFile))
 	if err != nil {
 		return nil, err
 	}

cmd/rebase_test.go

@@ -9,6 +9,8 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/github/gh-stack/internal/config"
@@ -953,6 +955,85 @@ func TestRebase_StateRoundTrip(t *testing.T) {
 	assert.Equal(t, original.OntoOldBase, loaded.OntoOldBase)
 }
 
+func TestRebase_StateReadErrors(t *testing.T) {
+	for _, kind := range []string{"missing", "directory", "invalid JSON"} {
+		t.Run(kind, func(t *testing.T) {
+			dir := t.TempDir()
+			path := filepath.Join(dir, rebaseStateFile)
+			switch kind {
+			case "directory":
+				require.NoError(t, os.Mkdir(path, 0700))
+			case "invalid JSON":
+				require.NoError(t, os.WriteFile(path, []byte("{incomplete"), 0600))
+			}
+			got, err := loadRebaseState(dir)
+			require.Error(t, err)
+			assert.Nil(t, got)
+			switch kind {
+			case "missing":
+				assert.ErrorIs(t, err, os.ErrNotExist)
+			case "directory":
+				var pathErr *os.PathError
+				require.ErrorAs(t, err, &pathErr)
+				assert.Equal(t, path, pathErr.Path)
+			case "invalid JSON":
+				var syntaxErr *json.SyntaxError
+				assert.ErrorAs(t, err, &syntaxErr)
+			}
+		})
+	}
+}
+
+func TestRebase_StateConcurrentReadWrite(t *testing.T) {
+	dir := t.TempDir()
+	state := &rebaseState{OriginalBranch: "initial", ConflictBranch: "initial"}
+	require.NoError(t, saveRebaseState(dir, state))
+	stop := make(chan struct{})
+	errs := make(chan error, 2)
+	var reads atomic.Int64
+	var readers sync.WaitGroup
+	for range 2 {
+		readers.Go(func() {
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				got, err := loadRebaseState(dir)
+				if err != nil {
+					errs <- err
+					return
+				}
+				if got == nil || got.OriginalBranch == "" || got.OriginalBranch != got.ConflictBranch {
+					errs <- fmt.Errorf("reader observed an incomplete rebase state")
+					return
+				}
+				reads.Add(1)
+			}
+		})
+	}
+	var writeErr error
+	for i := range 50 {
+		branch := strings.Repeat(fmt.Sprintf("%04d", i), 8192)
+		state.OriginalBranch, state.ConflictBranch = branch, branch
+		if writeErr = saveRebaseState(dir, state); writeErr != nil {
+			break
+		}
+	}
+	close(stop)
+	readers.Wait()
+	close(errs)
+	require.NoError(t, writeErr)
+	for err := range errs {
+		require.NoError(t, err)
+	}
+	assert.Positive(t, reads.Load())
+	got, err := loadRebaseState(dir)
+	require.NoError(t, err)
+	assert.Equal(t, state, got)
+}
+
 // TestRebase_Continue_RebasesRemainingBranches verifies the --continue success
 // path: RebaseContinue is called, remaining branches are rebased via RebaseOnto,
 // the state file is cleaned up, and the original branch is restored.

internal/modify/state.go

@@ -77,7 +77,7 @@ func StatePath(gitDir string) string {
 // LoadState reads the modify state file from the git directory.
 // Returns nil, nil if the file does not exist.
 func LoadState(gitDir string) (*StateFile, error) {
-	data, err := os.ReadFile(StatePath(gitDir))
+	data, err := stack.ReadStateFile(StatePath(gitDir))
 	if err != nil {
 		if errors.Is(err, os.ErrNotExist) {
 			return nil, nil

internal/stack/atomic.go

@@ -7,6 +7,12 @@ import (
 	"path/filepath"
 )
 
+// ReadStateFile reads a complete state file, allowing WriteAtomic to replace it
+// while the read is in progress. On Windows, its handle permits delete sharing.
+func ReadStateFile(path string) ([]byte, error) {
+	return readStateFile(path)
+}
+
 // WriteAtomic publishes data at path using a fully written temporary file in
 // the same directory. It preserves an existing regular file's permissions and
 // uses 0644 for a new file. The parent directory must already exist.

internal/stack/atomic_windows_test.go

@@ -57,3 +57,47 @@ func TestAtomicPublication_HeldWindowsReader(t *testing.T) {
 		})
 	}
 }
+
+func TestReadStateFile_WindowsSharing(t *testing.T) {
+	for _, shareDelete := range []bool{false, true} {
+		name := "exclusive handle"
+		if shareDelete {
+			name = "shared delete handle"
+		}
+		t.Run(name, func(t *testing.T) {
+			path := filepath.Join(t.TempDir(), "state")
+			original, replacement := []byte("original state"), []byte("new state")
+			require.NoError(t, WriteAtomic(path, original))
+			name, err := windowsFilePath(path)
+			require.NoError(t, err)
+			var shareMode uint32
+			if shareDelete {
+				shareMode = windows.FILE_SHARE_READ | windows.FILE_SHARE_WRITE | windows.FILE_SHARE_DELETE
+			}
+			handle, err := windows.CreateFile(&name[0], windows.GENERIC_READ|windows.DELETE,
+				shareMode, nil, windows.OPEN_EXISTING, windows.FILE_ATTRIBUTE_NORMAL, 0)
+			require.NoError(t, err)
+			reader := os.NewFile(uintptr(handle), path)
+			t.Cleanup(func() { assert.NoError(t, reader.Close()) })
+
+			// An existing DELETE-access handle requires new readers to share delete.
+			_, err = os.ReadFile(path)
+			require.ErrorIs(t, err, windows.ERROR_SHARING_VIOLATION)
+			got, err := ReadStateFile(path)
+			if !shareDelete {
+				require.ErrorIs(t, err, windows.ERROR_SHARING_VIOLATION)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, original, got)
+
+			require.NoError(t, WriteAtomic(path, replacement))
+			got, err = ReadStateFile(path)
+			require.NoError(t, err)
+			assert.Equal(t, replacement, got)
+			old, err := io.ReadAll(reader)
+			require.NoError(t, err)
+			assert.Equal(t, original, old)
+		})
+	}
+}

internal/stack/lock_test.go

@@ -504,6 +504,50 @@ func TestAtomicPublication_PreservesExistingFiles(t *testing.T) {
 	})
 }
 
+func TestReadStateFile(t *testing.T) {
+	t.Run("reads complete bytes", func(t *testing.T) {
+		path := filepath.Join(t.TempDir(), "state")
+		for _, want := range [][]byte{
+			{},
+			{0x00, 0xff, 0x0a},
+			[]byte(strings.Repeat("state\x00\xff\n", 16384)),
+		} {
+			require.NoError(t, os.WriteFile(path, want, 0600))
+			got, err := ReadStateFile(path)
+			require.NoError(t, err)
+			assert.Equal(t, want, got)
+		}
+	})
+
+	t.Run("missing file", func(t *testing.T) {
+		path := filepath.Join(t.TempDir(), "missing")
+		_, err := ReadStateFile(path)
+		require.ErrorIs(t, err, os.ErrNotExist)
+		assert.NoFileExists(t, path)
+	})
+
+	t.Run("directory", func(t *testing.T) {
+		path := t.TempDir()
+		_, err := ReadStateFile(path)
+		var pathErr *os.PathError
+		require.ErrorAs(t, err, &pathErr)
+		assert.Equal(t, path, pathErr.Path)
+		assert.DirExists(t, path)
+	})
+
+	t.Run("unreadable file", func(t *testing.T) {
+		if runtime.GOOS == "windows" || os.Geteuid() == 0 {
+			t.Skip("requires Unix permission enforcement")
+		}
+		path := filepath.Join(t.TempDir(), "state")
+		require.NoError(t, os.WriteFile(path, []byte("private state"), 0600))
+		require.NoError(t, os.Chmod(path, 0))
+		t.Cleanup(func() { assert.NoError(t, os.Chmod(path, 0600)) })
+		_, err := ReadStateFile(path)
+		require.ErrorIs(t, err, os.ErrPermission)
+	})
+}
+
 func TestWriteAtomic(t *testing.T) {
 	t.Run("creates and replaces exact bytes", func(t *testing.T) {
 		path := filepath.Join(t.TempDir(), "gh-stack-rebase-state")