feat(ws): release generic WebSocket client and server v0.31.0-preview

2026-10-29 19:20:12 +08:00
parent 19ee9746e9
commit 3272974bc1
58 changed files with 7654 additions and 21 deletions
--- a/ws/hub/hub.go
+++ b/ws/hub/hub.go
@@ -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
+}
--- a/ws/hub/hub_test.go
+++ b/ws/hub/hub_test.go
@@ -0,0 +1,271 @@
+package hub
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"git.yuntuo.hk/k2safe/go-omni/ws"
+	"github.com/gorilla/websocket"
+)
+
+type testConnection struct {
+	id   string
+	done chan struct{}
+	once sync.Once
+	full bool
+	sent atomic.Int64
+}
+
+func (c *testConnection) ID() string            { return c.id }
+func (c *testConnection) Done() <-chan struct{} { return c.done }
+func (c *testConnection) TrySend(ws.Message) error {
+	select {
+	case <-c.done:
+		return ws.ErrClosed
+	default:
+	}
+	if c.full {
+		return ws.ErrQueueFull
+	}
+	c.sent.Add(1)
+	return nil
+}
+func (c *testConnection) Close(int, string) error { c.once.Do(func() { close(c.done) }); return nil }
+func newConnection(id string) *testConnection {
+	return &testConnection{id: id, done: make(chan struct{})}
+}
+func testHub(t *testing.T, cfg Config) *Hub {
+	t.Helper()
+	h, err := New(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = h.Close() })
+	return h
+}
+
+func TestHubRoutingBackpressureAndCleanup(t *testing.T) {
+	h := testHub(t, Config{})
+	a, b, c := newConnection("a"), newConnection("b"), newConnection("c")
+	b.full = true
+	for _, connection := range []*testConnection{a, b, c} {
+		if err := h.Add(connection); err != nil {
+			t.Fatal(err)
+		}
+	}
+	for _, id := range []string{"a", "b"} {
+		if err := h.Join(id, "group"); err != nil {
+			t.Fatal(err)
+		}
+	}
+	message := ws.Message{Type: ws.Text, Data: []byte("hello")}
+	result, err := h.SendGroup("group", message)
+	if err != nil || result.Attempted != 2 || result.Queued != 1 || !errors.Is(result.Errors["b"], ws.ErrQueueFull) {
+		t.Fatalf("group: %#v %v", result, err)
+	}
+	if c.sent.Load() != 0 {
+		t.Fatal("group leaked")
+	}
+	result, err = h.Broadcast(message)
+	if err != nil || result.Attempted != 3 || result.Queued != 2 {
+		t.Fatalf("broadcast: %#v %v", result, err)
+	}
+	if err = h.Kick("a", 1000, ""); err != nil {
+		t.Fatal(err)
+	}
+	deadline := time.After(time.Second)
+	for h.Stats().Connections != 2 {
+		select {
+		case <-deadline:
+			t.Fatal("auto cleanup timeout")
+		default:
+			time.Sleep(time.Millisecond)
+		}
+	}
+	if got := h.Stats(); got.Memberships != 1 || got.Groups != 1 {
+		t.Fatalf("stats: %#v", got)
+	}
+	if !h.Remove("b") {
+		t.Fatal("remove")
+	}
+	if got := h.Stats(); got.Memberships != 0 || got.Groups != 0 {
+		t.Fatalf("group leak: %#v", got)
+	}
+	if _, err = h.Send("missing", message); !errors.Is(err, ErrNotFound) {
+		t.Fatal(err)
+	}
+	_ = h.Close()
+	if _, err = h.Broadcast(message); !errors.Is(err, ErrClosed) {
+		t.Fatal(err)
+	}
+	select {
+	case <-c.done:
+		t.Fatal("hub closed borrowed connection")
+	default:
+	}
+}
+
+func TestHubLimitsAndReusedID(t *testing.T) {
+	h := testHub(t, Config{MaxConnections: 2, MaxGroups: 1, MaxMemberships: 1, MaxGroupsPerConnection: 1})
+	first := newConnection("a")
+	if err := h.Add(first); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.Add(first); !errors.Is(err, ErrDuplicate) {
+		t.Fatal(err)
+	}
+	if err := h.Add(newConnection("b")); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.Add(newConnection("c")); !errors.Is(err, ErrLimit) {
+		t.Fatal(err)
+	}
+	if err := h.Join("a", "one"); err != nil {
+		t.Fatal(err)
+	}
+	if err := h.Join("a", "one"); err != nil {
+		t.Fatal("idempotent join", err)
+	}
+	for _, pair := range [][2]string{{"a", "two"}, {"b", "one"}} {
+		if err := h.Join(pair[0], pair[1]); !errors.Is(err, ErrLimit) {
+			t.Fatal(err)
+		}
+	}
+	h.Remove("a")
+	second := newConnection("a")
+	if err := h.Add(second); err != nil {
+		t.Fatal(err)
+	}
+	_ = first.Close(1000, "")
+	if err := h.Join("a", "replacement"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := h.Send("a", ws.Message{Type: ws.Text}); err != nil {
+		t.Fatal("old watcher removed replacement", err)
+	}
+	var nilConnection *testConnection
+	if err := h.Add(nilConnection); !errors.Is(err, ErrInvalid) {
+		t.Fatal(err)
+	}
+}
+
+func TestHubConcurrentLifecycle(t *testing.T) {
+	h := testHub(t, Config{MaxConnections: 100})
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			c := newConnection(fmt.Sprint(i))
+			if err := h.Add(c); err != nil {
+				t.Error(err)
+				return
+			}
+			for j := 0; j < 10; j++ {
+				_ = h.Join(c.ID(), "shared")
+				_, _ = h.Broadcast(ws.Message{Type: ws.Text})
+				h.Leave(c.ID(), "shared")
+			}
+			_ = c.Close(1000, "")
+		}(i)
+	}
+	wg.Wait()
+	_ = h.Close()
+	if h.Stats() != (Stats{}) {
+		t.Fatal("registry not released")
+	}
+}
+
+func TestHubRealSessionDeliveryAndDisconnect(t *testing.T) {
+	received := make(chan string, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		for {
+			_, data, err := conn.ReadMessage()
+			if err != nil {
+				return
+			}
+			received <- string(data)
+		}
+	}))
+	defer server.Close()
+	client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 100 * time.Millisecond}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		if err := client.Shutdown(ctx); err != nil {
+			t.Error(err)
+		}
+	}()
+	s, err := client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), ws.Handlers{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	h := testHub(t, Config{})
+	if err = h.Add(s); err != nil {
+		t.Fatal(err)
+	}
+	if err = h.Join(s.ID(), "room"); err != nil {
+		t.Fatal(err)
+	}
+	result, err := h.SendGroup("room", ws.Message{Type: ws.Text, Data: []byte("queued")})
+	if err != nil || result.Queued != 1 {
+		t.Fatalf("send: %#v %v", result, err)
+	}
+	select {
+	case got := <-received:
+		if got != "queued" {
+			t.Fatal(got)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("delivery timed out")
+	}
+	_ = s.Close(1000, "")
+	deadline := time.After(time.Second)
+	for h.Stats().Connections != 0 {
+		select {
+		case <-deadline:
+			t.Fatal("disconnected Session still registered")
+		default:
+			time.Sleep(time.Millisecond)
+		}
+	}
+}
+
+func BenchmarkHubBroadcast(b *testing.B) {
+	for _, count := range []int{1, 100, 1000} {
+		b.Run(fmt.Sprint(count), func(b *testing.B) {
+			h, _ := New(Config{})
+			defer h.Close()
+			for i := 0; i < count; i++ {
+				if err := h.Add(newConnection(fmt.Sprint(i))); err != nil {
+					b.Fatal(err)
+				}
+			}
+			message := ws.Message{Type: ws.Text, Data: []byte("benchmark")}
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				result, err := h.Broadcast(message)
+				if err != nil || result.Queued != count {
+					b.Fatal("broadcast failed")
+				}
+			}
+		})
+	}
+}