Skip to content

Commit 447609e

Browse files
DumbrisDumbris
authored
Merge pull request #61 from smart-mcp-proxy/feature/set-process-group
Implement process group management for Unix and Windows systems
2 parents 67000ed + 2d57276 commit 447609e

4 files changed

Lines changed: 274 additions & 5 deletions

File tree

‎internal/upstream/core/client.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ type Client struct {
6767

6868
// Process monitoring (for stdio transport)
6969
processCmd *exec.Cmd
70+
processGroupID int // Process group ID for proper cleanup
7071
processMonitorCtx context.Context
7172
processMonitorCancel context.CancelFunc
7273
processMonitorWG sync.WaitGroup

‎internal/upstream/core/connection.go‎

Lines changed: 71 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,21 @@ func (c *Client) Connect(ctx context.Context) error {
152152
}
153153
}
154154

155+
// CRITICAL FIX: Also cleanup process groups to prevent zombie processes on connection failure
156+
if c.processGroupID > 0 {
157+
c.logger.Warn("Connection failed - cleaning up process group to prevent zombie processes",
158+
zap.String("server", c.config.Name),
159+
zap.Int("pgid", c.processGroupID))
160+
161+
if err := killProcessGroup(c.processGroupID, c.logger, c.config.Name); err != nil {
162+
c.logger.Error("Failed to clean up process group after connection failure",
163+
zap.String("server", c.config.Name),
164+
zap.Int("pgid", c.processGroupID),
165+
zap.Error(err))
166+
}
167+
c.processGroupID = 0
168+
}
169+
155170
return fmt.Errorf("failed to connect: %w", err)
156171
}
157172

@@ -309,16 +324,18 @@ func (c *Client) connectStdio(ctx context.Context) error {
309324
}
310325
}
311326

312-
// Upstream transport with working directory support
327+
// Upstream transport with working directory support and process group management
313328
var stdioTransport *uptransport.Stdio
314329
if c.config.WorkingDir != "" {
315-
// Use custom CommandFunc to set working directory
316-
commandFunc := createWorkingDirCommandFunc(c.config.WorkingDir)
330+
// CRITICAL FIX: Use enhanced CommandFunc with process group management for proper cleanup
331+
commandFunc := createEnhancedWorkingDirCommandFunc(c.config.WorkingDir, c.logger)
317332
stdioTransport = uptransport.NewStdioWithOptions(finalCommand, envVars, finalArgs,
318333
uptransport.WithCommandFunc(commandFunc))
319334
} else {
320-
// Use default transport for backwards compatibility
321-
stdioTransport = uptransport.NewStdio(finalCommand, envVars, finalArgs...)
335+
// CRITICAL FIX: Use enhanced CommandFunc even without working directory to ensure process groups
336+
commandFunc := createEnhancedWorkingDirCommandFunc("", c.logger)
337+
stdioTransport = uptransport.NewStdioWithOptions(finalCommand, envVars, finalArgs,
338+
uptransport.WithCommandFunc(commandFunc))
322339
}
323340

324341
c.client = client.NewClient(stdioTransport)
@@ -370,6 +387,21 @@ func (c *Client) connectStdio(ctx context.Context) error {
370387
c.killDockerContainerByCommandWithContext(cleanupCtx)
371388
}
372389
}
390+
391+
// CRITICAL FIX: Also cleanup process groups to prevent zombie processes on initialization failure
392+
if c.processGroupID > 0 {
393+
c.logger.Warn("Initialization failed - cleaning up process group to prevent zombie processes",
394+
zap.String("server", c.config.Name),
395+
zap.Int("pgid", c.processGroupID))
396+
397+
if err := killProcessGroup(c.processGroupID, c.logger, c.config.Name); err != nil {
398+
c.logger.Error("Failed to clean up process group after initialization failure",
399+
zap.String("server", c.config.Name),
400+
zap.Int("pgid", c.processGroupID),
401+
zap.Error(err))
402+
}
403+
c.processGroupID = 0
404+
}
373405
return fmt.Errorf("MCP initialize failed for stdio transport: %w", err)
374406
}
375407

