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

@@ -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()
 }