diff --git a/replay/replay_test.go b/replay/replay_test.go index f83ddbf75b..26836473a9 100644 --- a/replay/replay_test.go +++ b/replay/replay_test.go @@ -571,7 +571,7 @@ func buildHeavyWorkload(t *testing.T) vfs.FS { require.NoError(t, b.Commit(pebble.NoSync)) require.NoError(t, d.Flush()) } - wc.Stop() + wc.WaitAndStop() defer d.Close() return destFS diff --git a/replay/workload_capture.go b/replay/workload_capture.go index 0acee9d0df..e01f6f8de7 100644 --- a/replay/workload_capture.go +++ b/replay/workload_capture.go @@ -383,6 +383,18 @@ func (w *WorkloadCollector) Start(destFS vfs.FS, destPath string) { go w.copyFiles() } +// WaitAndStop waits for all enqueued sstables to be copied over, and then +// calls Stop. Gracefully ensures that all sstables referenced in the collected +// manifest's latest version edit will exist in the copy directory. +func (w *WorkloadCollector) WaitAndStop() { + w.mu.Lock() + for w.mu.tablesEnqueued != w.mu.tablesCopied { + w.mu.copyCond.Wait() + } + w.mu.Unlock() + w.Stop() +} + // Stop stops collection of the workload. func (w *WorkloadCollector) Stop() { w.mu.Lock()