impl connection recovery

2026-08-15 17:31:35 +08:00
parent e63a34d53a
commit 7bd84af48d
33 changed files with 3858 additions and 200 deletions
@@ -17,16 +17,24 @@ import (
 // exactly how a whole server's worth of players used to drop at once.
 const StreamsBeforeGrowing = 1

+// allocateAttempts bounds how many times a caller retries Allocate when the
+// conn it was handed dies before the stream could be registered on it. The race
+// is narrow and each retry picks a different conn, so a small bound is enough;
+// an unbounded loop would spin against a hub that is refusing every connection.
+const allocateAttempts = 3
+
 // 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

+	connSeq atomic.Int64 // conn ids, for log correlation
+
 	mu      sync.Mutex
 	cond    *sync.Cond
 	conns   []*WorkerConn
-	dialing int   // dials currently in flight (foreground + background)
+	dialing int // dials currently in flight (foreground + background)
 	dialGen uint64
 	dialErr error // most recent dial failure
 }
@@ -147,11 +155,16 @@ func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
 		pool:        p,
 		fc:          sess.fc,
 		sendWndInit: sess.peerWnd,
-		recvWndInit: p.client.streamWnd,
+		resume:      sess.resume,
+		grace:       p.client.resumeGrace(sess.hubGrace),
+		id:          int(p.connSeq.Add(1)),
 		streams:     make(map[int]*Stream),
 		nextSid:     1,
 		done:        make(chan struct{}),
 	}
+	if p.client.statsOn() {
+		wc.stats = &connStats{opened: time.Now()}
+	}
 	wc.lastPong.Store(time.Now().UnixMilli())
 	go wc.readLoop()
 	if sess.heartbeat {
@@ -160,8 +173,8 @@ func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
 		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)
+	log.Printf("opened worker conn (send window %d, recv window %d, heartbeat %v, resume %v)",
+		wc.sendWndInit, p.client.streamWnd, sess.heartbeat, wc.resume)
 	return wc, nil
 }

@@ -195,19 +208,34 @@ func (p *WorkerPool) closeAll() {
 }

 // WorkerConn is one multiplexed worker connection to the hub.
+//
+// Only genuinely per-connection state lives here. Client-wide values (the
+// shaper, the advertised receive window, the DATA chunk cap) belong to Client:
+// reading them through a connection pointer would make every such read a
+// re-parenting hazard once a stream can migrate between conns.
 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)
+
+	// resume is whether this conn negotiated stream resumption, and grace how
+	// long a stream parked from it may keep trying to reattach. Both are
+	// per-conn: a reattach may land on a different (or restarted) hub, so the
+	// flag must be re-checked on the conn that will carry the RESUME.
+	resume bool
+	grace  time.Duration

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

+	id    int        // for log correlation only
+	stats *connStats // nil unless diagnostics are enabled
+
 	mu      sync.Mutex
 	streams map[int]*Stream
 	nextSid int
+	closed  bool // readLoop has exited; registerStream must refuse
 }

 // heartbeatLoop proves the worker conn is still carrying frames end to end. TCP
@@ -251,10 +279,23 @@ func (wc *WorkerConn) newSid() int {
 	return sid
 }

