go-omni/ws/observe/observe.go

// 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
}