redapricot/client/worker.go

package client

import (
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iceBear67/redapricot/client/wire"
)

// StreamsBeforeGrowing is how many streams a worker conn may carry before the
// pool starts opening another one. Set to 1 so the pool fans out to maxConn
// under load *before* stacking streams: concentrating every player on a single
// TCP connection makes that connection a shared point of failure, which is
// exactly how a whole server's worth of players used to drop at once.
const StreamsBeforeGrowing = 1

// WorkerPool manages up to maxConn worker connections and allocates streams
// using the least-loaded strategy (PROTOCOL.md §7.1).
type WorkerPool struct {
	client  *Client
	maxConn int

	mu      sync.Mutex
	cond    *sync.Cond
	conns   []*WorkerConn
	dialing int   // dials currently in flight (foreground + background)
	dialGen uint64
	dialErr error // most recent dial failure
}

func newWorkerPool(c *Client, maxConn int) *WorkerPool {
	p := &WorkerPool{client: c, maxConn: maxConn}
	p.cond = sync.NewCond(&p.mu)
	return p
}

// Allocate returns a worker conn and a fresh stream id to place a new stream on.
//
// A dial is never performed while holding p.mu: session establishment involves
// network I/O, and holding the pool lock across it would park every other
// player behind one unresponsive hub. When the pool is empty exactly one caller
// dials and the rest wait on the condition variable; when the pool is merely
// below maxConn, growth happens in the background and the caller is served
// immediately by an existing conn.
func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
	p.mu.Lock()
	for {
		best, bestCount := p.leastLoadedLocked()
		if best != nil {
			p.maybeGrowLocked(bestCount)
			p.mu.Unlock()
			return best, best.newSid(), nil
		}
		if p.dialing > 0 {
			// Someone is already dialing the first conn; wait for it rather
			// than piling up redundant connections.
			gen := p.dialGen
			p.cond.Wait()
			if len(p.conns) == 0 && p.dialGen != gen && p.dialErr != nil {
				err := p.dialErr
				p.mu.Unlock()
				return nil, 0, err
			}
			continue
		}
		p.dialing++
		p.mu.Unlock()
		wc, err := p.dialWorker()
		p.mu.Lock()
		p.dialing--
		p.dialGen++
		p.dialErr = err
		if err != nil {
			p.cond.Broadcast()
			p.mu.Unlock()
			return nil, 0, err
		}
		p.conns = append(p.conns, wc)
		p.cond.Broadcast()
	}
}

// leastLoadedLocked returns the worker conn carrying the fewest streams.
func (p *WorkerPool) leastLoadedLocked() (*WorkerConn, int) {
	var best *WorkerConn
	bestCount := 0
	for _, wc := range p.conns {
		n := wc.streamCount()
		if best == nil || n < bestCount {
			best = wc
			bestCount = n
		}
	}
	return best, bestCount
}

// maybeGrowLocked opens one more worker conn in the background when the pool is
// below maxConn and the least-loaded conn is already carrying streams. The
// caller does not wait for it: it keeps using the conn it already has, and the
// new one picks up subsequent players.
func (p *WorkerPool) maybeGrowLocked(bestCount int) {
	if bestCount < StreamsBeforeGrowing {
		return
	}
	if len(p.conns)+p.dialing >= p.maxConn {
		if bestCount > SaturationThreshold {
			log.Printf("worker pool at maxConn=%d with %d streams on the least-loaded conn", p.maxConn, bestCount)
		}
		return
	}
	p.dialing++
	go func() {
		wc, err := p.dialWorker()
		var surplus *WorkerConn
		p.mu.Lock()
		p.dialing--
		p.dialGen++
		p.dialErr = err
		switch {
		case err != nil:
		case len(p.conns) < p.maxConn:
			p.conns = append(p.conns, wc)
		default:
			surplus = wc // raced with another dial
		}
		p.cond.Broadcast()
		p.mu.Unlock()
		if err != nil {
			log.Printf("worker pool: background dial failed: %v", err)
		}
		if surplus != nil {
			_ = surplus.fc.Close()
		}
	}()
}

// dialWorker establishes one worker conn. It must be called without p.mu held.
func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
	sess, err := p.client.dialSession(MagicWorker)
	if err != nil {
		return nil, err
	}
	wc := &WorkerConn{
		pool:        p,
		fc:          sess.fc,
		sendWndInit: sess.peerWnd,
		recvWndInit: p.client.streamWnd,
		streams:     make(map[int]*Stream),
		nextSid:     1,
		done:        make(chan struct{}),
	}
	wc.lastPong.Store(time.Now().UnixMilli())
	go wc.readLoop()
	if sess.heartbeat {
		go wc.heartbeatLoop(p.client.cfg.pingInterval(), p.client.cfg.heartbeatTimeout())
	} else {
		log.Printf("worker conn: hub does not support the mux heartbeat; " +
			"a silently dropped path will only be caught by TCP keepalive")
	}
	log.Printf("opened worker conn (send window %d, recv window %d, heartbeat %v)",
		wc.sendWndInit, wc.recvWndInit, sess.heartbeat)
	return wc, nil
}

