Go-Zero项目开发17: IM私聊功能实现与消息存储设计

纲要

引言
私聊业务流程分析
- 消息发送与转发
- 核心处理步骤
消息模式：推与拉
- Push 模式
- Pull 模式
消息存储策略：读扩散与写扩散
- 读扩散
- 写扩散
- 本项目的选择：读扩散 + ChatLog + 用户会话记录
私聊功能实现
- 数据模型设计
  - ChatLog 消息记录
  - 消息类型枚举
- 会话 ID 生成
- MongoDB 连接与配置
- 消息结构体定义
  - 通用消息体 Message
  - 私聊业务数据 ChatData
  - 推送数据 PushData
- 错误消息构造
- 业务逻辑 ChatLogic
  - 处理私聊消息
  - 记录消息到 MongoDB
  - 实时推送给接收方
- 路由注册与 Conn 扩展
- 防止重复登录与连接管理
测试验证
总结

引言

在前几篇文章中，我们基于 go-zero 框架成功搭建了 WebSocket 服务，实现了 JWT 鉴权和心跳检测。现在，IM 服务的基础设施已经准备就绪，接下来就要进入即时通讯的核心业务------私聊。本文将首先分析私聊的实现流程，介绍消息传递中的"推拉"模式，以及消息存储中的"读扩散"与"写扩散"策略。随后，我们将结合 MongoDB，在已有的 WebSocket 服务基础上，完整实现用户与好友之间的私聊功能，包括消息持久化与实时推送。

私聊业务流程分析

私聊是最基本的一对一通信场景。其核心流程如下图所示：
接收方客户端 MongoDB WebSocket 服务发送方客户端接收方客户端 MongoDB WebSocket 服务发送方客户端 #mermaid-svg-VwQuJs63XKfwZv7t{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-VwQuJs63XKfwZv7t .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-VwQuJs63XKfwZv7t .error-icon{fill:#552222;}#mermaid-svg-VwQuJs63XKfwZv7t .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-VwQuJs63XKfwZv7t .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-VwQuJs63XKfwZv7t .marker{fill:#333333;stroke:#333333;}#mermaid-svg-VwQuJs63XKfwZv7t .marker.cross{stroke:#333333;}#mermaid-svg-VwQuJs63XKfwZv7t svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-VwQuJs63XKfwZv7t p{margin:0;}#mermaid-svg-VwQuJs63XKfwZv7t .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-VwQuJs63XKfwZv7t text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-VwQuJs63XKfwZv7t .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-VwQuJs63XKfwZv7t .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-VwQuJs63XKfwZv7t .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-VwQuJs63XKfwZv7t .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-VwQuJs63XKfwZv7t #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-VwQuJs63XKfwZv7t .sequenceNumber{fill:white;}#mermaid-svg-VwQuJs63XKfwZv7t #sequencenumber{fill:#333;}#mermaid-svg-VwQuJs63XKfwZv7t #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-VwQuJs63XKfwZv7t .messageText{fill:#333;stroke:none;}#mermaid-svg-VwQuJs63XKfwZv7t .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-VwQuJs63XKfwZv7t .labelText,#mermaid-svg-VwQuJs63XKfwZv7t .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-VwQuJs63XKfwZv7t .loopText,#mermaid-svg-VwQuJs63XKfwZv7t .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-VwQuJs63XKfwZv7t .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-VwQuJs63XKfwZv7t .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-VwQuJs63XKfwZv7t .noteText,#mermaid-svg-VwQuJs63XKfwZv7t .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-VwQuJs63XKfwZv7t .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-VwQuJs63XKfwZv7t .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-VwQuJs63XKfwZv7t .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-VwQuJs63XKfwZv7t .actorPopupMenu{position:absolute;}#mermaid-svg-VwQuJs63XKfwZv7t .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-VwQuJs63XKfwZv7t .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-VwQuJs63XKfwZv7t .actor-man circle,#mermaid-svg-VwQuJs63XKfwZv7t line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-VwQuJs63XKfwZv7t :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 发送私聊消息 (含接收者ID、内容等)从 JWT Token 获取发送者ID将消息写入对应会话若 B 在线，实时推送消息

