package sse import ( "encoding/json" "strconv" "sync" "sync/atomic" "time" ) const ( defaultUserBufferSize = 128 maxReplayEvents = 200 ) type Event struct { ID uint64 `json:"event_id"` Event string `json:"event"` TS int64 `json:"ts"` Payload interface{} `json:"payload"` } type subscriber struct { id uint64 ch chan Event quit chan struct{} } type Hub struct { seq uint64 mu sync.RWMutex subscribers map[string]map[uint64]*subscriber history map[string][]Event } func NewHub() *Hub { return &Hub{ subscribers: make(map[string]map[uint64]*subscriber), history: make(map[string][]Event), } } func (h *Hub) NextID() uint64 { return atomic.AddUint64(&h.seq, 1) } func ParseEventID(raw string) uint64 { if raw == "" { return 0 } id, err := strconv.ParseUint(raw, 10, 64) if err != nil { return 0 } return id } func (h *Hub) Subscribe(userID string, afterID uint64) (chan Event, func(), []Event) { subID := h.NextID() sub := &subscriber{ id: subID, ch: make(chan Event, defaultUserBufferSize), quit: make(chan struct{}), } h.mu.Lock() if _, ok := h.subscribers[userID]; !ok { h.subscribers[userID] = make(map[uint64]*subscriber) } h.subscribers[userID][subID] = sub replay := make([]Event, 0) for _, e := range h.history[userID] { if e.ID > afterID { replay = append(replay, e) } } h.mu.Unlock() cancel := func() { h.mu.Lock() defer h.mu.Unlock() if userSubs, ok := h.subscribers[userID]; ok { if s, exists := userSubs[subID]; exists { close(s.quit) delete(userSubs, subID) close(s.ch) } if len(userSubs) == 0 { delete(h.subscribers, userID) } } } return sub.ch, cancel, replay } func (h *Hub) HasSubscribers(userID string) bool { h.mu.RLock() defer h.mu.RUnlock() return len(h.subscribers[userID]) > 0 } func (h *Hub) PublishToUser(userID string, eventName string, payload interface{}) Event { ev := Event{ ID: h.NextID(), Event: eventName, TS: time.Now().UnixMilli(), Payload: payload, } h.publish(userID, ev) return ev } func (h *Hub) PublishToUsers(userIDs []string, eventName string, payload interface{}) { for _, uid := range userIDs { h.PublishToUser(uid, eventName, payload) } } func (h *Hub) publish(userID string, ev Event) { h.mu.Lock() history := append(h.history[userID], ev) if len(history) > maxReplayEvents { history = history[len(history)-maxReplayEvents:] } h.history[userID] = history targets := make([]*subscriber, 0, len(h.subscribers[userID])) for _, s := range h.subscribers[userID] { targets = append(targets, s) } h.mu.Unlock() for _, s := range targets { select { case <-s.quit: case s.ch <- ev: default: // 慢消费者丢弃单条消息,客户端可通过 Last-Event-ID + HTTP 同步补偿 } } } func EncodeData(ev Event) (string, error) { body, err := json.Marshal(ev.Payload) if err != nil { return "", err } return string(body), nil }