Repository: fatedier/frp (language: go)

Issue:
[Feature Request] Support Client Service Custom Context
### Describe the feature request

First of all, this is a fantastic project. We use this as a library, and I hope that in the future it will be better designed to support usage as a library. When creating a new client service, I do not have the option to pass in my own context, rather the service struct creates and maintains it's own context.

```go
func NewService(cfg config.ClientCommonConf, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf, cfgFile string) (svr *Service, err error) {
	ctx, cancel := context.WithCancel(context.Background())
	svr = &Service{
```

It would be easier if I could specify my own context like this

```go
svc := client.NewService(...)
svc.Run(ctx)
```

As it stands today, I have to hack my way around this

```go
svc := client.NewService(...)
go func() {
    <-ctx.Done()
   svc.Close()
}
svc.Run()
```

### Describe alternatives you've considered

_No response_

### Affected area

- [ ] Docs
- [ ] Installation
- [ ] Performance and Scalability
- [ ] Security
- [X] User Experience
- [ ] Test and Release
- [ ] Developer Infrastructure
- [X] Client Plugin
- [ ] Server Plugin
- [ ] Extensions
- [ ] Others


Discussion / hints:
Yes. We are doing some code refactoring. Used as a third-party library is also one of aims.
Unrelated but are there plans to clean up the logging implementation and use a simple interface rather than registering named loggers and then configuring which named logger to use?
> Unrelated but are there plans to clean up the logging implementation and use a simple interface rather than registering named loggers and then configuring which named logger to use?

Good point.

Fix patch:
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index 13b47c6950b..d24eaf818de 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -4,7 +4,3 @@ copilot:summary
 
 ### WHY
 <!-- author to complete -->
-
-### Walkthrough
-
-copilot:walkthrough
diff --git a/client/service.go b/client/service.go
index 2f7782d5e32..b5366b1e319 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 82c61720bba..1af6a29f2be 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 1ba1352f3a4..e1859376a9a 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 3c4f5a81e5b..7fbcb364be4 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
 }
 


Test patch: