mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 04:14:56 +00:00
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:
1 parent
86ef70906a
commit
de6bdf2290
3 files changed
+38
-82
No files matched your search
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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")
|
||||
}
|
||||
Reference in new issue
Block a user