@@ -405,6 +437,14 @@ func (c *Client) connectStdio(ctx context.Context) error {
405437
c.logger.Info("Successfully extracted process from stdio transport for lifecycle management",
406438
zap.String("server", c.config.Name),
407439
zap.Int("pid", c.processCmd.Process.Pid))
440+
441+
// CRITICAL FIX: Extract process group ID for proper cleanup
442+
c.processGroupID = extractProcessGroupID(cmd, c.logger, c.config.Name)
443+
if c.processGroupID > 0 {
444+
c.logger.Info("Process group ID tracked for cleanup",
445+
zap.String("server", c.config.Name),
446+
zap.Int("pgid", c.processGroupID))
447+
}
408448
}
409449
} else {
410450
c.logger.Warn("Could not extract process from stdio transport - will use alternative process tracking",
@@ -653,6 +693,11 @@ func createWorkingDirCommandFunc(workingDir string) uptransport.CommandFunc {
653693
}
654694
}
655695

696+
// createEnhancedWorkingDirCommandFunc creates a custom CommandFunc with process group management
697+
func createEnhancedWorkingDirCommandFunc(workingDir string, logger *zap.Logger) uptransport.CommandFunc {
698+
return createProcessGroupCommandFunc(workingDir, logger)
699+
}
700+
656701
// connectHTTP establishes HTTP transport connection with auth fallback
657702
func (c *Client) connectHTTP(ctx context.Context) error {
658703
// Try authentication strategies in order: headers -> no-auth -> OAuth
@@ -1435,6 +1480,27 @@ func (c *Client) DisconnectWithContext(_ context.Context) error {
14351480
zap.String("server", c.config.Name))
14361481
}
14371482

1483+
// CRITICAL FIX: Clean up process group to prevent zombie processes
1484+
if c.processGroupID > 0 {
1485+
c.logger.Info("Cleaning up process group to prevent zombie processes",
1486+
zap.String("server", c.config.Name),
1487+
zap.Int("pgid", c.processGroupID))
1488+
1489+
if err := killProcessGroup(c.processGroupID, c.logger, c.config.Name); err != nil {
1490+
c.logger.Error("Failed to clean up process group",
1491+
zap.String("server", c.config.Name),
1492+
zap.Int("pgid", c.processGroupID),
1493+
zap.Error(err))
1494+
} else {
1495+
c.logger.Info("Process group cleanup completed successfully",
1496+
zap.String("server", c.config.Name),
1497+
zap.Int("pgid", c.processGroupID))
1498+
}
1499+
1500+
// Reset process group ID after cleanup
1501+
c.processGroupID = 0
1502+
}
1503+
14381504
c.logger.Debug("Closing MCP client connection",
14391505
zap.String("server", c.config.Name))
14401506
c.client.Close()

‎internal/upstream/core/process_unix.go‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
//go:build unix
2+
3+
package core
4+
5+
import (
6+
"context"
7+
"os/exec"
8+
"syscall"
9+
"time"
10+
11+
"go.uber.org/zap"
12+
)
13+
14+
// ProcessGroup represents a Unix process group for proper child process management
15+
type ProcessGroup struct {
16+
PGID int
17+
}
18+
19+
// createProcessGroupCommandFunc creates a custom CommandFunc that sets process groups for Unix systems
20+
// This ensures that child processes can be properly cleaned up when the parent exits
21+
func createProcessGroupCommandFunc(workingDir string, logger *zap.Logger) func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error) {
22+
return func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error) {
23+
cmd := exec.CommandContext(ctx, command, args...)
24+
cmd.Env = env
25+
26+
// Set working directory if specified
27+
if workingDir != "" {
28+
cmd.Dir = workingDir
29+
}
30+
31+
// CRITICAL FIX: Set process group attributes for proper child process management
32+
// This ensures that when mcpproxy exits, all child processes can be terminated properly
33+
cmd.SysProcAttr = &syscall.SysProcAttr{
34+
// Create a new process group for this command and its children
35+
Setpgid: true,
36+
// Make this process the group leader
37+
Pgid: 0,
38+
}
39+
40+
logger.Debug("Process group configuration applied",
41+
zap.String("command", command),
42+
zap.Strings("args", args),
43+
zap.String("working_dir", workingDir))
44+
45+
return cmd, nil
46+
}
47+
}
48+
49+
// killProcessGroup terminates an entire process group on Unix systems
50+
// This is the proper way to clean up child processes and prevent zombies
51+
func killProcessGroup(pgid int, logger *zap.Logger, serverName string) error {
52+
if pgid <= 0 {
53+
return nil
54+
}
55+
56+
logger.Info("Terminating process group",
57+
zap.String("server", serverName),
58+
zap.Int("pgid", pgid))
59+
60+
// Step 1: Send SIGTERM to the entire process group
61+
err := syscall.Kill(-pgid, syscall.SIGTERM)
62+
if err != nil {
63+
logger.Warn("Failed to send SIGTERM to process group",
64+
zap.String("server", serverName),
65+
zap.Int("pgid", pgid),
66+
zap.Error(err))
67+
} else {
68+
logger.Debug("SIGTERM sent to process group",
69+
zap.String("server", serverName),
70+
zap.Int("pgid", pgid))
71+
}
72+
73+
// Step 2: Wait a bit for graceful termination
74+
time.Sleep(2 * time.Second)
75+
76+
// Step 3: Check if processes are still running and send SIGKILL if needed
77+
if err := syscall.Kill(-pgid, 0); err == nil {
78+
// Processes still exist, force kill them
79+
logger.Warn("Process group still running after SIGTERM, sending SIGKILL",
80+
zap.String("server", serverName),
81+
zap.Int("pgid", pgid))
82+
83+
if killErr := syscall.Kill(-pgid, syscall.SIGKILL); killErr != nil {
84+
logger.Error("Failed to send SIGKILL to process group",
85+
zap.String("server", serverName),
86+
zap.Int("pgid", pgid),
87+
zap.Error(killErr))
88+
return killErr
89+
}
90+
91+
logger.Info("SIGKILL sent to process group",
92+
zap.String("server", serverName),
93+
zap.Int("pgid", pgid))
94+
} else {
95+
logger.Info("Process group terminated successfully",
96+
zap.String("server", serverName),
97+
zap.Int("pgid", pgid))
98+
}
99+
100+
return nil
101+
}
102+
103+
// extractProcessGroupID extracts the process group ID from a running command
104+
func extractProcessGroupID(cmd *exec.Cmd, logger *zap.Logger, serverName string) int {
105+
if cmd == nil || cmd.Process == nil {
106+
return 0
107+
}
108+
109+
// Get the process group ID from the process
110+
pgid, err := syscall.Getpgid(cmd.Process.Pid)
111+
if err != nil {
112+
logger.Warn("Failed to get process group ID",
113+
zap.String("server", serverName),
114+
zap.Int("pid", cmd.Process.Pid),
115+
zap.Error(err))
116+
return 0
117+
}
118+
119+
logger.Debug("Process group ID extracted",
120+
zap.String("server", serverName),
121+
zap.Int("pid", cmd.Process.Pid),
122+
zap.Int("pgid", pgid))
123+
124+
return pgid
125+
}