服务端在处理私聊时，需要完成以下四件事：

解析消息类型：区分文本、图片、语音等不同类型，为后续扩展预留空间。
确定会话标识：根据发送者和接收者生成唯一的会话 ID，用于消息归集。
持久化消息：将聊天记录存入数据库，以便历史消息查询和离线消息拉取。
转发消息：从连接池中查找接收者的连接，并将消息实时推送给对方。

消息模式：推与拉

在即时通讯中，消息的投递方式分为两种：

Push（推） ：服务端主动将消息发送给客户端。接收方在线时，新消息由 WebSocket 立即推送，客户端被动接收，延时极低。
Pull（拉）：客户端主动向服务端请求数据。例如用户上线后拉取离线消息，或向上滑动加载更早的历史记录。

我们的系统采用 实时推 + 离线拉 的混合模式：正常聊天通过 Push 实时送达；当用户离线时，消息暂存在数据库，待其上线后通过 API 主动拉取。

消息存储策略：读扩散与写扩散

消息的存储方式直接影响系统性能。常见的两种模式如下：

策略	写入次数	读取方式	优点	缺点
读扩散	每条消息写一次（存于公共会话）	所有参与者都读同一份	写入简单，节省存储	读取压力大（尤其群聊），已读/未读状态难维护
写扩散	每条消息为每个接收者复制一份	每个用户只读自己的"收件箱"	读取快，个性化状态易实现	写入放大严重（大群尤为明显）

本项目的选择：读扩散 + `ChatLog` + 用户会话记录

在私聊场景中，消息量相对可控，我们采用 读扩散 作为基础方案。具体做法：

由发送方和接收方的 ID 经排序后拼接生成一个唯一的 ChatId，作为该对用户之间所有消息的标识。
所有聊天记录存入 chat_log 集合，每条文档都包含 chat_id 字段。
单独维护一个 用户会话列表 集合，记录每位用户在每个会话中的最后读取位置，从而支持未读计数和离线消息的精确拉取。

这种方案写操作仅一次，读操作通过 chat_id 快速过滤，性能满足私聊需求，同时可方便地扩展至群聊（需要额外的用户-会话状态维护）。

私聊功能实现

数据模型设计

`ChatLog` 消息记录

定义 ChatLog 结构体，映射到 MongoDB 的 chat_log 集合：

go 复制代码

// internal/model/chat_log.go 
package model 
 
import (
    "time"
    "go.mongodb.org/mongo-driver/bson/primitive"
)
 
// MessageType 消息类型 
type MessageType int 
 
const (
    TextMessage  MessageType = 1 
    ImageMessage MessageType = 2 
    VoiceMessage MessageType = 3 
)
 
// ChatLog 聊天记录 
type ChatLog struct {
    ID       primitive.ObjectID `bson:"_id,omitempty" json:"id"`
    ChatId   string             `bson:"chat_id" json:"chat_id"`   // 会话ID 
    FromId   string             `bson:"from_id" json:"from_id"`   // 发送者ID 
    ToId     string             `bson:"to_id" json:"to_id"`       // 接收者ID 
    Type     MessageType        `bson:"type" json:"type"`         // 消息类型 
    Content  string             `bson:"content" json:"content"`   // 消息内容 
    SendTime time.Time          `bson:"send_time" json:"send_time"` // 发送时间 
}

`ChatLogModel` 数据库操作

封装 Insert 方法用于写入消息：

go 复制代码

// internal/model/chat_log_model.go 
package model 
 
import (
    "context"
    "time"
    "go.mongodb.org/mongo-driver/mongo"
)
 
type ChatLogModel struct {
    collection *mongo.Collection 
}
 
func NewChatLogModel(db *mongo.Database, collectionName string) *ChatLogModel {
    return &ChatLogModel{
        collection: db.Collection(collectionName),
    }
}
 
func (m *ChatLogModel) Insert(ctx context.Context, log *ChatLog) error {
    log.SendTime = time.Now()
    _, err := m.collection.InsertOne(ctx, log)
    return err 
}

会话 ID 生成

私聊会话的 ChatId 由双方的 ID 按字典序拼接而成，保证无论哪一方发起，都使用同一个会话 ID。

