通过 Go 重写将 WebSocket 代理效率提升 100 倍

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

// ────────────────────────────
// FIRST DESIGN (naive, global mutex)
// Single slice + flags, one coarse-grained mutex.
// ────────────────────────────

type Conn struct {
    id      string
    ws      *websocket.Conn
    inUse   bool
    healthy bool
}

type Pool struct {
    mu       sync.Mutex
    conns    []*Conn
    maxSize  int
    dialURL  string
}

// newConn dials a backend and returns a connected *Conn.
// NOTE: In this first design we (incorrectly) call this under the global lock.
func (p *Pool) newConn(ctx context.Context) (*Conn, error) {
    d := websocket.Dialer{}
    c, _, err := d.DialContext(ctx, p.dialURL, nil)
    if err != nil {
        return nil, err
    }
    return &Conn{
        id:      uuid.NewString(),
        ws:      c,
        inUse:   false,
        healthy: true,
    }, nil
}

// ────────────────────────────
// AdjustPool
// Check capacity vs current size; if not full, fill it.
// BAD PATTERN: holds the global mutex across slow I/O (dial).
// ────────────────────────────
func (p *Pool) AdjustPool(ctx context.Context) error {
    p.mu.Lock()
    defer p.mu.Unlock()

    cur := len(p.conns)
    if cur >= p.maxSize {
        return nil
    }

    needed := p.maxSize - cur
    for i := 0; i < needed; i++ {
        conn, err := p.newConn(ctx)
        if err != nil {
            return fmt.Errorf("adjust: dial failed: %w", err)
        }
        p.conns = append(p.conns, conn)
    }
    return nil
}

// ────────────────────────────
// GetPoolStats
// Return total / in-use / available counts.
// BAD PATTERN: O(n) scan under global lock every call.
// ────────────────────────────
type PoolStats struct {
    Total     int
    InUse     int
    Available int
    Healthy   int
    Unhealthy int
}

func (p *Pool) GetPoolStats() PoolStats {
    p.mu.Lock()
    defer p.mu.Unlock()

    var inUse, healthy int
    for _, c := range p.conns {
        if c.inUse {
            inUse++
        }
        if c.healthy {
            healthy++
        }
    }
    total := len(p.conns)
    return PoolStats{
        Total:     total,
        InUse:     inUse,
        Available: total - inUse,
        Healthy:   healthy,
        Unhealthy: total - healthy,
    }
}

// ────────────────────────────
// HealthCheck
// Ping all connections and mark healthy=false on failures.
// BAD PATTERN: holds global lock during network I/O & mutates in place.
// ────────────────────────────
func (p *Pool) HealthCheck(ctx context.Context, timeout time.Duration) {
    p.mu.Lock()
    defer p.mu.Unlock()

    deadline := time.Now().Add(timeout)
    for _, c := range p.conns {
        if c.ws == nil {
            c.healthy = false
            continue
        }
        if err := c.ws.WriteControl(websocket.PingMessage, []byte("ping"), deadline); err != nil {
            c.healthy = false
            _ = c.ws.Close()
            c.ws = nil
            continue
        }
        c.healthy = true
    }
}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

package pool

import (
	"context"
	"errors"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

// ────────────────────────────
// Conn (a single reusable backend WebSocket connection)
// ────────────────────────────
//
// ✅ GOOD PATTERN:
// - Keep each connection self-contained and concurrent-safe using atomics.
// - Encapsulate health logic inside the Conn itself (no shared state mutation).
// - Avoid external locks and let each conn manage its own goroutine lifecycle.

type Conn struct {
	id       string
	ws       *websocket.Conn
	healthy  atomic.Bool   // ✅ lock-free health status flag
	lastPing atomic.Int64  // ✅ atomic timestamp for last heartbeat
}

// ✅ GOOD PATTERN: Explicit small helper methods (no external mutation)
func (c *Conn) ID() string { return c.id }

func (c *Conn) IsAlive() bool {
	return c.healthy.Load()
}

// ✅ GOOD PATTERN: Safe close, idempotent and isolated
func (c *Conn) Close() error {
	if c.ws != nil {
		return c.ws.Close()
	}
	return nil
}

// ✅ GOOD PATTERN: Health loop runs independently per connection
// - No shared/global lock.
// - Non-blocking heartbeat.
// - Fails fast and marks itself dead without blocking pool operations.
func (c *Conn) StartHealthLoop(ctx context.Context, interval time.Duration) {
	t := time.NewTicker(interval)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			deadline := time.Now().Add(interval / 2)
			if err := c.ws.WriteControl(websocket.PingMessage, []byte("ping"), deadline); err != nil {
				c.healthy.Store(false)
				_ = c.Close()
				return
			}
			c.lastPing.Store(time.Now().UnixNano())
			c.healthy.Store(true)
		}
	}
}

