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

This commit is contained in:
fatedier
2026-07-21 18:09:09 +08:00
committed by GitHub
parent fe79598ee4
commit d486018885
10 changed files with 2241 additions and 152 deletions

View File

@@ -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 = ""
info.Online = false
now := cr.clock.Now()
info.DisconnectedAt = now
setClientOffline(info, cr.clock.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.