7913f91 fix(snapshot): retain final snapshots when queue is full · vev

M internal/usecase/daemon/daemon.go +9 -1


@@ -110,6 +110,11 @@ type Daemon struct {
110	110		snapshotWorkerCancel context.CancelFunc
111	111		snapshotWorkerDone chan struct{}
112	112		snapshotWorkerFlush chan struct{}
	113	+	snapshotWorkerFinalWake chan struct{}
	114	+	// snapshotFinalJobs retains immutable terminal captures rejected by the
	115	+	// bounded regular queue. It is only populated during session teardown and
	116	+	// drained by the same worker, so persistence remains outside daemon locks.
	117	+	snapshotFinalJobs []*snapshotCapture
113	118		snapshotWorkerClosing bool
114	119		snapshotWorkerInFlight *snapshotCapture
115	120		restoreDone chan struct{}

@@ -417,8 +422,11 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error {
417	422		d.waitNotifies()
418	423		d.hardCancel()
419	424		d.serveCancel()
420		-	d.stopSnapshotEncodeWorker()
	425	+	// Wait for handlers before flushing snapshots: a handler that already
	426	+	// entered killSession may still submit its terminal capture after the
	427	+	// registry snapshot has been removed.
421	428		d.connWg.Wait()
	429	+	d.stopSnapshotEncodeWorker()
422	430		d.shutdownAll(ports.ReasonServerShutdown)
423	431		d.sessWg.Wait()
424	432		d.waitNotifies()

 	snapshotWorkerCancel    context.CancelFunc
 	snapshotWorkerDone      chan struct{}
 	snapshotWorkerFlush     chan struct{}
+	snapshotWorkerFinalWake chan struct{}
+	// snapshotFinalJobs retains immutable terminal captures rejected by the
+	// bounded regular queue. It is only populated during session teardown and
+	// drained by the same worker, so persistence remains outside daemon locks.
+	snapshotFinalJobs       []*snapshotCapture
 	snapshotWorkerClosing   bool
 	snapshotWorkerInFlight  *snapshotCapture
 	restoreDone             chan struct{}
 	d.waitNotifies()
 	d.hardCancel()
 	d.serveCancel()
-	d.stopSnapshotEncodeWorker()
+	// Wait for handlers before flushing snapshots: a handler that already
+	// entered killSession may still submit its terminal capture after the
+	// registry snapshot has been removed.
 	d.connWg.Wait()
+	d.stopSnapshotEncodeWorker()
 	d.shutdownAll(ports.ReasonServerShutdown)
 	d.sessWg.Wait()
 	d.waitNotifies()

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


@@ -679,7 +679,7 @@ func (d Daemon) killSession(sess session, reason uint8, purge bool) error {
679	679		}
680	680		} else {
681	681		markSnapshotDirty(sess)
682		-	d.scheduleSnapshot(sess)
	682	+	d.scheduleFinalSnapshot(sess)
683	683		}
684	684		}
685	685		d.mu.Lock()


@@ -679,7 +679,7 @@ func (d Daemon) killSession(sess session, reason uint8, purge bool) error {
679		}	679		}
680		} else {	680		} else {
681		markSnapshotDirty(sess)	681		markSnapshotDirty(sess)
682	-	d.scheduleSnapshot(sess)	682	+	d.scheduleFinalSnapshot(sess)
683		}	683		}
684		}	684		}
685		d.mu.Lock()	685		d.mu.Lock()

M internal/usecase/daemon/snapshot.go +63 -2


