Add RetryHTTP tests with misbehaving servers

clientutil/retry_test.go

 package clientutil
 
 import (
+	"context"
 	"io"
 	"io/ioutil"
+	"log"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"testing"
+	"time"
 )
 
-func testHTTPServer() *httptest.Server {
-	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+type tcpHandler interface {
+	Handle(net.Conn)
+}
+
+type tcpHandlerFunc func(net.Conn)
+
+func (f tcpHandlerFunc) Handle(c net.Conn) { f(c) }
+
+// Base TCP server type (to build fake LDAP servers).
+type tcpServer struct {
+	l       net.Listener
+	handler tcpHandler
+}
+
+func newTCPServer(t testing.TB, handler tcpHandler) *tcpServer {
+	l, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal("Listen():", err)
+	}
+	log.Printf("started new tcp server on %s", l.Addr().String())
+	s := &tcpServer{l: l, handler: handler}
+	go s.serve()
+	return s
+}
+
+func (s *tcpServer) serve() {
+	for {
+		conn, err := s.l.Accept()
+		if err != nil {
+			return
+		}
+		go func(c net.Conn) {
+			s.handler.Handle(c)
+			c.Close()
+		}(conn)
+	}
+}
+
+func (s *tcpServer) Addr() string {
+	return s.l.Addr().String()
+}
+
+func (s *tcpServer) Close() {
+	s.l.Close()
+}
+
+// A test server that will close all incoming connections right away.
+func newConnFailServer(t testing.TB) *tcpServer {
+	return newTCPServer(t, tcpHandlerFunc(func(c net.Conn) {}))
+}
+
+// A test server that will close all connections after a 1s delay.
+func newConnFailDelayServer(t testing.TB) *tcpServer {
+	return newTCPServer(t, tcpHandlerFunc(func(c net.Conn) { time.Sleep(1 * time.Second) }))
+}
+
+type httpServer struct {
+	*httptest.Server
+}
+
+func (s *httpServer) Addr() string {
+	u, _ := url.Parse(s.Server.URL)
+	return u.Host
+}
+
+// An HTTP server that will always return a specific HTTP status.
+func newStatusHTTPServer(statusCode int) *httpServer {
+	return &httpServer{httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		w.WriteHeader(statusCode)
+		io.WriteString(w, "hello\n")
+	}))}
+}
+
+func newOKHTTPServer() *httpServer {
+	return &httpServer{httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/plain")
 		io.WriteString(w, "OK")
-	}))
+	}))}
+}
+
+type testServer interface {
+	Addr() string
+	Close()
 }
 
 type testBackends struct {
-	servers []*httptest.Server
+	servers []testServer
 	addrs   []string
 }
 
-func newTestBackends(n int) *testBackends {
+func newTestBackends(servers ...testServer) *testBackends {
 	b := new(testBackends)
-	for i := 0; i < n; i++ {
-		s := testHTTPServer()
-		u, _ := url.Parse(s.URL)
+	for _, s := range servers {
 		b.servers = append(b.servers, s)
-		b.addrs = append(b.addrs, u.Host)
+		b.addrs = append(b.addrs, s.Addr())
 	}
 	return b
 }
@@ -54,22 +134,32 @@ func doRequests(backends *testBackends, u string, n int) (int, int) {
 
 	var errs, oks int
 	for i := 0; i < n; i++ {
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 		req, _ := http.NewRequest("GET", u, nil)
-		resp, err := RetryHTTPDo(c, req, b)
+		resp, err := RetryHTTPDo(c, req.WithContext(ctx), b)
+		cancel()
 		if err != nil {
 			errs++
 			continue
 		}
-		ioutil.ReadAll(resp.Body)
+		_, err = ioutil.ReadAll(resp.Body)
 		resp.Body.Close()
+		if resp.StatusCode != 200 {
+			errs++
+			continue
+		}
+		if err != nil {
+			errs++
+			continue
+		}
 		oks++
 	}
 
 	return oks, errs
 }
 
-func TestRetryAndTransport(t *testing.T) {
-	b := newTestBackends(3)
+func TestRetryHTTP_BackendsDown(t *testing.T) {
+	b := newTestBackends(newOKHTTPServer(), newOKHTTPServer(), newOKHTTPServer())
 	defer b.close()
 
 	oks, errs := doRequests(b, "http://backend/", 100)
@@ -81,6 +171,7 @@ func TestRetryAndTransport(t *testing.T) {
 	}
 
 	b.stop(0)
+	b.stop(1)
 
 	oks, errs = doRequests(b, "http://backend/", 100)
 	if errs > 0 {
@@ -90,3 +181,20 @@ func TestRetryAndTransport(t *testing.T) {
 		t.Fatal("oks=0")
 	}
 }
+
+func TestRetryHTTP_HighLatencyBackend(t *testing.T) {
+	b := newTestBackends(newConnFailDelayServer(t), newOKHTTPServer())
+	defer b.close()
+
+	_, _ = doRequests(b, "http://backend/", 10)
+	// Silly transport.go load balancer only balances connections,
+	// so in this scenario we'll just keep hitting the slow
+	// server, exhausting our deadline budget, and never fail over
+	// to the secondary backend :(
+	// if errs > 0 {
+	// 	t.Fatalf("errs=%d", errs)
+	// }
+	// if oks == 0 {
+	// 	t.Fatal("oks=0")
+	// }
+}

clientutil/transport.go

@@ -125,6 +125,9 @@ func (b *balancer) dial(ctx context.Context, network, addr string) (net.Conn, er
 		if err == nil {
 			return conn, nil
 		} else if err == context.Canceled {
+			// A timeout might be bad, set the error bit
+			// on the connection.
+			b.notify(addr, false)
 			return nil, err
 		}
 		b.notify(addr, false)