From 7c219068844803535e3b277326b7a3b0754d1ec7 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Tue, 11 Dec 2018 15:06:54 +0800
Subject: [PATCH] improve kcp shutdown

---
 client/control.go             | 1 +
 client/proxy/proxy_manager.go | 5 +++--
 client/service.go             | 1 +
 cmd/frpc/sub/root.go          | 9 ++++++++-
 server/service.go             | 1 +
 5 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/client/control.go b/client/control.go
index 376893f..750803f 100644
--- a/client/control.go
+++ b/client/control.go
@@ -145,6 +145,7 @@ func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
 }
 
 func (ctl *Control) Close() error {
+	ctl.pm.Close()
 	ctl.conn.Close()
 	return nil
 }
diff --git a/client/proxy/proxy_manager.go b/client/proxy/proxy_manager.go
index b65f37a..6f0bb41 100644
--- a/client/proxy/proxy_manager.go
+++ b/client/proxy/proxy_manager.go
@@ -50,11 +50,12 @@ func (pm *ProxyManager) StartProxy(name string, remoteAddr string, serverRespErr
 }
 
 func (pm *ProxyManager) Close() {
-	pm.mu.RLock()
-	defer pm.mu.RUnlock()
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
 	for _, pxy := range pm.proxies {
 		pxy.Stop()
 	}
+	pm.proxies = make(map[string]*ProxyWrapper)
 }
 
 func (pm *ProxyManager) HandleWorkConn(name string, workConn frpNet.Conn) {
diff --git a/client/service.go b/client/service.go
index 62cf151..93cf393 100644
--- a/client/service.go
+++ b/client/service.go
@@ -157,6 +157,7 @@ func (svr *Service) login() (conn frpNet.Conn, session *fmux.Session, err error)
 
 	if g.GlbClientCfg.TcpMux {
 		fmuxCfg := fmux.DefaultConfig()
+		fmuxCfg.KeepAliveInterval = 20 * time.Second
 		fmuxCfg.LogOutput = ioutil.Discard
 		session, err = fmux.Client(conn, fmuxCfg)
 		if err != nil {
diff --git a/cmd/frpc/sub/root.go b/cmd/frpc/sub/root.go
index c770b7a..301d34e 100644
--- a/cmd/frpc/sub/root.go
+++ b/cmd/frpc/sub/root.go
@@ -68,11 +68,15 @@ var (
 	serverName        string
 	bindAddr          string
 	bindPort          int
+
+	kcpDoneCh chan struct{}
 )
 
 func init() {
 	rootCmd.PersistentFlags().StringVarP(&cfgFile, "", "c", "./frpc.ini", "config file of frpc")
 	rootCmd.PersistentFlags().BoolVarP(&showVersion, "version", "v", false, "version of frpc")
+
+	kcpDoneCh = make(chan struct{})
 }
 
 var rootCmd = &cobra.Command{
@@ -106,7 +110,7 @@ func handleSignal(svr *client.Service) {
 	<-ch
 	svr.Close()
 	time.Sleep(250 * time.Millisecond)
-	os.Exit(0)
+	close(kcpDoneCh)
 }
 
 func parseClientCommonCfg(fileType int, content string) (err error) {
@@ -209,5 +213,8 @@ func startService(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]co
 	}
 
 	err = svr.Run()
+	if g.GlbClientCfg.Protocol == "kcp" {
+		<-kcpDoneCh
+	}
 	return
 }
diff --git a/server/service.go b/server/service.go
index 160f7e1..7826745 100644
--- a/server/service.go
+++ b/server/service.go
@@ -290,6 +290,7 @@ func (svr *Service) HandleListener(l frpNet.Listener) {
 
 			if g.GlbServerCfg.TcpMux {
 				fmuxCfg := fmux.DefaultConfig()
+				fmuxCfg.KeepAliveInterval = 20 * time.Second
 				fmuxCfg.LogOutput = ioutil.Discard
 				session, err := fmux.Server(frpConn, fmuxCfg)
 				if err != nil {