cursed gso

udp/conn.go

@@ -17,7 +17,8 @@ type Conn interface {
 	Rebind() error
 	LocalAddr() (netip.AddrPort, error)
 	ListenOut(pg PacketBufferGetter, pc chan *packet.Packet) error
-	WriteTo(b []byte, addr netip.AddrPort) error
+	WriteTo(p *packet.Packet) error
+	WriteDirect(b []byte, port netip.AddrPort) error
 	ReloadConfig(c *config.C)
 	Close() error
 }

udp/udp_linux.go

@@ -5,9 +5,11 @@ package udp
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
+	"sync"
 	"syscall"
 	"time"
 	"unsafe"
@@ -19,13 +21,136 @@ import (
 	"golang.org/x/sys/unix"
 )
 
+const (
+	defaultGSOMaxSegments  = 16
+	defaultGSOFlushTimeout = 150 * time.Microsecond
+	maxGSOBatchBytes       = 0xFFFF
+)
+
+var (
+	errGSOFallback = errors.New("udp gso fallback")
+	errGSODisabled = errors.New("udp gso disabled")
+)
+
 var readTimeout = unix.NsecToTimeval(int64(time.Millisecond * 500))
 
+type gsoState struct {
+	m            sync.Mutex
+	Buf          []byte
+	Addr         netip.AddrPort
+	SegSize      int
+	MaxSegments  int
+	MaxBytes     int
+	FlushTimeout time.Duration
+	Timer        *time.Timer
+
+	packets []*packet.Packet
+	msg     rawMessage
+	name    [unix.SizeofSockaddrInet6]byte
+	iov     []iovec
+	ctrl    []byte
+}
+
+func (g *gsoState) Init() {
+	g.iov = make([]iovec, g.MaxSegments)
+	for i := 0; i < g.MaxSegments; i++ {
+		g.iov[i] = iovec{}
+	}
+	g.msg.Hdr.Iov = &g.iov[0]
+	g.msg.Hdr.Iovlen = 1
+
+	g.packets = make([]*packet.Packet, 0, g.MaxSegments)
+	g.ctrl = make([]byte, unix.CmsgSpace(2))
+	hdr := (*unix.Cmsghdr)(unsafe.Pointer(&g.ctrl[0]))
+	hdr.Level = unix.SOL_UDP
+	hdr.Type = unix.UDP_SEGMENT
+	setCmsgLen(hdr, unix.CmsgLen(2))
+	g.msg.Hdr.Control = &g.ctrl[0]
+	g.msg.Hdr.Controllen = uint64(len(g.ctrl))
+
+	g.name = [unix.SizeofSockaddrInet6]byte{}
+	g.msg.Hdr.Name = &g.name[0]
+
+}
+
+func (g *gsoState) setSegSizeLocked(segSize int) {
+	g.SegSize = segSize
+	x := unix.CmsgLen(0)
+	binary.LittleEndian.PutUint16(g.ctrl[x:x+2], uint16(segSize))
+}
+
+func (g *gsoState) setNameLocked(x netip.AddrPort, isV4 bool) {
+	g.Addr = x
+	nameLen := encodeSockaddr(g.name[:], g.Addr, isV4)
+	g.msg.Hdr.Name = &g.name[0]
+	g.msg.Hdr.Namelen = nameLen
+}
+
+func encodeSockaddr(dst []byte, addr netip.AddrPort, isV4 bool) uint32 {
+	if isV4 {
+		//todo?
+		//if !addr.Addr().Is4() {
+		//	return 0, fmt.Errorf("Listener is IPv4, but writing to IPv6 remote")
+		//}
+		var sa unix.RawSockaddrInet4
+		sa.Family = unix.AF_INET
+		sa.Addr = addr.Addr().As4()
+		binary.BigEndian.PutUint16((*[2]byte)(unsafe.Pointer(&sa.Port))[:], addr.Port())
+		size := unix.SizeofSockaddrInet4
+		copy(dst[:size], (*(*[unix.SizeofSockaddrInet4]byte)(unsafe.Pointer(&sa)))[:])
+		return uint32(size)
+	}
+
+	var sa unix.RawSockaddrInet6
+	sa.Family = unix.AF_INET6
+	sa.Addr = addr.Addr().As16()
+	binary.BigEndian.PutUint16((*[2]byte)(unsafe.Pointer(&sa.Port))[:], addr.Port())
+	size := unix.SizeofSockaddrInet6
+	copy(dst[:size], (*(*[unix.SizeofSockaddrInet6]byte)(unsafe.Pointer(&sa)))[:])
+	return uint32(size)
+}
+
+func (g *gsoState) sendmsgLocked(fd int) error {
+	//name already set
+	//ctrl already set
+	//g.iov = g.iov[:0]
+	g.msg.Hdr.Iovlen = uint64(len(g.packets))
+	for i := range g.packets {
+		g.iov[i].Base = &g.packets[i].Payload[0]
+		g.iov[i].Len = uint64(len(g.packets[i].Payload))
+	}
+
+	const flags = 0
+	for {
+		_, _, err := unix.Syscall(
+			unix.SYS_SENDMSG,
+			uintptr(fd),
+			uintptr(unsafe.Pointer(&g.msg)),
+			uintptr(flags),
+		)
+		//todo no matter what, reset things
+		for i := range g.packets {
+			pool := packet.GetPool()
+			pool.Put(g.packets[i])
+		}
+		g.packets = g.packets[:0]
+
+		if err != 0 {
+			return &net.OpError{Op: "sendmsg", Err: err}
+		}
+
+		return nil
+	}
+}
+
 type StdConn struct {
-	sysFd int
-	isV4  bool
-	l     *logrus.Logger
-	batch int
+	sysFd     int
+	isV4      bool
+	l         *logrus.Logger
+	batch     int
+	enableGRO bool
+	enableGSO bool
+	gso       gsoState
 }
 
 func NewListener(l *logrus.Logger, ip netip.Addr, port int, multi bool, batch int) (Conn, error) {
@@ -145,15 +270,47 @@ func (u *StdConn) ListenOut(pg PacketBufferGetter, pc chan *packet.Packet) error
 				ip, _ = netip.AddrFromSlice(names[i][8:24])
 			}
 			out.Addr = netip.AddrPortFrom(ip.Unmap(), binary.BigEndian.Uint16(names[i][2:4]))