func (p *WorkerPool) count() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.conns)
}

func (p *WorkerPool) remove(wc *WorkerConn) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for i, c := range p.conns {
		if c == wc {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			// A waiter parked on an empty pool must re-evaluate: it may now
			// need to dial rather than keep waiting for this conn.
			p.cond.Broadcast()
			return
		}
	}
}

func (p *WorkerPool) closeAll() {
	p.mu.Lock()
	conns := append([]*WorkerConn(nil), p.conns...)
	p.mu.Unlock()
	for _, wc := range conns {
		_ = wc.fc.Close()
	}
}

// WorkerConn is one multiplexed worker connection to the hub.
type WorkerConn struct {
	pool *WorkerPool
	fc   *wire.FramedConn

	sendWndInit int // hub's advertised per-stream receive window (our send budget)
	recvWndInit int // our advertised per-stream receive window (bounds each recv queue)

	done     chan struct{} // closed when readLoop exits
	lastPong atomic.Int64  // unix ms of the most recent PONG

	mu      sync.Mutex
	streams map[int]*Stream
	nextSid int
}

// heartbeatLoop proves the worker conn is still carrying frames end to end. TCP
// alone cannot tell us: a middlebox that drops an established flow (conntrack
// expiry, firewall state loss) sends no FIN or RST, so the read loop would park
// forever, the dead conn would stay in the pool, and every player routed to it
// would silently fail until the process restarted.
func (wc *WorkerConn) heartbeatLoop(interval, timeout time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-wc.done:
			return
		case <-ticker.C:
			if silent := time.Since(time.UnixMilli(wc.lastPong.Load())); silent > timeout {
				log.Printf("worker conn silent for %s; dropping it and its %d stream(s)",
					silent.Round(time.Second), wc.streamCount())
				_ = wc.fc.Close() // readLoop unblocks and tears everything down
				return
			}
			msg := wire.NewWriter().U8(MuxPing).VarInt(MuxCtlSid).I64(time.Now().UnixMilli()).Out()
			if err := wc.fc.WriteFrame(msg); err != nil {
				return
			}
		}
	}
}

func (wc *WorkerConn) streamCount() int {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	return len(wc.streams)
}

func (wc *WorkerConn) newSid() int {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	sid := wc.nextSid
	wc.nextSid++
	return sid
}

func (wc *WorkerConn) registerStream(sid int, st *Stream) {
	wc.mu.Lock()
	wc.streams[sid] = st
	wc.mu.Unlock()
}

func (wc *WorkerConn) getStream(sid int) *Stream {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	return wc.streams[sid]
}

func (wc *WorkerConn) removeStream(sid int) *Stream {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	st := wc.streams[sid]
	delete(wc.streams, sid)
	return st
}

// readLoop dispatches inbound mux frames. It must never block on a stream's
// destination: DATA is only enqueued (the per-stream writeLoop does the actual
// destination writes), so one slow destination cannot stall other streams.
func (wc *WorkerConn) readLoop() {
	for {
		payload, err := wc.fc.ReadFrame()
		if err != nil {
			break
		}
		r := wire.NewReader(payload)
		ftype, err := r.U8()
		if err != nil {
			continue
		}
		sid, err := r.VarInt()
		if err != nil {
			continue
		}
		switch ftype {
		case MuxData:
			if st := wc.getStream(sid); st != nil {
				st.deliverFromHub(r.Remaining())
			}
		case MuxWnd:
			if delta, err := r.VarInt(); err == nil && delta > 0 {
				if st := wc.getStream(sid); st != nil {
					st.grantSendWnd(delta)
				}
			}
		case MuxFin:
			// Graceful: drain what is already queued to the destination first.
			if st := wc.removeStream(sid); st != nil {
				st.gracefulFin()
			}
		case MuxRst:
			if st := wc.removeStream(sid); st != nil {
				st.teardown(false)
			}
		case MuxPing:
			nonce, _ := r.I64()
			_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxPong).VarInt(MuxCtlSid).I64(nonce).Out())
		case MuxPong:
			wc.lastPong.Store(time.Now().UnixMilli())
		default:
			log.Printf("worker: unknown mux type %d", ftype)
		}
	}
	// Connection lost: tear down all streams and drop from pool.
	close(wc.done)
	wc.pool.remove(wc)
	wc.mu.Lock()
	streams := make([]*Stream, 0, len(wc.streams))
	for _, st := range wc.streams {
		streams = append(streams, st)
	}
	wc.streams = make(map[int]*Stream)
	wc.mu.Unlock()
	for _, st := range streams {
		st.teardown(false)
	}
}

