fix

client/worker.go

@@ -4,30 +4,87 @@ 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
-	conns []*WorkerConn
+	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 {
-	return &WorkerPool{client: c, maxConn: maxConn}
+	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()
-	defer p.mu.Unlock()
+	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 {
@@ -37,39 +94,74 @@ func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
 			bestCount = n
 		}
 	}
-
-	needNew := best == nil || (bestCount > SaturationThreshold && len(p.conns) < p.maxConn)
-	if needNew {
-		wc, err := p.dialWorker()
-		if err != nil {
-			if best == nil {
-				return nil, 0, err
-			}
-			log.Printf("worker dial failed, reusing existing conn: %v", err)
-		} else {
-			p.conns = append(p.conns, wc)
-			best = wc
-		}
-	}
-	return best, best.newSid(), nil
+	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) {
-	fc, peerWnd, err := p.client.dialSession(MagicWorker)
+	sess, err := p.client.dialSession(MagicWorker)
 	if err != nil {
 		return nil, err
 	}
 	wc := &WorkerConn{
 		pool:        p,
-		fc:          fc,
-		sendWndInit: peerWnd,
+		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()
-	log.Printf("opened worker conn (#%d in pool, send window %d, recv window %d)",
-		len(p.conns)+1, wc.sendWndInit, wc.recvWndInit)
+	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
 }
 
@@ -85,6 +177,9 @@ func (p *WorkerPool) remove(wc *WorkerConn) {
 	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
 		}
 	}
@@ -107,11 +202,41 @@ type WorkerConn struct {
 	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()
@@ -184,11 +309,17 @@ func (wc *WorkerConn) readLoop() {
 			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))
@@ -202,8 +333,8 @@ func (wc *WorkerConn) readLoop() {
 	}
 }
 
-func (wc *WorkerConn) sendSyn(sid int, cid []byte) {
-	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxSyn).VarInt(sid).Bytes(cid).Out())
+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 {
@@ -235,6 +366,7 @@ type Stream struct {
 	mapping Mapping
 	srcIP   string
 	srcPort int
+	vel     *velocityForwarder // non-nil when the mapping sets velocitySecret
 
 	mu         sync.Mutex
 	cond       *sync.Cond
@@ -242,14 +374,26 @@ type Stream struct {
 	connected  bool
 	closed     bool
 	finPending bool     // hub sent FIN; close the destination once the queue drains
-	q          [][]byte // hub -> destination, waiting for writeLoop
-	qBytes     int
-	sendWnd    int // flow control: budget for destination -> hub DATA
-	consumed   int // flow control: drained bytes not yet credited back to the hub
+	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
 }
@@ -285,6 +429,12 @@ func (s *Stream) run() {
 	}
 	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()
 
@@ -293,10 +443,15 @@ func (s *Stream) run() {
 	for {
 		n, err := dest.Read(buf)
 		if n > 0 {
-			if !s.acquireSendWnd(n) {
-				break
+			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 werr := s.wc.sendData(s.sid, buf[:n]); werr != nil {
+			if !s.sendToHub(chunk) {
 				break
 			}
 		}
@@ -307,6 +462,26 @@ func (s *Stream) run() {
 	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.
@@ -325,17 +500,21 @@ func (s *Stream) writeLoop() {
 			s.teardown(false)
 			return
 		}
-		data := s.q[0]
+		e := s.q[0]
 		s.q = s.q[1:]
-		s.qBytes -= len(data)
+		if e.fromHub {
+			s.qBytes -= len(e.data)
+		}
 		dest := s.dest
 		s.mu.Unlock()
 
-		if _, err := dest.Write(data); err != nil {
+		if _, err := dest.Write(e.data); err != nil {
 			s.teardown(true)
 			return
 		}
-		s.credit(len(data))
+		if e.fromHub {
+			s.credit(len(e.data))
+		}
 	}
 }
 
@@ -343,6 +522,9 @@ func (s *Stream) writeLoop() {
 // 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()
@@ -356,12 +538,25 @@ func (s *Stream) deliverFromHub(data []byte) {
 		s.teardown(false)
 		return
 	}
-	s.q = append(s.q, data)
+	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 {