a4c844c test(snapshot): synchronize worker completion · vev

M internal/usecase/daemon/session.go +1 -0


@@ -50,6 +50,7 @@ type session struct {
50	50		snapshotMu sync.Mutex
51	51		snapshotGeneration uint64
52	52		snapshotPending bool
	53	+	snapshotChanged chan struct{}
53	54		// syncGen makes synchronized-output watchdog generations unique across all
54	55		// panes in this session.
55	56		syncGen atomic.Uint64


@@ -50,6 +50,7 @@ type session struct {
50	snapshotMu sync.Mutex	50		snapshotMu sync.Mutex
51	snapshotGeneration uint64	51		snapshotGeneration uint64
52	snapshotPending bool	52		snapshotPending bool
		53	+	snapshotChanged chan struct{}
53	// syncGen makes synchronized-output watchdog generations unique across all	54		// syncGen makes synchronized-output watchdog generations unique across all
54	// panes in this session.	55		// panes in this session.
55	syncGen atomic.Uint64	56		syncGen atomic.Uint64

M internal/usecase/daemon/snapshot.go +19 -0


@@ -253,9 +253,26 @@ func markSnapshotDirty(sess *session) {
253	253		sess.snapshotMu.Lock()
254	254		sess.snapshotGeneration++
255	255		sess.snapDirty.Store(true)
	256	+	sess.signalSnapshotChangedLocked()
256	257		sess.snapshotMu.Unlock()
257	258		}
258	259
	260	+	// snapshotChangeLocked returns a channel closed by the next snapshot state
	261	+	// transition. snapshotMu must be held.
	262	+	func (sess *session) snapshotChangeLocked() chan struct{} {
	263	+	if sess.snapshotChanged == nil {
	264	+	sess.snapshotChanged = make(chan struct{})
	265	+	}
	266	+	return sess.snapshotChanged
	267	+	}
	268	+
	269	+	// signalSnapshotChangedLocked wakes waiters after a snapshot state transition.
	270	+	// snapshotMu must be held.
	271	+	func (sess *session) signalSnapshotChangedLocked() {
	272	+	close(sess.snapshotChangeLocked())
	273	+	sess.snapshotChanged = make(chan struct{})
	274	+	}
	275	+
259	276		func restorePaneTerminal(p *pane, snap snapcodec.Pane) error {
260	277		if len(snap.Tail) == 0 {
261	278		return fmt.Errorf("snapshot history: missing tail blob")

@@ -308,6 +325,7 @@ func (d Daemon) scheduleSnapshot(sess session) bool {
308	325		}
309	326		generation := sess.snapshotGeneration
310	327		sess.snapshotPending = true
	328	+	sess.signalSnapshotChangedLocked()
311	329		sess.snapshotMu.Unlock()
312	330
313	331		capture, ok := d.captureSnapshotState(sess, generation)

@@ -454,6 +472,7 @@ func (d Daemon) finishSnapshotCapture(capture snapshotCapture, succeeded bool)
454	472		} else if !succeeded {
455	473		capture.session.snapDirty.Store(true)
456	474		}
	475	+	capture.session.signalSnapshotChangedLocked()
457	476		capture.session.snapshotMu.Unlock()
458	477		})
459	478		}


@@ -253,9 +253,26 @@ func markSnapshotDirty(sess *session) {
253	sess.snapshotMu.Lock()	253		sess.snapshotMu.Lock()
254	sess.snapshotGeneration++	254		sess.snapshotGeneration++
255	sess.snapDirty.Store(true)	255		sess.snapDirty.Store(true)
		256	+	sess.signalSnapshotChangedLocked()
256	sess.snapshotMu.Unlock()	257		sess.snapshotMu.Unlock()
257	}	258		}
258		259
		260	+	// snapshotChangeLocked returns a channel closed by the next snapshot state
		261	+	// transition. snapshotMu must be held.
		262	+	func (sess *session) snapshotChangeLocked() chan struct{} {
		263	+	if sess.snapshotChanged == nil {
		264	+	sess.snapshotChanged = make(chan struct{})
		265	+	}
		266	+	return sess.snapshotChanged
		267	+	}
		268	+
		269	+	// signalSnapshotChangedLocked wakes waiters after a snapshot state transition.
		270	+	// snapshotMu must be held.
		271	+	func (sess *session) signalSnapshotChangedLocked() {
		272	+	close(sess.snapshotChangeLocked())
		273	+	sess.snapshotChanged = make(chan struct{})
		274	+	}
		275	+
259	func restorePaneTerminal(p *pane, snap snapcodec.Pane) error {	276		func restorePaneTerminal(p *pane, snap snapcodec.Pane) error {
260	if len(snap.Tail) == 0 {	277		if len(snap.Tail) == 0 {
261	return fmt.Errorf("snapshot history: missing tail blob")	278		return fmt.Errorf("snapshot history: missing tail blob")

@@ -308,6 +325,7 @@ func (d Daemon) scheduleSnapshot(sess session) bool {
308	}	325		}
309	generation := sess.snapshotGeneration	326		generation := sess.snapshotGeneration
310	sess.snapshotPending = true	327		sess.snapshotPending = true
		328	+	sess.signalSnapshotChangedLocked()
311	sess.snapshotMu.Unlock()	329		sess.snapshotMu.Unlock()
312		330
313	capture, ok := d.captureSnapshotState(sess, generation)	331		capture, ok := d.captureSnapshotState(sess, generation)

@@ -454,6 +472,7 @@ func (d Daemon) finishSnapshotCapture(capture snapshotCapture, succeeded bool)
454	} else if !succeeded {	472		} else if !succeeded {
455	capture.session.snapDirty.Store(true)	473		capture.session.snapDirty.Store(true)
456	}	474		}
		475	+	capture.session.signalSnapshotChangedLocked()
457	capture.session.snapshotMu.Unlock()	476		capture.session.snapshotMu.Unlock()
458	})	477		})
459	}	478		}

