Skip to content

[WIP] Fix WebSocket server close handling for graceful shutdown#51

Closed
Copilot wants to merge 1 commit intomasterfrom
copilot/fix-websocket-close-timeout
Closed

[WIP] Fix WebSocket server close handling for graceful shutdown#51
Copilot wants to merge 1 commit intomasterfrom
copilot/fix-websocket-close-timeout

Conversation

Copy link
Copy Markdown

Copilot AI commented Mar 25, 2026

Thanks for asking me to work on this. I will get started on it and keep this PR's description up to date as I form a plan and make progress.

Original prompt

Problem Statement: Graceful WebSocket Shutdown without Client Hang

Current Issue

The WebSocket server's Close() method in internal/websocket/server.go (lines 519-533) doesn't have proper timeout handling for gracefully shutting down WebSocket connections. This can cause:

  1. Client Hang: WebSocket clients may hang indefinitely waiting for close confirmation
  2. Resource Leak: Connected clients aren't properly notified of server shutdown
  3. Delayed Server Shutdown: Server shutdown can be blocked by slow/unresponsive clients
  4. No Graceful Message: Clients don't receive warning before disconnection

Current Implementation (Problematic)

func (s *Server) Close() error {
    s.mutex.Lock()
    defer s.mutex.Unlock()

    // Close all client connections
    for _, client := range s.clients {
        client.Close()
    }

    // Clear clients map
    s.clients = make(map[string]*Client)

    s.logger.Info("WebSocket server closed gracefully")
    return nil
}

Problems:

  • No timeout protection - infinite wait possible
  • No graceful disconnect message to clients
  • Client goroutines may still be running
  • No metrics on graceful shutdown

Solution Requirements

  1. Add Graceful Shutdown Method:

    • Shutdown(ctx context.Context) with configurable timeout
    • Send close message to all clients
    • Wait for all client handlers to exit (with timeout)
    • Log shutdown statistics
  2. Modify Server API Structure:

    • Add shutdown channel to track closing state
    • Add WaitGroup to track active client handlers
    • Add metrics for shutdown process
  3. Update API Server Integration (internal/api/server.go):

    • Pass WebSocket server reference to Shutdown method
    • Call ws shutdown before storage close
    • Use context timeout for graceful period
  4. Add Unit Tests:

    • Test graceful shutdown with active clients
    • Test timeout behavior
    • Test that clients receive close notification
    • Test concurrent client disconnection
  5. Add Integration Example:

    • Update main.go or cmd/api/main.go to demonstrate proper shutdown sequence
    • Document recommended shutdown timeout values

Expected Behavior After Fix

Shutdown sequence:
1. Server receives shutdown signal
2. Broadcast close notification to all clients
3. Close listener (no new connections)
4. Wait for existing clients to disconnect (with timeout)
5. Force close remaining connections after timeout
6. Log statistics: clients closed, timeout count, duration
7. Return with all goroutines cleaned up

Files to Modify

  1. internal/websocket/server.go - Add graceful shutdown with timeout
  2. internal/websocket/client.go - Add close notification flag
  3. internal/api/server.go - Integrate WebSocket shutdown into API server shutdown
  4. internal/websocket/server_test.go - Add comprehensive shutdown tests
  5. cmd/api/main.go or equivalent - Document shutdown sequence

Testing Approach

  • Unit tests with mock clients
  • Integration tests with real WebSocket connections
  • Stress tests with many concurrent clients
  • Timeout edge case tests

The following is the prior conversation context from the user's chat exploration (may be truncated):

User: Проанализируй репозиторий и найди плохие моменты по типу варианты для оптимизации и тд

Assistant: Я проанализирую репозиторий godofphonk/ServerEyeAPI. Начну с получения информации о его структуре и содержимом.

Assistant: Отлично! Теперь давайте подробнее изучим основные файлы проекта:

Assistant: Теперь получу более детальную информацию о коде:

Assistant: Let me search for more specific issues:

Assistant: Отлично! Я проанализировал репозиторий ServerEyeAPI. Вот подробный отчет о найденных проблемах и возможностях оптимизации:

🚨 Критические проблемы и "плохие моменты"

1. Отключенная аутентификация API (🔓 Уязвимость безопасности)

Файл: internal/api/server.go, строка 140

// Initialize API Key middleware (TODO: Fix and enable)
// apiKeyMiddleware := keyMiddleware.NewAPIKeyAuthMiddleware(apiKeyStorage, logger)

// Setup routes
router := SetupRoutes(
    ...
    nil, // TODO: apiKeyMiddleware

Проблема: Middleware для проверки API ключей закомментирована, что означает, что все защищенные эндпоинты не защищены.
Риск: Любой может получить доступ к confidential данным.
Решение: Активировать middleware и применить его к защищенным маршрутам.


2. CORS открыт для всех (⚠️ Security)

Файл: internal/api/middleware/cors.go

w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")

Проблема: CORS позволяет запросам с любых источников, что опасно для API.
Решение: Ограничить Access-Control-Allow-Origin известными доменами из конфигурации.


3. Неправильное преобразование server_key → server_id

Файл: internal/handlers/static_info.go, множественные места

// ...

</details>



<!-- START COPILOT CODING AGENT SUFFIX -->

*This pull request was created from Copilot chat.*
>

<!-- START COPILOT CODING AGENT TIPS -->
---Quickly spin up Copilot coding agent tasks from anywhere on your macOS or Windows machine with [Raycast](https://gh.io/cca-raycast-docs).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants