From de6bdf2290c3b533d9e28ccd29f6df6de883eedd Mon Sep 17 00:00:00 2001 From: viettranx Date: Fri, 10 Apr 2026 20:32:56 +0700 Subject: [PATCH] refactor(mcp): DRY reconnect logic, fix client pointer thread safety - Extract reconnectWithBackoff() shared by Manager and Pool reconnect paths - Use clientPtr.Load() in Stop/Evict/evict paths to avoid race with fullReconnect - Add tool staleness note to fullReconnect - Remove dead poolEntry.Client() --- internal/mcp/manager.go | 5 +- internal/mcp/manager_connect.go | 34 +++++++------- internal/mcp/pool.go | 81 +++++++-------------------------- 3 files changed, 38 insertions(+), 82 deletions(-) diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index df63d1d4..903e693a 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -516,8 +516,9 @@ func (m *Manager) Stop() { if ss.cancel != nil { ss.cancel() } - if ss.client != nil { - if err := ss.client.Close(); err != nil { + // Use atomic pointer — health loop may swap client via fullReconnect concurrently. + if client := ss.clientPtr.Load(); client != nil { + if err := client.Close(); err != nil { slog.Debug("mcp.server.close_error", "server", name, "error", err) } } diff --git a/internal/mcp/manager_connect.go b/internal/mcp/manager_connect.go index a71c3ecb..03a60257 100644 --- a/internal/mcp/manager_connect.go +++ b/internal/mcp/manager_connect.go @@ -309,17 +309,21 @@ func (m *Manager) healthLoop(ctx context.Context, ss *serverState) { } // tryReconnect attempts to reconnect with exponential backoff. -// For SSE/HTTP transports, a server-side restart invalidates the session, -// so pinging the old client will never succeed. In that case, we create a -// fresh client and atomically swap it via serverState.clientPtr — all -// BridgeTools see the new client on their next Execute call. func (m *Manager) tryReconnect(ctx context.Context, ss *serverState) { + reconnectWithBackoff(ctx, ss, "mcp.server") +} + +// reconnectWithBackoff implements the two-phase reconnect strategy shared by +// Manager.healthLoop and poolHealthLoop. Handles cooldown after exhausting +// max attempts, exponential backoff, fast-path ping (transient blips), and +// slow-path full reconnect (dead server-side session). +// logPrefix distinguishes log entries (e.g. "mcp.server" vs "mcp.pool"). +func reconnectWithBackoff(ctx context.Context, ss *serverState, logPrefix string) { ss.mu.Lock() if ss.reconnAttempts >= maxReconnectAttempts { ss.lastErr = fmt.Sprintf("max reconnect attempts (%d) reached, entering cooldown", maxReconnectAttempts) ss.mu.Unlock() - slog.Warn("mcp.server.reconnect_cooldown", "server", ss.name, "cooldown", reconnectCooldown) - // Cooldown: wait before resetting attempts so the server has time to recover. + slog.Warn(logPrefix+".reconnect_cooldown", "server", ss.name, "cooldown", reconnectCooldown) select { case <-ctx.Done(): return @@ -335,12 +339,7 @@ func (m *Manager) tryReconnect(ctx context.Context, ss *serverState) { ss.mu.Unlock() backoff := min(initialBackoff*time.Duration(1<<(attempt-1)), maxBackoff) - - slog.Info("mcp.server.reconnecting", - "server", ss.name, - "attempt", attempt, - "backoff", backoff, - ) + slog.Info(logPrefix+".reconnecting", "server", ss.name, "attempt", attempt, "backoff", backoff) select { case <-ctx.Done(): @@ -357,23 +356,26 @@ func (m *Manager) tryReconnect(ctx context.Context, ss *serverState) { ss.healthFailures = 0 ss.lastErr = "" ss.mu.Unlock() - slog.Info("mcp.server.reconnected", "server", ss.name) + slog.Info(logPrefix+".reconnected", "server", ss.name) return } // Slow path: server-side session is dead (container restart, OOM, etc.). if fullReconnect(ctx, ss) { - slog.Info("mcp.server.reconnected", "server", ss.name, "method", "full_reconnect") + slog.Info(logPrefix+".reconnected", "server", ss.name, "method", "full_reconnect") } } // fullReconnect creates a fresh MCP client, atomically swaps it into serverState, -// and closes the old one. Returns true on success. Shared by Manager.tryReconnect -// and poolTryReconnect to avoid duplication. +// and closes the old one. Returns true on success. Used by reconnectWithBackoff +// as the slow path when pinging the old client fails. // // The new client is created and validated FIRST, then swapped via clientPtr.Store() // so BridgeTools see the new client immediately. The old client is closed AFTER // the swap to avoid a window where ss.client points to a closed client. +// +// NOTE: Does not re-discover tools (ListTools). If the MCP server restarts with +// a different tool set, changes won't be reflected until the Manager reconnects. func fullReconnect(ctx context.Context, ss *serverState) bool { slog.Info("mcp.full_reconnect", "server", ss.name, "transport", ss.transport) diff --git a/internal/mcp/pool.go b/internal/mcp/pool.go index df86675b..b767f35f 100644 --- a/internal/mcp/pool.go +++ b/internal/mcp/pool.go @@ -131,8 +131,8 @@ func (p *Pool) Acquire(ctx context.Context, tenantID uuid.UUID, name, transportT if old.state.cancel != nil { old.state.cancel() } - if old.state.client != nil { - _ = old.state.client.Close() + if client := old.state.clientPtr.Load(); client != nil { + _ = client.Close() } delete(p.servers, key) // Return slot to semaphore @@ -216,8 +216,8 @@ func (p *Pool) AcquireUser(ctx context.Context, tenantID uuid.UUID, name, userID if old.state.cancel != nil { old.state.cancel() } - if old.state.client != nil { - _ = old.state.client.Close() + if client := old.state.clientPtr.Load(); client != nil { + _ = client.Close() } delete(p.userServers, key) // Return slot to per-server semaphore @@ -390,8 +390,8 @@ func (p *Pool) Stop() { if entry.state.cancel != nil { entry.state.cancel() } - if entry.state.client != nil { - _ = entry.state.client.Close() + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() } slog.Debug("mcp.pool.stopped", "key", key) } @@ -401,8 +401,8 @@ func (p *Pool) Stop() { if entry.state.cancel != nil { entry.state.cancel() } - if entry.state.client != nil { - _ = entry.state.client.Close() + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() } slog.Debug("mcp.pool.user.stopped", "key", key) } @@ -423,8 +423,8 @@ func (p *Pool) Evict(tenantID uuid.UUID, serverName string) { if entry.state.cancel != nil { entry.state.cancel() } - if entry.state.client != nil { - _ = entry.state.client.Close() + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() } delete(p.servers, key) select { @@ -481,8 +481,8 @@ func (p *Pool) evictIdle() { if entry.state.cancel != nil { entry.state.cancel() } - if entry.state.client != nil { - _ = entry.state.client.Close() + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() } delete(p.servers, key) select { @@ -499,8 +499,8 @@ func (p *Pool) evictIdle() { if entry.state.cancel != nil { entry.state.cancel() } - if entry.state.client != nil { - _ = entry.state.client.Close() + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() } delete(p.userServers, key) // Return slot to per-server semaphore @@ -588,9 +588,6 @@ func (p *Pool) evictOldestIdleLocked() bool { return true } -// Client returns the MCP client for this pool entry. -func (e *poolEntry) Client() *mcpclient.Client { return e.state.client } - // ClientPtr returns the atomic client pointer for this pool entry. // Used by BridgeTools to atomically load the current client during reconnect. func (e *poolEntry) ClientPtr() *atomic.Pointer[mcpclient.Client] { return &e.state.clientPtr } @@ -645,52 +642,8 @@ func poolHealthLoop(ctx context.Context, ss *serverState) { } } -// poolTryReconnect attempts a full reconnect for a pool-managed connection. -// Creates a fresh client, atomically swaps ss.clientPtr so all BridgeTools -// see the new client, then closes the old one. +// poolTryReconnect attempts reconnect for a pool-managed connection. +// Delegates to the shared reconnectWithBackoff with pool-specific log prefix. func poolTryReconnect(ctx context.Context, ss *serverState) { - ss.mu.Lock() - if ss.reconnAttempts >= maxReconnectAttempts { - ss.lastErr = fmt.Sprintf("max reconnect attempts (%d) reached, entering cooldown", maxReconnectAttempts) - ss.mu.Unlock() - slog.Warn("mcp.pool.reconnect_cooldown", "server", ss.name, "cooldown", reconnectCooldown) - select { - case <-ctx.Done(): - return - case <-time.After(reconnectCooldown): - } - ss.mu.Lock() - ss.reconnAttempts = 0 - ss.mu.Unlock() - return - } - ss.reconnAttempts++ - attempt := ss.reconnAttempts - ss.mu.Unlock() - - backoff := min(initialBackoff*time.Duration(1<<(attempt-1)), maxBackoff) - slog.Info("mcp.pool.reconnecting", "server", ss.name, "attempt", attempt, "backoff", backoff) - - select { - case <-ctx.Done(): - return - case <-time.After(backoff): - } - - // Fast path: transport may have auto-reconnected. - if err := ss.client.Ping(ctx); err == nil { - ss.connected.Store(true) - ss.mu.Lock() - ss.reconnAttempts = 0 - ss.healthFailures = 0 - ss.lastErr = "" - ss.mu.Unlock() - slog.Info("mcp.pool.reconnected", "server", ss.name) - return - } - - // Slow path: create fresh client, swap atomically, close old. - if fullReconnect(ctx, ss) { - slog.Info("mcp.pool.reconnected", "server", ss.name, "method", "full_reconnect") - } + reconnectWithBackoff(ctx, ss, "mcp.pool") }