service.Run supports passing in context (#3504)

.github/pull_request_template.md

@@ -4,7 +4,3 @@ copilot:summary
 
 ### WHY
 <!-- author to complete -->
-
-### Walkthrough
-
-copilot:walkthrough

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,11 +328,14 @@ func (svr *Service) GracefulClose(d time.Duration) {
 	svr.ctlMu.RLock()
 	if svr.ctl != nil {
 		svr.ctl.GracefulClose(d)
+		svr.ctl = nil
 	}
 	svr.ctlMu.RUnlock()
 
+	if svr.cancel != nil {
 		svr.cancel()
 	}
+}
 
 type ConnectionManager struct {
 	ctx context.Context

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
 	}

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
 }

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
 }