Explorar o código

tun: Keep relayed data on the reader transmit path

The sticky gate only drains socket submissions, so moving relay data between reader and staged sender CPUs can reorder packets on independent NIC queues. Keep relayed batches on the reader with separate, shared counter reservations while preserving the direct gate and sender affinity.

Cover gate switches, oversized relay credit, mixed direct placement and stalled shared-tunnel reservations. Add Forgejo build, test, race, crypto, lint and artifact jobs on incus-docker.
Ryan Huber hai 2 días
pai
achega
a784af4026
Modificáronse 6 ficheiros con 307 adicións e 26 borrados
  1. 5 4
      examples/config.yml
  2. 2 0
      flow.go
  3. 12 12
      flow_budget.go
  4. 55 10
      flow_tx.go
  5. 231 0
      flow_tx_ring_test.go
  6. 2 0
      flow_tx_test.go

+ 5 - 4
examples/config.yml

@@ -287,11 +287,12 @@ tun:
   # Small batches, below inline_below TX segments or RX packets (default 256), can skip the workers. On transmit each
   # Small batches, below inline_below TX segments or RX packets (default 256), can skip the workers. On transmit each
   # routine holds one mode, inline on the reader or fanout to the workers and the ordered sender, and switches only
   # routine holds one mode, inline on the reader or fanout to the workers and the ordered sender, and switches only
   # after gate_hold (default 64) reads in a row want the other one. Before going inline it waits for every batch in
   # after gate_hold (default 64) reads in a row want the other one. Before going inline it waits for every batch in
-  # flight to reach the wire, so later packets never overtake earlier ones; flow.tx.gate_switches counts the switches.
-  # Choosing per read instead let small inline sends leave from another CPU than staged ones, and TCP saw reordering.
+  # flight to be submitted to the socket; flow.tx.gate_switches counts the switches. This does not drain NIC queues.
+  # Relay data always stays on the reader, so gate switches and oversized reads cannot move it between NIC queues.
+  # Direct data retains the adaptive gate and independent sender placement.
   # On receive a listener with nothing in flight writes a small batch itself. With multiple routines, inline sends
   # On receive a listener with nothing in flight writes a small batch itself. With multiple routines, inline sends
-  # still take tunnel reservation credit. 0 hands every batch to the workers and the ordered sender or writer; only TX
-  # reads too large for a tunnel's whole credit are still sent by the reader, after it drains.
+  # still take tunnel reservation credit. 0 hands direct TX and all RX batches to the workers and ordered stages;
+  # relayed TX and reads too large for a tunnel's whole credit are still sent by the reader.
   # tx: false keeps transmit inline and stages only receive.
   # tx: false keeps transmit inline and stages only receive.
   # tx_slots and rx_slots bound the batches each routine has in flight per direction, minimum 2. A TX slot holds one
   # tx_slots and rx_slots bound the batches each routine has in flight per direction, minimum 2. A TX slot holds one
   # tun read and its ciphertext, about 1.1 MiB; the default is 8, or workers + 2 if more. An RX slot holds one
   # tun read and its ciphertext, about 1.1 MiB; the default is 8, or workers + 2 if more. An RX slot holds one

+ 2 - 0
flow.go

