mirror of
https://github.com/fatedier/frp.git
synced 2026-10-15 02:25:57 +08:00
fix(server): control work connection refill scheduling
This commit is contained in:
+5
-2
@@ -121,8 +121,9 @@ func (ctl *Control) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn,
|
||||
ctl.pm.SetInWorkConnCallback(cb)
|
||||
}
|
||||
|
||||
func (ctl *Control) handleReqWorkConn(_ msg.Message) {
|
||||
func (ctl *Control) handleReqWorkConn(raw msg.Message) {
|
||||
xl := ctl.xl
|
||||
req := raw.(*msg.ReqWorkConn)
|
||||
workConn, err := ctl.connectServer()
|
||||
if err != nil {
|
||||
xl.Warnf("start new connection to server error: %v", err)
|
||||
@@ -130,7 +131,9 @@ func (ctl *Control) handleReqWorkConn(_ msg.Message) {
|
||||
}
|
||||
|
||||
m := &msg.NewWorkConn{
|
||||
RunID: ctl.sessionCtx.RunID,
|
||||
RunID: ctl.sessionCtx.RunID,
|
||||
WorkConnType: req.WorkConnType,
|
||||
ControlID: req.ControlID,
|
||||
}
|
||||
if err = ctl.sessionCtx.Auth.Setter.SetNewWorkConn(m); err != nil {
|
||||
xl.Warnf("error during NewWorkConn authentication: %v", err)
|
||||
|
||||
+11
-1
@@ -147,9 +147,19 @@ type NewWorkConn struct {
|
||||
RunID string `json:"run_id,omitempty"`
|
||||
PrivilegeKey string `json:"privilege_key,omitempty"`
|
||||
Timestamp int64 `json:"timestamp,omitempty"`
|
||||
WorkConnType string `json:"work_conn_type,omitempty"`
|
||||
ControlID uint64 `json:"control_id,omitempty"`
|
||||
}
|
||||
|
||||
type ReqWorkConn struct{}
|
||||
const (
|
||||
WorkConnTypeDemand = "demand"
|
||||
WorkConnTypeReserve = "reserve"
|
||||
)
|
||||
|
||||
type ReqWorkConn struct {
|
||||
WorkConnType string `json:"work_conn_type,omitempty"`
|
||||
ControlID uint64 `json:"control_id,omitempty"`
|
||||
}
|
||||
|
||||
type StartWorkConn struct {
|
||||
ProxyName string `json:"proxy_name,omitempty"`
|
||||
|
||||
@@ -19,6 +19,8 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/fatedier/frp/pkg/proto/wire"
|
||||
)
|
||||
|
||||
func TestV1MessageTypeIDsAreStable(t *testing.T) {
|
||||
@@ -53,3 +55,22 @@ func TestMessageTypeMapIsCompleteAndUnique(t *testing.T) {
|
||||
msgTypes[msgType] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkConnRequestTypeRoundTrip(t *testing.T) {
|
||||
in := &ReqWorkConn{WorkConnType: WorkConnTypeReserve, ControlID: 7}
|
||||
out, err := DecodeV2MessageFrame(mustEncodeV2MessageFrame(t, in))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, in, out)
|
||||
|
||||
inNew := &NewWorkConn{RunID: "run", WorkConnType: WorkConnTypeDemand, ControlID: 7}
|
||||
outNew, err := DecodeV2MessageFrame(mustEncodeV2MessageFrame(t, inNew))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, inNew, outNew)
|
||||
}
|
||||
|
||||
func mustEncodeV2MessageFrame(t *testing.T, m Message) *wire.Frame {
|
||||
t.Helper()
|
||||
frame, err := EncodeV2MessageFrame(m)
|
||||
require.NoError(t, err)
|
||||
return frame
|
||||
}
|
||||
|
||||
@@ -33,6 +33,10 @@ type Manager struct {
|
||||
newUserConnPlugins []Plugin
|
||||
}
|
||||
|
||||
// ErrNewWorkConnControlIDChanged identifies an attempt to change the control
|
||||
// generation in a NewWorkConn plugin response.
|
||||
var ErrNewWorkConnControlIDChanged = errors.New("plugin changed work connection control generation")
|
||||
|
||||
func NewManager() *Manager {
|
||||
return &Manager{
|
||||
loginPlugins: make([]Plugin, 0),
|
||||
@@ -72,6 +76,7 @@ func handleMutableContent[T any](
|
||||
op string,
|
||||
content *T,
|
||||
logMode pluginErrorLogMode,
|
||||
afterChange func(*T) (*T, error),
|
||||
) (*T, error) {
|
||||
if len(plugins) == 0 {
|
||||
return content, nil
|
||||
@@ -100,6 +105,12 @@ func handleMutableContent[T any](
|
||||
// Preserve the existing Plugin contract: changed content must be *T.
|
||||
// Buggy Plugin implementations still panic here, by design.
|
||||
content = retContent.(*T)
|
||||
if afterChange != nil {
|
||||
content, err = afterChange(content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return content, nil
|
||||
@@ -127,11 +138,11 @@ func (m *Manager) Register(p Plugin) {
|
||||
}
|
||||
|
||||
func (m *Manager) Login(content *LoginContent) (*LoginContent, error) {
|
||||
return handleMutableContent(m.loginPlugins, OpLogin, content, pluginErrorLogWarn)
|
||||
return handleMutableContent(m.loginPlugins, OpLogin, content, pluginErrorLogWarn, nil)
|
||||
}
|
||||
|
||||
func (m *Manager) NewProxy(content *NewProxyContent) (*NewProxyContent, error) {
|
||||
return handleMutableContent(m.newProxyPlugins, OpNewProxy, content, pluginErrorLogWarn)
|
||||
return handleMutableContent(m.newProxyPlugins, OpNewProxy, content, pluginErrorLogWarn, nil)
|
||||
}
|
||||
|
||||
func (m *Manager) CloseProxy(content *CloseProxyContent) error {
|
||||
@@ -157,14 +168,30 @@ func (m *Manager) CloseProxy(content *CloseProxyContent) error {
|
||||
}
|
||||
|
||||
func (m *Manager) Ping(content *PingContent) (*PingContent, error) {
|
||||
return handleMutableContent(m.pingPlugins, OpPing, content, pluginErrorLogWarn)
|
||||
return handleMutableContent(m.pingPlugins, OpPing, content, pluginErrorLogWarn, nil)
|
||||
}
|
||||
|
||||
func (m *Manager) NewWorkConn(content *NewWorkConnContent) (*NewWorkConnContent, error) {
|
||||
return handleMutableContent(m.newWorkConnPlugins, OpNewWorkConn, content, pluginErrorLogWarn)
|
||||
if len(m.newWorkConnPlugins) == 0 {
|
||||
return content, nil
|
||||
}
|
||||
controlID, workConnType := content.ControlID, content.WorkConnType
|
||||
return handleMutableContent(m.newWorkConnPlugins, OpNewWorkConn, content, pluginErrorLogWarn,
|
||||
func(changed *NewWorkConnContent) (*NewWorkConnContent, error) {
|
||||
if changed.ControlID != 0 && changed.ControlID != controlID {
|
||||
return nil, fmt.Errorf("%w from %d to %d", ErrNewWorkConnControlIDChanged, controlID, changed.ControlID)
|
||||
}
|
||||
// Older plugins may omit scheduling metadata. Restore it before the
|
||||
// next plugin observes the replacement, leaving all other fields mutable.
|
||||
// Copy the response so the manager does not mutate plugin-owned content.
|
||||
restored := *changed
|
||||
restored.ControlID = controlID
|
||||
restored.WorkConnType = workConnType
|
||||
return &restored, nil
|
||||
})
|
||||
}
|
||||
|
||||
func (m *Manager) NewUserConn(content *NewUserConnContent) (*NewUserConnContent, error) {
|
||||
// Preserve the pre-refactor log level for NewUserConn plugin errors.
|
||||
return handleMutableContent(m.newUserConnPlugins, OpNewUserConn, content, pluginErrorLogInfo)
|
||||
return handleMutableContent(m.newUserConnPlugins, OpNewUserConn, content, pluginErrorLogInfo, nil)
|
||||
}
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
"time"
|
||||
|
||||
goliblog "github.com/fatedier/golib/log"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
frplog "github.com/fatedier/frp/pkg/util/log"
|
||||
@@ -259,6 +260,133 @@ func TestManagerMutableContentRejectStopsChain(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerNewWorkConnPreservesMetadataBetweenPlugins(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
id uint64
|
||||
kind string
|
||||
}{
|
||||
{name: "typed", id: 7, kind: msg.WorkConnTypeReserve},
|
||||
{name: "legacy"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
m := NewManager()
|
||||
observations := 0
|
||||
for range 2 {
|
||||
// Model an old plugin replacing the entire content, including
|
||||
// intentionally omitted ordinary fields. Only metadata is restored.
|
||||
partial := &NewWorkConnContent{NewWorkConn: msg.NewWorkConn{RunID: "mutated", Timestamp: 42}}
|
||||
m.Register(testPlugin{
|
||||
name: "partial", ops: map[string]bool{OpNewWorkConn: true},
|
||||
handler: func(context.Context, string, any) (*Response, any, error) {
|
||||
return &Response{Unchange: false}, partial, nil
|
||||
},
|
||||
})
|
||||
m.Register(testPlugin{
|
||||
name: "policy", ops: map[string]bool{OpNewWorkConn: true},
|
||||
handler: func(_ context.Context, _ string, raw any) (*Response, any, error) {
|
||||
content := raw.(NewWorkConnContent)
|
||||
observations++
|
||||
if content.ControlID != tc.id || content.WorkConnType != tc.kind {
|
||||
return &Response{Reject: true, RejectReason: "lost scheduling metadata"}, nil, nil
|
||||
}
|
||||
require.Equal(t, "mutated", content.RunID)
|
||||
require.Equal(t, int64(42), content.Timestamp)
|
||||
require.Empty(t, content.User.User, "ordinary omitted fields must not be merged back")
|
||||
require.Zero(t, partial.ControlID, "do not mutate the plugin's response")
|
||||
// Unchange must ignore even an invalid returned generation.
|
||||
content.ControlID = tc.id + 1
|
||||
content.WorkConnType = "ignored"
|
||||
return &Response{Unchange: true}, &content, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
got, err := m.NewWorkConn(&NewWorkConnContent{
|
||||
User: UserInfo{User: "original"},
|
||||
NewWorkConn: msg.NewWorkConn{RunID: "original", ControlID: tc.id, WorkConnType: tc.kind},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, observations)
|
||||
require.Equal(t, tc.id, got.ControlID)
|
||||
require.Equal(t, tc.kind, got.WorkConnType)
|
||||
require.Equal(t, "mutated", got.RunID)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerNewWorkConnMetadataValidationPreservesPluginContract(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
initialID uint64
|
||||
returnedID uint64
|
||||
response Response
|
||||
pluginErr error
|
||||
wantErr string
|
||||
wantIDErr bool
|
||||
}{
|
||||
{name: "changed generation", initialID: 7, returnedID: 8, wantErr: "from 7 to 8", wantIDErr: true},
|
||||
{name: "legacy cannot inject generation", returnedID: 8, wantErr: "from 0 to 8", wantIDErr: true},
|
||||
{name: "same generation", initialID: 7, returnedID: 7},
|
||||
{name: "omitted generation", initialID: 7},
|
||||
{name: "unchange ignores generation rewrite", initialID: 7, returnedID: 8, response: Response{Unchange: true}},
|
||||
{
|
||||
name: "reject stops before validation", initialID: 7, returnedID: 8,
|
||||
response: Response{Reject: true, RejectReason: "policy rejection"}, wantErr: "policy rejection",
|
||||
},
|
||||
{
|
||||
name: "error stops before validation", initialID: 7, returnedID: 8,
|
||||
pluginErr: errors.New("unavailable"), wantErr: "send NewWorkConn request to plugin error",
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
m := NewManager()
|
||||
m.Register(testPlugin{
|
||||
name: "mutate", ops: map[string]bool{OpNewWorkConn: true},
|
||||
handler: func(context.Context, string, any) (*Response, any, error) {
|
||||
return &tc.response, &NewWorkConnContent{NewWorkConn: msg.NewWorkConn{
|
||||
ControlID: tc.returnedID, WorkConnType: msg.WorkConnTypeDemand,
|
||||
}}, tc.pluginErr
|
||||
},
|
||||
})
|
||||
called := false
|
||||
m.Register(testPlugin{
|
||||
name: "observe", ops: map[string]bool{OpNewWorkConn: true},
|
||||
handler: func(_ context.Context, _ string, raw any) (*Response, any, error) {
|
||||
called = true
|
||||
content := raw.(NewWorkConnContent)
|
||||
require.Equal(t, tc.initialID, content.ControlID)
|
||||
require.Equal(t, msg.WorkConnTypeReserve, content.WorkConnType)
|
||||
return &Response{Unchange: true}, nil, nil
|
||||
},
|
||||
})
|
||||
got, err := m.NewWorkConn(&NewWorkConnContent{NewWorkConn: msg.NewWorkConn{
|
||||
ControlID: tc.initialID, WorkConnType: msg.WorkConnTypeReserve,
|
||||
}})
|
||||
if tc.wantErr != "" {
|
||||
require.ErrorContains(t, err, tc.wantErr)
|
||||
require.Nil(t, got)
|
||||
require.False(t, called)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got)
|
||||
require.True(t, called)
|
||||
}
|
||||
require.Equal(t, tc.wantIDErr, errors.Is(err, ErrNewWorkConnControlIDChanged))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerNewWorkConnChangedContentRequiresPointer(t *testing.T) {
|
||||
m := NewManager()
|
||||
m.Register(testPlugin{
|
||||
name: "invalid", ops: map[string]bool{OpNewWorkConn: true},
|
||||
handler: func(context.Context, string, any) (*Response, any, error) {
|
||||
return &Response{Unchange: false}, NewWorkConnContent{}, nil
|
||||
},
|
||||
})
|
||||
require.Panics(t, func() { _, _ = m.NewWorkConn(&NewWorkConnContent{}) })
|
||||
}
|
||||
|
||||
func TestManagerMutableContentPluginErrorLogLevel(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
+66
-85
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
"net"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -29,7 +28,6 @@ import (
|
||||
"github.com/fatedier/frp/pkg/auth"
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
v1 "github.com/fatedier/frp/pkg/config/v1"
|
||||
pkgerr "github.com/fatedier/frp/pkg/errors"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
plugin "github.com/fatedier/frp/pkg/plugin/server"
|
||||
"github.com/fatedier/frp/pkg/transport"
|
||||
@@ -46,8 +44,6 @@ type ControlID uint64
|
||||
|
||||
var nextControlID atomic.Uint64
|
||||
|
||||
const workConnPoolCapacityOffset = 10
|
||||
|
||||
type controlEntry struct {
|
||||
ctl *Control
|
||||
id ControlID
|
||||
@@ -305,35 +301,46 @@ func (cm *ControlManager) admitVisitorByRunID(runID string, admit func(user, wir
|
||||
// RegisterWorkConn transfers conn to ctl only if ctl is still the current
|
||||
// running generation. On error, ownership remains with the caller.
|
||||
func (cm *ControlManager) RegisterWorkConn(ctl *Control, conn *proxy.WorkConn) error {
|
||||
return cm.registerWorkConn(ctl, conn, workConnRequestLegacy)
|
||||
}
|
||||
|
||||
func (cm *ControlManager) RegisterWorkConnWithType(ctl *Control, conn *proxy.WorkConn, kind workConnRequestKind) error {
|
||||
return cm.registerWorkConn(ctl, conn, kind)
|
||||
}
|
||||
|
||||
func (cm *ControlManager) registerWorkConn(ctl *Control, conn *proxy.WorkConn, kind workConnRequestKind) error {
|
||||
entry, ok := cm.lockCurrentRun(ctl.runID, false)
|
||||
if !ok {
|
||||
cm.mu.RLock()
|
||||
closed := cm.closed
|
||||
cm.mu.RUnlock()
|
||||
if closed {
|
||||
return fmt.Errorf("control manager is closed")
|
||||
return fmt.Errorf("%w: control manager is closed", errWorkConnControlUnavailable)
|
||||
}
|
||||
return fmt.Errorf("client control for run id [%s] is no longer current", ctl.runID)
|
||||
return fmt.Errorf("%w: client control for run id [%s] is no longer current", errWorkConnControlUnavailable, ctl.runID)
|
||||
}
|
||||
defer entry.runMu.Unlock()
|
||||
if entry.ctl != ctl || entry.id != ctl.controlID {
|
||||
return fmt.Errorf("client control for run id [%s] is no longer current", ctl.runID)
|
||||
entry.runMu.Unlock()
|
||||
return fmt.Errorf("%w: client control for run id [%s] is no longer current", errWorkConnControlUnavailable, ctl.runID)
|
||||
}
|
||||
|
||||
ctl.lifecycleMu.Lock()
|
||||
defer ctl.lifecycleMu.Unlock()
|
||||
if ctl.state != controlStateRunning {
|
||||
return fmt.Errorf("client control for run id [%s] is not running", ctl.runID)
|
||||
ctl.lifecycleMu.Unlock()
|
||||
entry.runMu.Unlock()
|
||||
return fmt.Errorf("%w: client control for run id [%s] is not running", errWorkConnControlUnavailable, ctl.runID)
|
||||
}
|
||||
|
||||
select {
|
||||
case ctl.workConnCh <- conn:
|
||||
accepted, err := ctl.registerWorkConnState(conn, kind)
|
||||
ctl.lifecycleMu.Unlock()
|
||||
entry.runMu.Unlock()
|
||||
ctl.requestWorkConnReconcile()
|
||||
if accepted {
|
||||
ctl.xl.Debugf("new work connection registered")
|
||||
return nil
|
||||
default:
|
||||
ctl.xl.Debugf("work connection pool is full, discarding")
|
||||
return fmt.Errorf("work connection pool is full, discarding")
|
||||
}
|
||||
ctl.xl.Debugf("work connection pool is full, discarding")
|
||||
return err
|
||||
}
|
||||
|
||||
func (cm *ControlManager) Close() error {
|
||||
@@ -386,6 +393,8 @@ const (
|
||||
)
|
||||
|
||||
type Control struct {
|
||||
workConnScheduler
|
||||
|
||||
// session context
|
||||
sessionCtx *SessionContext
|
||||
|
||||
@@ -396,15 +405,9 @@ type Control struct {
|
||||
// It provides a channel for sending messages, and you can register handlers to process messages based on their respective types.
|
||||
msgDispatcher *msg.Dispatcher
|
||||
|
||||
// work connections
|
||||
workConnCh chan *proxy.WorkConn
|
||||
|
||||
// proxies in one client
|
||||
proxies map[string]proxy.Proxy
|
||||
|
||||
// pool count
|
||||
poolCount int
|
||||
|
||||
// ports used, for limitations
|
||||
portsUsedNum int
|
||||
|
||||
@@ -452,22 +455,44 @@ func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, erro
|
||||
)
|
||||
}
|
||||
poolCount := int(effectivePoolCount)
|
||||
leaseTimeout := time.Duration(sessionCtx.ServerCfg.UserConnTimeout) * time.Second
|
||||
if leaseTimeout <= 0 {
|
||||
leaseTimeout = 10 * time.Second
|
||||
}
|
||||
doneCh := make(chan struct{})
|
||||
xl := xlog.FromContextSafe(ctx)
|
||||
ctl := &Control{
|
||||
workConnScheduler: newWorkConnScheduler(
|
||||
poolCount,
|
||||
leaseTimeout,
|
||||
xl,
|
||||
doneCh,
|
||||
),
|
||||
sessionCtx: sessionCtx,
|
||||
workConnCh: make(chan *proxy.WorkConn, poolCount+workConnPoolCapacityOffset),
|
||||
proxies: make(map[string]proxy.Proxy),
|
||||
poolCount: poolCount,
|
||||
portsUsedNum: 0,
|
||||
runID: sessionCtx.LoginMsg.RunID,
|
||||
state: controlStateCreated,
|
||||
xl: xlog.FromContextSafe(ctx),
|
||||
xl: xl,
|
||||
ctx: ctx,
|
||||
doneCh: make(chan struct{}),
|
||||
doneCh: doneCh,
|
||||
serverMetrics: metrics.Server,
|
||||
}
|
||||
ctl.lastPing.Store(time.Now())
|
||||
|
||||
ctl.msgDispatcher = msg.NewDispatcher(sessionCtx.Conn)
|
||||
ctl.sendRequest = func(kind workConnRequestKind) error {
|
||||
ctl.lifecycleMu.Lock()
|
||||
controlID := ctl.controlID
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return ctl.msgDispatcher.Send(&msg.ReqWorkConn{
|
||||
WorkConnType: string(kind),
|
||||
ControlID: uint64(controlID),
|
||||
})
|
||||
}
|
||||
ctl.dispatcherDone = func() <-chan struct{} {
|
||||
return ctl.msgDispatcher.Done()
|
||||
}
|
||||
ctl.registerMsgHandlers()
|
||||
ctl.msgTransporter = transport.NewMessageTransporter(ctl.msgDispatcher)
|
||||
return ctl, nil
|
||||
@@ -531,18 +556,23 @@ func (ctl *Control) Close() error {
|
||||
switch ctl.state {
|
||||
case controlStateCreated, controlStatePending:
|
||||
ctl.state = controlStateClosing
|
||||
ctl.takeWorkConnPool()
|
||||
ctl.finishLocked()
|
||||
case controlStateRunning:
|
||||
ctl.state = controlStateClosing
|
||||
ctl.takeWorkConnPool()
|
||||
}
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return ctl.interruptReadAndClose()
|
||||
err := ctl.interruptReadAndClose()
|
||||
ctl.startWorkConnCleanup()
|
||||
return err
|
||||
}
|
||||
|
||||
func (ctl *Control) Replaced(newCtl *Control) {
|
||||
ctl.markReplaced()
|
||||
ctl.xl.Infof("replaced by client [%s] (control ID %d)", newCtl.runID, newCtl.ID())
|
||||
_ = ctl.interruptReadAndClose()
|
||||
ctl.startWorkConnCleanup()
|
||||
}
|
||||
|
||||
// markReplaced returns the transitive predecessor barrier. A pending control
|
||||
@@ -550,24 +580,31 @@ func (ctl *Control) Replaced(newCtl *Control) {
|
||||
// to the replacement. A running control is finished only by its worker.
|
||||
func (ctl *Control) markReplaced() <-chan struct{} {
|
||||
ctl.lifecycleMu.Lock()
|
||||
defer ctl.lifecycleMu.Unlock()
|
||||
|
||||
switch ctl.state {
|
||||
case controlStateCreated:
|
||||
ctl.state = controlStateClosing
|
||||
ctl.takeWorkConnPool()
|
||||
ctl.finishLocked()
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return nil
|
||||
case controlStatePending:
|
||||
barrier := ctl.handoffBarrier
|
||||
ctl.state = controlStateClosing
|
||||
ctl.takeWorkConnPool()
|
||||
ctl.finishLocked()
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return barrier
|
||||
case controlStateRunning:
|
||||
ctl.state = controlStateClosing
|
||||
ctl.takeWorkConnPool()
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return ctl.doneCh
|
||||
case controlStateClosing, controlStateClosed:
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return ctl.doneCh
|
||||
default:
|
||||
ctl.lifecycleMu.Unlock()
|
||||
return ctl.doneCh
|
||||
}
|
||||
}
|
||||
@@ -588,54 +625,6 @@ func (ctl *Control) finishLocked() {
|
||||
close(ctl.doneCh)
|
||||
}
|
||||
|
||||
// When frps get one user connection, we get one work connection from the pool and return it.
|
||||
// If no workConn available in the pool, send message to frpc to get one or more
|
||||
// and wait until it is available.
|
||||
// return an error if wait timeout
|
||||
func (ctl *Control) GetWorkConn() (workConn *proxy.WorkConn, err error) {
|
||||
xl := ctl.xl
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
xl.Errorf("panic error: %v", err)
|
||||
xl.Errorf(string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
|
||||
var ok bool
|
||||
// get a work connection from the pool
|
||||
select {
|
||||
case workConn, ok = <-ctl.workConnCh:
|
||||
if !ok {
|
||||
err = pkgerr.ErrCtlClosed
|
||||
return
|
||||
}
|
||||
xl.Debugf("get work connection from pool")
|
||||
default:
|
||||
// no work connections available in the poll, send message to frpc to get more
|
||||
if err := ctl.msgDispatcher.Send(&msg.ReqWorkConn{}); err != nil {
|
||||
return nil, fmt.Errorf("control is already closed")
|
||||
}
|
||||
|
||||
select {
|
||||
case workConn, ok = <-ctl.workConnCh:
|
||||
if !ok {
|
||||
err = pkgerr.ErrCtlClosed
|
||||
xl.Warnf("no work connections available, %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
case <-time.After(time.Duration(ctl.sessionCtx.ServerCfg.UserConnTimeout) * time.Second):
|
||||
err = fmt.Errorf("timeout trying to get work connection")
|
||||
xl.Warnf("%v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// When we get a work connection from pool, replace it with a new one.
|
||||
_ = ctl.msgDispatcher.Send(&msg.ReqWorkConn{})
|
||||
return
|
||||
}
|
||||
|
||||
func (ctl *Control) heartbeatWorker() {
|
||||
if ctl.sessionCtx.ServerCfg.Transport.HeartbeatTimeout <= 0 {
|
||||
return
|
||||
@@ -686,12 +675,7 @@ func (ctl *Control) worker() {
|
||||
|
||||
go ctl.heartbeatWorker()
|
||||
go ctl.msgDispatcher.Run()
|
||||
go func() {
|
||||
for i := 0; i < ctl.poolCount; i++ {
|
||||
// Ignore the error: it means this control is already closing.
|
||||
_ = ctl.msgDispatcher.Send(&msg.ReqWorkConn{})
|
||||
}
|
||||
}()
|
||||
ctl.requestWorkConnReconcile()
|
||||
|
||||
<-ctl.msgDispatcher.Done()
|
||||
ctl.lifecycleMu.Lock()
|
||||
@@ -701,11 +685,8 @@ func (ctl *Control) worker() {
|
||||
ctl.lifecycleMu.Unlock()
|
||||
_ = ctl.interruptReadAndClose()
|
||||
|
||||
ctl.closeWorkConnPool()
|
||||
ctl.mu.Lock()
|
||||
close(ctl.workConnCh)
|
||||
for workConn := range ctl.workConnCh {
|
||||
workConn.Close()
|
||||
}
|
||||
proxies := ctl.proxies
|
||||
ctl.proxies = make(map[string]proxy.Proxy)
|
||||
ctl.mu.Unlock()
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+25
-2
@@ -36,6 +36,7 @@ import (
|
||||
"github.com/fatedier/frp/pkg/auth"
|
||||
v1 "github.com/fatedier/frp/pkg/config/v1"
|
||||
"github.com/fatedier/frp/pkg/config/v1/validation"
|
||||
pkgerr "github.com/fatedier/frp/pkg/errors"
|
||||
modelmetrics "github.com/fatedier/frp/pkg/metrics"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/nathole"
|
||||
@@ -516,6 +517,13 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
|
||||
acceptedConn.wireProtocol,
|
||||
acceptedConn.clientHelloPresent,
|
||||
); err != nil {
|
||||
if errors.Is(err, errWorkConnPoolFull) || errors.Is(err, errWorkConnControlUnavailable) || errors.Is(err, pkgerr.ErrCtlClosed) {
|
||||
// Pool-full and control-close/replacement races are normal for
|
||||
// reserve or late work connections. Do not turn them into a
|
||||
// misleading StartWorkConn error on frpc.
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
_ = acceptedConn.conn.WriteMsg(&msg.StartWorkConn{
|
||||
Error: util.GenerateResponseErrorString("invalid NewWorkConn", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
|
||||
})
|
||||
@@ -860,11 +868,19 @@ func (svr *Service) RegisterWorkConn(
|
||||
ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
|
||||
if !exist {
|
||||
xl.Warnf("no client control found for run id [%s]", newMsg.RunID)
|
||||
return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
|
||||
return fmt.Errorf("%w: no client control found for run id [%s]", errWorkConnControlUnavailable, newMsg.RunID)
|
||||
}
|
||||
if workWireProtocol != ctl.sessionCtx.WireProtocol {
|
||||
return fmt.Errorf("work connection wire protocol mismatch: got %s want %s", workWireProtocol, ctl.sessionCtx.WireProtocol)
|
||||
}
|
||||
currentControlID := uint64(ctl.ID())
|
||||
requestedControlID := newMsg.ControlID
|
||||
if requestedControlID != 0 && requestedControlID != currentControlID {
|
||||
return fmt.Errorf(
|
||||
"%w: work connection belongs to control generation %d, current generation is %d",
|
||||
errWorkConnControlUnavailable, requestedControlID, currentControlID,
|
||||
)
|
||||
}
|
||||
|
||||
// server plugin hook
|
||||
content := &plugin.NewWorkConnContent{
|
||||
@@ -876,6 +892,9 @@ func (svr *Service) RegisterWorkConn(
|
||||
NewWorkConn: *newMsg,
|
||||
}
|
||||
retContent, err := svr.pluginManager.NewWorkConn(content)
|
||||
if errors.Is(err, plugin.ErrNewWorkConnControlIDChanged) {
|
||||
return fmt.Errorf("%w: %w", errWorkConnControlUnavailable, err)
|
||||
}
|
||||
if err == nil {
|
||||
newMsg = &retContent.NewWorkConn
|
||||
// Check auth.
|
||||
@@ -885,7 +904,11 @@ func (svr *Service) RegisterWorkConn(
|
||||
xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
|
||||
return err
|
||||
}
|
||||
return svr.ctlManager.RegisterWorkConn(ctl, proxy.NewWorkConn(workConn))
|
||||
return svr.ctlManager.RegisterWorkConnWithType(
|
||||
ctl,
|
||||
proxy.NewWorkConn(workConn),
|
||||
normalizeWorkConnRequestKind(workConnRequestKind(newMsg.WorkConnType)),
|
||||
)
|
||||
}
|
||||
|
||||
func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn, wireProtocol string) error {
|
||||
|
||||
+211
-3
@@ -18,6 +18,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net"
|
||||
"net/http"
|
||||
@@ -497,7 +498,7 @@ func TestServiceVisitorAdmissionSerializesReplacement(t *testing.T) {
|
||||
|
||||
func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
ctl, controlConn, err := registerLifecycleTestControl(svr)
|
||||
ctl, controlConn, err := registerLifecycleTestControlWithPoolCount(svr, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
pendingConn := newCountingCloseConn()
|
||||
@@ -524,6 +525,152 @@ func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {
|
||||
require.Equal(t, int64(1), runningConn.closeCount.Load())
|
||||
}
|
||||
|
||||
func TestServiceWorkConnRoutingRejectsStaleControlGeneration(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
ctl, controlConn, err := registerLifecycleTestControl(svr)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
|
||||
waitForSignal(t, controlConn.readStarted, "control reader to start")
|
||||
|
||||
workConn := newCountingCloseConn()
|
||||
workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
|
||||
err = svr.RegisterWorkConn(
|
||||
workMsgConn,
|
||||
&msg.NewWorkConn{RunID: "shared-run", ControlID: uint64(ctl.ID()) + 1},
|
||||
wire.ProtocolV1,
|
||||
false,
|
||||
)
|
||||
require.ErrorContains(t, err, "generation")
|
||||
require.Len(t, ctl.workConnCh, 0)
|
||||
_ = workMsgConn.Close()
|
||||
|
||||
require.NoError(t, ctl.Close())
|
||||
waitForControlDone(t, ctl)
|
||||
}
|
||||
|
||||
func TestServiceNewWorkConnPluginPreservesSchedulingMetadata(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
svr.pluginManager.Register(&partialNewWorkConnPlugin{runID: "plugin-modified"})
|
||||
ctl, controlConn, err := registerLifecycleTestControlWithPoolCount(svr, 2)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = ctl.Close() })
|
||||
policy := &workConnPolicyPlugin{
|
||||
controlID: uint64(ctl.ID()), workConnType: msg.WorkConnTypeReserve, runID: "plugin-modified",
|
||||
}
|
||||
svr.pluginManager.Register(policy)
|
||||
ctl.workConnMu.Lock()
|
||||
// Keep startup deterministic: the explicit leases below already represent
|
||||
// the supply this plugin-routing test needs.
|
||||
ctl.workConnReserveWindow = 0
|
||||
ctl.newWorkConnLeaseLocked(workConnRequestDemand)
|
||||
ctl.newWorkConnLeaseLocked(workConnRequestReserve)
|
||||
ctl.workConnMu.Unlock()
|
||||
require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
|
||||
waitForSignal(t, controlConn.readStarted, "control reader to start")
|
||||
|
||||
workConn := newCountingCloseConn()
|
||||
workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
|
||||
err = svr.RegisterWorkConn(
|
||||
workMsgConn,
|
||||
&msg.NewWorkConn{
|
||||
RunID: "shared-run",
|
||||
WorkConnType: msg.WorkConnTypeReserve,
|
||||
ControlID: uint64(ctl.ID()),
|
||||
},
|
||||
wire.ProtocolV1,
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, policy.called)
|
||||
require.Len(t, ctl.workConnCh, 1)
|
||||
ctl.workConnMu.Lock()
|
||||
demandPending := len(ctl.workConnPending[workConnRequestDemand])
|
||||
reservePending := len(ctl.workConnPending[workConnRequestReserve])
|
||||
ctl.workConnMu.Unlock()
|
||||
require.Equal(t, 1, demandPending)
|
||||
require.Zero(t, reservePending)
|
||||
|
||||
require.NoError(t, ctl.Close())
|
||||
waitForControlDone(t, ctl)
|
||||
require.Equal(t, int64(1), workConn.closeCount.Load())
|
||||
}
|
||||
|
||||
func TestServiceNewWorkConnPluginCannotRewriteControlID(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
ctl, controlConn, err := registerLifecycleTestControl(svr)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = ctl.Close() })
|
||||
svr.pluginManager.Register(&partialNewWorkConnPlugin{controlID: uint64(ctl.ID()) + 1})
|
||||
policy := &workConnPolicyPlugin{controlID: uint64(ctl.ID()), workConnType: msg.WorkConnTypeDemand, runID: "shared-run"}
|
||||
svr.pluginManager.Register(policy)
|
||||
require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
|
||||
waitForSignal(t, controlConn.readStarted, "control reader to start")
|
||||
|
||||
workConn := newCountingCloseConn()
|
||||
workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
|
||||
err = svr.RegisterWorkConn(
|
||||
workMsgConn,
|
||||
&msg.NewWorkConn{
|
||||
RunID: "shared-run",
|
||||
WorkConnType: msg.WorkConnTypeDemand,
|
||||
ControlID: uint64(ctl.ID()),
|
||||
},
|
||||
wire.ProtocolV1,
|
||||
false,
|
||||
)
|
||||
require.ErrorContains(t, err, "plugin changed work connection control generation")
|
||||
require.ErrorIs(t, err, errWorkConnControlUnavailable)
|
||||
require.ErrorIs(t, err, plugin.ErrNewWorkConnControlIDChanged)
|
||||
require.False(t, policy.called)
|
||||
require.Empty(t, ctl.workConnCh)
|
||||
_ = workMsgConn.Close()
|
||||
|
||||
require.NoError(t, ctl.Close())
|
||||
waitForControlDone(t, ctl)
|
||||
}
|
||||
|
||||
func TestServiceReservePoolFullClosesWithoutStartWorkConnError(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
ctl, controlConn, err := registerLifecycleTestControlWithPoolCount(svr, 1)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = ctl.Close() })
|
||||
require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
|
||||
waitForSignal(t, controlConn.readStarted, "control reader to start")
|
||||
|
||||
pooledConns := make([]*proxy.WorkConn, 0, cap(ctl.workConnCh))
|
||||
ctl.workConnMu.Lock()
|
||||
for i := 0; i < cap(ctl.workConnCh); i++ {
|
||||
conn := newCountingCloseConn()
|
||||
pooledConns = append(pooledConns, proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))))
|
||||
ctl.workConnCh <- pooledConns[len(pooledConns)-1]
|
||||
}
|
||||
ctl.workConnMu.Unlock()
|
||||
|
||||
clientConn, serverConn := net.Pipe()
|
||||
t.Cleanup(func() {
|
||||
_ = clientConn.Close()
|
||||
_ = serverConn.Close()
|
||||
})
|
||||
handleDone := make(chan struct{})
|
||||
go func() {
|
||||
svr.handleConnection(context.Background(), serverConn, false)
|
||||
close(handleDone)
|
||||
}()
|
||||
clientMsgConn := msg.NewConn(clientConn, msg.NewV1ReadWriter(clientConn))
|
||||
require.NoError(t, clientConn.SetReadDeadline(time.Now().Add(time.Second)))
|
||||
require.NoError(t, clientMsgConn.WriteMsg(&msg.NewWorkConn{
|
||||
RunID: "shared-run",
|
||||
WorkConnType: msg.WorkConnTypeReserve,
|
||||
}))
|
||||
_, err = clientMsgConn.ReadMsg()
|
||||
require.ErrorIs(t, err, io.EOF)
|
||||
waitForSignal(t, handleDone, "pool-full work connection handler to return")
|
||||
_ = clientConn.Close()
|
||||
|
||||
require.NoError(t, ctl.Close())
|
||||
waitForControlDone(t, ctl)
|
||||
}
|
||||
|
||||
func TestServiceWorkConnRoutingRejectsWireProtocolMismatch(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
ctl, controlConn, err := registerLifecycleTestControl(svr)
|
||||
@@ -570,9 +717,18 @@ func TestServiceWorkConnRoutingClientHelloPolicy(t *testing.T) {
|
||||
svr := newControlTestService(t)
|
||||
controlConn := newDeadlineReadConn()
|
||||
controlMsgConn := msg.NewConn(controlConn, msg.NewV2ReadWriter(controlConn))
|
||||
if tc.errorSubstring == "" {
|
||||
svr.cfg.Transport.MaxPoolCount = 1
|
||||
}
|
||||
ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
|
||||
RunID: "shared-run",
|
||||
ClientID: "client",
|
||||
PoolCount: func() int {
|
||||
if tc.errorSubstring == "" {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}(),
|
||||
ClientSpec: msg.ClientSpec{
|
||||
AlwaysAuthPass: true,
|
||||
},
|
||||
@@ -883,11 +1039,19 @@ func newControlTestService(t *testing.T) *Service {
|
||||
}
|
||||
|
||||
func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, error) {
|
||||
return registerLifecycleTestControlWithPoolCount(svr, 0)
|
||||
}
|
||||
|
||||
func registerLifecycleTestControlWithPoolCount(svr *Service, poolCount int) (*Control, *deadlineReadConn, error) {
|
||||
conn := newDeadlineReadConn()
|
||||
msgConn := msg.NewConn(conn, msg.NewReadWriter(conn, wire.ProtocolV1))
|
||||
if poolCount > 0 {
|
||||
svr.cfg.Transport.MaxPoolCount = int64(poolCount)
|
||||
}
|
||||
ctl, err := svr.RegisterControl(msgConn, &msg.Login{
|
||||
RunID: "shared-run",
|
||||
ClientID: "client",
|
||||
RunID: "shared-run",
|
||||
ClientID: "client",
|
||||
PoolCount: poolCount,
|
||||
ClientSpec: msg.ClientSpec{
|
||||
AlwaysAuthPass: true,
|
||||
},
|
||||
@@ -939,6 +1103,50 @@ type workConnBarrierPlugin struct {
|
||||
resume chan struct{}
|
||||
}
|
||||
|
||||
type partialNewWorkConnPlugin struct {
|
||||
controlID uint64
|
||||
workConnType string
|
||||
runID string
|
||||
}
|
||||
|
||||
type workConnPolicyPlugin struct {
|
||||
controlID uint64
|
||||
workConnType string
|
||||
runID string
|
||||
called bool
|
||||
}
|
||||
|
||||
func (*workConnPolicyPlugin) Name() string { return "work-conn-policy" }
|
||||
|
||||
func (*workConnPolicyPlugin) IsSupport(op string) bool { return op == plugin.OpNewWorkConn }
|
||||
|
||||
func (p *workConnPolicyPlugin) Handle(_ context.Context, _ string, rawContent any) (*plugin.Response, any, error) {
|
||||
p.called = true
|
||||
content := rawContent.(plugin.NewWorkConnContent)
|
||||
if content.ControlID != p.controlID || content.WorkConnType != p.workConnType || content.RunID != p.runID {
|
||||
return &plugin.Response{Reject: true, RejectReason: "unexpected scheduling metadata or mutable run ID"}, nil, nil
|
||||
}
|
||||
return &plugin.Response{Unchange: true}, nil, nil
|
||||
}
|
||||
|
||||
func (*partialNewWorkConnPlugin) Name() string { return "partial-new-work-conn" }
|
||||
|
||||
func (*partialNewWorkConnPlugin) IsSupport(op string) bool { return op == plugin.OpNewWorkConn }
|
||||
|
||||
func (p *partialNewWorkConnPlugin) Handle(
|
||||
_ context.Context,
|
||||
_ string,
|
||||
rawContent any,
|
||||
) (*plugin.Response, any, error) {
|
||||
content := rawContent.(plugin.NewWorkConnContent)
|
||||
content.ControlID = p.controlID
|
||||
content.WorkConnType = p.workConnType
|
||||
if p.runID != "" {
|
||||
content.RunID = p.runID
|
||||
}
|
||||
return &plugin.Response{Unchange: false}, &content, nil
|
||||
}
|
||||
|
||||
func newWorkConnBarrierPlugin() *workConnBarrierPlugin {
|
||||
return &workConnBarrierPlugin{
|
||||
entered: make(chan struct{}),
|
||||
|
||||
@@ -0,0 +1,739 @@
|
||||
// Copyright 2026 The frp Authors
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
pkgerr "github.com/fatedier/frp/pkg/errors"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
"github.com/fatedier/frp/server/proxy"
|
||||
)
|
||||
|
||||
const workConnPoolCapacityOffset = 10
|
||||
|
||||
const (
|
||||
workConnReserveWindowMax = 8
|
||||
workConnDemandWindow = 64
|
||||
workConnLegacyWindow = 64
|
||||
|
||||
// Canceled waiters remain as lazy tombstones so cancellation stays O(1).
|
||||
// Compact only after enough tombstones have accumulated and comprise at
|
||||
// least half the queue, making the occasional O(n) copy O(1) amortized.
|
||||
workConnWaiterCompactMinTombstones = 64
|
||||
)
|
||||
|
||||
var (
|
||||
errWorkConnPoolFull = errors.New("work connection pool is full, discarding")
|
||||
errWorkConnControlUnavailable = errors.New("work connection control unavailable")
|
||||
)
|
||||
|
||||
type workConnRequestKind string
|
||||
|
||||
const (
|
||||
workConnRequestDemand workConnRequestKind = msg.WorkConnTypeDemand
|
||||
workConnRequestReserve workConnRequestKind = msg.WorkConnTypeReserve
|
||||
workConnRequestLegacy workConnRequestKind = "legacy"
|
||||
)
|
||||
|
||||
type workConnWaiter struct {
|
||||
resultCh chan workConnResult
|
||||
active bool
|
||||
handoffConn *proxy.WorkConn
|
||||
handoffClosed bool
|
||||
handoffElement *list.Element
|
||||
}
|
||||
|
||||
type workConnResult struct {
|
||||
conn *proxy.WorkConn
|
||||
err error
|
||||
}
|
||||
|
||||
type workConnLease struct {
|
||||
kind workConnRequestKind
|
||||
timer *time.Timer
|
||||
|
||||
active bool
|
||||
dispatched bool
|
||||
late bool
|
||||
}
|
||||
|
||||
// workConnScheduler owns all work-connection supply state. It deliberately
|
||||
// communicates with Control through callbacks for request dispatch and
|
||||
// dispatcher shutdown; it never acquires Control lifecycle or manager locks.
|
||||
type workConnScheduler struct {
|
||||
workConnCh chan *proxy.WorkConn
|
||||
|
||||
workConnMu sync.Mutex
|
||||
workConnWaiters []*workConnWaiter
|
||||
workConnActiveWaiters int
|
||||
workConnWaiterTombstones int
|
||||
// Elements retain registration order for cleanup; each waiter owns its
|
||||
// element so consumption removes it in O(1), even during large bursts.
|
||||
workConnHandoffs list.List
|
||||
workConnPending map[workConnRequestKind][]*workConnLease
|
||||
workConnLate map[workConnRequestKind][]*workConnLease
|
||||
workConnPoolClosed bool
|
||||
|
||||
workConnReconcileCh chan struct{}
|
||||
workConnReconcileStop chan struct{}
|
||||
workConnReconcileDone chan struct{}
|
||||
workConnReconcileOnce sync.Once
|
||||
workConnReconcileStopOnce sync.Once
|
||||
|
||||
workConnCleanupMu sync.Mutex
|
||||
workConnCleanupPending []*proxy.WorkConn
|
||||
workConnCleanupReady bool
|
||||
workConnCleanupStarted bool
|
||||
workConnCleanupDone chan struct{}
|
||||
|
||||
workConnLegacyMode bool
|
||||
workConnReserveWindow int
|
||||
workConnLeaseTimeout time.Duration
|
||||
poolCount int
|
||||
|
||||
// sendRequest is called only after workConnMu has been released. The
|
||||
// callback captures the Control's immutable generation metadata and sends
|
||||
// the wire message without exposing lifecycle state to the scheduler.
|
||||
sendRequest func(workConnRequestKind) error
|
||||
dispatcherDone func() <-chan struct{}
|
||||
logger *xlog.Logger
|
||||
doneCh <-chan struct{}
|
||||
}
|
||||
|
||||
func newWorkConnScheduler(poolCount int, leaseTimeout time.Duration, logger *xlog.Logger, doneCh <-chan struct{}) workConnScheduler {
|
||||
return workConnScheduler{
|
||||
workConnCh: make(chan *proxy.WorkConn, poolCount+workConnPoolCapacityOffset),
|
||||
poolCount: poolCount,
|
||||
workConnPending: make(map[workConnRequestKind][]*workConnLease),
|
||||
workConnLate: make(map[workConnRequestKind][]*workConnLease),
|
||||
workConnReconcileCh: make(chan struct{}, 1),
|
||||
workConnReconcileStop: make(chan struct{}),
|
||||
workConnReconcileDone: make(chan struct{}),
|
||||
workConnCleanupDone: make(chan struct{}),
|
||||
workConnLegacyMode: true,
|
||||
workConnReserveWindow: min(poolCount, workConnReserveWindowMax),
|
||||
workConnLeaseTimeout: leaseTimeout,
|
||||
logger: logger,
|
||||
doneCh: doneCh,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeWorkConnRequestKind(kind workConnRequestKind) workConnRequestKind {
|
||||
if kind == workConnRequestDemand || kind == workConnRequestReserve {
|
||||
return kind
|
||||
}
|
||||
return workConnRequestLegacy
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) workConnLeaseTimeoutValue() time.Duration {
|
||||
if s.workConnLeaseTimeout > 0 {
|
||||
return s.workConnLeaseTimeout
|
||||
}
|
||||
return 10 * time.Second
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) startWorkConnReconciler() {
|
||||
s.workConnReconcileOnce.Do(func() {
|
||||
go func() {
|
||||
defer close(s.workConnReconcileDone)
|
||||
for {
|
||||
select {
|
||||
case <-s.workConnReconcileCh:
|
||||
s.reconcileWorkConnRequests()
|
||||
case <-s.workConnReconcileStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// requestWorkConnReconcile coalesces refill requests into one bounded trigger.
|
||||
// The reconciler may block on dispatcher backpressure, but it owns no control
|
||||
// lifecycle or generation locks and is stopped with the control pool.
|
||||
func (s *workConnScheduler) requestWorkConnReconcile() {
|
||||
s.workConnMu.Lock()
|
||||
closed := s.workConnPoolClosed
|
||||
s.workConnMu.Unlock()
|
||||
if closed {
|
||||
return
|
||||
}
|
||||
|
||||
s.startWorkConnReconciler()
|
||||
select {
|
||||
case <-s.workConnReconcileStop:
|
||||
return
|
||||
case s.workConnReconcileCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) activeWaiterCountLocked() int {
|
||||
return s.workConnActiveWaiters
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) appendWorkConnWaiterLocked(waiter *workConnWaiter) {
|
||||
s.workConnWaiters = append(s.workConnWaiters, waiter)
|
||||
if waiter.active {
|
||||
s.workConnActiveWaiters++
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) popWorkConnWaiterLocked() *workConnWaiter {
|
||||
for len(s.workConnWaiters) > 0 {
|
||||
waiter := s.workConnWaiters[0]
|
||||
s.workConnWaiters[0] = nil
|
||||
s.workConnWaiters = s.workConnWaiters[1:]
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
if len(s.workConnWaiters) == 0 {
|
||||
s.workConnWaiters = nil
|
||||
}
|
||||
return waiter
|
||||
}
|
||||
if s.workConnWaiterTombstones > 0 {
|
||||
s.workConnWaiterTombstones--
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = nil
|
||||
s.workConnWaiterTombstones = 0
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) compactWorkConnWaitersLocked() {
|
||||
if s.workConnWaiterTombstones == 0 {
|
||||
return
|
||||
}
|
||||
activeWaiters := make([]*workConnWaiter, 0, s.workConnActiveWaiters)
|
||||
for _, waiter := range s.workConnWaiters {
|
||||
if waiter.active {
|
||||
activeWaiters = append(activeWaiters, waiter)
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = activeWaiters
|
||||
s.workConnWaiterTombstones = 0
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) shouldCompactWorkConnWaitersLocked() bool {
|
||||
tombstones := s.workConnWaiterTombstones
|
||||
if tombstones < workConnWaiterCompactMinTombstones {
|
||||
return false
|
||||
}
|
||||
return tombstones*2 >= len(s.workConnWaiters)
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) removeWorkConnHandoffLocked(target *workConnWaiter) {
|
||||
if target.handoffElement != nil {
|
||||
s.workConnHandoffs.Remove(target.handoffElement)
|
||||
target.handoffElement = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) removeLeaseLocked(leases map[workConnRequestKind][]*workConnLease, target *workConnLease) {
|
||||
items := leases[target.kind]
|
||||
for i, lease := range items {
|
||||
if lease == target {
|
||||
items = append(items[:i], items[i+1:]...)
|
||||
if len(items) == 0 {
|
||||
delete(leases, target.kind)
|
||||
} else {
|
||||
leases[target.kind] = items
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) hasLeaseLocked(leases map[workConnRequestKind][]*workConnLease, target *workConnLease) bool {
|
||||
return slices.Contains(leases[target.kind], target)
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) popLeaseLocked(leases map[workConnRequestKind][]*workConnLease, kind workConnRequestKind) *workConnLease {
|
||||
items := leases[kind]
|
||||
for len(items) > 0 {
|
||||
lease := items[0]
|
||||
items = items[1:]
|
||||
if len(items) == 0 {
|
||||
delete(leases, kind)
|
||||
} else {
|
||||
leases[kind] = items
|
||||
}
|
||||
if lease.active {
|
||||
return lease
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) leaseCountLocked(leases map[workConnRequestKind][]*workConnLease, kind workConnRequestKind) int {
|
||||
count := 0
|
||||
for _, lease := range leases[kind] {
|
||||
if lease.active {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) workConnSupplyGapLocked() bool {
|
||||
target := s.activeWaiterCountLocked() + s.poolCount
|
||||
pending := s.leaseCountLocked(s.workConnPending, workConnRequestDemand) +
|
||||
s.leaseCountLocked(s.workConnPending, workConnRequestReserve)
|
||||
return len(s.workConnCh)+pending < target
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) consumeWorkConnLeaseLocked(kind workConnRequestKind) *workConnLease {
|
||||
kind = normalizeWorkConnRequestKind(kind)
|
||||
s.workConnLegacyMode = kind == workConnRequestLegacy
|
||||
kinds := []workConnRequestKind{kind}
|
||||
if s.workConnLegacyMode {
|
||||
// Untyped responses satisfy pending demand before pending reserve.
|
||||
kinds = []workConnRequestKind{workConnRequestDemand, workConnRequestReserve}
|
||||
}
|
||||
// There is no request ID: same-kind replies are interchangeable. A timely
|
||||
// retry must retire pending supply before late leases, otherwise its reply
|
||||
// is discarded while the retry remains pending and the pool never recovers.
|
||||
// For legacy clients, all pending kinds take precedence over late leases.
|
||||
for _, leases := range []map[workConnRequestKind][]*workConnLease{s.workConnPending, s.workConnLate} {
|
||||
for _, candidate := range kinds {
|
||||
if lease := s.popLeaseLocked(leases, candidate); lease != nil {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
return lease
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) expireLateWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if !lease.active || !lease.late {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
lease.active = false
|
||||
s.removeLeaseLocked(s.workConnLate, lease)
|
||||
s.workConnMu.Unlock()
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) expireWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if !lease.active || lease.late {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
s.removeLeaseLocked(s.workConnPending, lease)
|
||||
lease.late = true
|
||||
s.workConnLate[lease.kind] = append(s.workConnLate[lease.kind], lease)
|
||||
lease.timer = time.AfterFunc(s.workConnLeaseTimeoutValue(), func() {
|
||||
s.expireLateWorkConnLease(lease)
|
||||
})
|
||||
s.workConnMu.Unlock()
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) newWorkConnLeaseLocked(kind workConnRequestKind) *workConnLease {
|
||||
lease := &workConnLease{
|
||||
kind: kind,
|
||||
active: true,
|
||||
}
|
||||
s.workConnPending[kind] = append(s.workConnPending[kind], lease)
|
||||
return lease
|
||||
}
|
||||
|
||||
// startWorkConnLeaseTimer starts the response timeout only after the request
|
||||
// has been accepted by Dispatcher.Send. A response may race this call; the
|
||||
// active/pending check under workConnMu prevents installing a timer on a lease
|
||||
// that was already consumed, failed, or closed.
|
||||
// Send acknowledges queue admission, not WriteMsg completion. The dispatcher
|
||||
// has no per-message write acknowledgement, so this timeout includes queue time.
|
||||
func (s *workConnScheduler) startWorkConnLeaseTimer(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if lease.active && !lease.dispatched && !lease.late && lease.timer == nil && s.hasLeaseLocked(s.workConnPending, lease) {
|
||||
lease.dispatched = true
|
||||
lease.timer = time.AfterFunc(s.workConnLeaseTimeoutValue(), func() {
|
||||
s.expireWorkConnLease(lease)
|
||||
})
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) registerWorkConn(conn *proxy.WorkConn, kind workConnRequestKind) (bool, error) {
|
||||
accepted, err := s.registerWorkConnState(conn, kind)
|
||||
s.reconcileWorkConnRequests()
|
||||
return accepted, err
|
||||
}
|
||||
|
||||
// registerWorkConnState commits work connection ownership and lease accounting
|
||||
// without sending refill requests. Callers that hold lifecycle or generation
|
||||
// locks must release them before invoking reconcileWorkConnRequests.
|
||||
func (s *workConnScheduler) registerWorkConnState(conn *proxy.WorkConn, kind workConnRequestKind) (bool, error) {
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return false, pkgerr.ErrCtlClosed
|
||||
}
|
||||
s.consumeWorkConnLeaseLocked(kind)
|
||||
if waiter := s.popWorkConnWaiterLocked(); waiter != nil {
|
||||
// A reserve response may satisfy a demand waiter, but any other
|
||||
// pending lease is still an in-flight source of supply until it really
|
||||
// times out. Keep it accounted for instead of creating a replacement.
|
||||
waiter.handoffConn = conn
|
||||
waiter.handoffElement = s.workConnHandoffs.PushBack(waiter)
|
||||
waiter.resultCh <- workConnResult{conn: conn}
|
||||
s.workConnMu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
// Any response that did not satisfy a waiter must still prove a real supply
|
||||
// gap before entering idle. This covers unmatched responses, demand leases
|
||||
// whose waiter timed out, and late leases whose replacement is already
|
||||
// pending. A matching lease has been consumed above, so pending includes
|
||||
// only the remaining supply credits; replies cannot identify their own lease.
|
||||
if !s.workConnSupplyGapLocked() {
|
||||
s.workConnMu.Unlock()
|
||||
return false, errWorkConnPoolFull
|
||||
}
|
||||
|
||||
select {
|
||||
case s.workConnCh <- conn:
|
||||
s.workConnMu.Unlock()
|
||||
return true, nil
|
||||
default:
|
||||
s.workConnMu.Unlock()
|
||||
return false, errWorkConnPoolFull
|
||||
}
|
||||
}
|
||||
|
||||
// consumeWorkConnWaiterResult transfers a direct waiter handoff to its
|
||||
// consumer only while the control pool is still open. A close/replacement may
|
||||
// have already detached the connection for cleanup even though the buffered
|
||||
// result notification remains readable; in that case the caller must never
|
||||
// receive the stale connection.
|
||||
func (s *workConnScheduler) consumeWorkConnWaiterResult(waiter *workConnWaiter, result workConnResult) (*proxy.WorkConn, error) {
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed || waiter.handoffClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
conn := waiter.handoffConn
|
||||
if conn != nil {
|
||||
waiter.handoffConn = nil
|
||||
s.removeWorkConnHandoffLocked(waiter)
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
if result.err != nil {
|
||||
return nil, result.err
|
||||
}
|
||||
if conn == nil {
|
||||
conn = result.conn
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) cancelWorkConnWaiter(waiter *workConnWaiter) bool {
|
||||
canceled := false
|
||||
s.workConnMu.Lock()
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
s.workConnWaiterTombstones++
|
||||
if s.shouldCompactWorkConnWaitersLocked() {
|
||||
s.compactWorkConnWaitersLocked()
|
||||
}
|
||||
canceled = true
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
if canceled {
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
return canceled
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) startWorkConnCleanup() {
|
||||
s.workConnCleanupMu.Lock()
|
||||
if s.workConnCleanupStarted || !s.workConnCleanupReady {
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
conns := s.workConnCleanupPending
|
||||
s.workConnCleanupPending = nil
|
||||
s.workConnCleanupStarted = true
|
||||
s.workConnCleanupMu.Unlock()
|
||||
|
||||
if len(conns) == 0 {
|
||||
close(s.workConnCleanupDone)
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
for _, conn := range conns {
|
||||
_ = conn.Close()
|
||||
}
|
||||
close(s.workConnCleanupDone)
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) waitWorkConnCleanup() {
|
||||
s.workConnCleanupMu.Lock()
|
||||
started := s.workConnCleanupStarted
|
||||
s.workConnCleanupMu.Unlock()
|
||||
if started {
|
||||
<-s.workConnCleanupDone
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) takeWorkConnPool() {
|
||||
var conns []*proxy.WorkConn
|
||||
// Serialize pool extraction with cleanup startup. A concurrent worker may
|
||||
// observe the pool as closed before this caller has queued its streams;
|
||||
// keeping the cleanup gate held makes that ordering explicit and prevents
|
||||
// a premature workConnCleanupDone close from losing the extracted conns.
|
||||
s.workConnCleanupMu.Lock()
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
s.workConnPoolClosed = true
|
||||
s.workConnReconcileStopOnce.Do(func() { close(s.workConnReconcileStop) })
|
||||
for _, waiter := range s.workConnWaiters {
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
waiter.resultCh <- workConnResult{err: pkgerr.ErrCtlClosed}
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = nil
|
||||
s.workConnActiveWaiters = 0
|
||||
s.workConnWaiterTombstones = 0
|
||||
for _, leases := range s.workConnPending {
|
||||
for _, lease := range leases {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, leases := range s.workConnLate {
|
||||
for _, lease := range leases {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
for element := s.workConnHandoffs.Front(); element != nil; element = element.Next() {
|
||||
waiter := element.Value.(*workConnWaiter)
|
||||
if waiter.handoffConn != nil {
|
||||
conns = append(conns, waiter.handoffConn)
|
||||
waiter.handoffConn = nil
|
||||
}
|
||||
waiter.handoffClosed = true
|
||||
waiter.handoffElement = nil
|
||||
}
|
||||
s.workConnHandoffs.Init()
|
||||
for {
|
||||
select {
|
||||
case conn := <-s.workConnCh:
|
||||
conns = append(conns, conn)
|
||||
default:
|
||||
close(s.workConnCh)
|
||||
s.workConnMu.Unlock()
|
||||
s.workConnCleanupPending = append(s.workConnCleanupPending, conns...)
|
||||
s.workConnCleanupReady = true
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) closeWorkConnPool() {
|
||||
s.takeWorkConnPool()
|
||||
s.startWorkConnCleanup()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) reconcileWorkConnRequests() {
|
||||
type request struct {
|
||||
lease *workConnLease
|
||||
kind workConnRequestKind
|
||||
}
|
||||
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
var requests []request
|
||||
demandPending := s.leaseCountLocked(s.workConnPending, workConnRequestDemand)
|
||||
reservePending := s.leaseCountLocked(s.workConnPending, workConnRequestReserve)
|
||||
activeWaiters := s.activeWaiterCountLocked()
|
||||
demandNeed := activeWaiters - demandPending
|
||||
if demandNeed > 0 {
|
||||
budget := workConnDemandWindow - demandPending
|
||||
if s.workConnLegacyMode {
|
||||
legacyBudget := workConnLegacyWindow - demandPending - reservePending
|
||||
if budget > legacyBudget {
|
||||
budget = legacyBudget
|
||||
}
|
||||
}
|
||||
if budget < 0 {
|
||||
budget = 0
|
||||
}
|
||||
if demandNeed > budget {
|
||||
demandNeed = budget
|
||||
}
|
||||
for i := 0; i < demandNeed; i++ {
|
||||
requests = append(requests, request{lease: s.newWorkConnLeaseLocked(workConnRequestDemand), kind: workConnRequestDemand})
|
||||
}
|
||||
demandPending += demandNeed
|
||||
}
|
||||
|
||||
// Keep the invariant available + pending >= waiters + poolCount, while
|
||||
// limiting only reserve requests to the sliding window. Late leases do not
|
||||
// consume retry budget; their responses are admitted separately above only
|
||||
// when this projected supply still has a real gap.
|
||||
reserveNeed := max(activeWaiters+s.poolCount-len(s.workConnCh)-demandPending-reservePending, 0)
|
||||
reserveBudget := s.workConnReserveWindow - reservePending
|
||||
if s.workConnLegacyMode {
|
||||
legacyBudget := workConnLegacyWindow - demandPending - reservePending
|
||||
if reserveBudget > legacyBudget {
|
||||
reserveBudget = legacyBudget
|
||||
}
|
||||
}
|
||||
if reserveBudget < 0 {
|
||||
reserveBudget = 0
|
||||
}
|
||||
if reserveNeed > reserveBudget {
|
||||
reserveNeed = reserveBudget
|
||||
}
|
||||
for i := 0; i < reserveNeed; i++ {
|
||||
requests = append(requests, request{lease: s.newWorkConnLeaseLocked(workConnRequestReserve), kind: workConnRequestReserve})
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
|
||||
for _, req := range requests {
|
||||
if err := s.sendRequest(req.kind); err != nil {
|
||||
s.failWorkConnLease(req.lease)
|
||||
if s.dispatcherDone != nil {
|
||||
select {
|
||||
case <-s.dispatcherDone():
|
||||
// Wake waiters immediately when the control transport is already
|
||||
// closed; otherwise they would wait until UserConnTimeout.
|
||||
s.closeWorkConnPool()
|
||||
default:
|
||||
}
|
||||
}
|
||||
} else {
|
||||
s.startWorkConnLeaseTimer(req.lease)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) failWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if lease.active && !lease.late {
|
||||
lease.active = false
|
||||
s.removeLeaseLocked(s.workConnPending, lease)
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
|
||||
// When frps get one user connection, we get one work connection from the pool and return it.
|
||||
// If no workConn available in the pool, send message to frpc to get one or more
|
||||
// and wait until it is available.
|
||||
// return an error if wait timeout
|
||||
func (s *workConnScheduler) GetWorkConn() (workConn *proxy.WorkConn, err error) {
|
||||
xl := s.logger
|
||||
started := time.Now()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
xl.Errorf("panic error: %v", recovered)
|
||||
xl.Errorf(string(debug.Stack()))
|
||||
workConn = nil
|
||||
err = pkgerr.ErrCtlClosed
|
||||
}
|
||||
}()
|
||||
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
// Get an idle connection first.
|
||||
select {
|
||||
case workConn = <-s.workConnCh:
|
||||
s.workConnMu.Unlock()
|
||||
xl.Debugf("get work connection from pool")
|
||||
s.requestWorkConnReconcile()
|
||||
return workConn, nil
|
||||
default:
|
||||
}
|
||||
|
||||
waiter := &workConnWaiter{
|
||||
resultCh: make(chan workConnResult, 1),
|
||||
active: true,
|
||||
}
|
||||
s.appendWorkConnWaiterLocked(waiter)
|
||||
s.workConnMu.Unlock()
|
||||
|
||||
// Pool hits need no timer. On a miss, include time already spent acquiring
|
||||
// the pool lock and install the timer before requesting a refill.
|
||||
timer := time.NewTimer(s.workConnLeaseTimeoutValue() - time.Since(started))
|
||||
defer timer.Stop()
|
||||
s.requestWorkConnReconcile()
|
||||
select {
|
||||
case result := <-waiter.resultCh:
|
||||
return s.consumeWorkConnWaiterResult(waiter, result)
|
||||
case <-timer.C:
|
||||
if !s.cancelWorkConnWaiter(waiter) {
|
||||
select {
|
||||
case result := <-waiter.resultCh:
|
||||
return s.consumeWorkConnWaiterResult(waiter, result)
|
||||
default:
|
||||
}
|
||||
}
|
||||
s.workConnMu.Lock()
|
||||
closed := s.workConnPoolClosed
|
||||
s.workConnMu.Unlock()
|
||||
if closed {
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
err = fmt.Errorf("timeout trying to get work connection")
|
||||
xl.Warnf("%v", err)
|
||||
return nil, err
|
||||
case <-s.doneCh:
|
||||
// doneCh is terminal for this control generation. Do not consume a
|
||||
// buffered successful handoff here: takeWorkConnPool detaches that
|
||||
// connection for cleanup, and returning it would leak it past shutdown.
|
||||
s.cancelWorkConnWaiter(waiter)
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
// Copyright 2026 The frp Authors
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
pkgerr "github.com/fatedier/frp/pkg/errors"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
"github.com/fatedier/frp/server/proxy"
|
||||
)
|
||||
|
||||
func TestWorkConnReplacementRestoresPoolWithoutOriginalResponse(t *testing.T) {
|
||||
for _, kind := range []workConnRequestKind{workConnRequestReserve, workConnRequestLegacy} {
|
||||
t.Run(string(kind), func(t *testing.T) {
|
||||
ctl, serverConn, clientConn := newWorkConnSchedulingControl(t, 1)
|
||||
ctl.workConnLeaseTimeout = time.Minute
|
||||
t.Cleanup(func() {
|
||||
ctl.closeWorkConnPool()
|
||||
_ = serverConn.Close()
|
||||
_ = clientConn.Close()
|
||||
})
|
||||
ctl.reconcileWorkConnRequests()
|
||||
readWorkConnRequest(t, clientConn)
|
||||
// The original request never returns. Trigger its timeout transition
|
||||
// without sleeping, then deliver only the replacement's response.
|
||||
moveReserveWorkConnLeaseToLateForTest(t, ctl)
|
||||
ctl.reconcileWorkConnRequests()
|
||||
readWorkConnRequest(t, clientConn)
|
||||
conn := newCountingCloseConn()
|
||||
workConn := proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn)))
|
||||
accepted, err := ctl.registerWorkConnState(workConn, kind)
|
||||
if !accepted {
|
||||
_ = workConn.Close()
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.True(t, accepted)
|
||||
ctl.workConnMu.Lock()
|
||||
pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve)
|
||||
late := ctl.leaseCountLocked(ctl.workConnLate, workConnRequestReserve)
|
||||
idle := len(ctl.workConnCh)
|
||||
ctl.workConnMu.Unlock()
|
||||
require.Zero(t, pending, "the successful retry must not time out and retry again")
|
||||
require.Equal(t, 1, late)
|
||||
require.Equal(t, 1, idle)
|
||||
got, err := ctl.GetWorkConn()
|
||||
if got != nil {
|
||||
t.Cleanup(func() { _ = got.Close() })
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Same(t, workConn, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkConnLeaseMatchingPrefersPendingSupply(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
responseKind workConnRequestKind
|
||||
pendingKind workConnRequestKind
|
||||
lateKind workConnRequestKind
|
||||
}{
|
||||
{"typed demand", workConnRequestDemand, workConnRequestDemand, workConnRequestDemand},
|
||||
{"typed reserve", workConnRequestReserve, workConnRequestReserve, workConnRequestReserve},
|
||||
{"legacy demand", workConnRequestLegacy, workConnRequestDemand, workConnRequestDemand},
|
||||
{"legacy reserve before late demand", workConnRequestLegacy, workConnRequestReserve, workConnRequestDemand},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := newWorkConnScheduler(1, time.Minute, xlog.New(), nil)
|
||||
t.Cleanup(s.closeWorkConnPool)
|
||||
s.workConnMu.Lock()
|
||||
late := &workConnLease{kind: tc.lateKind, active: true, late: true}
|
||||
s.workConnLate[tc.lateKind] = []*workConnLease{late}
|
||||
pending := s.newWorkConnLeaseLocked(tc.pendingKind)
|
||||
first := s.consumeWorkConnLeaseLocked(tc.responseKind)
|
||||
lateStillActive := late.active
|
||||
second := s.consumeWorkConnLeaseLocked(tc.responseKind)
|
||||
s.workConnMu.Unlock()
|
||||
require.Same(t, pending, first)
|
||||
require.True(t, lateStillActive)
|
||||
require.Same(t, late, second)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkConnQueuedLeaseTimerUsesAdmissionBoundary(t *testing.T) {
|
||||
conn := newBlockingWriteConn()
|
||||
dispatcher := msg.NewDispatcher(msg.NewConn(conn, msg.NewV1ReadWriter(conn)))
|
||||
s := newWorkConnScheduler(1, time.Minute, xlog.New(), nil)
|
||||
s.sendRequest = func(kind workConnRequestKind) error {
|
||||
return dispatcher.Send(&msg.ReqWorkConn{WorkConnType: string(kind)})
|
||||
}
|
||||
s.dispatcherDone = func() <-chan struct{} { return dispatcher.Done() }
|
||||
t.Cleanup(func() {
|
||||
s.closeWorkConnPool()
|
||||
_ = conn.Close()
|
||||
})
|
||||
dispatcher.Run()
|
||||
require.NoError(t, dispatcher.Send(&msg.Ping{}))
|
||||
waitForSignal(t, conn.writeStarted, "initial dispatcher write to block")
|
||||
|
||||
// Unlike the full-queue test, Send can enqueue here. This records the
|
||||
// remaining API boundary: a timer starts while the writer is still blocked
|
||||
// on Ping, because Dispatcher offers no per-message write acknowledgement.
|
||||
s.reconcileWorkConnRequests()
|
||||
s.workConnMu.Lock()
|
||||
leases := s.workConnPending[workConnRequestReserve]
|
||||
pending := len(leases)
|
||||
timerInstalled := pending == 1 && leases[0].timer != nil && leases[0].dispatched
|
||||
s.workConnMu.Unlock()
|
||||
require.Equal(t, 1, pending)
|
||||
require.True(t, timerInstalled)
|
||||
select {
|
||||
case <-conn.allowWrite:
|
||||
t.Fatal("writer was unexpectedly released")
|
||||
default:
|
||||
}
|
||||
_ = conn.Close()
|
||||
waitForSignal(t, dispatcher.Done(), "dispatcher shutdown")
|
||||
}
|
||||
|
||||
type orderedWorkConnClose struct {
|
||||
*countingCloseConn
|
||||
id int
|
||||
closed chan<- int
|
||||
}
|
||||
|
||||
func (c *orderedWorkConnClose) Close() error {
|
||||
c.closed <- c.id
|
||||
return c.countingCloseConn.Close()
|
||||
}
|
||||
|
||||
func TestWorkConnHandoffsPreserveFIFOAndCleanupOrder(t *testing.T) {
|
||||
s := newWorkConnScheduler(0, time.Minute, xlog.New(), nil)
|
||||
t.Cleanup(s.closeWorkConnPool)
|
||||
const count = 6
|
||||
closed := make(chan int, count)
|
||||
waiters := make([]*workConnWaiter, count)
|
||||
conns := make([]*orderedWorkConnClose, count)
|
||||
results := make([]workConnResult, count)
|
||||
for i := range waiters {
|
||||
waiters[i] = &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true}
|
||||
s.workConnMu.Lock()
|
||||
s.appendWorkConnWaiterLocked(waiters[i])
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
for i := range waiters {
|
||||
conns[i] = &orderedWorkConnClose{countingCloseConn: newCountingCloseConn(), id: i, closed: closed}
|
||||
workConn := proxy.NewWorkConn(msg.NewConn(conns[i], msg.NewV1ReadWriter(conns[i])))
|
||||
accepted, err := s.registerWorkConnState(workConn, workConnRequestDemand)
|
||||
require.NoError(t, err)
|
||||
require.True(t, accepted)
|
||||
results[i] = waitForResult(t, waiters[i].resultCh, "FIFO waiter delivery")
|
||||
require.Same(t, workConn, results[i].conn)
|
||||
}
|
||||
// Consumers may run out of order. Removing middle/head/tail nodes must
|
||||
// preserve the remaining registration order used by shutdown cleanup.
|
||||
for _, i := range []int{3, 0, 5} {
|
||||
conn, err := s.consumeWorkConnWaiterResult(waiters[i], results[i])
|
||||
require.NoError(t, err)
|
||||
require.Same(t, results[i].conn, conn)
|
||||
require.NoError(t, conn.Close())
|
||||
require.Equal(t, i, waitForResult(t, closed, "consumer close"))
|
||||
}
|
||||
s.closeWorkConnPool()
|
||||
waitForSignal(t, s.workConnCleanupDone, "ordered handoff cleanup")
|
||||
for _, i := range []int{1, 2, 4} {
|
||||
require.Equal(t, i, waitForResult(t, closed, "cleanup close"))
|
||||
conn, err := s.consumeWorkConnWaiterResult(waiters[i], results[i])
|
||||
require.Nil(t, conn)
|
||||
require.ErrorIs(t, err, pkgerr.ErrCtlClosed)
|
||||
}
|
||||
for _, conn := range conns {
|
||||
require.Equal(t, int64(1), conn.closeCount.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkConnConcurrentHandoffConsumptionAndClose(t *testing.T) {
|
||||
s := newWorkConnScheduler(0, time.Minute, xlog.New(), nil)
|
||||
t.Cleanup(s.closeWorkConnPool)
|
||||
const count = 512
|
||||
waiters := make([]*workConnWaiter, count)
|
||||
conns := make([]*countingCloseConn, count)
|
||||
results := make([]workConnResult, count)
|
||||
for i := range waiters {
|
||||
waiters[i] = &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true}
|
||||
s.workConnMu.Lock()
|
||||
s.appendWorkConnWaiterLocked(waiters[i])
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
for i := range waiters {
|
||||
conns[i] = newCountingCloseConn()
|
||||
accepted, err := s.registerWorkConnState(
|
||||
proxy.NewWorkConn(msg.NewConn(conns[i], msg.NewV1ReadWriter(conns[i]))), workConnRequestDemand,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, accepted)
|
||||
results[i] = waitForResult(t, waiters[i].resultCh, "burst waiter delivery")
|
||||
}
|
||||
// Guarantee coverage of both ownership outcomes in addition to the race:
|
||||
// one quarter claims before close, one quarter only attempts after close.
|
||||
for i := range count / 4 {
|
||||
conn, err := s.consumeWorkConnWaiterResult(waiters[i], results[i])
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, conn.Close())
|
||||
}
|
||||
start := make(chan struct{})
|
||||
closed := make(chan struct{})
|
||||
outcomes := make(chan workConnResult, count)
|
||||
var ready sync.WaitGroup
|
||||
ready.Add(count - count/4)
|
||||
for i := count / 4; i < count; i++ {
|
||||
go func() {
|
||||
ready.Done()
|
||||
if i >= 3*count/4 {
|
||||
<-closed
|
||||
} else {
|
||||
<-start
|
||||
}
|
||||
conn, err := s.consumeWorkConnWaiterResult(waiters[i], results[i])
|
||||
if conn != nil {
|
||||
_ = conn.Close()
|
||||
}
|
||||
outcomes <- workConnResult{conn: conn, err: err}
|
||||
}()
|
||||
}
|
||||
ready.Wait()
|
||||
close(start)
|
||||
s.closeWorkConnPool()
|
||||
close(closed)
|
||||
for i := count / 4; i < count; i++ {
|
||||
outcome := waitForResult(t, outcomes, "handoff consumer to finish")
|
||||
if outcome.err != nil {
|
||||
require.ErrorIs(t, outcome.err, pkgerr.ErrCtlClosed)
|
||||
require.Nil(t, outcome.conn)
|
||||
} else {
|
||||
require.NotNil(t, outcome.conn)
|
||||
}
|
||||
}
|
||||
waitForSignal(t, s.workConnCleanupDone, "concurrent handoff cleanup")
|
||||
s.workConnMu.Lock()
|
||||
remaining := s.workConnHandoffs.Len()
|
||||
retained := false
|
||||
for _, waiter := range waiters {
|
||||
retained = retained || waiter.handoffConn != nil || waiter.handoffElement != nil
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
require.Zero(t, remaining)
|
||||
require.False(t, retained)
|
||||
for _, conn := range conns {
|
||||
require.Equal(t, int64(1), conn.closeCount.Load(), "consumer and cleanup must never both close the same connection")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkWorkConnPoolHit(b *testing.B) {
|
||||
s := newWorkConnScheduler(1, time.Minute, xlog.New(), nil)
|
||||
// Isolate acquisition from background refill work. The coalesced trigger
|
||||
// remains buffered; every iteration restores the same idle connection.
|
||||
s.workConnReconcileOnce.Do(func() {})
|
||||
conn := &proxy.WorkConn{}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
s.workConnCh <- conn
|
||||
got, err := s.GetWorkConn()
|
||||
if err != nil || got != conn {
|
||||
b.Fatalf("unexpected pool result: conn=%p err=%v", got, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user