GODRIVER-1540 fix deadlock in connection (#348)

Author: iwysiu

x/mongo/driver/topology/connection.go

@@ -290,16 +290,17 @@ func (c *connection) close() error {
 	if atomic.LoadInt32(&c.connected) != connected {
 		return nil
 	}
-	if c.pool == nil {
-		var err error
 
-		if c.nc != nil {
-			err = c.nc.Close()
-		}
-		atomic.StoreInt32(&c.connected, disconnected)
-		return err
+	var err error
+	if c.nc != nil {
+		err = c.nc.Close()
 	}
-	return c.pool.closeConnection(c)
+	atomic.StoreInt32(&c.connected, disconnected)
+
+	if c.pool != nil {
+		_ = c.pool.removeConnection(c)
+	}
+	return err
 }
 
 func (c *connection) expired() bool {

x/mongo/driver/topology/connection_test.go

@@ -492,6 +492,18 @@ func (nc *netconn) Close() error {
 	return nc.Conn.Close()
 }
 
+type writeFailConn struct {
+	net.Conn
+}
+
+func (wfc *writeFailConn) Write([]byte) (int, error) {
+	return 0, errors.New("Write error")
+}
+
+func (wfc *writeFailConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
 type dialer struct {
 	Dialer
 	opened        map[*netconn]struct{}

x/mongo/driver/topology/pool.go

@@ -431,6 +431,18 @@ func (p *pool) closeConnection(c *connection) error {
 	return nil
 }
 
+// removeConnection removes a connection from the pool.
+func (p *pool) removeConnection(c *connection) error {
+	if c.pool != p {
+		return ErrWrongPool
+	}
+	p.Lock()
+	delete(p.opened, c.poolID)
+	p.Unlock()
+
+	return nil
+}
+
 // put returns a connection to this pool. If the pool is connected, the connection is not
 // stale, and there is space in the cache, the connection is returned to the cache.
 func (p *pool) put(c *connection) error {

x/mongo/driver/topology/pool_test.go

@@ -4,11 +4,13 @@ import (
 	"context"
 	"errors"
 	"net"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
 
 	"go.mongodb.org/mongo-driver/x/mongo/driver/address"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/operation"
 )
 
 func TestPool(t *testing.T) {
@@ -518,6 +520,34 @@ func TestPool(t *testing.T) {
 			}
 			close(cleanup)
 		})
+		t.Run("handshaker i/o fails", func(t *testing.T) {
+			want := "unable to write wire message to network: Write error"
+
+			pc := poolConfig{
+				Address: address.Address(""),
+			}
+			p, err := newPool(pc, WithHandshaker(func(Handshaker) Handshaker {
+				return operation.NewIsMaster()
+			}),
+				WithDialer(func(Dialer) Dialer {
+					return DialerFunc(func(context.Context, string, string) (net.Conn, error) {
+						return &writeFailConn{&net.TCPConn{}}, nil
+					})
+				}),
+			)
+			noerr(t, err)
+			err = p.connect()
+			noerr(t, err)
+
+			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+			defer cancel()
+
+			_, err = p.get(ctx)
+			connErr := err.(ConnectionError)
+			if !strings.Contains(connErr.Error(), want) {
+				t.Errorf("Incorrect error. got %v; error should contain %v", connErr.Wrapped, want)
+			}
+		})
 	})
 	t.Run("Connection", func(t *testing.T) {
 		t.Run("Connection Close Does Not Error After Pool Is Disconnected", func(t *testing.T) {