Browse Source

improve kcp shutdown

fatedier 6 years ago
parent
commit
7c21906884
5 changed files with 14 additions and 3 deletions
  1. 1 0
      client/control.go
  2. 3 2
      client/proxy/proxy_manager.go
  3. 1 0
      client/service.go
  4. 8 1
      cmd/frpc/sub/root.go
  5. 1 0
      server/service.go

+ 1 - 0
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
 }

+ 3 - 2
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) {

+ 1 - 0
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 {

+ 8 - 1
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
 }

+ 1 - 0
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 {