GODRIVER-1663 setting cancelConext to nil to avoid context being pinned (#436)

Author: KritiRav

x/mongo/driver/topology/connection.go

@@ -51,6 +51,7 @@ type connection struct {
 	connectContextMade   chan struct{}
 	canStream            bool
 	currentlyStreaming   bool
+	connectContextMutex  sync.Mutex
 
 	// pool related fields
 	pool       *pool
@@ -108,7 +109,23 @@ func (c *connection) connect(ctx context.Context) {
 	}
 	defer close(c.connectDone)
 
+	c.connectContextMutex.Lock()
 	ctx, c.cancelConnectContext = context.WithCancel(ctx)
+	c.connectContextMutex.Unlock()
+
+	defer func() {
+		var cancelFn context.CancelFunc
+
+		c.connectContextMutex.Lock()
+		cancelFn = c.cancelConnectContext
+		c.cancelConnectContext = nil
+		c.connectContextMutex.Unlock()
+
+		if cancelFn != nil {
+			cancelFn()
+		}
+	}()
+
 	close(c.connectContextMade)
 
 	// Assign the result of DialContext to a temporary net.Conn to ensure that c.nc is not set in an error case.
@@ -195,7 +212,16 @@ func (c *connection) wait() error {
 
 func (c *connection) closeConnectContext() {
 	<-c.connectContextMade
-	c.cancelConnectContext()
+	var cancelFn context.CancelFunc
+
+	c.connectContextMutex.Lock()
+	cancelFn = c.cancelConnectContext
+	c.cancelConnectContext = nil
+	c.connectContextMutex.Unlock()
+
+	if cancelFn != nil {
+		cancelFn()
+	}
 }
 
 func transformNetworkError(originalError error, contextDeadlineUsed bool) error {

x/mongo/driver/topology/connection_test.go

@@ -110,6 +110,22 @@ func TestConnection(t *testing.T) {
 				assert.NotNil(t, err, "expected connect error %v, got nil", want)
 				assert.Equal(t, want, got, "expected error %v, got %v", want, got)
 			})
+			t.Run("cancelConnectContext is nil after connect", func(t *testing.T) {
+				conn, err := newConnection(context.Background(), address.Address(""))
+				assert.Nil(t, err, "newConnection shouldn't error. got %v; want nil", err)
+				var wg sync.WaitGroup
+				wg.Add(1)
+
+				go func() {
+					defer wg.Done()
+					conn.connect(context.Background())
+					assert.Nil(t, conn.cancelConnectContext, "expected nil, got context.CancelFunc")
+				}()
+
+				conn.closeConnectContext()
+				assert.Nil(t, conn.cancelConnectContext, "expected nil, got context.CancelFunc")
+				wg.Wait()
+			})
 		})
 		t.Run("writeWireMessage", func(t *testing.T) {
 			t.Run("closed connection", func(t *testing.T) {