merge(next): integrate history snapshot review fixes
7 files · +164 -43
M internal/usecase/copy/mode.go +2 -1
| @@ -53,7 +53,8 @@ func NewSnapshotFromRows(rows [][]renderer.Cell, width, height int) Snapshot { | |||
| 53 | 53 | // Len returns the number of document rows. | |
| 54 | 54 | func (s Snapshot) Len() int { return s.history.Len() + s.screen.Height } | |
| 55 | 55 | ||
| 56 | - | // Row returns document row i, or nil when i is out of range. | |
| 56 | + | // Row returns document row i, or nil when i is out of range. History rows | |
| 57 | + | // borrow immutable history storage and must be treated as read-only. | |
| 57 | 58 | func (s Snapshot) Row(i int) []renderer.Cell { | |
| 58 | 59 | if i < 0 { | |
| 59 | 60 | return nil | |
| @@ -53,7 +53,8 @@ func NewSnapshotFromRows(rows [][]renderer.Cell, width, height int) Snapshot { | |||||
| 53 | // Len returns the number of document rows. | 53 | // Len returns the number of document rows. | ||
| 54 | func (s Snapshot) Len() int { return s.history.Len() + s.screen.Height } | 54 | func (s Snapshot) Len() int { return s.history.Len() + s.screen.Height } | ||
| 55 | 55 | ||||
| 56 | - | // Row returns document row i, or nil when i is out of range. | 56 | + | // Row returns document row i, or nil when i is out of range. History rows |
| 57 | + | // borrow immutable history storage and must be treated as read-only. | |||
| 57 | func (s Snapshot) Row(i int) []renderer.Cell { | 58 | func (s Snapshot) Row(i int) []renderer.Cell { | ||
| 58 | if i < 0 { | 59 | if i < 0 { | ||
| 59 | return nil | 60 | return nil | ||
M internal/usecase/copy/mode_test.go +3 -3
| @@ -610,9 +610,9 @@ func BenchmarkNewSnapshot10KRows(b *testing.B) { | |||
| 610 | 610 | historyRows = 10_000 | |
| 611 | 611 | ) | |
| 612 | 612 | ||
| 613 | - | scrollback := newHistory(historyRows) | |
| 613 | + | history := newHistory(historyRows) | |
| 614 | 614 | for range historyRows { | |
| 615 | - | scrollback.Append(row(strings.Repeat("x", width))) | |
| 615 | + | history.Append(row(strings.Repeat("x", width))) | |
| 616 | 616 | } | |
| 617 | 617 | screen := renderer.NewFrame(width, height) | |
| 618 | 618 | for y := range height { | |
| @@ -622,7 +622,7 @@ func BenchmarkNewSnapshot10KRows(b *testing.B) { | |||
| 622 | 622 | b.ReportAllocs() | |
| 623 | 623 | b.ResetTimer() | |
| 624 | 624 | for b.Loop() { | |
| 625 | - | benchmarkSnapshotSink = NewSnapshot(scrollback, screen) | |
| 625 | + | benchmarkSnapshotSink = NewSnapshot(history, screen) | |
| 626 | 626 | } | |
| 627 | 627 | } | |
| 628 | 628 | ||
| @@ -610,9 +610,9 @@ func BenchmarkNewSnapshot10KRows(b *testing.B) { | |||||
| 610 | historyRows = 10_000 | 610 | historyRows = 10_000 | ||
| 611 | ) | 611 | ) | ||
| 612 | 612 | ||||
| 613 | - | scrollback := newHistory(historyRows) | 613 | + | history := newHistory(historyRows) |
| 614 | for range historyRows { | 614 | for range historyRows { | ||
| 615 | - | scrollback.Append(row(strings.Repeat("x", width))) | 615 | + | history.Append(row(strings.Repeat("x", width))) |
| 616 | } | 616 | } | ||
| 617 | screen := renderer.NewFrame(width, height) | 617 | screen := renderer.NewFrame(width, height) | ||
| 618 | for y := range height { | 618 | for y := range height { | ||
| @@ -622,7 +622,7 @@ func BenchmarkNewSnapshot10KRows(b *testing.B) { | |||||
| 622 | b.ReportAllocs() | 622 | b.ReportAllocs() | ||
| 623 | b.ResetTimer() | 623 | b.ResetTimer() | ||
| 624 | for b.Loop() { | 624 | for b.Loop() { | ||
| 625 | - | benchmarkSnapshotSink = NewSnapshot(scrollback, screen) | 625 | + | benchmarkSnapshotSink = NewSnapshot(history, screen) |
| 626 | } | 626 | } | ||
| 627 | } | 627 | } | ||
| 628 | 628 | ||||
M internal/usecase/daemon/daemon.go +10 -2
| @@ -111,6 +111,11 @@ type Daemon struct { | |||
| 111 | 111 | snapshotWorkerCancel context.CancelFunc | |
| 112 | 112 | snapshotWorkerDone chan struct{} | |
| 113 | 113 | snapshotWorkerFlush chan struct{} | |
| 114 | + | snapshotWorkerFinalWake chan struct{} | |
| 115 | + | // snapshotFinalJobs retains immutable terminal captures rejected by the | |
| 116 | + | // bounded regular queue. It is only populated during session teardown and | |
| 117 | + | // drained by the same worker, so persistence remains outside daemon locks. | |
| 118 | + | snapshotFinalJobs []*snapshotCapture | |
| 114 | 119 | snapshotWorkerClosing bool | |
| 115 | 120 | snapshotWorkerInFlight *snapshotCapture | |
| 116 | 121 | restoreDone chan struct{} | |
| @@ -375,7 +380,7 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||
| 375 | 380 | }) | |
| 376 | 381 | } | |
| 377 | 382 | if d.snapsEnabled { | |
| 378 | - | d.startSnapshotEncodeWorker(d.serveCtx) | |
| 383 | + | d.startSnapshotEncodeWorker() | |
| 379 | 384 | d.sessWg.Go(func() { | |
| 380 | 385 | d.snapshotSaver(d.serveCtx) | |
| 381 | 386 | }) | |
| @@ -424,8 +429,11 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||
| 424 | 429 | d.waitNotifies() | |
| 425 | 430 | d.hardCancel() | |
| 426 | 431 | d.serveCancel() | |
| 427 | - | d.stopSnapshotEncodeWorker() | |
| 432 | + | // Wait for handlers before flushing snapshots: a handler that already | |
| 433 | + | // entered killSession may still submit its terminal capture after the | |
| 434 | + | // registry snapshot has been removed. | |
| 428 | 435 | d.connWg.Wait() | |
| 436 | + | d.stopSnapshotEncodeWorker() | |
| 429 | 437 | d.shutdownAll(ports.ReasonServerShutdown) | |
| 430 | 438 | d.sessWg.Wait() | |
| 431 | 439 | d.waitNotifies() | |
| @@ -111,6 +111,11 @@ type Daemon struct { | |||||
| 111 | snapshotWorkerCancel context.CancelFunc | 111 | snapshotWorkerCancel context.CancelFunc | ||
| 112 | snapshotWorkerDone chan struct{} | 112 | snapshotWorkerDone chan struct{} | ||
| 113 | snapshotWorkerFlush chan struct{} | 113 | snapshotWorkerFlush chan struct{} | ||
| 114 | + | snapshotWorkerFinalWake chan struct{} | |||
| 115 | + | // snapshotFinalJobs retains immutable terminal captures rejected by the | |||
| 116 | + | // bounded regular queue. It is only populated during session teardown and | |||
| 117 | + | // drained by the same worker, so persistence remains outside daemon locks. | |||
| 118 | + | snapshotFinalJobs []*snapshotCapture | |||
| 114 | snapshotWorkerClosing bool | 119 | snapshotWorkerClosing bool | ||
| 115 | snapshotWorkerInFlight *snapshotCapture | 120 | snapshotWorkerInFlight *snapshotCapture | ||
| 116 | restoreDone chan struct{} | 121 | restoreDone chan struct{} | ||
| @@ -375,7 +380,7 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||||
| 375 | }) | 380 | }) | ||
| 376 | } | 381 | } | ||
| 377 | if d.snapsEnabled { | 382 | if d.snapsEnabled { | ||
| 378 | - | d.startSnapshotEncodeWorker(d.serveCtx) | 383 | + | d.startSnapshotEncodeWorker() |
| 379 | d.sessWg.Go(func() { | 384 | d.sessWg.Go(func() { | ||
| 380 | d.snapshotSaver(d.serveCtx) | 385 | d.snapshotSaver(d.serveCtx) | ||
| 381 | }) | 386 | }) | ||
| @@ -424,8 +429,11 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||||
| 424 | d.waitNotifies() | 429 | d.waitNotifies() | ||
| 425 | d.hardCancel() | 430 | d.hardCancel() | ||
| 426 | d.serveCancel() | 431 | d.serveCancel() | ||
| 427 | - | d.stopSnapshotEncodeWorker() | 432 | + | // Wait for handlers before flushing snapshots: a handler that already |
| 433 | + | // entered killSession may still submit its terminal capture after the | |||
| 434 | + | // registry snapshot has been removed. | |||
| 428 | d.connWg.Wait() | 435 | d.connWg.Wait() | ||
| 436 | + | d.stopSnapshotEncodeWorker() | |||
| 429 | d.shutdownAll(ports.ReasonServerShutdown) | 437 | d.shutdownAll(ports.ReasonServerShutdown) | ||
| 430 | d.sessWg.Wait() | 438 | d.sessWg.Wait() | ||
| 431 | d.waitNotifies() | 439 | d.waitNotifies() | ||
M internal/usecase/daemon/performance_benchmark_test.go +2 -7
| @@ -1,7 +1,6 @@ | |||
| 1 | 1 | package daemon | |
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | - | "context" | |
| 5 | 4 | "fmt" | |
| 6 | 5 | "io" | |
| 7 | 6 | "runtime" | |
| @@ -736,12 +735,8 @@ func newPerformanceFixture(t testing.TB, config performanceConfig) *performanceF | |||
| 736 | 735 | resizeSizes: [2]domain.Size{{Cols: 100, Rows: 30}, config.size}, | |
| 737 | 736 | } | |
| 738 | 737 | WithSnapshotStore(fixture.snaps)(d) | |
| 739 | - | snapshotCtx, cancelSnapshotWorker := context.WithCancel(context.Background()) | |
| 740 | - | d.startSnapshotEncodeWorker(snapshotCtx) | |
| 741 | - | t.Cleanup(func() { | |
| 742 | - | cancelSnapshotWorker() | |
| 743 | - | d.stopSnapshotEncodeWorker() | |
| 744 | - | }) | |
| 738 | + | d.startSnapshotEncodeWorker() | |
| 739 | + | t.Cleanup(d.stopSnapshotEncodeWorker) | |
| 745 | 740 | for tabIndex, tb := range sess.tabs { | |
| 746 | 741 | fixture.configureTab(tb, tabIndex, config.panes, config.historyRows, tabSize(config.size)) | |
| 747 | 742 | } | |
| @@ -1,7 +1,6 @@ | |||||
| 1 | package daemon | 1 | package daemon | ||
| 2 | 2 | ||||
| 3 | import ( | 3 | import ( | ||
| 4 | - | "context" | |||
| 5 | "fmt" | 4 | "fmt" | ||
| 6 | "io" | 5 | "io" | ||
| 7 | "runtime" | 6 | "runtime" | ||
| @@ -736,12 +735,8 @@ func newPerformanceFixture(t testing.TB, config performanceConfig) *performanceF | |||||
| 736 | resizeSizes: [2]domain.Size{{Cols: 100, Rows: 30}, config.size}, | 735 | resizeSizes: [2]domain.Size{{Cols: 100, Rows: 30}, config.size}, | ||
| 737 | } | 736 | } | ||
| 738 | WithSnapshotStore(fixture.snaps)(d) | 737 | WithSnapshotStore(fixture.snaps)(d) | ||
| 739 | - | snapshotCtx, cancelSnapshotWorker := context.WithCancel(context.Background()) | 738 | + | d.startSnapshotEncodeWorker() |
| 740 | - | d.startSnapshotEncodeWorker(snapshotCtx) | 739 | + | t.Cleanup(d.stopSnapshotEncodeWorker) |
| 741 | - | t.Cleanup(func() { | |||
| 742 | - | cancelSnapshotWorker() | |||
| 743 | - | d.stopSnapshotEncodeWorker() | |||
| 744 | - | }) | |||
| 745 | for tabIndex, tb := range sess.tabs { | 740 | for tabIndex, tb := range sess.tabs { | ||
| 746 | fixture.configureTab(tb, tabIndex, config.panes, config.historyRows, tabSize(config.size)) | 741 | fixture.configureTab(tb, tabIndex, config.panes, config.historyRows, tabSize(config.size)) | ||
| 747 | } | 742 | } | ||
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 +65 -4
| @@ -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. | |
| @@ -477,8 +504,8 @@ func (d *Daemon) finishSnapshotCapture(capture *snapshotCapture, succeeded bool) | |||
| 477 | 504 | }) | |
| 478 | 505 | } | |
| 479 | 506 | ||
| 480 | - | func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |
| 481 | - | if d == nil || ctx.Err() != nil { | |
| 507 | + | func (d *Daemon) startSnapshotEncodeWorker() { | |
| 508 | + | if d == nil { | |
| 482 | 509 | return | |
| 483 | 510 | } | |
| 484 | 511 | d.snapshotWorkerMu.Lock() | |
| @@ -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: | |
| @@ -315,6 +315,17 @@ func restorePTYSize(contentRect domain.Rect, tabSize domain.Size) domain.Size { | |||||
| 315 | // scheduleSnapshot captures a session only once while an immutable capture is | 315 | // scheduleSnapshot captures a session only once while an immutable capture is | ||
| 316 | // queued or being written. It never waits for the bounded worker queue. | 316 | // queued or being written. It never waits for the bounded worker queue. | ||
| 317 | func (d *Daemon) scheduleSnapshot(sess *session) bool { | 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 | if d == nil || sess == nil || !d.snapsEnabled || d.snaps == nil || !sess.snapEligible.Load() { | 329 | if d == nil || sess == nil || !d.snapsEnabled || d.snaps == nil || !sess.snapEligible.Load() { | ||
| 319 | return true | 330 | return true | ||
| 320 | } | 331 | } | ||
| @@ -333,7 +344,7 @@ func (d *Daemon) scheduleSnapshot(sess *session) bool { | |||||
| 333 | d.finishSnapshotCapture(capture, false) | 344 | d.finishSnapshotCapture(capture, false) | ||
| 334 | return false | 345 | return false | ||
| 335 | } | 346 | } | ||
| 336 | - | if d.enqueueSnapshotCapture(capture) { | 347 | + | if d.enqueueSnapshotCapture(capture) || (final && d.enqueueFinalSnapshotCapture(capture)) { |
| 337 | return true | 348 | return true | ||
| 338 | } | 349 | } | ||
| 339 | // Coalesce under saturation or shutdown: the latest state stays dirty and | 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 | // captureSession rotates history tails and clones visible frames while holding | 389 | // captureSession rotates history tails and clones visible frames while holding | ||
| 363 | // each pane lock. The returned capture contains only immutable state; encoding | 390 | // each pane lock. The returned capture contains only immutable state; encoding | ||
| 364 | // and persistence are deliberately deferred to snapshotEncodeWorker. | 391 | // and persistence are deliberately deferred to snapshotEncodeWorker. | ||
| @@ -477,8 +504,8 @@ func (d *Daemon) finishSnapshotCapture(capture *snapshotCapture, succeeded bool) | |||||
| 477 | }) | 504 | }) | ||
| 478 | } | 505 | } | ||
| 479 | 506 | ||||
| 480 | - | func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | 507 | + | func (d *Daemon) startSnapshotEncodeWorker() { |
| 481 | - | if d == nil || ctx.Err() != nil { | 508 | + | if d == nil { |
| 482 | return | 509 | return | ||
| 483 | } | 510 | } | ||
| 484 | d.snapshotWorkerMu.Lock() | 511 | d.snapshotWorkerMu.Lock() | ||
| @@ -495,10 +522,13 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||||
| 495 | d.snapshotWorkerCancel = cancel | 522 | d.snapshotWorkerCancel = cancel | ||
| 496 | d.snapshotWorkerDone = make(chan struct{}) | 523 | d.snapshotWorkerDone = make(chan struct{}) | ||
| 497 | d.snapshotWorkerFlush = make(chan struct{}) | 524 | d.snapshotWorkerFlush = make(chan struct{}) | ||
| 525 | + | d.snapshotWorkerFinalWake = make(chan struct{}, 1) | |||
| 526 | + | d.snapshotFinalJobs = nil | |||
| 498 | d.snapshotWorkerClosing = false | 527 | d.snapshotWorkerClosing = false | ||
| 499 | d.snapshotWorkerInFlight = nil | 528 | d.snapshotWorkerInFlight = nil | ||
| 500 | done := d.snapshotWorkerDone | 529 | done := d.snapshotWorkerDone | ||
| 501 | flush := d.snapshotWorkerFlush | 530 | flush := d.snapshotWorkerFlush | ||
| 531 | + | finalWake := d.snapshotWorkerFinalWake | |||
| 502 | d.snapshotWorkerMu.Unlock() | 532 | d.snapshotWorkerMu.Unlock() | ||
| 503 | go func() { | 533 | go func() { | ||
| 504 | defer close(done) | 534 | defer close(done) | ||
| @@ -518,6 +548,17 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||||
| 518 | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | 548 | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | ||
| 519 | return workerCtx.Err() == nil | 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 | flushQueued := func() { | 562 | flushQueued := func() { | ||
| 522 | for { | 563 | for { | ||
| 523 | select { | 564 | select { | ||
| @@ -526,6 +567,7 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||||
| 526 | return | 567 | return | ||
| 527 | } | 568 | } | ||
| 528 | default: | 569 | default: | ||
| 570 | + | _ = drainFinal() | |||
| 529 | return | 571 | return | ||
| 530 | } | 572 | } | ||
| 531 | } | 573 | } | ||
| @@ -537,6 +579,10 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||||
| 537 | case <-flush: | 579 | case <-flush: | ||
| 538 | flushQueued() | 580 | flushQueued() | ||
| 539 | return | 581 | return | ||
| 582 | + | case <-finalWake: | |||
| 583 | + | if !drainFinal() { | |||
| 584 | + | return | |||
| 585 | + | } | |||
| 540 | case capture := <-d.snapshotJobs: | 586 | case capture := <-d.snapshotJobs: | ||
| 541 | if !write(capture) { | 587 | if !write(capture) { | ||
| 542 | return | 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 | func (d *Daemon) setSnapshotWorkerInFlight(workerID uint64, capture *snapshotCapture) bool { | 607 | func (d *Daemon) setSnapshotWorkerInFlight(workerID uint64, capture *snapshotCapture) bool { | ||
| 550 | d.snapshotWorkerMu.Lock() | 608 | d.snapshotWorkerMu.Lock() | ||
| 551 | defer d.snapshotWorkerMu.Unlock() | 609 | defer d.snapshotWorkerMu.Unlock() | ||
| @@ -601,9 +659,12 @@ func (d *Daemon) finishStoppedSnapshotWorker(abandoned bool) { | |||||
| 601 | d.snapshotWorkerCancel = nil | 659 | d.snapshotWorkerCancel = nil | ||
| 602 | d.snapshotWorkerDone = nil | 660 | d.snapshotWorkerDone = nil | ||
| 603 | d.snapshotWorkerFlush = nil | 661 | d.snapshotWorkerFlush = nil | ||
| 662 | + | d.snapshotWorkerFinalWake = nil | |||
| 604 | d.snapshotWorkerClosing = false | 663 | d.snapshotWorkerClosing = false | ||
| 605 | d.snapshotWorkerInFlight = nil | 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 | for { | 668 | for { | ||
| 608 | select { | 669 | select { | ||
| 609 | case capture := <-d.snapshotJobs: | 670 | case capture := <-d.snapshotJobs: | ||
M internal/usecase/daemon/snapshot_test.go +81 -25
| @@ -25,12 +25,8 @@ func TestSnapshotEncodeWorkerDefersMarshalUntilAfterCapture(t *testing.T) { | |||
| 25 | 25 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |
| 26 | 26 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| 27 | 27 | WithSnapshotStore(store)(d) | |
| 28 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 29 | - | t.Cleanup(func() { | |
| 30 | - | cancel() | |
| 31 | - | d.stopSnapshotEncodeWorker() | |
| 32 | - | }) | |
| 33 | - | d.startSnapshotEncodeWorker(ctx) | |
| 28 | + | t.Cleanup(d.stopSnapshotEncodeWorker) | |
| 29 | + | d.startSnapshotEncodeWorker() | |
| 34 | 30 | ||
| 35 | 31 | entered := make(chan struct{}) | |
| 36 | 32 | release := make(chan struct{}) | |
| @@ -72,11 +68,9 @@ func exerciseBlockedSnapshotWorker(t *testing.T, blockEncode bool) { | |||
| 72 | 68 | store := newGatedSnapshotStore() | |
| 73 | 69 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| 74 | 70 | WithSnapshotStore(store)(d) | |
| 75 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 76 | - | d.startSnapshotEncodeWorker(ctx) | |
| 71 | + | d.startSnapshotEncodeWorker() | |
| 77 | 72 | t.Cleanup(func() { | |
| 78 | 73 | store.unblock() | |
| 79 | - | cancel() | |
| 80 | 74 | d.stopSnapshotEncodeWorker() | |
| 81 | 75 | }) | |
| 82 | 76 | if blockEncode { | |
| @@ -157,9 +151,7 @@ func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *tes | |||
| 157 | 151 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |
| 158 | 152 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |
| 159 | 153 | WithSnapshotStore(store)(d) | |
| 160 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 161 | - | defer cancel() | |
| 162 | - | d.startSnapshotEncodeWorker(ctx) | |
| 154 | + | d.startSnapshotEncodeWorker() | |
| 163 | 155 | ||
| 164 | 156 | entered := make(chan struct{}) | |
| 165 | 157 | release := make(chan struct{}) | |
| @@ -169,7 +161,8 @@ func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *tes | |||
| 169 | 161 | return snapcodec.Marshal(s) | |
| 170 | 162 | } | |
| 171 | 163 | sess := newSnapshotTestSession(t, "final", false, "/work") | |
| 172 | - | pty := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |
| 164 | + | pty, ok := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |
| 165 | + | require.True(t, ok) | |
| 173 | 166 | pty.EXPECT().Close().Return(nil).Once() | |
| 174 | 167 | sess.tabs[0].panes["pane-1"].screen.Write([]byte("\rnewest")) | |
| 175 | 168 | d.sessions[sess.id] = sess | |
| @@ -202,9 +195,7 @@ func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) | |||
| 202 | 195 | clock := newFinalFlushClock() | |
| 203 | 196 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |
| 204 | 197 | WithSnapshotStore(store)(d) | |
| 205 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 206 | - | defer cancel() | |
| 207 | - | d.startSnapshotEncodeWorker(ctx) | |
| 198 | + | d.startSnapshotEncodeWorker() | |
| 208 | 199 | ||
| 209 | 200 | sess := newSnapshotTestSession(t, "stuck-write", false, "/work") | |
| 210 | 201 | require.True(t, d.captureSession(sess)) | |
| @@ -240,10 +231,8 @@ func TestStoppedSnapshotWorkerLeavesCaptureRetryable(t *testing.T) { | |||
| 240 | 231 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |
| 241 | 232 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| 242 | 233 | WithSnapshotStore(store)(d) | |
| 243 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 244 | - | d.startSnapshotEncodeWorker(ctx) | |
| 234 | + | d.startSnapshotEncodeWorker() | |
| 245 | 235 | d.stopSnapshotEncodeWorker() | |
| 246 | - | t.Cleanup(cancel) | |
| 247 | 236 | ||
| 248 | 237 | sess := newSnapshotTestSession(t, "stopped", false, "/work") | |
| 249 | 238 | markSnapshotDirty(sess) | |
| @@ -287,6 +276,77 @@ func TestSnapshotQueueSaturationLeavesCaptureRetryable(t *testing.T) { | |||
| 287 | 276 | awaitSnapshotClean(t, third) | |
| 288 | 277 | } | |
| 289 | 278 | ||
| 279 | + | func TestNamedFinalSnapshotSurvivesSaturatedQueue(t *testing.T) { | |
| 280 | + | for _, tc := range []struct { | |
| 281 | + | name string | |
| 282 | + | stop func(t *testing.T, d *Daemon, sess *session) | |
| 283 | + | }{ | |
| 284 | + | { | |
| 285 | + | name: "kill session", | |
| 286 | + | stop: func(t *testing.T, d *Daemon, sess *session) { | |
| 287 | + | t.Helper() | |
| 288 | + | require.NoError(t, d.killSession(sess, ports.ReasonSessionKilled, false)) | |
| 289 | + | }, | |
| 290 | + | }, | |
| 291 | + | { | |
| 292 | + | name: "shutdown all", | |
| 293 | + | stop: func(t *testing.T, d *Daemon, _ *session) { | |
| 294 | + | t.Helper() | |
| 295 | + | d.shutdownAll(ports.ReasonServerShutdown) | |
| 296 | + | }, | |
| 297 | + | }, | |
| 298 | + | } { | |
| 299 | + | t.Run(tc.name, func(t *testing.T) { | |
| 300 | + | store := &channelSnapshotStore{writes: make(chan []byte, 3)} | |
| 301 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| 302 | + | WithSnapshotStore(store)(d) | |
| 303 | + | startSnapshotEncodeWorker(t, d) | |
| 304 | + | ||
| 305 | + | entered := make(chan struct{}) | |
| 306 | + | release := make(chan struct{}) | |
| 307 | + | var calls atomic.Int32 | |
| 308 | + | d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) { | |
| 309 | + | if calls.Add(1) == 1 { | |
| 310 | + | close(entered) | |
| 311 | + | <-release | |
| 312 | + | } | |
| 313 | + | return snapcodec.Marshal(s) | |
| 314 | + | } | |
| 315 | + | first := newSnapshotTestSession(t, "first", false, "/work") | |
| 316 | + | second := newSnapshotTestSession(t, "second", false, "/work") | |
| 317 | + | final := newSnapshotTestSession(t, "final", false, "/work") | |
| 318 | + | finalPTY, ok := final.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |
| 319 | + | require.True(t, ok) | |
| 320 | + | finalPTY.EXPECT().Close().Return(nil).Once() | |
| 321 | + | markSnapshotDirty(first) | |
| 322 | + | markSnapshotDirty(second) | |
| 323 | + | require.True(t, d.scheduleSnapshot(first)) | |
| 324 | + | <-entered | |
| 325 | + | require.True(t, d.scheduleSnapshot(second)) | |
| 326 | + | d.sessions[final.id] = final | |
| 327 | + | ||
| 328 | + | tc.stop(t, d, final) | |
| 329 | + | require.NotContains(t, d.sessions, final.id, "final capture must survive session removal") | |
| 330 | + | close(release) | |
| 331 | + | ||
| 332 | + | written := make(map[string]bool, 3) | |
| 333 | + | deadline := time.NewTimer(time.Second) | |
| 334 | + | defer deadline.Stop() | |
| 335 | + | for len(written) < 3 { | |
| 336 | + | select { | |
| 337 | + | case data := <-store.writes: | |
| 338 | + | snap, err := snapcodec.Unmarshal(data) | |
| 339 | + | require.NoError(t, err) | |
| 340 | + | written[snap.Name] = true | |
| 341 | + | case <-deadline.C: | |
| 342 | + | t.Fatalf("final snapshot was dropped: wrote %v", written) | |
| 343 | + | } | |
| 344 | + | } | |
| 345 | + | require.True(t, written[final.name]) | |
| 346 | + | }) | |
| 347 | + | } | |
| 348 | + | } | |
| 349 | + | ||
| 290 | 350 | func TestSnapshotCompletionSignalsWaiters(t *testing.T) { | |
| 291 | 351 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |
| 292 | 352 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| @@ -846,12 +906,8 @@ func awaitSnapshotClean(t *testing.T, sess *session) { | |||
| 846 | 906 | ||
| 847 | 907 | func startSnapshotEncodeWorker(t *testing.T, d *Daemon) { | |
| 848 | 908 | t.Helper() | |
| 849 | - | ctx, cancel := context.WithCancel(t.Context()) | |
| 850 | - | d.startSnapshotEncodeWorker(ctx) | |
| 851 | - | t.Cleanup(func() { | |
| 852 | - | cancel() | |
| 853 | - | d.stopSnapshotEncodeWorker() | |
| 854 | - | }) | |
| 909 | + | d.startSnapshotEncodeWorker() | |
| 910 | + | t.Cleanup(d.stopSnapshotEncodeWorker) | |
| 855 | 911 | } | |
| 856 | 912 | ||
| 857 | 913 | type gatedSnapshotStore struct { | |
| @@ -25,12 +25,8 @@ func TestSnapshotEncodeWorkerDefersMarshalUntilAfterCapture(t *testing.T) { | |||||
| 25 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | 25 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | ||
| 26 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | 26 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | ||
| 27 | WithSnapshotStore(store)(d) | 27 | WithSnapshotStore(store)(d) | ||
| 28 | - | ctx, cancel := context.WithCancel(t.Context()) | 28 | + | t.Cleanup(d.stopSnapshotEncodeWorker) |
| 29 | - | t.Cleanup(func() { | 29 | + | d.startSnapshotEncodeWorker() |
| 30 | - | cancel() | |||
| 31 | - | d.stopSnapshotEncodeWorker() | |||
| 32 | - | }) | |||
| 33 | - | d.startSnapshotEncodeWorker(ctx) | |||
| 34 | 30 | ||||
| 35 | entered := make(chan struct{}) | 31 | entered := make(chan struct{}) | ||
| 36 | release := make(chan struct{}) | 32 | release := make(chan struct{}) | ||
| @@ -72,11 +68,9 @@ func exerciseBlockedSnapshotWorker(t *testing.T, blockEncode bool) { | |||||
| 72 | store := newGatedSnapshotStore() | 68 | store := newGatedSnapshotStore() | ||
| 73 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | 69 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | ||
| 74 | WithSnapshotStore(store)(d) | 70 | WithSnapshotStore(store)(d) | ||
| 75 | - | ctx, cancel := context.WithCancel(t.Context()) | 71 | + | d.startSnapshotEncodeWorker() |
| 76 | - | d.startSnapshotEncodeWorker(ctx) | |||
| 77 | t.Cleanup(func() { | 72 | t.Cleanup(func() { | ||
| 78 | store.unblock() | 73 | store.unblock() | ||
| 79 | - | cancel() | |||
| 80 | d.stopSnapshotEncodeWorker() | 74 | d.stopSnapshotEncodeWorker() | ||
| 81 | }) | 75 | }) | ||
| 82 | if blockEncode { | 76 | if blockEncode { | ||
| @@ -157,9 +151,7 @@ func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *tes | |||||
| 157 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | 151 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | ||
| 158 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | 152 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | ||
| 159 | WithSnapshotStore(store)(d) | 153 | WithSnapshotStore(store)(d) | ||
| 160 | - | ctx, cancel := context.WithCancel(t.Context()) | 154 | + | d.startSnapshotEncodeWorker() |
| 161 | - | defer cancel() | |||
| 162 | - | d.startSnapshotEncodeWorker(ctx) | |||
| 163 | 155 | ||||
| 164 | entered := make(chan struct{}) | 156 | entered := make(chan struct{}) | ||
| 165 | release := make(chan struct{}) | 157 | release := make(chan struct{}) | ||
| @@ -169,7 +161,8 @@ func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *tes | |||||
| 169 | return snapcodec.Marshal(s) | 161 | return snapcodec.Marshal(s) | ||
| 170 | } | 162 | } | ||
| 171 | sess := newSnapshotTestSession(t, "final", false, "/work") | 163 | sess := newSnapshotTestSession(t, "final", false, "/work") | ||
| 172 | - | pty := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | 164 | + | pty, ok := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) |
| 165 | + | require.True(t, ok) | |||
| 173 | pty.EXPECT().Close().Return(nil).Once() | 166 | pty.EXPECT().Close().Return(nil).Once() | ||
| 174 | sess.tabs[0].panes["pane-1"].screen.Write([]byte("\rnewest")) | 167 | sess.tabs[0].panes["pane-1"].screen.Write([]byte("\rnewest")) | ||
| 175 | d.sessions[sess.id] = sess | 168 | d.sessions[sess.id] = sess | ||
| @@ -202,9 +195,7 @@ func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) | |||||
| 202 | clock := newFinalFlushClock() | 195 | clock := newFinalFlushClock() | ||
| 203 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | 196 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | ||
| 204 | WithSnapshotStore(store)(d) | 197 | WithSnapshotStore(store)(d) | ||
| 205 | - | ctx, cancel := context.WithCancel(t.Context()) | 198 | + | d.startSnapshotEncodeWorker() |
| 206 | - | defer cancel() | |||
| 207 | - | d.startSnapshotEncodeWorker(ctx) | |||
| 208 | 199 | ||||
| 209 | sess := newSnapshotTestSession(t, "stuck-write", false, "/work") | 200 | sess := newSnapshotTestSession(t, "stuck-write", false, "/work") | ||
| 210 | require.True(t, d.captureSession(sess)) | 201 | require.True(t, d.captureSession(sess)) | ||
| @@ -240,10 +231,8 @@ func TestStoppedSnapshotWorkerLeavesCaptureRetryable(t *testing.T) { | |||||
| 240 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | 231 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | ||
| 241 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | 232 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | ||
| 242 | WithSnapshotStore(store)(d) | 233 | WithSnapshotStore(store)(d) | ||
| 243 | - | ctx, cancel := context.WithCancel(t.Context()) | 234 | + | d.startSnapshotEncodeWorker() |
| 244 | - | d.startSnapshotEncodeWorker(ctx) | |||
| 245 | d.stopSnapshotEncodeWorker() | 235 | d.stopSnapshotEncodeWorker() | ||
| 246 | - | t.Cleanup(cancel) | |||
| 247 | 236 | ||||
| 248 | sess := newSnapshotTestSession(t, "stopped", false, "/work") | 237 | sess := newSnapshotTestSession(t, "stopped", false, "/work") | ||
| 249 | markSnapshotDirty(sess) | 238 | markSnapshotDirty(sess) | ||
| @@ -287,6 +276,77 @@ func TestSnapshotQueueSaturationLeavesCaptureRetryable(t *testing.T) { | |||||
| 287 | awaitSnapshotClean(t, third) | 276 | awaitSnapshotClean(t, third) | ||
| 288 | } | 277 | } | ||
| 289 | 278 | ||||
| 279 | + | func TestNamedFinalSnapshotSurvivesSaturatedQueue(t *testing.T) { | |||
| 280 | + | for _, tc := range []struct { | |||
| 281 | + | name string | |||
| 282 | + | stop func(t *testing.T, d *Daemon, sess *session) | |||
| 283 | + | }{ | |||
| 284 | + | { | |||
| 285 | + | name: "kill session", | |||
| 286 | + | stop: func(t *testing.T, d *Daemon, sess *session) { | |||
| 287 | + | t.Helper() | |||
| 288 | + | require.NoError(t, d.killSession(sess, ports.ReasonSessionKilled, false)) | |||
| 289 | + | }, | |||
| 290 | + | }, | |||
| 291 | + | { | |||
| 292 | + | name: "shutdown all", | |||
| 293 | + | stop: func(t *testing.T, d *Daemon, _ *session) { | |||
| 294 | + | t.Helper() | |||
| 295 | + | d.shutdownAll(ports.ReasonServerShutdown) | |||
| 296 | + | }, | |||
| 297 | + | }, | |||
| 298 | + | } { | |||
| 299 | + | t.Run(tc.name, func(t *testing.T) { | |||
| 300 | + | store := &channelSnapshotStore{writes: make(chan []byte, 3)} | |||
| 301 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |||
| 302 | + | WithSnapshotStore(store)(d) | |||
| 303 | + | startSnapshotEncodeWorker(t, d) | |||
| 304 | + | ||||
| 305 | + | entered := make(chan struct{}) | |||
| 306 | + | release := make(chan struct{}) | |||
| 307 | + | var calls atomic.Int32 | |||
| 308 | + | d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) { | |||
| 309 | + | if calls.Add(1) == 1 { | |||
| 310 | + | close(entered) | |||
| 311 | + | <-release | |||
| 312 | + | } | |||
| 313 | + | return snapcodec.Marshal(s) | |||
| 314 | + | } | |||
| 315 | + | first := newSnapshotTestSession(t, "first", false, "/work") | |||
| 316 | + | second := newSnapshotTestSession(t, "second", false, "/work") | |||
| 317 | + | final := newSnapshotTestSession(t, "final", false, "/work") | |||
| 318 | + | finalPTY, ok := final.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |||
| 319 | + | require.True(t, ok) | |||
| 320 | + | finalPTY.EXPECT().Close().Return(nil).Once() | |||
| 321 | + | markSnapshotDirty(first) | |||
| 322 | + | markSnapshotDirty(second) | |||
| 323 | + | require.True(t, d.scheduleSnapshot(first)) | |||
| 324 | + | <-entered | |||
| 325 | + | require.True(t, d.scheduleSnapshot(second)) | |||
| 326 | + | d.sessions[final.id] = final | |||
| 327 | + | ||||
| 328 | + | tc.stop(t, d, final) | |||
| 329 | + | require.NotContains(t, d.sessions, final.id, "final capture must survive session removal") | |||
| 330 | + | close(release) | |||
| 331 | + | ||||
| 332 | + | written := make(map[string]bool, 3) | |||
| 333 | + | deadline := time.NewTimer(time.Second) | |||
| 334 | + | defer deadline.Stop() | |||
| 335 | + | for len(written) < 3 { | |||
| 336 | + | select { | |||
| 337 | + | case data := <-store.writes: | |||
| 338 | + | snap, err := snapcodec.Unmarshal(data) | |||
| 339 | + | require.NoError(t, err) | |||
| 340 | + | written[snap.Name] = true | |||
| 341 | + | case <-deadline.C: | |||
| 342 | + | t.Fatalf("final snapshot was dropped: wrote %v", written) | |||
| 343 | + | } | |||
| 344 | + | } | |||
| 345 | + | require.True(t, written[final.name]) | |||
| 346 | + | }) | |||
| 347 | + | } | |||
| 348 | + | } | |||
| 349 | + | ||||
| 290 | func TestSnapshotCompletionSignalsWaiters(t *testing.T) { | 350 | func TestSnapshotCompletionSignalsWaiters(t *testing.T) { | ||
| 291 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | 351 | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | ||
| 292 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | 352 | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | ||
| @@ -846,12 +906,8 @@ func awaitSnapshotClean(t *testing.T, sess *session) { | |||||
| 846 | 906 | ||||
| 847 | func startSnapshotEncodeWorker(t *testing.T, d *Daemon) { | 907 | func startSnapshotEncodeWorker(t *testing.T, d *Daemon) { | ||
| 848 | t.Helper() | 908 | t.Helper() | ||
| 849 | - | ctx, cancel := context.WithCancel(t.Context()) | 909 | + | d.startSnapshotEncodeWorker() |
| 850 | - | d.startSnapshotEncodeWorker(ctx) | 910 | + | t.Cleanup(d.stopSnapshotEncodeWorker) |
| 851 | - | t.Cleanup(func() { | |||
| 852 | - | cancel() | |||
| 853 | - | d.stopSnapshotEncodeWorker() | |||
| 854 | - | }) | |||
| 855 | } | 911 | } | ||
| 856 | 912 | ||||
| 857 | type gatedSnapshotStore struct { | 913 | type gatedSnapshotStore struct { | ||