func (wc *WorkerConn) sendSyn(sid int, cid []byte) error {
	return wc.fc.WriteFrame(wire.NewWriter().U8(MuxSyn).VarInt(sid).Bytes(cid).Out())
}

func (wc *WorkerConn) sendData(sid int, data []byte) error {
	return wc.fc.WriteFrame(wire.NewWriter().U8(MuxData).VarInt(sid).Bytes(data).Out())
}

func (wc *WorkerConn) sendFin(sid int) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxFin).VarInt(sid).Out())
}

func (wc *WorkerConn) sendRst(sid int) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxRst).VarInt(sid).Out())
}

func (wc *WorkerConn) sendWndUpdate(sid, delta int) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxWnd).VarInt(sid).VarInt(delta).Out())
}

// Stream bridges one player (via the hub) to one destination connection.
//
// Data from the hub is queued and written to the destination by a dedicated
// writeLoop goroutine. The queue is bounded by the advertised receive window —
// the hub never sends more un-credited bytes, so overflow is a protocol
// violation and resets the stream.
type Stream struct {
	wc      *WorkerConn
	sid     int
	cid     []byte
	mapping Mapping
	srcIP   string
	srcPort int
	vel     *velocityForwarder // non-nil when the mapping sets velocitySecret

	mu         sync.Mutex
	cond       *sync.Cond
	dest       net.Conn
	connected  bool
	closed     bool
	finPending bool     // hub sent FIN; close the destination once the queue drains
	q          []qentry // hub/local -> destination, waiting for writeLoop
	qBytes     int      // hub bytes only: bounds the peer against its window
	sendWnd    int      // flow control: budget for destination -> hub DATA
	consumed   int      // flow control: drained bytes not yet credited back to the hub
}

// qentry is one queued write towards the destination. Only hub-originated
// entries take part in flow control; locally injected bytes (the velocity
// login response) are neither counted against the hub's window nor credited
// back when drained.
type qentry struct {
	data    []byte
	fromHub bool
}

func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
	s := &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port, sendWnd: wc.sendWndInit}
	if m.VelocitySecret != "" {
		s.vel = newVelocityForwarder(m.VelocitySecret, ip)
	}
	s.cond = sync.NewCond(&s.mu)
	return s
}

// run dials the destination, optionally writes the PROXY v2 header, then pumps
// destination -> hub (respecting the stream send window when negotiated).
func (s *Stream) run() {
	dest, err := net.DialTimeout("tcp", s.mapping.Destination, 10*time.Second)
	if err != nil {
		log.Printf("stream %d: dial %s failed: %v", s.sid, s.mapping.Destination, err)
		s.wc.removeStream(s.sid)
		s.wc.sendRst(s.sid)
		s.teardown(false)
		return
	}
	if tcp, ok := dest.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
	}

	if s.mapping.ProxyProtocol {
		if hdr := s.buildProxyHeader(dest); hdr != nil {
			if _, err := dest.Write(hdr); err != nil {
				log.Printf("stream %d: proxy header write: %v", s.sid, err)
			}
		}
	}

	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		_ = dest.Close()
		return
	}
	s.dest = dest
	s.connected = true
	if s.finPending {
		// The hub FIN'd while we were still dialing, so gracefulFin could not
		// arm the drain deadline (there was no destination yet). Arm it now,
		// otherwise writeLoop can block on an unresponsive destination forever.
		_ = dest.SetWriteDeadline(time.Now().Add(finDrainTimeout))
	}
	s.cond.Broadcast() // wake writeLoop: queued hub bytes can flow now
	s.mu.Unlock()

	// destination -> hub
	buf := make([]byte, DataChunkSize)
	for {
		n, err := dest.Read(buf)
		if n > 0 {
			chunk := buf[:n]
			if s.vel != nil && !s.vel.Passthrough() {
				fwd, inject := s.vel.ProcessS2C(chunk)
				if len(inject) > 0 {
					s.injectToDest(inject)
				}
				chunk = fwd
			}
			if !s.sendToHub(chunk) {
				break
			}
		}
		if err != nil {
			break
		}
	}
	s.teardown(true)
}

// sendToHub forwards destination bytes to the hub in DATA frames of at most
// DataChunkSize, honoring the stream send window. Returns false once the
// stream closed or the worker conn failed.
func (s *Stream) sendToHub(data []byte) bool {
	for len(data) > 0 {
		n := len(data)
		if n > DataChunkSize {
			n = DataChunkSize
		}
		if !s.acquireSendWnd(n) {
			return false
		}
		if err := s.wc.sendData(s.sid, data[:n]); err != nil {
			return false
		}
		data = data[n:]
	}
	return true
}

