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
	// Type is the client-declared device type (browser / macos / windows / linux / ios).
	// It lets presence label a live device that has no devices-table row (a global-SSO
	// browser or a broker-authenticated desktop).
	Type 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.
// deviceType is the caller's declared device type, surfaced in presence for a live device
// that has no devices-table row.
func (h *Hub) Connect(ctx context.Context, userID, deviceID, deviceType string) *Client {
	c := &Client{
		UserID:   userID,
		DeviceID: deviceID,
		Type:     deviceType,
		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.
//
// The send happens while still holding the read lock so it can never race a Kick / Connect-
// replace / Close that closes the channel (those hold the write lock): a send on a closed
// channel panics even inside a select, so close-vs-send must be mutually exclusive. The send
// is non-blocking (select default), so holding the read lock across it is brief.
func (h *Hub) SendTo(userID, deviceID string, ev Event) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	c, ok := h.users[userID][deviceID]
	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. The non-blocking sends run
// under the read lock so they can't race a concurrent channel close (see SendTo).
func (h *Hub) Broadcast(userID string, ev Event) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, c := range h.users[userID] {
		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. The close happens UNDER the write lock — the
// same discipline as Connect-replace and Close — so it can never race a send from
// publishPresence / Broadcast / SendTo (those hold the read lock), which would otherwise
// panic on a send to a closed channel. revokeDevice now Kicks on every logout / device
// delete / session revoke, so this path is hot, not rare.
func (h *Hub) Kick(userID, deviceID string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if c, ok := h.users[userID][deviceID]; ok {
		delete(h.users[userID], deviceID)
		if len(h.users[userID]) == 0 {
			delete(h.users, userID)
		}
		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) publishPresence(ctx context.Context, userID string) {
	devs, err := h.devices.ListDevicesByUser(ctx, userID)
	if err != nil {
		// Log but continue with an empty managed-device set: a transient DB error must
		// not blank out the presence of live, unmanaged devices that need no row.
		slog.Error("presence: list devices failed", "user", userID, "err", err)
	}

	now := time.Now().Unix()
	// Hold the read lock across the build AND the sends: the non-blocking sends below must
	// be mutually exclusive with any channel close (Kick / Connect-replace / Close hold the
	// write lock), or a send could hit a closed channel and panic.
	h.mu.RLock()
	defer h.mu.RUnlock()
	live := h.users[userID]
	seen := make(map[string]bool, len(devs))
	items := make([]PresenceDevice, 0, len(devs)+len(live))
	for _, d := range devs {
		_, online := live[d.Name]
		items = append(items, PresenceDevice{
			Name: d.Name, Type: d.Type, Online: online, LastSeen: d.LastSeen,
		})
		seen[d.Name] = true
	}
	// Live connections without a devices-table row — a global-SSO browser before 代铸, or a
	// device whose cache row hasn't landed yet. They are reachable on the hub (clipboard /
	// signaling already route to them), so they must appear as online peers too.
	for name, c := range live {
		if seen[name] {
			continue
		}
		items = append(items, PresenceDevice{
			Name: name, Type: c.Type, Online: true, LastSeen: now,
		})
	}

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