GODRIVER-2348 Use context listener to cancel HB check

x/mongo/driver/driver.go

@@ -24,20 +24,16 @@ import (
 	"go.mongodb.org/mongo-driver/x/mongo/driver/mnet"
 )
 
-// ServerSelectionTimeoutGetter returns a timeout that should be used to set a
-// deadline for server selection. This logic is not handleded internally by the
-// ServerSelector, as a resulting deadline may be applicable by follow-up
-// operations, such as checking out a connection.
-type ServerSelectionTimeoutGetter interface {
-	GetServerSelectionTimeout() time.Duration
-}
-
 // Deployment is implemented by types that can select a server from a deployment.
 type Deployment interface {
-	ServerSelectionTimeoutGetter
-
 	SelectServer(context.Context, description.ServerSelector) (Server, error)
 	Kind() description.TopologyKind
+
+	// GetServerSelectionTimeout returns a timeout that should be used to set a
+	// deadline for server selection. This logic is not handleded internally by
+	// the ServerSelector, as a resulting deadline may be applicable by follow-up
+	// operations such as checking out a connection.
+	GetServerSelectionTimeout() time.Duration
 }
 
 // Connector represents a type that can connect to a server.

x/mongo/driver/topology/connection.go

@@ -65,9 +65,10 @@ type connection struct {
 	connectContextMade   chan struct{}
 	canStream            bool
 	currentlyStreaming   bool
-	cancellationListener cancellationListener
+	cancellationListener contextListener
 	serverConnectionID   *int64 // the server's ID for this client's connection
 	cancelConnSig        chan struct{}
+	prevCanceled         atomic.Value
 
 	// pool related fields
 	pool *pool
@@ -89,7 +90,7 @@ func newConnection(addr address.Address, opts ...ConnectionOption) *connection {
 		connectDone:          make(chan struct{}),
 		config:               cfg,
 		connectContextMade:   make(chan struct{}),
-		cancellationListener: newCancellListener(),
+		cancellationListener: newContextDoneListener(),
 	}
 	// Connections to non-load balanced deployments should eagerly set the generation numbers so errors encountered
 	// at any point during connection establishment can be processed without the connection being considered stale.
@@ -527,6 +528,14 @@ func (c *connection) ServerConnectionID() *int64 {
 	return c.serverConnectionID
 }
 
+func (c *connection) previousCanceled() bool {
+	if val := c.prevCanceled.Load(); val != nil {
+		return val.(bool)
+	}
+
+	return false
+}
+
 // initConnection is an adapter used during connection initialization. It has the minimum
 // functionality necessary to implement the driver.Connection interface, which is required to pass a
 // *connection to a Handshaker.
@@ -830,47 +839,3 @@ func configureTLS(ctx context.Context,
 	}
 	return client, nil
 }
-
-// TODO: Naming?
-
-// cancellListener listens for context cancellation and notifies listeners via a
-// callback function.
-type cancellListener struct {
-	aborted bool
-	done    chan struct{}
-}
-
-// newCancellListener constructs a cancellListener.
-func newCancellListener() *cancellListener {
-	return &cancellListener{
-		done: make(chan struct{}),
-	}
-}
-
-// Listen blocks until the provided context is cancelled or listening is aborted
-// via the StopListening function. If this detects that the context has been
-// cancelled (i.e. errors.Is(ctx.Err(), context.Canceled), the provided callback is
-// called to abort in-progress work. Even if the context expires, this function
-// will block until StopListening is called.
-func (c *cancellListener) Listen(ctx context.Context, abortFn func()) {
-	c.aborted = false
-
-	select {
-	case <-ctx.Done():
-		if errors.Is(ctx.Err(), context.Canceled) {
-			c.aborted = true
-			abortFn()
-		}
-
-		<-c.done
-	case <-c.done:
-	}
-}
-
-// StopListening stops the in-progress Listen call. This blocks if there is no
-// in-progress Listen call. This function will return true if the provided abort
-// callback was called when listening for cancellation on the previous context.
-func (c *cancellListener) StopListening() bool {
-	c.done <- struct{}{}
-	return c.aborted
-}

x/mongo/driver/topology/connection_test.go

@@ -1048,7 +1048,7 @@ func (d *dialer) lenclosed() int {
 }
 
 type testCancellationListener struct {
-	listener         *cancellListener
+	listener         *contextDoneListener
 	numListen        int
 	numStopListening int
 	aborted          bool
@@ -1058,7 +1058,7 @@ type testCancellationListener struct {
 // returned by the StopListening method.
 func newTestCancellationListener(aborted bool) *testCancellationListener {
 	return &testCancellationListener{
-		listener: newCancellListener(),
+		listener: newContextDoneListener(),
 		aborted:  aborted,
 	}
 }

x/mongo/driver/topology/context_listener.go

@@ -0,0 +1,91 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package topology
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+)
+
+type contextListener interface {
+	Listen(context.Context, func())
+	StopListening() bool
+}
+
+// contextDoneListener listens for context-ending eventsand notifies listeners
+// via a callback function.
+type contextDoneListener struct {
+	aborted     atomic.Value
+	done        chan struct{}
+	blockOnDone bool
+}
+
+var _ contextListener = &contextDoneListener{}
+
+// newContextDoneListener constructs a contextDoneListener that will block
+// when a context is done until StopListening is called.
+func newContextDoneListener() *contextDoneListener {
+	return &contextDoneListener{
+		done:        make(chan struct{}),
+		blockOnDone: true,
+	}
+}
+
+// newNonBlockingContextDoneLIstener constructs a contextDoneListener that
+// will not block when a context is done. In this case there are two ways to
+// unblock the listener: a finished context or a call to StopListening.
+func newNonBlockingContextDoneListener() *contextDoneListener {
+	return &contextDoneListener{
+		done:        make(chan struct{}),
+		blockOnDone: false,
+	}
+}
+
+// Listen blocks until the provided context is cancelled or listening is aborted
+// via the StopListening function. If this detects that the context has been
+// cancelled (i.e. errors.Is(ctx.Err(), context.Canceled), the provided callback
+// is called to abort in-progress work. If blockOnDone is true, this function
+// will block until StopListening is called, even if the context expires.
+func (c *contextDoneListener) Listen(ctx context.Context, abortFn func()) {
+	c.aborted.Store(false)
+
+	select {
+	case <-ctx.Done():
+		if errors.Is(ctx.Err(), context.Canceled) {
+			c.aborted.Store(true)
+
+			abortFn()
+		}
+
+		if c.blockOnDone {
+			<-c.done
+		}
+	case <-c.done:
+	}
+}
+
+// StopListening stops the in-progress Listen call. If blockOnDone is true, then
+// this blocks if there is no in-progress Listen call. This function will return
+// true if the provided abort callback was called when listening for
+// cancellation on the previous context.
+func (c *contextDoneListener) StopListening() bool {
+	if c.blockOnDone {
+		c.done <- struct{}{}
+	} else {
+		select {
+		case c.done <- struct{}{}:
+		default:
+		}
+	}
+
+	if aborted := c.aborted.Load(); aborted != nil {
+		return aborted.(bool)
+	}
+
+	return false
+}