add credit based mux window

client/worker.go

@@ -55,18 +55,21 @@ func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
 }
 
 func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
-	fc, err := p.client.dialSession(MagicWorker)
+	fc, peerWnd, err := p.client.dialSession(MagicWorker)
 	if err != nil {
 		return nil, err
 	}
 	wc := &WorkerConn{
-		pool:    p,
-		fc:      fc,
-		streams: make(map[int]*Stream),
-		nextSid: 1,
+		pool:        p,
+		fc:          fc,
+		sendWndInit: peerWnd,
+		recvWndInit: p.client.streamWnd,
+		streams:     make(map[int]*Stream),
+		nextSid:     1,
 	}
 	go wc.readLoop()
-	log.Printf("opened worker conn (#%d in pool)", len(p.conns)+1)
+	log.Printf("opened worker conn (#%d in pool, send window %d, recv window %d)",
+		len(p.conns)+1, wc.sendWndInit, wc.recvWndInit)
 	return wc, nil
 }
 
@@ -101,6 +104,9 @@ 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)
+
 	mu      sync.Mutex
 	streams map[int]*Stream
 	nextSid int
@@ -140,6 +146,9 @@ func (wc *WorkerConn) removeStream(sid int) *Stream {
 	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()
@@ -160,9 +169,20 @@ func (wc *WorkerConn) readLoop() {
 			if st := wc.getStream(sid); st != nil {
 				st.deliverFromHub(r.Remaining())
 			}
-		case MuxFin, MuxRst:
+		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.shutdown(false)
+				st.gracefulFin()
+			}
+		case MuxRst:
+			if st := wc.removeStream(sid); st != nil {
+				st.teardown(false)
 			}
 		default:
 			log.Printf("worker: unknown mux type %d", ftype)
@@ -178,7 +198,7 @@ func (wc *WorkerConn) readLoop() {
 	wc.streams = make(map[int]*Stream)
 	wc.mu.Unlock()
 	for _, st := range streams {
-		st.shutdown(false)
+		st.teardown(false)
 	}
 }
 
@@ -198,7 +218,16 @@ 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
@@ -207,25 +236,33 @@ type Stream struct {
 	srcIP   string
 	srcPort int
 
-	mu        sync.Mutex
-	dest      net.Conn
-	connected bool
-	preBuf    []byte
-	closed    bool
+	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          [][]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
 }
 
 func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
-	return &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port}
+	s := &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port, sendWnd: wc.sendWndInit}
+	s.cond = sync.NewCond(&s.mu)
+	return s
 }
 
-// run dials the destination, optionally writes the PROXY v2 header, flushes any
-// buffered hub bytes, then pumps destination -> hub.
+// 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 {
@@ -240,7 +277,6 @@ func (s *Stream) run() {
 		}
 	}
 
-	// Atomically flush pre-connect buffer and enable direct writes.
 	s.mu.Lock()
 	if s.closed {
 		s.mu.Unlock()
@@ -248,18 +284,18 @@ func (s *Stream) run() {
 		return
 	}
 	s.dest = dest
-	if len(s.preBuf) > 0 {
-		_, _ = dest.Write(s.preBuf)
-		s.preBuf = nil
-	}
 	s.connected = true
+	s.cond.Broadcast() // wake writeLoop: queued hub bytes can flow now
 	s.mu.Unlock()
 
 	// destination -> hub
-	buf := make([]byte, 32*1024)
+	buf := make([]byte, DataChunkSize)
 	for {
 		n, err := dest.Read(buf)
 		if n > 0 {
+			if !s.acquireSendWnd(n) {
+				break
+			}
 			if werr := s.wc.sendData(s.sid, buf[:n]); werr != nil {
 				break
 			}
@@ -268,7 +304,99 @@ func (s *Stream) run() {
 			break
 		}
 	}
-	s.shutdown(true)
+	s.teardown(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
+		}
+		data := s.q[0]
+		s.q = s.q[1:]
+		s.qBytes -= len(data)
+		dest := s.dest
+		s.mu.Unlock()
+
+		if _, err := dest.Write(data); err != nil {
+			s.teardown(true)
+			return
+		}
+		s.credit(len(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) {
+	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, data)
+	s.qBytes += len(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 {
@@ -283,28 +411,29 @@ func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
 	return BuildProxyV2(srcIP, s.srcPort, dstTCP.IP, dstTCP.Port)
 }
 
-// deliverFromHub writes bytes coming from the hub to the destination, buffering
-// until the destination connection is established.
-func (s *Stream) deliverFromHub(data []byte) {
+// 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 {
+	if s.closed || s.finPending {
 		s.mu.Unlock()
 		return
 	}
-	if !s.connected {
-		s.preBuf = append(s.preBuf, data...)
-		s.mu.Unlock()
-		return
+	s.finPending = true
+	if s.dest != nil {
+		_ = s.dest.SetWriteDeadline(time.Now().Add(finDrainTimeout))
 	}
-	dest := s.dest
+	s.cond.Broadcast()
 	s.mu.Unlock()
-	if _, err := dest.Write(data); err != nil {
-		s.shutdown(true)
-	}
 }
 
-// shutdown closes the stream; notifyHub sends a FIN to the hub when true.
-func (s *Stream) shutdown(notifyHub bool) {
+// 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()
@@ -312,6 +441,7 @@ func (s *Stream) shutdown(notifyHub bool) {
 	}
 	s.closed = true
 	dest := s.dest
+	s.cond.Broadcast()
 	s.mu.Unlock()
 
 	if dest != nil {