Improve darwin tun performance

internal/fdbased_darwin/processors.go

@@ -0,0 +1,275 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package fdbased
+
+import (
+	"context"
+	"encoding/binary"
+
+	"github.com/sagernet/gvisor/pkg/rand"
+	"github.com/sagernet/gvisor/pkg/sleep"
+	"github.com/sagernet/gvisor/pkg/sync"
+	"github.com/sagernet/gvisor/pkg/tcpip"
+	"github.com/sagernet/gvisor/pkg/tcpip/hash/jenkins"
+	"github.com/sagernet/gvisor/pkg/tcpip/header"
+	"github.com/sagernet/gvisor/pkg/tcpip/stack"
+	"github.com/sagernet/gvisor/pkg/tcpip/stack/gro"
+)
+
+// +stateify savable
+type processor struct {
+	mu processorMutex `state:"nosave"`
+	// +checklocks:mu
+	pkts stack.PacketBufferList
+
+	e           *endpoint
+	gro         gro.GRO
+	sleeper     sleep.Sleeper
+	packetWaker sleep.Waker
+	closeWaker  sleep.Waker
+}
+
+func (p *processor) start(wg *sync.WaitGroup) {
+	defer wg.Done()
+	defer p.sleeper.Done()
+	for {
+		switch w := p.sleeper.Fetch(true); {
+		case w == &p.packetWaker:
+			p.deliverPackets()
+		case w == &p.closeWaker:
+			p.mu.Lock()
+			p.pkts.Reset()
+			p.mu.Unlock()
+			return
+		}
+	}
+}
+
+func (p *processor) deliverPackets() {
+	p.e.mu.RLock()
+	p.gro.Dispatcher = p.e.dispatcher
+	p.e.mu.RUnlock()
+	if p.gro.Dispatcher == nil {
+		p.mu.Lock()
+		p.pkts.Reset()
+		p.mu.Unlock()
+		return
+	}
+
+	p.mu.Lock()
+	for p.pkts.Len() > 0 {
+		pkt := p.pkts.PopFront()
+		p.mu.Unlock()
+		p.gro.Enqueue(pkt)
+		pkt.DecRef()
+		p.mu.Lock()
+	}
+	p.mu.Unlock()
+	p.gro.Flush()
+}
+
+// processorManager handles starting, closing, and queuing packets on processor
+// goroutines.
+//
+// +stateify savable
+type processorManager struct {
+	processors []processor
+	seed       uint32
+	wg         sync.WaitGroup `state:"nosave"`
+	e          *endpoint
+	ready      []bool
+}
+
+// newProcessorManager creates a new processor manager.
+func newProcessorManager(opts *Options, e *endpoint) *processorManager {
+	m := &processorManager{}
+	m.seed = rand.Uint32()
+	m.ready = make([]bool, opts.ProcessorsPerChannel)
+	m.processors = make([]processor, opts.ProcessorsPerChannel)
+	m.e = e
+	m.wg.Add(opts.ProcessorsPerChannel)
+
+	for i := range m.processors {
+		p := &m.processors[i]
+		p.sleeper.AddWaker(&p.packetWaker)
+		p.sleeper.AddWaker(&p.closeWaker)
+		p.gro.Init(false)
+		p.e = e
+	}
+
+	return m
+}
+
+// start starts the processor goroutines if the processor manager is configured
+// with more than one processor.
+func (m *processorManager) start() {
+	for i := range m.processors {
+		p := &m.processors[i]
+		// Only start processor in a separate goroutine if we have multiple of them.
+		if len(m.processors) > 1 {
+			go p.start(&m.wg)
+		}
+	}
+}
+
+// afterLoad is invoked by stateify.
+func (m *processorManager) afterLoad(context.Context) {
+	m.wg.Add(len(m.processors))
+	m.start()
+}
+
+func (m *processorManager) connectionHash(cid *connectionID) uint32 {
+	var payload [4]byte
+	binary.LittleEndian.PutUint16(payload[0:], cid.srcPort)
+	binary.LittleEndian.PutUint16(payload[2:], cid.dstPort)
+
+	h := jenkins.Sum32(m.seed)
+	h.Write(payload[:])
+	h.Write(cid.srcAddr)
+	h.Write(cid.dstAddr)
+	return h.Sum32()
+}
+
+// queuePacket queues a packet to be delivered to the appropriate processor.
+func (m *processorManager) queuePacket(pkt *stack.PacketBuffer, hasEthHeader bool) {
+	var pIdx uint32
+	cid, nonConnectionPkt := tcpipConnectionID(pkt)
+	if !hasEthHeader {
+		if nonConnectionPkt {
+			// If there's no eth header this should be a standard tcpip packet. If
+			// it isn't the packet is invalid so drop it.
+			return
+		}
+		pkt.NetworkProtocolNumber = cid.proto
+	}
+	if len(m.processors) == 1 || nonConnectionPkt {
+		// If the packet is not associated with an active connection, use the
+		// first processor.
+		pIdx = 0
+	} else {
+		pIdx = m.connectionHash(&cid) % uint32(len(m.processors))
+	}
+	p := &m.processors[pIdx]
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.pkts.PushBack(pkt.IncRef())
+	m.ready[pIdx] = true
+}
+
+type connectionID struct {
+	srcAddr, dstAddr []byte
+	srcPort, dstPort uint16
+	proto            tcpip.NetworkProtocolNumber
+}
+
+// tcpipConnectionID returns a tcpip connection id tuple based on the data found
+// in the packet. It returns true if the packet is not associated with an active
+// connection (e.g ARP, NDP, etc). The method assumes link headers have already
+// been processed if they were present.
+func tcpipConnectionID(pkt *stack.PacketBuffer) (connectionID, bool) {
+	var cid connectionID
+	h, ok := pkt.Data().PullUp(1)
+	if !ok {
+		// Skip this packet.
+		return cid, true
+	}
+
+	const tcpSrcDstPortLen = 4
+	switch header.IPVersion(h) {
+	case header.IPv4Version:
+		hdrLen := header.IPv4(h).HeaderLength()
+		h, ok = pkt.Data().PullUp(int(hdrLen) + tcpSrcDstPortLen)
+		if !ok {
+			return cid, true
+		}
+		ipHdr := header.IPv4(h[:hdrLen])
+		tcpHdr := header.TCP(h[hdrLen:][:tcpSrcDstPortLen])
+
+		cid.srcAddr = ipHdr.SourceAddressSlice()
+		cid.dstAddr = ipHdr.DestinationAddressSlice()
+		// All fragment packets need to be processed by the same goroutine, so
+		// only record the TCP ports if this is not a fragment packet.
+		if ipHdr.IsValid(pkt.Data().Size()) && !ipHdr.More() && ipHdr.FragmentOffset() == 0 {
+			cid.srcPort = tcpHdr.SourcePort()
+			cid.dstPort = tcpHdr.DestinationPort()
+		}
+		cid.proto = header.IPv4ProtocolNumber
+	case header.IPv6Version:
+		h, ok = pkt.Data().PullUp(header.IPv6FixedHeaderSize + tcpSrcDstPortLen)
+		if !ok {
+			return cid, true
+		}
+		ipHdr := header.IPv6(h)
+
+		var tcpHdr header.TCP
+		if tcpip.TransportProtocolNumber(ipHdr.NextHeader()) == header.TCPProtocolNumber {
+			tcpHdr = header.TCP(h[header.IPv6FixedHeaderSize:][:tcpSrcDstPortLen])
+		} else {
+			// Slow path for IPv6 extension headers :(.
+			dataBuf := pkt.Data().ToBuffer()
+			dataBuf.TrimFront(header.IPv6MinimumSize)
+			it := header.MakeIPv6PayloadIterator(header.IPv6ExtensionHeaderIdentifier(ipHdr.NextHeader()), dataBuf)
+			defer it.Release()
+			for {
+				hdr, done, err := it.Next()
+				if done || err != nil {
+					break
+				}
+				hdr.Release()
+			}
+			h, ok = pkt.Data().PullUp(int(it.HeaderOffset()) + tcpSrcDstPortLen)
+			if !ok {
+				return cid, true
+			}
+			tcpHdr = header.TCP(h[it.HeaderOffset():][:tcpSrcDstPortLen])
+		}
+		cid.srcAddr = ipHdr.SourceAddressSlice()
+		cid.dstAddr = ipHdr.DestinationAddressSlice()
+		cid.srcPort = tcpHdr.SourcePort()
+		cid.dstPort = tcpHdr.DestinationPort()
+		cid.proto = header.IPv6ProtocolNumber
+	default:
+		return cid, true
+	}
+	return cid, false
+}
+
+func (m *processorManager) close() {
+	if len(m.processors) < 2 {
+		return
+	}
+	for i := range m.processors {
+		p := &m.processors[i]
+		p.closeWaker.Assert()
+	}
+}
+
+// wakeReady wakes up all processors that have a packet queued. If there is only
+// one processor, the method delivers the packet inline without waking a
+// goroutine.
+func (m *processorManager) wakeReady() {
+	for i, ready := range m.ready {
+		if !ready {
+			continue
+		}
+		p := &m.processors[i]
+		if len(m.processors) > 1 {
+			p.packetWaker.Assert()
+		} else {
+			p.deliverPackets()
+		}
+		m.ready[i] = false
+	}
+}