diff --git a/client/control.go b/client/control.go index 49dfb290..c7020580 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 a360e5b8..d961befa 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 f2ee04af..d6b2c2c6 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 +}