如何使用WebSocket实现一个公域聊天室？ --Go

我叫黑大帅
0 阅读4分钟

如何使用WebSocket实现一个公域聊天室？

什么是「公域聊天室」？

所有连接到服务端的用户，都在同一个公共房间里：

  1. 任意用户发消息，所有人都能看到
  2. 用户进入 / 离开，系统会全员通知
  3. 支持多人同时在线，互不阻塞

核心原理

  • WebSocket 长连接：客户端和服务端永久连通，随时收发消息

  • Go goroutine 并发：每个用户独占一个协程，多人连接不阻塞

  • 全局客户端列表 + 广播：服务端存储所有在线用户，一人发消息，全员推送

  • sync.Mutex 互斥锁：多人同时操作客户端列表，保证线程安全

代码具体实现

定义全局变量（存储用户 + 保证并发安全）

var (
	// 存储所有在线的WebSocket连接（核心：全员广播的基础）
	clients = make(map[*websocket.Conn]bool)
	// 存储连接对应的用户名（区分谁发的消息）
	clientNames = make(map[*websocket.Conn]string)
	// 互斥锁：多个协程同时修改map，必须加锁，否则程序崩溃
	clientsMutex sync.Mutex
	// HTTP 升级为 WebSocket 的工具
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true // 允许跨域（浏览器访问必备）
		},
	}
)
  • clients：记录谁在线，用于广播消息
  • clientNames：给连接绑定用户名
  • clientsMutex：多人并发操作共享数据的安全保障
  • upgrader：把普通 HTTP 请求 → 长连接 WebSocket

启动服务，注册路由

func main() {
	// 注册路由：访问 /ws 就触发聊天处理逻辑
	http.HandleFunc("/ws", wsHandler)

	fmt.Println("WebSocket 服务启动：ws://127.0.0.1:8081/ws")
	// 核心：启动HTTP服务，自动为每个客户端开goroutine（不阻塞）
	if err := http.ListenAndServe(":8081", nil); err != nil {
		fmt.Println("服务启动失败：", err)
	}
}

wsHandler（用户连接全流程）

这是聊天室最核心的函数，处理用户从「连接→登录→聊天→断开」的全生命周期：

升级 HTTP 为 WebSocket 连接
// 把普通HTTP请求，升级为长连接WebSocket
conn, err := upgrader.Upgrade(w, r, nil)
// 函数结束自动关闭连接（用户断开时触发）
defer conn.Close()
业务交互：让用户输入用户名
// 服务端主动发消息：欢迎语
conn.WriteMessage(websocket.TextMessage, []byte("欢迎连接！请输入你的名字："))
// 读取用户输入的名字
_, msg, err := conn.ReadMessage()
name := string(msg)
业务交互：注册用户 + 广播进入通知
// 把当前用户加入在线列表
registerClient(conn, name)
// 广播：XXX 进入聊天室（全员可见）
broadcast("系统: " + name + " 进入了聊天室")
用户断开时自动清理
// defer：函数最后执行（用户关闭网页时触发）
defer func() {
	unregisterClient(conn) // 从在线列表删除
	broadcast("系统: " + name + " 离开了聊天室") // 广播离开消息
}()
主要业务：循环读取消息 + 全员广播
// 死循环：持续监听用户发的消息
for {
	_, msg, err := conn.ReadMessage()
	if err != nil {
		break // 出错/断开，退出循环
	}
	// 拼接消息：用户名: 内容
	chatMsg := fmt.Sprintf("%s: %s", name, string(msg))
	// 广播给所有人
	broadcast(chatMsg)
}

线程安全的用户注册 / 注销

因为多个协程同时修改全局 map，必须加锁，否则程序崩溃：

// 注册用户：加锁 → 修改 → 解锁
func registerClient(conn *websocket.Conn, name string) {
	clientsMutex.Lock()   // 加锁
	defer clientsMutex.Unlock() // 自动解锁
	clients[conn] = true
	clientNames[conn] = name
}

