From 870e4c2681376242ea3ec487501b944280dbfa09 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Thu, 11 Feb 2016 15:20:19 -0800
Subject: [PATCH] *: fix schedule.Wait race

---
 pkg/schedule/schedule.go      | 12 ++++--------
 pkg/schedule/schedule_test.go |  2 +-
 storage/kvstore_test.go       |  2 +-
 3 files changed, 6 insertions(+), 10 deletions(-)

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 b616545d7..157880b51 100644
--- a/storage/kvstore_test.go
+++ b/storage/kvstore_test.go
@@ -286,7 +286,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)