Files
redapricot/client/worker.go
T
iceBear67 da17140583 fix: harden resume/shutdown paths, tighten Intent-18 and PSK handshake handling
Client (Go) — resume correctness
- C1: completeResume now re-parks the stream when replay fails mid-conn-loss.
  parked was cleared before the replay loop, so the dying conn's teardown
  would start a second resumeLoop and the two loops could strand the stream
  with neither alive. Resume stats are counted only after the replay lands.
- C2: RST(ALREADY_BOUND) is retryable instead of terminating the loop. With
  C1 fixed there is never a genuine second attempt, so "already bound" means
  the hub still holds the stream on a half-open conn; the retry waits out
  that bind (bounded by the grace deadline, teardown on expiry) instead of
  returning and leaving the destination socket hung forever.

Client (Go) — shutdown semantics
- C3: Close() sets a closing flag and cancels an internal context; dialSession
  takes a ctx (DialContext + AfterFunc so shutdown aborts in-flight
  handshakes); the worker pool refuses new conns after closeAll (Allocate,
  background growth, cond waiters); serveControl's reconnect loop is gated by
  closing so Close works even when the caller's Start context is not
  cancelled; conn-loss teardown closes streams outright during shutdown
  instead of parking them for a reattach that will never come.

Client (Go) — hygiene
- C4: pingInterval() clamps at the single point a duration is derived, so a
  hand-built Config with PingIntervalMs <= 0 can no longer panic
  time.NewTicker (added DefaultPingIntervalMs).
- E6: shaperStall is sampled right after shaper.Acquire, before the socket
  write, so a hub that is not reading is no longer charged to the bandwidth
  cap in the stats.
- E7: stream log lines now carry conn%d/sid%d (leg.String()), making streams
  traceable across reattaches.
- P5: mirror constants IntentReserved/RegisterOk/RegisterErrPattern added;
  RegisterAck dispatch logs rejection reasons via the named codes.

Hub (Java) + PROTOCOL.md
- P3: Intent 18 replies with a Minecraft status-response packet
  ([Len: VarInt][0x00][JSON: String]) and closes (socket.end, so the write
  always lands) instead of closing silently; documented in PROTOCOL.md §2.
- P4: PSK address check is strict equality with the lowercase hex address;
  an uppercase/case-folded variant is now rejected per PROTOCOL.md §2.
- P7: PROTOCOL.md §5 SessionReady row lists its real fields
  (Flags/RecvWindow/ResumeGraceMs) instead of "(none)".

Verified: go vet, go test -race ./client/..., gradle test, full e2e suite
(twice), resume e2e 3x, plus live probes of the hub with the real client
codec (Intent-18 status reply, strict-lowercase PSK acceptance/rejection).%
2026-08-15 17:47:39 +08:00

967 lines
30 KiB
Go

package client
import (
"errors"
"fmt"
"log"
"net"
"sync"
"sync/atomic"
"time"
"github.com/iceBear67/redapricot/client/wire"
)
// errPoolClosed is returned by Allocate after Close: the pool is shutting down
// and must not start new dials, so a caller (handleControlRequest,
// allocateForResume) gives up rather than wait on a cond no one will satisfy.
var errPoolClosed = errors.New("worker pool closed")
// 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
// 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)
dialGen uint64
dialErr error // most recent dial failure
closed bool // closeAll ran; no new conns may join the pool
}
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 {
if p.closed {
// Close won. No new conn may join the pool, so no stream may be
// placed; waiting on the cond could only be satisfied by a dial we
// must not start.
p.mu.Unlock()
return nil, 0, errPoolClosed
}
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
}
if p.closed {
// Close raced this dial: the conn must not enter the pool. Closing
// it here, under p.mu, is a raw socket close — fine, and it makes
// the shutdown atomic with the pool state.
p.cond.Broadcast()
p.mu.Unlock()
_ = wc.fc.Close()
return nil, 0, errPoolClosed
}
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 p.closed {
return // shutdown; do not start dials nobody will join the pool
}
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 p.closed:
surplus = wc // Close raced this background dial
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.
// The dial runs on the client's context so Close aborts it mid-handshake.
func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
sess, err := p.client.dialSession(p.client.ctx, MagicWorker)
if err != nil {
return nil, err
}
wc := &WorkerConn{
pool: p,
fc: sess.fc,
sendWndInit: sess.peerWnd,
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 {
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, resume %v)",
wc.sendWndInit, p.client.streamWnd, sess.heartbeat, wc.resume)
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()
p.closed = true
conns := append([]*WorkerConn(nil), p.conns...)
// Wake waiters parked in Allocate: the closed flag they re-check is the
// only thing that can release them now that no conn will ever join.
p.cond.Broadcast()
p.mu.Unlock()
for _, wc := range conns {
_ = wc.fc.Close()
}
}
// 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)
// 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
// 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
}
// 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
return true
}
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
}
if wc.stats != nil {
wc.stats.framesIn.Add(1)
}
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:
// 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.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:
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)
}
}
// Connection lost: tear down all streams and drop from pool.
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.
// During Close there is no reattach to come: a stream that parked now would
// hold its destination socket open past shutdown, so teardown instead.
if wc.pool.client.closing.Load() {
for _, st := range streams {
st.teardown(false)
}
return
}
for _, st := range streams {
if !st.park(wc.grace) {
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 {
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) {
_ = 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())
}
// 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
}
// String formats one (conn, sid) snapshot for log correlation. Stream ids
// restart at 1 per conn, so a bare sid cannot be traced across a reattach —
// the conn id is what ties the log lines together.
func (lg *leg) String() string { return fmt.Sprintf("conn%d/sid%d", lg.wc.id, lg.sid) }
// 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 {
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
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
// 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(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)
}
s.cond = sync.NewCond(&s.mu)
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 {
lg := s.conn()
log.Printf("stream %s: dial %s failed: %v", lg, s.mapping.Destination, err)
lg.wc.removeStream(lg.sid)
lg.wc.sendRst(lg.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 %s: proxy header write: %v", s.conn(), 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 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 > 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
}
var shaperStall stallClock
shaperStall.begin(s.stats != nil)
if !s.client.shaper.Acquire(&s.share, n, s.done) {
return false
}
// Sampled here, before the socket write: emit's write to the hub can
// block when the hub is not reading, and charging that time to the
// bandwidth cap would blame maxBandwidth for a hub that is not draining.
waited := shaperStall.elapsed()
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 += waited
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.
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.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))
}
}
}
// 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.client.streamWnd {
s.mu.Unlock()
lg := s.conn()
log.Printf("stream %s: peer exceeded flow-control window; resetting", lg)
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()
}
// 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()
// 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
}
s.sendWnd -= n
return true
}
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()
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.
// 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.parked || s.consumed*2 < s.client.streamWnd {
s.mu.Unlock()
return
}
delta := s.consumed
s.consumed = 0
s.mu.Unlock()
lg := s.conn()
lg.wc.sendWndUpdate(lg.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) {
// 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()
return
}
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()
}
lg := s.conn()
lg.wc.removeStream(lg.sid)
if notifyHub {
lg.wc.sendFin(lg.sid)
}
s.logSummary()
}