Fix connection leak caused by rapid context cancellation (mysqljs#1024)

Fixes mysqljs#1023

Author: mherr-google

connector.go

@@ -64,6 +64,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	// Call startWatcher for context support (From Go 1.8)
 	mc.startWatcher()
 	if err := mc.watchCancel(ctx); err != nil {
+		mc.cleanup()
 		return nil, err
 	}
 	defer mc.finish()

driver_go110_test.go

@@ -135,3 +135,56 @@ func TestConnectorTimeoutsDuringOpen(t *testing.T) {
 		t.Fatalf("(*Connector).Connect should have timed out")
 	}
 }
+
+// A connection which can only be closed.
+type dummyConnection struct {
+	net.Conn
+	closed bool
+}
+
+func (d *dummyConnection) Close() error {
+	d.closed = true
+	return nil
+}
+
+func TestConnectorTimeoutsWatchCancel(t *testing.T) {
+	var (
+		cancel  func()           // Used to cancel the context just after connecting.
+		created *dummyConnection // The created connection.
+	)
+
+	RegisterDialContext("TestConnectorTimeoutsWatchCancel", func(ctx context.Context, addr string) (net.Conn, error) {
+		// Canceling at this time triggers the watchCancel error branch in Connect().
+		cancel()
+		created = &dummyConnection{}
+		return created, nil
+	})
+
+	mycnf := NewConfig()
+	mycnf.User = "root"
+	mycnf.Addr = "foo"
+	mycnf.Net = "TestConnectorTimeoutsWatchCancel"
+
+	conn, err := NewConnector(mycnf)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	db := sql.OpenDB(conn)
+	defer db.Close()
+
+	var ctx context.Context
+	ctx, cancel = context.WithCancel(context.Background())
+	defer cancel()
+
+	if _, err := db.Conn(ctx); err != context.Canceled {
+		t.Errorf("got %v, want context.Canceled", err)
+	}
+
+	if created == nil {
+		t.Fatal("no connection created")
+	}
+	if !created.closed {
+		t.Errorf("connection not closed")
+	}
+}