M internal/usecase/daemon/snapshot_test.go +35 -8


@@ -4,7 +4,6 @@ import (
4	4		"context"
5	5		"errors"
6	6		"io"
7		-	"runtime"
8	7		"sync"
9	8		"sync/atomic"
10	9		"testing"

@@ -288,6 +287,31 @@ func TestSnapshotQueueSaturationLeavesCaptureRetryable(t *testing.T) {
288	287		awaitSnapshotClean(t, third)
289	288		}
290	289
	290	+	func TestSnapshotCompletionSignalsWaiters(t *testing.T) {
	291	+	store := &channelSnapshotStore{writes: make(chan []byte, 1)}
	292	+	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})
	293	+	WithSnapshotStore(store)(d)
	294	+	startSnapshotEncodeWorker(t, d)
	295	+
	296	+	entered := make(chan struct{})
	297	+	release := make(chan struct{})
	298	+	d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) {
	299	+	close(entered)
	300	+	<-release
	301	+	return snapcodec.Marshal(s)
	302	+	}
	303	+	sess := newSnapshotTestSession(t, "completion", false, "/work")
	304	+	require.True(t, d.captureSession(sess))
	305	+	<-entered
	306	+
	307	+	sess.snapshotMu.Lock()
	308	+	changed := sess.snapshotChangeLocked()
	309	+	sess.snapshotMu.Unlock()
	310	+	close(release)
	311	+	<-changed
	312	+	awaitSnapshotClean(t, sess)
	313	+	}
	314	+
291	315		func TestSnapshotSuccessDoesNotClearNewerGeneration(t *testing.T) {
292	316		store := &channelSnapshotStore{writes: make(chan []byte, 2)}
293	317		d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})

@@ -794,27 +818,30 @@ func newSnapshotTestSession(t *testing.T, name string, ephemeral bool, cwd strin
794	818
795	819		func awaitSnapshotIdle(t testing.T, sess session) {
796	820		t.Helper()
797		-	for range 4096 {
	821	+	for {
798	822		sess.snapshotMu.Lock()
799	823		pending := sess.snapshotPending
	824	+	changed := sess.snapshotChangeLocked()
800	825		sess.snapshotMu.Unlock()
801	826		if !pending {
802	827		return
803	828		}
804		-	runtime.Gosched()
	829	+	<-changed
805	830		}
806		-	t.Fatal("snapshot worker did not become idle")
807	831		}
808	832
809	833		func awaitSnapshotClean(t testing.T, sess session) {
810	834		t.Helper()
811		-	for range 4096 {
812		-	if !sess.snapDirty.Load() {
	835	+	for {
	836	+	sess.snapshotMu.Lock()
	837	+	clean := !sess.snapDirty.Load()
	838	+	changed := sess.snapshotChangeLocked()
	839	+	sess.snapshotMu.Unlock()
	840	+	if clean {
813	841		return
814	842		}
815		-	runtime.Gosched()
	843	+	<-changed
816	844		}
817		-	t.Fatal("snapshot worker did not finish")
818	845		}
819	846
820	847		func startSnapshotEncodeWorker(t testing.T, d Daemon) {

 	"context"
 	"errors"
 	"io"
-	"runtime"
 	"sync"
 	"sync/atomic"
 	"testing"
 	awaitSnapshotClean(t, third)
 }
+func TestSnapshotCompletionSignalsWaiters(t *testing.T) {
+	store := &channelSnapshotStore{writes: make(chan []byte, 1)}
+	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})
+	WithSnapshotStore(store)(d)
+	startSnapshotEncodeWorker(t, d)
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) {
+		close(entered)
+		<-release
+		return snapcodec.Marshal(s)
+	}
+	sess := newSnapshotTestSession(t, "completion", false, "/work")
+	require.True(t, d.captureSession(sess))
+	<-entered
+	sess.snapshotMu.Lock()
+	changed := sess.snapshotChangeLocked()
+	sess.snapshotMu.Unlock()
+	close(release)
+	<-changed
+	awaitSnapshotClean(t, sess)
+}
 func TestSnapshotSuccessDoesNotClearNewerGeneration(t *testing.T) {
 	store := &channelSnapshotStore{writes: make(chan []byte, 2)}
 	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})
 func awaitSnapshotIdle(t *testing.T, sess *session) {
 	t.Helper()
-	for range 4096 {
+	for {
 		sess.snapshotMu.Lock()
 		pending := sess.snapshotPending
+		changed := sess.snapshotChangeLocked()
 		sess.snapshotMu.Unlock()
 		if !pending {
 			return
 		}
-		runtime.Gosched()
+		<-changed
 	}
-	t.Fatal("snapshot worker did not become idle")
 }
 func awaitSnapshotClean(t *testing.T, sess *session) {
 	t.Helper()
-	for range 4096 {
+	for {
-		if !sess.snapDirty.Load() {
+		sess.snapshotMu.Lock()
+		clean := !sess.snapDirty.Load()
+		changed := sess.snapshotChangeLocked()
+		sess.snapshotMu.Unlock()
+		if clean {
 			return
 		}
-		runtime.Gosched()
+		<-changed
 	}
-	t.Fatal("snapshot worker did not finish")
 }
 func startSnapshotEncodeWorker(t *testing.T, d *Daemon) {