conn: support not shutting down the storage when closing the connecti…

…on (#185)

Co-authored-by: 3pointer <[email protected]>

pkg/conn/conn.go

@@ -57,6 +57,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)
@@ -167,10 +168,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
@@ -388,11 +390,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()
 }

pkg/glue/glue.go

@@ -17,6 +17,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.

pkg/gluetidb/glue.go

@@ -52,6 +52,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)