-func (wc *WorkerConn) registerStream(sid int, st *Stream) {
+// registerStream publishes a stream in the conn's table, or reports false if
+// the conn has already died.
+//
+// The check is not advisory. Allocate hands out a (conn, sid) pair under the
+// pool lock, and the conn's readLoop can exit before the caller gets here — it
+// has then already swapped the stream map, so a blind insert would land in a map
+// nothing iterates and the stream would never be torn down. That normally hides
+// behind a failing SYN, but not on a half-open conn whose readLoop died on a
+// framing error while the socket is still writable. Callers must re-Allocate.
+func (wc *WorkerConn) registerStream(sid int, st *Stream) bool {
 	wc.mu.Lock()
+	defer wc.mu.Unlock()
+	if wc.closed {
+		return false
+	}
 	wc.streams[sid] = st
-	wc.mu.Unlock()
+	return true
 }

 func (wc *WorkerConn) getStream(sid int) *Stream {
@@ -280,6 +321,9 @@ func (wc *WorkerConn) readLoop() {
 		if err != nil {
 			break
 		}
+		if wc.stats != nil {
+			wc.stats.framesIn.Add(1)
+		}
 		r := wire.NewReader(payload)
 		ftype, err := r.U8()
 		if err != nil {
@@ -306,14 +350,37 @@ func (wc *WorkerConn) readLoop() {
 				st.gracefulFin()
 			}
 		case MuxRst:
+			// The reason is an optional trailing byte; older peers send none.
+			reason := RstUnspecified
+			if b, err := r.U8(); err == nil {
+				reason = int(b)
+			}
 			if st := wc.removeStream(sid); st != nil {
-				st.teardown(false)
+				st.onRst(reason)
+			}
+		case MuxResumeAck:
+			accepted, aerr := r.I64()
+			delivered, derr := r.I64()
+			cid, cerr := r.Bytes(CIDLen)
+			if aerr != nil || derr != nil || cerr != nil {
+				continue
+			}
+			if st := wc.getStream(sid); st != nil {
+				st.deliverResume(resumeResult{accepted: accepted, delivered: delivered, cid: cid})
 			}
 		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())
+			now := time.Now()
+			wc.lastPong.Store(now.UnixMilli())
+			// The probe's nonce is the timestamp we sent, echoed back, so the
+			// round trip is free to measure and nobody was reading it.
+			if wc.stats != nil {
+				if sent, err := r.I64(); err == nil {
+					wc.stats.observeRTT(now.Sub(time.UnixMilli(sent)))
+				}
+			}
 		default:
 			log.Printf("worker: unknown mux type %d", ftype)
 		}
@@ -322,14 +389,22 @@ func (wc *WorkerConn) readLoop() {
 	close(wc.done)
 	wc.pool.remove(wc)
 	wc.mu.Lock()
+	// Marked before the map is swapped, under the same lock, so a concurrent
+	// registerStream either lands in the map we are about to drain or is refused.
+	wc.closed = true
 	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()
+	// Only the tunnel leg died. Where the session negotiated resumption the
+	// destination sockets are kept open and each stream reattaches over a fresh
+	// conn (§7.5); otherwise this is the old, unconditional teardown.
 	for _, st := range streams {
-		st.teardown(false)
+		if !st.park(wc.grace) {
+			st.teardown(false)
+		}
 	}
 }

@@ -338,7 +413,14 @@ func (wc *WorkerConn) sendSyn(sid int, cid []byte) error {
 }

 func (wc *WorkerConn) sendData(sid int, data []byte) error {
-	return wc.fc.WriteFrame(wire.NewWriter().U8(MuxData).VarInt(sid).Bytes(data).Out())
+	err := wc.fc.WriteFrame(wire.NewWriter().U8(MuxData).VarInt(sid).Bytes(data).Out())
+	if wc.stats != nil {
+		wc.stats.framesOut.Add(1)
+		if err != nil {
+			wc.stats.writeErrs.Add(1)
+		}
+	}
+	return err
 }

 func (wc *WorkerConn) sendFin(sid int) {
@@ -353,6 +435,16 @@ func (wc *WorkerConn) sendWndUpdate(sid, delta int) {
 	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxWnd).VarInt(sid).VarInt(delta).Out())
 }

+// leg binds a stream to one worker conn. The two fields are only meaningful
+// together: stream ids are per-conn and restart at 1, so conn A's sid 3 and
+// conn B's sid 3 belong to different players. A torn read across the two would
+// credit, reset or FIN a stranger's stream, so the pair is swapped as one
+// immutable value rather than as two fields.
+type leg struct {
+	wc  *WorkerConn
+	sid int
+}
+
 // 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
@@ -360,24 +452,74 @@ func (wc *WorkerConn) sendWndUpdate(sid, delta int) {
 // 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
+	client  *Client
+	leg     atomic.Pointer[leg]
 	mapping Mapping
 	srcIP   string
 	srcPort int
 	vel     *velocityForwarder // non-nil when the mapping sets velocitySecret

+	share shaperShare   // this stream's position in the egress fair queue
+	done  chan struct{} // closed on teardown; unparks a shaper wait
+
+	// resumable is fixed at creation from the conn's negotiated flag. With it
+	// false none of the bookkeeping below runs and no buffer is ever allocated,
+	// so a client with resumption disabled pays exactly what it used to.
+	resumable bool
+
+	stats *streamStats // nil unless diagnostics are enabled
+
+	// ackedOffset is the running sum of WND deltas received. Credit is granted
+	// only as bytes reach the peer's terminal socket, so it is a sound lower
+	// bound on what has been delivered. Atomic because the worker readLoop
+	// advances it and must never block behind the send path.
+	ackedOffset atomic.Int64
+
+	// sendMu serializes the send path — buffer the chunk, advance the offset,
+	// write the frame — against a reattach's retransmit, so replayed bytes can
+	// never interleave with live ones.
+	//
+	// Deliberately not s.mu: deliverFromHub takes s.mu from the worker readLoop,
+	// and holding s.mu across a WriteFrame would stall frame dispatch for every
+	// other stream on the connection.
+	sendMu sync.Mutex
+	un     unackedBuf // guarded by sendMu
+
 	mu         sync.Mutex
 	cond       *sync.Cond
+	cid        []byte // takeover capability; re-minted by the hub on each resume
 	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
+	parked     bool      // worker conn died; awaiting reattach on a fresh one
+	parkedAt   time.Time // when the current hang began; diagnostics only
+	finPending bool      // hub sent FIN; close the destination once the queue drains
+	finToHub   bool      // destination closed while parked; FIN owed once reattached
+	q          []qentry  // hub/local -> destination, waiting for writeLoop
+	qBytes     int       // hub bytes only: bounds the peer against its window
+	// acceptedOffset counts hub bytes enqueued toward the destination. This, not
+	// "bytes written", is what a reattach reports: acceptance is synchronous and
+	// stable at park time, whereas delivery is signalled asynchronously and goes
+	// silent exactly when the connection dies — which would under-report and make
+	// the hub retransmit bytes the player already has.
+	acceptedOffset int64
+	// deliveredOffset counts hub bytes actually written to the destination.
+	// Distinct from acceptedOffset and needed for a different job: a reattach
+	// replays from what the peer *accepted*, but sizes the window from what it
+	// *delivered*, because the window is a promise about undelivered bytes.
+	deliveredOffset int64
+	sendWnd         int // flow control: budget for destination -> hub DATA
+	consumed        int // flow control: drained bytes not yet credited back to the hub
+	resumeWait      chan resumeResult
+}
+
+// resumeResult is the hub's answer to a RESUME: how far it got in both senses,
+// plus a fresh CID — or the reason the reattach was refused.
+type resumeResult struct {
+	accepted  int64
+	delivered int64
+	cid       []byte
+	err       error
 }

 // qentry is one queued write towards the destination. Only hub-originated
@@ -389,8 +531,13 @@ type qentry struct {
 	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}
+func newStream(c *Client, wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
+	s := &Stream{client: c, cid: cid, mapping: m, srcIP: ip, srcPort: port,
+		resumable: wc.resume, sendWnd: wc.sendWndInit, done: make(chan struct{})}
+	s.leg.Store(&leg{wc: wc, sid: sid})
+	if c.statsOn() {
+		s.stats = &streamStats{opened: time.Now()}
+	}
 	if m.VelocitySecret != "" {
 		s.vel = newVelocityForwarder(m.VelocitySecret, ip)
 	}
@@ -398,14 +545,20 @@ func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port i
 	return s
 }

+// conn returns the stream's current binding. Every caller must take exactly one
+// snapshot and use both fields from it; re-loading mid-operation reintroduces
+// the torn-pair hazard the leg exists to prevent.
+func (s *Stream) conn() *leg { return s.leg.Load() }
+
 // 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)
+		lg := s.conn()
+		log.Printf("stream %d: dial %s failed: %v", lg.sid, s.mapping.Destination, err)
+		lg.wc.removeStream(lg.sid)
+		lg.wc.sendRst(lg.sid)
 		s.teardown(false)
 		return
 	}
