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()
This commit is contained in:
viettranx committed 2026-04-10 20:32:56 +07:00
1 parent 86ef70906a
commit de6bdf2290
3 files changed
+38 -82

No files matched your search

+3 -2
View File
@@ -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)
}
}
+18 -16
View File
@@ -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)
+17 -64
View File
@@ -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")
}