go-omni/ws/hub/hub.go

// Package hub provides a bounded, process-local WebSocket registry. Connection
// and group IDs are opaque application keys; callers authorize membership.
// It does not persist messages or provide cross-process delivery.
package hub

import (
	"errors"
	"reflect"
	"strings"
	"sync"

	"git.yuntuo.hk/k2safe/go-omni/ws"
)

var (
	ErrClosed    = errors.New("ws/hub: closed")
	ErrNotFound  = errors.New("ws/hub: connection not found")
	ErrDuplicate = errors.New("ws/hub: duplicate connection ID")
	ErrLimit     = errors.New("ws/hub: resource limit reached")
	ErrInvalid   = errors.New("ws/hub: invalid connection or group")
)

// Connection is implemented by *ws.Session. Implementations must be safe for
// concurrent use, have a stable ID, and make TrySend nonblocking with bounded
// memory. Done closes only after connection work finishes.
type Connection interface {
	ID() string
	Done() <-chan struct{}
	TrySend(ws.Message) error
	Close(int, string) error
}

type Config struct {
	MaxConnections         int `json:"maxConnections"`
	MaxGroups              int `json:"maxGroups"`
	MaxMemberships         int `json:"maxMemberships"`
	MaxGroupsPerConnection int `json:"maxGroupsPerConnection"`
}

func (cfg Config) Validate() error {
	if cfg.MaxConnections < 0 || cfg.MaxConnections > 1000000 || cfg.MaxGroups < 0 || cfg.MaxGroups > 1000000 || cfg.MaxMemberships < 0 || cfg.MaxMemberships > 10000000 || cfg.MaxGroupsPerConnection < 0 || cfg.MaxGroupsPerConnection > 10000 {
		return errors.New("ws/hub: invalid limits")
	}
	return nil
}

type entry struct {
	connection Connection
	groups     map[string]struct{}
	removed    chan struct{}
}
type Hub struct {
	mu          sync.RWMutex
	config      Config
	connections map[string]*entry
	groups      map[string]map[string]*entry
	memberships int
	closed      bool
	done        chan struct{}
	watchers    sync.WaitGroup
}

func New(cfg Config) (*Hub, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	if cfg.MaxConnections == 0 {
		cfg.MaxConnections = 10000
	}
	if cfg.MaxGroups == 0 {
		cfg.MaxGroups = 1024
	}
	if cfg.MaxMemberships == 0 {
		cfg.MaxMemberships = 100000
	}
	if cfg.MaxGroupsPerConnection == 0 {
		cfg.MaxGroupsPerConnection = 32
	}
	return &Hub{config: cfg, connections: make(map[string]*entry), groups: make(map[string]map[string]*entry), done: make(chan struct{})}, nil
}

func validKey(key string) bool {
	return key != "" && len(key) <= 256 && !strings.ContainsAny(key, "\r\n\x00")
}

// Add borrows a connection and starts one bounded watcher to remove it when
// Done closes. Remove and Close stop those watchers without closing connections.
func (h *Hub) Add(connection Connection) error {
	if connection == nil {
		return ErrInvalid
	}
	v := reflect.ValueOf(connection)
	if (v.Kind() == reflect.Pointer || v.Kind() == reflect.Interface || v.Kind() == reflect.Map || v.Kind() == reflect.Func || v.Kind() == reflect.Slice || v.Kind() == reflect.Chan) && v.IsNil() {
		return ErrInvalid
	}
	id := connection.ID()
	done := connection.Done()
	if !validKey(id) || done == nil {
		return ErrInvalid
	}
	select {
	case <-done:
		return ws.ErrClosed
	default:
	}
	h.mu.Lock()
	if h.closed {
		h.mu.Unlock()
		return ErrClosed
	}
	if _, ok := h.connections[id]; ok {
		h.mu.Unlock()
		return ErrDuplicate
	}
	if len(h.connections) >= h.config.MaxConnections {
		h.mu.Unlock()
		return ErrLimit
	}
	e := &entry{connection: connection, groups: make(map[string]struct{}), removed: make(chan struct{})}
	h.connections[id] = e
	h.watchers.Add(1)
	h.mu.Unlock()
	go func() {
		defer h.watchers.Done()
		select {
		case <-done:
			h.mu.Lock()
			if h.connections[id] == e {
				h.removeLocked(id, e)
			}
			h.mu.Unlock()
		case <-e.removed:
		case <-h.done:
		}
	}()
	return nil
}
func (h *Hub) removeLocked(id string, e *entry) {
	delete(h.connections, id)
	close(e.removed)
	for name := range e.groups {
		delete(h.groups[name], id)
		h.memberships--
		if len(h.groups[name]) == 0 {
			delete(h.groups, name)
		}
	}
}
func (h *Hub) Remove(id string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	e := h.connections[id]
	if e == nil {
		return false
	}
	h.removeLocked(id, e)
	return true
}