// writeLoop is the only writer to the destination. It drains the receive queue,
// credits the hub as bytes land on the destination socket, and performs the
// deferred graceful close when a FIN arrived with data still queued.
func (s *Stream) writeLoop() {
	for {
		s.mu.Lock()
		for !s.closed && (!s.connected || (len(s.q) == 0 && !s.finPending)) {
			s.cond.Wait()
		}
		if s.closed {
			s.mu.Unlock()
			return
		}
		if len(s.q) == 0 { // finPending and fully drained
			s.mu.Unlock()
			s.teardown(false)
			return
		}
		e := s.q[0]
		s.q = s.q[1:]
		if e.fromHub {
			s.qBytes -= len(e.data)
		}
		dest := s.dest
		s.mu.Unlock()

		if _, err := dest.Write(e.data); err != nil {
			s.teardown(true)
			return
		}
		if e.fromHub {
			s.credit(len(e.data))
		}
	}
}

// deliverFromHub enqueues hub bytes for the destination. Called from the worker
// readLoop; it never blocks — a peer that exceeds the advertised window is a
// protocol violator and gets the stream reset.
func (s *Stream) deliverFromHub(data []byte) {
	if s.vel != nil {
		s.vel.ObserveC2S(data) // observation only; bytes still forwarded verbatim
	}
	s.mu.Lock()
	if s.closed || s.finPending {
		s.mu.Unlock()
		return
	}
	if s.qBytes+len(data) > s.wc.recvWndInit {
		s.mu.Unlock()
		log.Printf("stream %d: peer exceeded flow-control window; resetting", s.sid)
		s.wc.removeStream(s.sid)
		s.wc.sendRst(s.sid)
		s.teardown(false)
		return
	}
	s.q = append(s.q, qentry{data: data, fromHub: true})
	s.qBytes += len(data)
	s.cond.Broadcast()
	s.mu.Unlock()
}

// injectToDest queues locally generated bytes (the velocity login response)
// for the destination, outside flow-control accounting.
func (s *Stream) injectToDest(data []byte) {
	s.mu.Lock()
	if s.closed || s.finPending {
		s.mu.Unlock()
		return
	}
	s.q = append(s.q, qentry{data: data})
	s.cond.Broadcast()
	s.mu.Unlock()
}

// acquireSendWnd blocks until the stream may send n more bytes to the hub.
// Returns false if the stream closed while waiting.
func (s *Stream) acquireSendWnd(n int) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	for !s.closed && s.sendWnd < n {
		s.cond.Wait()
	}
	if s.closed {
		return false
	}
	s.sendWnd -= n
	return true
}

func (s *Stream) grantSendWnd(delta int) {
	s.mu.Lock()
	s.sendWnd += delta
	s.cond.Broadcast()
	s.mu.Unlock()
}

// credit accounts bytes drained to the destination and grants the hub more
// window once half of our receive window has been consumed.
func (s *Stream) credit(n int) {
	s.mu.Lock()
	s.consumed += n
	if s.closed || s.consumed*2 < s.wc.recvWndInit {
		s.mu.Unlock()
		return
	}
	delta := s.consumed
	s.consumed = 0
	s.mu.Unlock()
	s.wc.sendWndUpdate(s.sid, delta)
}

func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
	srcIP := net.ParseIP(s.srcIP)
	if srcIP == nil {
		return nil
	}
	dstTCP, ok := dest.RemoteAddr().(*net.TCPAddr)
	if !ok {
		return nil
	}
	return BuildProxyV2(srcIP, s.srcPort, dstTCP.IP, dstTCP.Port)
}

// finDrainTimeout bounds how long a FIN'd stream may keep draining its queue
// into the destination, so a dead destination cannot hold the stream forever.
const finDrainTimeout = 10 * time.Second

// gracefulFin marks the hub-side close; writeLoop finishes the queue and then
// tears the stream down.
func (s *Stream) gracefulFin() {
	s.mu.Lock()
	if s.closed || s.finPending {
		s.mu.Unlock()
		return
	}
	s.finPending = true
	if s.dest != nil {
		_ = s.dest.SetWriteDeadline(time.Now().Add(finDrainTimeout))
	}
	s.cond.Broadcast()
	s.mu.Unlock()
}

// teardown closes the stream immediately; notifyHub sends a FIN when true.
// Idempotent; wakes every goroutine parked on the stream.
func (s *Stream) teardown(notifyHub bool) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	dest := s.dest
	s.cond.Broadcast()
	s.mu.Unlock()

	if dest != nil {
		_ = dest.Close()
	}
	s.wc.removeStream(s.sid)
	if notifyHub {
		s.wc.sendFin(s.sid)
	}
}