GODRIVER-1345 - Respect context.Context cancellation in Disconnect (#199)

Author: edaniels

x/mongo/driver/topology/server.go

@@ -89,9 +89,10 @@ type Server struct {
 	sem  *semaphore.Weighted
 
 	// goroutine management fields
-	done     chan struct{}
-	checkNow chan struct{}
-	closewg  sync.WaitGroup
+	done          chan struct{}
+	checkNow      chan struct{}
+	disconnecting chan struct{}
+	closewg       sync.WaitGroup
 
 	// description related fields
 	desc                   atomic.Value // holds a description.Server
@@ -139,8 +140,9 @@ func NewServer(addr address.Address, opts ...ServerOption) (*Server, error) {
 
 		sem: semaphore.NewWeighted(int64(maxConns)),
 
-		done:     make(chan struct{}),
-		checkNow: make(chan struct{}, 1),
+		done:          make(chan struct{}),
+		checkNow:      make(chan struct{}, 1),
+		disconnecting: make(chan struct{}),
 
 		subscribers: make(map[uint64]chan description.Server),
 	}
@@ -193,7 +195,14 @@ func (s *Server) Disconnect(ctx context.Context) error {
 
 	// For every call to Connect there must be at least 1 goroutine that is
 	// waiting on the done channel.
-	s.done <- struct{}{}
+	select {
+	case <-ctx.Done():
+		// signal a disconnect and still wait for receiver of done
+		// to finish.
+		close(s.disconnecting)
+		s.done <- struct{}{}
+	case s.done <- struct{}{}:
+	}
 	err := s.pool.disconnect(ctx)
 	if err != nil {
 		return err
@@ -398,6 +407,13 @@ func (s *Server) update() {
 		conn.nc.Close()
 	}
 	for {
+		select {
+		case <-done:
+			closeServer()
+			return
+		default:
+		}
+
 		select {
 		case <-heartbeatTicker.C:
 		case <-checkNow:
@@ -463,7 +479,15 @@ func (s *Server) heartbeat(conn *connection) (description.Server, *connection) {
 	var desc description.Server
 	var set bool
 	var err error
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-s.disconnecting:
+			cancel()
+		}
+	}()
 
 	for i := 1; i <= maxRetry; i++ {
 		var now time.Time
@@ -499,7 +523,7 @@ func (s *Server) heartbeat(conn *connection) (description.Server, *connection) {
 
 			conn.connect(ctx)
 
-			err := conn.wait()
+			err = conn.wait()
 			if err == nil {
 				descPtr = &conn.desc
 			}