go + websocket实现即时通信之私聊/群聊功能示例

Go  /  管理员 发布于 4个月前   268

公司项目中后面需要开发一个即时通讯的聊天功能,

所以我先自己写一个简单的demo ,以实现效果为主, 后面在根据业务场景调整之.

还是那套环境

go + gin + websocket

废话不多少直接上代码:

数据库设计一个简单的表:

CREATE TABLE `chat_sls` (
  `id` bigint unsigned NOT NULL AUTO_INCREMENT,
  `user_id` varchar(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NOT NULL COMMENT 'from 用户',
  `to` varchar(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT 'to  用户',
  `type` tinyint DEFAULT NULL,
  `text` text CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci COMMENT 'msg',
  `created_at` datetime DEFAULT NULL,
  `updated_at` datetime DEFAULT NULL,
  `deleted_at` datetime DEFAULT NULL,
  PRIMARY KEY (`id`)
) ENGINE=InnoDB AUTO_INCREMENT=15 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='chat';

model:

package model

import (
    "time"
    "gorm.io/gorm"
)

var (
    ChatSlAll = []string{"id", "user_id", "to", "type", "text", "created_at"}
)
type ChatSl struct {
    Model
    UserId    string `json:"userId,omitempty" gorm:"comment:from用户id"`
    To        string `json:"to,omitempty" gorm:"comment:to用户id"`
    Type      int8   `json:"type" gorm:"comment:聊天类型 2私聊"`
    Text      string `json:"text,omitempty" gorm:"comment:msg;default:''`
    CreatedAt time.Time
    UpdatedAt time.Time      `json:"updateAt,omitempty" gorm:"autoUpdateTime:false"`
    DeletedAt gorm.DeletedAt `gorm:"index"`
}

service:

package service

import (
    "sokogate-go/model"
)

type ChatSlService struct {
    Property
}

func NewChatSlService() *ChatSlService {
    return &ChatSlService{
        Property: Property{
            DB: SqlDB(),
        },
    }
}

// 创建
func (s *ChatSlService) Create(m *model.ChatSl) (err error) {
    err = s.DB.Create(m).Error
    return
}

// 查询私聊记录
func (s *ChatSlService) GetSlList(userId string, to string) (list []model.ChatSl, err error) {
    err = s.DB.Select(model.ChatSlAll).Where("user_id = ? and `to` = ?", userId, to).Order("created_at DESC").Limit(10).Find(&list).Error
    return
}

路由:

package router

import (
    v1 "sokogate-go/api/v1"
    "github.com/gin-gonic/gin"
)

func InitChatRouter(r *gin.RouterGroup) {
    g := r.Group("chat")
    {
        g.GET("addsl", v1.WebSocketHandler)
        g.GET("getsllist", v1.GetChatSllist)
    }
}

api控制器:

package v1

import (
	"encoding/json"
	"fmt"

	"sokogate-go/model"
	"sokogate-go/model/resp"
	"sokogate-go/service"

	//"strconv"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// websocket配置
var clients = make(map[string]*websocket.Conn)
var clientMutex sync.Mutex

func WebSocketHandler(c *gin.Context) {
	userID := c.Query("userID")
	fmt.Println("用户id", userID)
	ws, err := websocket.Upgrade(c.Writer, c.Request, nil, 1024, 1024)
	if err != nil {
		panic(err)
	}

	// 加入新客户端
	clientMutex.Lock()
	clients[userID] = ws // 将新客户端添加到客户端集合中
	clientMutex.Unlock()
	// 处理WebSocket消息
	go handleWebSocketMessages(userID, ws) // 开启一个 goroutine 处理该客户端的消息
}

func handleWebSocketMessages(userID string, ws *websocket.Conn) {
	for {
		messageType, p, err := ws.ReadMessage() // 阻塞等待读取客户端发送的消息
		fmt.Println("messageType", messageType, "p", string(p), "clients", clients)

		if err != nil {
			break // 无法读取信息，退出循环
		}
		// 解析消息，识别目标客户端或群组
		// 根据消息内容选择性发送给特定客户端或群组
		// 例如，如果消息格式为私聊消息，则根据消息中指定的目标用户ID发送给特定客户端
		// 如果消息格式为群聊消息，则将消息发送给所有群组成员
		//recipientIDs := []string{"1", "2"} // 限定发送人群

		// 解析消息，这里假设消息是JSON格式，包含"to"字段指定接收者ID
		//var msg Message
		msg := model.ChatSl{}
		err = json.Unmarshal(p, &msg)
		if err != nil {
			fmt.Println("Error parsing message:", err)
			continue
		}
		//from id
		msg.UserId = userID
		//fmt.Println("msg:", msg)

		// 检查目标用户ID是否合法
		if msg.To == "" || msg.To == userID {
			fmt.Println("Invalid recipient ID or self-messaging detected")
			continue
		}

		var recipientIDs []string // 在这里声明变量，但不初始化
		//私聊 type= 2
		if msg.Type == 2 {
			recipientIDs = []string{userID, msg.To} // 直接赋值，无需再次声明

			//入库
			err = service.NewChatSlService().Create(&msg)
			if err != nil {
				fmt.Println("type=1 add error")
				continue
			}
		} else {
			//群聊
			// 创建一个空的整数切片来存储键  用户群
			recipientIDs = make([]string, 0, len(clients)) // 初始化切片，并分配初始容量

			// 遍历map，并将键添加到切片中
			for recipientID := range clients {
				recipientIDs = append(recipientIDs, recipientID)
			}
		}

		fmt.Println("recipientIDs", recipientIDs)
		broadcastMessage(recipientIDs, messageType, p) // 广播消息给所有客户端
	}

	clientMutex.Lock()
	defer clientMutex.Unlock()

	delete(clients, userID) // 删除客户端连接
	ws.Close()              // 关闭 WebSocket 连接
}

func broadcastMessage(recipientIDs []string, messageType int, p []byte) {
	clientMutex.Lock()
	defer clientMutex.Unlock()

	// 遍历所有客户端连接
	for id, client := range clients {
		// 检查当前客户端是否在接收者列表中  -- 私聊功能
		for _, recipientID := range recipientIDs {
			if id == recipientID {
				// 向当前客户端发送消息
				err := client.WriteMessage(messageType, p)
				if err != nil {
					// 发送消息失败，关闭连接并删除客户端
					client.Close()
					delete(clients, id)
				}
				break
			}
		}
	}
}

//历史记录
func GetChatSllist(c *gin.Context) {
	var param model.ChatSl
	err := c.ShouldBindJSON(&param)
	if err != nil {
		fmt.Println("GetChatSllist error")
		return
	}

	list, err := service.NewChatSlService().GetSlList(param.UserId, param.To)
	if err != nil {
		resp.ERRWithMsg(c, resp.USER_FIND_ERR, err.Error())
		return
	}

	resp.OkWithData(list, c)
}

看看效果

用户1:

用户2:

入库数据:

简单分享一下, 有兴趣可自行测试.