make batchConn more generic

Author: zjx20

batchconn.go

@@ -10,3 +10,8 @@ type batchConn interface {
 	WriteBatch(ms []ipv4.Message, flags int) (int, error)
 	ReadBatch(ms []ipv4.Message, flags int) (int, error)
 }
+
+type batchErrDetector interface {
+	ReadBatchUnavailable(err error) bool
+	WriteBatchUnavailable(err error) bool
+}

batchconn_generic.go

@@ -0,0 +1,28 @@
+// +build !linux
+
+package kcp
+
+import (
+	"net"
+)
+
+func toBatchConn(c net.PacketConn) batchConn {
+	if xconn, ok := c.(batchConn); ok {
+		return xconn
+	}
+	return nil
+}
+
+func readBatchUnavailable(xconn batchConn, err error) bool {
+	if detector, ok := xconn.(batchErrDetector); ok {
+		return detector.ReadBatchUnavailable(err)
+	}
+	return false
+}
+
+func writeBatchUnavailable(xconn batchConn, err error) bool {
+	if detector, ok := xconn.(batchErrDetector); ok {
+		return detector.WriteBatchUnavailable(err)
+	}
+	return false
+}

batchconn_linux.go

@@ -0,0 +1,80 @@
+// +build linux
+
+package kcp
+
+import (
+	"net"
+	"os"
+
+	"golang.org/x/net/ipv4"
+	"golang.org/x/net/ipv6"
+)
+
+func toBatchConn(c net.PacketConn) batchConn {
+	if xconn, ok := c.(batchConn); ok {
+		return xconn
+	}
+	if _, ok := c.(*net.UDPConn); ok {
+		var xconn batchConn
+		addr, err := net.ResolveUDPAddr("udp", c.LocalAddr().String())
+		if err == nil {
+			if addr.IP.To4() != nil {
+				xconn = ipv4.NewPacketConn(c)
+			} else {
+				xconn = ipv6.NewPacketConn(c)
+			}
+		}
+		return xconn
+	}
+	return nil
+}
+
+func isPacketConn(xconn batchConn) bool {
+	if _, ok := xconn.(*ipv4.PacketConn); ok {
+		return true
+	}
+	if _, ok := xconn.(*ipv6.PacketConn); ok {
+		return true
+	}
+	return false
+}
+
+func readBatchUnavailable(xconn batchConn, err error) bool {
+	if isPacketConn(xconn) {
+		// compatibility issue:
+		// for linux kernel<=2.6.32, support for sendmmsg is not available
+		// an error of type os.SyscallError will be returned
+		if operr, ok := err.(*net.OpError); ok {
+			if se, ok := operr.Err.(*os.SyscallError); ok {
+				if se.Syscall == "recvmmsg" {
+					return true
+				}
+			}
+		}
+		return false
+	}
+	if detector, ok := xconn.(batchErrDetector); ok {
+		return detector.ReadBatchUnavailable(err)
+	}
+	return false
+}
+
+func writeBatchUnavailable(xconn batchConn, err error) bool {
+	if isPacketConn(xconn) {
+		// compatibility issue:
+		// for linux kernel<=2.6.32, support for sendmmsg is not available
+		// an error of type os.SyscallError will be returned
+		if operr, ok := err.(*net.OpError); ok {
+			if se, ok := operr.Err.(*os.SyscallError); ok {
+				if se.Syscall == "sendmmsg" {
+					return true
+				}
+			}
+		}
+		return false
+	}
+	if detector, ok := xconn.(batchErrDetector); ok {
+		return detector.WriteBatchUnavailable(err)
+	}
+	return false
+}

readloop.go

@@ -4,19 +4,42 @@ import (
 	"sync/atomic"
 
 	"github.com/pkg/errors"
+	"golang.org/x/net/ipv4"
 )
 
+func (s *UDPSession) readLoop() {
+	// default version
+	if s.xconn == nil {
+		s.defaultReadLoop()
+		return
+	}
+	s.batchReadLoop()
+}
+
+func (l *Listener) monitor() {
+	xconn := toBatchConn(l.conn)
+
+	// default version
+	if xconn == nil {
+		l.defaultMonitor()
+		return
+	}
+	l.batchMonitor(xconn)
+}
+
 func (s *UDPSession) defaultReadLoop() {
 	buf := make([]byte, mtuLimit)
 	var src string
 	for {
 		if n, addr, err := s.conn.ReadFrom(buf); err == nil {
 			// make sure the packet is from the same source
-			if src == "" { // set source address
-				src = addr.String()
-			} else if addr.String() != src {
-				atomic.AddUint64(&DefaultSnmp.InErrs, 1)
-				continue
+			if addr.String() != src {
+				if len(src) == 0 { // set source address
+					src = addr.String()
+				} else {
+					atomic.AddUint64(&DefaultSnmp.InErrs, 1)
+					continue
+				}
 			}
 			s.packetInput(buf[:n])
 		} else {
@@ -37,3 +60,63 @@ func (l *Listener) defaultMonitor() {
 		}
 	}
 }
+
+func (s *UDPSession) batchReadLoop() {
+	// x/net version
+	var src string
+	msgs := make([]ipv4.Message, batchSize)
+	for k := range msgs {
+		msgs[k].Buffers = [][]byte{make([]byte, mtuLimit)}
+	}
+
+	for {
+		if count, err := s.xconn.ReadBatch(msgs, 0); err == nil {
+			for i := 0; i < count; i++ {
+				msg := &msgs[i]
+				// make sure the packet is from the same source
+				if msg.Addr.String() != src {
+					if len(src) == 0 { // set source address if nil
+						src = msg.Addr.String()
+					} else {
+						atomic.AddUint64(&DefaultSnmp.InErrs, 1)
+						continue
+					}
+				}
+
+				// source and size has validated
+				s.packetInput(msg.Buffers[0][:msg.N])
+			}
+		} else {
+			if readBatchUnavailable(s.xconn, err) {
+				s.defaultReadLoop()
+				return
+			}
+			s.notifyReadError(errors.WithStack(err))
+			return
+		}
+	}
+}
+
+func (l *Listener) batchMonitor(xconn batchConn) {
+	// x/net version
+	msgs := make([]ipv4.Message, batchSize)
+	for k := range msgs {
+		msgs[k].Buffers = [][]byte{make([]byte, mtuLimit)}
+	}
+
+	for {
+		if count, err := xconn.ReadBatch(msgs, 0); err == nil {
+			for i := 0; i < count; i++ {
+				msg := &msgs[i]
+				l.packetInput(msg.Buffers[0][:msg.N], msg.Addr)
+			}
+		} else {
+			if readBatchUnavailable(xconn, err) {
+				l.defaultMonitor()
+				return
+			}
+			l.notifyReadError(errors.WithStack(err))
+			return
+		}
+	}
+}

readloop_generic.go

@@ -139,17 +139,8 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn
 	sess.block = block
 	sess.recvbuf = make([]byte, mtuLimit)
 
-	// cast to writebatch conn
-	if _, ok := conn.(*net.UDPConn); ok {
-		addr, err := net.ResolveUDPAddr("udp", conn.LocalAddr().String())
-		if err == nil {
-			if addr.IP.To4() != nil {
-				sess.xconn = ipv4.NewPacketConn(conn)
-			} else {
-				sess.xconn = ipv6.NewPacketConn(conn)
-			}
-		}
-	}
+	// cast to batchConn, can be nil
+	sess.xconn = toBatchConn(conn)
 
 	// FEC codec initialization
 	sess.fecDecoder = newFECDecoder(dataShards, parityShards)