go + websocket實現(xiàn)即時通信之私聊/群聊功能示例

Go  /  管理員 發(fā)布于 10個月前   511

公司項目中后面需要開發(fā)一個即時通訊的聊天功能,

所以我先自己寫一個簡單的demo ,以實現(xiàn)效果為主, 后面在根據業(yè)務場景調整之.

還是那套環(huán)境

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 (
    "zongscan-go/model"
)

type ChatSlService struct {
    Property
}

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

// 創(chuàng)建
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 "zongscan-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"

	"zongscan-go/model"
	"zongscan-go/model/resp"
	"zongscan-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() // 阻塞等待讀取客戶端發(fā)送的消息
		fmt.Println("messageType", messageType, "p", string(p), "clients", clients)

		if err != nil {
			break // 無法讀取信息，退出循環(huán)
		}
		// 解析消息，識別目標客戶端或群組
		// 根據消息內容選擇性發(fā)送給特定客戶端或群組
		// 例如，如果消息格式為私聊消息，則根據消息中指定的目標用戶ID發(fā)送給特定客戶端
		// 如果消息格式為群聊消息，則將消息發(fā)送給所有群組成員
		//recipientIDs := []string{"1", "2"} // 限定發(fā)送人群

		// 解析消息，這里假設消息是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 {
			//群聊
			// 創(chuàng)建一個空的整數切片來存儲鍵  用戶群
			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 {
				// 向當前客戶端發(fā)送消息
				err := client.WriteMessage(messageType, p)
				if err != nil {
					// 發(fā)送消息失敗，關閉連接并刪除客戶端
					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:

入庫數據:

簡單分享一下, 有興趣可自行測試.