WebSocket 实时通信（二）

WebSocket 即时消息推送系统

1. 项目概述

1.1 项目背景

在现代 Web 应用中，实时通信功能越来越重要，例如在线聊天、实时通知、股票行情更新等。本项目基于 WebSocket 技术，构建一个高效、稳定的即时消息推送系统。

1.2 技术选型

• 后端：Go（使用 gorilla/websocket 实现 WebSocket 服务器）

• 前端：Python（websockets 库作为客户端示例）

• 通信协议：WebSocket

• 并发管理：sync.Map 维护在线客户端，支持高并发

• 异常处理：超时检测、心跳检测、自动断线处理

---

2. 代码实现

2.1 WebSocket 服务器（Go 语言）

服务器监听 :8808/chat 端口，支持多客户端连接、消息广播和心跳检测。

package main

import (
    "context"
    "github.com/gorilla/websocket"
    "log"
    "net/http"
    "sync"
    "time"
)

// Client 结构体封装 WebSocket 连接，保证写操作的线程安全
type Client struct {
    Conn *websocket.Conn
    mu   sync.Mutex
}

var (
    clients   sync.Map            // 存储所有在线的客户端
    broadcast = make(chan []byte) // 消息广播通道
)

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

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    http.HandleFunc("/chat", handleChat)
    go handleMessages(ctx)

    log.Println("Server started on localhost:8808")
    if err := http.ListenAndServe(":8808", nil); err != nil {
        log.Fatal("ListenAndServe: ", err)
    }
}

// 处理 WebSocket 连接
func handleChat(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("Upgrade error:", err)
        return
    }
    client := &Client{Conn: conn}
    clientID := conn.RemoteAddr().String()
    clients.Store(clientID, client)
    log.Printf("New connection: %s", clientID)

    defer func() {
        clients.Delete(clientID)
        conn.Close()
        log.Printf("Connection closed: %s", clientID)
    }()

    // 设置 Pong 处理，避免超时断连
    conn.SetReadDeadline(time.Now().Add(60 * time.Second))
    conn.SetPongHandler(func(string) error {
        conn.SetReadDeadline(time.Now().Add(60 * time.Second))
        return nil
    })

    // 发送心跳 Ping
    go func() {
        ticker := time.NewTicker(30 * time.Second)
        defer ticker.Stop()
        for {
            select {
            case <-ticker.C:
                client.mu.Lock()
                err := client.Conn.WriteMessage(websocket.PingMessage, nil)
                client.mu.Unlock()
                if err != nil {
                    log.Println("Ping error:", err)
                    conn.Close()
                    clients.Delete(clientID)
                    return
                }
            }
        }
    }()

    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            break
        }
        log.Printf("Received message from %s: %s", clientID, string(message))
        broadcast <- message
    }
}

// 处理消息广播
func handleMessages(ctx context.Context) {
    for {
        select {
        case message := <-broadcast:
            distributeMessage(message)
        case <-ctx.Done():
            return
        }
    }
}

// 分发消息给所有在线客户端
func distributeMessage(message []byte) {
    clients.Range(func(key, value interface{}) bool {
        client := value.(*Client)
        client.mu.Lock()
        defer client.mu.Unlock()

        if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
            log.Println("WriteMessage error:", err)
            client.Conn.Close()
            clients.Delete(key)
        }
        return true
    })
}

2.2 WebSocket 客户端 A（Python）

客户端 A 负责发送消息，并接收 WebSocket 服务器的响应。

import json
import time
import datetime
import logging
import asyncio
import websockets

_logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
async def send_message(uri: str, message: dict, timeout: int = 10) -> dict:
    """
    发送 WebSocket 消息，并在 `timeout` 秒内等待响应
    """
    try:
        async with websockets.connect(uri) as websocket:
            await asyncio.wait_for(websocket.send(json.dumps(message)), timeout)
            response = await asyncio.wait_for(websocket.recv(), timeout)
            return json.loads(response)

    except asyncio.TimeoutError:
        return {"error": "Connection timed out"}
    except websockets.exceptions.ConnectionClosedError as e:
        return {"error": f"Connection closed: {e}"}
    except json.JSONDecodeError as e:
        return {"error": f"JSON decode error: {e}"}
    except Exception as e:
        return {"error": f"Unexpected error: {e}"}

def send_message_sync(message: dict) -> dict:
    """同步调用 WebSocket 消息发送"""
    return asyncio.run(send_message("ws://127.0.0.1:8808/chat", message))

if __name__ == "__main__":
    """
       注意这里的：版本：websockets==15.0
    """

    for i in range(1000):
        test_message = {"message": "Hello WebSocket", "time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
        response = send_message_sync(test_message)
        print(f"Received response: {response}")
        time.sleep(1)

2.3 WebSocket 客户端 B（Go 语言）

客户端 B 只接收服务器的消息。

package main

import (
    "fmt"
    "log"
    "os"
    "os/signal"
    "syscall"
    "time"

    "github.com/gorilla/websocket"
)

func main() {
    // 创建WebSocket连接
    u := "ws://127.0.0.1:8808/chat"
    log.Printf("连接到 %s", u)
    c, _, err := websocket.DefaultDialer.Dial(u, nil)
    if err != nil {
        log.Fatal("连接错误：", err)
    }
    defer c.Close()

    // 处理中断信号
    interrupt := make(chan os.Signal, 1)
    signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)

    // 启动协程接收消息
    go func() {
        for {
            _, message, err := c.ReadMessage()
            if err != nil {
                log.Println("读取消息错误：", err)
                break
            }
            fmt.Println("WebSocket服务 Get message:", string(message))
        }
    }()

    // 等待中断信号
    <-interrupt
    log.Println("接收到中断信号，关闭连接")
    err = c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
    if err != nil {
        log.Println("关闭连接错误：", err)
    }
    time.Sleep(time.Second)
}

3. 测试数据

3.1、客户端A:产生数据

3.2、服务端:接收数据

3.2、客户端B:消费数据

3. 总结

本项目实现了一个 WebSocket 即时消息推送系统，支持： 多客户端连接
消息广播
心跳检测
异常处理

此 WebSocket 方案可用于 在线聊天、实时推送、协同编辑等场景