go 复制代码

// internal/logic/chat_utils.go 
package logic 
 
func GenerateChatId(uid1, uid2 string) string {
    if uid1 < uid2 {
        return uid1 + "_" + uid2 
    }
    return uid2 + "_" + uid1 
}

MongoDB 连接与配置

在 internal/config/config.go 中添加 MongoDB 配置：

go 复制代码

type Config struct {
    rest.RestConf 
    JwtAuth     JwtAuthConf 
    MaxIdleTime time.Duration 
    MongoDB     MongoDBConf 
}
 
type MongoDBConf struct {
    URI      string 
    Database string 
}

etc/im-ws.yaml 配置示例：

yaml 复制代码

Name: im-ws 
Host: 0.0.0.0 
Port: 8888 
MaxIdleTime: 30s 
JwtAuth:
  Secret: "your-secret-key"
MongoDB:
  URI: "mongodb://root:password@192.168.1.100:27017"
  Database: "im_chat"

在 ServiceContext 中初始化 MongoDB 客户端并创建 ChatLogModel：

go 复制代码

// internal/svc/service_context.go 
package svc 
 
import (
    "context"
    "time"
    "strconv"
    "im-ws/internal/config"
    "im-ws/internal/logic"
    "im-ws/internal/model"
    "github.com/zeromicro/go-zero/core/logx"
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
)
 
type ServiceContext struct {
    Config       config.Config 
    WsServer     *logic.Server 
    MongoClient  *mongo.Client 
    ChatLogModel *model.ChatLogModel 
}
 
func NewServiceContext(c config.Config, auth logic.Auth) *ServiceContext {
    clientOptions := options.Client().ApplyURI(c.MongoDB.URI)
    client, err := mongo.Connect(context.Background(), clientOptions)
    if err != nil {
        logx.Must(err)
    }
    if err = client.Ping(context.Background(), nil); err != nil {
        logx.Must(err)
    }
 
    db := client.Database(c.MongoDB.Database)
    chatLogModel := model.NewChatLogModel(db, "chat_log")
 
    server := logic.NewServer(c.Host+":"+strconv.Itoa(c.Port),
        logic.WithAuth(auth),
        logic.WithMaxIdleTime(c.MaxIdleTime),
    )
 
    return &ServiceContext{
        Config:       c,
        WsServer:     server,
        MongoClient:  client,
        ChatLogModel: chatLogModel,
    }
}

消息结构体定义

我们沿用并扩展之前的 Message 结构，并新定义了 ChatData 和 PushData。

go 复制代码

// internal/types/message.go 
package types 
 
// Message 通用消息结构（客户端与服务端通信）
type Message struct {
    Method string      `json:"method"`
    FromID string      `json:"from_id"`
    Data   interface{} `json:"data"` // 业务数据，如 ChatData 
}
 
// ChatData 私聊/群聊消息的具体业务数据 
type ChatData struct {
    ChatType int    `json:"chat_type"` // 1:私聊, 2:群聊 
    ToId     string `json:"to_id"`
    Content  string `json:"content"`
    MsgType  int    `json:"msg_type"`   // 1:text, 2:image, 3:voice 
}
 
// PushData 推送给接收方的消息格式 
type PushData struct {
    FromId   string `json:"from_id"`
    ToId     string `json:"to_id"`
    Content  string `json:"content"`
    MsgType  int    `json:"msg_type"`
    SendTime int64  `json:"send_time"`
}

错误消息构造

为方便向客户端返回统一格式的错误，我们在 Server 上增加 SendError 方法：

go 复制代码

// internal/logic/server.go 补充 
func (s *Server) SendError(conn *Conn, method, errMsg string) {
    errResp := types.Message{
        Method: method + "_err",
        FromID: "server",
        Data:   errMsg,
    }
    data, _ := json.Marshal(errResp)
    conn.WriteMessage(websocket.TextMessage, data)
}

业务逻辑 `ChatLogic`

ChatLogic 负责处理聊天消息的存储与转发。

go 复制代码

// internal/logic/chat_logic.go 
package logic 
 
