GODRIVER-1298 fix panic from closing disconnected connection

Change-Id: I3ae9b0aca1c40c4a1ee3a672b249fabc4ec8b3ad

Author: iwysiu

x/mongo/driver/topology/connection.go

@@ -37,7 +37,7 @@ type connection struct {
 	nc               net.Conn // When nil, the connection is closed.
 	addr             address.Address
 	idleTimeout      time.Duration
-	idleDeadline     time.Time
+	idleDeadline     atomic.Value // Stores a time.Time
 	lifetimeDeadline time.Time
 	readTimeout      time.Duration
 	writeTimeout     time.Duration
@@ -87,6 +87,7 @@ 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
 	}
@@ -151,8 +152,10 @@ func (c *connection) connect(ctx context.Context) {
 	}
 }
 
-func (c *connection) connectWait() error {
-	<-c.connectDone
+func (c *connection) wait() error {
+	if c.connectDone != nil {
+		<-c.connectDone
+	}
 	return c.connectErr
 }
 
@@ -259,7 +262,11 @@ func (c *connection) close() error {
 		return nil
 	}
 	if c.pool == nil {
-		err := c.nc.Close()
+		var err error
+
+		if c.nc != nil {
+			err = c.nc.Close()
+		}
 		atomic.StoreInt32(&c.connected, disconnected)
 		return err
 	}
@@ -268,7 +275,8 @@ func (c *connection) close() error {
 
 func (c *connection) expired() bool {
 	now := time.Now()
-	if !c.idleDeadline.IsZero() && now.After(c.idleDeadline) {
+	idleDeadline, ok := c.idleDeadline.Load().(time.Time)
+	if ok && now.After(idleDeadline) {
 		return true
 	}
 
@@ -281,7 +289,7 @@ func (c *connection) expired() bool {
 
 func (c *connection) bumpIdleDeadline() {
 	if c.idleTimeout > 0 {
-		c.idleDeadline = time.Now().Add(c.idleTimeout)
+		c.idleDeadline.Store(time.Now().Add(c.idleTimeout))
 	}
 }
 

x/mongo/driver/topology/connection_test.go

@@ -57,7 +57,7 @@ func TestConnection(t *testing.T) {
 					t.Errorf("newConnection shouldn't error. got %v; want nil", got)
 				}
 				conn.connect(context.Background())
-				got = conn.connectWait()
+				got = conn.wait()
 				if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 					t.Errorf("errors do not match. got %v; want %v", got, want)
 				}
@@ -81,7 +81,7 @@ func TestConnection(t *testing.T) {
 					t.Errorf("newConnection shouldn't error. got %v; want nil", got)
 				}
 				conn.connect(context.Background())
-				got = conn.connectWait()
+				got = conn.wait()
 				if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 					t.Errorf("errors do not match. got %v; want %v", got, want)
 				}
@@ -105,7 +105,7 @@ func TestConnection(t *testing.T) {
 				)
 				noerr(t, err)
 				conn.connect(context.Background())
-				err = conn.connectWait()
+				err = conn.wait()
 				noerr(t, err)
 				if !cmp.Equal(got, want) {
 					t.Errorf("Server descriptions do not match. got %v; want %v", got, want)

x/mongo/driver/topology/pool.go

@@ -116,7 +116,11 @@ func connectionCloseFunc(v interface{}) {
 		return
 	}
 
-	go func() { _ = c.pool.closeConnection(c) }()
+	go func() {
+		// wait for connection to finish trying to connect
+		_ = c.wait()
+		_ = c.pool.closeConnection(c)
+	}()
 }
 
 // connectionInitFunc returns an init function for the resource pool that will make new connections for this pool
@@ -329,7 +333,7 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 			c.connect(ctx)
 		}
 
-		err := c.connectWait()
+		err := c.wait()
 		if err != nil {
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
@@ -377,7 +381,7 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 
 		c.connect(ctx)
 		// wait for conn to be connected
-		err = c.connectWait()
+		err = c.wait()
 		if err != nil {
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
@@ -413,10 +417,14 @@ func (p *pool) closeConnection(c *connection) error {
 	if !atomic.CompareAndSwapInt32(&c.connected, connected, disconnected) {
 		return nil // We're closing an already closed connection
 	}
-	err := c.nc.Close()
-	if err != nil {
-		return ConnectionError{ConnectionID: c.id, Wrapped: err, message: "failed to closeConnection net.Conn"}
+
+	if c.nc != nil {
+		err := c.nc.Close()
+		if err != nil {
+			return ConnectionError{ConnectionID: c.id, Wrapped: err, message: "failed to close net.Conn"}
+		}
 	}
+
 	return nil
 }
 

x/mongo/driver/topology/pool_test.go

@@ -3,12 +3,12 @@ package topology
 import (
 	"context"
 	"errors"
-	"go.mongodb.org/mongo-driver/x/mongo/driver/address"
 	"net"
-
 	"sync/atomic"
 	"testing"
 	"time"
+
+	"go.mongodb.org/mongo-driver/x/mongo/driver/address"
 )
 
 func TestPool(t *testing.T) {
@@ -531,5 +531,41 @@ func TestPool(t *testing.T) {
 				t.Errorf("Should not return connection to pool twice. got %d; want %d", p.conns.size, 1)
 			}
 		})
+		t.Run("close does not panic if expires before connected", func(t *testing.T) {
+			cleanup := make(chan struct{})
+			defer close(cleanup)
+			addr := bootstrapConnections(t, 3, func(nc net.Conn) {
+				<-cleanup
+				_ = nc.Close()
+			})
+			d := newSleepDialer(&net.Dialer{})
+			pc := poolConfig{
+				Address:     address.Address(addr.String()),
+				MinPoolSize: 1,
+			}
+			maintainInterval = time.Second
+			p, err := newPool(pc, WithDialer(func(Dialer) Dialer { return d }),
+				WithLifeTimeout(func(time.Duration) time.Duration { return 10 * time.Millisecond }),
+			)
+			maintainInterval = time.Minute
+			noerr(t, err)
+			err = p.connect()
+			noerr(t, err)
+			_, err = p.get(context.Background())
+			noerr(t, err)
+		})
 	})
 }
+
+type sleepDialer struct {
+	Dialer
+}
+
+func newSleepDialer(d Dialer) *sleepDialer {
+	return &sleepDialer{d}
+}
+
+func (d *sleepDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	time.Sleep(5 * time.Second)
+	return d.Dialer.DialContext(ctx, network, address)
+}

x/mongo/driver/topology/server.go

@@ -499,7 +499,7 @@ func (s *Server) heartbeat(conn *connection) (description.Server, *connection) {
 
 			conn.connect(ctx)
 
-			err := conn.connectWait()
+			err := conn.wait()
 			if err == nil {
 				descPtr = &conn.desc
 			}