@@ -18,6 +18,8 @@ import (
 // times 64 KiB with GRO, and the count keeps up to 384 entries in flight in four to six slots: six and 24 MiB at
 // times 64 KiB with GRO, and the count keeps up to 384 entries in flight in four to six slots: six and 24 MiB at
 // listen.batch=64.
 // listen.batch=64.
 // tun.staged.tx_slots and rx_slots override the counts. Small-MSS TX arenas can grow.
 // tun.staged.tx_slots and rx_slots override the counts. Small-MSS TX arenas can grow.
+// Relay data stays on the tun reader, regardless of the TX gate, so it never switches between reader and sender
+// NIC queues. Direct data keeps its adaptive placement; relay and direct batches both share tunnel counter credit.
 //
 //
 // Metrics are registered only for active pipelines. flow.{tx,rx}.batches counts reads handed to the ordered stage and
 // Metrics are registered only for active pipelines. flow.{tx,rx}.batches counts reads handed to the ordered stage and
 // flow.{tx,rx}.inline those the reader handled itself (TX inline also counts oversized reads); their rates show
 // flow.{tx,rx}.inline those the reader handled itself (TX inline also counts oversized reads); their rates show

+ 12 - 12
flow_budget.go

@@ -110,10 +110,10 @@ func (r *txReserver) tryBudget() bool {
 // reader first flushes its own oldest batch, flushes need no credit; with none, another reader holds all outstanding
 // reader first flushes its own oldest batch, flushes need no credit; with none, another reader holds all outstanding
 // work and can flush it independently. So there is no circular wait, even with relay tunnels acquired in opposite
 // work and can flush it independently. So there is no circular wait, even with relay tunnels acquired in opposite
 // orders. The sleep is only for readers with no local work, not the uncontended hot path. Acquisition is
 // orders. The sleep is only for readers with no local work, not the uncontended hot path. Acquisition is
-// all-or-nothing per read, so pressure on one tunnel also delays packets for idle tunnels in the read; splitting
-// would need separate slot ownership and coupled inner/relay reservations. Reader fairness is not guaranteed.
-func (p *stagedTx) waitBudget() bool {
-	if p.tryBudget() {
+// all-or-nothing within each direct or relay batch, so pressure on one tunnel also delays packets for idle tunnels
+// in that batch. Reader fairness is not guaranteed.
+func (p *stagedTx) waitBudget(r *txReserver) bool {
+	if r.tryBudget() {
 		return true
 		return true
 	}
 	}
 	for {
 	for {
@@ -123,7 +123,7 @@ func (p *stagedTx) waitBudget() bool {
 		if !p.flushOldest() {
 		if !p.flushOldest() {
 			time.Sleep(50 * time.Microsecond)
 			time.Sleep(50 * time.Microsecond)
 		}
 		}
-		if p.tryBudget() {
+		if r.tryBudget() {
 			return true
 			return true
 		}
 		}
 	}
 	}
@@ -141,7 +141,7 @@ func (s *txSlot) releaseBudget() {
 // the pipeline, then segments and sends them synchronously with credit for each segment, never exceeding the cap and
 // the pipeline, then segments and sends them synchronously with credit for each segment, never exceeding the cap and
 // never filtering or running connection manager side effects twice. This deliberately slow path avoids a lone-batch
 // never filtering or running connection manager side effects twice. This deliberately slow path avoids a lone-batch
 // exception to the shared invariant.
 // exception to the shared invariant.
-func (p *stagedTx) sendLarge(items []txItem) {
+func (p *stagedTx) sendLarge(r *txReserver, items []txItem) {
 	for p.flushOldest() {
 	for p.flushOldest() {
 	}
 	}
 	for _, it := range items {
 	for _, it := range items {
@@ -151,20 +151,20 @@ func (p *stagedTx) sendLarge(items []txItem) {
 			}
 			}
 			part := it
 			part := it
 			part.pkt, part.segs = tio.Packet{Bytes: seg}, 1
 			part.pkt, part.segs = tio.Packet{Bytes: seg}, 1
-			p.tunnels = p.tunnels[:0]
-			part.counter = p.need(part.hostinfo, part.ci, 1)
+			r.tunnels = r.tunnels[:0]
+			part.counter = r.need(part.hostinfo, part.ci, 1)
 			if part.relayHI != nil {
 			if part.relayHI != nil {
-				part.relayCounter = p.need(part.relayHI, part.relayHI.ConnectionState, 1)
+				part.relayCounter = r.need(part.relayHI, part.relayHI.ConnectionState, 1)
 			}
 			}
-			if !p.waitBudget() {
+			if !p.waitBudget(r) {
 				return nil
 				return nil
 			}
 			}
 			one := [1]txItem{part}
 			one := [1]txItem{part}
-			if reserved := p.reserve(one[:]); len(reserved) != 0 {
+			if reserved := r.reserve(one[:]); len(reserved) != 0 {
 				p.f.encryptTxItem(&reserved[0], p.nb, p.sb)
 				p.f.encryptTxItem(&reserved[0], p.nb, p.sb)
 				p.f.flushSendBatch(p.sb, p.q)
 				p.f.flushSendBatch(p.sb, p.q)
 			}
 			}
-			for _, t := range p.tunnels {
+			for _, t := range r.tunnels {
 				t.ci.txBudget.release(t.n, false)
 				t.ci.txBudget.release(t.n, false)
 			}
 			}
 			return nil
 			return nil

+ 55 - 10
flow_tx.go

@@ -134,6 +134,8 @@ type txReserver struct {
 	f       *Interface
 	f       *Interface
 	tunnels []txTunnel
 	tunnels []txTunnel
 	owner   uint64
 	owner   uint64
+	// The staged reader separates relay data before reserving direct work. A nil pointer collects both paths.
+	inlineRelays *[]txItem
 }
 }
 
 
 // collect runs the filter over pkts and appends the sendable ones to items, tallying each tunnel's counter needs.
 // collect runs the filter over pkts and appends the sendable ones to items, tallying each tunnel's counter needs.
@@ -150,13 +152,17 @@ func (r *txReserver) collect(items []txItem, pkts []tio.Packet, fwPacket *firewa
 			continue
 			continue
 		}
 		}
 		it.segs = uint64(pkt.SegmentCount())
 		it.segs = uint64(pkt.SegmentCount())
-		it.counter = r.need(it.hostinfo, it.ci, it.segs)
 		if it.relayHI != nil {
 		if it.relayHI != nil {
 			// The relay bookkeeping inline sends do per segment in prepareSendVia, once per item
 			// The relay bookkeeping inline sends do per segment in prepareSendVia, once per item
 			f.connectionManager.OutNoRebind(it.relayHI)
 			f.connectionManager.OutNoRebind(it.relayHI)
 			f.connectionManager.RelayUsed(it.relay.LocalIndex)
 			f.connectionManager.RelayUsed(it.relay.LocalIndex)
+			if r.inlineRelays != nil {
+				*r.inlineRelays = append(*r.inlineRelays, it)
+				continue
+			}
 			it.relayCounter = r.need(it.relayHI, it.relayHI.ConnectionState, it.segs)
 			it.relayCounter = r.need(it.relayHI, it.relayHI.ConnectionState, it.segs)
 		}
 		}
+		it.counter = r.need(it.hostinfo, it.ci, it.segs)
 		items = append(items, it)
 		items = append(items, it)
 	}
 	}
 	return items
 	return items
@@ -257,11 +263,15 @@ func (s *txSlot) run(nb []byte) {
 // reader goroutine.
 // reader goroutine.
 type stagedTx struct {
 type stagedTx struct {
 	txReserver
 	txReserver
-	q         int
-	cpu       int // reader CPU, inline sends stay here
-	senderCPU int // -1 leaves the sender within its inherited affinity
-	queue     handoffQueue
-	pool      *flowPool
+	// Relay data always leaves on the reader, including small and oversized reads. Its independent reservation
+	// scratch keeps mixed reads from changing the direct pipeline's counters, gate or sender placement.
+	relayReserver txReserver
+	relaySlot     txSlot
+	q             int
+	cpu           int // reader CPU, inline sends stay here
+	senderCPU     int // -1 leaves the sender within its inherited affinity
+	queue         handoffQueue
+	pool          *flowPool
 	// sb and nb are the reader's own, for the sends it does itself
 	// sb and nb are the reader's own, for the sends it does itself
 	sb *batch.SendBatch
 	sb *batch.SendBatch
 	nb []byte
 	nb []byte
@@ -302,6 +312,10 @@ func newStagedTx(f *Interface, queue handoffQueue, q int) *stagedTx {
 		gateInline:         f.staged.inlineBelow > 0,
 		gateInline:         f.staged.inlineBelow > 0,
 		metricGateSwitches: metrics.GetOrRegisterCounter("flow.tx.gate_switches", nil),
 		metricGateSwitches: metrics.GetOrRegisterCounter("flow.tx.gate_switches", nil),
 	}
 	}
+	// Inline relay sends can share a tunnel with older direct work on this routine. They are not the ordered
+	// sender's sole owner, so they must not recycle credit around a stalled direct reservation.
+	p.relayReserver = txReserver{f: f, tunnels: make([]txTunnel, 0, 4)}
+	p.inlineRelays = &p.relaySlot.items
 	if f.pinThreads {
 	if f.pinThreads {
 		p.cpu = f.threadCPU(q)
 		p.cpu = f.threadCPU(q)
 		// Only an explicit list pins the sender, and never onto its reader's CPU, where the two would serialise.
 		// Only an explicit list pins the sender, and never onto its reader's CPU, where the two would serialise.
@@ -361,6 +375,9 @@ func (p *stagedTx) handle(pkts []tio.Packet, fwPacket *firewall.ParsedPacket, re
 	slot := p.free[len(p.free)-1]
 	slot := p.free[len(p.free)-1]
 	p.free = p.free[:len(p.free)-1]
 	p.free = p.free[:len(p.free)-1]
 	items := p.collect(slot.items[:0], pkts, fwPacket, rejectBuf, p.q, localCache)
 	items := p.collect(slot.items[:0], pkts, fwPacket, rejectBuf, p.q, localCache)
+	if len(p.relaySlot.items) != 0 {
+		p.sendRelays()
+	}
 	if len(items) == 0 {
 	if len(items) == 0 {
 		slot.items = items
 		slot.items = items
 		p.free = append(p.free, slot)
 		p.free = append(p.free, slot)
@@ -369,12 +386,12 @@ func (p *stagedTx) handle(pkts []tio.Packet, fwPacket *firewall.ParsedPacket, re
 
 
 	if p.oversized() {
 	if p.oversized() {
 		p.metricInline.Inc(1)
 		p.metricInline.Inc(1)
-		p.sendLarge(items)
+		p.sendLarge(&p.txReserver, items)
 		slot.items = items[:0]
 		slot.items = items[:0]
 		p.free = append(p.free, slot)
 		p.free = append(p.free, slot)
 		return
 		return
 	}
 	}
-	if !p.waitBudget() {
+	if !p.waitBudget(&p.txReserver) {
 		clear(items)
 		clear(items)
 		slot.items = items[:0]
 		slot.items = items[:0]
 		p.free = append(p.free, slot)
 		p.free = append(p.free, slot)
@@ -406,12 +423,40 @@ func (p *stagedTx) handle(pkts []tio.Packet, fwPacket *firewall.ParsedPacket, re
 	p.pool.jobs <- slot
 	p.pool.jobs <- slot
 }
 }
 
 
+// sendRelays keeps relay data on one send path for the lifetime of this routine. The gate's software drain only
+// waits for socket submission; switching CPUs can still overtake packets on another NIC queue. Direct data retains
+// the existing gate and sender placement. Relay credit is still shared with all other users of both tunnels.
+func (p *stagedTx) sendRelays() {
+	r, slot := &p.relayReserver, &p.relaySlot
+	r.tunnels = r.tunnels[:0]
+	for i := range slot.items {
+		it := &slot.items[i]
+		it.counter = r.need(it.hostinfo, it.ci, it.segs)
+		it.relayCounter = r.need(it.relayHI, it.relayHI.ConnectionState, it.segs)
+	}
+	if r.oversized() {
+		p.metricInline.Inc(1)
+		p.sendLarge(r, slot.items)
+		slot.items = slot.items[:0]
+		return
+	}
+	if !p.waitBudget(r) {
+		clear(slot.items)
+		slot.items = slot.items[:0]
+		return
+	}
+	slot.reservations = append(slot.reservations[:0], r.tunnels...)
+	slot.items = r.reserve(slot.items)
+	p.inline(slot)
+}
+
 // gate reports whether a read goes inline. Each routine holds one
 // gate reports whether a read goes inline. Each routine holds one
 // transmit mode, inline on the reader or fanout to the workers and the sender, until gateHold reads in a row want the
 // transmit mode, inline on the reader or fanout to the workers and the sender, until gateHold reads in a row want the
 // other one, the last of which goes the new way. Choosing per read let small inline sends leave from the reader's CPU
 // other one, the last of which goes the new way. Choosing per read let small inline sends leave from the reader's CPU
 // between staged ones leaving from the sender's, and the receiver's TCP saw them reordered. Going inline the reader
 // between staged ones leaving from the sender's, and the receiver's TCP saw them reordered. Going inline the reader
-// first waits until everything in flight is on the wire and the sender is idle, so no inline send overtakes a staged
-// one. Going to fanout nothing is in flight: inline sends finish before the reader moves on.
+// first waits until everything in flight has been submitted to the socket and the sender is idle. This does not
+// drain independent NIC queues; relay data therefore always uses the reader. Going to fanout nothing is in flight:
+// inline sends finish before the reader moves on.
 func (p *stagedTx) gate(small bool) bool {
 func (p *stagedTx) gate(small bool) bool {
 	if small == p.gateInline {
 	if small == p.gateInline {
 		p.gateStreak = 0
 		p.gateStreak = 0

+ 231 - 0
flow_tx_ring_test.go

@@ -0,0 +1,231 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"fmt"
+	"net/netip"
+	"sync"
+	"time"
+
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/slackhq/nebula/overlay/tio"
+
+	"github.com/flynn/noise"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/udp"
+	"github.com/slackhq/nebula/util"
+	"github.com/stretchr/testify/require"
+)
+
+// queuedWire models two CPU-selected NIC queues with independent completion: the reader's
+// queue is congested, while another CPU's queue drains immediately. WriteBatch returning only
+// means the kernel accepted the packets; draining staged slots does not drain a NIC queue.
+// The sampled affinity makes the schedule deterministic without a real socket or NIC. An
+// unpinned sender alternates between allowed CPUs, even if it happened not to in this run.
+func (c *affinityConn) queuedWire(readerCPU int) [][]byte {
+	var slow, fast [][]byte
+	wire := c.captured()
+	off := 0
+	for i, n := range c.sizes {
+		mask := c.masks[i]
+		if mask.IsSet(readerCPU) && (mask.Count() == 1 || i%2 == 0) {
+			slow = append(slow, wire[off:off+n]...)
+		} else {
+			fast = append(fast, wire[off:off+n]...)
+		}
+		off += n
+	}
+	return append(fast, slow...)
+}
+
+// TestFlowStagedTXRelayQueuedWire reproduces reorder after successful, ordered socket
+// submissions. Mixed reads switch the sticky gate with an empty software pipeline; relayed
+// data must stay on the reader throughout, without moving the direct sender. Each routine owns a flow.
+func TestFlowStagedTXRelayQueuedWire(t *testing.T) {
+	if os.Getenv("NEBULA_FLOW_RING_TEST") != "1" {
+		cmd := exec.Command(os.Args[0], "-test.run=^TestFlowStagedTXRelayQueuedWire$")
+		cmd.Env = append(os.Environ(), "NEBULA_FLOW_RING_TEST=1")
+		out, err := cmd.CombinedOutput()
+		require.NoError(t, err, "%s", out)
+		return
+	}
+	allowed, err := util.AllowedCPUs()
+	require.NoError(t, err)
+	if len(allowed) < 2 {
+		t.Skip("requires two allowed CPUs")
+	}
+	for _, mode := range []string{"gate", "always-staged", "oversized"} {
+		for _, explicit := range []bool{false, true} {
+			for _, routines := range []int{1, 2} {
+				for q := range routines {
+					t.Run(fmt.Sprintf("%s/explicit=%t/routines=%d/q=%d", mode, explicit, routines, q), func(t *testing.T) {
+						sc := gated(stagedGateHold)
+						if mode != "gate" {
+							sc.inlineBelow = 0
+						}
+						sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+						f := sender.f
+						f.pinThreads, f.routines = true, routines
+						if explicit {
+							f.cpuAffinity = allowed[:2]
+						}
+						_, relayHI, relayPeer := addTestRelay(t, sender, ReplayWindow)
+						conn := &affinityConn{}
+						f.writers = make([]udp.Conn, routines)
+						for i := range f.writers {
+							f.writers[i] = conn
+						}
+						shape := append(repeatBool(false, 2), repeatBool(true, stagedGateHold+2)...)
+						shape = append(shape, repeatBool(false, stagedGateHold+2)...)
+						if mode != "gate" {
+							shape = []bool{false, true, false, true}
+						}
+						reads, want := gateReads(shape...)
+						if mode == "oversized" {
+							// A deliberately tiny outer budget alone forces the oversized fallback.
+							relayHI.ConnectionState.window = NewBits(16)
+						}
+						queue := &affinityQueue{flowBatchQueue: newFlowBatchQueue(reads), conn: conn}
+						// Each Read waits for all previous socket submissions, making software drain
+						// irrelevant to the independently queued packets in queuedWire.
+						f.listenIn(queue, q)
+						unwrap := func(wire [][]byte) [][]byte {
+							inner := make([][]byte, 0, len(wire))
+							for _, pkt := range wire {
+								inner = append(inner, pkt[header.Len:len(pkt)-relayPeer.dKey.Overhead()])
+							}
+							return inner
+						}
+						require.Equal(t, want, decryptWire(t, receiver, unwrap(conn.captured())), "socket submission order")
+						require.Equal(t, want, decryptWire(t, receiver, unwrap(conn.queuedWire(f.threadCPU(q)))), "NIC queue completion order")
+					})
+				}
+			}
+		}
+	}
+}
+
+// Mixed reads must keep direct sends on the original staged sender, including when relay data is present.
+func TestFlowStagedTXMixedPlacement(t *testing.T) {
+	if os.Getenv("NEBULA_FLOW_MIXED_TEST") != "1" {
+		cmd := exec.Command(os.Args[0], "-test.run=^TestFlowStagedTXMixedPlacement$")
+		cmd.Env = append(os.Environ(), "NEBULA_FLOW_MIXED_TEST=1")
+		out, err := cmd.CombinedOutput()
+		require.NoError(t, err, "%s", out)
+		return
+	}
+	allowed, err := util.AllowedCPUs()
+	require.NoError(t, err)
+	if len(allowed) < 2 {
+		t.Skip("requires two allowed CPUs")
+	}
+	for _, explicit := range []bool{false, true} {
+		t.Run(fmt.Sprintf("explicit=%t", explicit), func(t *testing.T) {
+			sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+			_, _, relayPeer := addTestRelay(t, sender, ReplayWindow)
+			direct, directReceiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+			directAddr := netip.MustParseAddr("10.0.0.4")
+			direct.hostinfo.vpnAddrs = []netip.Addr{directAddr}
+			f := sender.f
+			f.hostMap.Hosts[directAddr] = direct.hostinfo
+			f.pinThreads, f.routines = true, 2
+			sendCPU := -1
+			if explicit {
+				f.cpuAffinity = allowed[:2]
+				sendCPU = allowed[1]
+			}
+			conn := &affinityConn{}
+			f.writers = []udp.Conn{conn, conn}
+			var reads [][]tio.Packet
+			var wantRelay, wantDirect []decodedSeg
+			for seq := range uint32(6) {
+				r := flowTCPSuperpacket(0, seq, 3)
+				d := flowTCPSuperpacket(1, seq, 9)
+				copy(d.Bytes[16:20], directAddr.AsSlice())
+				reads = append(reads, []tio.Packet{r, d})
+				for i := range 3 {
+					wantRelay = append(wantRelay, decodedSeg{seq: seq, idx: uint16(i)})
+				}
+				for i := range 9 {
+					wantDirect = append(wantDirect, decodedSeg{flow: 1, seq: seq, idx: uint16(i)})
+				}
+			}
+			queue := &affinityQueue{flowBatchQueue: newFlowBatchQueue(reads), conn: conn}
+			f.listenIn(queue, 0)
+			wire := conn.captured()
+			var directWire, relayWire [][]byte
+			off := 0
+			for i, n := range conn.sizes {
+				for _, pkt := range wire[off : off+n] {
+					var h header.H
+					require.NoError(t, h.Parse(pkt))
+					cpu := sendCPU
+					if h.Subtype == header.MessageRelay {
+						cpu = f.threadCPU(0)
+						require.NoError(t, relayPeer.VerifyRelay(nil, h.MessageCounter, pkt, make([]byte, 12)))
+						relayWire = append(relayWire, pkt[header.Len:len(pkt)-relayPeer.dKey.Overhead()])
+					} else {
+						directWire = append(directWire, pkt)
+					}
+					requireFlowCPU(t, conn.masks[i:i+1], cpu, allowed)
+				}
+				off += n
+			}
+			require.Equal(t, wantDirect, decryptWire(t, directReceiver, directWire))
+			require.Equal(t, wantRelay, decryptWire(t, receiver, relayWire))
+		})
+	}
+}
+
+// A relayed batch shares its outer tunnel with older staged direct data on the SAME routine. Inline completion
+// must not refund that direct sender's credit while its oldest reservation is stalled.
+func TestFlowStagedTXRelaySharedBudget(t *testing.T) {
+	sender, _ := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	_, via, _ := addTestRelay(t, sender, ReplayWindow)
+	// Give the relay host the same ordinary peer metadata as a direct tunnel.
+	direct := newFlowTestPeer(t, flowTestLocal, via.vpnAddrs[0], via.ConnectionState, 300, 400)
+	via.buildNetworks(sender.f.myVpnNetworksTable, direct.hostinfo.ConnectionState.peerCert.Certificate)
+	via.remoteIndexId = direct.hostinfo.remoteIndexId
+	ci := via.ConnectionState
+	ci.window = NewBits(32)
+	start := ci.messageCounter.Load()
+	entered, unblock := make(chan struct{}), make(chan struct{})
+	var release sync.Once
+	delayCrypto(direct, func(n uint64) {
+		if n == start+1 {
+			close(entered)
+			<-unblock
+		}
+	})
+	p := newStagedTx(sender.f, newFlowBatchQueue(nil), 0)
+	defer p.close()
+	defer release.Do(func() { close(unblock) })
+	d := flowTCPSuperpacket(1, 0, 9)
+	copy(d.Bytes[16:20], via.vpnAddrs[0].AsSlice())
+	p.send([]tio.Packet{d})
+	select {
+	case <-entered:
+	case <-time.After(5 * time.Second):
+		t.Fatal("direct worker did not start")
+	}
+	for seq := range uint32(7) {
+		p.send([]tio.Packet{flowTCPSuperpacket(0, seq, 1)})
+	}
+	require.Equal(t, uint64(16), (ci.txBudget.state.Load()>>24)&txBudgetMask)
+	done := make(chan struct{})
+	go func() { defer close(done); p.send([]tio.Packet{flowTCPSuperpacket(0, 7, 1)}) }()
+	select {
+	case <-done:
+		t.Error("relay recycled credit around a stalled direct reservation")
+	case <-time.After(50 * time.Millisecond):
+	}
+	require.Equal(t, start+16, ci.messageCounter.Load())
+	release.Do(func() { close(unblock) })
+	<-done
+	p.drain()
+	require.Zero(t, ci.txBudget.state.Load())
+}

+ 2 - 0
flow_tx_test.go

@@ -226,12 +226,14 @@ func (q *affinityQueue) Read() ([]tio.Packet, error) {
 type affinityConn struct {
 type affinityConn struct {
 	flowCaptureConn
 	flowCaptureConn
 	masks []unix.CPUSet
 	masks []unix.CPUSet
+	sizes []int
 }
 }
 
 
 func (c *affinityConn) WriteBatch(p [][]byte, addrs []netip.AddrPort) (int, error) {
 func (c *affinityConn) WriteBatch(p [][]byte, addrs []netip.AddrPort) (int, error) {
 	mask := threadMask()
 	mask := threadMask()
 	c.mu.Lock()
 	c.mu.Lock()
 	c.masks = append(c.masks, mask)
 	c.masks = append(c.masks, mask)
+	c.sizes = append(c.sizes, len(p))
 	c.mu.Unlock()
 	c.mu.Unlock()
 	return c.flowCaptureConn.WriteBatch(p, addrs)
 	return c.flowCaptureConn.WriteBatch(p, addrs)
 }
 }