diff --git a/pkg/schedule/schedule.go b/pkg/schedule/schedule.go index e680f8554..e2bf2fef2 100644 --- a/pkg/schedule/schedule.go +++ b/pkg/schedule/schedule.go @@ -37,8 +37,8 @@ type Scheduler interface { // Finished returns the number of finished jobs Finished() int - // WaitFinish waits all pending jobs to finish. - WaitFinish() + // WaitFinish waits until at least n job are finished and all pending jobs are finished. + WaitFinish(n int) // Stop stops the scheduler. Stop() @@ -110,13 +110,9 @@ func (f *fifo) Finished() int { return f.finished } -func (f *fifo) WaitFinish() { +func (f *fifo) WaitFinish(n int) { f.finishCond.L.Lock() - finish := f.finished - f.finishCond.L.Unlock() - - f.finishCond.L.Lock() - for f.finished == finish || len(f.pendings) != 0 { + for f.finished < n || len(f.pendings) != 0 { f.finishCond.Wait() } f.finishCond.L.Unlock() diff --git a/pkg/schedule/schedule_test.go b/pkg/schedule/schedule_test.go index 9c61ec6e2..9d714ac99 100644 --- a/pkg/schedule/schedule_test.go +++ b/pkg/schedule/schedule_test.go @@ -43,7 +43,7 @@ func TestFIFOSchedule(t *testing.T) { s.Schedule(j) } - s.WaitFinish() + s.WaitFinish(100) if s.Scheduled() != 100 { t.Errorf("scheduled = %d, want %d", s.Scheduled(), 100) } diff --git a/storage/kvstore_test.go b/storage/kvstore_test.go index 41b3c076b..a6dce9a11 100644 --- a/storage/kvstore_test.go +++ b/storage/kvstore_test.go @@ -337,7 +337,7 @@ func TestStoreCompact(t *testing.T) { b.tx.rangeRespc <- rangeResp{[][]byte{key1, key2}, nil} s.Compact(3) - s.fifoSched.WaitFinish() + s.fifoSched.WaitFinish(1) if s.compactMainRev != 3 { t.Errorf("compact main rev = %d, want 3", s.compactMainRev)