@@ -416,7 +569,7 @@ func (s *Stream) run() {
 	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)
+				log.Printf("stream %d: proxy header write: %v", s.conn().sid, err)
 			}
 		}
 	}
@@ -462,26 +615,88 @@ 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.
+// sendToHub forwards destination bytes to the hub in bounded DATA frames,
+// honoring both the stream send window and the client-wide bandwidth cap.
+// 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 n > s.client.chunk {
+			n = s.client.chunk
 		}
+		// Credit first, bandwidth second. The reverse order would charge the
+		// budget for bytes still parked on an exhausted window, so the client
+		// would throttle itself below the configured rate. Holding credit while
+		// waiting for tokens is free — credit is per-stream, and the hub returns
+		// it as it drains data to the player, independent of our pacing.
 		if !s.acquireSendWnd(n) {
 			return false
 		}
-		if err := s.wc.sendData(s.sid, data[:n]); err != nil {
+		var shaperStall stallClock
+		shaperStall.begin(s.stats != nil)
+		if !s.client.shaper.Acquire(&s.share, n, s.done) {
 			return false
 		}
+		if !s.emit(data[:n]) {
+			return false
+		}
+		if s.stats != nil {
+			s.mu.Lock()
+			// Separated from the window stall on purpose: this one says the
+			// configured cap is the binding constraint, and raising maxBandwidth
+			// is the fix. The window stall says the opposite.
+			s.stats.shaperStall += shaperStall.elapsed()
+			s.stats.bytesUp += int64(n)
+			s.mu.Unlock()
+		}
 		data = data[n:]
 	}
 	return true
 }