// 注销用户：从map删除
func unregisterClient(conn *websocket.Conn) {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()
	delete(clients, conn)
	delete(clientNames, conn)
}

广播函数（一人发消息，全员收到）

func broadcast(message string) {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()

	// 遍历所有在线用户
	for client := range clients {
		// 给每个用户发送消息
		client.WriteMessage(websocket.TextMessage, []byte(message))
	}
}

给你返回全部代码（doge）🐶

package main

import (
    "fmt"
    "net/http"
    "sync"

    "github.com/gorilla/websocket" // go get github.com/gorilla/websocket
)

var (
    // clients 存储所有在线的 WebSocket 连接
    clients = make(map[*websocket.Conn]bool)
    // clientNames 存储连接对应的用户名
    clientNames = make(map[*websocket.Conn]string)
    // clientsMutex 保护 clients 和 clientNames 的并发读写
    clientsMutex sync.Mutex
    // upgrader 用于将 HTTP 请求升级为 WebSocket 请求
    upgrader = websocket.Upgrader{
       CheckOrigin: func(r *http.Request) bool {
          return true // 允许跨域
       },
    }
)

func main() {
    // 注册 WebSocket 路由
    http.HandleFunc("/ws", wsHandler)

    fmt.Println("WebSocket 服务启动：ws://127.0.0.1:8081/ws")
    if err := http.ListenAndServe(":8081", nil); err != nil {
       fmt.Println("服务启动失败：", err)
    }
}

// wsHandler 处理 WebSocket 连接请求
func wsHandler(w http.ResponseWriter, r *http.Request) {
    // 1. 升级 HTTP 连接为 WebSocket 连接
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
       fmt.Println("升级失败：", err)
       return
    }
    // 确保连接最终被关闭
    defer conn.Close()

    // 2. 询问并读取用户名
    if err := conn.WriteMessage(websocket.TextMessage, []byte("欢迎连接！请输入你的名字：")); err != nil {
       return
    }
    _, msg, err := conn.ReadMessage()
    if err != nil {
       return
    }
    name := string(msg)

    // 3. 注册新用户并广播进入消息
    registerClient(conn, name)
    broadcast("系统: " + name + " 进入了聊天室")
    fmt.Printf("用户 [%s] 已加入\n", name)

    // 4. 设置断开连接时的清理操作
    defer func() {
       unregisterClient(conn)
       broadcast("系统: " + name + " 离开了聊天室")
       fmt.Printf("用户 [%s] 已断开\n", name)
    }()

    // 5. 循环读取消息并广播
    for {
       _, msg, err := conn.ReadMessage()
       if err != nil {
          break // 读取错误或连接关闭
       }

       // 构造并广播消息
       chatMsg := fmt.Sprintf("%s: %s", name, string(msg))
       fmt.Println(chatMsg) // 服务端日志
       broadcast(chatMsg)
    }
}

// registerClient 线程安全地注册新连接
func registerClient(conn *websocket.Conn, name string) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    clients[conn] = true
    clientNames[conn] = name
}

// unregisterClient 线程安全地注销连接
func unregisterClient(conn *websocket.Conn) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    delete(clients, conn)
    delete(clientNames, conn)
}

// broadcast 向所有在线客户端广播消息
func broadcast(message string) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    msgBytes := []byte(message)
    for client := range clients {
       // 注意：如果某个客户端网络阻塞，这里可能会阻塞较长时间
       // 实际生产中应将发送逻辑放入独立的 goroutine 或使用 channel
       err := client.WriteMessage(websocket.TextMessage, msgBytes)
       if err != nil {
          fmt.Printf("广播消息失败: %v\n", err)
          client.Close()
          delete(clients, client)
          delete(clientNames, client)
       }
    }
}

相关文章

Go 标准库 net/http 包都能干嘛？Go 标准库 net/http 包是干嘛的？ net/http 是 Go 官 - 掘金

Golang中实时推送的功臣 - WebSocketGolang中实时推送的功臣 - WebSocket WebSock - 掘金

avatar
文章
阅读
粉丝