From 53d0a53a36c0863b3d9ff7bf3de4f79916edfca5 Mon Sep 17 00:00:00 2001 From: fatedier Date: Mon, 14 Sep 2026 21:46:51 +0800 Subject: [PATCH] fix(server): control work connection refill scheduling --- client/control.go | 7 +- pkg/msg/msg.go | 12 +- pkg/msg/msg_test.go | 21 + pkg/plugin/server/manager.go | 37 +- pkg/plugin/server/manager_test.go | 128 +++ server/control.go | 151 ++- server/control_test.go | 1536 +++++++++++++++++++++++++++++ server/service.go | 27 +- server/service_test.go | 214 +++- server/workconn_scheduler.go | 739 ++++++++++++++ server/workconn_scheduler_test.go | 287 ++++++ 11 files changed, 3061 insertions(+), 98 deletions(-) create mode 100644 server/workconn_scheduler.go create mode 100644 server/workconn_scheduler_test.go diff --git a/client/control.go b/client/control.go index 4a354ff6..807a2f5d 100644 --- a/client/control.go +++ b/client/control.go @@ -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) diff --git a/pkg/msg/msg.go b/pkg/msg/msg.go index 52948bd5..206025f7 100644 --- a/pkg/msg/msg.go +++ b/pkg/msg/msg.go @@ -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"` diff --git a/pkg/msg/msg_test.go b/pkg/msg/msg_test.go index 06272faf..2420087b 100644 --- a/pkg/msg/msg_test.go +++ b/pkg/msg/msg_test.go @@ -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 +} diff --git a/pkg/plugin/server/manager.go b/pkg/plugin/server/manager.go index 6f1061c6..a93cd59e 100644 --- a/pkg/plugin/server/manager.go +++ b/pkg/plugin/server/manager.go @@ -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) } diff --git a/pkg/plugin/server/manager_test.go b/pkg/plugin/server/manager_test.go index 92391d4b..69696cbb 100644 --- a/pkg/plugin/server/manager_test.go +++ b/pkg/plugin/server/manager_test.go @@ -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 diff --git a/server/control.go b/server/control.go index 8009d8ac..1b4a4fc3 100644 --- a/server/control.go +++ b/server/control.go @@ -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() diff --git a/server/control_test.go b/server/control_test.go index 965ea5ea..2513742c 100644 --- a/server/control_test.go +++ b/server/control_test.go @@ -17,10 +17,12 @@ package server import ( "context" "errors" + "fmt" "math" "net" "os" "sync" + "sync/atomic" "testing" "time" @@ -28,6 +30,7 @@ import ( "github.com/fatedier/frp/pkg/auth" 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/server/controller" @@ -104,6 +107,1402 @@ func TestNewControlPoolCountBoundaries(t *testing.T) { } } +func TestControlReserveRefillUsesSlidingWindow(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 20) + ctl.workConnLeaseTimeout = 100 * time.Millisecond + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + ctl.reconcileWorkConnRequests() + for range workConnReserveWindowMax { + req := readWorkConnRequest(t, clientMsgConn) + require.Equal(t, msg.WorkConnTypeReserve, req.WorkConnType) + } + + for i := range 20 { + conn := newCountingCloseConn() + workConn := proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))) + accepted, err := ctl.registerWorkConn(workConn, workConnRequestReserve) + require.NoError(t, err) + require.True(t, accepted) + if i < 12 { + req := readWorkConnRequest(t, clientMsgConn) + require.Equal(t, msg.WorkConnTypeReserve, req.WorkConnType) + } + } + require.Len(t, ctl.workConnCh, 20) +} + +func TestControlDemandRefillPrecedesReserve(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 2) + ctl.controlID = 7 + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + resultCh := make(chan workConnResult, 1) + go func() { + conn, err := ctl.GetWorkConn() + resultCh <- workConnResult{conn: conn, err: err} + }() + + demandReq := readWorkConnRequest(t, clientMsgConn) + require.Equal(t, msg.WorkConnTypeDemand, demandReq.WorkConnType) + require.Equal(t, uint64(7), demandReq.ControlID) + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestDemand, + ) + require.NoError(t, err) + require.True(t, accepted) + result := <-resultCh + require.NoError(t, result.err) + require.NotNil(t, result.conn) +} + +func TestControlLegacySchedulingCapsMixedDemandAndReserveWindow(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 20) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + // Startup reserve requests consume part of the legacy aggregate window. + ctl.reconcileWorkConnRequests() + for range workConnReserveWindowMax { + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + } + + var firstWaiter *workConnWaiter + ctl.workConnMu.Lock() + for range 128 { + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.appendWorkConnWaiterLocked(waiter) + if firstWaiter == nil { + firstWaiter = waiter + } + } + ctl.workConnMu.Unlock() + + ctl.reconcileWorkConnRequests() + ctl.workConnMu.Lock() + legacyMode := ctl.workConnLegacyMode + demandPending := len(ctl.workConnPending[workConnRequestDemand]) + reservePending := len(ctl.workConnPending[workConnRequestReserve]) + ctl.workConnMu.Unlock() + require.True(t, legacyMode) + require.Equal(t, workConnLegacyWindow-workConnReserveWindowMax, demandPending) + require.Equal(t, workConnReserveWindowMax, reservePending) + + // A legacy response keeps the aggregate in-flight count bounded while + // refilling demand for the remaining waiters. + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestLegacy, + ) + require.NoError(t, err) + require.True(t, accepted) + result := <-firstWaiter.resultCh + require.NotNil(t, result.conn) + _ = result.conn.Close() + ctl.workConnMu.Lock() + pending := len(ctl.workConnPending[workConnRequestDemand]) + len(ctl.workConnPending[workConnRequestReserve]) + ctl.workConnMu.Unlock() + require.LessOrEqual(t, pending, workConnLegacyWindow) +} + +func TestControlLegacyResponsesDeliverConcurrentWaiters(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 20) + ctl.workConnLeaseTimeout = time.Second + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + // The net.Pipe peer must keep draining ReqWorkConn messages so the real + // dispatcher send loop can exercise refill pacing without blocking. + go func() { + for { + if _, err := clientMsgConn.ReadMsg(); err != nil { + return + } + } + }() + + results := make(chan workConnResult, 128) + for i := 0; i < cap(results); i++ { + go func() { + conn, err := ctl.GetWorkConn() + results <- workConnResult{conn: conn, err: err} + }() + } + + deadline := time.Now().Add(time.Second) + for { + ctl.workConnMu.Lock() + waiters := ctl.activeWaiterCountLocked() + demandPending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + reservePending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + legacyMode := ctl.workConnLegacyMode + ctl.workConnMu.Unlock() + if waiters == cap(results) { + require.True(t, legacyMode) + require.LessOrEqual(t, demandPending+reservePending, workConnLegacyWindow) + break + } + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %d work connection waiters, got %d", cap(results), waiters) + } + time.Sleep(time.Millisecond) + } + + for i := 0; i < cap(results); i++ { + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestLegacy, + ) + require.NoError(t, err) + require.True(t, accepted) + } + + for i := 0; i < cap(results); i++ { + result := <-results + require.NoError(t, result.err) + require.NotNil(t, result.conn) + require.NoError(t, result.conn.Close()) + } +} + +func TestControlTypedReserveResponsesDoNotOvershootPool(t *testing.T) { + const ( + poolCount = 10 + waiterCount = 100 + ) + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, poolCount) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + go func() { + for { + if _, err := clientMsgConn.ReadMsg(); err != nil { + return + } + } + }() + + waiters := make([]*workConnWaiter, 0, waiterCount) + ctl.workConnMu.Lock() + ctl.workConnLegacyMode = false + for range waiterCount { + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + waiters = append(waiters, waiter) + ctl.appendWorkConnWaiterLocked(waiter) + } + ctl.workConnMu.Unlock() + ctl.reconcileWorkConnRequests() + + ctl.workConnMu.Lock() + demandPending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + reservePending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + lateKinds := len(ctl.workConnLate) + ctl.workConnMu.Unlock() + require.Equal(t, workConnDemandWindow, demandPending) + require.Equal(t, workConnReserveWindowMax, reservePending) + require.Zero(t, lateKinds) + + registerResponse := func(kind workConnRequestKind) bool { + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + kind, + ) + if accepted { + require.NoError(t, err) + return true + } + require.ErrorIs(t, err, errWorkConnPoolFull) + require.NoError(t, conn.Close()) + return false + } + + reserveResponses := 0 + for { + ctl.workConnMu.Lock() + activeWaiters := ctl.activeWaiterCountLocked() + reservePending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + ctl.workConnMu.Unlock() + if activeWaiters == 0 || reservePending == 0 { + break + } + require.True(t, registerResponse(workConnRequestReserve)) + reserveResponses++ + require.LessOrEqual(t, reserveResponses, waiterCount) + } + + discarded := 0 + totalResponses := reserveResponses + for totalResponses <= 3*(waiterCount+poolCount) { + ctl.workConnMu.Lock() + demandSupply := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + + ctl.leaseCountLocked(ctl.workConnLate, workConnRequestDemand) + reserveSupply := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + + ctl.leaseCountLocked(ctl.workConnLate, workConnRequestReserve) + ctl.workConnMu.Unlock() + if demandSupply+reserveSupply == 0 { + break + } + kind := workConnRequestReserve + if demandSupply > 0 { + kind = workConnRequestDemand + } + if !registerResponse(kind) { + discarded++ + } + totalResponses++ + } + + require.Equal(t, waiterCount+poolCount, totalResponses, "each requested connection should contribute exactly once to waiter demand or idle reserve") + require.Equal(t, waiterCount+poolCount-workConnDemandWindow, reserveResponses) + require.Zero(t, discarded) + require.Len(t, ctl.workConnCh, poolCount) + ctl.workConnMu.Lock() + activeWaiters := ctl.activeWaiterCountLocked() + pendingKinds := len(ctl.workConnPending) + lateKinds = len(ctl.workConnLate) + ctl.workConnMu.Unlock() + require.Zero(t, activeWaiters) + require.Zero(t, pendingKinds) + require.Zero(t, lateKinds) + + for _, waiter := range waiters { + result := <-waiter.resultCh + require.NoError(t, result.err) + require.NotNil(t, result.conn) + require.NoError(t, result.conn.Close()) + } +} + +func TestControlLegacyResponsesKeepExactSupply(t *testing.T) { + const ( + poolCount = 10 + waiterCount = 100 + ) + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, poolCount) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + go func() { + for { + if _, err := clientMsgConn.ReadMsg(); err != nil { + return + } + } + }() + + waiters := make([]*workConnWaiter, 0, waiterCount) + ctl.workConnMu.Lock() + for range waiterCount { + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + waiters = append(waiters, waiter) + ctl.appendWorkConnWaiterLocked(waiter) + } + ctl.workConnMu.Unlock() + ctl.reconcileWorkConnRequests() + + ctl.workConnMu.Lock() + legacyMode := ctl.workConnLegacyMode + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + + ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + ctl.workConnMu.Unlock() + require.True(t, legacyMode) + require.Equal(t, workConnLegacyWindow, pending) + + for range waiterCount + poolCount { + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestLegacy, + ) + require.NoError(t, err) + require.True(t, accepted) + } + + require.Len(t, ctl.workConnCh, poolCount) + ctl.workConnMu.Lock() + activeWaiters := ctl.activeWaiterCountLocked() + pendingKinds := len(ctl.workConnPending) + lateKinds := len(ctl.workConnLate) + ctl.workConnMu.Unlock() + require.Zero(t, activeWaiters) + require.Zero(t, pendingKinds) + require.Zero(t, lateKinds) + + for _, waiter := range waiters { + result := <-waiter.resultCh + require.NoError(t, result.err) + require.NotNil(t, result.conn) + require.NoError(t, result.conn.Close()) + } +} + +func TestControlReserveTimeoutReleasesWindowBudget(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 4) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + ctl.reconcileWorkConnRequests() + for range 4 { + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + } + + // Move the four dispatched leases through the first timeout transition + // deterministically, then let reconcile create their replacement budget. + for range 4 { + moveReserveWorkConnLeaseToLateForTest(t, ctl) + } + ctl.reconcileWorkConnRequests() + for range 4 { + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + } + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestReserve, + ) + require.True(t, accepted) + require.NoError(t, err) + require.Len(t, ctl.workConnCh, 1) + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + ctl.workConnMu.Unlock() + require.Equal(t, 3, pending) +} + +func TestControlLateReserveResponseRejectedWhenReplacementCoversPool(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 1) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + ctl.reconcileWorkConnRequests() + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + moveReserveWorkConnLeaseToLateForTest(t, ctl) + ctl.reconcileWorkConnRequests() + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + + // The replacement fills the pool even when the original never responded. + replacementConn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(replacementConn, msg.NewV1ReadWriter(replacementConn))), + workConnRequestReserve, + ) + require.True(t, accepted) + require.NoError(t, err) + require.Len(t, ctl.workConnCh, 1) + + // If the original eventually arrives, it cannot grow an already full pool. + conn := newCountingCloseConn() + accepted, err = ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestReserve, + ) + require.False(t, accepted) + require.ErrorIs(t, err, errWorkConnPoolFull) + require.NoError(t, conn.Close()) + ctl.workConnMu.Lock() + late := len(ctl.workConnLate[workConnRequestReserve]) + pending := len(ctl.workConnPending[workConnRequestReserve]) + ctl.workConnMu.Unlock() + require.Zero(t, late) + require.Zero(t, pending) +} + +func TestControlLegacyResponseKeepsDemandPriority(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + resultCh := make(chan workConnResult, 1) + go func() { + conn, err := ctl.GetWorkConn() + resultCh <- workConnResult{conn: conn, err: err} + }() + require.Equal(t, msg.WorkConnTypeDemand, readWorkConnRequest(t, clientMsgConn).WorkConnType) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestLegacy, + ) + require.NoError(t, err) + require.True(t, accepted) + result := <-resultCh + require.NoError(t, result.err) + require.NotNil(t, result.conn) +} + +func TestControlReserveResponseKeepsOutstandingDemandSupply(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 2) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + resultCh := make(chan workConnResult, 1) + go func() { + conn, err := ctl.GetWorkConn() + resultCh <- workConnResult{conn: conn, err: err} + }() + require.Equal(t, msg.WorkConnTypeDemand, readWorkConnRequest(t, clientMsgConn).WorkConnType) + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestReserve, + ) + require.NoError(t, err) + require.True(t, accepted) + result := <-resultCh + require.NoError(t, result.err) + require.NotNil(t, result.conn) + _ = result.conn.Close() + + ctl.workConnMu.Lock() + demandPending := len(ctl.workConnPending[workConnRequestDemand]) + reservePending := len(ctl.workConnPending[workConnRequestReserve]) + lateKinds := len(ctl.workConnLate) + ctl.workConnMu.Unlock() + require.Equal(t, 1, demandPending) + require.Equal(t, 1, reservePending) + require.Zero(t, lateKinds) + + conn = newCountingCloseConn() + accepted, err = ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestDemand, + ) + require.NoError(t, err) + require.True(t, accepted) + + conn = newCountingCloseConn() + accepted, err = ctl.registerWorkConn( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestReserve, + ) + require.NoError(t, err) + require.True(t, accepted) + require.Len(t, ctl.workConnCh, 2) +} + +func TestControlCloseWakesWorkConnWaiter(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + resultCh := make(chan error, 1) + go func() { + _, err := ctl.GetWorkConn() + resultCh <- err + }() + require.Equal(t, msg.WorkConnTypeDemand, readWorkConnRequest(t, clientMsgConn).WorkConnType) + ctl.closeWorkConnPool() + require.ErrorIs(t, <-resultCh, pkgerr.ErrCtlClosed) +} + +func TestControlManagerRegisterWorkConnDoesNotHoldLifecycleDuringRefill(t *testing.T) { + for _, action := range []string{"close", "replace"} { + t.Run(action, func(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + conn := newBlockingWriteConn() + ctl, err := NewControl(context.Background(), &SessionContext{ + RC: &controller.ResourceController{}, + PxyManager: proxy.NewManager(), + PluginManager: plugin.NewManager(), + AuthVerifier: auth.AlwaysPassVerifier, + Conn: msg.NewConn(conn, msg.NewV1ReadWriter(conn)), + LoginMsg: &msg.Login{ + RunID: "stalled-refill", + ClientID: "stalled-refill-client", + PoolCount: 10, + }, + ServerCfg: &v1.ServerConfig{}, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = ctl.Close() }) + require.NoError(t, manager.Add(ctl)) + active, err := manager.Activate(ctl) + require.NoError(t, err) + require.True(t, active) + require.True(t, ctl.Start()) + + // Block the dispatcher writer, then leave enough queued messages for + // the refill batch to block in Dispatcher.Send. + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + waitForSignal(t, conn.writeStarted, "dispatcher write to stall") + for range 40 { + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + } + + const waiterCount = 100 + waiters := make([]*workConnWaiter, 0, waiterCount) + ctl.workConnMu.Lock() + ctl.workConnLegacyMode = false + for range waiterCount { + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + waiters = append(waiters, waiter) + ctl.appendWorkConnWaiterLocked(waiter) + } + ctl.workConnMu.Unlock() + + registerDone := make(chan error, 1) + workConn := newCountingCloseConn() + go func() { + registerDone <- manager.RegisterWorkConnWithType( + ctl, + proxy.NewWorkConn(msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))), + workConnRequestDemand, + ) + }() + result := waitForResult(t, waiters[0].resultCh, "registered work connection waiter") + require.NoError(t, result.err) + require.NotNil(t, result.conn) + deadline := time.Now().Add(time.Second) + for { + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + + ctl.leaseCountLocked(ctl.workConnPending, workConnRequestReserve) + ctl.workConnMu.Unlock() + if pending > 0 { + break + } + if time.Now().After(deadline) { + t.Fatal("refill reconciler did not create a lease") + } + time.Sleep(time.Millisecond) + } + + if action == "close" { + closeDone := make(chan error, 1) + go func() { closeDone <- ctl.Close() }() + select { + case err := <-closeDone: + require.NoError(t, err) + case <-time.After(time.Second): + _ = ctl.interruptReadAndClose() + t.Errorf("Close blocked while refill Send was stalled") + } + } else { + replacement, _ := newLifecycleTestControl(t, "stalled-refill", "replacement", newCountingServerMetrics()) + addDone := make(chan error, 1) + go func() { addDone <- manager.Add(replacement) }() + select { + case err := <-addDone: + require.NoError(t, err) + case <-time.After(time.Second): + _ = ctl.interruptReadAndClose() + t.Errorf("replacement blocked while refill Send was stalled") + } + } + + select { + case err := <-registerDone: + require.NoError(t, err) + case <-time.After(time.Second): + _ = ctl.interruptReadAndClose() + t.Fatal("RegisterWorkConn did not finish after control unblocked") + } + _ = result.conn.Close() + waitForControlDone(t, ctl) + select { + case <-ctl.workConnReconcileDone: + case <-time.After(time.Second): + t.Fatal("work connection reconciler leaked after control close/replacement") + } + }) + } +} + +func TestControlCloseAndReplacementDoNotBlockOnPooledWorkConnClose(t *testing.T) { + for _, action := range []string{"close", "replace"} { + t.Run(action, func(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + metrics := newCountingServerMetrics() + ctl, controlConn := newLifecycleTestControl(t, "blocked-cleanup", "blocked-client", metrics) + mustAddAndActivate(t, manager, ctl) + require.True(t, ctl.Start()) + waitForSignal(t, controlConn.readStarted, "control reader to start") + + pooledConn := newBlockingCloseConn() + pooledWorkConn := proxy.NewWorkConn(msg.NewConn(pooledConn, msg.NewV1ReadWriter(pooledConn))) + ctl.workConnMu.Lock() + ctl.workConnCh <- pooledWorkConn + ctl.workConnMu.Unlock() + + if action == "close" { + closeDone := make(chan error, 1) + go func() { closeDone <- ctl.Close() }() + select { + case err := <-closeDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("Control.Close blocked on pooled work connection Close") + } + } else { + replacement, _ := newLifecycleTestControl(t, "blocked-cleanup", "replacement", newCountingServerMetrics()) + addDone := make(chan error, 1) + go func() { addDone <- manager.Add(replacement) }() + select { + case err := <-addDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("same-run replacement blocked on pooled work connection Close") + } + } + + waitForSignal(t, pooledConn.closeStarted, "pooled work connection cleanup to start") + close(pooledConn.allowClose) + waitForSignal(t, ctl.workConnCleanupDone, "pooled work connection cleanup to finish") + require.Equal(t, int64(1), pooledConn.closeCount.Load()) + waitForControlDone(t, ctl) + }) + } +} + +func TestControlUndeliveredWaiterHandoffIsClosedOnShutdown(t *testing.T) { + for _, action := range []string{"close", "replace"} { + t.Run(action, func(t *testing.T) { + metrics := newCountingServerMetrics() + ctl, controlConn := newLifecycleTestControl(t, "undelivered-handoff", "old-client", metrics) + manager := NewControlManager(registry.NewClientRegistry()) + mustAddAndActivate(t, manager, ctl) + require.True(t, ctl.Start()) + waitForSignal(t, controlConn.readStarted, "control reader to start") + + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.workConnMu.Lock() + ctl.appendWorkConnWaiterLocked(waiter) + ctl.workConnMu.Unlock() + + pooledConn := newBlockingCloseConn() + workConn := proxy.NewWorkConn(msg.NewConn(pooledConn, msg.NewV1ReadWriter(pooledConn))) + accepted, err := ctl.registerWorkConnState(workConn, workConnRequestDemand) + require.NoError(t, err) + require.True(t, accepted) + // Leave the buffered result unread to model a consumer preempted + // between the channel wakeup and its ownership claim. + require.Len(t, waiter.resultCh, 1) + + if action == "close" { + closeDone := make(chan error, 1) + go func() { closeDone <- ctl.Close() }() + select { + case err := <-closeDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("Control.Close blocked on an undelivered waiter handoff") + } + } else { + replacement, _ := newLifecycleTestControl(t, "undelivered-handoff", "new-client", newCountingServerMetrics()) + addDone := make(chan error, 1) + go func() { addDone <- manager.Add(replacement) }() + select { + case err := <-addDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("same-run replacement blocked on an undelivered waiter handoff") + } + } + + waitForSignal(t, pooledConn.closeStarted, "undelivered work connection cleanup to start") + result := <-waiter.resultCh + conn, err := ctl.consumeWorkConnWaiterResult(waiter, result) + require.Nil(t, conn) + require.ErrorIs(t, err, pkgerr.ErrCtlClosed) + close(pooledConn.allowClose) + waitForSignal(t, ctl.workConnCleanupDone, "undelivered work connection cleanup to finish") + require.Equal(t, int64(1), pooledConn.closeCount.Load()) + waitForControlDone(t, ctl) + }) + } +} + +func TestControlGetWorkConnTimeoutStartsBeforeBackpressuredRefill(t *testing.T) { + conn := newBlockingWriteConn() + cfg := &v1.ServerConfig{UserConnTimeout: 1} + cfg.Transport.MaxPoolCount = 0 + ctl, err := NewControl(context.Background(), &SessionContext{ + RC: &controller.ResourceController{}, + PxyManager: proxy.NewManager(), + PluginManager: plugin.NewManager(), + AuthVerifier: auth.AlwaysPassVerifier, + Conn: msg.NewConn(conn, msg.NewV1ReadWriter(conn)), + LoginMsg: &msg.Login{ + RunID: "backpressured-timeout", + PoolCount: 0, + }, + ServerCfg: cfg, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = ctl.Close() }) + ctl.state = controlStateRunning + ctl.msgDispatcher.Run() + ctl.workConnLeaseTimeout = 30 * time.Millisecond + + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + waitForSignal(t, conn.writeStarted, "dispatcher write to stall") + for range 100 { + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + } + + resultCh := make(chan error, 1) + started := time.Now() + go func() { + _, err := ctl.GetWorkConn() + resultCh <- err + }() + select { + case err := <-resultCh: + require.ErrorContains(t, err, "timeout trying to get work connection") + require.GreaterOrEqual(t, time.Since(started), 20*time.Millisecond) + require.Less(t, time.Since(started), 300*time.Millisecond) + case <-time.After(time.Second): + t.Fatal("GetWorkConn ignored timeout while refill was backpressured") + } + + _ = ctl.Close() + select { + case <-ctl.workConnReconcileDone: + case <-time.After(time.Second): + t.Fatal("work connection reconciler did not stop after close") + } +} + +func TestControlUnsentWorkConnLeaseStartsTimeoutAfterDispatch(t *testing.T) { + conn := newBlockingWriteConn() + cfg := &v1.ServerConfig{UserConnTimeout: 1} + cfg.Transport.MaxPoolCount = 0 + ctl, err := NewControl(context.Background(), &SessionContext{ + RC: &controller.ResourceController{}, + PxyManager: proxy.NewManager(), + PluginManager: plugin.NewManager(), + AuthVerifier: auth.AlwaysPassVerifier, + Conn: msg.NewConn(conn, msg.NewV1ReadWriter(conn)), + LoginMsg: &msg.Login{ + RunID: "queued-lease-timeout", + PoolCount: 0, + }, + ServerCfg: cfg, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = ctl.Close() }) + ctl.state = controlStateRunning + ctl.msgDispatcher.Run() + ctl.workConnLeaseTimeout = 25 * time.Millisecond + + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + waitForSignal(t, conn.writeStarted, "dispatcher write to stall") + for range 100 { + require.NoError(t, ctl.msgDispatcher.Send(&msg.Ping{})) + } + + resultCh := make(chan error, 1) + go func() { + _, err := ctl.GetWorkConn() + resultCh <- err + }() + deadline := time.Now().Add(time.Second) + for { + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + late := ctl.leaseCountLocked(ctl.workConnLate, workConnRequestDemand) + var timer *time.Timer + dispatched := false + if leases := ctl.workConnPending[workConnRequestDemand]; len(leases) > 0 { + timer = leases[0].timer + dispatched = leases[0].dispatched + } + ctl.workConnMu.Unlock() + if pending == 1 { + require.Zero(t, late) + require.Nil(t, timer) + require.False(t, dispatched) + break + } + if time.Now().After(deadline) { + t.Fatal("demand lease was not queued") + } + time.Sleep(time.Millisecond) + } + time.Sleep(80 * time.Millisecond) + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + late := ctl.leaseCountLocked(ctl.workConnLate, workConnRequestDemand) + var leaseTimer *time.Timer + if leases := ctl.workConnPending[workConnRequestDemand]; len(leases) > 0 { + leaseTimer = leases[0].timer + } + ctl.workConnMu.Unlock() + require.Equal(t, 1, pending) + require.Zero(t, late) + require.Nil(t, leaseTimer) + + conn.releaseWrite() + deadline = time.Now().Add(time.Second) + for { + ctl.workConnMu.Lock() + var timer *time.Timer + if leases := ctl.workConnPending[workConnRequestDemand]; len(leases) > 0 { + timer = leases[0].timer + } + ctl.workConnMu.Unlock() + if timer != nil { + break + } + if time.Now().After(deadline) { + t.Fatal("queued request did not start its response timeout after dispatch") + } + time.Sleep(time.Millisecond) + } + _ = ctl.Close() + select { + case <-ctl.workConnReconcileDone: + case <-time.After(time.Second): + t.Fatal("work connection reconciler did not stop after close") + } + select { + case err := <-resultCh: + require.ErrorContains(t, err, "timeout trying to get work connection") + case <-time.After(time.Second): + t.Fatal("GetWorkConn did not finish") + } +} + +func TestControlLeaseTimerDoesNotStartAfterResponseWinsRace(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 1) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + ctl.workConnMu.Lock() + lease := ctl.newWorkConnLeaseLocked(workConnRequestDemand) + ctl.workConnMu.Unlock() + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestDemand, + ) + require.True(t, accepted) + require.NoError(t, err) + ctl.startWorkConnLeaseTimer(lease) + + ctl.workConnMu.Lock() + dispatched := lease.dispatched + leaseTimer := lease.timer + pendingKinds := len(ctl.workConnPending) + ctl.workConnMu.Unlock() + require.False(t, dispatched) + require.Nil(t, leaseTimer) + require.Zero(t, pendingKinds) +} + +func moveReserveWorkConnLeaseToLateForTest(t *testing.T, ctl *Control) { + t.Helper() + const kind = workConnRequestReserve + ctl.workConnMu.Lock() + leases := ctl.workConnPending[kind] + if len(leases) == 0 { + ctl.workConnMu.Unlock() + t.Fatal("expected a pending work connection lease") + return + } + lease := leases[0] + ctl.removeLeaseLocked(ctl.workConnPending, lease) + lease.late = true + if lease.timer != nil { + lease.timer.Stop() + lease.timer = nil + } + ctl.workConnLate[kind] = append(ctl.workConnLate[kind], lease) + ctl.workConnMu.Unlock() +} + +func TestControlLateResponseAdmissionAvoidsReplacementOversupply(t *testing.T) { + for _, responseKind := range []workConnRequestKind{workConnRequestReserve, workConnRequestLegacy} { + t.Run(string(responseKind), func(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 1) + ctl.workConnLeaseTimeout = time.Minute + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + ctl.reconcileWorkConnRequests() + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + moveReserveWorkConnLeaseToLateForTest(t, ctl) + ctl.reconcileWorkConnRequests() + require.Equal(t, msg.WorkConnTypeReserve, readWorkConnRequest(t, clientMsgConn).WorkConnType) + + // Replies carry no request ID. Even if the original arrives first, + // it retires pending supply and the subsequent duplicate is rejected. + lateConn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(lateConn, msg.NewV1ReadWriter(lateConn))), + responseKind, + ) + require.True(t, accepted) + require.NoError(t, err) + require.Len(t, ctl.workConnCh, 1) + + replacementConn := newCountingCloseConn() + accepted, err = ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(replacementConn, msg.NewV1ReadWriter(replacementConn))), + responseKind, + ) + require.False(t, accepted) + require.ErrorIs(t, err, errWorkConnPoolFull) + require.NoError(t, replacementConn.Close()) + require.Len(t, ctl.workConnCh, 1) + }) + } +} + +func TestControlUnmatchedResponseCannotCacheBeyondDemandOrPool(t *testing.T) { + for _, responseKind := range []workConnRequestKind{ + workConnRequestDemand, + workConnRequestReserve, + workConnRequestLegacy, + } { + t.Run(string(responseKind), func(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + responseKind, + ) + require.False(t, accepted) + require.ErrorIs(t, err, errWorkConnPoolFull) + require.NoError(t, conn.Close()) + require.Len(t, ctl.workConnCh, 0) + }) + } + + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.workConnMu.Lock() + ctl.appendWorkConnWaiterLocked(waiter) + lease := ctl.newWorkConnLeaseLocked(workConnRequestDemand) + ctl.workConnMu.Unlock() + require.True(t, ctl.cancelWorkConnWaiter(waiter)) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestDemand, + ) + require.False(t, accepted) + require.ErrorIs(t, err, errWorkConnPoolFull) + require.NoError(t, conn.Close()) + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + leaseTimer := lease.timer + ctl.workConnMu.Unlock() + require.Zero(t, pending) + require.Nil(t, leaseTimer) +} + +func TestControlLateResponseCanSatisfyWaiterOrRealGap(t *testing.T) { + for _, responseKind := range []workConnRequestKind{workConnRequestDemand, workConnRequestLegacy} { + t.Run(string(responseKind), func(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.workConnMu.Lock() + ctl.appendWorkConnWaiterLocked(waiter) + ctl.workConnLate[workConnRequestDemand] = append(ctl.workConnLate[workConnRequestDemand], &workConnLease{ + kind: workConnRequestDemand, + active: true, + late: true, + }) + ctl.workConnMu.Unlock() + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + responseKind, + ) + require.True(t, accepted) + require.NoError(t, err) + result := <-waiter.resultCh + handoff, err := ctl.consumeWorkConnWaiterResult(waiter, result) + require.NoError(t, err) + require.NotNil(t, handoff) + require.NoError(t, handoff.Close()) + }) + } + + for _, responseKind := range []workConnRequestKind{workConnRequestReserve, workConnRequestLegacy} { + t.Run("gap_"+string(responseKind), func(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 2) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + ctl.workConnMu.Lock() + ctl.workConnLate[workConnRequestReserve] = append(ctl.workConnLate[workConnRequestReserve], &workConnLease{ + kind: workConnRequestReserve, + active: true, + late: true, + }) + ctl.workConnPending[workConnRequestReserve] = append(ctl.workConnPending[workConnRequestReserve], &workConnLease{ + kind: workConnRequestReserve, + active: true, + }) + ctl.workConnMu.Unlock() + + lateConn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(lateConn, msg.NewV1ReadWriter(lateConn))), + responseKind, + ) + require.True(t, accepted) + require.NoError(t, err) + require.Len(t, ctl.workConnCh, 1) + + pendingConn := newCountingCloseConn() + accepted, err = ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(pendingConn, msg.NewV1ReadWriter(pendingConn))), + responseKind, + ) + require.True(t, accepted) + require.NoError(t, err) + require.Len(t, ctl.workConnCh, 2) + }) + } +} + +func TestControlActiveWaiterCountTracksFIFOCancelHandoffAndClose(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + waiters := []*workConnWaiter{ + {resultCh: make(chan workConnResult, 1), active: true}, + {resultCh: make(chan workConnResult, 1), active: true}, + {resultCh: make(chan workConnResult, 1), active: true}, + } + ctl.workConnMu.Lock() + for _, waiter := range waiters { + ctl.appendWorkConnWaiterLocked(waiter) + } + activeWaiters := ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Equal(t, 3, activeWaiters) + + require.True(t, ctl.cancelWorkConnWaiter(waiters[1])) + ctl.workConnMu.Lock() + activeWaiters = ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Equal(t, 2, activeWaiters) + + conn := newCountingCloseConn() + accepted, err := ctl.registerWorkConnState( + proxy.NewWorkConn(msg.NewConn(conn, msg.NewV1ReadWriter(conn))), + workConnRequestDemand, + ) + require.True(t, accepted) + require.NoError(t, err) + ctl.workConnMu.Lock() + activeWaiters = ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Equal(t, 1, activeWaiters) + + result := <-waiters[0].resultCh + handoff, err := ctl.consumeWorkConnWaiterResult(waiters[0], result) + require.NoError(t, err) + require.NotNil(t, handoff) + require.NoError(t, handoff.Close()) + ctl.workConnMu.Lock() + activeWaiters = ctl.activeWaiterCountLocked() + handoffs := ctl.workConnHandoffs.Len() + ctl.workConnMu.Unlock() + require.Equal(t, 1, activeWaiters) + require.Zero(t, handoffs) + + ctl.closeWorkConnPool() + ctl.workConnMu.Lock() + activeWaiters = ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Zero(t, activeWaiters) +} + +func TestControlActiveWaiterCountLargeConcurrentCancellation(t *testing.T) { + const waiterCount = 1024 + + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + waiters := make([]*workConnWaiter, waiterCount) + ctl.workConnMu.Lock() + for i := range waiters { + waiters[i] = &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.appendWorkConnWaiterLocked(waiters[i]) + } + initialCount := ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Equal(t, waiterCount, initialCount) + + var ( + wg sync.WaitGroup + canceledCount atomic.Int64 + ) + for i := 0; i < waiterCount; i += 2 { + wg.Add(1) + go func(waiter *workConnWaiter) { + defer wg.Done() + if ctl.cancelWorkConnWaiter(waiter) { + canceledCount.Add(1) + } + }(waiters[i]) + } + wg.Wait() + require.Equal(t, int64(waiterCount/2), canceledCount.Load()) + + ctl.workConnMu.Lock() + activeCount := ctl.activeWaiterCountLocked() + queueLen := len(ctl.workConnWaiters) + // Canceled waiters remain as lazy tombstones; reconcile reads the O(1) + // active count and does not compact or rescan the full queue on every + // cancellation. A threshold-triggered compact may already have removed + // the tombstones. + ctl.workConnMu.Unlock() + require.Equal(t, waiterCount/2, activeCount) + require.LessOrEqual(t, queueLen, waiterCount) + + ctl.workConnMu.Lock() + popped := make([]*workConnWaiter, 0, waiterCount/2) + for i := 1; i < waiterCount; i += 2 { + popped = append(popped, ctl.popWorkConnWaiterLocked()) + } + nilPop := ctl.popWorkConnWaiterLocked() + remainingCount := ctl.activeWaiterCountLocked() + remainingQueueLen := len(ctl.workConnWaiters) + remainingTombstones := ctl.workConnWaiterTombstones + ctl.workConnMu.Unlock() + for i, waiter := range popped { + require.Same(t, waiters[2*i+1], waiter) + } + require.Nil(t, nilPop) + require.Zero(t, remainingCount) + require.Zero(t, remainingQueueLen) + require.Zero(t, remainingTombstones) +} + +func TestControlCanceledWaitersAreCompactedAndBounded(t *testing.T) { + const canceledCount = 10000 + + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + for range canceledCount { + waiter := &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.workConnMu.Lock() + ctl.appendWorkConnWaiterLocked(waiter) + ctl.workConnMu.Unlock() + require.True(t, ctl.cancelWorkConnWaiter(waiter)) + } + + ctl.workConnMu.Lock() + activeCount := ctl.activeWaiterCountLocked() + tombstoneCount := ctl.workConnWaiterTombstones + queueLen := len(ctl.workConnWaiters) + ctl.workConnMu.Unlock() + require.Zero(t, activeCount) + require.Less(t, tombstoneCount, workConnWaiterCompactMinTombstones) + require.Less(t, queueLen, workConnWaiterCompactMinTombstones) + require.Equal(t, activeCount+tombstoneCount, queueLen) + + // Closing the control clears both active waiters and residual tombstones. + ctl.closeWorkConnPool() + ctl.workConnMu.Lock() + activeCount = ctl.activeWaiterCountLocked() + tombstoneCount = ctl.workConnWaiterTombstones + queueLen = len(ctl.workConnWaiters) + ctl.workConnMu.Unlock() + require.Zero(t, activeCount) + require.Zero(t, tombstoneCount) + require.Zero(t, queueLen) +} + +func TestControlCanceledWaiterCompactionPreservesFIFO(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControl(t, 0) + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + const waiterCount = 256 + waiters := make([]*workConnWaiter, waiterCount) + ctl.workConnMu.Lock() + for i := range waiters { + waiters[i] = &workConnWaiter{resultCh: make(chan workConnResult, 1), active: true} + ctl.appendWorkConnWaiterLocked(waiters[i]) + } + ctl.workConnMu.Unlock() + + for i := 1; i < waiterCount; i += 2 { + require.True(t, ctl.cancelWorkConnWaiter(waiters[i])) + } + + ctl.workConnMu.Lock() + activeCount := ctl.activeWaiterCountLocked() + ctl.workConnMu.Unlock() + require.Equal(t, waiterCount/2, activeCount) + + ctl.workConnMu.Lock() + popped := make([]*workConnWaiter, 0, waiterCount/2) + for i := 0; i < waiterCount; i += 2 { + popped = append(popped, ctl.popWorkConnWaiterLocked()) + } + nilPop := ctl.popWorkConnWaiterLocked() + remainingCount := ctl.activeWaiterCountLocked() + remainingQueueLen := len(ctl.workConnWaiters) + remainingTombstones := ctl.workConnWaiterTombstones + ctl.workConnMu.Unlock() + for i, waiter := range popped { + require.Same(t, waiters[2*i], waiter) + } + require.Nil(t, nilPop) + require.Zero(t, remainingCount) + require.Zero(t, remainingQueueLen) + require.Zero(t, remainingTombstones) +} + +func TestControlNonPositiveWorkConnTimeoutIsSharedByWaiterAndLease(t *testing.T) { + for _, configuredTimeout := range []int64{0, -1} { + t.Run(fmt.Sprintf("timeout_%d", configuredTimeout), func(t *testing.T) { + ctl, serverMsgConn, clientMsgConn := newWorkConnSchedulingControlWithTimeout(t, 0, configuredTimeout) + require.Equal(t, 10*time.Second, ctl.workConnLeaseTimeoutValue()) + // Shorten the effective value for a deterministic test while retaining + // the non-positive configuration that would otherwise fall back to 10s. + ctl.workConnLeaseTimeout = 30 * time.Millisecond + t.Cleanup(func() { + ctl.closeWorkConnPool() + _ = serverMsgConn.Close() + _ = clientMsgConn.Close() + }) + + resultCh := make(chan error, 1) + started := time.Now() + go func() { + _, err := ctl.GetWorkConn() + resultCh <- err + }() + require.Equal(t, msg.WorkConnTypeDemand, readWorkConnRequest(t, clientMsgConn).WorkConnType) + select { + case err := <-resultCh: + require.ErrorContains(t, err, "timeout trying to get work connection") + require.GreaterOrEqual(t, time.Since(started), 20*time.Millisecond) + require.Less(t, time.Since(started), 300*time.Millisecond) + case <-time.After(time.Second): + t.Fatal("waiter did not use the effective non-positive timeout") + } + + deadline := time.Now().Add(time.Second) + for { + ctl.workConnMu.Lock() + pending := ctl.leaseCountLocked(ctl.workConnPending, workConnRequestDemand) + late := ctl.leaseCountLocked(ctl.workConnLate, workConnRequestDemand) + ctl.workConnMu.Unlock() + if pending == 0 && late == 1 { + break + } + if time.Now().After(deadline) { + t.Fatalf("lease timeout did not match waiter timeout: pending=%d late=%d", pending, late) + } + time.Sleep(time.Millisecond) + } + ctl.closeWorkConnPool() + select { + case <-ctl.workConnReconcileDone: + case <-time.After(time.Second): + t.Fatal("work connection reconciler did not stop after pool close") + } + }) + } +} + func TestControlRunningReplacementFinishesInWorker(t *testing.T) { clientRegistry := registry.NewClientRegistry() manager := NewControlManager(clientRegistry) @@ -423,6 +1822,7 @@ func waitForControlDone(t *testing.T, ctl *Control) { close(done) }() waitForSignal(t, done, "control to finish") + ctl.waitWorkConnCleanup() } func currentControlForTest(manager *ControlManager, runID string) *Control { @@ -468,6 +1868,64 @@ func waitForSignal(t *testing.T, ch <-chan struct{}, description string) { } } +func newWorkConnSchedulingControl(t *testing.T, poolCount int) (*Control, *msg.Conn, *msg.Conn) { + return newWorkConnSchedulingControlWithTimeout(t, poolCount, 1) +} + +func newWorkConnSchedulingControlWithTimeout(t *testing.T, poolCount int, userConnTimeout int64) (*Control, *msg.Conn, *msg.Conn) { + t.Helper() + serverConn, clientConn := net.Pipe() + serverMsgConn := msg.NewConn(serverConn, msg.NewV1ReadWriter(serverConn)) + clientMsgConn := msg.NewConn(clientConn, msg.NewV1ReadWriter(clientConn)) + cfg := &v1.ServerConfig{} + cfg.Transport.MaxPoolCount = int64(poolCount) + cfg.UserConnTimeout = userConnTimeout + ctl, err := NewControl(context.Background(), &SessionContext{ + RC: &controller.ResourceController{}, + PxyManager: proxy.NewManager(), + PluginManager: plugin.NewManager(), + AuthVerifier: auth.AlwaysPassVerifier, + Conn: serverMsgConn, + LoginMsg: &msg.Login{ + RunID: "work-conn-scheduling", + PoolCount: poolCount, + }, + ServerCfg: cfg, + }) + require.NoError(t, err) + ctl.state = controlStateRunning + ctl.msgDispatcher.Run() + return ctl, serverMsgConn, clientMsgConn +} + +func readWorkConnRequest(t *testing.T, conn *msg.Conn) *msg.ReqWorkConn { + t.Helper() + resultCh := make(chan struct { + msg *msg.ReqWorkConn + err error + }, 1) + go func() { + m, err := conn.ReadMsg() + req, ok := m.(*msg.ReqWorkConn) + if !ok && err == nil { + err = errors.New("unexpected work connection request message") + } + resultCh <- struct { + msg *msg.ReqWorkConn + err error + }{msg: req, err: err} + }() + select { + case result := <-resultCh: + require.NoError(t, result.err) + require.NotNil(t, result.msg) + return result.msg + case <-time.After(time.Second): + t.Fatal("timed out waiting for ReqWorkConn") + return nil + } +} + type deadlineReadConn struct { readStarted chan struct{} unblockRead chan struct{} @@ -539,6 +1997,84 @@ func (c *deadlineReadConn) eventsSnapshot() []string { return append([]string(nil), c.events...) } +type blockingWriteConn struct { + writeStarted chan struct{} + closed chan struct{} + allowWrite chan struct{} + writeOnce sync.Once + closeOnce sync.Once + allowOnce sync.Once +} + +func newBlockingWriteConn() *blockingWriteConn { + return &blockingWriteConn{ + writeStarted: make(chan struct{}), + closed: make(chan struct{}), + allowWrite: make(chan struct{}), + } +} + +type blockingCloseConn struct { + closeStarted chan struct{} + allowClose chan struct{} + closeCount atomic.Int64 + closeOnce sync.Once +} + +func newBlockingCloseConn() *blockingCloseConn { + return &blockingCloseConn{ + closeStarted: make(chan struct{}), + allowClose: make(chan struct{}), + } +} + +func (*blockingCloseConn) Read([]byte) (int, error) { return 0, net.ErrClosed } +func (*blockingCloseConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *blockingCloseConn) Close() error { + c.closeOnce.Do(func() { close(c.closeStarted) }) + <-c.allowClose + c.closeCount.Add(1) + return nil +} +func (*blockingCloseConn) LocalAddr() net.Addr { return lifecycleTestAddr("local") } +func (*blockingCloseConn) RemoteAddr() net.Addr { return lifecycleTestAddr("remote") } +func (*blockingCloseConn) SetDeadline(time.Time) error { return nil } +func (*blockingCloseConn) SetReadDeadline(time.Time) error { return nil } +func (*blockingCloseConn) SetWriteDeadline(time.Time) error { return nil } + +func (c *blockingWriteConn) Read([]byte) (int, error) { + <-c.closed + return 0, net.ErrClosed +} + +func (c *blockingWriteConn) Write(p []byte) (int, error) { + c.writeOnce.Do(func() { close(c.writeStarted) }) + select { + case <-c.allowWrite: + return len(p), nil + case <-c.closed: + return 0, net.ErrClosed + } +} + +func (c *blockingWriteConn) Close() error { + c.closeOnce.Do(func() { + close(c.closed) + c.allowOnce.Do(func() { close(c.allowWrite) }) + }) + return nil +} + +func (c *blockingWriteConn) releaseWrite() { + c.allowOnce.Do(func() { close(c.allowWrite) }) +} + +func (*blockingWriteConn) LocalAddr() net.Addr { return lifecycleTestAddr("local") } +func (*blockingWriteConn) RemoteAddr() net.Addr { return lifecycleTestAddr("remote") } +func (*blockingWriteConn) SetDeadline(time.Time) error { return nil } +func (*blockingWriteConn) SetReadDeadline(time.Time) error { return nil } +func (*blockingWriteConn) SetWriteDeadline(time.Time) error { return nil } + type lifecycleTestAddr string func (a lifecycleTestAddr) Network() string { return string(a) } diff --git a/server/service.go b/server/service.go index 6ea24cc4..ba1d23a7 100644 --- a/server/service.go +++ b/server/service.go @@ -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 { diff --git a/server/service_test.go b/server/service_test.go index d4fcbe77..1316406f 100644 --- a/server/service_test.go +++ b/server/service_test.go @@ -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{}), diff --git a/server/workconn_scheduler.go b/server/workconn_scheduler.go new file mode 100644 index 00000000..c922be48 --- /dev/null +++ b/server/workconn_scheduler.go @@ -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 + } +} diff --git a/server/workconn_scheduler_test.go b/server/workconn_scheduler_test.go new file mode 100644 index 00000000..3d02a966 --- /dev/null +++ b/server/workconn_scheduler_test.go @@ -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) + } + } +}