import (
    "context"
    "encoding/json"
    "time"
 
    "im-ws/internal/model"
    "im-ws/internal/types"
 
    "github.com/gorilla/websocket"
    "github.com/zeromicro/go-zero/core/logx"
)
 
type ChatLogic struct {
    server       *Server 
    chatLogModel *model.ChatLogModel 
}
 
func NewChatLogic(s *Server, chatLogModel *model.ChatLogModel) *ChatLogic {
    return &ChatLogic{
        server:       s,
        chatLogModel: chatLogModel,
    }
}
 
// HandleChat 统一处理聊天消息（私聊或群聊）
func (l *ChatLogic) HandleChat(conn *Conn, msg *types.Message) {
    // 将 interface{} 转为具体结构 
    dataBytes, err := json.Marshal(msg.Data)
    if err != nil {
        l.server.SendError(conn, msg.Method, "invalid data format")
        return 
    }
    var chatData types.ChatData 
    if err = json.Unmarshal(dataBytes, &chatData); err != nil {
        l.server.SendError(conn, msg.Method, "invalid chat data")
        return 
    }
 
    // 获取发送者ID，已在鉴权时通过 context 传递给连接 
    fromID := conn.UserID()
    toID := chatData.ToId 
    if fromID == "" || toID == "" {
        l.server.SendError(conn, msg.Method, "missing user id")
        return 
    }
 
    // 根据聊天类型分别处理 
    if chatData.ChatType == 1 {
        l.handlePrivateChat(fromID, toID, chatData, conn)
    } else if chatData.ChatType == 2 {
        // 群聊逻辑后续补充 
        l.server.SendError(conn, msg.Method, "group chat not implemented yet")
    } else {
        l.server.SendError(conn, msg.Method, "unsupported chat type")
    }
}
 
func (l *ChatLogic) handlePrivateChat(fromID, toID string, chatData types.ChatData, conn *Conn) {
    // 生成会话ID 
    chatID := GenerateChatId(fromID, toID)
 
    // 构建 ChatLog 并写入 MongoDB 
    chatLog := &model.ChatLog{
        ChatId:  chatID,
        FromId:  fromID,
        ToId:    toID,
        Type:    model.MessageType(chatData.MsgType),
        Content: chatData.Content,
    }
    if err := l.chatLogModel.Insert(context.Background(), chatLog); err != nil {
        logx.Errorf("Failed to insert chat log: %v", err)
        l.server.SendError(conn, "chat", "message store failed")
        return 
    }
 
    // 构建推送给接收方的消息 
    pushData := types.PushData{
        FromId:   fromID,
        ToId:     toID,
        Content:  chatData.Content,
        MsgType:  chatData.MsgType,
        SendTime: time.Now().Unix(),
    }
    pushMsg := types.Message{
        Method: "push_message",
        FromID: fromID,
        Data:   pushData,
    }
 
    // 实时推送：查找接收方连接并发送 
    if targetConn, ok := l.server.connPool.GetConnByUserID(toID); ok {
        payload, _ := json.Marshal(pushMsg)
        if err := targetConn.WriteMessage(websocket.TextMessage, payload); err != nil {
            logx.Errorf("Push to user %s failed: %v", toID, err)
        }
    } else {
        // 用户不在线，消息已存入数据库，待上线后拉取 
        logx.Infof("User %s offline, message stored", toID)
    }
}

路由注册与 `Conn` 扩展

为方便业务逻辑中获取当前连接的用户 ID，我们为 Conn 结构体增加 userID 字段，并在鉴权成功后进行设置。

首先，在 Conn 结构体中添加字段和方法：

go 复制代码

// internal/logic/connection.go 补充 
type Conn struct {
    *websocket.Conn 
    server   *Server 
    lastRead time.Time 
    maxIdle  time.Duration 
    closeCh  chan struct{}
    mu       sync.Mutex 
    once     sync.Once 
    userID   string   // 新增：绑定的用户ID 
}
 
// SetUserID 设置用户ID 
func (c *Conn) SetUserID(uid string) {
    c.userID = uid 
}
 
// UserID 获取用户ID 
func (c *Conn) UserID() string {
    return c.userID 
}

然后，在 ServeWS 方法中，鉴权通过后立即设置用户 ID。下面是 ServeWS 的关键片段：

