Commilitia-Drop/internal/hub/hub.go

package hub

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"commilitia.net/cdrop/internal/db"
)

// DeviceLister is the subset of *db.Queries the hub needs.
// Declared as interface so tests can inject a fake.
type DeviceLister interface {
	ListDevicesByUser(ctx context.Context, userID string) ([]db.Device, error)
}

// Event is what the hub fan-outs to SSE clients.
type Event struct {
	Type string `json:"-"`
	Data any    `json:"-"`
}

// PresenceDevice is the per-device entry in a `presence` event payload.
// Field names align with brief §5: { name, type, online }; last_seen 是
// 设置页展示「上次活跃」需要的扩展字段。
type PresenceDevice struct {
	Name     string `json:"name"`
	Type     string `json:"type"`
	Online   bool   `json:"online"`
	LastSeen int64  `json:"last_seen"`
}

const clientBuffer = 32

// Client is a single live SSE connection.
type Client struct {
	UserID   string
	DeviceID string
	ch       chan Event
}

func (c *Client) Events() <-chan Event { return c.ch }

// Hub is the in-memory presence + signaling fan-out.
//
// Concurrency: a single sync.RWMutex protects the user→device map.
// Sends to client channels are non-blocking — slow consumers drop events
// rather than stall the broadcaster.
type Hub struct {
	mu    sync.RWMutex
	users map[string]map[string]*Client

	devices DeviceLister
	grace   time.Duration

	closed bool
}

func New(devices DeviceLister) *Hub {
	return &Hub{
		users:   map[string]map[string]*Client{},
		devices: devices,
		grace:   30 * time.Second,
	}
}

// Connect registers a new SSE client and announces presence to the user's other devices.
// If a client for (userID, deviceID) already exists (e.g., a tab refresh), its channel is closed.
func (h *Hub) Connect(ctx context.Context, userID, deviceID string) *Client {
	c := &Client{
		UserID:   userID,
		DeviceID: deviceID,
		ch:       make(chan Event, clientBuffer),
	}

	h.mu.Lock()
	if h.users[userID] == nil {
		h.users[userID] = map[string]*Client{}
	}
	if old, ok := h.users[userID][deviceID]; ok {
		close(old.ch)
	}
	h.users[userID][deviceID] = c
	h.mu.Unlock()

	go h.publishPresence(ctx, userID)
	return c
}

// Disconnect removes a client and, after a grace period, re-publishes presence
// to peers if the client has not reconnected.
func (h *Hub) Disconnect(c *Client) {
	h.mu.Lock()
	if active, ok := h.users[c.UserID][c.DeviceID]; ok && active == c {
		delete(h.users[c.UserID], c.DeviceID)
		if len(h.users[c.UserID]) == 0 {
			delete(h.users, c.UserID)
		}
	}
	h.mu.Unlock()

	time.AfterFunc(h.grace, func() {
		h.mu.RLock()
		_, stillOnline := h.users[c.UserID][c.DeviceID]
		h.mu.RUnlock()
		if stillOnline {
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		h.publishPresence(ctx, c.UserID)
	})
}

// SendTo routes an event to a specific (userID, deviceID). Reports whether
// the target was online and the event was queued.
func (h *Hub) SendTo(userID, deviceID string, ev Event) bool {
	h.mu.RLock()
	c, ok := h.users[userID][deviceID]
	h.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case c.ch <- ev:
		return true
	default:
		slog.Warn("client buffer full; dropping event",
			"user", userID, "device", deviceID, "type", ev.Type)
		return false
	}
}

// Broadcast fans an event out to every live client of a user.
func (h *Hub) Broadcast(userID string, ev Event) {
	clients := h.snapshotClients(userID)
	for _, c := range clients {
		select {
		case c.ch <- ev:
		default:
			slog.Warn("client buffer full; dropping event",
				"user", c.UserID, "device", c.DeviceID, "type", ev.Type)
		}
	}
}

// Online reports whether a (userID, deviceID) currently has a live SSE.
func (h *Hub) Online(userID, deviceID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	_, ok := h.users[userID][deviceID]
	return ok
}

// Kick force-removes a (userID, deviceID) entry from the hub and closes its
// event channel; the SSE handler exits on the next iteration. 与 Connect 的
// 旧通道关闭模式一致（与并发 SendTo 之间存在极窄竞态，但 Kick 罕用，可接受）。
func (h *Hub) Kick(userID, deviceID string) {
	h.mu.Lock()
	c, ok := h.users[userID][deviceID]
	if ok {
		delete(h.users[userID], deviceID)
		if len(h.users[userID]) == 0 {
			delete(h.users, userID)
		}
	}
	h.mu.Unlock()
	if ok {
		close(c.ch)
	}
}

// PublishPresence broadcasts the user's current device list to all live
// clients. 用于在变更设备集合后（如 DELETE /api/devices/{name}）立刻广播，
// 而不是等待 Disconnect 的宽限期。
func (h *Hub) PublishPresence(ctx context.Context, userID string) {
	h.publishPresence(ctx, userID)
}

// Close drops every connected client. Used on graceful shutdown.
func (h *Hub) Close() {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.closed {
		return
	}
	h.closed = true
	for _, devices := range h.users {
		for _, c := range devices {
			close(c.ch)
		}
	}
	h.users = map[string]map[string]*Client{}
}

func (h *Hub) snapshotClients(userID string) []*Client {
	h.mu.RLock()
	defer h.mu.RUnlock()
	src := h.users[userID]
	out := make([]*Client, 0, len(src))
	for _, c := range src {
		out = append(out, c)
	}
	return out
}

func (h *Hub) publishPresence(ctx context.Context, userID string) {
	devs, err := h.devices.ListDevicesByUser(ctx, userID)
	if err != nil {
		slog.Error("presence: list devices failed", "user", userID, "err", err)
		return
	}

	h.mu.RLock()
	live := h.users[userID]
	items := make([]PresenceDevice, 0, len(devs))
	for _, d := range devs {
		_, online := live[d.Name]
		items = append(items, PresenceDevice{
			Name: d.Name, Type: d.Type, Online: online, LastSeen: d.LastSeen,
		})
	}
	clients := make([]*Client, 0, len(live))
	for _, c := range live {
		clients = append(clients, c)
	}
	h.mu.RUnlock()

	ev := Event{
		Type: "presence",
		Data: map[string]any{"devices": items},
	}
	for _, c := range clients {
		select {
		case c.ch <- ev:
		default:
			slog.Warn("presence: client buffer full; dropping",
				"user", c.UserID, "device", c.DeviceID)
		}
	}
}