+			ctrlLen := getRawMessageControlLen(&msgs[i])
+			if ctrlLen > 0 {
+				packets[i].SegSize = parseGROControl(packets[i].Control[:ctrlLen])
+			} else {
+				packets[i].SegSize = 0
+			}
+
 			pc <- out
 
 			//rotate this packet out so we don't overwrite it
 			packets[i] = pg()
 			msgs[i].Hdr.Iov.Base = &packets[i].Payload[0]
+			if u.enableGRO {
+				msgs[i].Hdr.Control = &packets[i].Control[0]
+				msgs[i].Hdr.Controllen = uint64(cap(packets[i].Control))
+			}
+
 		}
 	}
 }
 
+func parseGROControl(control []byte) int {
+	if len(control) == 0 {
+		return 0
+	}
+
+	cmsgs, err := unix.ParseSocketControlMessage(control)
+	if err != nil {
+		return 0
+	}
+
+	for _, c := range cmsgs {
+		if c.Header.Level == unix.SOL_UDP && c.Header.Type == unix.UDP_GRO && len(c.Data) >= 2 {
+			segSize := int(binary.LittleEndian.Uint16(c.Data[:2]))
+			return segSize
+		}
+	}
+
+	return 0
+}
+
 func (u *StdConn) ReadSingle(msgs []rawMessage) (int, error) {
 	for {
 		n, _, err := unix.Syscall6(
@@ -201,11 +358,123 @@ func (u *StdConn) ReadMulti(msgs []rawMessage) (int, error) {
 	}
 }
 
-func (u *StdConn) WriteTo(b []byte, ip netip.AddrPort) error {
-	if u.isV4 {
-		return u.writeTo4(b, ip)
+func (u *StdConn) WriteTo(p *packet.Packet) error {
+	if u.enableGSO && p.Addr.IsValid() {
+		if err := u.queueGSOPacket(p); err == nil {
+			return nil
+		} else if !errors.Is(err, errGSOFallback) {
+			return err
+		}
 	}
-	return u.writeTo6(b, ip)
+
+	var err error
+	if u.isV4 {
+		err = u.writeTo4(p.Payload, p.Addr)
+	} else {
+		err = u.writeTo4(p.Payload, p.Addr)
+	}
+	packet.GetPool().Put(p)
+	return err
+}
+
+func (u *StdConn) WriteDirect(b []byte, addr netip.AddrPort) error {
+	if u.isV4 {
+		return u.writeTo4(b, addr)
+	}
+	return u.writeTo6(b, addr)
+}
+
+func (u *StdConn) scheduleGSOFlushLocked() {
+	if u.gso.Timer == nil {
+		u.gso.Timer = time.AfterFunc(u.gso.FlushTimeout, u.gsoFlushTimer)
+		return
+	}
+	u.gso.Timer.Reset(u.gso.FlushTimeout)
+}
+
+func (u *StdConn) stopGSOTimerLocked() {
+	if u.gso.Timer != nil {
+		u.gso.Timer.Stop()
+		u.gso.Timer = nil //todo I also don't like this
+	}
+}
+
+func (u *StdConn) queueGSOPacket(p *packet.Packet) error {
+	if len(p.Payload) == 0 {
+		return nil
+	}
+
+	u.gso.m.Lock()
+	defer u.gso.m.Unlock()
+
+	if !u.enableGSO || !p.Addr.IsValid() || len(p.Payload) > u.gso.MaxBytes {
+		if err := u.flushGSOlocked(); err != nil {
+			return err
+		}
+		return errGSOFallback
+	}
+
+	if len(u.gso.packets) == 0 {
+		u.gso.setNameLocked(p.Addr, u.isV4)
+		u.gso.SegSize = len(p.Payload)
+		u.gso.packets = append(u.gso.packets, p)
+	} else if p.Addr != u.gso.Addr || len(p.Payload) != u.gso.SegSize {
+		if err := u.flushGSOlocked(); err != nil {
+			return err
+		} //todo deal with "one small packet" case
+		u.gso.setNameLocked(p.Addr, u.isV4)
+		u.gso.SegSize = len(p.Payload)
+		u.gso.packets = append(u.gso.packets, p)
+	} else {
+		u.gso.packets = append(u.gso.packets, p)
+	}
+
+	//big todo
+	//if len(u.gso.Buf)+len(p.Payload) > u.gso.MaxBytes {
+	//	if err := u.flushGSOlocked(); err != nil {
+	//		return err
+	//	}
+	//	u.gso.setNameLocked(p.Addr, u.isV4)
+	//	u.gso.SegSize = len(p.Payload)
+	//	u.gso.packets = append(u.gso.packets, p)
+	//}
+
+	if len(u.gso.packets) >= u.gso.MaxSegments || u.gso.FlushTimeout <= 0 {
+		return u.flushGSOlocked()
+	}
+
+	u.scheduleGSOFlushLocked()
+	return nil
+}
+
+func (u *StdConn) flushGSOlocked() error {
+	if len(u.gso.packets) == 0 {
+		u.stopGSOTimerLocked()
+		return nil
+	}
+
+	u.stopGSOTimerLocked()
+
+	if u.gso.SegSize <= 0 {
+		return errGSOFallback
+	}
+
+	err := u.gso.sendmsgLocked(u.sysFd)
+	if errors.Is(err, errGSODisabled) {
+		u.l.WithField("addr", u.gso.Addr).Warn("UDP GSO disabled by kernel, falling back to sendto")
+		u.enableGSO = false
+		//todo!
+		//return u.sendSegmentsIndividually(payload, addr, segSize)
+	}
+	u.gso.SegSize = 0
+
+	return err
+}
+
+func (u *StdConn) gsoFlushTimer() {
+	u.gso.m.Lock()
+	_ = u.flushGSOlocked()
+	u.gso.m.Unlock()
 }
 
 func (u *StdConn) writeTo6(b []byte, ip netip.AddrPort) error {
@@ -308,6 +577,72 @@ func (u *StdConn) ReloadConfig(c *config.C) {
 			u.l.WithError(err).Error("Failed to set listen.so_mark")
 		}
 	}
+	u.configureGRO(true)
+	u.configureGSO(c)
+}
+
+func (u *StdConn) configureGRO(enable bool) {
+	if enable == u.enableGRO {
+		return
+	}
+
+	if enable {
+		if err := unix.SetsockoptInt(u.sysFd, unix.SOL_UDP, unix.UDP_GRO, 1); err != nil {
+			u.l.WithError(err).Warn("Failed to enable UDP GRO")
+			return
+		}
+		u.enableGRO = true
+		u.l.Info("UDP GRO enabled")
+		return
+	}
+
+	if err := unix.SetsockoptInt(u.sysFd, unix.SOL_UDP, unix.UDP_GRO, 0); err != nil && err != unix.ENOPROTOOPT {
+		u.l.WithError(err).Warn("Failed to disable UDP GRO")
+	}
+	u.enableGRO = false
+}
+
+func (u *StdConn) configureGSO(c *config.C) {
+	enable := c.GetBool("listen.enable_gso", true)
+	if !enable {
+		u.disableGSO()
+	} else {
+		u.enableGSO = true
+	}
+
+	segments := c.GetInt("listen.gso_max_segments", defaultGSOMaxSegments)
+	if segments < 1 {
+		segments = 1
+	}
+	u.gso.MaxSegments = segments
+
+	maxBytes := c.GetInt("listen.gso_max_bytes", 0)
+	if maxBytes <= 0 {
+		maxBytes = MTU * segments
+	}
+	if maxBytes > maxGSOBatchBytes {
+		u.l.WithField("requested", maxBytes).Warn("listen.gso_max_bytes larger than UDP limit; clamping")
+		maxBytes = maxGSOBatchBytes
+	}
+	u.gso.MaxBytes = maxBytes
+
+	timeout := c.GetDuration("listen.gso_flush_timeout", defaultGSOFlushTimeout)
+	if timeout < 0 {
+		timeout = 0
+	}
+	u.gso.FlushTimeout = timeout
+	u.gso.Init()
+}
+
+func (u *StdConn) disableGSO() {
+	u.gso.m.Lock()
+	defer u.gso.m.Unlock()
+	u.enableGSO = false
+	_ = u.flushGSOlocked()
+	u.gso.Buf = nil
+	u.gso.packets = u.gso.packets[:0]
+	u.gso.SegSize = 0
+	u.stopGSOTimerLocked()
 }
 
 func (u *StdConn) getMemInfo(meminfo *[unix.SK_MEMINFO_VARS]uint32) error {

udp/udp_linux_64.go

@@ -34,6 +34,24 @@ type rawMessage struct {
 	Pad0 [4]byte
 }
 
+func setRawMessageControl(msg *rawMessage, buf []byte) {
+	if len(buf) == 0 {
+		msg.Hdr.Control = nil
+		msg.Hdr.Controllen = 0
+		return
+	}
+	msg.Hdr.Control = &buf[0]
+	msg.Hdr.Controllen = uint64(len(buf))
+}
+
+func getRawMessageControlLen(msg *rawMessage) int {
+	return int(msg.Hdr.Controllen)
+}
+
+func setCmsgLen(h *unix.Cmsghdr, l int) {
+	h.Len = uint64(l)
+}
+
 func (u *StdConn) PrepareRawMessages(n int, pg PacketBufferGetter) ([]rawMessage, []*packet.Packet, [][]byte) {
 	msgs := make([]rawMessage, n)
 	names := make([][]byte, n)
@@ -42,6 +60,7 @@ func (u *StdConn) PrepareRawMessages(n int, pg PacketBufferGetter) ([]rawMessage
 	for i := range packets {
 		packets[i] = pg()
 	}
+	//todo?
 
 	for i := range msgs {
 		names[i] = make([]byte, unix.SizeofSockaddrInet6)
@@ -55,6 +74,13 @@ func (u *StdConn) PrepareRawMessages(n int, pg PacketBufferGetter) ([]rawMessage
 
 		msgs[i].Hdr.Name = &names[i][0]
 		msgs[i].Hdr.Namelen = uint32(len(names[i]))
+		if u.enableGRO {
+			msgs[i].Hdr.Control = &packets[i].Control[0]
+			msgs[i].Hdr.Controllen = uint64(len(packets[i].Control))
+		} else {
+			msgs[i].Hdr.Control = nil
+			msgs[i].Hdr.Controllen = 0
+		}
 	}
 
 	return msgs, packets, names