Browse Source

sender: support elegant stop

Yicheng Qin 11 years ago
parent
commit
12dba7d413
2 changed files with 4 additions and 6 deletions
  1. 4 0
      etcdserver/sender.go
  2. 0 6
      etcdserver/sender_test.go

+ 4 - 0
etcdserver/sender.go

@@ -134,6 +134,7 @@ type sender struct {
 	fs  *stats.FollowerStats
 	q   chan []byte
 	mu  sync.RWMutex
+	wg  sync.WaitGroup
 }
 
 func newSender(tr http.RoundTripper, u string, cid types.ID, fs *stats.FollowerStats) *sender {
@@ -144,6 +145,7 @@ func newSender(tr http.RoundTripper, u string, cid types.ID, fs *stats.FollowerS
 		fs:  fs,
 		q:   make(chan []byte),
 	}
+	s.wg.Add(connPerSender)
 	for i := 0; i < connPerSender; i++ {
 		go s.handle()
 	}
@@ -162,9 +164,11 @@ func (s *sender) send(data []byte) error {
 
 func (s *sender) stop() {
 	close(s.q)
+	s.wg.Wait()
 }
 
 func (s *sender) handle() {
+	defer s.wg.Done()
 	for d := range s.q {
 		start := time.Now()
 		err := s.post(d)

+ 0 - 6
etcdserver/sender_test.go

@@ -102,9 +102,6 @@ func TestSenderSend(t *testing.T) {
 		t.Fatalf("unexpect send error: %v", err)
 	}
 	s.stop()
-	// wait for goroutines end
-	// TODO: elegant stop
-	time.Sleep(10 * time.Millisecond)
 
 	if tr.Request() == nil {
 		t.Errorf("sender fails to post the data")
@@ -155,9 +152,6 @@ func TestSenderSendFailed(t *testing.T) {
 		t.Fatalf("unexpect send error: %v", err)
 	}
 	s.stop()
-	// wait for goroutines end
-	// TODO: elegant stop
-	time.Sleep(10 * time.Millisecond)
 
 	fs.Lock()
 	defer fs.Unlock()