|
@@ -29,7 +29,6 @@ import (
|
|
|
|
|
|
"github.com/fatedier/frp/pkg/auth"
|
|
|
"github.com/fatedier/frp/pkg/config"
|
|
|
- "github.com/fatedier/frp/pkg/consts"
|
|
|
pkgerr "github.com/fatedier/frp/pkg/errors"
|
|
|
"github.com/fatedier/frp/pkg/msg"
|
|
|
plugin "github.com/fatedier/frp/pkg/plugin/server"
|
|
@@ -55,13 +54,14 @@ func NewControlManager() *ControlManager {
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-func (cm *ControlManager) Add(runID string, ctl *Control) (oldCtl *Control) {
|
|
|
+func (cm *ControlManager) Add(runID string, ctl *Control) (old *Control) {
|
|
|
cm.mu.Lock()
|
|
|
defer cm.mu.Unlock()
|
|
|
|
|
|
- oldCtl, ok := cm.ctlsByRunID[runID]
|
|
|
+ var ok bool
|
|
|
+ old, ok = cm.ctlsByRunID[runID]
|
|
|
if ok {
|
|
|
- oldCtl.Replaced(ctl)
|
|
|
+ old.Replaced(ctl)
|
|
|
}
|
|
|
cm.ctlsByRunID[runID] = ctl
|
|
|
return
|
|
@@ -141,14 +141,13 @@ type Control struct {
|
|
|
// replace old controller instantly.
|
|
|
runID string
|
|
|
|
|
|
- // control status
|
|
|
- status string
|
|
|
-
|
|
|
readerShutdown *shutdown.Shutdown
|
|
|
writerShutdown *shutdown.Shutdown
|
|
|
managerShutdown *shutdown.Shutdown
|
|
|
allShutdown *shutdown.Shutdown
|
|
|
|
|
|
+ started bool
|
|
|
+
|
|
|
mu sync.RWMutex
|
|
|
|
|
|
// Server configuration information
|
|
@@ -187,7 +186,6 @@ func NewControl(
|
|
|
portsUsedNum: 0,
|
|
|
lastPing: time.Now(),
|
|
|
runID: loginMsg.RunID,
|
|
|
- status: consts.Working,
|
|
|
readerShutdown: shutdown.New(),
|
|
|
writerShutdown: shutdown.New(),
|
|
|
managerShutdown: shutdown.New(),
|
|
@@ -208,11 +206,19 @@ func (ctl *Control) Start() {
|
|
|
Error: "",
|
|
|
}
|
|
|
_ = msg.WriteMsg(ctl.conn, loginRespMsg)
|
|
|
+ ctl.mu.Lock()
|
|
|
+ ctl.started = true
|
|
|
+ ctl.mu.Unlock()
|
|
|
|
|
|
go ctl.writer()
|
|
|
- for i := 0; i < ctl.poolCount; i++ {
|
|
|
- ctl.sendCh <- &msg.ReqWorkConn{}
|
|
|
- }
|
|
|
+ go func() {
|
|
|
+ for i := 0; i < ctl.poolCount; i++ {
|
|
|
+ // ignore error here, that means that this control is closed
|
|
|
+ _ = errors.PanicToError(func() {
|
|
|
+ ctl.sendCh <- &msg.ReqWorkConn{}
|
|
|
+ })
|
|
|
+ }
|
|
|
+ }()
|
|
|
|
|
|
go ctl.manager()
|
|
|
go ctl.reader()
|
|
@@ -418,6 +424,14 @@ func (ctl *Control) stoper() {
|
|
|
|
|
|
// block until Control closed
|
|
|
func (ctl *Control) WaitClosed() {
|
|
|
+ ctl.mu.RLock()
|
|
|
+ started := ctl.started
|
|
|
+ ctl.mu.RUnlock()
|
|
|
+
|
|
|
+ if !started {
|
|
|
+ ctl.allShutdown.Done()
|
|
|
+ return
|
|
|
+ }
|
|
|
ctl.allShutdown.WaitDone()
|
|
|
}
|
|
|
|