go 复制代码

// 鉴权成功后，获取用户ID并设置到连接上 
userID := s.auth.UserID(r)
conn.SetUserID(userID)
 
// 防止重复登录：踢掉旧连接 
if oldConn, ok := s.connPool.GetConnByUserID(userID); ok {
    oldConn.Close()
}
s.connPool.Add(userID, conn)
defer s.connPool.Remove(userID, conn)

这样，后续的所有业务处理方法（如 ChatLogic.HandleChat）都可以通过 conn.UserID() 安全地获取当前操作的用户身份，无需再从请求中重复解析。

相应的路由注册保持不变，只需在初始化时创建 ChatLogic 并绑定路由：

go 复制代码

// internal/handler/routes.go 
package handler 
 
import (
    "im-ws/internal/logic"
    "im-ws/internal/svc"
    "im-ws/internal/types"
)
 
func RegisterRoutes(ctx *svc.ServiceContext) {
    server := ctx.WsServer 
    chatLogic := logic.NewChatLogic(server, ctx.ChatLogModel)
 
    server.AddRoutes([]logic.Route{
        {
            Method: "chat",
            Handler: func(s *logic.Server, conn *logic.Conn, msg *types.Message) {
                chatLogic.HandleChat(conn, msg)
            },
        },
        // ... 其他路由 
    })
}

注意：路由处理函数的第二个参数类型已由 *websocket.Conn 改为 *logic.Conn，这与我们前文对 Server 路由的改造保持一致，确保类型统一。

防止重复登录与连接管理

在 ServeWS 中我们已经加入了踢掉旧连接的逻辑。当同一个用户再次连接时，旧的连接会被关闭并从连接池中移除，然后添加新连接。这一机制保证了每个用户同时只保留一个有效连接，避免消息重复投递或投递到已失效的连接。

测试验证

分别用两个用户登录，获取 JWT Token。
使用 websocat 或 ApiPost 连接到 WebSocket 服务，并携带 Token。

用户 A 发送私聊消息：

json 复制代码

{
  "method": "chat",
  "from_id": "userA",
  "data": {
    "chat_type": 1,
    "to_id": "userB",
    "content": "Hello, B!",
    "msg_type": 1 
  }
}

用户 B 实时收到推送消息：

json 复制代码

{
  "method": "push_message",
  "from_id": "userA",
  "data": {
    "from_id": "userA",
    "to_id": "userB",
    "content": "Hello, B!",
    "msg_type": 1,
    "send_time": 1698765432 
  }
}

断开 B 的连接，A 再次发送消息，服务端日志提示 User userB offline, message stored，消息已成功写入 MongoDB 的 chat_log 集合。B 重新上线后可通过 API 拉取离线消息。

总结

本文从理论到实践，完整阐述了私聊功能的实现过程。我们首先分析了私聊的业务流程，回顾了实时通讯中"推"与"拉"两种模式，进而对比了"读扩散"与"写扩散"两种存储方案，并选择了适合当前项目规模的读扩散策略。

在代码实现层面，我们设计了 ChatLog 数据模型，集成了 MongoDB，定义了清晰的消息结构体，最终在 WebSocket 服务中通过 ChatLogic 实现了消息的存储和实时转发。同时，通过对 Conn 的扩展，使连接对象携带用户身份，优化了业务逻辑的调用方式。

Go-Zero项目开发17: IM私聊功能实现与消息存储设计

纲要

引言

私聊业务流程分析

消息模式：推与拉

消息存储策略：读扩散与写扩散

本项目的选择：读扩散 + ChatLog + 用户会话记录

私聊功能实现

数据模型设计

ChatLog 消息记录

ChatLogModel 数据库操作

会话 ID 生成

MongoDB 连接与配置

消息结构体定义

错误消息构造

业务逻辑 ChatLogic

路由注册与 Conn 扩展

防止重复登录与连接管理

测试验证

总结

本项目的选择：读扩散 + `ChatLog` + 用户会话记录

`ChatLog` 消息记录

`ChatLogModel` 数据库操作

业务逻辑 `ChatLogic`

路由注册与 `Conn` 扩展