How a Go Rewrite Made Our WebSocket Proxy 100x More Efficient

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

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(),
	}
}

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()
			}
		}
	}()
}