fix(server): control work connection refill scheduling

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)
@@ -131,6 +132,8 @@ func (ctl *Control) handleReqWorkConn(_ msg.Message) {
 
 	m := &msg.NewWorkConn{
 		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)

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"`

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
+}

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)
 }

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

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()

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 {

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",
+		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{}),

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
+	}
+}

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)
+		}
+	}
+}