+// emit records a chunk for possible retransmission and writes it to the current
+// worker conn. Returns false once the stream is finished with.
+//
+// The record is taken first and unconditionally. A write that fails on a dying
+// connection has already spent window and may have put part of the frame on the
+// wire, so the only trustworthy account of what the peer still owes us is the
+// one taken before the attempt. That also makes a failure survivable: while the
+// stream can still be resumed the bytes are already safe, and the reattach
+// replays them from wherever the hub says it got to.
+//
+// Note the old code let a failed write drop the rest of the chunk on the floor —
+// the caller's slice advance sat after the error return.
+func (s *Stream) emit(chunk []byte) bool {
+	s.sendMu.Lock()
+	if s.resumable {
+		// Reclaim what the hub has credited before growing the buffer, so the
+		// outstanding region stays bounded by one window.
+		s.un.advance(s.ackedOffset.Load())
+		s.un.append(chunk)
+	}
+	lg := s.conn()
+	err := lg.wc.sendData(lg.sid, chunk)
+	s.sendMu.Unlock()
+
+	if err == nil {
+		return true
+	}
+	if !s.resumable {
+		return false
+	}
+	// The conn is gone but the stream is not: readLoop parks it and a reattach
+	// replays the buffer. Keep pumping the destination — acquireSendWnd stops us
+	// once a full window is outstanding, so nothing is lost and nothing grows
+	// without bound.
+	return !s.isClosed()
+}
+
+func (s *Stream) isClosed() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.closed
+}
+
 // 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.
@@ -513,6 +728,12 @@ func (s *Stream) writeLoop() {
 			return
 		}
 		if e.fromHub {
+			s.mu.Lock()
+			s.deliveredOffset += int64(len(e.data))
+			if s.stats != nil {
+				s.stats.bytesDown += int64(len(e.data))
+			}
+			s.mu.Unlock()
 			s.credit(len(e.data))
 		}
 	}
@@ -530,16 +751,26 @@ func (s *Stream) deliverFromHub(data []byte) {
 		s.mu.Unlock()
 		return
 	}
-	if s.qBytes+len(data) > s.wc.recvWndInit {
+	if s.qBytes+len(data) > s.client.streamWnd {
 		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)
+		lg := s.conn()
+		log.Printf("stream %d: peer exceeded flow-control window; resetting", lg.sid)
+		lg.wc.removeStream(lg.sid)
+		lg.wc.sendRst(lg.sid)
 		s.teardown(false)
 		return
 	}
 	s.q = append(s.q, qentry{data: data, fromHub: true})
 	s.qBytes += len(data)
