fix(mount): join background writers before shutdown

cmd/relayfile-cli/main.go

@@ -14353,6 +14353,10 @@ func runMountLoop(rootCtx context.Context, syncer *mountsync.Syncer, localDir, w
 }
 
 func runMountLoopWithAuthLock(rootCtx context.Context, syncer *mountsync.Syncer, localDir, workspaceID, serverURL, delegatedCredsFile string, timeout, interval time.Duration, intervalJitter float64, websocketEnabled, once, daemonized bool, pidFile, logFile string, authMu *sync.Mutex) error {
+	// The loop owns every deferred receipt/checkpoint task admitted by this
+	// Syncer. Join them on every exit path before the mount directory can be
+	// removed or handed to another process.
+	defer syncer.Close()
 	interval = enforcePollIntervalFloor(interval)
 	httpClient, _ := syncerClient(syncer)
 	record, _ := workspaceRecordByID(workspaceID)

internal/mountsync/realtime_collaboration_test.go

@@ -26,9 +26,11 @@ type delayedIncrementalReadClient struct {
 
 type blockingReceiptClient struct {
 	*fakeClient
-	started chan struct{}
-	release chan struct{}
-	once    sync.Once
+	started  chan struct{}
+	canceled chan struct{}
+	release  chan struct{}
+	once     sync.Once
+	cancel   sync.Once
 }
 
 func TestUntrackedAtomicSaveStagingFileNeverWritesBack(t *testing.T) {
@@ -72,12 +74,12 @@ func TestUntrackedAtomicSaveStagingFileNeverWritesBack(t *testing.T) {
 
 func (c *blockingReceiptClient) GetOperation(ctx context.Context, workspaceID, opID string) (OperationStatus, error) {
 	c.once.Do(func() { close(c.started) })
-	select {
-	case <-c.release:
-		return c.fakeClient.GetOperation(ctx, workspaceID, opID)
-	case <-ctx.Done():
-		return OperationStatus{}, ctx.Err()
+	if c.canceled != nil {
+		<-ctx.Done()
+		c.cancel.Do(func() { close(c.canceled) })
 	}
+	<-c.release
+	return c.fakeClient.GetOperation(context.Background(), workspaceID, opID)
 }
 
 func (c *delayedIncrementalReadClient) ReadFile(ctx context.Context, workspaceID, path string) (RemoteFile, error) {
@@ -414,8 +416,9 @@ func TestHandleLocalChangesBatchesElevenFilesAndDefersPendingReceipts(t *testing
 			files:      map[string]RemoteFile{},
 			operations: map[string]OperationStatus{},
 		},
-		started: make(chan struct{}),
-		release: make(chan struct{}),
+		started:  make(chan struct{}),
+		canceled: make(chan struct{}),
+		release:  make(chan struct{}),
 	}
 	client.bulkWriteResponseFunc = func(_ context.Context, _ string, files []BulkWriteFile) (BulkWriteResponse, error) {
 		results := make([]BulkWriteResult, 0, len(files))
@@ -473,16 +476,87 @@ func TestHandleLocalChangesBatchesElevenFilesAndDefersPendingReceipts(t *testing
 	if client.getOperationCalls != 0 {
 		t.Fatalf("blocked receipt unexpectedly completed: calls=%d", client.getOperationCalls)
 	}
+
+	// Shutdown must cancel and join the blocked receipt worker as well as the
+	// checkpoint callback it would otherwise leave behind. Returning from Close
+	// is the boundary that makes immediate TempDir cleanup safe.
+	closed := make(chan struct{})
+	go func() {
+		syncer.Close()
+		close(closed)
+	}()
+	select {
+	case <-client.canceled:
+	case <-time.After(time.Second):
+		t.Fatal("syncer Close did not cancel deferred receipt work")
+	}
+	select {
+	case <-closed:
+		t.Fatal("syncer Close returned before the receipt writer terminated")
+	default:
+	}
 	close(client.release)
-	deadline := time.Now().Add(time.Second)
-	for time.Now().Before(deadline) {
-		syncer.receiptMu.Lock()
-		active := len(syncer.receiptActive)
-		syncer.receiptMu.Unlock()
-		if active == 0 {
-			break
+	select {
+	case <-closed:
+	case <-time.After(time.Second):
+		t.Fatal("syncer Close did not join the terminated receipt writer")
+	}
+	syncer.receiptMu.Lock()
+	active := len(syncer.receiptActive)
+	syncer.receiptMu.Unlock()
+	if active != 0 {
+		t.Fatalf("receipt workers still active after Close: %d", active)
+	}
+}
+
+func TestSyncerCloseJoinsRunningCheckpointWriter(t *testing.T) {
+	localDir := t.TempDir()
+	syncer, err := NewSyncer(&fakeClient{files: map[string]RemoteFile{}}, SyncerOptions{
+		WorkspaceID: "ws_checkpoint_close",
+		RemoteRoot:  "/",
+		LocalRoot:   localDir,
+	})
+	if err != nil {
+		t.Fatalf("new syncer: %v", err)
+	}
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+	syncer.checkpointTestHook = func(stage string) {
+		if stage == "local-write-checkpoint-before-save" {
+			close(started)
+			<-release
 		}
-		time.Sleep(time.Millisecond)
+	}
+	syncer.scheduleLocalWriteCheckpoint()
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("checkpoint writer did not start")
+	}
+
+	closed := make(chan struct{})
+	go func() {
+		syncer.Close()
+		close(closed)
+	}()
+	select {
+	case <-closed:
+		t.Fatal("syncer Close returned while checkpoint writer was running")
+	default:
+	}
+	close(release)
+	select {
+	case <-closed:
+	case <-time.After(time.Second):
+		t.Fatal("syncer Close did not join checkpoint writer")
+	}
+
+	if err := os.RemoveAll(localDir); err != nil {
+		t.Fatalf("remove mount after Close: %v", err)
+	}
+	if _, err := os.Stat(localDir); !errors.Is(err, os.ErrNotExist) {
+		t.Fatalf("mount root recreated after Close: %v", err)
 	}
 }
 

internal/mountsync/syncer.go

@@ -1805,6 +1805,11 @@ type Syncer struct {
 	localMutationMu           sync.Mutex
 	websocket                 bool
 	rootCtx                   context.Context
+	backgroundCtx             context.Context
+	backgroundCancel          context.CancelFunc
+	backgroundMu              sync.Mutex
+	backgroundClosed          bool
+	backgroundWG              sync.WaitGroup
 	wsConn                    *websocket.Conn
 	wsCancel                  context.CancelFunc
 	wsNextAttempt             time.Time
@@ -2633,6 +2638,10 @@ func NewSyncer(client RemoteClient, opts SyncerOptions) (*Syncer, error) {
 	} else if _, ok := client.(*HTTPClient); ok {
 		cache = defaultObjectCache()
 	}
+	// Create the owned lifecycle context only after every constructor path that
+	// can still return an error. From here it is transferred directly to the
+	// Syncer and released by Close.
+	backgroundCtx, backgroundCancel := context.WithCancel(rootCtx)
 	return &Syncer{
 		client:                    client,
 		objectCache:               cache,
@@ -2653,6 +2662,8 @@ func NewSyncer(client RemoteClient, opts SyncerOptions) (*Syncer, error) {
 		websocket:                 websocketEnabled,
 		recoverStartupDrift:       true,
 		rootCtx:                   rootCtx,
+		backgroundCtx:             backgroundCtx,
+		backgroundCancel:          backgroundCancel,
 		logger:                    opts.Logger,
 		denialLogPath:             filepath.Join(localRoot, ".relay", "permissions-denied.log"),
 		bulkFlushThreshold:        bulkFlushThreshold,
@@ -2693,6 +2704,60 @@ func NewSyncer(client RemoteClient, opts SyncerOptions) (*Syncer, error) {
 	}, nil
 }
 
+// Close cancels and joins every deferred receipt and checkpoint writer owned by
+// the Syncer. Callers must stop admitting foreground work before calling Close.
+// Close is idempotent and does not return until no owned background task can
+// write beneath the mount root.
+func (s *Syncer) Close() {
+	if s == nil {
+		return
+	}
+	s.backgroundMu.Lock()
+	if !s.backgroundClosed {
+		s.backgroundClosed = true
+		if s.backgroundCancel != nil {
+			s.backgroundCancel()
+		}
+	}
+	s.backgroundMu.Unlock()
+	// A websocket reader owns a second receive goroutine and may be blocked in
+	// the transport even after its context is canceled. Closing the connection
+	// wakes that read; both goroutines are joined through backgroundWG below.
+	s.ResetWebSocket()
+
+	// Prevent a pending checkpoint timer from becoming a writer after shutdown.
+	// Timers whose callbacks already started remain counted in backgroundWG and
+	// are joined below.
+	s.checkpointMu.Lock()
+	s.checkpointVersion++
+	if s.checkpointTimer != nil && s.checkpointTimer.Stop() {
+		s.backgroundWG.Done()
+	}
+	s.checkpointTimer = nil
+	s.checkpointStarted = time.Time{}
+	s.checkpointMu.Unlock()
+
+	s.backgroundWG.Wait()
+}
+
+func (s *Syncer) startBackground(run func(context.Context)) bool {
+	s.backgroundMu.Lock()
+	defer s.backgroundMu.Unlock()
+	if s.backgroundClosed {
+		return false
+	}
+	ctx := s.backgroundCtx
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	s.backgroundWG.Add(1)
+	go func() {
+		defer s.backgroundWG.Done()
+		run(ctx)
+	}()
+	return true
+}
+
 // Circuit returns the cloud-error breaker for tests and status reporters.
 func (s *Syncer) Circuit() *CloudErrorCircuit { return s.circuit }
 
@@ -4382,16 +4447,12 @@ func (s *Syncer) scheduleOutboxReceiptSettlements(records []outboxRecord) {
 		sem := s.receiptSem
 		s.receiptMu.Unlock()
 
-		go func() {
+		started := s.startBackground(func(root context.Context) {
 			defer func() {
 				s.receiptMu.Lock()
 				delete(s.receiptActive, opID)
 				s.receiptMu.Unlock()
 			}()
-			root := s.rootCtx
-			if root == nil {
-				root = context.Background()
-			}
 			select {
 			case sem <- struct{}{}:
 				defer func() { <-sem }()
@@ -4430,7 +4491,12 @@ func (s *Syncer) scheduleOutboxReceiptSettlements(records []outboxRecord) {
 			// derived state document instead of making every receipt worker hold
 			// the main state mutex through a full outbox summary scan.
 			s.scheduleLocalWriteCheckpoint()
-		}()
+		})
+		if !started {
+			s.receiptMu.Lock()
+			delete(s.receiptActive, opID)
+			s.receiptMu.Unlock()
+		}
 	}
 }
 
@@ -4452,9 +4518,21 @@ func (s *Syncer) scheduleLocalWriteCheckpoint() {
 		delay = time.Nanosecond
 	}
 	if s.checkpointTimer != nil {
-		s.checkpointTimer.Stop()
+		if s.checkpointTimer.Stop() {
+			s.backgroundWG.Done()
+		}
+	}
+	s.backgroundMu.Lock()
+	if s.backgroundClosed {
+		s.backgroundMu.Unlock()
+		s.checkpointTimer = nil
+		s.checkpointStarted = time.Time{}
+		return
 	}
+	s.backgroundWG.Add(1)
+	s.backgroundMu.Unlock()
 	s.checkpointTimer = time.AfterFunc(delay, func() {
+		defer s.backgroundWG.Done()
 		s.checkpointMu.Lock()
 		if s.checkpointVersion != version {
 			s.checkpointMu.Unlock()
@@ -4464,6 +4542,7 @@ func (s *Syncer) scheduleLocalWriteCheckpoint() {
 		s.checkpointStarted = time.Time{}
 		s.checkpointMu.Unlock()
 
+		s.runCheckpointTestHook("local-write-checkpoint-before-save")
 		s.mu.Lock()
 		// Match the WebSocket checkpoint contract: persist only the private
 		// recovery cursor/state on the burst path. The public .relay/state.json
@@ -5487,7 +5566,7 @@ func (s *Syncer) connectWebSocket(ctx context.Context) error {
 	}
 	conn.SetReadLimit(maxWebSocketMessageBytes)
 
-	readCtx, cancel := context.WithCancel(s.rootCtx)
+	readCtx, cancel := context.WithCancel(s.backgroundCtx)
 
 	s.mu.Lock()
 	if s.wsGeneration != generation || s.wsConn != nil {
@@ -5507,16 +5586,25 @@ func (s *Syncer) connectWebSocket(ctx context.Context) error {
 	s.wsLastConnectedAt = time.Now().UTC()
 	s.mu.Unlock()
 
-	go s.readWebSocketLoop(readCtx, conn)
+	if !s.startBackground(func(context.Context) {
+		s.readWebSocketLoop(readCtx, conn)
+	}) {
+		cancel()
+		s.ResetWebSocket()
+	}
 	return nil
 }
 
 func (s *Syncer) readWebSocketLoop(ctx context.Context, conn *websocket.Conn) {
 	defer s.handleWebSocketDisconnect(conn)
+	ctx, cancelReader := context.WithCancel(ctx)
 
 	eventCh := make(chan websocketEvent, webSocketApplyQueueSize)
 	readErrCh := make(chan error, 1)
+	var readerWG sync.WaitGroup
+	readerWG.Add(1)
 	go func() {
+		defer readerWG.Done()
 		for {
 			var event websocketEvent
 			if err := wsjson.Read(ctx, conn, &event); err != nil {
@@ -5533,6 +5621,13 @@ func (s *Syncer) readWebSocketLoop(ctx context.Context, conn *websocket.Conn) {
 			}
 		}
 	}()
+	defer func() {
+		// The apply loop can stop before the transport read does (for example,
+		// after a persistence failure). Cancel first, then join, so the nested
+		// reader cannot survive the lifecycle owner or deadlock its return path.
+		cancelReader()
+		readerWG.Wait()
+	}()
 
 	var checkpointTimer *time.Timer
 	var checkpointC <-chan time.Time
@@ -6167,7 +6262,9 @@ func (s *Syncer) bootstrapContext(parent context.Context) (context.Context, cont
 		pollEvery = 10 * time.Millisecond
 	}
 	done := make(chan struct{})
+	stopped := make(chan struct{})
 	go func() {
+		defer close(stopped)
 		ticker := time.NewTicker(pollEvery)
 		defer ticker.Stop()
 		for {
@@ -6192,6 +6289,7 @@ func (s *Syncer) bootstrapContext(parent context.Context) (context.Context, cont
 	wrapped := func() {
 		close(done)
 		cancel()
+		<-stopped
 	}
 	return ctx, wrapped, prog, nil
 }

internal/mountsync/watcher.go

@@ -336,7 +336,9 @@ func (fw *FileWatcher) Start(ctx context.Context) error {
 	fw.healthy.Store(true)
 
 	// Event loop
+	fw.wg.Add(1)
 	go func() {
+		defer fw.wg.Done()
 		defer fw.healthy.Store(false)
 		for {
 			select {
@@ -658,6 +660,9 @@ func (fw *FileWatcher) Close() error {
 	fw.mu.Unlock()
 
 	err := fw.watcher.Close()
+	// Closing the backend terminates the fsnotify event loop. Join both that
+	// loop and every admitted debounce callback before returning so callers can
+	// safely remove the watched tree immediately after Close.
 	fw.wg.Wait()
 	return err
 }