@@ -315,6 +315,17 @@ func restorePTYSize(contentRect domain.Rect, tabSize domain.Size) domain.Size {
315	315		// scheduleSnapshot captures a session only once while an immutable capture is
316	316		// queued or being written. It never waits for the bounded worker queue.
317	317		func (d Daemon) scheduleSnapshot(sess session) bool {
	318	+	return d.scheduleSnapshotWithFinalFallback(sess, false)
	319	+	}
	320	+
	321	+	// scheduleFinalSnapshot retains a terminal capture rejected by the regular
	322	+	// bounded queue. A removed session cannot be retried by snapshotSaver, so the
	323	+	// worker drains this fallback before it stops.
	324	+	func (d Daemon) scheduleFinalSnapshot(sess session) bool {
	325	+	return d.scheduleSnapshotWithFinalFallback(sess, true)
	326	+	}
	327	+
	328	+	func (d Daemon) scheduleSnapshotWithFinalFallback(sess session, final bool) bool {
318	329		if d == nil \|\| sess == nil \|\| !d.snapsEnabled \|\| d.snaps == nil \|\| !sess.snapEligible.Load() {
319	330		return true
320	331		}

@@ -333,7 +344,7 @@ func (d Daemon) scheduleSnapshot(sess session) bool {
333	344		d.finishSnapshotCapture(capture, false)
334	345		return false
335	346		}
336		-	if d.enqueueSnapshotCapture(capture) {
	347	+	if d.enqueueSnapshotCapture(capture) \|\| (final && d.enqueueFinalSnapshotCapture(capture)) {
337	348		return true
338	349		}
339	350		// Coalesce under saturation or shutdown: the latest state stays dirty and

@@ -359,6 +370,22 @@ func (d Daemon) enqueueSnapshotCapture(capture snapshotCapture) bool {
359	370		}
360	371		}
361	372
	373	+	// enqueueFinalSnapshotCapture hands terminal captures to the existing worker
	374	+	// without making session teardown wait behind its normal bounded queue.
	375	+	func (d Daemon) enqueueFinalSnapshotCapture(capture snapshotCapture) bool {
	376	+	d.snapshotWorkerMu.Lock()
	377	+	defer d.snapshotWorkerMu.Unlock()
	378	+	if d.snapshotWorkerClosing \|\| d.snapshotWorkerCancel == nil \|\| d.snapshotWorkerCtx == nil \|\| d.snapshotWorkerCtx.Err() != nil {
	379	+	return false
	380	+	}
	381	+	d.snapshotFinalJobs = append(d.snapshotFinalJobs, capture)
	382	+	select {
	383	+	case d.snapshotWorkerFinalWake <- struct{}{}:
	384	+	default:
	385	+	}
	386	+	return true
	387	+	}
	388	+
362	389		// captureSession rotates history tails and clones visible frames while holding
363	390		// each pane lock. The returned capture contains only immutable state; encoding
364	391		// and persistence are deliberately deferred to snapshotEncodeWorker.

@@ -495,10 +522,13 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) {
495	522		d.snapshotWorkerCancel = cancel
496	523		d.snapshotWorkerDone = make(chan struct{})
497	524		d.snapshotWorkerFlush = make(chan struct{})
	525	+	d.snapshotWorkerFinalWake = make(chan struct{}, 1)
	526	+	d.snapshotFinalJobs = nil
498	527		d.snapshotWorkerClosing = false
499	528		d.snapshotWorkerInFlight = nil
500	529		done := d.snapshotWorkerDone
501	530		flush := d.snapshotWorkerFlush
	531	+	finalWake := d.snapshotWorkerFinalWake
502	532		d.snapshotWorkerMu.Unlock()
503	533		go func() {
504	534		defer close(done)

@@ -518,6 +548,17 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) {
518	548		d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil)
519	549		return workerCtx.Err() == nil
520	550		}
	551	+	drainFinal := func() bool {
	552	+	for {
	553	+	capture := d.takeFinalSnapshotCapture()
	554	+	if capture == nil {
	555	+	return true
	556	+	}
	557	+	if !write(capture) {
	558	+	return false
	559	+	}
	560	+	}
	561	+	}
521	562		flushQueued := func() {
522	563		for {
523	564		select {

@@ -526,6 +567,7 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) {
526	567		return
527	568		}
528	569		default:
	570	+	_ = drainFinal()
529	571		return
530	572		}
531	573		}

@@ -537,6 +579,10 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) {
537	579		case <-flush:
538	580		flushQueued()
539	581		return
	582	+	case <-finalWake:
	583	+	if !drainFinal() {
	584	+	return
	585	+	}
540	586		case capture := <-d.snapshotJobs:
541	587		if !write(capture) {
542	588		return

@@ -546,6 +592,18 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) {
546	592		}()
547	593		}
548	594
	595	+	func (d Daemon) takeFinalSnapshotCapture() snapshotCapture {
	596	+	d.snapshotWorkerMu.Lock()
	597	+	defer d.snapshotWorkerMu.Unlock()
	598	+	if len(d.snapshotFinalJobs) == 0 {
	599	+	return nil
	600	+	}
	601	+	capture := d.snapshotFinalJobs[0]
	602	+	d.snapshotFinalJobs[0] = nil
	603	+	d.snapshotFinalJobs = d.snapshotFinalJobs[1:]
	604	+	return capture
	605	+	}
	606	+
