fix(server): prevent control replacement lifecycle leaks (#5424)

Release.md

@@ -1,9 +1,3 @@
-## Features
-
-* Expanded the frps dashboard API v2 migration across Clients, Proxies, Server Overview, Client Detail, and Proxy Detail, covering paginated users/clients/proxies, detail data, proxy traffic history, server system info, offline proxy statistics pruning, server-side pagination, search, and proxy type filtering.
-
 ## Fixes
 
-* WebSocket and WSS tunnel payloads are now sent as binary frames, avoiding disconnects through RFC-compliant intermediaries that validate text frames as UTF-8.
+* Fixed control-session replacement leaks when frpc reconnects through a half-open TCP multiplexed connection.
-* The `tls2raw` client plugin now writes the proxy protocol header to the local raw connection when proxy protocol is enabled.
-* frpc now rejects duplicate proxy and visitor names in config files instead of silently overwriting earlier entries.

server/control.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"net"
 	"runtime/debug"
 	"sync"
 	"sync/atomic"
@@ -40,55 +41,311 @@ import (
 	"github.com/fatedier/frp/server/registry"
 )
 
+type ControlID uint64
+
+var nextControlID atomic.Uint64
+
+type controlEntry struct {
+	ctl *Control
+	id  ControlID
+	// runMu serializes lifecycle and routing decisions for one run ID.
+	// Replacements inherit it; removing the entry releases the manager's reference.
+	runMu *sync.Mutex
+
+	registryOnline    bool
+	registryControlID ControlID
+}
+
 type ControlManager struct {
 	// controls indexed by run id
-	ctlsByRunID map[string]*Control
+	ctlsByRunID map[string]*controlEntry
+	registry    *registry.ClientRegistry
+	closed      bool
 
 	mu sync.RWMutex
 }
 
-func NewControlManager() *ControlManager {
+func NewControlManager(clientRegistry *registry.ClientRegistry) *ControlManager {
 	return &ControlManager{
-		ctlsByRunID: make(map[string]*Control),
+		ctlsByRunID: make(map[string]*controlEntry),
+		registry:    clientRegistry,
 	}
 }
 
-func (cm *ControlManager) Add(runID string, ctl *Control) (old *Control) {
+// lockCurrentRun returns the current entry with its run gate held. It never
-	cm.mu.Lock()
+// waits for the gate while holding cm.mu and revalidates the gate after waiting.
-	defer cm.mu.Unlock()
+// The global order is runMu, cm.mu, ctl.lifecycleMu, then registry locks.
-
+func (cm *ControlManager) lockCurrentRun(runID string, allowClosed bool) (*controlEntry, bool) {
-	var ok bool
+	cm.mu.RLock()
-	old, ok = cm.ctlsByRunID[runID]
+	entry, ok := cm.ctlsByRunID[runID]
-	if ok {
+	if cm.closed && !allowClosed {
-		old.Replaced(ctl)
+		ok = false
 	}
-	cm.ctlsByRunID[runID] = ctl
+	cm.mu.RUnlock()
-	return
+	if !ok {
+		return nil, false
+	}
+
+	runMu := entry.runMu
+	runMu.Lock()
+	cm.mu.RLock()
+	entry, ok = cm.ctlsByRunID[runID]
+	if (cm.closed && !allowClosed) || !ok || entry.runMu != runMu {
+		ok = false
+	}
+	cm.mu.RUnlock()
+	if !ok {
+		runMu.Unlock()
+		return nil, false
+	}
+	return entry, true
 }
 
-// we should make sure if it's the same control to prevent delete a new one
+// Add makes ctl the pending current generation and records the predecessor
-func (cm *ControlManager) Del(runID string, ctl *Control) {
+// finalization barrier it must wait for before activation.
+func (cm *ControlManager) Add(ctl *Control) error {
+	for {
+		// Never wait for a run gate while holding cm.mu.
+		cm.mu.RLock()
+		old := cm.ctlsByRunID[ctl.runID]
+		cm.mu.RUnlock()
+		if old != nil {
+			old.runMu.Lock()
+		}
+
+		cm.mu.Lock()
+		if cm.closed {
+			cm.mu.Unlock()
+			if old != nil {
+				old.runMu.Unlock()
+			}
+			return fmt.Errorf("control manager is closed")
+		}
+		if cm.ctlsByRunID[ctl.runID] != old {
+			cm.mu.Unlock()
+			if old != nil {
+				old.runMu.Unlock()
+			}
+			continue
+		}
+
+		id := ControlID(nextControlID.Add(1))
+		if err := ctl.admit(cm, id); err != nil {
+			cm.mu.Unlock()
+			if old != nil {
+				old.runMu.Unlock()
+			}
+			return err
+		}
+
+		runMu := &sync.Mutex{}
+		if old != nil {
+			runMu = old.runMu
+		}
+		entry := &controlEntry{ctl: ctl, id: id, runMu: runMu}
+		var (
+			oldCtl  *Control
+			barrier <-chan struct{}
+		)
+		if old != nil {
+			oldCtl = old.ctl
+			barrier = oldCtl.markReplaced()
+			ctl.setHandoffBarrier(barrier)
+			entry.registryOnline = old.registryOnline
+			entry.registryControlID = old.registryControlID
+		}
+		cm.ctlsByRunID[ctl.runID] = entry
+		cm.mu.Unlock()
+		if old != nil {
+			old.runMu.Unlock()
+		}
+
+		if oldCtl != nil {
+			oldCtl.Replaced(ctl)
+		}
+		return nil
+	}
+}
+
+// Activate registers ctl as online only if it is still the pending current
+// generation.
+func (cm *ControlManager) Activate(ctl *Control) (bool, error) {
+	entry, ok := cm.lockCurrentRun(ctl.runID, false)
+	if !ok {
+		return false, nil
+	}
+	defer entry.runMu.Unlock()
 	cm.mu.Lock()
 	defer cm.mu.Unlock()
-	if c, ok := cm.ctlsByRunID[runID]; ok && c == ctl {
+
-		delete(cm.ctlsByRunID, runID)
+	if cm.closed || cm.ctlsByRunID[ctl.runID] != entry || entry.ctl != ctl || entry.id != ctl.controlID {
+		return false, nil
 	}
+
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	if ctl.state != controlStatePending {
+		return false, nil
+	}
+	if ctl.activated {
+		return true, nil
+	}
+
+	loginMsg := ctl.sessionCtx.LoginMsg
+	remoteAddr := ctl.sessionCtx.Conn.RemoteAddr().String()
+	if host, _, err := net.SplitHostPort(remoteAddr); err == nil {
+		remoteAddr = host
+	}
+	_, conflict := cm.registry.RegisterWithControlID(
+		loginMsg.User,
+		loginMsg.ClientID,
+		ctl.runID,
+		loginMsg.Hostname,
+		loginMsg.Version,
+		remoteAddr,
+		ctl.sessionCtx.WireProtocol,
+		uint64(entry.id),
+	)
+	if conflict {
+		return true, fmt.Errorf("client_id [%s] for user [%s] is already online", loginMsg.ClientID, loginMsg.User)
+	}
+
+	entry.registryOnline = true
+	entry.registryControlID = entry.id
+	ctl.activated = true
+	return true, nil
+}
+
+// completeLogin reserves ctl's current ownership with its run gate while the
+// bounded successful LoginResp write runs, then transitions it to running.
+// The callback must only perform that bounded write; it must not call back into
+// the control manager or the same control lifecycle.
+func (cm *ControlManager) completeLogin(ctl *Control, writeSuccess func() error) (bool, error) {
+	entry, ok := cm.lockCurrentRun(ctl.runID, false)
+	if !ok {
+		return false, nil
+	}
+	defer entry.runMu.Unlock()
+	if entry.ctl != ctl || entry.id != ctl.controlID {
+		return false, nil
+	}
+
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	if ctl.state != controlStatePending || !ctl.activated {
+		return false, nil
+	}
+	if err := writeSuccess(); err != nil {
+		return false, err
+	}
+	if !ctl.startLocked() {
+		return false, nil
+	}
+	return true, nil
+}
+
+// Remove deletes and offlines ctl only if it is still the current generation.
+func (cm *ControlManager) Remove(ctl *Control) bool {
+	entry, ok := cm.lockCurrentRun(ctl.runID, true)
+	if !ok {
+		return false
+	}
+	defer entry.runMu.Unlock()
+	cm.mu.Lock()
+	defer cm.mu.Unlock()
+
+	if cm.ctlsByRunID[ctl.runID] != entry || entry.ctl != ctl || entry.id != ctl.controlID {
+		return false
+	}
+	delete(cm.ctlsByRunID, ctl.runID)
+	if entry.registryOnline {
+		cm.registry.MarkOfflineByRunIDAndControlID(ctl.runID, uint64(entry.registryControlID))
+	}
+	return true
 }
 
 func (cm *ControlManager) GetByID(runID string) (ctl *Control, ok bool) {
-	cm.mu.RLock()
+	entry, ok := cm.lockCurrentRun(runID, false)
-	defer cm.mu.RUnlock()
+	if !ok {
-	ctl, ok = cm.ctlsByRunID[runID]
+		return nil, false
-	return
+	}
+	defer entry.runMu.Unlock()
+	ctl = entry.ctl
+
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	if ctl.state != controlStateRunning {
+		return nil, false
+	}
+	return ctl, true
+}
+
+// admitVisitorByRunID commits a visitor admission against the current running
+// control while its run and lifecycle ownership are held. The callback must
+// only perform the in-memory, buffered visitor admission.
+func (cm *ControlManager) admitVisitorByRunID(runID string, admit func(user string) error) (bool, error) {
+	entry, ok := cm.lockCurrentRun(runID, false)
+	if !ok {
+		return false, nil
+	}
+	defer entry.runMu.Unlock()
+	ctl := entry.ctl
+
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	if ctl.state != controlStateRunning {
+		return false, nil
+	}
+	return true, admit(ctl.sessionCtx.LoginMsg.User)
+}
+
+// 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 {
+	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("client control for run id [%s] is no longer current", 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)
+	}
+
+	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)
+	}
+
+	select {
+	case ctl.workConnCh <- conn:
+		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")
+	}
 }
 
 func (cm *ControlManager) Close() error {
 	cm.mu.Lock()
-	defer cm.mu.Unlock()
+	cm.closed = true
-	for _, ctl := range cm.ctlsByRunID {
+	ctls := make([]*Control, 0, len(cm.ctlsByRunID))
-		ctl.Close()
+	for _, entry := range cm.ctlsByRunID {
+		ctls = append(ctls, entry.ctl)
+	}
+	cm.mu.Unlock()
+
+	for _, ctl := range ctls {
+		cm.Remove(ctl)
+		_ = ctl.Close()
 	}
-	cm.ctlsByRunID = make(map[string]*Control)
 	return nil
 }
 
@@ -110,12 +367,20 @@ type SessionContext struct {
 	LoginMsg *msg.Login
 	// server configuration
 	ServerCfg *v1.ServerConfig
-	// client registry
-	ClientRegistry *registry.ClientRegistry
 	// negotiated wire protocol for this client session
 	WireProtocol string
 }
 
+type controlState uint8
+
+const (
+	controlStateCreated controlState = iota
+	controlStatePending
+	controlStateRunning
+	controlStateClosing
+	controlStateClosed
+)
+
 type Control struct {
 	// session context
 	sessionCtx *SessionContext
@@ -142,30 +407,42 @@ type Control struct {
 	// last time got the Ping message
 	lastPing atomic.Value
 
-	// A new run id will be generated when a new client login.
+	// runID never changes during the lifetime of a control. controlID is assigned
-	// If run id got from login message has same run id, it means it's the same client, so we can
+	// once by ControlManager and distinguishes same-runID generations.
-	// replace old controller instantly.
+	runID     string
-	runID string
+	controlID ControlID
+	manager   *ControlManager
+
+	lifecycleMu    sync.Mutex
+	state          controlState
+	activated      bool
+	handoffBarrier <-chan struct{}
+
+	interruptOnce sync.Once
+	interruptErr  error
 
 	mu sync.RWMutex
 
-	xl     *xlog.Logger
+	xl            *xlog.Logger
-	ctx    context.Context
+	ctx           context.Context
-	doneCh chan struct{}
+	doneCh        chan struct{}
+	serverMetrics metrics.ServerMetrics
 }
 
 func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, error) {
 	poolCount := min(sessionCtx.LoginMsg.PoolCount, int(sessionCtx.ServerCfg.Transport.MaxPoolCount))
 	ctl := &Control{
-		sessionCtx:   sessionCtx,
+		sessionCtx:    sessionCtx,
-		workConnCh:   make(chan *proxy.WorkConn, poolCount+10),
+		workConnCh:    make(chan *proxy.WorkConn, poolCount+10),
-		proxies:      make(map[string]proxy.Proxy),
+		proxies:       make(map[string]proxy.Proxy),
-		poolCount:    poolCount,
+		poolCount:     poolCount,
-		portsUsedNum: 0,
+		portsUsedNum:  0,
-		runID:        sessionCtx.LoginMsg.RunID,
+		runID:         sessionCtx.LoginMsg.RunID,
-		xl:           xlog.FromContextSafe(ctx),
+		state:         controlStateCreated,
-		ctx:          ctx,
+		xl:            xlog.FromContextSafe(ctx),
-		doneCh:       make(chan struct{}),
+		ctx:           ctx,
+		doneCh:        make(chan struct{}),
+		serverMetrics: metrics.Server,
 	}
 	ctl.lastPing.Store(time.Now())
 
@@ -175,48 +452,121 @@ func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, erro
 	return ctl, nil
 }
 
-// Start starts the control session workers after login succeeds.
+func (ctl *Control) RunID() string {
-func (ctl *Control) Start() {
+	return ctl.runID
-	go func() {
-		for i := 0; i < ctl.poolCount; i++ {
-			// ignore error here, that means that this control is closed
-			_ = ctl.msgDispatcher.Send(&msg.ReqWorkConn{})
-		}
-	}()
-	go ctl.worker()
 }
 
-func (ctl *Control) Close() error {
+func (ctl *Control) ID() ControlID {
-	ctl.sessionCtx.Conn.Close()
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	return ctl.controlID
+}
+
+func (ctl *Control) admit(manager *ControlManager, id ControlID) error {
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	if ctl.state != controlStateCreated {
+		return fmt.Errorf("control [%s] is not in created state", ctl.runID)
+	}
+	ctl.manager = manager
+	ctl.controlID = id
+	ctl.state = controlStatePending
 	return nil
 }
 
-func (ctl *Control) Replaced(newCtl *Control) {
+func (ctl *Control) setHandoffBarrier(barrier <-chan struct{}) {
-	xl := ctl.xl
+	ctl.lifecycleMu.Lock()
-	xl.Infof("replaced by client [%s]", newCtl.runID)
+	ctl.handoffBarrier = barrier
-	ctl.runID = ""
+	ctl.lifecycleMu.Unlock()
-	ctl.sessionCtx.Conn.Close()
 }
 
-func (ctl *Control) RegisterWorkConn(conn *proxy.WorkConn) error {
+func (ctl *Control) WaitForHandoff() {
-	xl := ctl.xl
+	ctl.lifecycleMu.Lock()
-	defer func() {
+	barrier := ctl.handoffBarrier
-		if err := recover(); err != nil {
+	ctl.lifecycleMu.Unlock()
-			xl.Errorf("panic error: %v", err)
+	if barrier != nil {
-			xl.Errorf(string(debug.Stack()))
+		<-barrier
-		}
-	}()
-
-	select {
-	case ctl.workConnCh <- conn:
-		xl.Debugf("new work connection registered")
-		return nil
-	default:
-		xl.Debugf("work connection pool is full, discarding")
-		return fmt.Errorf("work connection pool is full, discarding")
 	}
 }
 
+// Start starts the control session workers after login succeeds.
+func (ctl *Control) Start() bool {
+	ctl.lifecycleMu.Lock()
+	defer ctl.lifecycleMu.Unlock()
+	return ctl.startLocked()
+}
+
+func (ctl *Control) startLocked() bool {
+	if ctl.state != controlStatePending || !ctl.activated {
+		return false
+	}
+	ctl.state = controlStateRunning
+	go ctl.worker()
+	return true
+}
+
+func (ctl *Control) Close() error {
+	ctl.lifecycleMu.Lock()
+	switch ctl.state {
+	case controlStateCreated, controlStatePending:
+		ctl.state = controlStateClosing
+		ctl.finishLocked()
+	case controlStateRunning:
+		ctl.state = controlStateClosing
+	}
+	ctl.lifecycleMu.Unlock()
+	return ctl.interruptReadAndClose()
+}
+
+func (ctl *Control) Replaced(newCtl *Control) {
+	ctl.markReplaced()
+	ctl.xl.Infof("replaced by client [%s] (control ID %d)", newCtl.runID, newCtl.ID())
+	_ = ctl.interruptReadAndClose()
+}
+
+// markReplaced returns the transitive predecessor barrier. A pending control
+// has no worker, so it finishes immediately and passes its inherited barrier
+// 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.finishLocked()
+		return nil
+	case controlStatePending:
+		barrier := ctl.handoffBarrier
+		ctl.state = controlStateClosing
+		ctl.finishLocked()
+		return barrier
+	case controlStateRunning:
+		ctl.state = controlStateClosing
+		return ctl.doneCh
+	case controlStateClosing, controlStateClosed:
+		return ctl.doneCh
+	default:
+		return ctl.doneCh
+	}
+}
+
+func (ctl *Control) interruptReadAndClose() error {
+	ctl.interruptOnce.Do(func() {
+		_ = ctl.sessionCtx.Conn.SetReadDeadline(time.Now())
+		ctl.interruptErr = ctl.sessionCtx.Conn.Close()
+	})
+	return ctl.interruptErr
+}
+
+func (ctl *Control) finishLocked() {
+	if ctl.state == controlStateClosed {
+		return
+	}
+	ctl.state = controlStateClosed
+	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.
@@ -271,10 +621,10 @@ func (ctl *Control) heartbeatWorker() {
 	}
 
 	xl := ctl.xl
-	go wait.Until(func() {
+	wait.Until(func() {
 		if time.Since(ctl.lastPing.Load().(time.Time)) > time.Duration(ctl.sessionCtx.ServerCfg.Transport.HeartbeatTimeout)*time.Second {
 			xl.Warnf("heartbeat timeout")
-			ctl.sessionCtx.Conn.Close()
+			_ = ctl.Close()
 			return
 		}
 	}, time.Second, ctl.doneCh)
@@ -289,14 +639,14 @@ func (ctl *Control) loginUserInfo() plugin.UserInfo {
 	return plugin.UserInfo{
 		User:  ctl.sessionCtx.LoginMsg.User,
 		Metas: ctl.sessionCtx.LoginMsg.Metas,
-		RunID: ctl.sessionCtx.LoginMsg.RunID,
+		RunID: ctl.runID,
 	}
 }
 
 func (ctl *Control) closeProxy(pxy proxy.Proxy) {
 	pxy.Close()
 	ctl.sessionCtx.PxyManager.Del(pxy.GetName())
-	metrics.Server.CloseProxy(pxy.GetName(), pxy.GetConfigurer().GetBaseConfig().Type)
+	ctl.serverMetrics.CloseProxy(pxy.GetName(), pxy.GetConfigurer().GetBaseConfig().Type)
 
 	notifyContent := &plugin.CloseProxyContent{
 		User: ctl.loginUserInfo(),
@@ -311,12 +661,24 @@ func (ctl *Control) closeProxy(pxy proxy.Proxy) {
 
 func (ctl *Control) worker() {
 	xl := ctl.xl
+	ctl.serverMetrics.NewClient()
 
 	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.msgDispatcher.Done()
-	ctl.sessionCtx.Conn.Close()
+	ctl.lifecycleMu.Lock()
+	if ctl.state == controlStateRunning {
+		ctl.state = controlStateClosing
+	}
+	ctl.lifecycleMu.Unlock()
+	_ = ctl.interruptReadAndClose()
 
 	ctl.mu.Lock()
 	close(ctl.workConnCh)
@@ -331,10 +693,14 @@ func (ctl *Control) worker() {
 		ctl.closeProxy(pxy)
 	}
 
-	metrics.Server.CloseClient()
+	ctl.serverMetrics.CloseClient()
-	ctl.sessionCtx.ClientRegistry.MarkOfflineByRunID(ctl.runID)
+	if ctl.manager != nil {
+		ctl.manager.Remove(ctl)
+	}
 	xl.Infof("client exit success")
-	close(ctl.doneCh)
+	ctl.lifecycleMu.Lock()
+	ctl.finishLocked()
+	ctl.lifecycleMu.Unlock()
 }
 
 func (ctl *Control) registerMsgHandlers() {
@@ -374,9 +740,9 @@ func (ctl *Control) handleNewProxy(m msg.Message) {
 		xl.Infof("new proxy [%s] type [%s] success", inMsg.ProxyName, inMsg.ProxyType)
 		clientID := ctl.sessionCtx.LoginMsg.ClientID
 		if clientID == "" {
-			clientID = ctl.sessionCtx.LoginMsg.RunID
+			clientID = ctl.runID
 		}
-		metrics.Server.NewProxy(inMsg.ProxyName, inMsg.ProxyType, ctl.sessionCtx.LoginMsg.User, clientID)
+		ctl.serverMetrics.NewProxy(inMsg.ProxyName, inMsg.ProxyType, ctl.sessionCtx.LoginMsg.User, clientID)
 	}
 	_ = ctl.msgDispatcher.Send(resp)
 }

server/control_test.go

@@ -0,0 +1,543 @@
+// 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 (
+	"context"
+	"errors"
+	"net"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/fatedier/frp/pkg/auth"
+	v1 "github.com/fatedier/frp/pkg/config/v1"
+	"github.com/fatedier/frp/pkg/msg"
+	plugin "github.com/fatedier/frp/pkg/plugin/server"
+	"github.com/fatedier/frp/server/controller"
+	"github.com/fatedier/frp/server/proxy"
+	"github.com/fatedier/frp/server/registry"
+)
+
+func TestControlPendingReplacementFinishesWithoutStarting(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	oldCtl, oldConn := newLifecycleTestControl(t, "same-run", "client", metrics)
+	newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	mustAddAndActivate(t, manager, oldCtl)
+
+	err := manager.Add(newCtl)
+	require.NoError(t, err)
+	waitForControlDone(t, oldCtl)
+	require.False(t, oldCtl.Start())
+	require.Equal(t, []string{"deadline", "close"}, oldConn.eventsSnapshot())
+	require.Equal(t, int64(0), metrics.newClients())
+	require.Equal(t, int64(0), metrics.closedClients())
+}
+
+func TestControlRunningReplacementFinishesInWorker(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	oldCtl, oldConn := newLifecycleTestControl(t, "same-run", "client", metrics)
+	newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	mustAddAndActivate(t, manager, oldCtl)
+	require.True(t, oldCtl.Start())
+	waitForSignal(t, oldConn.readStarted, "control reader to start")
+
+	err := manager.Add(newCtl)
+	require.NoError(t, err)
+	waitForControlDone(t, oldCtl)
+	require.Equal(t, []string{"deadline", "close"}, oldConn.eventsSnapshot())
+	require.Equal(t, int64(1), metrics.newClients())
+	require.Equal(t, int64(1), metrics.closedClients())
+
+	_, ok := manager.GetByID("same-run")
+	require.False(t, ok)
+	require.Same(t, newCtl, currentControlForTest(manager, "same-run"))
+	info, ok := clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.True(t, info.Online)
+	require.Equal(t, uint64(oldCtl.ID()), info.ControlID)
+
+	active, err := manager.Activate(newCtl)
+	require.NoError(t, err)
+	require.True(t, active)
+	_, ok = manager.GetByID("same-run")
+	require.False(t, ok)
+	info, ok = clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.Equal(t, uint64(newCtl.ID()), info.ControlID)
+}
+
+func TestControlClosePendingAndRunning(t *testing.T) {
+	t.Run("pending", func(t *testing.T) {
+		manager := NewControlManager(registry.NewClientRegistry())
+		metrics := newCountingServerMetrics()
+		ctl, conn := newLifecycleTestControl(t, "pending", "pending", metrics)
+		err := manager.Add(ctl)
+		require.NoError(t, err)
+
+		require.NoError(t, ctl.Close())
+		waitForControlDone(t, ctl)
+		require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot())
+		require.Equal(t, int64(0), metrics.newClients())
+		require.Equal(t, int64(0), metrics.closedClients())
+	})
+
+	t.Run("running", func(t *testing.T) {
+		manager := NewControlManager(registry.NewClientRegistry())
+		metrics := newCountingServerMetrics()
+		ctl, conn := newLifecycleTestControl(t, "running", "running", metrics)
+		mustAddAndActivate(t, manager, ctl)
+		require.True(t, ctl.Start())
+		waitForSignal(t, conn.readStarted, "control reader to start")
+
+		require.NoError(t, ctl.Close())
+		waitForControlDone(t, ctl)
+		require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot())
+		require.Equal(t, int64(1), metrics.newClients())
+		require.Equal(t, int64(1), metrics.closedClients())
+	})
+}
+
+func TestControlCloseAndReplacedAreIdempotent(t *testing.T) {
+	manager := NewControlManager(registry.NewClientRegistry())
+	metrics := newCountingServerMetrics()
+	ctl, conn := newLifecycleTestControl(t, "same-run", "client", metrics)
+	replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	err := manager.Add(ctl)
+	require.NoError(t, err)
+	err = manager.Add(replacement)
+	require.NoError(t, err)
+	require.NoError(t, ctl.Close())
+	ctl.Replaced(replacement)
+	require.NoError(t, ctl.Close())
+	waitForControlDone(t, ctl)
+
+	require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot())
+	require.Equal(t, int64(0), metrics.newClients())
+	require.Equal(t, int64(0), metrics.closedClients())
+}
+
+func TestControlHeartbeatTimeoutInterruptsRead(t *testing.T) {
+	manager := NewControlManager(registry.NewClientRegistry())
+	metrics := newCountingServerMetrics()
+	ctl, conn := newLifecycleTestControl(t, "heartbeat", "heartbeat", metrics)
+	ctl.sessionCtx.ServerCfg.Transport.HeartbeatTimeout = 1
+	ctl.lastPing.Store(time.Now().Add(-2 * time.Second))
+
+	mustAddAndActivate(t, manager, ctl)
+	require.True(t, ctl.Start())
+	waitForSignal(t, conn.readStarted, "control reader to start")
+	waitForControlDone(t, ctl)
+
+	require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot())
+	require.Equal(t, int64(1), metrics.newClients())
+	require.Equal(t, int64(1), metrics.closedClients())
+}
+
+func TestControlStartReplacementRacePairsMetrics(t *testing.T) {
+	for range 100 {
+		clientRegistry := registry.NewClientRegistry()
+		manager := NewControlManager(clientRegistry)
+		metrics := newCountingServerMetrics()
+		ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+		replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+		mustAddAndActivate(t, manager, ctl)
+
+		startGate := make(chan struct{})
+		startedCh := make(chan bool, 1)
+		addErrCh := make(chan error, 1)
+		go func() {
+			<-startGate
+			startedCh <- ctl.Start()
+		}()
+		go func() {
+			<-startGate
+			addErr := manager.Add(replacement)
+			addErrCh <- addErr
+		}()
+		close(startGate)
+
+		started := <-startedCh
+		require.NoError(t, <-addErrCh)
+		waitForControlDone(t, ctl)
+		if started {
+			require.Equal(t, int64(1), metrics.newClients())
+			require.Equal(t, int64(1), metrics.closedClients())
+		} else {
+			require.Equal(t, int64(0), metrics.newClients())
+			require.Equal(t, int64(0), metrics.closedClients())
+		}
+	}
+}
+
+func TestControlManagerRejectsStaleActivateAndRemove(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	oldCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+	newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	mustAddAndActivate(t, manager, oldCtl)
+	err := manager.Add(newCtl)
+	require.NoError(t, err)
+	require.Greater(t, uint64(newCtl.ID()), uint64(oldCtl.ID()))
+
+	active, err := manager.Activate(oldCtl)
+	require.NoError(t, err)
+	require.False(t, active)
+	require.False(t, manager.Remove(oldCtl))
+
+	_, ok := manager.GetByID("same-run")
+	require.False(t, ok)
+	require.Same(t, newCtl, currentControlForTest(manager, "same-run"))
+	info, ok := clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.True(t, info.Online)
+	require.Equal(t, uint64(oldCtl.ID()), info.ControlID)
+
+	active, err = manager.Activate(newCtl)
+	require.NoError(t, err)
+	require.True(t, active)
+	info, ok = clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.True(t, info.Online)
+	require.Equal(t, uint64(newCtl.ID()), info.ControlID)
+}
+
+func TestControlManagerPreservesClientIDConflict(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	first, _ := newLifecycleTestControl(t, "run-one", "shared-client", metrics)
+	conflicting, _ := newLifecycleTestControl(t, "run-two", "shared-client", metrics)
+
+	mustAddAndActivate(t, manager, first)
+	err := manager.Add(conflicting)
+	require.NoError(t, err)
+	active, err := manager.Activate(conflicting)
+	require.True(t, active)
+	require.ErrorContains(t, err, "already online")
+
+	require.True(t, manager.Remove(conflicting))
+	info, ok := clientRegistry.GetByKey("shared-client")
+	require.True(t, ok)
+	require.True(t, info.Online)
+	require.Equal(t, "run-one", info.RunID)
+}
+
+func TestControlManagerFailedLoginWriteReleasesRunWithoutStarting(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+	replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	mustAddAndActivate(t, manager, ctl)
+
+	writeErr := errors.New("write failed")
+	committed, err := manager.completeLogin(ctl, func() error { return writeErr })
+	require.ErrorIs(t, err, writeErr)
+	require.False(t, committed)
+
+	err = manager.Add(replacement)
+	require.NoError(t, err)
+	waitForControlDone(t, ctl)
+	require.Same(t, replacement, currentControlForTest(manager, "same-run"))
+	require.Equal(t, int64(0), metrics.newClients())
+	require.Equal(t, int64(0), metrics.closedClients())
+	require.True(t, manager.Remove(replacement))
+	info, ok := clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.False(t, info.Online)
+	require.Empty(t, info.RunID)
+	require.Zero(t, info.ControlID)
+	require.False(t, info.DisconnectedAt.IsZero())
+	require.NoError(t, replacement.Close())
+}
+
+func TestControlManagerCloseWaitsForInFlightLoginRun(t *testing.T) {
+	clientRegistry := registry.NewClientRegistry()
+	manager := NewControlManager(clientRegistry)
+	metrics := newCountingServerMetrics()
+	ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics)
+
+	mustAddAndActivate(t, manager, ctl)
+
+	writeEntered := make(chan struct{})
+	resumeWrite := make(chan struct{})
+	loginDone := make(chan struct {
+		committed bool
+		err       error
+	}, 1)
+	go func() {
+		committed, loginErr := manager.completeLogin(ctl, func() error {
+			close(writeEntered)
+			<-resumeWrite
+			return nil
+		})
+		loginDone <- struct {
+			committed bool
+			err       error
+		}{committed: committed, err: loginErr}
+	}()
+	waitForSignal(t, writeEntered, "LoginResp write")
+
+	closeDone := make(chan error, 1)
+	go func() { closeDone <- manager.Close() }()
+	waitForManagerClosed(t, manager)
+	select {
+	case err := <-closeDone:
+		t.Fatalf("manager close completed during LoginResp write: %v", err)
+	default:
+	}
+
+	close(resumeWrite)
+	result := <-loginDone
+	require.NoError(t, result.err)
+	require.True(t, result.committed)
+	require.NoError(t, <-closeDone)
+	waitForControlDone(t, ctl)
+	require.Nil(t, currentControlForTest(manager, "same-run"))
+	require.Equal(t, int64(1), metrics.newClients())
+	require.Equal(t, int64(1), metrics.closedClients())
+	info, ok := clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.False(t, info.Online)
+}
+
+func newLifecycleTestControl(
+	t *testing.T,
+	runID string,
+	clientID string,
+	serverMetrics *countingServerMetrics,
+) (*Control, *deadlineReadConn) {
+	t.Helper()
+	conn := newDeadlineReadConn()
+	msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
+	ctl, err := NewControl(context.Background(), &SessionContext{
+		RC:            &controller.ResourceController{},
+		PxyManager:    proxy.NewManager(),
+		PluginManager: plugin.NewManager(),
+		AuthVerifier:  auth.AlwaysPassVerifier,
+		Conn:          msgConn,
+		LoginMsg: &msg.Login{
+			RunID:    runID,
+			ClientID: clientID,
+		},
+		ServerCfg: &v1.ServerConfig{},
+	})
+	require.NoError(t, err)
+	ctl.serverMetrics = serverMetrics
+	t.Cleanup(func() { _ = ctl.Close() })
+	return ctl, conn
+}
+
+func mustAddAndActivate(t *testing.T, manager *ControlManager, ctl *Control) {
+	t.Helper()
+	require.NoError(t, manager.Add(ctl))
+	active, err := manager.Activate(ctl)
+	require.NoError(t, err)
+	require.True(t, active)
+}
+
+func waitForControlDone(t *testing.T, ctl *Control) {
+	t.Helper()
+	done := make(chan struct{})
+	go func() {
+		ctl.WaitClosed()
+		close(done)
+	}()
+	waitForSignal(t, done, "control to finish")
+}
+
+func currentControlForTest(manager *ControlManager, runID string) *Control {
+	manager.mu.RLock()
+	defer manager.mu.RUnlock()
+	entry := manager.ctlsByRunID[runID]
+	if entry == nil {
+		return nil
+	}
+	return entry.ctl
+}
+
+func currentRunGateForTest(manager *ControlManager, runID string) *sync.Mutex {
+	manager.mu.RLock()
+	defer manager.mu.RUnlock()
+	entry := manager.ctlsByRunID[runID]
+	if entry == nil {
+		return nil
+	}
+	return entry.runMu
+}
+
+func waitForManagerClosed(t *testing.T, manager *ControlManager) {
+	t.Helper()
+	deadline := time.Now().Add(3 * time.Second)
+	for time.Now().Before(deadline) {
+		manager.mu.RLock()
+		closed := manager.closed
+		manager.mu.RUnlock()
+		if closed {
+			return
+		}
+	}
+	t.Fatal("timed out waiting for control manager to close")
+}
+
+func waitForSignal(t *testing.T, ch <-chan struct{}, description string) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(3 * time.Second):
+		t.Fatalf("timed out waiting for %s", description)
+	}
+}
+
+type deadlineReadConn struct {
+	readStarted chan struct{}
+	unblockRead chan struct{}
+
+	readOnce     sync.Once
+	unblockOnce  sync.Once
+	deadlineOnce sync.Once
+	closeOnce    sync.Once
+
+	eventsMu sync.Mutex
+	events   []string
+}
+
+func newDeadlineReadConn() *deadlineReadConn {
+	return &deadlineReadConn{
+		readStarted: make(chan struct{}),
+		unblockRead: make(chan struct{}),
+	}
+}
+
+func (c *deadlineReadConn) Read([]byte) (int, error) {
+	c.readOnce.Do(func() { close(c.readStarted) })
+	<-c.unblockRead
+	return 0, os.ErrDeadlineExceeded
+}
+
+func (*deadlineReadConn) Write(p []byte) (int, error) { return len(p), nil }
+
+func (c *deadlineReadConn) Close() error {
+	c.closeOnce.Do(func() {
+		c.recordEvent("close")
+		c.unblockOnce.Do(func() { close(c.unblockRead) })
+	})
+	return nil
+}
+
+func (*deadlineReadConn) LocalAddr() net.Addr  { return lifecycleTestAddr("local") }
+func (*deadlineReadConn) RemoteAddr() net.Addr { return lifecycleTestAddr("remote") }
+
+func (c *deadlineReadConn) SetDeadline(deadline time.Time) error {
+	if err := c.SetReadDeadline(deadline); err != nil {
+		return err
+	}
+	return c.SetWriteDeadline(deadline)
+}
+
+func (c *deadlineReadConn) SetReadDeadline(deadline time.Time) error {
+	if deadline.IsZero() {
+		return nil
+	}
+	c.deadlineOnce.Do(func() {
+		c.recordEvent("deadline")
+		c.unblockOnce.Do(func() { close(c.unblockRead) })
+	})
+	return nil
+}
+
+func (*deadlineReadConn) SetWriteDeadline(time.Time) error { return nil }
+
+func (c *deadlineReadConn) recordEvent(event string) {
+	c.eventsMu.Lock()
+	c.events = append(c.events, event)
+	c.eventsMu.Unlock()
+}
+
+func (c *deadlineReadConn) eventsSnapshot() []string {
+	c.eventsMu.Lock()
+	defer c.eventsMu.Unlock()
+	return append([]string(nil), c.events...)
+}
+
+type lifecycleTestAddr string
+
+func (a lifecycleTestAddr) Network() string { return string(a) }
+func (a lifecycleTestAddr) String() string  { return string(a) }
+
+type countingServerMetrics struct {
+	mu          sync.Mutex
+	newCount    int64
+	closeCount  int64
+	closeEnter  chan struct{}
+	closeResume chan struct{}
+	closeOnce   sync.Once
+}
+
+func newCountingServerMetrics() *countingServerMetrics {
+	return &countingServerMetrics{}
+}
+
+func (m *countingServerMetrics) NewClient() {
+	m.mu.Lock()
+	m.newCount++
+	m.mu.Unlock()
+}
+
+func (m *countingServerMetrics) CloseClient() {
+	m.mu.Lock()
+	m.closeCount++
+	closeEnter := m.closeEnter
+	closeResume := m.closeResume
+	m.mu.Unlock()
+	if closeEnter != nil {
+		m.closeOnce.Do(func() { close(closeEnter) })
+		<-closeResume
+	}
+}
+
+func (*countingServerMetrics) NewProxy(string, string, string, string) {}
+func (*countingServerMetrics) CloseProxy(string, string)               {}
+func (*countingServerMetrics) OpenConnection(string, string)           {}
+func (*countingServerMetrics) CloseConnection(string, string)          {}
+func (*countingServerMetrics) AddTrafficIn(string, string, int64)      {}
+func (*countingServerMetrics) AddTrafficOut(string, string, int64)     {}
+
+func (m *countingServerMetrics) newClients() int64 {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.newCount
+}
+
+func (m *countingServerMetrics) closedClients() int64 {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.closeCount
+}

server/registry/registry.go

@@ -28,6 +28,7 @@ type ClientInfo struct {
 	User             string
 	RawClientID      string
 	RunID            string
+	ControlID        uint64
 	Hostname         string
 	IP               string
 	Version          string
@@ -64,6 +65,16 @@ func newClientRegistryWithClock(clk clock.PassiveClock) *ClientRegistry {
 
 // Register stores/updates metadata for a client and returns the registry key plus whether it conflicts with an online client.
 func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol string) (key string, conflict bool) {
+	return cr.RegisterWithControlID(user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol, 0)
+}
+
+// RegisterWithControlID is the generation-aware form used by ControlManager.
+// A control ID is process-local and prevents an older control generation from
+// changing the registry entry now owned by a newer generation with the same run ID.
+func (cr *ClientRegistry) RegisterWithControlID(
+	user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol string,
+	controlID uint64,
+) (key string, conflict bool) {
 	if runID == "" {
 		return "", false
 	}
@@ -83,6 +94,16 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
 	if enforceUnique && exists && info.Online && info.RunID != "" && info.RunID != runID {
 		return key, true
 	}
+	if previousKey, ok := cr.runIndex[runID]; ok && previousKey != key {
+		if previous, ok := cr.clients[previousKey]; ok && previous.RunID == runID {
+			if previous.RawClientID == "" {
+				delete(cr.clients, previousKey)
+			} else {
+				setClientOffline(previous, now)
+			}
+		}
+		delete(cr.runIndex, runID)
+	}
 
 	if !exists {
 		info = &ClientInfo{
@@ -97,6 +118,7 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
 
 	info.RawClientID = rawClientID
 	info.RunID = runID
+	info.ControlID = controlID
 	info.Hostname = hostname
 	info.IP = remoteAddr
 	info.Version = version
@@ -114,6 +136,16 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
 
 // MarkOfflineByRunID marks the client as offline when the corresponding control disconnects.
 func (cr *ClientRegistry) MarkOfflineByRunID(runID string) {
+	cr.markOfflineByRunID(runID, 0, false)
+}
+
+// MarkOfflineByRunIDAndControlID marks a client offline only when the registry
+// entry still belongs to the supplied control generation.
+func (cr *ClientRegistry) MarkOfflineByRunIDAndControlID(runID string, controlID uint64) {
+	cr.markOfflineByRunID(runID, controlID, true)
+}
+
+func (cr *ClientRegistry) markOfflineByRunID(runID string, controlID uint64, matchControlID bool) {
 	cr.mu.Lock()
 	defer cr.mu.Unlock()
 
@@ -121,17 +153,23 @@ func (cr *ClientRegistry) MarkOfflineByRunID(runID string) {
 	if !ok {
 		return
 	}
-	if info, ok := cr.clients[key]; ok && info.RunID == runID {
+	if info, ok := cr.clients[key]; ok && info.RunID == runID && (!matchControlID || info.ControlID == controlID) {
 		if info.RawClientID == "" {
 			delete(cr.clients, key)
 		} else {
-			info.RunID = ""
+			setClientOffline(info, cr.clock.Now())
-			info.Online = false
-			now := cr.clock.Now()
-			info.DisconnectedAt = now
 		}
 	}
-	delete(cr.runIndex, runID)
+	if info, ok := cr.clients[key]; !ok || info.RunID != runID {
+		delete(cr.runIndex, runID)
+	}
+}
+
+func setClientOffline(info *ClientInfo, now time.Time) {
+	info.RunID = ""
+	info.ControlID = 0
+	info.Online = false
+	info.DisconnectedAt = now
 }
 
 // List returns a snapshot of all known clients.

server/registry/registry_test.go

@@ -72,3 +72,89 @@ func TestClientRegistryUsesClockForTimestamps(t *testing.T) {
 		t.Fatalf("disconnected time mismatch, want %s got %s", disconnectedAt, info.DisconnectedAt)
 	}
 }
+
+func TestClientRegistryControlIDPreventsStaleOffline(t *testing.T) {
+	registry := NewClientRegistry()
+	key, conflict := registry.RegisterWithControlID(
+		"user", "client-id", "run-id", "old-host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
+	)
+	if conflict {
+		t.Fatal("unexpected client conflict")
+	}
+	_, conflict = registry.RegisterWithControlID(
+		"user", "client-id", "run-id", "new-host", "1.0.1", "127.0.0.2", wire.ProtocolV2, 2,
+	)
+	if conflict {
+		t.Fatal("same run ID replacement should not conflict")
+	}
+
+	registry.MarkOfflineByRunIDAndControlID("run-id", 1)
+	info, ok := registry.GetByKey(key)
+	if !ok {
+		t.Fatalf("client %q not found", key)
+	}
+	if !info.Online || info.ControlID != 2 || info.Hostname != "new-host" {
+		t.Fatalf("stale offline changed current generation: %+v", info)
+	}
+
+	registry.MarkOfflineByRunIDAndControlID("run-id", 2)
+	info, ok = registry.GetByKey(key)
+	if !ok {
+		t.Fatalf("client %q not found after disconnect", key)
+	}
+	if info.Online || info.ControlID != 0 || info.RunID != "" {
+		t.Fatalf("current generation was not marked offline: %+v", info)
+	}
+}
+
+func TestClientRegistryClientIDConflictSemantics(t *testing.T) {
+	registry := NewClientRegistry()
+	_, conflict := registry.RegisterWithControlID(
+		"user", "client-id", "run-one", "host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
+	)
+	if conflict {
+		t.Fatal("unexpected initial client conflict")
+	}
+	_, conflict = registry.RegisterWithControlID(
+		"user", "client-id", "run-two", "host", "1.0.0", "127.0.0.2", wire.ProtocolV1, 2,
+	)
+	if !conflict {
+		t.Fatal("different online run IDs with the same explicit client ID must conflict")
+	}
+
+	registry.MarkOfflineByRunIDAndControlID("run-one", 1)
+	_, conflict = registry.RegisterWithControlID(
+		"user", "client-id", "run-two", "host", "1.0.0", "127.0.0.2", wire.ProtocolV1, 2,
+	)
+	if conflict {
+		t.Fatal("offline explicit client ID should be reusable")
+	}
+}
+
+func TestClientRegistrySameRunIDMovesBetweenClientKeys(t *testing.T) {
+	registry := NewClientRegistry()
+	oldKey, conflict := registry.RegisterWithControlID(
+		"user", "old-client", "run-id", "old-host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
+	)
+	if conflict {
+		t.Fatal("unexpected initial client conflict")
+	}
+	newKey, conflict := registry.RegisterWithControlID(
+		"user", "new-client", "run-id", "new-host", "1.0.1", "127.0.0.2", wire.ProtocolV2, 2,
+	)
+	if conflict {
+		t.Fatal("same run ID moving to a new client key should not conflict")
+	}
+
+	oldInfo, ok := registry.GetByKey(oldKey)
+	if !ok {
+		t.Fatalf("old explicit client %q should remain as offline history", oldKey)
+	}
+	if oldInfo.Online || oldInfo.RunID != "" || oldInfo.ControlID != 0 {
+		t.Fatalf("old client key remained online: %+v", oldInfo)
+	}
+	newInfo, ok := registry.GetByKey(newKey)
+	if !ok || !newInfo.Online || newInfo.RunID != "run-id" || newInfo.ControlID != 2 {
+		t.Fatalf("new client key was not registered: %+v", newInfo)
+	}
+}

server/service.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -51,7 +52,6 @@ import (
 	"github.com/fatedier/frp/pkg/util/xlog"
 	"github.com/fatedier/frp/server/controller"
 	"github.com/fatedier/frp/server/group"
-	"github.com/fatedier/frp/server/metrics"
 	"github.com/fatedier/frp/server/ports"
 	"github.com/fatedier/frp/server/proxy"
 	"github.com/fatedier/frp/server/registry"
@@ -64,6 +64,8 @@ const (
 	vhostReadWriteTimeout time.Duration = 30 * time.Second
 )
 
+var errControlReplaced = errors.New("control was replaced during login")
+
 func init() {
 	crypto.DefaultSalt = "frp"
 	// Disable quic-go's receive buffer warning.
@@ -161,9 +163,10 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) {
 		return nil, err
 	}
 
+	clientRegistry := registry.NewClientRegistry()
 	svr := &Service{
-		ctlManager:     NewControlManager(),
+		ctlManager:     NewControlManager(clientRegistry),
-		clientRegistry: registry.NewClientRegistry(),
+		clientRegistry: clientRegistry,
 		pxyManager:     proxy.NewManager(),
 		pluginManager:  plugin.NewManager(),
 		rc: &controller.ResourceController{
@@ -469,6 +472,9 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
 
 		if err != nil {
 			xl.Warnf("register control error: %v", err)
+			if ctl != nil {
+				svr.ctlManager.Remove(ctl)
+			}
 			if writeErr := writeWithDeadline(conn, connWriteTimeout, func() error {
 				return acceptedConn.conn.WriteMsg(&msg.LoginResp{
 					Version: version.Full(),
@@ -477,29 +483,27 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
 			}); writeErr != nil {
 				xl.Warnf("write login error response error: %v", writeErr)
 			}
-			conn.Close()
+			if ctl != nil {
+				_ = ctl.Close()
+			} else {
+				conn.Close()
+			}
 			return
 		}
-		if err = writeWithDeadline(conn, connWriteTimeout, func() error {
+		if err = svr.completeControlLogin(ctl, func() error {
-			return acceptedConn.conn.WriteMsg(&msg.LoginResp{
+			return writeWithDeadline(conn, connWriteTimeout, func() error {
-				Version: version.Full(),
+				return acceptedConn.conn.WriteMsg(&msg.LoginResp{
-				RunID:   ctl.runID,
+					Version: version.Full(),
-				Error:   "",
+					RunID:   ctl.runID,
+					Error:   "",
+				})
 			})
 		}); err != nil {
-			xl.Warnf("write login response error: %v", err)
+			xl.Warnf("complete control login error: %v", err)
-			svr.ctlManager.Del(m.RunID, ctl)
+			svr.ctlManager.Remove(ctl)
-			svr.clientRegistry.MarkOfflineByRunID(m.RunID)
+			_ = ctl.Close()
-			conn.Close()
 			return
 		}
-		ctl.Start()
-		metrics.Server.NewClient()
-		go func() {
-			// block until control closed
-			ctl.WaitClosed()
-			svr.ctlManager.Del(m.RunID, ctl)
-		}()
 	case *msg.NewWorkConn:
 		if err := svr.RegisterWorkConn(acceptedConn.conn, m); err != nil {
 			_ = acceptedConn.conn.WriteMsg(&msg.StartWorkConn{
@@ -527,6 +531,17 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
 	}
 }
 
+func (svr *Service) completeControlLogin(ctl *Control, writeSuccess func() error) error {
+	committed, err := svr.ctlManager.completeLogin(ctl, writeSuccess)
+	if err != nil {
+		return err
+	}
+	if !committed {
+		return errControlReplaced
+	}
+	return nil
+}
+
 type acceptedConnection struct {
 	conn          *msg.Conn
 	wireProtocol  string
@@ -768,16 +783,15 @@ func (svr *Service) RegisterControl(
 	}
 
 	ctl, err := NewControl(ctx, &SessionContext{
-		RC:             svr.rc,
+		RC:            svr.rc,
-		PxyManager:     svr.pxyManager,
+		PxyManager:    svr.pxyManager,
-		PluginManager:  svr.pluginManager,
+		PluginManager: svr.pluginManager,
-		AuthVerifier:   authVerifier,
+		AuthVerifier:  authVerifier,
-		EncryptionKey:  svr.auth.EncryptionKey(),
+		EncryptionKey: svr.auth.EncryptionKey(),
-		Conn:           ctlConn,
+		Conn:          ctlConn,
-		LoginMsg:       loginMsg,
+		LoginMsg:      loginMsg,
-		ServerCfg:      svr.cfg,
+		ServerCfg:     svr.cfg,
-		ClientRegistry: svr.clientRegistry,
+		WireProtocol:  wireProtocol,
-		WireProtocol:   wireProtocol,
 	})
 	if err != nil {
 		xl.Warnf("create new controller error: %v", err)
@@ -785,18 +799,17 @@ func (svr *Service) RegisterControl(
 		return nil, fmt.Errorf("unexpected error when creating new controller")
 	}
 
-	if oldCtl := svr.ctlManager.Add(loginMsg.RunID, ctl); oldCtl != nil {
+	if err := svr.ctlManager.Add(ctl); err != nil {
-		oldCtl.WaitClosed()
+		return ctl, err
 	}
+	ctl.WaitForHandoff()
 
-	remoteAddr := ctlConn.RemoteAddr().String()
+	active, err := svr.ctlManager.Activate(ctl)
-	if host, _, err := net.SplitHostPort(remoteAddr); err == nil {
+	if err != nil {
-		remoteAddr = host
+		return ctl, err
 	}
-	_, conflict := svr.clientRegistry.Register(loginMsg.User, loginMsg.ClientID, loginMsg.RunID, loginMsg.Hostname, loginMsg.Version, remoteAddr, wireProtocol)
+	if !active {
-	if conflict {
+		return ctl, errControlReplaced
-		svr.ctlManager.Del(loginMsg.RunID, ctl)
-		return nil, fmt.Errorf("client_id [%s] for user [%s] is already online", loginMsg.ClientID, loginMsg.User)
 	}
 
 	return ctl, nil
@@ -830,20 +843,25 @@ func (svr *Service) RegisterWorkConn(workConn *msg.Conn, newMsg *msg.NewWorkConn
 		xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
 		return err
 	}
-	return ctl.RegisterWorkConn(proxy.NewWorkConn(workConn))
+	return svr.ctlManager.RegisterWorkConn(ctl, proxy.NewWorkConn(workConn))
 }
 
 func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn, wireProtocol string) error {
-	visitorUser := ""
+	admit := func(visitorUser string) error {
+		return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
+			newMsg.UseEncryption, newMsg.UseCompression, visitorUser, wireProtocol)
+	}
 	// TODO(deprecation): Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID.
 	// If runID is required, it is not compatible with versions prior to v0.50.0.
 	if newMsg.RunID != "" {
-		ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
+		admitted, err := svr.ctlManager.admitVisitorByRunID(newMsg.RunID, admit)
-		if !exist {
+		if err != nil {
+			return err
+		}
+		if !admitted {
 			return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
 		}
-		visitorUser = ctl.sessionCtx.LoginMsg.User
+		return nil
 	}
-	return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
+	return admit("")
-		newMsg.UseEncryption, newMsg.UseCompression, visitorUser, wireProtocol)
 }

server/service_test.go

@@ -15,12 +15,27 @@
 package server
 
 import (
+	"context"
 	"errors"
 	"net"
+	"runtime"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/fatedier/frp/pkg/auth"
+	v1 "github.com/fatedier/frp/pkg/config/v1"
+	"github.com/fatedier/frp/pkg/msg"
+	plugin "github.com/fatedier/frp/pkg/plugin/server"
+	"github.com/fatedier/frp/pkg/proto/wire"
+	"github.com/fatedier/frp/pkg/util/util"
+	"github.com/fatedier/frp/server/controller"
+	"github.com/fatedier/frp/server/proxy"
+	"github.com/fatedier/frp/server/registry"
+	"github.com/fatedier/frp/server/visitor"
 )
 
 func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) {
@@ -61,3 +76,504 @@ func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) {
 		t.Fatal("timed out waiting for write after deadline reset")
 	}
 }
+
+func TestServiceControlHandoffSkipsStalePendingGeneration(t *testing.T) {
+	svr := newControlTestService(t)
+	metrics := newCountingServerMetrics()
+	metrics.closeEnter = make(chan struct{})
+	metrics.closeResume = make(chan struct{})
+
+	ctlA, connA, err := registerLifecycleTestControl(svr)
+	require.NoError(t, err)
+	ctlA.serverMetrics = metrics
+	require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil }))
+	waitForSignal(t, connA.readStarted, "A reader to start")
+
+	require.NoError(t, ctlA.Close())
+	waitForSignal(t, metrics.closeEnter, "A finalization barrier")
+
+	type registerResult struct {
+		ctl  *Control
+		conn *deadlineReadConn
+		err  error
+	}
+	resultB := make(chan registerResult, 1)
+	go func() {
+		ctl, conn, registerErr := registerLifecycleTestControl(svr)
+		resultB <- registerResult{ctl: ctl, conn: conn, err: registerErr}
+	}()
+	ctlB := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlA)
+	ctlB.serverMetrics = metrics
+
+	resultC := make(chan registerResult, 1)
+	go func() {
+		ctl, conn, registerErr := registerLifecycleTestControl(svr)
+		resultC <- registerResult{ctl: ctl, conn: conn, err: registerErr}
+	}()
+	ctlC := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlB)
+	ctlC.serverMetrics = metrics
+	waitForControlDone(t, ctlB)
+
+	select {
+	case result := <-resultB:
+		t.Fatalf("B returned before A finalized: %v", result.err)
+	default:
+	}
+	select {
+	case result := <-resultC:
+		t.Fatalf("C returned before A finalized: %v", result.err)
+	default:
+	}
+
+	close(metrics.closeResume)
+	waitForControlDone(t, ctlA)
+
+	b := <-resultB
+	require.Same(t, ctlB, b.ctl)
+	require.ErrorIs(t, b.err, errControlReplaced)
+	require.False(t, svr.ctlManager.Remove(ctlB))
+	require.NoError(t, ctlB.Close())
+
+	c := <-resultC
+	require.NoError(t, c.err)
+	require.Same(t, ctlC, c.ctl)
+	_, ok := svr.ctlManager.GetByID("shared-run")
+	require.False(t, ok)
+	require.Same(t, ctlC, currentControlForTest(svr.ctlManager, "shared-run"))
+
+	info, ok := svr.clientRegistry.GetByKey("client")
+	require.True(t, ok)
+	require.True(t, info.Online)
+	require.Equal(t, uint64(ctlC.ID()), info.ControlID)
+
+	var staleWrites atomic.Int64
+	err = svr.completeControlLogin(ctlB, func() error {
+		staleWrites.Add(1)
+		return nil
+	})
+	require.ErrorIs(t, err, errControlReplaced)
+	require.Equal(t, int64(0), staleWrites.Load())
+
+	require.NoError(t, svr.completeControlLogin(ctlC, func() error { return nil }))
+	waitForSignal(t, c.conn.readStarted, "C reader to start")
+	current, ok := svr.ctlManager.GetByID("shared-run")
+	require.True(t, ok)
+	require.Same(t, ctlC, current)
+	require.Equal(t, int64(2), metrics.newClients())
+	require.Equal(t, int64(1), metrics.closedClients())
+
+	require.NoError(t, ctlC.Close())
+	waitForControlDone(t, ctlC)
+	require.Equal(t, int64(2), metrics.newClients())
+	require.Equal(t, int64(2), metrics.closedClients())
+	_, ok = svr.ctlManager.GetByID("shared-run")
+	require.False(t, ok)
+}
+
+func TestServiceLoginResponseSynchronizationIsScopedToRun(t *testing.T) {
+	svr := newControlTestService(t)
+	metrics := newCountingServerMetrics()
+	ctlA, connA, err := registerLifecycleTestControl(svr)
+	require.NoError(t, err)
+	ctlA.serverMetrics = metrics
+
+	writeEntered := make(chan struct{})
+	resumeWrite := make(chan struct{})
+	var resumeWriteOnce sync.Once
+	resume := func() {
+		resumeWriteOnce.Do(func() { close(resumeWrite) })
+	}
+	t.Cleanup(resume)
+	writeCount := atomic.Int64{}
+	loginDone := make(chan error, 1)
+	go func() {
+		loginDone <- svr.completeControlLogin(ctlA, func() error {
+			close(writeEntered)
+			<-resumeWrite
+			writeCount.Add(1)
+			return nil
+		})
+	}()
+	waitForSignal(t, writeEntered, "A LoginResp write")
+
+	runMu := currentRunGateForTest(svr.ctlManager, "shared-run")
+	require.NotNil(t, runMu)
+	if !svr.ctlManager.mu.TryLock() {
+		t.Fatal("ControlManager mutex was held while LoginResp write was in progress")
+	}
+	svr.ctlManager.mu.Unlock()
+
+	ctlB, connB := newLifecycleTestControl(t, "shared-run", "client", metrics)
+	gateAvailable := make(chan bool)
+	addDone := make(chan error, 1)
+	go func() {
+		if runMu.TryLock() {
+			runMu.Unlock()
+			gateAvailable <- true
+		} else {
+			gateAvailable <- false
+		}
+		addErr := svr.ctlManager.Add(ctlB)
+		addDone <- addErr
+	}()
+	available := waitForResult(t, gateAvailable, "same-run replacement gate probe")
+	require.False(t, available, "same-run gate was available to replacement during LoginResp write")
+	select {
+	case addErr := <-addDone:
+		t.Fatalf("same-run replacement completed during LoginResp write: %v", addErr)
+	case <-time.After(20 * time.Millisecond):
+	}
+	require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run"))
+
+	otherMetrics := newCountingServerMetrics()
+	otherCtl, otherConn := newLifecycleTestControl(t, "other-run", "other-client", otherMetrics)
+	type unrelatedResult struct {
+		addErr      error
+		active      bool
+		activateErr error
+		loginErr    error
+		current     *Control
+		found       bool
+	}
+	unrelatedDone := make(chan unrelatedResult, 1)
+	go func() {
+		result := unrelatedResult{}
+		result.addErr = svr.ctlManager.Add(otherCtl)
+		if result.addErr == nil {
+			result.active, result.activateErr = svr.ctlManager.Activate(otherCtl)
+		}
+		if result.activateErr == nil && result.active {
+			result.loginErr = svr.completeControlLogin(otherCtl, func() error { return nil })
+		}
+		result.current, result.found = svr.ctlManager.GetByID("other-run")
+		unrelatedDone <- result
+	}()
+	result := waitForResult(t, unrelatedDone, "unrelated run lifecycle")
+	require.NoError(t, result.addErr)
+	require.NoError(t, result.activateErr)
+	require.True(t, result.active)
+	require.NoError(t, result.loginErr)
+	require.True(t, result.found)
+	require.Same(t, otherCtl, result.current)
+	waitForSignal(t, otherConn.readStarted, "unrelated control reader to start")
+	require.Equal(t, int64(1), otherMetrics.newClients())
+
+	resume()
+	require.NoError(t, waitForResult(t, loginDone, "LoginResp completion"))
+	require.NoError(t, waitForResult(t, addDone, "replacement"))
+	waitForControlDone(t, ctlA)
+	require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run"))
+	require.Equal(t, int64(1), writeCount.Load())
+	require.Equal(t, int64(1), metrics.newClients())
+	require.Equal(t, int64(1), metrics.closedClients())
+	require.Equal(t, []string{"deadline", "close"}, connA.eventsSnapshot())
+
+	require.False(t, svr.ctlManager.Remove(ctlA))
+	require.NoError(t, ctlA.Close())
+	require.True(t, svr.ctlManager.Remove(ctlB))
+	require.NoError(t, ctlB.Close())
+	require.Equal(t, []string{"deadline", "close"}, connB.eventsSnapshot())
+
+	require.NoError(t, otherCtl.Close())
+	waitForControlDone(t, otherCtl)
+	require.Equal(t, int64(1), otherMetrics.newClients())
+	require.Equal(t, int64(1), otherMetrics.closedClients())
+}
+
+func TestServiceVisitorAdmissionSerializesReplacement(t *testing.T) {
+	svr := newControlTestService(t)
+	ctlA, controlConn, err := registerLifecycleTestControl(svr)
+	require.NoError(t, err)
+	ctlA.sessionCtx.LoginMsg.User = "old-user"
+	require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil }))
+	waitForSignal(t, controlConn.readStarted, "A reader to start")
+
+	admissionEntered := make(chan struct{})
+	resumeAdmission := make(chan struct{})
+	var resumeOnce sync.Once
+	resume := func() {
+		resumeOnce.Do(func() { close(resumeAdmission) })
+	}
+	t.Cleanup(resume)
+	type admissionResult struct {
+		admitted bool
+		user     string
+		err      error
+	}
+	admissionDone := make(chan admissionResult, 1)
+	go func() {
+		var admittedUser string
+		admitted, admitErr := svr.ctlManager.admitVisitorByRunID("shared-run", func(user string) error {
+			admittedUser = user
+			close(admissionEntered)
+			<-resumeAdmission
+			return nil
+		})
+		admissionDone <- admissionResult{admitted: admitted, user: admittedUser, err: admitErr}
+	}()
+	waitForSignal(t, admissionEntered, "visitor admission callback")
+	runMu := currentRunGateForTest(svr.ctlManager, "shared-run")
+	require.NotNil(t, runMu)
+
+	type registerResult struct {
+		ctl *Control
+		err error
+	}
+	gateAvailable := make(chan bool)
+	replacementDone := make(chan registerResult, 1)
+	go func() {
+		if runMu.TryLock() {
+			runMu.Unlock()
+			gateAvailable <- true
+		} else {
+			gateAvailable <- false
+		}
+		ctl, _, registerErr := registerLifecycleTestControl(svr)
+		replacementDone <- registerResult{ctl: ctl, err: registerErr}
+	}()
+	available := waitForResult(t, gateAvailable, "visitor replacement gate probe")
+	require.False(t, available, "same-run gate was available during visitor admission")
+	select {
+	case result := <-replacementDone:
+		t.Fatalf("replacement completed during visitor admission: %v", result.err)
+	case <-time.After(20 * time.Millisecond):
+	}
+	require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run"))
+
+	resume()
+	admission := waitForResult(t, admissionDone, "visitor admission")
+	require.NoError(t, admission.err)
+	require.True(t, admission.admitted)
+	require.Equal(t, "old-user", admission.user)
+	replacement := waitForResult(t, replacementDone, "replacement")
+	require.NoError(t, replacement.err)
+	ctlB := replacement.ctl
+	require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run"))
+	waitForControlDone(t, ctlA)
+	require.True(t, svr.ctlManager.Remove(ctlB))
+	require.NoError(t, ctlB.Close())
+}
+
+func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {
+	svr := newControlTestService(t)
+	ctl, controlConn, err := registerLifecycleTestControl(svr)
+	require.NoError(t, err)
+
+	pendingConn := newCountingCloseConn()
+	pendingMsgConn := msg.NewConn(pendingConn, msg.NewV1ReadWriter(pendingConn))
+	err = registerWorkConnAsCaller(svr, pendingMsgConn, &msg.NewWorkConn{RunID: "shared-run"})
+	require.Error(t, err)
+	require.Equal(t, int64(1), pendingConn.closeCount.Load())
+	require.Len(t, ctl.workConnCh, 0)
+
+	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
+	waitForSignal(t, controlConn.readStarted, "control reader to start")
+	current, ok := svr.ctlManager.GetByID("shared-run")
+	require.True(t, ok)
+	require.Same(t, ctl, current)
+	require.Len(t, ctl.workConnCh, 0)
+
+	runningConn := newCountingCloseConn()
+	runningMsgConn := msg.NewConn(runningConn, msg.NewV1ReadWriter(runningConn))
+	require.NoError(t, svr.RegisterWorkConn(runningMsgConn, &msg.NewWorkConn{RunID: "shared-run"}))
+	require.Len(t, ctl.workConnCh, 1)
+
+	require.NoError(t, ctl.Close())
+	waitForControlDone(t, ctl)
+	require.Equal(t, int64(1), runningConn.closeCount.Load())
+}
+
+func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) {
+	for _, action := range []string{"replace", "close"} {
+		t.Run(action, func(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")
+
+			barrier := newWorkConnBarrierPlugin()
+			svr.pluginManager.Register(barrier)
+			workConn := newCountingCloseConn()
+			workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
+			routeDone := make(chan error, 1)
+			go func() {
+				routeDone <- registerWorkConnAsCaller(svr, workMsgConn, &msg.NewWorkConn{RunID: "shared-run"})
+			}()
+			waitForSignal(t, barrier.entered, "work connection plugin barrier")
+
+			var replacement *Control
+			switch action {
+			case "replace":
+				replacement, _, err = registerLifecycleTestControl(svr)
+				require.NoError(t, err)
+			case "close":
+				require.NoError(t, ctl.Close())
+				waitForControlDone(t, ctl)
+			}
+
+			close(barrier.resume)
+			require.Error(t, waitForResult(t, routeDone, "work connection route to finish"))
+			require.Equal(t, int64(1), workConn.closeCount.Load())
+			require.Len(t, ctl.workConnCh, 0)
+
+			if replacement != nil {
+				require.Len(t, replacement.workConnCh, 0)
+				require.True(t, svr.ctlManager.Remove(replacement))
+				require.NoError(t, replacement.Close())
+				waitForControlDone(t, replacement)
+			}
+		})
+	}
+}
+
+func TestServiceVisitorRoutingExcludesPendingUser(t *testing.T) {
+	svr := newControlTestService(t)
+	listener, err := svr.rc.VisitorManager.Listen("visitor", "secret", []string{"pending-user"})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = listener.Close() })
+
+	controlConn := newDeadlineReadConn()
+	controlMsgConn := msg.NewConn(controlConn, msg.NewV1ReadWriter(controlConn))
+	ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
+		RunID:    "visitor-run",
+		User:     "pending-user",
+		ClientID: "visitor-client",
+		ClientSpec: msg.ClientSpec{
+			AlwaysAuthPass: true,
+		},
+	}, true, wire.ProtocolV1)
+	require.NoError(t, err)
+
+	timestamp := time.Now().Unix()
+	visitorMsg := &msg.NewVisitorConn{
+		RunID:     "visitor-run",
+		ProxyName: "visitor",
+		Timestamp: timestamp,
+		SignKey:   util.GetAuthKey("secret", timestamp),
+	}
+	pendingConn := newCountingCloseConn()
+	err = svr.RegisterVisitorConn(pendingConn, visitorMsg, wire.ProtocolV1)
+	require.ErrorContains(t, err, "no client control found")
+	require.NoError(t, pendingConn.Close())
+	require.Equal(t, int64(1), pendingConn.closeCount.Load())
+
+	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
+	waitForSignal(t, controlConn.readStarted, "control reader to start")
+	runningConn := newCountingCloseConn()
+	require.NoError(t, svr.RegisterVisitorConn(runningConn, visitorMsg, wire.ProtocolV1))
+	accepted, err := listener.Accept()
+	require.NoError(t, err)
+	require.NoError(t, accepted.Close())
+	require.Equal(t, int64(1), runningConn.closeCount.Load())
+
+	require.NoError(t, ctl.Close())
+	waitForControlDone(t, ctl)
+}
+
+func newControlTestService(t *testing.T) *Service {
+	t.Helper()
+	cfg := &v1.ServerConfig{}
+	cfg.Auth.Method = v1.AuthMethodToken
+	authRuntime, err := auth.BuildServerAuth(&cfg.Auth)
+	require.NoError(t, err)
+	clientRegistry := registry.NewClientRegistry()
+	return &Service{
+		ctlManager:     NewControlManager(clientRegistry),
+		clientRegistry: clientRegistry,
+		pxyManager:     proxy.NewManager(),
+		pluginManager:  plugin.NewManager(),
+		rc: &controller.ResourceController{
+			VisitorManager: visitor.NewManager(),
+		},
+		auth: authRuntime,
+		cfg:  cfg,
+	}
+}
+
+func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, error) {
+	conn := newDeadlineReadConn()
+	msgConn := msg.NewConn(conn, msg.NewReadWriter(conn, wire.ProtocolV1))
+	ctl, err := svr.RegisterControl(msgConn, &msg.Login{
+		RunID:    "shared-run",
+		ClientID: "client",
+		ClientSpec: msg.ClientSpec{
+			AlwaysAuthPass: true,
+		},
+	}, true, wire.ProtocolV1)
+	return ctl, conn, err
+}
+
+func waitForDifferentCurrentControl(t *testing.T, manager *ControlManager, runID string, old *Control) *Control {
+	t.Helper()
+	deadline := time.Now().Add(3 * time.Second)
+	for time.Now().Before(deadline) {
+		if ctl := currentControlForTest(manager, runID); ctl != nil && ctl != old {
+			return ctl
+		}
+		runtime.Gosched()
+	}
+	t.Fatalf("timed out waiting for a new current control after ID %d", old.ID())
+	return nil
+}
+
+func registerWorkConnAsCaller(svr *Service, workConn *msg.Conn, newMsg *msg.NewWorkConn) error {
+	err := svr.RegisterWorkConn(workConn, newMsg)
+	if err != nil {
+		_ = workConn.Close()
+	}
+	return err
+}
+
+func waitForResult[T any](t *testing.T, ch <-chan T, description string) T {
+	t.Helper()
+	select {
+	case result := <-ch:
+		return result
+	case <-time.After(3 * time.Second):
+		t.Fatalf("timed out waiting for %s", description)
+		var zero T
+		return zero
+	}
+}
+
+type workConnBarrierPlugin struct {
+	entered chan struct{}
+	resume  chan struct{}
+}
+
+func newWorkConnBarrierPlugin() *workConnBarrierPlugin {
+	return &workConnBarrierPlugin{
+		entered: make(chan struct{}),
+		resume:  make(chan struct{}),
+	}
+}
+
+func (*workConnBarrierPlugin) Name() string { return "work-conn-barrier" }
+
+func (*workConnBarrierPlugin) IsSupport(op string) bool { return op == plugin.OpNewWorkConn }
+
+func (p *workConnBarrierPlugin) Handle(
+	context.Context,
+	string,
+	any,
+) (*plugin.Response, any, error) {
+	close(p.entered)
+	<-p.resume
+	return &plugin.Response{Unchange: true}, nil, nil
+}
+
+type countingCloseConn struct {
+	closeCount atomic.Int64
+}
+
+func newCountingCloseConn() *countingCloseConn { return &countingCloseConn{} }
+
+func (*countingCloseConn) Read([]byte) (int, error)         { return 0, net.ErrClosed }
+func (*countingCloseConn) Write(p []byte) (int, error)      { return len(p), nil }
+func (c *countingCloseConn) Close() error                   { c.closeCount.Add(1); return nil }
+func (*countingCloseConn) LocalAddr() net.Addr              { return lifecycleTestAddr("local") }
+func (*countingCloseConn) RemoteAddr() net.Addr             { return lifecycleTestAddr("remote") }
+func (*countingCloseConn) SetDeadline(time.Time) error      { return nil }
+func (*countingCloseConn) SetReadDeadline(time.Time) error  { return nil }
+func (*countingCloseConn) SetWriteDeadline(time.Time) error { return nil }

test/e2e/framework/process.go

@@ -94,11 +94,9 @@ func (f *Framework) RunProcessesWithBinaries(
 }
 
 func (f *Framework) RunFrps(args ...string) (*process.Process, string, error) {
-	p := process.NewWithEnvs(TestContext.FRPServerPath, args, f.osEnvs)
+	p, output, err := f.StartFrps(args...)
-	f.serverProcesses = append(f.serverProcesses, p)
-	err := p.Start()
 	if err != nil {
-		return p, p.Output(), err
+		return p, output, err
 	}
 	select {
 	case <-p.Done():
@@ -107,17 +105,39 @@ func (f *Framework) RunFrps(args ...string) (*process.Process, string, error) {
 	return p, p.Output(), nil
 }
 
+// StartFrps starts frps without an implicit sleep so tests can wait on an
+// explicit readiness event.
+func (f *Framework) StartFrps(args ...string) (*process.Process, string, error) {
+	p := process.NewWithEnvs(TestContext.FRPServerPath, args, f.osEnvs)
+	f.serverProcesses = append(f.serverProcesses, p)
+	err := p.Start()
+	if err != nil {
+		return p, p.Output(), err
+	}
+	return p, p.Output(), nil
+}
+
 func (f *Framework) RunFrpc(args ...string) (*process.Process, string, error) {
+	p, output, err := f.StartFrpc(args...)
+	if err != nil {
+		return p, output, err
+	}
+	select {
+	case <-p.Done():
+	case <-time.After(1500 * time.Millisecond):
+	}
+	return p, p.Output(), nil
+}
+
+// StartFrpc starts frpc without an implicit sleep so tests can wait on an
+// explicit login or proxy-readiness event.
+func (f *Framework) StartFrpc(args ...string) (*process.Process, string, error) {
 	p := process.NewWithEnvs(TestContext.FRPClientPath, args, f.osEnvs)
 	f.clientProcesses = append(f.clientProcesses, p)
 	err := p.Start()
 	if err != nil {
 		return p, p.Output(), err
 	}
-	select {
-	case <-p.Done():
-	case <-time.After(1500 * time.Millisecond):
-	}
 	return p, p.Output(), nil
 }
 

test/e2e/pkg/relay/halfopen.go

@@ -0,0 +1,208 @@
+// 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 relay
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"strconv"
+	"sync"
+	"time"
+)
+
+// HalfOpen forwards TCP connections until Blackhole is called. A blackholed
+// pair stops forwarding but deliberately retains the upstream socket so the
+// peer sees a real half-open connection until the relay is closed.
+type HalfOpen struct {
+	bindAddr     string
+	bindPort     int
+	upstreamAddr string
+
+	listener net.Listener
+	done     chan struct{}
+	accepted chan struct{}
+
+	mu    sync.Mutex
+	pairs []*connectionPair
+
+	wg        sync.WaitGroup
+	closeOnce sync.Once
+}
+
+type connectionPair struct {
+	downstream net.Conn
+	upstream   net.Conn
+
+	mu         sync.Mutex
+	blackholed bool
+}
+
+func New(upstreamAddr string) *HalfOpen {
+	return &HalfOpen{
+		bindAddr:     "127.0.0.1",
+		upstreamAddr: upstreamAddr,
+		done:         make(chan struct{}),
+		accepted:     make(chan struct{}, 1),
+	}
+}
+
+func (r *HalfOpen) Run() error {
+	listener, err := net.Listen("tcp", net.JoinHostPort(r.bindAddr, strconv.Itoa(r.bindPort)))
+	if err != nil {
+		return err
+	}
+	r.listener = listener
+	r.bindPort = listener.Addr().(*net.TCPAddr).Port
+
+	r.wg.Add(1)
+	go r.acceptLoop()
+	return nil
+}
+
+func (r *HalfOpen) acceptLoop() {
+	defer r.wg.Done()
+	for {
+		downstream, err := r.listener.Accept()
+		if err != nil {
+			return
+		}
+		upstream, err := net.DialTimeout("tcp", r.upstreamAddr, 3*time.Second)
+		if err != nil {
+			_ = downstream.Close()
+			continue
+		}
+
+		pair := &connectionPair{downstream: downstream, upstream: upstream}
+		r.mu.Lock()
+		r.pairs = append(r.pairs, pair)
+		r.mu.Unlock()
+		select {
+		case r.accepted <- struct{}{}:
+		default:
+		}
+
+		r.wg.Add(1)
+		go r.servePair(pair)
+	}
+}
+
+func (r *HalfOpen) servePair(pair *connectionPair) {
+	defer r.wg.Done()
+	copyDone := make(chan struct{}, 2)
+	go func() {
+		_, _ = io.Copy(pair.upstream, pair.downstream)
+		copyDone <- struct{}{}
+	}()
+	go func() {
+		_, _ = io.Copy(pair.downstream, pair.upstream)
+		copyDone <- struct{}{}
+	}()
+
+	completed := 0
+	select {
+	case <-copyDone:
+		completed = 1
+		if !pair.isBlackholed() {
+			_ = pair.downstream.Close()
+			_ = pair.upstream.Close()
+		}
+	case <-r.done:
+		_ = pair.downstream.Close()
+		_ = pair.upstream.Close()
+	}
+	for completed < 2 {
+		<-copyDone
+		completed++
+	}
+
+	if pair.isBlackholed() {
+		<-r.done
+		_ = pair.downstream.Close()
+		_ = pair.upstream.Close()
+	}
+}
+
+func (r *HalfOpen) WaitForConnections(count int, timeout time.Duration) error {
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+	for {
+		r.mu.Lock()
+		accepted := len(r.pairs)
+		r.mu.Unlock()
+		if accepted >= count {
+			return nil
+		}
+		select {
+		case <-r.accepted:
+		case <-r.done:
+			return fmt.Errorf("relay closed after accepting %d of %d connections", accepted, count)
+		case <-timer.C:
+			return fmt.Errorf("timed out after accepting %d of %d connections", accepted, count)
+		}
+	}
+}
+
+// Blackhole uses a one-based connection index in accept order.
+func (r *HalfOpen) Blackhole(index int) error {
+	r.mu.Lock()
+	if index <= 0 || index > len(r.pairs) {
+		accepted := len(r.pairs)
+		r.mu.Unlock()
+		return fmt.Errorf("connection %d is unavailable; accepted %d", index, accepted)
+	}
+	pair := r.pairs[index-1]
+	r.mu.Unlock()
+
+	pair.mu.Lock()
+	if pair.blackholed {
+		pair.mu.Unlock()
+		return nil
+	}
+	pair.blackholed = true
+	pair.mu.Unlock()
+
+	now := time.Now()
+	_ = pair.downstream.SetDeadline(now)
+	_ = pair.upstream.SetDeadline(now)
+	return nil
+}
+
+func (r *HalfOpen) Close() error {
+	r.closeOnce.Do(func() {
+		close(r.done)
+		if r.listener != nil {
+			_ = r.listener.Close()
+		}
+		r.mu.Lock()
+		pairs := append([]*connectionPair(nil), r.pairs...)
+		r.mu.Unlock()
+		for _, pair := range pairs {
+			_ = pair.downstream.Close()
+			_ = pair.upstream.Close()
+		}
+		r.wg.Wait()
+	})
+	return nil
+}
+
+func (r *HalfOpen) BindAddr() string { return r.bindAddr }
+func (r *HalfOpen) BindPort() int    { return r.bindPort }
+
+func (p *connectionPair) isBlackholed() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.blackholed
+}

test/e2e/v1/features/control_replacement.go

@@ -0,0 +1,300 @@
+// 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 features
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/onsi/ginkgo/v2"
+
+	"github.com/fatedier/frp/test/e2e/framework"
+	"github.com/fatedier/frp/test/e2e/framework/consts"
+	"github.com/fatedier/frp/test/e2e/pkg/relay"
+	"github.com/fatedier/frp/test/e2e/pkg/request"
+)
+
+var _ = ginkgo.Describe("[Feature: ControlReplacement]", func() {
+	f := framework.NewDefaultFramework()
+
+	for _, wireProtocol := range []string{"v1", "v2"} {
+		for _, tcpMux := range []bool{true, false} {
+			ginkgo.It(fmt.Sprintf("recovers a %s control through a half-open relay with tcpMux=%t", wireProtocol, tcpMux), func() {
+				runHalfOpenControlReplacement(f, wireProtocol, tcpMux)
+			})
+		}
+	}
+})
+
+func runHalfOpenControlReplacement(f *framework.Framework, wireProtocol string, tcpMux bool) {
+	serverPort := f.AllocPort()
+	dashboardPort := f.AllocPort()
+	remotePort := f.AllocPort()
+	heartbeatTimeout := int64(-1)
+	if !tcpMux {
+		heartbeatTimeout = 3
+	}
+
+	serverConfig := fmt.Sprintf(`
+bindAddr = "127.0.0.1"
+bindPort = %d
+log.level = "trace"
+transport.tcpMux = %t
+transport.tcpMuxKeepaliveInterval = 30
+transport.heartbeatTimeout = %d
+webServer.addr = "127.0.0.1"
+webServer.port = %d
+webServer.pprofEnable = true
+enablePrometheus = true
+`, serverPort, tcpMux, heartbeatTimeout, dashboardPort)
+	serverConfigPath := f.WriteTempFile("issue-5391-frps.toml", serverConfig)
+	serverProcess, _, err := f.StartFrps("-c", serverConfigPath)
+	framework.ExpectNoError(err)
+	framework.ExpectNoError(framework.WaitForTCPReady(fmt.Sprintf("127.0.0.1:%d", serverPort), 5*time.Second))
+
+	halfOpenRelay := relay.New(fmt.Sprintf("127.0.0.1:%d", serverPort))
+	f.RunServer("", halfOpenRelay)
+
+	heartbeatInterval := int64(-1)
+	clientHeartbeatTimeout := int64(-1)
+	if !tcpMux {
+		heartbeatInterval = 1
+		clientHeartbeatTimeout = 3
+	}
+	clientConfig := fmt.Sprintf(`
+serverAddr = "127.0.0.1"
+serverPort = %d
+clientID = "issue-5391"
+loginFailExit = false
+log.level = "trace"
+transport.wireProtocol = %q
+transport.tcpMux = %t
+transport.tcpMuxKeepaliveInterval = 1
+transport.heartbeatInterval = %d
+transport.heartbeatTimeout = %d
+transport.tls.enable = false
+
+[[proxies]]
+name = "issue-5391-tcp"
+type = "tcp"
+localPort = %d
+remotePort = %d
+`, halfOpenRelay.BindPort(), wireProtocol, tcpMux, heartbeatInterval, clientHeartbeatTimeout,
+		f.PortByName(framework.TCPEchoServerPort), remotePort)
+	clientConfigPath := f.WriteTempFile("issue-5391-frpc.toml", clientConfig)
+	clientProcess, _, err := f.StartFrpc("-c", clientConfigPath)
+	framework.ExpectNoError(err)
+	framework.ExpectNoError(halfOpenRelay.WaitForConnections(1, 5*time.Second))
+	framework.ExpectNoError(clientProcess.WaitForOutput("login to server success", 1, 10*time.Second))
+	framework.ExpectNoError(clientProcess.WaitForOutput("[issue-5391-tcp] start proxy success", 1, 10*time.Second))
+	framework.ExpectNoError(waitForReplacementState(dashboardPort, remotePort, 10*time.Second))
+
+	replacementCount := 1
+	if tcpMux {
+		replacementCount = 3
+	}
+	for i := 0; i < replacementCount; i++ {
+		connectionIndex := 1
+		if tcpMux {
+			connectionIndex = i + 1
+		}
+		framework.ExpectNoError(halfOpenRelay.Blackhole(connectionIndex))
+		if tcpMux {
+			framework.ExpectNoError(halfOpenRelay.WaitForConnections(connectionIndex+1, 15*time.Second))
+		}
+		framework.ExpectNoError(clientProcess.WaitForOutput("login to server success", i+2, 15*time.Second))
+		framework.ExpectNoError(clientProcess.WaitForOutput("[issue-5391-tcp] start proxy success", i+2, 15*time.Second))
+		framework.ExpectNoError(waitForReplacementState(dashboardPort, remotePort, 10*time.Second))
+	}
+
+	_ = clientProcess.Stop()
+	select {
+	case <-clientProcess.Done():
+	case <-time.After(5 * time.Second):
+		framework.Failf("frpc did not exit")
+	}
+	framework.ExpectNoError(waitForReplacementShutdown(dashboardPort, 10*time.Second))
+	framework.ExpectNoError(halfOpenRelay.Close())
+	framework.ExpectNoError(waitForNoHandoffWaiters(dashboardPort, 5*time.Second))
+
+	_ = serverProcess.Stop()
+	select {
+	case <-serverProcess.Done():
+	case <-time.After(5 * time.Second):
+		framework.Failf("frps did not exit")
+	}
+}
+
+func waitForReplacementState(dashboardPort, remotePort int, timeout time.Duration) error {
+	return waitForLifecycleCondition(timeout, func() error {
+		metricsBody, err := getLifecycleEndpoint(dashboardPort, "/metrics")
+		if err != nil {
+			return err
+		}
+		if err := expectMetricValue(metricsBody, "frp_server_client_counts", "", 1); err != nil {
+			return err
+		}
+		if err := expectMetricValue(metricsBody, "frp_server_proxy_counts", `type="tcp"`, 1); err != nil {
+			return err
+		}
+		clients, err := getOnlineLifecycleClients(dashboardPort)
+		if err != nil {
+			return err
+		}
+		if len(clients) != 1 || clients[0].ClientID != "issue-5391" {
+			return fmt.Errorf("expected one online client, got %+v", clients)
+		}
+		profile, err := getLifecycleEndpoint(dashboardPort, "/debug/pprof/goroutine?debug=2")
+		if err != nil {
+			return err
+		}
+		if err := expectNoHandoffWaiter(profile, "after replacement"); err != nil {
+			return err
+		}
+		resp, err := request.New().
+			TCP().
+			Port(remotePort).
+			Timeout(time.Second).
+			Body([]byte(consts.TestString)).
+			Do()
+		if err != nil {
+			return err
+		}
+		if string(resp.Content) != consts.TestString {
+			return fmt.Errorf("unexpected proxy response %q", resp.Content)
+		}
+		return nil
+	})
+}
+
+func waitForReplacementShutdown(dashboardPort int, timeout time.Duration) error {
+	return waitForLifecycleCondition(timeout, func() error {
+		metricsBody, err := getLifecycleEndpoint(dashboardPort, "/metrics")
+		if err != nil {
+			return err
+		}
+		if err := expectMetricValue(metricsBody, "frp_server_client_counts", "", 0); err != nil {
+			return err
+		}
+		if err := expectMetricValue(metricsBody, "frp_server_proxy_counts", `type="tcp"`, 0); err != nil {
+			return err
+		}
+		clients, err := getOnlineLifecycleClients(dashboardPort)
+		if err != nil {
+			return err
+		}
+		if len(clients) != 0 {
+			return fmt.Errorf("expected no online clients, got %+v", clients)
+		}
+		return nil
+	})
+}
+
+func waitForNoHandoffWaiters(dashboardPort int, timeout time.Duration) error {
+	return waitForLifecycleCondition(timeout, func() error {
+		profile, err := getLifecycleEndpoint(dashboardPort, "/debug/pprof/goroutine?debug=2")
+		if err != nil {
+			return err
+		}
+		return expectNoHandoffWaiter(profile, "after relay shutdown")
+	})
+}
+
+type lifecycleClient struct {
+	ClientID string `json:"clientID"`
+}
+
+func expectNoHandoffWaiter(profile, phase string) error {
+	if strings.Contains(profile, "(*Control).WaitForHandoff") {
+		return fmt.Errorf("control handoff waiter remained %s", phase)
+	}
+	return nil
+}
+
+func getOnlineLifecycleClients(dashboardPort int) ([]lifecycleClient, error) {
+	body, err := getLifecycleEndpoint(dashboardPort, "/api/clients?status=online")
+	if err != nil {
+		return nil, err
+	}
+	var clients []lifecycleClient
+	if err := json.Unmarshal([]byte(body), &clients); err != nil {
+		return nil, err
+	}
+	return clients, nil
+}
+
+func getLifecycleEndpoint(port int, path string) (string, error) {
+	client := &http.Client{Timeout: time.Second}
+	resp, err := client.Get(fmt.Sprintf("http://127.0.0.1:%d%s", port, path))
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("GET %s returned %s", path, resp.Status)
+	}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return string(body), nil
+}
+
+func expectMetricValue(body, name, labels string, want float64) error {
+	prefix := name
+	if labels != "" {
+		prefix += "{" + labels + "}"
+	}
+	for line := range strings.SplitSeq(body, "\n") {
+		fields := strings.Fields(line)
+		if len(fields) != 2 || fields[0] != prefix {
+			continue
+		}
+		got, err := strconv.ParseFloat(fields[1], 64)
+		if err != nil {
+			return err
+		}
+		if got != want {
+			return fmt.Errorf("metric %s = %v, want %v", prefix, got, want)
+		}
+		return nil
+	}
+	return fmt.Errorf("metric %s not found", prefix)
+}
+
+func waitForLifecycleCondition(timeout time.Duration, condition func() error) error {
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+	ticker := time.NewTicker(25 * time.Millisecond)
+	defer ticker.Stop()
+	var lastErr error
+	for {
+		err := condition()
+		if err == nil {
+			return nil
+		}
+		lastErr = err
+		select {
+		case <-ticker.C:
+		case <-timer.C:
+			return fmt.Errorf("condition was not met: %w", lastErr)
+		}
+	}
+}