Implementing Ping-Pong for WebSocket Health Checks in Go

WebSockets are a flexible protocol that enable persistent connections between server and client. To ensure these connections remain healthy, it’s beneficial to implement health checks. One common method is using ping-pong frames to verify that the server and client are still active and responsive.

Setting Up a Basic WebSocket Server

Before we dive into implementing the ping-pong functionality, let’s set up a basic WebSocket server using the Gorilla WebSocket package in Go. First, ensure you have the package installed:

go get -u github.com/gorilla/websocket

Next, we will create a simple WebSocket server:

package main

import (
    "net/http"
    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

func handleConnections(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Println(err)
        return
    }

    defer ws.Close()

    for {
        msgType, msg, err := ws.ReadMessage()
        if err != nil {
            fmt.Println(err)
            break
        }
        ws.WriteMessage(msgType, msg)
    }
}

func main() {
    http.HandleFunc("/ws", handleConnections)
    http.ListenAndServe(":8080", nil)
}

This code establishes a basic WebSocket server that echoes back messages received from the client.

Implementing Ping-Pong

To include health checks using ping-pong frames, we need to periodically send a "ping" frame to the client. The client should automatically respond with a "pong", and if it doesn’t, we can consider the connection unresponsive. Here are the necessary modifications:

First, enhance the server with a ping handler:

import (
    "time"
    "log"
)

func handleConnections(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Println(err)
        return
    }

    defer ws.Close()

    ws.SetPongHandler(func(appData string) error {
        log.Println("pong received")
        return nil
    })

    ticker := time.NewTicker(10 * time.Second)
    defer ticker.Stop()

    done := make(chan struct{})

    go func() {
        defer close(done)
        for {
            _, _, err := ws.ReadMessage()
            if err != nil {
                fmt.Println(err)
                break
            }
        }
    }()

    for {
        select {
        case <-done:
            return
        case t := <-ticker.C:
            log.Println("Ping: ", t)
            err := ws.WriteMessage(websocket.PingMessage, nil)
            if err != nil {
                fmt.Println("Error sending Ping: ", err)
                return
            }
        }
    }
}

This code uses a ticker to regularly send ping frames to the client every 10 seconds. The SetPongHandler registers a callback for when a pong frame is received, enabling efficient monitoring of connection health.

Testing the Implementation

Run the server:

go run main.go

With the server running, you can use a WebSocket client to connect to ws://localhost:8080/ws and observe the logs to see ping and pong messages, indicating the connection health.

Conclusion

Implementing ping-pong for WebSocket health checks in Go involves periodic ping frames and pong frame handlers. This ensures that our WebSocket connections remain functional and responsive, providing a safeguard for real-time communications over the Web.