Add graceful shutdown handling for MCP server

[sgaunet]

Problem

The MCP server in main.go:521 blocks indefinitely on server.ServeStdio() without signal handling for graceful shutdown.

Current implementation:

func main() {
    // ... initialization ...
    
    if err := server.ServeStdio(); err != nil {
        logger.Log().Fatal().Err(err).Msg("Server error")
    }
}

Issues

1. No cleanup on SIGTERM/SIGINT: Database connections not closed properly
2. Hanging resources: In-flight operations may not complete
3. Abrupt termination: No opportunity to finish pending work
4. Lost logs: Buffered logs may not flush
5. Docker unfriendly: Containers can't gracefully stop

Impact Scenarios

Scenario 1: User presses Ctrl+C

Current behavior:

• Process terminates immediately
• Database connections may be left open
• No cleanup logs

Desired behavior:

• Catch SIGINT signal
• Close database connections
• Flush logs
• Exit cleanly

Scenario 2: Container shutdown (Docker/Kubernetes)

Current behavior:

• Container receives SIGTERM
• Process doesn't handle it
• After timeout, receives SIGKILL (forced termination)

Desired behavior:

• Handle SIGTERM gracefully
• Complete in-flight operations (with timeout)
• Close resources
• Exit before SIGKILL

Proposed Solution

Option 1: Context with Signal Handling (Recommended)

func main() {
    // ... initialization ...
    
    // Create context that cancels on interrupt signals
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    
    // Handle interrupt signals
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
    
    // Goroutine to handle shutdown
    go func() {
        sig := <-sigCh
        logger.Log().Info().Str("signal", sig.String()).Msg("Received shutdown signal")
        cancel()
    }()
    
    // Serve with context
    if err := server.ServeStdioWithContext(ctx); err != nil && err != context.Canceled {
        logger.Log().Fatal().Err(err).Msg("Server error")
    }
    
    // Cleanup
    cleanup()
    logger.Log().Info().Msg("Server shutdown complete")
}

func cleanup() {
    logger.Log().Info().Msg("Starting cleanup...")
    
    // Close database connections
    if appInstance != nil && appInstance.client != nil {
        if err := appInstance.client.Close(); err != nil {
            logger.Log().Error().Err(err).Msg("Error closing database connection")
        } else {
            logger.Log().Info().Msg("Database connection closed")
        }
    }
    
    // Flush logs
    // Add any other cleanup tasks
}

Option 2: Timeout for Graceful Shutdown

const shutdownTimeout = 30 * time.Second

func main() {
    // ... initialization ...
    
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
    
    errCh := make(chan error, 1)
    go func() {
        errCh <- server.ServeStdioWithContext(ctx)
    }()
    
    select {
    case sig := <-sigCh:
        logger.Log().Info().Str("signal", sig.String()).Msg("Received shutdown signal")
        
        // Give server time to finish current operations
        shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
        defer shutdownCancel()
        
        cancel() // Stop accepting new requests
        
        // Wait for server to finish or timeout
        select {
        case <-shutdownCtx.Done():
            logger.Log().Warn().Msg("Shutdown timeout exceeded, forcing exit")
        case err := <-errCh:
            if err != nil && err != context.Canceled {
                logger.Log().Error().Err(err).Msg("Server error during shutdown")
            }
        }
        
        cleanup()
        
    case err := <-errCh:
        if err != nil {
            logger.Log().Fatal().Err(err).Msg("Server error")
        }
    }
    
    logger.Log().Info().Msg("Server shutdown complete")
}

Option 3: Minimal Signal Handling

If the MCP library doesn't support context-based shutdown:

func main() {
    // ... initialization ...
    
    // Setup cleanup on exit
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
    
    go func() {
        <-sigCh
        logger.Log().Info().Msg("Received shutdown signal, cleaning up...")
        cleanup()
        os.Exit(0)
    }()
    
    if err := server.ServeStdio(); err != nil {
        cleanup()
        logger.Log().Fatal().Err(err).Msg("Server error")
    }
}

Recommended Approach

Use Option 1 if the MCP server library supports context, otherwise use Option 3 for basic signal handling.

Benefits

• ✅ Clean database connection closure
• ✅ Proper log flushing
• ✅ Docker/Kubernetes friendly
• ✅ Better resource management
• ✅ Prevents connection leaks

Testing Graceful Shutdown

# Start server
./postgresql-mcp

# In another terminal, send signals
kill -SIGTERM <pid>   # Should shutdown gracefully
kill -SIGINT <pid>    # Should shutdown gracefully (Ctrl+C)

# Check logs for cleanup messages

Impact

• Severity: LOW
• Type: Enhancement
• Location: main.go:521
• Benefits: Proper cleanup, Docker-friendly, better resource management

Checklist

• Check if MCP server library supports context-based shutdown
• Implement signal handling (SIGTERM, SIGINT)
• Add cleanup function to close database connections
• Add graceful shutdown timeout (30 seconds recommended)
• Add logging for shutdown events
• Test with Ctrl+C and kill signals
• Test in Docker container
• Document shutdown behavior in README
• Consider adding health check endpoint for orchestration systems