Fix timeout in etcdwatcher tests

madhusudancs · Dec 22, 2015 · 9f35eeb · 9f35eeb
1 parent 1f26fa4
commit 9f35eeb
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/pkg/storage/etcd/etcd_watcher.go b/pkg/storage/etcd/etcd_watcher.go
@@ -150,8 +150,7 @@ func (w *etcdWatcher) etcdWatch(ctx context.Context, client etcd.KeysAPI, key st
 
 	// All calls to etcd are coming from this function - once it is finished
 	// no other call to etcd should be generated by this watcher.
-	w.wg.Add(1)
-	defer w.wg.Done()
+	done := func() {}
 
 	// We need to be prepared, that Stop() can be called at any time.
 	// It can potentially also be called, even before this function is called.
@@ -165,6 +164,8 @@ func (w *etcdWatcher) etcdWatch(ctx context.Context, client etcd.KeysAPI, key st
 			// Watcher has already been stopped - don't event initiate it here.
 			return true
 		}
+		w.wg.Add(1)
+		done = w.wg.Done
 		// Perform initialization of watcher under lock - we want to avoid situation when
 		// Stop() is called in the meantime (which in tests can cause etcd termination and
 		// strange behavior here).
@@ -185,6 +186,7 @@ func (w *etcdWatcher) etcdWatch(ctx context.Context, client etcd.KeysAPI, key st
 		w.ctx, w.cancel = context.WithCancel(ctx)
 		return false
 	}()
+	defer done()
 	if returned {
 		return
 	}
@@ -455,7 +457,6 @@ func (w *etcdWatcher) ResultChan() <-chan watch.Event {
 // Stop implements watch.Interface.
 func (w *etcdWatcher) Stop() {
 	w.stopLock.Lock()
-	defer w.stopLock.Unlock()
 	if w.cancel != nil {
 		w.cancel()
 		w.cancel = nil
@@ -464,6 +465,8 @@ func (w *etcdWatcher) Stop() {
 		w.stopped = true
 		close(w.userStop)
 	}
+	w.stopLock.Unlock()
+
 	// Wait until all calls to etcd are finished and no other
 	// will be issued.
 	w.wg.Wait()