Mxmilu666/frp
1
1
Fork 1
mirror of https://github.com/fatedier/frp.git synced 2026-10-16 09:55:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: synchronize service startup and shutdown (#5560)

Browse Source
This commit is contained in:
fatedier
2026-10-09 12:27:12 +08:00
committed by GitHub
parent 4cd4747525
commit 55107d003d
6 changed files with 485 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/service.go
+41 -9
View File
@@ -25,6 +25,7 @@ import (
"net/http"
"os"
"strconv"
"sync"
"time"
"github.com/fatedier/golib/crypto"
@@ -120,6 +121,9 @@ type Service struct {
// web server for dashboard UI and apis
webServer *httppkg.Server
// HTTP vhost resources are started and owned by NewService.
httpVhostServer *http.Server
httpVhostListener net.Listener
sshTunnelGateway *ssh.Gateway
@@ -133,7 +137,10 @@ type Service struct {
// service context
ctx context.Context
// call cancel to stop service
cancel context.CancelFunc
cancel context.CancelFunc
lifecycleMu sync.Mutex
closed bool
closeOnce sync.Once
}
func NewService(cfg *v1.ServerConfig) (*Service, error) {
@@ -319,6 +326,8 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) {
return nil, fmt.Errorf("create vhost http listener error, %v", err)
}
}
svr.httpVhostServer = server
svr.httpVhostListener = l
go func() {
_ = server.Serve(l)
}()
@@ -365,8 +374,22 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) {
func (svr *Service) Run(ctx context.Context) {
ctx, cancel := context.WithCancel(ctx)
svr.lifecycleMu.Lock()
if svr.closed {
svr.lifecycleMu.Unlock()
cancel()
return
}
svr.ctx = ctx
svr.cancel = cancel
svr.lifecycleMu.Unlock()
defer svr.Close()
if ctx.Err() != nil {
return
}
// Cancellation must unblock the synchronous main listener's Accept.
stopClose := context.AfterFunc(ctx, func() { _ = svr.Close() })
defer stopClose()
// run dashboard web server.
if svr.webServer != nil {
@@ -400,13 +423,21 @@ func (svr *Service) Run(ctx context.Context) {
svr.HandleListener(svr.listener, false)
<-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 {
svr.lifecycleMu.Lock()
svr.closed = true
cancel := svr.cancel
svr.lifecycleMu.Unlock()
if cancel != nil {
cancel()
}
svr.closeOnce.Do(svr.closeResources)
return nil
}
func (svr *Service) closeResources() {
if svr.kcpListener != nil {
svr.kcpListener.Close()
}
@@ -428,16 +459,17 @@ func (svr *Service) Close() error {
if svr.webServer != nil {
svr.webServer.Close()
}
if svr.httpVhostServer != nil {
svr.httpVhostServer.Close()
// Close even if Serve has not registered its listener yet.
svr.httpVhostListener.Close()
}
if svr.sshTunnelGateway != nil {
svr.sshTunnelGateway.Close()
}
svr.rc.Close()
svr.muxer.Close()
svr.ctlManager.Close()
if svr.cancel != nil {
svr.cancel()
}
return nil
}
func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) {
server/service_lifecycle_test.go
+151
View File
@@ -0,0 +1,151 @@
package server
import (
"context"
"fmt"
"net"
"sync"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
v1 "github.com/fatedier/frp/pkg/config/v1"
"github.com/fatedier/frp/pkg/config/v1/validation"
)
type lifecycleListener struct {
net.Listener
accepted, closeEntered, closeResume chan struct{}
acceptOnce, closeOnce sync.Once
accepts, closes atomic.Int64
}
func (l *lifecycleListener) Accept() (net.Conn, error) {
l.accepts.Add(1)
l.acceptOnce.Do(func() { close(l.accepted) })
return l.Listener.Accept()
}
func (l *lifecycleListener) Close() error {
l.closes.Add(1)
err := l.Listener.Close()
if l.closeEntered != nil {
l.closeOnce.Do(func() { close(l.closeEntered) })
<-l.closeResume
}
return err
}
func freeLifecycleTCPPort(t *testing.T) int {
t.Helper()
l, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer l.Close()
return l.Addr().(*net.TCPAddr).Port
}
func newLifecycleTestService(t *testing.T, web bool) (*Service, *lifecycleListener, []string) {
t.Helper()
cfg := &v1.ServerConfig{
BindAddr: "127.0.0.1", ProxyBindAddr: "127.0.0.1",
}
if web {
cfg.WebServer = v1.WebServerConfig{Addr: "127.0.0.1", Port: freeLifecycleTCPPort(t)}
cfg.VhostHTTPPort = freeLifecycleTCPPort(t)
}
require.NoError(t, cfg.Complete())
cfg.BindPort = 0
_, err := validation.NewConfigValidator(nil).ValidateServerConfig(cfg)
require.NoError(t, err)
svr, err := NewService(cfg)
require.NoError(t, err)
l := &lifecycleListener{Listener: svr.listener, accepted: make(chan struct{})}
svr.listener = l
addresses := []string{l.Addr().String()}
if web {
addresses = append(addresses, svr.webServer.Address())
addresses = append(addresses, net.JoinHostPort(cfg.ProxyBindAddr, fmt.Sprint(cfg.VhostHTTPPort)))
}
assertLifecyclePorts(t, addresses, false)
return svr, l, addresses
}
func assertLifecyclePorts(t *testing.T, addresses []string, released bool) {
t.Helper()
for _, address := range addresses {
ln, err := net.Listen("tcp", address)
if ln != nil {
_ = ln.Close()
}
if released {
require.NoError(t, err, "listener %s was not released", address)
} else {
require.Error(t, err, "listener %s was not bound", address)
}
}
}
func TestServiceRunAfterEarlyShutdown(t *testing.T) {
for _, mode := range []string{"close before Run", "close before Run with web", "parent already cancelled"} {
t.Run(mode, func(t *testing.T) {
svr, listener, addresses := newLifecycleTestService(t, mode == "close before Run with web")
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
defer svr.Close()
if mode != "parent already cancelled" {
require.NoError(t, svr.Close())
assertLifecyclePorts(t, addresses, true)
} else {
cancel()
}
done := make(chan struct{})
go func() { svr.Run(ctx); close(done) }()
waitForSignal(t, done, "Run after early shutdown")
require.Zero(t, listener.accepts.Load(), "closed service started accepting")
require.Equal(t, int64(1), listener.closes.Load())
assertLifecyclePorts(t, addresses, true)
})
}
}
func TestServiceCloseMarksClosedBeforeCleanup(t *testing.T) {
svr, listener, addresses := newLifecycleTestService(t, false)
listener.closeEntered, listener.closeResume = make(chan struct{}), make(chan struct{})
var releaseOnce sync.Once
release := func() { releaseOnce.Do(func() { close(listener.closeResume) }) }
defer release()
closeDone := make(chan struct{})
go func() { _ = svr.Close(); close(closeDone) }()
waitForSignal(t, listener.closeEntered, "Close cleanup barrier")
runDone := make(chan struct{})
go func() { svr.Run(context.Background()); close(runDone) }()
waitForSignal(t, runDone, "Run to observe closed before cleanup finishes")
require.Zero(t, listener.accepts.Load(), "Run started during cleanup")
assertLifecyclePorts(t, addresses, false)
release()
waitForSignal(t, closeDone, "Close cleanup to finish")
assertLifecyclePorts(t, addresses, true)
}
func TestServiceParentCancellationAndCloseShareCleanup(t *testing.T) {
svr, listener, addresses := newLifecycleTestService(t, false)
listener.closeEntered, listener.closeResume = make(chan struct{}), make(chan struct{})
var releaseOnce sync.Once
release := func() { releaseOnce.Do(func() { close(listener.closeResume) }) }
defer release()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
runDone := make(chan struct{})
go func() { svr.Run(ctx); close(runDone) }()
waitForSignal(t, listener.accepted, "main listener Accept")
cancel()
waitForSignal(t, listener.closeEntered, "parent cancellation to close main listener")
closeDone := make(chan struct{})
go func() { _ = svr.Close(); close(closeDone) }()
release()
waitForSignal(t, closeDone, "external Close to finish")
waitForSignal(t, runDone, "Run cleanup to finish")
require.Equal(t, int64(1), listener.closes.Load(), "cleanup ran more than once")
assertLifecyclePorts(t, addresses, true)
}