GODRIVER-1411 fix data race between connection.connect and pool.disconnect (#284)

Author: iwysiu

x/mongo/driver/topology/connection.go

@@ -30,22 +30,24 @@ var globalConnectionID uint64 = 1
 func nextConnectionID() uint64 { return atomic.AddUint64(&globalConnectionID, 1) }
 
 type connection struct {
-	id               string
-	nc               net.Conn // When nil, the connection is closed.
-	addr             address.Address
-	idleTimeout      time.Duration
-	idleDeadline     atomic.Value // Stores a time.Time
-	lifetimeDeadline time.Time
-	readTimeout      time.Duration
-	writeTimeout     time.Duration
-	desc             description.Server
-	compressor       wiremessage.CompressorID
-	zliblevel        int
-	zstdLevel        int
-	connected        int32 // must be accessed using the sync/atomic package
-	connectDone      chan struct{}
-	connectErr       error
-	config           *connectionConfig
+	id                   string
+	nc                   net.Conn // When nil, the connection is closed.
+	addr                 address.Address
+	idleTimeout          time.Duration
+	idleDeadline         atomic.Value // Stores a time.Time
+	lifetimeDeadline     time.Time
+	readTimeout          time.Duration
+	writeTimeout         time.Duration
+	desc                 description.Server
+	compressor           wiremessage.CompressorID
+	zliblevel            int
+	zstdLevel            int
+	connected            int32 // must be accessed using the sync/atomic package
+	connectDone          chan struct{}
+	connectErr           error
+	config               *connectionConfig
+	cancelConnectContext context.CancelFunc
+	connectContextMade   chan struct{}
 
 	// pool related fields
 	pool       *pool
@@ -68,14 +70,15 @@ func newConnection(ctx context.Context, addr address.Address, opts ...Connection
 	id := fmt.Sprintf("%s[-%d]", addr, nextConnectionID())
 
 	c := &connection{
-		id:               id,
-		addr:             addr,
-		idleTimeout:      cfg.idleTimeout,
-		lifetimeDeadline: lifetimeDeadline,
-		readTimeout:      cfg.readTimeout,
-		writeTimeout:     cfg.writeTimeout,
-		connectDone:      make(chan struct{}),
-		config:           cfg,
+		id:                 id,
+		addr:               addr,
+		idleTimeout:        cfg.idleTimeout,
+		lifetimeDeadline:   lifetimeDeadline,
+		readTimeout:        cfg.readTimeout,
+		writeTimeout:       cfg.writeTimeout,
+		connectDone:        make(chan struct{}),
+		config:             cfg,
+		connectContextMade: make(chan struct{}),
 	}
 	atomic.StoreInt32(&c.connected, initialized)
 
@@ -85,12 +88,14 @@ func newConnection(ctx context.Context, addr address.Address, opts ...Connection
 // connect handles the I/O for a connection. It will dial, configure TLS, and perform
 // initialization handshakes.
 func (c *connection) connect(ctx context.Context) {
-
 	if !atomic.CompareAndSwapInt32(&c.connected, initialized, connected) {
 		return
 	}
 	defer close(c.connectDone)
 
+	ctx, c.cancelConnectContext = context.WithCancel(ctx)
+	close(c.connectContextMade)
+
 	var err error
 	c.nc, err = c.config.dialer.DialContext(ctx, c.addr.Network(), c.addr.String())
 	if err != nil {
@@ -178,6 +183,11 @@ func (c *connection) wait() error {
 	return c.connectErr
 }
 
+func (c *connection) closeConnectContext() {
+	<-c.connectContextMade
+	c.cancelConnectContext()
+}
+
 func (c *connection) writeWireMessage(ctx context.Context, wm []byte) error {
 	var err error
 	if atomic.LoadInt32(&c.connected) != connected {

x/mongo/driver/topology/pool.go

@@ -117,8 +117,6 @@ func connectionCloseFunc(v interface{}) {
 	}
 
 	go func() {
-		// wait for connection to finish trying to connect
-		_ = c.wait()
 		_ = c.pool.closeConnection(c)
 	}()
 }
@@ -414,6 +412,11 @@ func (p *pool) closeConnection(c *connection) error {
 	delete(p.opened, c.poolID)
 	p.Unlock()
 
+	if atomic.LoadInt32(&c.connected) == connected {
+		c.closeConnectContext()
+		_ = c.wait() // Make sure that the connection has finished connecting
+	}
+
 	if !atomic.CompareAndSwapInt32(&c.connected, connected, disconnected) {
 		return nil // We're closing an already closed connection
 	}

x/mongo/driver/topology/pool_test.go

@@ -175,6 +175,51 @@ func TestPool(t *testing.T) {
 				t.Errorf("Should have set the connection state on return. got %d; want %d", state, disconnected)
 			}
 		})
+		t.Run("no race if connections are also connecting", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			addr := bootstrapConnections(t, 3, func(nc net.Conn) {
+				<-cleanup
+				_ = nc.Close()
+			})
+			d := newdialer(&net.Dialer{})
+			pc := poolConfig{
+				Address: address.Address(addr.String()),
+			}
+			p, err := newPool(pc, WithDialer(func(Dialer) Dialer { return d }))
+			noerr(t, err)
+			err = p.connect()
+			noerr(t, err)
+			getDone := make(chan struct{})
+			disconnectDone := make(chan struct{})
+			_, err = p.get(context.Background())
+			noerr(t, err)
+			getCtx, getCancel := context.WithCancel(context.Background())
+			defer getCancel()
+			go func() {
+				defer close(getDone)
+				for {
+					select {
+					case <-disconnectDone:
+						return
+					default:
+						_, _ = p.get(getCtx)
+						noerr(t, err)
+						time.Sleep(time.Microsecond)
+					}
+				}
+			}()
+			go func() {
+				_, err := p.get(getCtx)
+				noerr(t, err)
+				ctx, cancel := context.WithTimeout(context.Background(), 3*time.Microsecond)
+				defer cancel()
+				err = p.disconnect(ctx)
+				noerr(t, err)
+				close(disconnectDone)
+			}()
+			<-getDone
+			close(cleanup)
+		})
 	})
 	t.Run("connect", func(t *testing.T) {
 		t.Run("can reconnect a disconnected pool", func(t *testing.T) {