relay: keep a short batch's remainder; per-source cap to 4000/s

CHANGELOG.md

@@ -5,6 +5,22 @@ All notable changes to this project are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Fixed
+- **A partly written relay batch is no longer dropped.** When the batched
+  send wrote fewer messages than it was given without returning an error (a
+  short `sendmmsg`, or any platform without it, where one message is sent per
+  call), the rest of the batch was discarded. The remainder is now sent
+  per-packet.
+
+### Changed
+- **The per-source relay limit is 4000 packets a second, up from 1000.** The
+  limit counts packets. Daemons are moving from ~4.2 KB stream datagrams
+  (which IP fragmented) to ~1.2 KB ones that fit a single packet, so the same
+  traffic is about 3.5 times as many packets; 4000 keeps the limit where it
+  was in bytes (about 4 MB/s per source).
+
 ## [v0.1.0]
 
 Initial release.

server.go

@@ -155,6 +155,13 @@ func relayWorkerCount() int {
 	return n
 }
 
+// relayQueueFor returns the relay queue (and with it the worker) that carries
+// a packet for destID. Every packet to one destination takes the same queue,
+// so the packets of a flow leave the beacon in the order they arrived.
+func (s *Server) relayQueueFor(destID uint32) chan relayJob {
+	return s.relayChs[destID%uint32(len(s.relayChs))]
+}
+
 // maxRelayPayload caps the relay payload size. UDP itself limits datagrams to ~65KB,
 // but this provides defense-in-depth against future transport changes.
 const maxRelayPayload = 65535
@@ -187,12 +194,18 @@ type relaySourceWindow struct {
 }
 
 // maxRelaysPerSourcePerSecond caps relays a single sender can push through