‎internal/upstream/core/process_windows.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
//go:build windows
2+
3+
package core
4+
5+
import (
6+
"context"
7+
"os/exec"
8+
9+
"go.uber.org/zap"
10+
)
11+
12+
// ProcessGroup represents a Windows process group for proper child process management
13+
type ProcessGroup struct {
14+
PGID int
15+
logger *zap.Logger
16+
}
17+
18+
// createProcessGroupCommandFunc creates a custom CommandFunc for Windows systems
19+
// Note: Windows process group management is different from Unix and requires different approaches
20+
func createProcessGroupCommandFunc(workingDir string, logger *zap.Logger) func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error) {
21+
return func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error) {
22+
cmd := exec.CommandContext(ctx, command, args...)
23+
cmd.Env = env
24+
25+
// Set working directory if specified
26+
if workingDir != "" {
27+
cmd.Dir = workingDir
28+
}
29+
30+
// TODO: Implement Windows-specific process group management
31+
// Windows uses Job Objects instead of process groups
32+
// For now, we'll use the standard command creation
33+
34+
logger.Debug("Process group configuration applied (Windows)",
35+
zap.String("command", command),
36+
zap.Strings("args", args),
37+
zap.String("working_dir", workingDir))
38+
39+
return cmd, nil
40+
}
41+
}
42+
43+
// killProcessGroup terminates processes on Windows systems
44+
// This is a simplified implementation for Windows compatibility
45+
func killProcessGroup(pgid int, logger *zap.Logger, serverName string) error {
46+
// TODO: Implement proper Windows process termination
47+
// For now, this is a placeholder that does nothing
48+
// Windows process management would require Win32 API calls or Job Objects
49+
50+
logger.Debug("Process group termination requested (Windows placeholder)",
51+
zap.String("server", serverName),
52+
zap.Int("pgid", pgid))
53+
54+
return nil
55+
}
56+
57+
// extractProcessGroupID extracts the process group ID from a running command on Windows
58+
func extractProcessGroupID(cmd *exec.Cmd, logger *zap.Logger, serverName string) int {
59+
// Windows doesn't have Unix-style process groups
60+
// Return the PID as a fallback identifier
61+
if cmd == nil || cmd.Process == nil {
62+
return 0
63+
}
64+
65+
logger.Debug("Process group ID extracted (Windows - using PID)",
66+
zap.String("server", serverName),
67+
zap.Int("pid", cmd.Process.Pid))
68+
69+
return cmd.Process.Pid
70+
}
71+
72+
// isProcessGroupAlive checks if processes are still running on Windows
73+
func isProcessGroupAlive(pgid int) bool {
74+
// TODO: Implement Windows-specific process checking
75+
// For now, return false as a safe default
76+
return false
77+
}

0 commit comments

Comments
 (0)