backend: closed before/after handshake (#157)

pingcap · Dec 22, 2022 · b5d2300 · b5d2300
1 parent 8aee2ce
commit b5d2300
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
@@ -90,6 +90,7 @@ type BackendConnManager struct {
 	handshakeHandler HandshakeHandler
 	getBackendIO     backendIOGetter
 	connectionID     uint64
+	handshaked       bool
 }
 
 // NewBackendConnManager creates a BackendConnManager.
@@ -165,6 +166,7 @@ func (mgr *BackendConnManager) Connect(ctx context.Context, clientIO *pnet.Packe
 	mgr.wg.Run(func() {
 		mgr.processSignals(childCtx, clientIO)
 	})
+	mgr.handshaked = true
 	return nil
 }
 
@@ -409,7 +411,7 @@ func (mgr *BackendConnManager) Close() error {
 	}
 	mgr.processLock.Unlock()
 
-	handErr := mgr.handshakeHandler.OnConnClose(mgr.authenticator)
+	handErr := mgr.handshakeHandler.OnConnClose(mgr.authenticator, mgr.handshaked)
 
 	eventReceiver := mgr.getEventReceiver()
 	if eventReceiver != nil {

diff --git a/pkg/proxy/backend/handshake_handler.go b/pkg/proxy/backend/handshake_handler.go
@@ -42,7 +42,7 @@ type ConnContext interface {
 type HandshakeHandler interface {
 	HandleHandshakeResp(ctx ConnContext, resp *pnet.HandshakeResp) error
 	GetRouter(ctx ConnContext, resp *pnet.HandshakeResp) (router.Router, error)
-	OnConnClose(ctx ConnContext) error
+	OnConnClose(ctx ConnContext, handshaked bool) error
 	GetCapability() pnet.Capability
 }
 
@@ -71,7 +71,7 @@ func (handler *DefaultHandshakeHandler) GetRouter(ctx ConnContext, resp *pnet.Ha
 	return ns.GetRouter(), nil
 }
 
-func (handler *DefaultHandshakeHandler) OnConnClose(ConnContext) error {
+func (handler *DefaultHandshakeHandler) OnConnClose(ConnContext, bool) error {
 	return nil
 }
 

diff --git a/pkg/proxy/backend/mock_proxy_test.go b/pkg/proxy/backend/mock_proxy_test.go
@@ -111,7 +111,7 @@ func (handler *CustomHandshakeHandler) GetRouter(ctx ConnContext, resp *pnet.Han
 	return nil, nil
 }
 
-func (handler *CustomHandshakeHandler) OnConnClose(ctx ConnContext) error {
+func (handler *CustomHandshakeHandler) OnConnClose(ctx ConnContext, _ bool) error {
 	return nil
 }