package api import ( "context" "encoding/json" "log" "sync" "time" "github.com/gorilla/websocket" "github.com/redis/go-redis/v9" ) const pushChannel = "pargar:push" // Hub manages per-user websocket connections and fans out messages. // Messages are published on a shared Redis channel so all instances deliver to their own local clients. type Hub struct { mu sync.RWMutex conns map[int64]map[*Client]struct{} rdb *redis.Client onSeen func(userID int64) } type Client struct { hub *Hub conn *websocket.Conn user int64 send chan []byte once sync.Once } func NewHub(rdb *redis.Client) *Hub { h := &Hub{ conns: map[int64]map[*Client]struct{}{}, rdb: rdb, } go h.subscribe() return h } func (h *Hub) subscribe() { ctx := context.Background() sub := h.rdb.Subscribe(ctx, pushChannel) defer sub.Close() ch := sub.Channel() for msg := range ch { var p struct { User int64 `json:"user"` Data json.RawMessage `json:"data"` } if err := json.Unmarshal([]byte(msg.Payload), &p); err != nil { continue } h.deliver(p.User, p.Data) } } // Push fans a message out to a user across all instances. func (h *Hub) Push(userID int64, msg any) { b, err := json.Marshal(msg) if err != nil { return } payload, _ := json.Marshal(map[string]any{"user": userID, "data": json.RawMessage(b)}) h.rdb.Publish(context.Background(), pushChannel, payload) } func (h *Hub) deliver(userID int64, data []byte) { h.mu.RLock() clients := h.conns[userID] for c := range clients { select { case c.send <- data: default: } } h.mu.RUnlock() } // Online reports whether a user currently has at least one live connection. func (h *Hub) Online(userID int64) bool { h.mu.RLock() defer h.mu.RUnlock() return len(h.conns[userID]) > 0 } // SetSeen records a callback invoked when a user connects, stays connected, or disconnects. func (h *Hub) SetSeen(fn func(userID int64)) { h.mu.Lock() h.onSeen = fn h.mu.Unlock() } func (h *Hub) noteSeen(userID int64) { if userID == 0 { return } h.mu.RLock() fn := h.onSeen h.mu.RUnlock() if fn == nil { return } go fn(userID) } // Register adds a new client connection for a user. func (h *Hub) Register(userID int64, conn *websocket.Conn) *Client { c := &Client{hub: h, conn: conn, user: userID, send: make(chan []byte, 16)} h.mu.Lock() if h.conns[userID] == nil { h.conns[userID] = map[*Client]struct{}{} } h.conns[userID][c] = struct{}{} h.mu.Unlock() h.noteSeen(userID) go c.writePump() go c.readPump() return c } func (h *Hub) Unregister(c *Client) { offline := false h.mu.Lock() if set, ok := h.conns[c.user]; ok { if _, ok := set[c]; ok { delete(set, c) if len(set) == 0 { delete(h.conns, c.user) offline = true } } } h.mu.Unlock() close(c.send) if offline { h.noteSeen(c.user) } } // close unregisters the client and closes its connection exactly once, // even though both readPump and writePump may return independently. func (c *Client) close() { c.once.Do(func() { c.hub.Unregister(c) c.conn.Close() }) } func (c *Client) writePump() { defer func() { if r := recover(); r != nil { log.Printf("ws writePump recovered user=%d: %v", c.user, r) } }() ticker := time.NewTicker(30 * time.Second) defer func() { ticker.Stop() c.close() }() for { select { case msg, ok := <-c.send: if !ok { c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil { return } case <-ticker.C: c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } c.hub.noteSeen(c.user) } } } func (c *Client) readPump() { defer func() { if r := recover(); r != nil { log.Printf("ws readPump recovered user=%d: %v", c.user, r) } }() defer c.close() c.conn.SetReadLimit(4096) c.conn.SetReadDeadline(time.Now().Add(70 * time.Second)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(70 * time.Second)) return nil }) for { _, data, err := c.conn.ReadMessage() if err != nil { // close errors are expected; log genuinely unexpected ones only. if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { log.Printf("ws read err user=%d: %v", c.user, err) } return } var incoming struct { Type string `json:"type"` To int64 `json:"to"` } if json.Unmarshal(data, &incoming) == nil && incoming.Type == "typing" && incoming.To > 0 { c.hub.Push(incoming.To, map[string]any{"type": "typing", "from": c.user}) } } }