549	607		func (d Daemon) setSnapshotWorkerInFlight(workerID uint64, capture snapshotCapture) bool {
550	608		d.snapshotWorkerMu.Lock()
551	609		defer d.snapshotWorkerMu.Unlock()

@@ -601,9 +659,12 @@ func (d *Daemon) finishStoppedSnapshotWorker(abandoned bool) {
601	659		d.snapshotWorkerCancel = nil
602	660		d.snapshotWorkerDone = nil
603	661		d.snapshotWorkerFlush = nil
	662	+	d.snapshotWorkerFinalWake = nil
604	663		d.snapshotWorkerClosing = false
605	664		d.snapshotWorkerInFlight = nil
606		-	queued := make([]*snapshotCapture, 0, len(d.snapshotJobs))
	665	+	queued := make([]*snapshotCapture, 0, len(d.snapshotJobs)+len(d.snapshotFinalJobs))
	666	+	queued = append(queued, d.snapshotFinalJobs...)
	667	+	d.snapshotFinalJobs = nil
607	668		for {
608	669		select {
609	670		case capture := <-d.snapshotJobs:

 // scheduleSnapshot captures a session only once while an immutable capture is
 // queued or being written. It never waits for the bounded worker queue.
 func (d *Daemon) scheduleSnapshot(sess *session) bool {
+	return d.scheduleSnapshotWithFinalFallback(sess, false)
+}
+// scheduleFinalSnapshot retains a terminal capture rejected by the regular
+// bounded queue. A removed session cannot be retried by snapshotSaver, so the
+// worker drains this fallback before it stops.
+func (d *Daemon) scheduleFinalSnapshot(sess *session) bool {
+	return d.scheduleSnapshotWithFinalFallback(sess, true)
+}
+func (d *Daemon) scheduleSnapshotWithFinalFallback(sess *session, final bool) bool {
 	if d == nil || sess == nil || !d.snapsEnabled || d.snaps == nil || !sess.snapEligible.Load() {
 		return true
 	}
 		d.finishSnapshotCapture(capture, false)
 		return false
 	}
-	if d.enqueueSnapshotCapture(capture) {
+	if d.enqueueSnapshotCapture(capture) || (final && d.enqueueFinalSnapshotCapture(capture)) {
 		return true
 	}
 	// Coalesce under saturation or shutdown: the latest state stays dirty and
 	}
 }
+// enqueueFinalSnapshotCapture hands terminal captures to the existing worker
+// without making session teardown wait behind its normal bounded queue.
+func (d *Daemon) enqueueFinalSnapshotCapture(capture *snapshotCapture) bool {
+	d.snapshotWorkerMu.Lock()
+	defer d.snapshotWorkerMu.Unlock()
+	if d.snapshotWorkerClosing || d.snapshotWorkerCancel == nil || d.snapshotWorkerCtx == nil || d.snapshotWorkerCtx.Err() != nil {
+		return false
+	}
+	d.snapshotFinalJobs = append(d.snapshotFinalJobs, capture)
+	select {
+	case d.snapshotWorkerFinalWake <- struct{}{}:
+	default:
+	}
+	return true
+}
 // captureSession rotates history tails and clones visible frames while holding
 // each pane lock. The returned capture contains only immutable state; encoding
 // and persistence are deliberately deferred to snapshotEncodeWorker.
 	d.snapshotWorkerCancel = cancel
 	d.snapshotWorkerDone = make(chan struct{})
 	d.snapshotWorkerFlush = make(chan struct{})
+	d.snapshotWorkerFinalWake = make(chan struct{}, 1)
+	d.snapshotFinalJobs = nil
 	d.snapshotWorkerClosing = false
 	d.snapshotWorkerInFlight = nil
 	done := d.snapshotWorkerDone
 	flush := d.snapshotWorkerFlush
+	finalWake := d.snapshotWorkerFinalWake
 	d.snapshotWorkerMu.Unlock()
 	go func() {
 		defer close(done)
 			d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil)
 			return workerCtx.Err() == nil
 		}
+		drainFinal := func() bool {
+			for {
+				capture := d.takeFinalSnapshotCapture()
+				if capture == nil {
+					return true
+				}
+				if !write(capture) {
+					return false
+				}
+			}
+		}
 		flushQueued := func() {
 			for {
 				select {
 						return
 					}
 				default:
+					_ = drainFinal()
 					return
 				}
 			}
 			case <-flush:
 				flushQueued()
 				return
+			case <-finalWake:
+				if !drainFinal() {
+					return
+				}
 			case capture := <-d.snapshotJobs:
 				if !write(capture) {
 					return
 	}()
 }
+func (d *Daemon) takeFinalSnapshotCapture() *snapshotCapture {
+	d.snapshotWorkerMu.Lock()
+	defer d.snapshotWorkerMu.Unlock()
+	if len(d.snapshotFinalJobs) == 0 {
+		return nil
+	}
+	capture := d.snapshotFinalJobs[0]
+	d.snapshotFinalJobs[0] = nil
+	d.snapshotFinalJobs = d.snapshotFinalJobs[1:]
+	return capture
+}
 func (d *Daemon) setSnapshotWorkerInFlight(workerID uint64, capture *snapshotCapture) bool {
 	d.snapshotWorkerMu.Lock()
 	defer d.snapshotWorkerMu.Unlock()
 	d.snapshotWorkerCancel = nil
 	d.snapshotWorkerDone = nil
 	d.snapshotWorkerFlush = nil
+	d.snapshotWorkerFinalWake = nil
 	d.snapshotWorkerClosing = false
 	d.snapshotWorkerInFlight = nil
-	queued := make([]*snapshotCapture, 0, len(d.snapshotJobs))
+	queued := make([]*snapshotCapture, 0, len(d.snapshotJobs)+len(d.snapshotFinalJobs))
+	queued = append(queued, d.snapshotFinalJobs...)
+	d.snapshotFinalJobs = nil
 	for {
 		select {
 		case capture := <-d.snapshotJobs:

M internal/usecase/daemon/snapshot_test.go +69 -0


@@ -287,6 +287,75 @@ func TestSnapshotQueueSaturationLeavesCaptureRetryable(t *testing.T) {
287	287		awaitSnapshotClean(t, third)
288	288		}
289	289
	290	+	func TestNamedFinalSnapshotSurvivesSaturatedQueue(t *testing.T) {
	291	+	for _, tc := range []struct {
	292	+	name string
	293	+	stop func(t testing.T, d Daemon, sess *session)
	294	+	}{
	295	+	{
	296	+	name: "kill session",
	297	+	stop: func(t testing.T, d Daemon, sess *session) {
	298	+	t.Helper()
	299	+	require.NoError(t, d.killSession(sess, ports.ReasonSessionKilled, false))
	300	+	},
	301	+	},
	302	+	{
	303	+	name: "shutdown all",
	304	+	stop: func(t testing.T, d Daemon, _ *session) {
	305	+	t.Helper()
	306	+	d.shutdownAll(ports.ReasonServerShutdown)
	307	+	},
	308	+	},
	309	+	} {
	310	+	t.Run(tc.name, func(t *testing.T) {
	311	+	store := &channelSnapshotStore{writes: make(chan []byte, 3)}
	312	+	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})
	313	+	WithSnapshotStore(store)(d)
	314	+	startSnapshotEncodeWorker(t, d)
	315	+
	316	+	entered := make(chan struct{})
	317	+	release := make(chan struct{})
	318	+	var calls atomic.Int32
	319	+	d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) {
	320	+	if calls.Add(1) == 1 {
	321	+	close(entered)
	322	+	<-release
	323	+	}
	324	+	return snapcodec.Marshal(s)
	325	+	}
	326	+	first := newSnapshotTestSession(t, "first", false, "/work")
	327	+	second := newSnapshotTestSession(t, "second", false, "/work")
	328	+	final := newSnapshotTestSession(t, "final", false, "/work")
	329	+	final.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY).EXPECT().Close().Return(nil).Once()
	330	+	markSnapshotDirty(first)
	331	+	markSnapshotDirty(second)
	332	+	require.True(t, d.scheduleSnapshot(first))
	333	+	<-entered
	334	+	require.True(t, d.scheduleSnapshot(second))
	335	+	d.sessions[final.id] = final
	336	+
	337	+	tc.stop(t, d, final)
	338	+	require.NotContains(t, d.sessions, final.id, "final capture must survive session removal")
	339	+	close(release)
	340	+
	341	+	written := make(map[string]bool, 3)
	342	+	deadline := time.NewTimer(time.Second)
	343	+	defer deadline.Stop()
	344	+	for len(written) < 3 {
	345	+	select {
	346	+	case data := <-store.writes:
	347	+	snap, err := snapcodec.Unmarshal(data)
	348	+	require.NoError(t, err)
	349	+	written[snap.Name] = true
	350	+	case <-deadline.C:
	351	+	t.Fatalf("final snapshot was dropped: wrote %v", written)
	352	+	}
	353	+	}
	354	+	require.True(t, written[final.name])
	355	+	})
	356	+	}
	357	+	}
	358	+