+	// Accepted, not delivered: from here the bytes are ours to write, and the
+	// only way we fail to is by destroying the stream — which also ends any
+	// prospect of resuming it. That makes this a sound reattach coordinate.
+	s.acceptedOffset += int64(len(data))
+	if s.stats != nil && s.qBytes > s.stats.qPeak {
+		// How close the receive queue came to the advertised window: near it
+		// means the destination is the slow party.
+		s.stats.qPeak = s.qBytes
+	}
 	s.cond.Broadcast()
 	s.mu.Unlock()
 }
@@ -562,9 +793,18 @@ func (s *Stream) injectToDest(data []byte) {
 func (s *Stream) acquireSendWnd(n int) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	// Timed only when it actually blocks, so a stream that never runs out of
+	// credit never reads the clock. A large windowStall is the signal that the
+	// peer is not draining to its terminal socket — the bottleneck is past the
+	// tunnel, not in it.
+	var stall stallClock
 	for !s.closed && s.sendWnd < n {
+		stall.begin(s.stats != nil)
 		s.cond.Wait()
 	}
+	if s.stats != nil {
+		s.stats.windowStall += stall.elapsed()
+	}
 	if s.closed {
 		return false
 	}
@@ -573,6 +813,11 @@ func (s *Stream) acquireSendWnd(n int) bool {
 }

 func (s *Stream) grantSendWnd(delta int) {
+	// The running sum doubles as the acked offset: the hub grants credit exactly
+	// as bytes reach the player socket, so a byte that has been credited can
+	// never need retransmitting. Advanced without a lock so the worker readLoop
+	// never blocks behind a send in progress.
+	s.ackedOffset.Add(int64(delta))
 	s.mu.Lock()
 	s.sendWnd += delta
 	s.cond.Broadcast()
@@ -581,17 +826,22 @@ func (s *Stream) grantSendWnd(delta int) {

 // credit accounts bytes drained to the destination and grants the hub more
 // window once half of our receive window has been consumed.
+// While parked the grant is only withheld, never dropped: consumed keeps
+// accumulating and a reattach flushes it on the new conn. Resetting it would
+// destroy up to half a window of credit per outage, and after a few flaps the
+// stream would throttle to a crawl.
 func (s *Stream) credit(n int) {
 	s.mu.Lock()
 	s.consumed += n
-	if s.closed || s.consumed*2 < s.wc.recvWndInit {
+	if s.closed || s.parked || s.consumed*2 < s.client.streamWnd {
 		s.mu.Unlock()
 		return
 	}
 	delta := s.consumed
 	s.consumed = 0
 	s.mu.Unlock()
-	s.wc.sendWndUpdate(s.sid, delta)
+	lg := s.conn()
+	lg.wc.sendWndUpdate(lg.sid, delta)
 }

 func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
@@ -629,6 +879,13 @@ func (s *Stream) gracefulFin() {
 // 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) {
+	// A parked stream owes the hub a FIN it cannot send: the only conn it has is
+	// the one that just failed. Keep it alive so the reattach can deliver it and
+	// the player gets a clean disconnect, rather than hanging until the hub's
+	// grace expires.
+	if notifyHub && s.noteFinWhileParked() {
+		return
+	}
 	s.mu.Lock()
 	if s.closed {
 		s.mu.Unlock()
@@ -636,14 +893,23 @@ func (s *Stream) teardown(notifyHub bool) {
 	}
 	s.closed = true
 	dest := s.dest
+	close(s.done) // guarded by the idempotence check above, so exactly once
 	s.cond.Broadcast()
 	s.mu.Unlock()

+	// The stream can no longer be resumed, so the retransmit buffer is dead
+	// weight — up to a full window of it per stream.
+	s.sendMu.Lock()
+	s.un.reset()
+	s.sendMu.Unlock()
+
 	if dest != nil {
 		_ = dest.Close()
 	}
-	s.wc.removeStream(s.sid)
+	lg := s.conn()
+	lg.wc.removeStream(lg.sid)
 	if notifyHub {
-		s.wc.sendFin(s.sid)
+		lg.wc.sendFin(lg.sid)
 	}
+	s.logSummary()
 }