From 29051d4aa7979a23bbb5fc794e64f92dd9d2d7ac Mon Sep 17 00:00:00 2001
From: kennytm <kennytm@gmail.com>
Date: Wed, 11 Mar 2020 09:00:28 +0800
Subject: [PATCH] conn: support not shutting down the storage when closing the
 connection

---
 pkg/conn/conn.go     | 21 +++++++++++++--------
 pkg/glue/glue.go     |  4 ++++
 pkg/gluetidb/glue.go |  5 +++++
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/pkg/conn/conn.go b/pkg/conn/conn.go
index 2ab0a0232..45368a5b0 100644
--- a/pkg/conn/conn.go
+++ b/pkg/conn/conn.go
@@ -55,6 +55,7 @@ type Mgr struct {
 		mu   sync.Mutex
 		clis map[uint64]*grpc.ClientConn
 	}
+	ownsStorage bool
 }
 
 type pdHTTPRequest func(context.Context, string, string, *http.Client, string, io.Reader) ([]byte, error)
@@ -165,10 +166,11 @@ func NewMgr(
 	}
 
 	mgr := &Mgr{
-		pdClient: pdClient,
-		storage:  storage,
-		dom:      dom,
-		tlsConf:  tlsConf,
+		pdClient:    pdClient,
+		storage:     storage,
+		dom:         dom,
+		tlsConf:     tlsConf,
+		ownsStorage: g.OwnsStorage(),
 	}
 	mgr.pdHTTP.addrs = processedAddrs
 	mgr.pdHTTP.cli = cli
@@ -386,11 +388,14 @@ func (mgr *Mgr) Close() {
 
 	// Gracefully shutdown domain so it does not affect other TiDB DDL.
 	// Must close domain before closing storage, otherwise it gets stuck forever.
-	if mgr.dom != nil {
-		mgr.dom.Close()
+	if mgr.ownsStorage {
+		if mgr.dom != nil {
+			mgr.dom.Close()
+		}
+
+		atomic.StoreUint32(&tikv.ShuttingDown, 1)
+		mgr.storage.Close()
 	}
 
-	atomic.StoreUint32(&tikv.ShuttingDown, 1)
-	mgr.storage.Close()
 	mgr.pdClient.Close()
 }
diff --git a/pkg/glue/glue.go b/pkg/glue/glue.go
index b680370aa..2746090a3 100644
--- a/pkg/glue/glue.go
+++ b/pkg/glue/glue.go
@@ -15,6 +15,10 @@ type Glue interface {
 	BootstrapSession(store kv.Storage) (*domain.Domain, error)
 	CreateSession(store kv.Storage) (Session, error)
 	Open(path string, option pd.SecurityOption) (kv.Storage, error)
+
+	// OwnsStorage returns whether the storage returned by Open() is owned
+	// If this method returns false, the connection manager will never close the storage.
+	OwnsStorage() bool
 }
 
 // Session is an abstraction of the session.Session interface.
diff --git a/pkg/gluetidb/glue.go b/pkg/gluetidb/glue.go
index 27ae01c37..3feb1ab7f 100644
--- a/pkg/gluetidb/glue.go
+++ b/pkg/gluetidb/glue.go
@@ -50,6 +50,11 @@ func (Glue) Open(path string, option pd.SecurityOption) (kv.Storage, error) {
 	return tikv.Driver{}.Open(path)
 }
 
+// OwnsStorage implements glue.Glue
+func (Glue) OwnsStorage() bool {
+	return true
+}
+
 // Execute implements glue.Session
 func (gs *tidbSession) Execute(ctx context.Context, sql string) error {
 	_, err := gs.se.Execute(ctx, sql)