tio: Count superpacket segments on Packet

Move the staged transmit segment count onto tio.Packet, next to
IsSuperpacket and SegmentSuperpacket, which it predicts.

flow_tx.go

@@ -120,15 +120,6 @@ func (f *Interface) encryptTxSegmentVia(it *txItem, c, rc uint64, seg, nb []byte
 	sb.Commit(toSend, remote)
 }
 
-// segmentCount reports how many segments tio.SegmentSuperpacket will cut pkt into, without touching its bytes.
-func segmentCount(pkt tio.Packet) uint64 {
-	if !pkt.GSO.IsSuperpacket() {
-		return 1
-	}
-	n := (len(pkt.Bytes) - int(pkt.GSO.HdrLen) + int(pkt.GSO.Size) - 1) / int(pkt.GSO.Size)
-	return uint64(max(n, 1))
-}
-
 // txTunnel accumulates one tunnel's counter needs across a batch so the batch reserves them with one atomic add.
 type txTunnel struct {
 	ci        *ConnectionState
@@ -159,7 +150,7 @@ func (r *txReserver) collect(items []txItem, pkts []tio.Packet, fwPacket *firewa
 		if !f.prepareTxItem(&it, hostinfo, pkt) {
 			continue
 		}
-		it.segs = segmentCount(pkt)
+		it.segs = uint64(pkt.SegmentCount())
 		it.counter = r.need(it.hostinfo, it.ci, it.segs)
 		if it.relayHI != nil {
 			// The relay bookkeeping inline sends do per segment in prepareSendVia, once per item
@@ -496,7 +487,7 @@ func (p *stagedTx) small(pkts []tio.Packet) bool {
 	}
 	var n uint64
 	for _, pkt := range pkts {
-		if n += segmentCount(pkt); n >= below {
+		if n += uint64(pkt.SegmentCount()); n >= below {
 			return false
 		}
 	}

flow_tx_test.go

@@ -43,15 +43,15 @@ func TestReservedEncryptionMatchesInline(t *testing.T) {
 	require.Equal(t, want, decryptWire(t, receiver, sender.conn.captured()))
 }
 
-// TestSegmentCount checks segmentCount predicts exactly the segments tio.SegmentSuperpacket cuts, odd tails included.
+// TestSegmentCount checks tio.Packet.SegmentCount predicts exactly the segments tio.SegmentSuperpacket cuts, odd tails included.
 func TestSegmentCount(t *testing.T) {
-	require.Equal(t, uint64(1), segmentCount(tio.Packet{Bytes: make([]byte, 60)}))
+	require.Equal(t, 1, tio.Packet{Bytes: make([]byte, 60)}.SegmentCount())
 	for _, segs := range []int{1, 2, 3, 45} {
 		for _, tail := range []int{0, 1, flowTestMSS - 1} {
 			pkt := flowTCPSuperpacket(0, 0, segs)
 			pkt.Bytes = pkt.Bytes[:len(pkt.Bytes)-tail]
 			pkt.GSO = tio.GSOInfo{Size: flowTestMSS, HdrLen: 40, CsumStart: 20, Proto: tio.GSOProtoTCP}
-			want, got := segmentCount(pkt), uint64(0)
+			want, got := pkt.SegmentCount(), 0
 			require.NoError(t, tio.SegmentSuperpacket(pkt, func([]byte) error { got++; return nil }))
 			require.Equal(t, want, got, "segs %d tail %d", segs, tail)
 		}
@@ -212,7 +212,7 @@ type affinityQueue struct {
 func (q *affinityQueue) Read() ([]tio.Packet, error) {
 	if n := q.reads.Load(); n > 0 {
 		for _, p := range q.batches[n-1] {
-			q.sent += int(segmentCount(p))
+			q.sent += p.SegmentCount()
 		}
 	}
 	q.boundaries = append(q.boundaries, q.conn.waitFor(q.sent))

overlay/tio/tio.go

@@ -75,6 +75,15 @@ type GSOInfo struct {
 // superpacket that needs segmentation before its bytes can be encrypted and sent on the wire.
 func (g GSOInfo) IsSuperpacket() bool { return g.Size > 0 }
 
+// SegmentCount reports how many segments SegmentSuperpacket will cut p into, without touching its bytes.
+func (p Packet) SegmentCount() int {
+	if !p.GSO.IsSuperpacket() {
+		return 1
+	}
+	n := (len(p.Bytes) - int(p.GSO.HdrLen) + int(p.GSO.Size) - 1) / int(p.GSO.Size)
+	return max(n, 1)
+}
+
 // Clone returns a Packet whose Bytes is a freshly allocated copy of p.Bytes,
 // safe to retain past the next Read or Close on the originating Queue.
 // GSO metadata is copied verbatim.