3b3a3d7 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: prevent goroutine orphaning in mcp.Close() and shell.KillAll() (#2159)

3b3a3d7M1xA8 months ago

internal/agent/tools/mcp/init.go+16−13

compare with main

unifiedsplit

@@ -124,26 +124,29 @@ func GetState(name string) (ClientInfo, bool) {

124

125// Close closes all MCP clients. This should be called during application shutdown.

126func Close() error {

127+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

128+ defer cancel()

129+

127 var wg sync.WaitGroup

130 var wg sync.WaitGroup

128− done := make(chan struct{}, 1)

129− go func() {

130− for name, session := range sessions.Seq2() {

131+ for name, session := range sessions.Seq2() {

131− wg.Go(func() {

132+ wg.Go(func() {

133+ done := make(chan error, 1)

134+ go func() {

135+ done <- session.Close()

136+ }()

137+ select {

138+ case err := <-done:

132− if err := session.Close(); err != nil &&

139+ if err != nil &&

133 !errors.Is(err, io.EOF) &&

140 !errors.Is(err, io.EOF) &&

134 !errors.Is(err, context.Canceled) &&

141 !errors.Is(err, context.Canceled) &&

135 err.Error() != "signal: killed" {

142 err.Error() != "signal: killed" {

136 slog.Warn("Failed to shutdown MCP client", "name", name, "error", err)

143 slog.Warn("Failed to shutdown MCP client", "name", name, "error", err)

137 }

144 }

138− })

145+ case <-ctx.Done():

139− }

146+ }

140− wg.Wait()

141− done <- struct{}{}

142− }()

147+ })

143− select {

144− case <-done:

145− case <-time.After(5 * time.Second):

146 }

148 }

149+ wg.Wait()

147 broker.Shutdown()

150 broker.Shutdown()

148 return nil

151 return nil

149}

152}

internal/app/app.go+5−3

@@ -540,14 +540,16 @@ func (app *App) Shutdown() {

540 // Now run remaining cleanup tasks in parallel.

541 var wg sync.WaitGroup

542

543+ // Shared shutdown context for all timeout-bounded cleanup.

544+ shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)

545+ defer cancel()

546+

543 // Kill all background shells.

547 // Kill all background shells.

544 wg.Go(func() {

548 wg.Go(func() {

545− shell.GetBackgroundShellManager().KillAll()

549+ shell.GetBackgroundShellManager().KillAll(shutdownCtx)

546 })

550 })

547

551

548 // Shutdown all LSP clients.

552 // Shutdown all LSP clients.

549− shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)

550− defer cancel()

551 wg.Go(func() {

553 wg.Go(func() {

552 app.LSPManager.StopAll(shutdownCtx)

554 app.LSPManager.StopAll(shutdownCtx)

553 })

555 })

internal/shell/background.go+13−19

@@ -191,29 +191,23 @@ func (m *BackgroundShellManager) Cleanup() int {

191 return len(toRemove)

192}

193

194−// KillAll terminates all background shells.

194+// KillAll terminates all background shells. The provided context bounds how

195+// long the function waits for each shell to exit.

195−func (m *BackgroundShellManager) KillAll() {

196+func (m *BackgroundShellManager) KillAll(ctx context.Context) {

196 shells := slices.Collect(m.shells.Seq())

197 shells := slices.Collect(m.shells.Seq())

197 m.shells.Reset(map[string]*BackgroundShell{})

198 m.shells.Reset(map[string]*BackgroundShell{})

198− done := make(chan struct{}, 1)

199− go func() {

200− var wg sync.WaitGroup

201− for _, shell := range shells {

202− wg.Go(func() {

203− shell.cancel()

204− <-shell.done

205− })

206− }

207− wg.Wait()

208− done <- struct{}{}

209− }()

210

199

200+ var wg sync.WaitGroup

201+ for _, shell := range shells {

202+ wg.Go(func() {

203+ shell.cancel()

211− select {

204+ select {

212− case <-done:

205+ case <-shell.done:

213− return

206+ case <-ctx.Done():

214− case <-time.After(time.Second * 5):

207+ }

215− return

208+ })

216 }

209 }

210+ wg.Wait()

217}

211}

218

212

219// GetOutput returns the current output of a background shell.

213// GetOutput returns the current output of a background shell.

internal/shell/background_test.go+25−1

@@ -6,6 +6,8 @@ import (

6 "strings"

7 "testing"

8 "time"

10+ "github.com/stretchr/testify/require"

11)

11func TestBackgroundShellManager_Start(t *testing.T) {

13func TestBackgroundShellManager_Start(t *testing.T) {

@@ -248,7 +250,7 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {

248 }

250 }

249

251

250 // Kill all shells

252 // Kill all shells

251− manager.KillAll()

253+ manager.KillAll(context.Background())

252

254

253 // Verify all shells are done

255 // Verify all shells are done

254 if !shell1.IsDone() {

256 if !shell1.IsDone() {

@@ -280,3 +282,25 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {

280 }

282 }

281 }

283 }

282}

284}

285+

286+func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {

287+ t.Parallel()

288+

289+ workingDir := t.TempDir()

290+ manager := newBackgroundShellManager()

291+

292+ // Start a shell that traps signals and ignores cancellation.

293+ _, err := manager.Start(context.Background(), workingDir, nil, "trap '' TERM INT; sleep 60", "")

294+ require.NoError(t, err)

295+

296+ // Short timeout to test the timeout path.

297+ ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)

298+ defer cancel()

299+

300+ start := time.Now()

301+ manager.KillAll(ctx)

302+ elapsed := time.Since(start)

303+

304+ // Must return promptly after timeout, not hang for 60 seconds.

305+ require.Less(t, elapsed, 2*time.Second)

306+}