From 4c4d5f0d0daa2e88be1d90945d9fa9ad52ed5911 Mon Sep 17 00:00:00 2001 From: fatedier Date: Thu, 29 Jun 2023 18:04:20 +0800 Subject: [PATCH] service.Run supports passing in context (#3504) --- .github/pull_request_template.md | 4 ---- client/service.go | 21 +++++++++++++++------ cmd/frpc/sub/root.go | 3 ++- cmd/frps/root.go | 3 ++- server/service.go | 26 ++++++++++++++++++++------ 5 files changed, 39 insertions(+), 18 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 13b47c6..d24eaf8 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -4,7 +4,3 @@ copilot:summary ### WHY - -### Walkthrough - -copilot:walkthrough diff --git a/client/service.go b/client/service.go index 2f7782d..b5366b1 100644 --- a/client/service.go +++ b/client/service.go @@ -86,16 +86,14 @@ func NewService( visitorCfgs map[string]config.VisitorConf, cfgFile string, ) (svr *Service, err error) { - ctx, cancel := context.WithCancel(context.Background()) svr = &Service{ authSetter: auth.NewAuthSetter(cfg.ClientConfig), cfg: cfg, cfgFile: cfgFile, pxyCfgs: pxyCfgs, visitorCfgs: visitorCfgs, + ctx: context.Background(), exit: 0, - ctx: xlog.NewContext(ctx, xlog.New()), - cancel: cancel, } return } @@ -106,7 +104,11 @@ func (svr *Service) GetController() *Control { return svr.ctl } -func (svr *Service) Run() error { +func (svr *Service) Run(ctx context.Context) error { + ctx, cancel := context.WithCancel(ctx) + svr.ctx = xlog.NewContext(ctx, xlog.New()) + svr.cancel = cancel + xl := xlog.FromContextSafe(svr.ctx) // set custom DNSServer @@ -161,6 +163,10 @@ func (svr *Service) Run() error { log.Info("admin server listen on %s:%d", svr.cfg.AdminAddr, svr.cfg.AdminPort) } <-svr.ctx.Done() + // service context may not be canceled by svr.Close(), we should call it here to release resources + if atomic.LoadUint32(&svr.exit) == 0 { + svr.Close() + } return nil } @@ -182,7 +188,7 @@ func (svr *Service) keepControllerWorking() { return } - // the first three retry with no delay + // the first three attempts with a low delay if reconnectCounts > 3 { util.RandomSleep(reconnectDelay, 0.9, 1.1) xl.Info("wait %v to reconnect", reconnectDelay) @@ -322,10 +328,13 @@ func (svr *Service) GracefulClose(d time.Duration) { svr.ctlMu.RLock() if svr.ctl != nil { svr.ctl.GracefulClose(d) + svr.ctl = nil } svr.ctlMu.RUnlock() - svr.cancel() + if svr.cancel != nil { + svr.cancel() + } } type ConnectionManager struct { diff --git a/cmd/frpc/sub/root.go b/cmd/frpc/sub/root.go index 82c6172..1af6a29 100644 --- a/cmd/frpc/sub/root.go +++ b/cmd/frpc/sub/root.go @@ -15,6 +15,7 @@ package sub import ( + "context" "fmt" "io/fs" "net" @@ -233,7 +234,7 @@ func startService( go handleSignal(svr, closedDoneCh) } - err = svr.Run() + err = svr.Run(context.Background()) if err == nil && shouldGracefulClose { <-closedDoneCh } diff --git a/cmd/frps/root.go b/cmd/frps/root.go index 1ba1352..e185937 100644 --- a/cmd/frps/root.go +++ b/cmd/frps/root.go @@ -15,6 +15,7 @@ package main import ( + "context" "fmt" "os" @@ -210,6 +211,6 @@ func runServer(cfg config.ServerCommonConf) (err error) { return err } log.Info("frps started successfully") - svr.Run() + svr.Run(context.Background()) return } diff --git a/server/service.go b/server/service.go index 3c4f5a8..7fbcb36 100644 --- a/server/service.go +++ b/server/service.go @@ -115,7 +115,6 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) { return } - ctx, cancel := context.WithCancel(context.Background()) svr = &Service{ ctlManager: NewControlManager(), pxyManager: proxy.NewManager(), @@ -129,8 +128,7 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) { authVerifier: auth.NewAuthVerifier(cfg.ServerConfig), tlsConfig: tlsConfig, cfg: cfg, - ctx: ctx, - cancel: cancel, + ctx: context.Background(), } // Create tcpmux httpconnect multiplexer. @@ -329,7 +327,11 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) { return } -func (svr *Service) Run() { +func (svr *Service) Run(ctx context.Context) { + ctx, cancel := context.WithCancel(ctx) + svr.ctx = ctx + svr.cancel = cancel + if svr.kcpListener != nil { go svr.HandleListener(svr.kcpListener) } @@ -343,27 +345,39 @@ func (svr *Service) Run() { go svr.rc.NatHoleController.CleanWorker(svr.ctx) } svr.HandleListener(svr.listener) + + <-svr.ctx.Done() + // service context may not be canceled by svr.Close(), we should call it here to release resources + if svr.listener != nil { + svr.Close() + } } func (svr *Service) Close() error { if svr.kcpListener != nil { svr.kcpListener.Close() + svr.kcpListener = nil } if svr.quicListener != nil { svr.quicListener.Close() + svr.quicListener = nil } if svr.websocketListener != nil { svr.websocketListener.Close() + svr.websocketListener = nil } if svr.tlsListener != nil { svr.tlsListener.Close() + svr.tlsConfig = nil } if svr.listener != nil { svr.listener.Close() + svr.listener = nil } - svr.cancel() - svr.ctlManager.Close() + if svr.cancel != nil { + svr.cancel() + } return nil }