Merge pull request #1464 from gravitational/sasha/fixclose

fixes issue with race on closed listener

lib/utils/loadbalancer.go

@@ -33,11 +33,14 @@ func NewLoadBalancer(ctx context.Context, frontend NetAddr, backends ...NetAddr)
 	if ctx == nil {
 		return nil, trace.BadParameter("missing parameter context")
 	}
+	waitCtx, waitCancel := context.WithCancel(ctx)
 	return &LoadBalancer{
 		frontend:     frontend,
 		ctx:          ctx,
 		backends:     backends,
 		currentIndex: -1,
+		waitCtx:      waitCtx,
+		waitCancel:   waitCancel,
 		Entry: log.WithFields(log.Fields{
 			trace.Component: "loadbalancer",
 			trace.ComponentFields: log.Fields{
@@ -61,6 +64,8 @@ type LoadBalancer struct {
 	listener       net.Listener
 	listenerClosed bool
 	connections    map[NetAddr]map[int64]net.Conn
+	waitCtx        context.Context
+	waitCancel     context.CancelFunc
 }
 
 // trackeConnection adds connection to the connection tracker
@@ -174,6 +179,7 @@ func (l *LoadBalancer) Listen() error {
 
 // Serve starts accepting connections
 func (l *LoadBalancer) Serve() error {
+	defer l.waitCancel()
 	backoffTimer := time.NewTicker(5 * time.Second)
 	defer backoffTimer.Stop()
 	for {
@@ -200,6 +206,12 @@ func (l *LoadBalancer) forwardConnection(conn net.Conn) {
 	}
 }
 
+// this is to workaround issue https://github.com/golang/go/issues/10527
+// in tests
+func (l *LoadBalancer) Wait() {
+	<-l.waitCtx.Done()
+}
+
 func (l *LoadBalancer) forward(conn net.Conn) error {
 	defer conn.Close()
 

lib/utils/loadbalancer_test.go

@@ -167,9 +167,11 @@ func (s *LBSuite) TestClose(c *check.C) {
 	// second close works
 	lb.Close()
 
+	lb.Wait()
+
 	// requests are failing
 	out, err = roundtrip(frontend.String())
-	c.Assert(err, check.NotNil)
+	c.Assert(err, check.NotNil, check.Commentf("output: %v, err: %v", string(out), err))
 }
 
 func (s *LBSuite) TestDropConnections(c *check.C) {