// ────────────────────────────
// ConnPool (fast path only — no dialing, no blocking I/O)
// ────────────────────────────
//
// ✅ GOOD PATTERN:
// - Separate responsibilities: pool only manages available connections.
// - No dialing / blocking I/O under locks.
// - Use channel buffering and atomic counters for concurrency safety.
// - Eliminates coarse-grained global mutex.

type ConnPool struct {
	available chan *Conn   // ✅ buffered channel for ready conns (lock-free)
	inUse     sync.Map     // ✅ concurrent map for tracking active conns
	statsIn   atomic.Int64 // ✅ atomic counters (no need for locks)
	statsOut  atomic.Int64
}

// ✅ GOOD PATTERN: Explicit, fixed-capacity pool construction
func NewConnPool(capacity int) *ConnPool {
	return &ConnPool{
		available: make(chan *Conn, capacity),
	}
}

func (p *ConnPool) Capacity() int { return cap(p.available) }

// ✅ GOOD PATTERN: Non-blocking Acquire
// - Never holds locks while waiting for I/O.
// - Returns instantly if no conn available.
func (p *ConnPool) Acquire(ctx context.Context) (*Conn, error) {
	select {
	case c := <-p.available:
		p.inUse.Store(c.id, c)
		p.statsIn.Add(1)
		return c, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
		return nil, errors.New("no connection available")
	}
}

// ✅ GOOD PATTERN: Non-blocking Release
// - Never waits for space in the channel.
// - Drops unhealthy or excess conns immediately.
// - No global mutex.
func (p *ConnPool) Release(c *Conn) {
	p.inUse.Delete(c.id)
	p.statsOut.Add(1)

	if c.IsAlive() {
		select {
		case p.available <- c:
			// ✅ returned to pool safely
		default:
			// ✅ pool full → discard stale conn safely
			_ = c.Close()
		}
	} else {
		// ✅ unhealthy → close immediately
		_ = c.Close()
	}
}

// ✅ GOOD PATTERN: Offer is used by reconciler (external goroutine)
// - Keeps dialing/repair logic out of hot path.
// - Backpressure-safe with non-blocking insert.
func (p *ConnPool) Offer(c *Conn) bool {
	select {
	case p.available <- c:
		return true
	default:
		return false
	}
}

// ────────────────────────────
// Stats and Metrics
// ────────────────────────────
//
// ✅ GOOD PATTERN:
// - Lock-free snapshot using atomics.
// - Avoids holding locks during metrics collection.

type PoolStats struct {
	Capacity  int
	Available int
	InUse     int
	Acquired  int64
	Released  int64
}

// ✅ GOOD PATTERN: Snapshot safely aggregates pool state without blocking
func (p *ConnPool) Snapshot() PoolStats {
	inUseCount := 0
	p.inUse.Range(func(_, _ any) bool {
		inUseCount++
		return true
	})
	return PoolStats{
		Capacity:  cap(p.available),
		Available: len(p.available),
		InUse:     inUseCount,
		Acquired:  p.statsIn.Load(),
		Released:  p.statsOut.Load(),
	}
}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

type ServerConnectionPool struct {
	ctx         context.Context
	cancel      context.CancelFunc
	reconcileCh chan struct{}
	logger      *zap.Logger
	// ... other fields ...
}

