58 lines
1.1 KiB
Go
58 lines
1.1 KiB
Go
package wshub
|
|
|
|
import "sync"
|
|
|
|
type Event struct {
|
|
Type string `json:"type"`
|
|
TaskID int64 `json:"task_id"`
|
|
Data any `json:"data,omitempty"`
|
|
}
|
|
|
|
type Hub struct {
|
|
mu sync.Mutex
|
|
nextID int64
|
|
subs map[int64]map[int64]chan Event // taskID -> subID -> ch
|
|
}
|
|
|
|
func New() *Hub {
|
|
return &Hub{subs: make(map[int64]map[int64]chan Event)}
|
|
}
|
|
|
|
func (h *Hub) Subscribe(taskID int64) (int64, <-chan Event) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.nextID++
|
|
id := h.nextID
|
|
ch := make(chan Event, 64)
|
|
if h.subs[taskID] == nil {
|
|
h.subs[taskID] = make(map[int64]chan Event)
|
|
}
|
|
h.subs[taskID][id] = ch
|
|
return id, ch
|
|
}
|
|
|
|
func (h *Hub) Unsubscribe(taskID, id int64) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if m := h.subs[taskID]; m != nil {
|
|
if ch, ok := m[id]; ok {
|
|
close(ch)
|
|
delete(m, id)
|
|
}
|
|
if len(m) == 0 {
|
|
delete(h.subs, taskID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Publish(ev Event) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
for _, ch := range h.subs[ev.TaskID] {
|
|
select {
|
|
case ch <- ev:
|
|
default: // медленный подписчик — событие дропаем, не блокируем воркер
|
|
}
|
|
}
|
|
}
|