tcp: simplified Policy implementation based on @MDr164 (#190)

* begin prepping policy refactor manually * implement tcp.Policy and refactor rto to use it * fix CI * chatting with claude gave me idea to reformulate Policy * tcp.Policy: add newTransmitLimit output * merge with main and fix failing tests * remove fix.patch * answer my own comments
2026-09-07 23:39:04 +00:00 · 2026-09-07 10:46:24 -03:00
parent d7f3924489
commit 07afcfd924
14 changed files with 1518 additions and 806 deletions
@@ -31,14 +31,8 @@ type Handler struct {
 	optcodec OptionCodec
 	// reasm tracks out-of-order segments staged in bufRx's free region. Always
 	// enabled once buffers are set (see [Handler.SetBuffers]).
-	reasm reassembly
-	// loss is the optional packet-loss recovery algorithm (RTO, congestion
-	// control, ...) driven from the rx/tx hooks. nil disables loss recovery, in
-	// which case the connection behaves as if no timing existed. nanotime is the
-	// monotonic time source (nanoseconds) passed to those hooks; it is non-nil
-	// whenever loss is non-nil (enforced by [Conn.Configure]). See [LossRecovery].
-	loss     LossRecovery
-	nanotime func() int64
+	reasm  reassembly
+	policy Policy

 	closing    bool
 	shutdownRx bool
@@ -79,27 +73,16 @@ func (h *Handler) SetBuffers(txbuf, rxbuf []byte, packets int) error {
 	return h.bufTx.ResetOrReuse(txbuf, packets, 0)
 }

-// SetLossRecovery installs the packet-loss recovery algorithm and the monotonic
-// time source (nanoseconds, the func() int64 convention used across lneto) that
-// drives it. The tcp package keeps no clock of its own; nanotime is read only to
-// stamp the rx/tx hooks (see [LossRecovery]). Passing loss == nil disables loss
-// recovery. It should be set before the connection is opened.
-func (h *Handler) SetLossRecovery(loss LossRecovery, nanotime func() int64) {
-	h.loss = loss
-	h.nanotime = nanotime
+// SetPolicy installs the transmit-steering algorithm. nil disables it.
+// It should be set before the connection is opened. See [Policy].
+func (h *Handler) SetPolicy(policy Policy) {
+	h.policy = policy
 }
+func (h *Handler) policyEnabled() bool { return h.policy != nil }

-func (h *Handler) lossEnabled() bool { return h.loss != nil }
-
-// NextDeadline returns the monotonic-nanosecond instant at which the connection
-// must next be serviced by a transmit attempt (e.g. an RTO expiry), or 0 when
-// there is no deadline or no loss recovery is configured. See [LossRecovery].
-func (h *Handler) NextDeadline() int64 {
-	if h.loss == nil {
-		return 0
-	}
-	return h.loss.NextDeadline()
-}
+// ControlBlock returns the state machine underlying the Handler, mainly so a
+// [Policy] can read the sequence spaces. Not for modification.
+func (h *Handler) ControlBlock() *ControlBlock { return &h.scb }

 // LocalPort returns the local port of the connection. Returns 0 if the connection is closed and uninitialized.
 func (h *Handler) LocalPort() uint16 {
@@ -165,16 +148,15 @@ func (h *Handler) reset(localPort, remotePort uint16, iss Value) {
 		shutdownRx: false,
 		// Persist configuration across reopen:
 		validator: h.validator,
-		loss:      h.loss,
-		nanotime:  h.nanotime,
+		policy:    h.policy,
 		logger:    h.logger,
 		// persist memory across repoen:
 		bufTx: h.bufTx,
 		bufRx: h.bufRx,
 		reasm: h.reasm,
 	}
-	if h.lossEnabled() {
-		h.loss.Reset()
+	if h.policyEnabled() {
+		h.policy.Reset()
 	}
 	h.reasm.clear() // preserve metadata capacity across reopen, drop held segments.
 	h.bufTx.ResetOrReuse(nil, 0, iss)
@@ -212,9 +194,7 @@ func (h *Handler) Recv(incomingPacket []byte) error {
 		return nil
 	}

-	// Notify loss recovery of the received segment (RTT sampling, timer
-	// management) and let it drop the segment before processing if it asks to.
-	if h.lossEnabled() && !h.loss.PreRx(segIncoming, h.nanotime()).Keep {
+	if h.policyEnabled() && !h.policy.PreRx(h, tfrm) {
 		return nil
 	}

@@ -245,6 +225,9 @@ func (h *Handler) Recv(incomingPacket []byte) error {
 	if prevState != h.scb.State() {
 		h.info("tcp.Handler:rx-statechange", slog.Uint64("port", uint64(h.localPort)), slog.String("old", prevState.String()), slog.String("new", h.scb.State().String()), slog.String("rxflags", segIncoming.Flags.String()))
 	}
+	if h.policyEnabled() {
+		h.policy.PostRx(h, prevState, tfrm)
+	}
 	if segIncoming.DATALEN != 0 && h.shutdownRx && (h.scb.State() == StateFinWait1 || h.scb.State() == StateFinWait2) {
 		// soypat/lneto#50: the application is done in both directions — read side
 		// shut down (CloseRead) and our FIN sent (Close) — so inbound data has no
@@ -380,16 +363,31 @@ func (h *Handler) Send(b []byte) (int, error) {
 	if h.IsTxOver() {
 		return 0, net.ErrClosed
 	}
-	var now int64
-	if h.lossEnabled() {
-		now = h.nanotime()
-		if h.loss.PreTx(now).RetransmitAll {
-			// Go-back-N retransmission directed by loss recovery: rewind the
-			// send sequence and transmit buffer so unacknowledged data is resent
-			// from snd.UNA. Done before the early short-circuit below so an
-			// expired RTO retransmits even with no new data queued.
-			h.scb.RetransmitAll()
-			h.bufTx.RetransmitFromUNA()
+	tfrm, err := NewFrame(b)
+	if err != nil {
+		return 0, err
+	}
+	offset := uint8(5)
+	txLimit := TransmitUnlimited
+	if h.policyEnabled() {
+		// Hand the Policy a defined frame: zeroed header at the minimum offset.
+		// It may append options and raise the offset, which is read back below.
+		tfrm.ClearHeader()
+		tfrm.SetOffsetAndFlags(offset, 0)
+		limit, rtxFrom, doRtx := h.policy.PreTx(h, tfrm)
+		txLimit = limit
+		if limit == 0 {
+			h.info("tcp.Policy:newTxLimit=0") // Can cause headaches for users.
+		}
+		if doRtx && h.scb.RetransmitFrom(rtxFrom) {
+			// Retransmission directed by the Policy: rewind the transmit buffer
+			// to match the send sequence so unacknowledged data is resent. Done
+			// before the early short-circuit below so an expired RTO
+			// retransmits even with no new data queued.
+			h.bufTx.RetransmitFrom(rtxFrom)
+		}
+		if o, _ := tfrm.OffsetAndFlags(); o > offset && int(o)*4 < len(b) {
+			offset = o
 		}
 	}
 	awaitingSyn := h.AwaitingSynSend()
@@ -405,29 +403,27 @@ func (h *Handler) Send(b []byte) (int, error) {
 		// Early nop short circuit.
 		return 0, nil
 	}
-	tfrm, err := NewFrame(b)
-	if err != nil {
-		return 0, err
-	}
 	if buffered == 0 && h.closing && (h.scb.State() != StateCloseWait || !h.scb.HasPending()) {
 		// If Close called and no more data to be sent, terminate connection.
 		// In CLOSE-WAIT: wait until the pending ACK is sent first, since scb.Close()
 		// overwrites pending with [FIN|ACK] (unlike ESTABLISHED which merges via bitmask).
 		h.closing = false
-		err = h.scb.Close()
+		err := h.scb.Close()
 		if err != nil {
 			h.logerr("tcp.Handler.Close", slog.String("err", errstr(err)), slog.String("state", h.State().String()))
 			h.Abort()
 			return 0, io.EOF
 		}
 	}
-	offset := uint8(5)
-	mss := uint16(len(b) - sizeHeaderTCP)
+	// optHead is where the Handler's own options begin: after the fixed header
+	// and after any options the Policy already wrote, so neither clobbers the other.
+	optHead := int(offset) * 4
+	mss := uint16(len(b) - optHead)
 	var segment Segment
 	if awaitingSyn || requeueControl && h.scb.State() == StateSynSent {
 		// Handling init syn segment.
 		segment = ClientSynSegment(h.bufTx.iss, Size(h.bufRx.Size()))
-		h.optcodec.PutOption16(b[sizeHeaderTCP:], OptMaxSegmentSize, mss)
+		h.optcodec.PutOption16(b[optHead:], OptMaxSegmentSize, mss)
 		offset++
 		if requeueControl {
 			h.info("tcp.Handler:requeue-syn", slog.Uint64("port", uint64(h.localPort)), slog.Uint64("rport", uint64(h.remotePort)))
@@ -439,7 +435,7 @@ func (h *Handler) Send(b []byte) (int, error) {
 			WND:   Size(h.bufRx.Free()),
 			Flags: synack,
 		}
-		h.optcodec.PutOption16(b[sizeHeaderTCP:], OptMaxSegmentSize, mss)
+		h.optcodec.PutOption16(b[optHead:], OptMaxSegmentSize, mss)
 		offset++
 		h.info("tcp.Handler:requeue-synack", slog.Uint64("port", uint64(h.localPort)), slog.Uint64("rport", uint64(h.remotePort)))
 	} else if requeueControl {
@@ -447,17 +443,21 @@ func (h *Handler) Send(b []byte) (int, error) {
 		return 0, nil
 	} else {
 		var ok bool
-		maxPayload := len(b) - sizeHeaderTCP
+		maxPayload := len(b) - optHead
+		if txLimit < Size(maxPayload) && !h.nextSegmentIsRetransmit() {
+			// Policy clamped new data.
+			maxPayload = int(txLimit)
+		}
 		segment, ok = h.scb.PendingSegment(maxPayload)
 		segment.WND = h.recvWindow()
 		if !ok {
 			// No pending control segment or data to send. Yield.
 			return 0, nil
 		} else if segment.Flags == synack {
-			h.optcodec.PutOption16(b[sizeHeaderTCP:], OptMaxSegmentSize, mss)
+			h.optcodec.PutOption16(b[optHead:], OptMaxSegmentSize, mss)
 			offset++
 		} else if segment.DATALEN > 0 {
-			n, err := h.bufTx.MakePacket(b[sizeHeaderTCP:sizeHeaderTCP+segment.DATALEN], segment.SEQ)
+			n, err := h.bufTx.MakePacket(b[optHead:optHead+int(segment.DATALEN)], segment.SEQ)
 			if err != nil {
 				return 0, err
 			}
@@ -474,15 +474,19 @@ func (h *Handler) Send(b []byte) (int, error) {
 	} else if prevState != h.scb.State() && h.logenabled(slog.LevelInfo) {
 		h.info("tcp.Handler:tx-statechange", slog.Uint64("port", uint64(h.localPort)), slog.String("oldState", prevState.String()), slog.String("newState", h.scb.State().String()), slog.String("txflags", segment.Flags.String()))
 	}
-	if h.lossEnabled() {
-		h.loss.PostTx(segment, now)
-	}
 	h.requeueControl = false
 	tfrm.SetSourcePort(h.localPort)
 	tfrm.SetDestinationPort(h.remotePort)
 	tfrm.SetSegment(segment, offset)
 	tfrm.SetUrgentPtr(0)
 	datalen := int(offset)*4 + int(segment.DATALEN)
+	if h.policyEnabled() {
+		// Frame trimmed to what is actually emitted so the Policy's Payload()
+		// is the segment data and nothing more.
+		if sent, err := NewFrame(b[:datalen]); err == nil {
+			h.policy.PostTx(h, sent)
+		}
+	}
 	closedSuccess := prevState == StateTimeWait && segment.Flags.HasAny(FlagACK)
 	if closedSuccess {
 		h.reset(0, 0, 0)
@@ -494,6 +498,27 @@ func (h *Handler) Send(b []byte) (int, error) {
 	return datalen, nil
 }

+// nextSegmentIsRetransmit reports whether the next data segment would resend
+// already-transmitted bytes rather than open new sequence space. Used to let a
+// retransmission through while a [Policy] holds new data back.
+func (h *Handler) nextSegmentIsRetransmit() bool {
+	endSeq, hasSent := h.bufTx.sentEndSeq()
+	return hasSent && h.scb.snd.NXT.LessThan(endSeq)
+}
+
+// NextSegmentSYN returns syn=true if next outgoing segment is a handshake SYN.
+// This method is exported for use by [Policy] implementations to decide handshake-only options (window scale, SACK-permitted, timestamps).
+func (h *Handler) NextSegmentSYN() (syn, ack bool) {
+	state := h.scb.State()
+	if h.AwaitingSynSend() || h.requeueControl && state == StateSynSent {
+		return true, false // SYN initial/requeue.
+	} else if h.requeueControl && state == StateSynRcvd {
+		return true, true // SYNACK requeue.
+	}
+	pending := h.scb.pending[0]
+	return pending.HasAny(FlagSYN), pending.HasAny(FlagACK)
+}
+
 // Write implements [io.Writer] by copying b to a internal buffer to be sent over the network on the next
 // [Handler.Send] call that can send data to remote peer. Use [Handler.Free] to know the maximum length the argument slice can be before erroring.
 func (h *Handler) Write(b []byte) (int, error) {