From fca7f42b370ad16f51c8155e433b06e266661ec7 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Sun, 11 Jun 2017 17:22:05 +0800
Subject: [PATCH] msg: new message CloseProxy

---
 client/control.go |  3 +++
 models/msg/msg.go |  6 ++++++
 server/control.go | 29 ++++++++++++++++++++++++++---
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/client/control.go b/client/control.go
index 49dfb29..c702058 100644
--- a/client/control.go
+++ b/client/control.go
@@ -358,6 +358,9 @@ func (ctl *Control) manager() {
 				pxy := NewProxy(ctl, cfg)
 				if err := pxy.Run(); err != nil {
 					ctl.Warn("[%s] proxy start running error: %v", m.ProxyName, err)
+					ctl.sendCh <- &msg.CloseProxy{
+						ProxyName: m.ProxyName,
+					}
 					continue
 				}
 				ctl.proxies[m.ProxyName] = pxy
diff --git a/models/msg/msg.go b/models/msg/msg.go
index a360e5b..d961bef 100644
--- a/models/msg/msg.go
+++ b/models/msg/msg.go
@@ -24,6 +24,7 @@ const (
 	TypeLoginResp     = '1'
 	TypeNewProxy      = 'p'
 	TypeNewProxyResp  = '2'
+	TypeCloseProxy    = 'c'
 	TypeNewWorkConn   = 'w'
 	TypeReqWorkConn   = 'r'
 	TypeStartWorkConn = 's'
@@ -45,6 +46,7 @@ func init() {
 	TypeMap[TypeLoginResp] = reflect.TypeOf(LoginResp{})
 	TypeMap[TypeNewProxy] = reflect.TypeOf(NewProxy{})
 	TypeMap[TypeNewProxyResp] = reflect.TypeOf(NewProxyResp{})
+	TypeMap[TypeCloseProxy] = reflect.TypeOf(CloseProxy{})
 	TypeMap[TypeNewWorkConn] = reflect.TypeOf(NewWorkConn{})
 	TypeMap[TypeReqWorkConn] = reflect.TypeOf(ReqWorkConn{})
 	TypeMap[TypeStartWorkConn] = reflect.TypeOf(StartWorkConn{})
@@ -105,6 +107,10 @@ type NewProxyResp struct {
 	Error     string `json:"error"`
 }
 
+type CloseProxy struct {
+	ProxyName string `json:"proxy_name"`
+}
+
 type NewWorkConn struct {
 	RunId string `json:"run_id"`
 }
diff --git a/server/control.go b/server/control.go
index f2ee04a..d6b2c2c 100644
--- a/server/control.go
+++ b/server/control.go
@@ -50,7 +50,7 @@ type Control struct {
 	workConnCh chan net.Conn
 
 	// proxies in one client
-	proxies []Proxy
+	proxies map[string]Proxy
 
 	// pool count
 	poolCount int
@@ -82,7 +82,7 @@ func NewControl(svr *Service, ctlConn net.Conn, loginMsg *msg.Login) *Control {
 		sendCh:          make(chan msg.Message, 10),
 		readCh:          make(chan msg.Message, 10),
 		workConnCh:      make(chan net.Conn, loginMsg.PoolCount+10),
-		proxies:         make([]Proxy, 0),
+		proxies:         make(map[string]Proxy),
 		poolCount:       loginMsg.PoolCount,
 		lastPing:        time.Now(),
 		runId:           loginMsg.RunId,
@@ -265,6 +265,8 @@ func (ctl *Control) stoper() {
 		workConn.Close()
 	}
 
+	ctl.mu.Lock()
+	defer ctl.mu.Unlock()
 	for _, pxy := range ctl.proxies {
 		pxy.Close()
 		ctl.svr.DelProxy(pxy.GetName())
@@ -317,6 +319,9 @@ func (ctl *Control) manager() {
 					StatsNewProxy(m.ProxyName, m.ProxyType)
 				}
 				ctl.sendCh <- resp
+			case *msg.CloseProxy:
+				ctl.CloseProxy(m)
+				ctl.conn.Info("close proxy [%s] success", m.ProxyName)
 			case *msg.Ping:
 				ctl.lastPing = time.Now()
 				ctl.conn.Debug("receive heartbeat")
@@ -355,6 +360,24 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (err error) {
 	if err != nil {
 		return err
 	}
-	ctl.proxies = append(ctl.proxies, pxy)
+
+	ctl.mu.Lock()
+	ctl.proxies[pxy.GetName()] = pxy
+	ctl.mu.Unlock()
 	return nil
 }
+
+func (ctl *Control) CloseProxy(closeMsg *msg.CloseProxy) (err error) {
+	ctl.mu.Lock()
+	defer ctl.mu.Unlock()
+
+	pxy, ok := ctl.proxies[closeMsg.ProxyName]
+	if !ok {
+		return
+	}
+
+	pxy.Close()
+	ctl.svr.DelProxy(pxy.GetName())
+	StatsCloseProxy(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)
+	return
+}