proxy: add OnConnClose for handler (#154)

pingcap · Dec 16, 2022 · 1a5316d · 1a5316d
1 parent 5490197
commit 1a5316d
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 8 deletions.
diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
@@ -36,6 +36,10 @@ import (
 	"go.uber.org/zap"
 )
 
+var (
+	ErrCloseConnMgr = errors.New("failed to close connection manager")
+)
+
 const (
 	sqlQueryState    = "SHOW SESSION_STATES"
 	sqlSetState      = "SET SESSION_STATES '%s'"
@@ -387,16 +391,18 @@ func (mgr *BackendConnManager) Close() error {
 	}
 	mgr.wg.Wait()
 
-	var err error
+	var connErr error
 	var addr string
 	mgr.processLock.Lock()
 	if mgr.backendConn != nil {
 		addr = mgr.backendConn.address
-		err = mgr.backendConn.Close()
+		connErr = mgr.backendConn.Close()
 		mgr.backendConn = nil
 	}
 	mgr.processLock.Unlock()
 
+	handErr := mgr.handshakeHandler.OnConnClose(mgr.authenticator)
+
 	eventReceiver := mgr.getEventReceiver()
 	if eventReceiver != nil {
 		// Notify the receiver if there's any event.
@@ -410,5 +416,5 @@ func (mgr *BackendConnManager) Close() error {
 			}
 		}
 	}
-	return err
+	return errors.Collect(ErrCloseConnMgr, connErr, handErr)
 }
diff --git a/pkg/proxy/backend/handshake_handler.go b/pkg/proxy/backend/handshake_handler.go
@@ -41,8 +41,9 @@ type ConnContext interface {
 
 type HandshakeHandler interface {
 	HandleHandshakeResp(ctx ConnContext, resp *pnet.HandshakeResp) error
-	GetCapability() pnet.Capability
 	GetRouter(ctx ConnContext, resp *pnet.HandshakeResp) (router.Router, error)
+	OnConnClose(ctx ConnContext) error
+	GetCapability() pnet.Capability
 }
 
 type DefaultHandshakeHandler struct {
@@ -59,10 +60,6 @@ func (handler *DefaultHandshakeHandler) HandleHandshakeResp(ConnContext, *pnet.H
 	return nil
 }
 
-func (handler *DefaultHandshakeHandler) GetCapability() pnet.Capability {
-	return SupportedServerCapabilities
-}
-
 func (handler *DefaultHandshakeHandler) GetRouter(ctx ConnContext, resp *pnet.HandshakeResp) (router.Router, error) {
 	ns, ok := handler.nsManager.GetNamespaceByUser(resp.User)
 	if !ok {
@@ -73,3 +70,11 @@ func (handler *DefaultHandshakeHandler) GetRouter(ctx ConnContext, resp *pnet.Ha
 	}
 	return ns.GetRouter(), nil
 }
+
+func (handler *DefaultHandshakeHandler) OnConnClose(ConnContext) error {
+	return nil
+}
+
+func (handler *DefaultHandshakeHandler) GetCapability() pnet.Capability {
+	return SupportedServerCapabilities
+}
diff --git a/pkg/proxy/backend/mock_proxy_test.go b/pkg/proxy/backend/mock_proxy_test.go
@@ -111,6 +111,10 @@ func (handler *CustomHandshakeHandler) GetRouter(ctx ConnContext, resp *pnet.Han
 	return nil, nil
 }
 
+func (handler *CustomHandshakeHandler) OnConnClose(ctx ConnContext) error {
+	return nil
+}
+
 func (handler *CustomHandshakeHandler) HandleHandshakeResp(ctx ConnContext, resp *pnet.HandshakeResp) error {
 	handler.inUsername = resp.User
 	resp.User = handler.outUsername