-// dispatchRelay per second. Set generously (1000 relays/sec) — a legitimate
-// source behind NAT may relay traffic for many agents, but a DoS source
-// flooding relays to a known target can saturate the 524288-deep relayCh
-// and cause queue-full drops for everyone. This cap prevents one source
-// from consuming more than ~0.2% of total queue capacity per second.
-const maxRelaysPerSourcePerSecond = 1000
+// dispatchRelay per second — a legitimate source behind NAT may relay
+// traffic for many agents, but a DoS source flooding relays to a known
+// target can fill a relay queue and cause queue-full drops for the flows
+// that share it.
+//
+// The cap counts packets, so what it allows in bytes depends on the size of
+// the sender's datagrams. Daemons used to send stream segments as ~4.2 KB
+// datagrams (fragmented by IP), and 1000 a second was about 4 MB/s. They now
+// send ~1.2 KB datagrams that fit one packet; 4000 a second keeps the same
+// ceiling in bytes. At 1000 a relayed transfer was held to about 1.1 MB/s,
+// with everything past the cap dropped and retransmitted.
+const maxRelaysPerSourcePerSecond = 4000
 const relaySourceCleanupInterval = 5 * time.Minute
 
 func New() *Server {
@@ -964,7 +977,7 @@ func (s *Server) dispatchRelay(data []byte, src relaySourceKey) {
 	}
 	copy(buf, payload)
 
-	ch := s.relayChs[destID%uint32(len(s.relayChs))]
+	ch := s.relayQueueFor(destID)
 	select {
 	case ch <- relayJob{senderID: senderID, destID: destID, payload: buf}:
 	default:
@@ -1077,12 +1090,16 @@ func (s *Server) relayWorkerLoop(sendConn *ipv4.PacketConn, rawConn *net.UDPConn
 		} else {
 			var err error
 			n, err = sendConn.WriteBatch(msgs, 0)
-			if err != nil {
-				// Unexpected on an IPv4-bound socket, but fall back to
-				// per-packet for the messages the batch didn't send rather
-				// than dropping them.
-				slog.Debug("relay write batch failed, falling back to per-packet send",
-					"n", n, "of", len(msgs), "err", err)
+			if n < len(msgs) {
+				// The batch call sent only part of the batch: an error
+				// (unexpected on an IPv4-bound socket), a short sendmmsg,
+				// or a platform without sendmmsg, where WriteBatch sends
+				// one message per call and reports no error. Send the
+				// rest per-packet rather than dropping them.
+				if err != nil {
+					slog.Debug("relay write batch failed, falling back to per-packet send",
+						"n", n, "of", len(msgs), "err", err)
+				}
 				for i := n; i < len(msgs); i++ {
 					dst, ok := msgs[i].Addr.(*net.UDPAddr)
 					if !ok {

zz_compat_wss_test.go

@@ -15,8 +15,8 @@ import (
 
 	"github.com/coder/websocket"
 
-	"github.com/pilot-protocol/common/protocol"
 	"github.com/pilot-protocol/common/crypto"
+	"github.com/pilot-protocol/common/protocol"
 )
 
 // authChallengeMsg / authReplyMsg / authOKMsg mirror the unexported
@@ -120,8 +120,6 @@ func TestEnableCompatWSS_HappyPath(t *testing.T) {
 	nodeID := uint32(4242)
 
 	s := New()
-	go s.ListenAndServe("127.0.0.1:0")
-	<-s.Ready()
 	defer s.Close()
 
 	lookup := func(qid uint32) (ed25519.PublicKey, bool) {
@@ -131,9 +129,16 @@ func TestEnableCompatWSS_HappyPath(t *testing.T) {
 		return nil, false
 	}
 
+	// Before ListenAndServe, as EnableCompatWSS documents: the read loop
+	// reads the WSS server it installs on every relayed packet. Enabling it
+	// after the loop had started was a data race, which the race detector
+	// reported intermittently in CI against this test and whichever tests
+	// were running beside it.
 	if err := s.EnableCompatWSS("127.0.0.1:0", lookup); err != nil {
 		t.Fatalf("EnableCompatWSS: %v", err)
 	}
+	go s.ListenAndServe("127.0.0.1:0")
+	<-s.Ready()
 
 	wsAddr := s.WSSAddr()
 	if wsAddr == "" {

zz_relay_order_test.go

@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+package beacon_test
+
+import (
+	"encoding/binary"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/pilot-protocol/common/protocol"
+)
+
+// The packets of one relayed flow must leave the beacon in the order they
+// arrived. The relay workers used to share one queue and batch separately, so
+// consecutive packets left in different batches, out of order; a stream
+// receiver reports each early arrival, and three reports make the sender
+// retransmit and halve its window with nothing actually lost.
+func TestRelayKeepsAFlowsPacketsInOrder(t *testing.T) {
+	t.Parallel()
+	s, addr := startTestBeacon(t)
+	defer s.Close()
+
+	const senderID, destID = 4101, 4102
+
+	dest, err := net.DialUDP("udp", nil, addr)
+	if err != nil {
+		t.Fatalf("dial dest: %v", err)
+	}
+	defer dest.Close()
+	discover := make([]byte, 5)
+	discover[0] = protocol.BeaconMsgDiscover
+	binary.BigEndian.PutUint32(discover[1:], destID)
+	if _, err := dest.Write(discover); err != nil {
+		t.Fatalf("discover: %v", err)
+	}
+	// The discover reply; after it, everything on this socket is a delivery.
+	dest.SetReadDeadline(time.Now().Add(2 * time.Second))
+	buf := make([]byte, 2048)
+	if _, err := dest.Read(buf); err != nil {
+		t.Fatalf("discover reply: %v", err)
+	}
+
+	sender, err := net.DialUDP("udp", nil, addr)
+	if err != nil {
+		t.Fatalf("dial sender: %v", err)
+	}
+	defer sender.Close()
+
+	// Well under the per-source limit, in bursts short enough that the
+	// loopback socket buffers hold them.
+	const packets = 600
+	go func() {
+		msg := make([]byte, 9+4)
+		msg[0] = protocol.BeaconMsgRelay
+		binary.BigEndian.PutUint32(msg[1:], senderID)
+		binary.BigEndian.PutUint32(msg[5:], destID)
+		for i := 0; i < packets; i++ {
+			binary.BigEndian.PutUint32(msg[9:], uint32(i))
+			sender.Write(msg)
+			if i%50 == 49 {
+				time.Sleep(time.Millisecond)
+			}
+		}
+	}()
+
+	got, outOfOrder, last := 0, 0, -1
+	for got < packets {
+		dest.SetReadDeadline(time.Now().Add(time.Second))
+		n, err := dest.Read(buf)
+		if err != nil {
+			break // whatever is left was dropped on loopback; ordering is what is under test
+		}
+		if n != 1+4+4 || buf[0] != protocol.BeaconMsgRelayDeliver {
+			continue
+		}
+		seq := int(binary.BigEndian.Uint32(buf[5:9]))
+		if seq < last {
+			outOfOrder++
+		} else {
+			last = seq
+		}
+		got++
+	}
+	if got < packets*9/10 {
+		t.Fatalf("only %d of %d relayed packets arrived", got, packets)
+	}
+	if outOfOrder != 0 {
+		t.Fatalf("%d of %d relayed packets arrived after a later one", outOfOrder, got)
+	}
+}

zz_relay_queue_test.go

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+package beacon
+
+import "testing"
+
+// Packets to one destination keep to one relay queue, so one worker sends a
+// flow in order; different destinations spread over all the queues.
+func TestRelayQueueIsStablePerDestAndSpreadsDests(t *testing.T) {
+	t.Parallel()
+	s := New()
+	if len(s.relayChs) != relayWorkerCount() {
+		t.Fatalf("%d relay queues, want one per worker (%d)", len(s.relayChs), relayWorkerCount())
+	}
+	total := 0
+	for _, q := range s.relayChs {
+		total += cap(q)
+	}
+	if total > relayQueueSize || total < relayQueueSize-len(s.relayChs) {
+		t.Fatalf("relay queues hold %d jobs in total, want relayQueueSize (%d) split across them", total, relayQueueSize)
+	}
+
+	q := s.relayQueueFor(2002)
+	for i := 0; i < 100; i++ {
+		if s.relayQueueFor(2002) != q {
+			t.Fatal("the same destination was given a different queue")
+		}
+	}
+
+	// 8000 destinations spread within 2x of an even split over the workers.
+	used := make(map[chan relayJob]int)
+	const dests = 8000
+	for i := uint32(0); i < dests; i++ {
+		used[s.relayQueueFor(100000+i)]++
+	}
+	if len(used) != len(s.relayChs) {
+		t.Fatalf("%d destinations used %d of %d queues", dests, len(used), len(s.relayChs))
+	}
+	fair := dests / len(s.relayChs)
+	for _, n := range used {
+		if n < fair/2 || n > fair*2 {
+			t.Fatalf("a queue got %d of %d destinations, want about %d each", n, dests, fair)
+		}
+	}
+}