Go and WebSocket Cluster - Enhancing Distributed User Support!

In this episode, we will discuss creating a WebSocket Server that can operate in a Cluster mode to support distributed usage, increasing the capacity to handle a large number of users and improving system stability!

What is a Cluster?

A Cluster is a combination of multiple servers that work together as a single system, with a Load Balancer distributing tasks (such as WebSocket connections) to each server in the Cluster.

Advantages of Working in a Cluster :

1. Supports a large number of users : Distributes connections to multiple servers.

2. Increases stability : If some servers fail, the system continues to operate.

3. Easily scalable : Add or reduce the number of servers as needed.

Managing WebSocket in a Cluster

The main issues to address in a Cluster :

1. Session Stickiness : Ensures that a user's connections are sent to the same server.

2. Message Broadcasting : Allows messages sent from one server to be broadcast to all users in the Cluster.

3. State Synchronization : Maintains the state of users across multiple servers.

Building a WebSocket Cluster with Redis Pub/Sub

Redis Pub/Sub allows multiple servers to communicate with each other, where each server subscribes to messages from the same channel.

Example of Using Redis Pub/Sub for WebSocket Cluster

1. Install Redis and Go Redis Library

Install Redis on your server.

sudo apt install redis

add Library Go Redis: 

go get github.com/redis/go-redis/v9

2. Example Code for WebSocket Cluster

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
    "github.com/redis/go-redis/v9"
)

var (
    upgrader = websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
    clients  = make(map[*websocket.Conn]bool)
    broadcast = make(chan []byte)
    rdb      = redis.NewClient(&redis.Options{
        Addr: "localhost:6379",
    })
    ctx = context.Background()
)

func handleConnections(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("Error upgrading connection:", err)
        return
    }
    defer conn.Close()

    clients[conn] = true
    defer delete(clients, conn)

    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("Error reading message:", err)
            break
        }

        // ส่งข้อความไปยัง Redis Channel
        if err := rdb.Publish(ctx, "websocketChannel", msg).Err(); err != nil {
            log.Println("Error publishing to Redis:", err)
        }
    }
}

func handleRedisMessages() {
    sub := rdb.Subscribe(ctx, "websocketChannel")
    ch := sub.Channel()

    for msg := range ch {
        log.Printf("Received message from Redis: %s", msg.Payload)
        broadcast <- []byte(msg.Payload)
    }
}

func handleBroadcasts() {
    for {
        msg := <-broadcast
        for client := range clients {
            if err := client.WriteMessage(websocket.TextMessage, msg); err != nil {
                log.Println("Error writing message:", err)
                client.Close()
                delete(clients, client)
            }
        }
    }
}

func main() {
    http.HandleFunc("/ws", handleConnections)

    go handleBroadcasts()
    go handleRedisMessages()

    log.Println("WebSocket server started at :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

Code Explanation:

1. Connect to Redis : All servers subscribe to websocketChannel.

2. Send Messages to Redis : When a user sends a message, it is published to Redis.

3. Broadcast Messages from Redis : Messages received from Redis are sent to all users connected to the server.

Testing the Cluster System:

1. Run multiple WebSocket Servers (on different ports, e.g., :8080, :8081).
2. Use NGINX or a Load Balancer to distribute requests to each server.
3. Connect multiple clients and send messages to observe the results.

Further Improvements:

1. Add Sticky Sessions : Use a Load Balancer like NGINX to track user sessions.
2. Monitor Redis : Use tools like RedisInsight to check messages and Redis channels.
3. Enable Auto Scaling : Use Kubernetes to manage the number of servers.

Summary:

• Use Redis Pub/Sub for message distribution in a WebSocket Cluster.
• Use NGINX or a Load Balancer to distribute connections.
• Improve state synchronization and add Sticky Sessions for enhanced performance.