func (h *Hub) Join(id, group string) error {
	if !validKey(group) {
		return ErrInvalid
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.closed {
		return ErrClosed
	}
	e := h.connections[id]
	if e == nil {
		return ErrNotFound
	}
	if _, ok := e.groups[group]; ok {
		return nil
	}
	if len(e.groups) >= h.config.MaxGroupsPerConnection || h.memberships >= h.config.MaxMemberships {
		return ErrLimit
	}
	if h.groups[group] == nil {
		if len(h.groups) >= h.config.MaxGroups {
			return ErrLimit
		}
		h.groups[group] = make(map[string]*entry)
	}
	h.groups[group][id] = e
	e.groups[group] = struct{}{}
	h.memberships++
	return nil
}
func (h *Hub) Leave(id, group string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	e := h.connections[id]
	if e == nil {
		return false
	}
	if _, ok := e.groups[group]; !ok {
		return false
	}
	delete(e.groups, group)
	delete(h.groups[group], id)
	h.memberships--
	if len(h.groups[group]) == 0 {
		delete(h.groups, group)
	}
	return true
}

// Result reports queue acceptance, never remote receipt or business success.
// Errors contains one entry per rejected destination; no messages are silently
// dropped. Concurrent removal may race a snapshot and yield a closed error.
type Result struct {
	Attempted int
	Queued    int
	Errors    map[string]error
}

func send(entries map[string]*entry, message ws.Message) Result {
	result := Result{Attempted: len(entries)}
	for id, e := range entries {
		if err := e.connection.TrySend(message); err != nil {
			if result.Errors == nil {
				result.Errors = make(map[string]error)
			}
			result.Errors[id] = err
		} else {
			result.Queued++
		}
	}
	return result
}
func (h *Hub) Send(id string, message ws.Message) (Result, error) {
	h.mu.RLock()
	if h.closed {
		h.mu.RUnlock()
		return Result{}, ErrClosed
	}
	e := h.connections[id]
	h.mu.RUnlock()
	if e == nil {
		return Result{}, ErrNotFound
	}
	return send(map[string]*entry{id: e}, message), nil
}
func (h *Hub) snapshot(group string) (map[string]*entry, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	if h.closed {
		return nil, ErrClosed
	}
	source := h.connections
	if group != "" {
		source = h.groups[group]
	}
	out := make(map[string]*entry, len(source))
	for id, e := range source {
		out[id] = e
	}
	return out, nil
}
func (h *Hub) Broadcast(message ws.Message) (Result, error) {
	entries, err := h.snapshot("")
	if err != nil {
		return Result{}, err
	}
	return send(entries, message), nil
}
func (h *Hub) SendGroup(group string, message ws.Message) (Result, error) {
	if !validKey(group) {
		return Result{}, ErrInvalid
	}
	entries, err := h.snapshot(group)
	if err != nil {
		return Result{}, err
	}
	return send(entries, message), nil
}
func (h *Hub) Kick(id string, code int, reason string) error {
	h.mu.RLock()
	if h.closed {
		h.mu.RUnlock()
		return ErrClosed
	}
	e := h.connections[id]
	h.mu.RUnlock()
	if e == nil {
		return ErrNotFound
	}
	return e.connection.Close(code, reason)
}

type Stats struct {
	Connections int
	Groups      int
	Memberships int
}

func (h *Hub) Stats() Stats {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return Stats{len(h.connections), len(h.groups), h.memberships}
}

// Close unregisters borrowed connections, releases membership memory and waits
// for internal watchers only. Use Session/Server shutdown to close connections.
func (h *Hub) Close() error {
	h.mu.Lock()
	if !h.closed {
		h.closed = true
		close(h.done)
		for id, e := range h.connections {
			h.removeLocked(id, e)
		}
		h.groups = make(map[string]map[string]*entry)
	}
	h.mu.Unlock()
	h.watchers.Wait()
	return nil
}