290	359		func TestSnapshotCompletionSignalsWaiters(t *testing.T) {
291	360		store := &channelSnapshotStore{writes: make(chan []byte, 1)}
292	361		d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})


@@ -287,6 +287,75 @@ func TestSnapshotQueueSaturationLeavesCaptureRetryable(t *testing.T) {
287	awaitSnapshotClean(t, third)	287		awaitSnapshotClean(t, third)
288	}	288		}
289		289
		290	+	func TestNamedFinalSnapshotSurvivesSaturatedQueue(t *testing.T) {
		291	+	for _, tc := range []struct {
		292	+	name string
		293	+	stop func(t testing.T, d Daemon, sess *session)
		294	+	}{
		295	+	{
		296	+	name: "kill session",
		297	+	stop: func(t testing.T, d Daemon, sess *session) {
		298	+	t.Helper()
		299	+	require.NoError(t, d.killSession(sess, ports.ReasonSessionKilled, false))
		300	+	},
		301	+	},
		302	+	{
		303	+	name: "shutdown all",
		304	+	stop: func(t testing.T, d Daemon, _ *session) {
		305	+	t.Helper()
		306	+	d.shutdownAll(ports.ReasonServerShutdown)
		307	+	},
		308	+	},
		309	+	} {
		310	+	t.Run(tc.name, func(t *testing.T) {
		311	+	store := &channelSnapshotStore{writes: make(chan []byte, 3)}
		312	+	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})
		313	+	WithSnapshotStore(store)(d)
		314	+	startSnapshotEncodeWorker(t, d)
		315	+
		316	+	entered := make(chan struct{})
		317	+	release := make(chan struct{})
		318	+	var calls atomic.Int32
		319	+	d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) {
		320	+	if calls.Add(1) == 1 {
		321	+	close(entered)
		322	+	<-release
		323	+	}
		324	+	return snapcodec.Marshal(s)
		325	+	}
		326	+	first := newSnapshotTestSession(t, "first", false, "/work")
		327	+	second := newSnapshotTestSession(t, "second", false, "/work")
		328	+	final := newSnapshotTestSession(t, "final", false, "/work")
		329	+	final.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY).EXPECT().Close().Return(nil).Once()
		330	+	markSnapshotDirty(first)
		331	+	markSnapshotDirty(second)
		332	+	require.True(t, d.scheduleSnapshot(first))
		333	+	<-entered
		334	+	require.True(t, d.scheduleSnapshot(second))
		335	+	d.sessions[final.id] = final
		336	+
		337	+	tc.stop(t, d, final)
		338	+	require.NotContains(t, d.sessions, final.id, "final capture must survive session removal")
		339	+	close(release)
		340	+
		341	+	written := make(map[string]bool, 3)
		342	+	deadline := time.NewTimer(time.Second)
		343	+	defer deadline.Stop()
		344	+	for len(written) < 3 {
		345	+	select {
		346	+	case data := <-store.writes:
		347	+	snap, err := snapcodec.Unmarshal(data)
		348	+	require.NoError(t, err)
		349	+	written[snap.Name] = true
		350	+	case <-deadline.C:
		351	+	t.Fatalf("final snapshot was dropped: wrote %v", written)
		352	+	}
		353	+	}
		354	+	require.True(t, written[final.name])
		355	+	})
		356	+	}
		357	+	}
		358	+
290	func TestSnapshotCompletionSignalsWaiters(t *testing.T) {	359		func TestSnapshotCompletionSignalsWaiters(t *testing.T) {
291	store := &channelSnapshotStore{writes: make(chan []byte, 1)}	360		store := &channelSnapshotStore{writes: make(chan []byte, 1)}
292	d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})	361		d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{})