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/observe/README.md
+++ b/ws/observe/README.md
@@ -0,0 +1,44 @@
+# WebSocket 可选观测
+
+`ws/observe` 为 `ws.Client`、`ws.Server` 的会话提供连接日志和聚合快照。核心 `ws` 不依赖全局 Logger；观测器不启动后台 goroutine，不给每条消息增加日志、锁或 span。
+
+```go
+logger, err := logx.New(logx.Default())
+if err != nil { return err }
+defer logger.Close() // 必须晚于 Client/Server 的 Shutdown 完成
+
+observer := observe.New(logger) // nil 表示只统计，不记日志
+handlers := observer.Wrap(ws.Handlers{
+    OnMessage: func(ctx context.Context, session *ws.Session, message ws.Message) error {
+        return session.Send(ctx, message)
+    },
+})
+server, err := ws.NewServer(ws.ServerConfig{})
+if err != nil { return err }
+mux.Handle("/ws", server.Handler(handlers))
+
+// 由项目按自己的采样周期读取，接入 metrics、Prometheus 或 OTel。
+snapshot := observer.Snapshot()
+```
+
+导入路径为 `git.yuntuo.hk/k2safe/go-omni/ws/observe`。同一个 Observer 和同一组包装后的 Handlers 可供多个连接并发复用；每个会话有独立的开始时间。每个会话仅包装一次。使用客户端重连和 `ClientConfig.Prepare` 时，在每次 Prepare 返回前调用 `observer.Wrap(handlers)`。
+
+快照包含：
+
+- `Opened / Closed / Active`：已观察到的连接总数、关闭回调结束数、当前处于受观察生命周期中的连接数。连接可能在 OnOpen 执行前断开，此时仍在 OnClose 计入 Opened 和 Closed，日志时长为 0。
+- `MessagesSent / MessagesReceived / BytesSent / BytesReceived`：已关闭会话的最终统计，加上当前会话的即时统计。发送成功表示写入完成；不是对端收到消息或业务处理成功的确认。统计口径与 `Session.Stats()` 一致。
+- `WriteQueuedMessages / WriteQueuedBytes / ReadQueuedMessages / ReadQueuedBytes`：调用快照时，各活跃连接队列水位之和。接收侧包括仍在处理、尚未释放的消息。
+- `SampledPeak*`：历次 Snapshot 观察到的聚合峰值；不是精确的历史峰值，采样间隔内的突增可能未被记录。
+- `CloseCallbackPanics`：下游 OnClose 回调发生 panic 的次数。
+
+Snapshot 为 O(活跃会话数)，只在采样时读取 `Session.Stats()`，不会替换消息回调。避免在每条消息上采样；采样周期由项目根据连接规模决定。快照不会保存已关闭连接的 ID，也不引入高基数指标标签。自行导出 OTel 时，建议记录握手、连接时长、消息数与字节数；业务消息追踪应显式传递上下文，不自动延长握手 span 到整条连接的寿命。
+
+日志只由适配器写入固定事件名、Session ID、时长、关闭码和固定阶段。不会展开错误文本、Cause、URL、请求头、消息体或对端关闭原因；业务自定义 `OpError.Op` 也通过固定白名单筛选。用户传入 Logger 自身已经附加的属性仍由用户负责。默认日志级别为 Info。
+
+OnClose 的收尾使用 defer，即使业务回调 panic，也会归档消息统计、移除活跃记录、记录固定的 `on_close_panic` 阶段，再把 panic 交还核心处理。此处“关闭”表示 I/O 已停止且下游 OnClose 已返回或 panic；核心随后才关闭 `Session.Done()`。因此 `Snapshot.Active == 0` 不是生命周期完成屏障，停机必须等待 `Client/Server.Shutdown` 或 `Session.Done`，再关闭 Logger。观测器不会在 OnClose 中等待自身 Done，也不会在 Done 后继续写日志。
+
+运行真实 TCP 与并发测试：
+
+```sh
+go test -race ./ws/observe
+```
--- a/ws/observe/observe.go
+++ b/ws/observe/observe.go
@@ -0,0 +1,194 @@
+// Package observe adds opt-in, payload-free WebSocket lifecycle logs and
+// aggregate snapshots. The ws core has no logger dependency. Observers do not
+// own loggers, start goroutines, or create per-message tracing spans.
+package observe
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"git.yuntuo.hk/k2safe/go-omni/logx"
+	"git.yuntuo.hk/k2safe/go-omni/ws"
+)
+
+// Snapshot contains process-local totals for sessions observed by this Observer.
+// Active covers observed callback lifetimes: removal happens after the user's
+// OnClose returns, immediately before Session.Done. It is not a drain barrier.
+// Queue peaks are sampled only when Snapshot is called, not lifetime maxima.
+type Snapshot struct {
+	Opened                         uint64 `json:"opened"`
+	Closed                         uint64 `json:"closed"`
+	Active                         int    `json:"active"`
+	CloseCallbackPanics            uint64 `json:"closeCallbackPanics"`
+	MessagesSent                   uint64 `json:"messagesSent"`
+	MessagesReceived               uint64 `json:"messagesReceived"`
+	BytesSent                      uint64 `json:"bytesSent"`
+	BytesReceived                  uint64 `json:"bytesReceived"`
+	WriteQueuedMessages            int    `json:"writeQueuedMessages"`
+	WriteQueuedBytes               int64  `json:"writeQueuedBytes"`
+	ReadQueuedMessages             int    `json:"readQueuedMessages"`
+	ReadQueuedBytes                int64  `json:"readQueuedBytes"`
+	SampledPeakWriteQueuedMessages int    `json:"sampledPeakWriteQueuedMessages"`
+	SampledPeakWriteQueuedBytes    int64  `json:"sampledPeakWriteQueuedBytes"`
+	SampledPeakReadQueuedMessages  int    `json:"sampledPeakReadQueuedMessages"`
+	SampledPeakReadQueuedBytes     int64  `json:"sampledPeakReadQueuedBytes"`
+}
+
+type observation struct{ started time.Time }
+
+type Observer struct {
+	log    *logx.Logger
+	mu     sync.Mutex
+	active map[*ws.Session]observation
+	totals Snapshot
+}
+
+// New borrows log without changing or closing it. nil disables logs and retains
+// metrics. Keep the logger open until Client/Server shutdown completes.
+func New(log *logx.Logger) *Observer {
+	return &Observer{log: log, active: make(map[*ws.Session]observation)}
+}
+
+// Wrap can be reused across concurrent sessions. It leaves OnMessage unchanged;
+// message accounting comes from Session.Stats during snapshots and closure.
+// Apply it once per session. With ClientConfig.Prepare, wrap each returned set
+// of handlers so each reconnect uses the same observer and its own start time.
+func (o *Observer) Wrap(next ws.Handlers) ws.Handlers {
+	if o == nil {
+		return next
+	}
+	return ws.Handlers{
+		OnOpen: func(ctx context.Context, s *ws.Session) error {
+			o.opened(s)
+			if next.OnOpen != nil {
+				return next.OnOpen(ctx, s)
+			}
+			return nil
+		},
+		OnMessage: next.OnMessage,
+		OnClose: func(s *ws.Session, err error) {
+			defer func() {
+				p := recover()
+				o.closed(s, err, p != nil)
+				if p != nil {
+					panic(p)
+				}
+			}()
+			if next.OnClose != nil {
+				next.OnClose(s, err)
+			}
+		},
+	}
+}
+
+func (o *Observer) opened(s *ws.Session) {
+	if s == nil {
+		return
+	}
+	o.mu.Lock()
+	_, exists := o.active[s]
+	if !exists {
+		o.active[s] = observation{started: time.Now()}
+		o.totals.Opened++
+	}
+	o.mu.Unlock()
+	if !exists && o.log != nil {
+		o.log.Info("ws.open", "sessionId", s.ID(), "phase", "open")
+	}
+}
+
+func (o *Observer) closed(s *ws.Session, err error, panicked bool) {
+	if s == nil {
+		return
+	}
+	stats := s.Stats()
+	o.mu.Lock()
+	observed, exists := o.active[s]
+	delete(o.active, s)
+	// A connection can close before its dispatch loop reaches OnOpen. Count
+	// that observed physical connection too, with zero measured duration.
+	if !exists {
+		o.totals.Opened++
+	}
+	o.totals.Closed++
+	if panicked {
+		o.totals.CloseCallbackPanics++
+	}
+	addMessages(&o.totals, stats)
+	o.mu.Unlock()
+	duration := time.Duration(0)
+	if !observed.started.IsZero() {
+		duration = time.Since(observed.started)
+	}
+	if o.log != nil {
+		code, phase := closeFields(err, panicked)
+		o.log.Info("ws.close", "sessionId", s.ID(), "durationMs", duration.Seconds()*1000, "closeCode", code, "phase", phase)
+	}
+}
+
+func closeFields(err error, panicked bool) (int, string) {
+	if panicked {
+		return 0, "on_close_panic"
+	}
+	if err == nil {
+		return 1000, "close"
+	}
+	var closed *ws.CloseError
+	if errors.As(err, &closed) {
+		return closed.Code, "close"
+	}
+	var op *ws.OpError
+	if errors.As(err, &op) {
+		// OpError is exported. Never write an arbitrary user-created Op to logs.
+		switch op.Op {
+		case "read", "write", "on_open", "on_message", "on_close", "ping", "pong", "close", "encode":
+			return 0, op.Op
+		}
+	}
+	if errors.Is(err, context.Canceled) {
+		return 0, "canceled"
+	}
+	if errors.Is(err, context.DeadlineExceeded) {
+		return 0, "timeout"
+	}
+	return 0, "error"
+}
+
+func addMessages(out *Snapshot, stats ws.SessionStats) {
+	out.MessagesSent += stats.MessagesSent
+	out.MessagesReceived += stats.MessagesReceived
+	out.BytesSent += stats.BytesSent
+	out.BytesReceived += stats.BytesReceived
+}
+
+// Snapshot samples live Session counters and adds immutable closed totals. It
+// costs O(active sessions), acquires no locks on the message path, and keeps no
+// per-connection records after OnClose. Message totals count I/O, not peer ACKs.
+func (o *Observer) Snapshot() Snapshot {
+	if o == nil {
+		return Snapshot{}
+	}
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	out := o.totals
+	out.Active = len(o.active)
+	for session := range o.active {
+		stats := session.Stats()
+		addMessages(&out, stats)
+		out.WriteQueuedMessages += stats.QueuedMessages
+		out.WriteQueuedBytes += stats.QueuedBytes
+		out.ReadQueuedMessages += stats.ReadQueuedMessages
+		out.ReadQueuedBytes += stats.ReadQueuedBytes
+	}
+	o.totals.SampledPeakWriteQueuedMessages = max(o.totals.SampledPeakWriteQueuedMessages, out.WriteQueuedMessages)
+	o.totals.SampledPeakWriteQueuedBytes = max(o.totals.SampledPeakWriteQueuedBytes, out.WriteQueuedBytes)
+	o.totals.SampledPeakReadQueuedMessages = max(o.totals.SampledPeakReadQueuedMessages, out.ReadQueuedMessages)
+	o.totals.SampledPeakReadQueuedBytes = max(o.totals.SampledPeakReadQueuedBytes, out.ReadQueuedBytes)
+	out.SampledPeakWriteQueuedMessages = o.totals.SampledPeakWriteQueuedMessages
+	out.SampledPeakWriteQueuedBytes = o.totals.SampledPeakWriteQueuedBytes
+	out.SampledPeakReadQueuedMessages = o.totals.SampledPeakReadQueuedMessages
+	out.SampledPeakReadQueuedBytes = o.totals.SampledPeakReadQueuedBytes
+	return out
+}
--- a/ws/observe/observe_test.go
+++ b/ws/observe/observe_test.go
@@ -0,0 +1,300 @@
+package observe
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"git.yuntuo.hk/k2safe/go-omni/logx"
+	"git.yuntuo.hk/k2safe/go-omni/ws"
+	"github.com/gorilla/websocket"
+)
+
+func echo(w http.ResponseWriter, r *http.Request) {
+	conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+	for {
+		kind, data, err := conn.ReadMessage()
+		if err != nil {
+			return
+		}
+		if err = conn.WriteMessage(kind, data); err != nil {
+			return
+		}
+	}
+}
+func testClient(t *testing.T) *ws.Client {
+	t.Helper()
+	c, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 50 * time.Millisecond}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		if err := c.Shutdown(ctx); err != nil {
+			t.Error(err)
+		}
+	})
+	return c
+}
+func endpoint(s *httptest.Server) string { return "ws" + strings.TrimPrefix(s.URL, "http") }
+func wait(t *testing.T, ch <-chan struct{}) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout")
+	}
+}
+
+func TestObserverConcurrentSessionsAndTotals(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(echo))
+	defer server.Close()
+	client := testClient(t)
+	o := New(nil)
+	const connections = 8
+	const count = 10
+	var workers sync.WaitGroup
+	stopSamples := make(chan struct{})
+	samplesDone := make(chan struct{})
+	go func() {
+		defer close(samplesDone)
+		for {
+			select {
+			case <-stopSamples:
+				return
+			default:
+				_ = o.Snapshot()
+			}
+		}
+	}()
+	for i := 0; i < connections; i++ {
+		workers.Add(1)
+		go func() {
+			defer workers.Done()
+			opened := make(chan struct{})
+			received := make(chan struct{}, count)
+			s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { received <- struct{}{}; return nil }}))
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			wait(t, opened)
+			for j := 0; j < count; j++ {
+				if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("test")}); err != nil {
+					t.Error(err)
+					return
+				}
+				wait(t, received)
+			}
+			_ = s.Close(1000, "")
+			wait(t, s.Done())
+		}()
+	}
+	workers.Wait()
+	close(stopSamples)
+	wait(t, samplesDone)
+	got := o.Snapshot()
+	if got.Opened != connections || got.Closed != connections || got.Active != 0 || got.MessagesSent != connections*count || got.MessagesReceived != connections*count || got.BytesSent != 4*connections*count || got.BytesReceived != 4*connections*count {
+		t.Fatalf("snapshot: %+v", got)
+	}
+	if again := o.Snapshot(); again != got {
+		t.Fatalf("closed totals changed: %+v -> %+v", got, again)
+	}
+}
+
+func TestObserverCloseCallbackBoundaryAndPanic(t *testing.T) {
+	for _, panics := range []bool{false, true} {
+		t.Run(fmt.Sprint(panics), func(t *testing.T) {
+			server := httptest.NewServer(http.HandlerFunc(echo))
+			defer server.Close()
+			client := testClient(t)
+			o := New(nil)
+			opened := make(chan struct{})
+			closing := make(chan struct{})
+			release := make(chan struct{})
+			s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnClose: func(*ws.Session, error) {
+				close(closing)
+				<-release
+				if panics {
+					panic("SECRET CALLBACK PANIC")
+				}
+			}}))
+			if err != nil {
+				t.Fatal(err)
+			}
+			wait(t, opened)
+			_ = s.Close(1000, "")
+			wait(t, closing)
+			if got := o.Snapshot(); got.Active != 1 || got.Closed != 0 {
+				t.Fatalf("closed before callback finished: %+v", got)
+			}
+			select {
+			case <-s.Done():
+				t.Fatal("Done closed before user callback")
+			default:
+			}
+			close(release)
+			wait(t, s.Done())
+			got := o.Snapshot()
+			if got.Active != 0 || got.Closed != 1 {
+				t.Fatalf("callback did not release observation: %+v", got)
+			}
+			if panics && got.CloseCallbackPanics != 1 {
+				t.Fatalf("panic not counted: %+v", got)
+			}
+		})
+	}
+}
+
+func TestObserverPayloadFreeLogsAndIndependentDurations(t *testing.T) {
+	var output bytes.Buffer
+	logger, err := logx.New(logx.Config{Writer: &output, Format: "json", Level: "info"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer logger.Close()
+	o := New(logger)
+	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()
+		if _, _, err = conn.ReadMessage(); err == nil {
+			_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1008, "SECRET CLOSE REASON"), time.Now().Add(time.Second))
+			_, _, _ = conn.ReadMessage()
+		}
+	}))
+	defer server.Close()
+	client := testClient(t)
+	var sessions []*ws.Session
+	for i := 0; i < 2; i++ {
+		opened := make(chan struct{})
+		s, err := client.Dial(context.Background(), endpoint(server)+"/?token=SECRET-URL", o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }}))
+		if err != nil {
+			t.Fatal(err)
+		}
+		wait(t, opened)
+		sessions = append(sessions, s)
+		if i == 0 {
+			time.Sleep(40 * time.Millisecond)
+		}
+	}
+	for _, s := range sessions {
+		if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("SECRET MESSAGE")}); err != nil {
+			t.Fatal(err)
+		}
+		wait(t, s.Done())
+	}
+	if strings.Contains(output.String(), "SECRET") {
+		t.Fatalf("log leaked source data: %s", output.String())
+	}
+	var durations []float64
+	for _, line := range strings.Split(strings.TrimSpace(output.String()), "\n") {
+		var row map[string]any
+		if err = json.Unmarshal([]byte(line), &row); err != nil {
+			t.Fatal(err)
+		}
+		for key := range row {
+			switch key {
+			case "time", "level", "msg", "sessionId", "phase", "durationMs", "closeCode":
+			default:
+				t.Fatalf("unexpected log field %q", key)
+			}
+		}
+		if row["msg"] == "ws.close" {
+			if row["closeCode"] != float64(1008) {
+				t.Fatalf("close code: %#v", row)
+			}
+			durations = append(durations, row["durationMs"].(float64))
+		}
+	}
+	if len(durations) != 2 || durations[0] < 30 {
+		t.Fatalf("timers not per session: %v", durations)
+	}
+}
+
+func TestObserverQueueSnapshotsAndSampledPeaks(t *testing.T) {
+	startSending := make(chan struct{})
+	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()
+		<-startSending
+		for i := 0; i < 5; i++ {
+			if conn.WriteMessage(websocket.TextMessage, []byte("queue")) != nil {
+				return
+			}
+		}
+		for {
+			if _, _, err = conn.ReadMessage(); err != nil {
+				return
+			}
+		}
+	}))
+	defer server.Close()
+	client := testClient(t)
+	o := New(nil)
+	opened := make(chan struct{})
+	release := make(chan struct{})
+	s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { <-release; return nil }}))
+	if err != nil {
+		t.Fatal(err)
+	}
+	wait(t, opened)
+	close(startSending)
+	deadline := time.After(2 * time.Second)
+	var sampled Snapshot
+	for {
+		sampled = o.Snapshot()
+		if sampled.MessagesReceived == 5 {
+			break
+		}
+		select {
+		case <-deadline:
+			close(release)
+			t.Fatal("receive timeout")
+		default:
+			time.Sleep(time.Millisecond)
+		}
+	}
+	if sampled.ReadQueuedMessages != 5 || sampled.ReadQueuedBytes != 25 || sampled.SampledPeakReadQueuedMessages != 5 {
+		t.Fatalf("queue snapshot: %+v", sampled)
+	}
+	close(release)
+	_ = s.Close(1000, "")
+	wait(t, s.Done())
+	final := o.Snapshot()
+	if final.ReadQueuedMessages != 0 || final.SampledPeakReadQueuedMessages != 5 {
+		t.Fatalf("final: %+v", final)
+	}
+}
+
+type opaqueError struct{}
+
+func (opaqueError) Error() string { panic("must not expand error") }
+func TestObserverCloseFieldsNeverExpandsErrors(t *testing.T) {
+	for _, err := range []error{opaqueError{}, &ws.OpError{Op: "SECRET STAGE", Cause: opaqueError{}}, &ws.CloseError{Code: 1008, Cause: opaqueError{}}} {
+		_, phase := closeFields(err, false)
+		if strings.Contains(phase, "SECRET") {
+			t.Fatal("untrusted phase")
+		}
+	}
+	if code, phase := closeFields(opaqueError{}, true); code != 0 || phase != "on_close_panic" {
+		t.Fatalf("panic fields: %d %s", code, phase)
+	}
+}