Improve UDP batch support

common/interrupt/conn.go

@@ -3,6 +3,7 @@ package interrupt
 import (
 	"net"
 
+	"github.com/sagernet/sing/common/bufio"
 	"github.com/sagernet/sing/common/x/list"
 )
 
@@ -71,5 +72,5 @@ func (c *PacketConn) WriterReplaceable() bool {
 }
 
 func (c *PacketConn) Upstream() any {
-	return c.PacketConn
+	return bufio.NewPacketConn(c.PacketConn)
 }

common/listener/listener.go

@@ -25,9 +25,9 @@ type Listener struct {
 	logger                   logger.ContextLogger
 	network                  []string
 	listenOptions            option.ListenOptions
-	connHandler              adapter.ConnectionHandlerEx
-	packetHandler            adapter.PacketHandlerEx
-	oobPacketHandler         adapter.OOBPacketHandlerEx
+	connHandler              adapter.ConnectionHandler
+	packetHandler            adapter.PacketHandler
+	oobPacketHandler         adapter.OOBPacketHandler
 	threadUnsafePacketWriter bool
 	disablePacketOutput      bool
 	setSystemProxy           bool
@@ -48,9 +48,9 @@ type Options struct {
 	Logger                   logger.ContextLogger
 	Network                  []string
 	Listen                   option.ListenOptions
-	ConnectionHandler        adapter.ConnectionHandlerEx
-	PacketHandler            adapter.PacketHandlerEx
-	OOBPacketHandler         adapter.OOBPacketHandlerEx
+	ConnectionHandler        adapter.ConnectionHandler
+	PacketHandler            adapter.PacketHandler
+	OOBPacketHandler         adapter.OOBPacketHandler
 	ThreadUnsafePacketWriter bool
 	DisablePacketOutput      bool
 	SetSystemProxy           bool

common/listener/listener_tcp.go

@@ -106,6 +106,6 @@ func (l *Listener) loopTCPIn() {
 		metadata.OriginDestination = M.SocksaddrFromNet(conn.LocalAddr()).Unwrap()
 		ctx := log.ContextWithNewID(l.ctx)
 		l.logger.InfoContext(ctx, "inbound connection from ", metadata.Source)
-		go l.connHandler.NewConnectionEx(ctx, conn, metadata, nil)
+		go l.connHandler.NewConnection(ctx, conn, metadata, nil)
 	}
 }

common/listener/listener_udp.go

@@ -11,6 +11,7 @@ import (
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/common/redir"
 	"github.com/sagernet/sing/common/buf"
+	sBufio "github.com/sagernet/sing/common/bufio"
 	"github.com/sagernet/sing/common/control"
 	E "github.com/sagernet/sing/common/exceptions"
 	M "github.com/sagernet/sing/common/metadata"
@@ -18,6 +19,8 @@ import (
 	"github.com/sagernet/sing/service"
 )
 
+const udpOutputBatchSize = 128
+
 func (l *Listener) ListenUDP() (net.PacketConn, error) {
 	bindAddr := M.SocksaddrFrom(l.listenOptions.Listen.Build(netip.AddrFrom4([4]byte{127, 0, 0, 1})), l.listenOptions.ListenPort)
 	var listenConfig net.ListenConfig
@@ -98,6 +101,15 @@ func (l *Listener) PacketWriter() N.PacketWriter {
 
 func (l *Listener) loopUDPIn() {
 	defer close(l.packetOutboundClosed)
+	if l.oobPacketHandler == nil {
+		if batchHandler, isBatchHandler := l.packetHandler.(adapter.PacketBatchHandler); isBatchHandler {
+			packetConn := sBufio.NewPacketConn(l.udpConn)
+			if readWaiter, created := sBufio.CreatePacketBatchReadWaiter(packetConn); created {
+				l.loopUDPInBatch(batchHandler, readWaiter)
+				return
+			}
+		}
+	}
 	var buffer *buf.Buffer
 	if !l.threadUnsafePacketWriter {
 		buffer = buf.NewPacket()
@@ -126,7 +138,7 @@ func (l *Listener) loopUDPIn() {
 				return
 			}
 			buffer.Truncate(n)
-			l.oobPacketHandler.NewPacketEx(buffer, oob[:oobN], M.SocksaddrFromNetIP(addr).Unwrap())
+			l.oobPacketHandler.NewPacket(buffer, oob[:oobN], M.SocksaddrFromNetIP(addr).Unwrap())
 		}
 	} else {
 		for {
@@ -148,37 +160,82 @@ func (l *Listener) loopUDPIn() {
 				return
 			}
 			buffer.Truncate(n)
-			l.packetHandler.NewPacketEx(buffer, M.SocksaddrFromNetIP(addr).Unwrap())
+			l.packetHandler.NewPacket(buffer, M.SocksaddrFromNetIP(addr).Unwrap())
 		}
 	}
 }
 
+func (l *Listener) loopUDPInBatch(handler adapter.PacketBatchHandler, readWaiter N.PacketBatchReadWaiter) {
+	readWaitOptions := N.ReadWaitOptions{
+		BatchSize: sBufio.DefaultPacketReadBatchSize,
+	}
+	readWaiter.InitializeReadWaiter(readWaitOptions)
+	for {
+		buffers, sources, err := readWaiter.WaitReadPackets()
+		if err != nil {
+			buf.ReleaseMulti(buffers)
+			if l.shutdown.Load() && E.IsClosed(err) {
+				return
+			}
+			l.udpConn.Close()
+			l.logger.Error("udp listener closed: ", err)
+			return
+		}
+		handler.NewPacketBatch(buffers, sources)
+	}
+}
+
 func (l *Listener) loopUDPOut() {
+	packetConn := sBufio.NewPacketConn(l.udpConn)
+	batchWriter := sBufio.NewPacketBatchWriter(packetConn)
+	packets := make([]*N.PacketBuffer, 0, udpOutputBatchSize)
+	buffers := make([]*buf.Buffer, 0, udpOutputBatchSize)
+	destinations := make([]M.Socksaddr, 0, udpOutputBatchSize)
 	for {
 		select {
 		case packet := <-l.packetOutbound:
-			destination := packet.Destination.AddrPort()
-			_, err := l.udpConn.WriteToUDPAddrPort(packet.Buffer.Bytes(), destination)
-			packet.Buffer.Release()
-			N.PutPacketBuffer(packet)
-			if err != nil {
-				if l.shutdown.Load() && E.IsClosed(err) {
-					return
-				}
-				l.logger.Error("udp listener write back: ", destination, ": ", err)
-				continue
-			}
-			continue
+			packets = append(packets, packet)
 		case <-l.packetOutboundClosed:
+			l.releasePacketOutbound()
+			return
 		}
-		for {
+	drain:
+		for len(packets) < udpOutputBatchSize {
 			select {
 			case packet := <-l.packetOutbound:
-				packet.Buffer.Release()
-				N.PutPacketBuffer(packet)
+				packets = append(packets, packet)
 			default:
+				break drain
+			}
+		}
+		for _, packet := range packets {
+			buffers = append(buffers, packet.Buffer)
+			destinations = append(destinations, packet.Destination)
+		}
+		err := batchWriter.WritePacketBatch(buffers, destinations)
+		for _, packet := range packets {
+			N.PutPacketBuffer(packet)
+		}
+		packets = packets[:0]
+		buffers = buffers[:0]
+		destinations = destinations[:0]
+		if err != nil {
+			if l.shutdown.Load() && E.IsClosed(err) {
 				return
 			}
+			l.logger.Error("udp listener write back: ", err)
+		}
+	}
+}
+
+func (l *Listener) releasePacketOutbound() {
+	for {
+		select {
+		case packet := <-l.packetOutbound:
+			packet.Buffer.Release()
+			N.PutPacketBuffer(packet)
+		default:
+			return
 		}
 	}
 }
@@ -203,5 +260,30 @@ func (w *packetWriter) WritePacket(buffer *buf.Buffer, destination M.Socksaddr)
 	}
 }
 
+func (w *packetWriter) WritePacketBatch(buffers []*buf.Buffer, destinations []M.Socksaddr) error {
+	if len(buffers) == 0 || len(buffers) != len(destinations) {
+		buf.ReleaseMulti(buffers)
+		return os.ErrInvalid
+	}
+	for index, buffer := range buffers {
+		packet := N.NewPacketBuffer()
+		packet.Buffer = buffer
+		packet.Destination = destinations[index]
+		select {
+		case w.packetOutbound <- packet:
+		default:
+			buffer.Release()
+			N.PutPacketBuffer(packet)
+			buf.ReleaseMulti(buffers[index+1:])
+			if w.shutdown.Load() {
+				return os.ErrClosed
+			}
+			w.logger.Trace("dropped packet batch to ", destinations[index])
+			return nil
+		}
+	}
+	return nil
+}
+
 func (w *packetWriter) WriteIsThreadUnsafe() {
 }

common/mux/router.go

@@ -42,7 +42,7 @@ func NewRouterWithOptions(router adapter.ConnectionRouterEx, logger logger.Conte
 			return log.ContextWithNewID(ctx)
 		},
 		Logger:    logger,
-		HandlerEx: adapter.NewRouteContextHandlerEx(router),
+		HandlerEx: adapter.NewRouteContextHandler(router),
 		Padding:   options.Padding,
 		Brutal:    brutalOptions,
 	})