From 14f15c33fed43a90d21a6e4e7c1958f11b0928ef Mon Sep 17 00:00:00 2001
From: Brandon Philips <brandon@ifup.co>
Date: Thu, 6 Feb 2014 17:25:41 -0800
Subject: [PATCH] fix(tests/server_utils): use a WaitGroup for RunServer

Make sure that all of the servers exit before we return by using a
WaitGroup in each handler. This was used to help track down the issue
with the TestDiscoverySecondPeerUp test and the hung go-etcd watcher.
---
 tests/server_utils.go | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/tests/server_utils.go b/tests/server_utils.go
index 8c1f360f7..63dc9a3f0 100644
--- a/tests/server_utils.go
+++ b/tests/server_utils.go
@@ -4,6 +4,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/coreos/etcd/third_party/github.com/coreos/raft"
@@ -69,19 +70,23 @@ func RunServer(f func(*server.Server)) {
 
 	ps.SetServer(s)
 
+	w := &sync.WaitGroup{}
+
 	// Start up peer server.
 	c := make(chan bool)
 	go func() {
 		c <- true
 		ps.Start(false, []string{})
-		http.Serve(psListener, ps.HTTPHandler())
+		h := waitHandler{w, ps.HTTPHandler()}
+		http.Serve(psListener, &h)
 	}()
 	<-c
 
 	// Start up etcd server.
 	go func() {
 		c <- true
-		http.Serve(sListener, s.HTTPHandler())
+		h := waitHandler{w, s.HTTPHandler()}
+		http.Serve(sListener, &h)
 	}()
 	<-c
 
@@ -95,4 +100,20 @@ func RunServer(f func(*server.Server)) {
 	ps.Stop()
 	psListener.Close()
 	sListener.Close()
+	w.Wait()
+}
+
+type waitHandler struct {
+        wg *sync.WaitGroup
+        handler http.Handler
+}
+
+func (h *waitHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
+        h.wg.Add(1)
+        defer h.wg.Done()
+        h.handler.ServeHTTP(w, r)
+
+        //important to flush before decrementing the wait group.
+        //we won't get a chance to once main() ends.
+        w.(http.Flusher).Flush()
 }