dockerfile修改

internal/ws/connection.go

@@ -3,6 +3,7 @@ package ws
 import (
 	"ChatRoom/internal/models"
 	"ChatRoom/internal/rabbitmq"
+	"ChatRoom/internal/redis"
 	"encoding/json"
 	"log"
 	"net/http"
@@ -30,14 +31,15 @@ var upgrader = websocket.Upgrader{
 
 // 连接
 type Connection struct {
-	wsConn    *websocket.Conn  //websocket连接
-	rmqClient *rabbitmq.Client //rabbitmq客户端
-	queueName string           //队列名称
-	userID    string           //用户ID
-	send      chan []byte      //发送通道
+	wsConn      *websocket.Conn  //websocket连接
+	rmqClient   *rabbitmq.Client //rabbitmq客户端
+	redisClient *redis.Client    //redis客户端
+	queueName   string           //队列名称
+	userID      string           //用户ID
+	send        chan []byte      //发送通道
 }
 
-func NewConnection(w http.ResponseWriter, r *http.Request, rmq *rabbitmq.Client) {
+func NewConnection(w http.ResponseWriter, r *http.Request, rmq *rabbitmq.Client, redisClient *redis.Client) {
 	// 1. 升级 HTTP 到 WebSocket
 	wsConn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -51,6 +53,13 @@ func NewConnection(w http.ResponseWriter, r *http.Request, rmq *rabbitmq.Client)
 		wsConn.Close()
 		return
 	}
+	//redis记录用户
+	if err := redisClient.AddUserToZSet(redis.UserZSet, userID, time.Now().Unix()); err != nil {
+		log.Printf("Redis 用户记录失败: %v", err)
+		wsConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(5000, "redis error"))
+		wsConn.Close()
+		return
+	}
 	// 3. 为用户创建 RabbitMQ 队列
 	queueName, err := rmq.DeclareQueue()
 	if err != nil {
@@ -75,11 +84,12 @@ func NewConnection(w http.ResponseWriter, r *http.Request, rmq *rabbitmq.Client)
 
 	// 5. 创建 Connection 对象
 	conn := &Connection{
-		wsConn:    wsConn,
-		rmqClient: rmq,
-		queueName: queueName,
-		userID:    userID,
-		send:      make(chan []byte, 256),
+		wsConn:      wsConn,
+		rmqClient:   rmq,
+		redisClient: redisClient,
+		queueName:   queueName,
+		userID:      userID,
+		send:        make(chan []byte, 256),
 	}
 
 	// 发送登录成功消息
@@ -93,13 +103,20 @@ func NewConnection(w http.ResponseWriter, r *http.Request, rmq *rabbitmq.Client)
 		conn.send <- data
 	}
 
+	// 广播在线人数
+	conn.broadcastUserCount()
+
 	go conn.writePump()
 	go conn.readPump()
 	go conn.consumeFromRabbitMQ()
 }
 
 func (c *Connection) readPump() {
-	defer c.wsConn.Close()
+	defer func() {
+		c.redisClient.RemoveUserFromZSet(redis.UserZSet, c.userID)
+		c.broadcastUserCount()
+		c.wsConn.Close()
+	}()
 
 	c.wsConn.SetReadLimit(maxMessageSize)
 	c.wsConn.SetReadDeadline(time.Now().Add(pongWait))
@@ -231,3 +248,24 @@ func (c *Connection) consumeFromRabbitMQ() {
 		}
 	}
 }
+
+// 广播在线人数
+func (c *Connection) broadcastUserCount() {
+	count, err := c.redisClient.CountUsers(redis.UserZSet)
+	if err != nil {
+		log.Printf("获取在线人数失败: %v", err)
+		return
+	}
+
+	msg := &models.Message{
+		Type:  models.MsgTypeUserCount,
+		User:  "system",
+		Count: count,
+		Time:  time.Now().UTC().Format(time.RFC3339),
+	}
+
+	body, _ := json.Marshal(msg)
+	if err := c.rmqClient.Publish(c.rmqClient.ExchangeName, "chat.global", body); err != nil {
+		log.Printf("广播在线人数失败: %v", err)
+	}
+}