feat(ws): release generic WebSocket client and server v0.31.0-preview
This commit is contained in:
319
ws/hub/hub.go
Normal file
319
ws/hub/hub.go
Normal file
@@ -0,0 +1,319 @@
|
||||
// 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
|
||||
}
|
||||
Reference in New Issue
Block a user