From c8f29496b5fb6aeec17df6b3de6e67e8ce456944 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Fri, 28 Aug 2020 13:25:49 +0800
Subject: [PATCH] basedb: close db when fail ping (#942) (#943)

---
 pkg/conn/basedb.go      | 23 +++++++++++++++++++++--
 pkg/conn/basedb_test.go | 17 +++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/pkg/conn/basedb.go b/pkg/conn/basedb.go
index ae63fd1d0f..f157a6f530 100644
--- a/pkg/conn/basedb.go
+++ b/pkg/conn/basedb.go
@@ -16,12 +16,16 @@ package conn
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"net/url"
 	"strconv"
 	"sync"
 	"sync/atomic"
 
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/pingcap/failpoint"
+
 	"github.com/pingcap/dm/dm/config"
 	"github.com/pingcap/dm/pkg/retry"
 	"github.com/pingcap/dm/pkg/terror"
@@ -48,6 +52,9 @@ func init() {
 	DefaultDBProvider = &DefaultDBProviderImpl{}
 }
 
+// mock is used in unit test
+var mock sqlmock.Sqlmock
+
 // Apply will build BaseDB with DBConfig
 func (d *DefaultDBProviderImpl) Apply(config config.DBConfig) (*BaseDB, error) {
 	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?charset=utf8mb4&interpolateParams=true&maxAllowedPacket=%d",
@@ -96,8 +103,20 @@ func (d *DefaultDBProviderImpl) Apply(config config.DBConfig) (*BaseDB, error) {
 	if err != nil {
 		return nil, terror.DBErrorAdapt(err, terror.ErrDBDriverError)
 	}
-
-	if err = db.Ping(); err != nil {
+	failpoint.Inject("failDBPing", func(_ failpoint.Value) {
+		db.Close()
+		db, mock, _ = sqlmock.New()
+		mock.ExpectPing()
+		mock.ExpectClose()
+	})
+
+	err = db.Ping()
+	failpoint.Inject("failDBPing", func(_ failpoint.Value) {
+		err = errors.New("injected error")
+	})
+	if err != nil {
+		db.Close()
+		doFuncInClose()
 		return nil, terror.DBErrorAdapt(err, terror.ErrDBDriverError)
 	}
 
diff --git a/pkg/conn/basedb_test.go b/pkg/conn/basedb_test.go
index 2f365aaf18..9c621cb7e7 100644
--- a/pkg/conn/basedb_test.go
+++ b/pkg/conn/basedb_test.go
@@ -15,8 +15,11 @@ package conn
 
 import (
 	. "github.com/pingcap/check"
+	"github.com/pingcap/failpoint"
 
 	sqlmock "github.com/DATA-DOG/go-sqlmock"
+
+	"github.com/pingcap/dm/dm/config"
 	tcontext "github.com/pingcap/dm/pkg/context"
 )
 
@@ -57,3 +60,17 @@ func (t *testBaseDBSuite) TestGetBaseConn(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(affected, Equals, 1)
 }
+
+func (t *testBaseDBSuite) TestFailDBPing(c *C) {
+	c.Assert(failpoint.Enable("github.com/pingcap/dm/pkg/conn/failDBPing", "return"), IsNil)
+	defer failpoint.Disable("github.com/pingcap/dm/pkg/conn/failDBPing")
+
+	cfg := config.DBConfig{User: "root", Host: "127.0.0.1", Port: 3306}
+	cfg.Adjust()
+	db, err := DefaultDBProvider.Apply(cfg)
+	c.Assert(db, IsNil)
+	c.Assert(err, NotNil)
+
+	err = mock.ExpectationsWereMet()
+	c.Assert(err, IsNil)
+}