// ✅ GOOD PATTERN: Constructor wires a buffered (size=1) signal channel to enable coalescing.
func NewServerConnectionPool(logger *zap.Logger /* ... */) *ServerConnectionPool {
	ctx, cancel := context.WithCancel(context.Background())
	return &ServerConnectionPool{
		ctx:         ctx,
		cancel:      cancel,
		reconcileCh: make(chan struct{}, 1), // ✅ coalescing buffer
		logger:      logger,
		// ... init other fields ...
	}
}

// ✅ GOOD PATTERN: Non-blocking signal helper; bursts coalesce into a single pending signal.
func (bp *ServerConnectionPool) trySend(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
		// already queued; coalesced
	}
}

// triggerReconcile only *requests* reconciliation; never performs it inline.
// ✅ GOOD PATTERN: no work on the caller's goroutine, prevents stampedes.
func (bp *ServerConnectionPool) triggerReconcile() {
	bp.trySend(bp.reconcileCh)
}

// ✅ GOOD PATTERN: Public starter that owns the worker lifecycle.
func (bp *ServerConnectionPool) Start() {
	go bp.reconcileWorker()
}

// ✅ GOOD PATTERN: Graceful shutdown.
func (bp *ServerConnectionPool) Stop() {
	bp.cancel()
}

// reconcileWorker serializes reconciliation and coalesces bursts.
// ✅ GOOD PATTERN:
//   - Single-threaded worker → no concurrent ensureCapacity() runs
//   - Periodic safety net with light jitter to avoid thundering herds
//   - Drain queue before each run to collapse multiple signals into one
func (bp *ServerConnectionPool) reconcileWorker() {
	jitter := func(base time.Duration) time.Duration {
		// small ±10% jitter
		n := time.Duration(float64(base) * (0.9 + 0.2*rand.Float64()))
		return n
	}

	ticker := time.NewTicker(jitter(5 * time.Second))
	defer ticker.Stop()

	bp.logger.Debug("Reconciliation worker started", zap.String("pool", bp.GetName()))
	defer bp.logger.Debug("Reconciliation worker stopped", zap.String("pool", bp.GetName()))

	// Optional: run once immediately on startup
	bp.ensureCapacity()

	for {
		select {
		case <-bp.ctx.Done():
			return

		case <-ticker.C:
			// ✅ Periodic reconciliation as a safety net
			bp.ensureCapacity()
			// reset ticker with jitter to spread load
			ticker.Reset(jitter(5 * time.Second))

		case <-bp.reconcileCh:
			// ✅ Drain any queued signals: burst -> single reconciliation
			for {
				select {
				case <-bp.reconcileCh:
					// keep draining
				default:
					bp.ensureCapacity()
					goto CONTINUE
				}
			}
		}
	CONTINUE:
	}
}

// ensureCapacity is the single authoritative place that:
// 1) Cleans unhealthy queued conns
// 2) Computes current vs target
// 3) Grows or shrinks the pool
// 4) Emits metrics/logs
func (bp *ServerConnectionPool) ensureCapacity() {}

// ────────────────────────────
// Event sources that *request* reconciliation
// ────────────────────────────

// Health watcher: when a server flips health, request reconcile.
// ✅ GOOD PATTERN: do not call ensureCapacity() here; just signal.
func (bp *ServerConnectionPool) startHealthWatcher(healthCh <-chan HealthEvent) {
	go func() {
		for {
			select {
			case <-bp.ctx.Done():
				return
			case ev := <-healthCh:
				if ev.Changed {
					bp.logger.Debug("health change → reconcile",
						zap.String("server", ev.ServerID))
					bp.triggerReconcile()
				}
			}
		}
	}()
}

// Config watcher: when scale target changes, request reconcile.
// ✅ GOOD PATTERN: apply config change, then signal worker.
func (bp *ServerConnectionPool) startConfigWatcher(cfgCh <-chan ScaleTarget) {
	go func() {
		for {
			select {
			case <-bp.ctx.Done():
				return
			case target := <-cfgCh:
				// update internal target state...
				bp.logger.Debug("scale target changed → reconcile",
					zap.Int("target", target.Connections))
				bp.triggerReconcile()
			}
		}
	}()
}