reduce heap allocations #15 (#368)

Summary:
Pull Request resolved: #368

Use highly optimized `ReadPacketBuf` instead of `ReadFromUDP` which is very similar to what we do in `ptp4u`.

Reviewed By: abulimov

Differential Revision: D58703243

fbshipit-source-id: 286f552762616b074f18f52d01c614a70869f8a6

ptp/sptp/client/connection.go

@@ -28,18 +28,24 @@ import (
 	"github.com/facebook/time/timestamp"
 )
 
-// UDPConn describes what functionality we expect from UDP connection
-type UDPConn interface {
-	ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
+// UDPConnNoTS describes what functionality we expect from UDP connection
+type UDPConnNoTS interface {
 	WriteTo(b []byte, addr net.Addr) (int, error)
+	ReadPacketBuf(buf []byte) (int, string, error)
 	Close() error
 }
 
 // UDPConnWithTS describes what functionality we expect from UDP connection that allows us to read TX timestamps
 type UDPConnWithTS interface {
-	UDPConn
 	WriteToWithTS(b []byte, addr net.Addr) (int, time.Time, error)
 	ReadPacketWithRXTimestampBuf(buf, oob []byte) (int, unix.Sockaddr, time.Time, error)
+	Close() error
+}
+
+// UDPConn is a wrapper around udp connection and a corresponding fd
+type UDPConn struct {
+	*net.UDPConn
+	connFd int
 }
 
 // UDPConnTS is a wrapper around udp connection and a corresponding fd
@@ -49,6 +55,25 @@ type UDPConnTS struct {
 	l      sync.Mutex
 }
 
+// NewUDPConn initialises a new struct UDPConn
+func NewUDPConn(conn *net.UDPConn) (*UDPConn, error) {
+	// get FD of the connection. Can be optimized by doing this when connection is created
+	connFd, err := timestamp.ConnFd(conn)
+	if err != nil {
+		return nil, err
+	}
+
+	// set it to blocking mode, otherwise recvmsg will just return with nothing most of the time
+	if err = unix.SetNonblock(connFd, false); err != nil {
+		return nil, fmt.Errorf("failed to set event socket to blocking: %w", err)
+	}
+
+	return &UDPConn{
+		UDPConn: conn,
+		connFd:  connFd,
+	}, nil
+}
+
 // NewUDPConnTS initialises a new struct UDPConnTS
 func NewUDPConnTS(conn *net.UDPConn, connFd int) *UDPConnTS {
 	return &UDPConnTS{
@@ -114,3 +139,13 @@ func (c *UDPConnTS) WriteToWithTS(b []byte, addr net.Addr) (int, time.Time, erro
 func (c *UDPConnTS) ReadPacketWithRXTimestampBuf(buf, oob []byte) (int, unix.Sockaddr, time.Time, error) {
 	return timestamp.ReadPacketWithRXTimestampBuf(c.connFd, buf, oob)
 }
+
+// ReadPacketBuf reads bytes from underlying fd
+func (c *UDPConn) ReadPacketBuf(buf []byte) (int, string, error) {
+	n, saddr, err := unix.Recvfrom(c.connFd, buf, 0)
+	if err != nil {
+		return 0, "", err
+	}
+
+	return n, timestamp.SockaddrToIP(saddr).String(), err
+}

ptp/sptp/client/sptp.go

@@ -62,7 +62,7 @@ type SPTP struct {
 	lastTick   time.Time
 
 	clockID ptp.ClockIdentity
-	genConn UDPConn
+	genConn UDPConnNoTS
 	// listening connection on port 319
 	eventConn UDPConnWithTS
 }
@@ -132,7 +132,10 @@ func (p *SPTP) init() error {
 	if err != nil {
 		return err
 	}
-	p.genConn = genConn
+	p.genConn, err = NewUDPConn(genConn)
+	if err != nil {
+		return err
+	}
 	// bind to event port
 	eventConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("::"), Port: ptp.PortEvent})
 	if err != nil {
@@ -219,17 +222,17 @@ func (p *SPTP) RunListener(ctx context.Context) error {
 			announce := &ptp.Announce{}
 			buf := make([]byte, timestamp.PayloadSizeBytes)
 			for {
-				bbuf, addr, err := p.genConn.ReadFromUDP(buf)
+				bbuf, addr, err := p.genConn.ReadPacketBuf(buf)
 				if err != nil {
 					doneChan <- err
 					return
 				}
-				if addr == nil {
+				if addr == "" {
 					doneChan <- fmt.Errorf("received packet on port 320 with nil source address")
 					return
 				}
 				log.Debugf("got packet on port 320, n = %v, addr = %v", bbuf, addr)
-				cc, found := p.clients[addr.IP.String()]
+				cc, found := p.clients[addr]
 				if !found {
 					log.Warningf("ignoring packets from server %v", addr)
 					continue

ptp/sptp/client/sptp_test.go

@@ -388,8 +388,8 @@ func TestRunListenerNoAddr(t *testing.T) {
 	defer ctrl.Finish()
 	mockEventConn := NewMockUDPConnWithTS(ctrl)
 	mockEventConn.EXPECT().ReadPacketWithRXTimestampBuf(gomock.Any(), gomock.Any()).AnyTimes()
-	mockGenConn := NewMockUDPConn(ctrl)
-	mockGenConn.EXPECT().ReadFromUDP(gomock.Any()).AnyTimes()
+	mockGenConn := NewMockUDPConnNoTS(ctrl)
+	mockGenConn.EXPECT().ReadPacketBuf(gomock.Any()).AnyTimes()
 	mockClock := NewMockClock(ctrl)
 	mockServo := NewMockServo(ctrl)
 	mockStatsServer := NewMockStatsServer(ctrl)
@@ -421,8 +421,8 @@ func TestRunListenerError(t *testing.T) {
 	defer ctrl.Finish()
 	mockEventConn := NewMockUDPConnWithTS(ctrl)
 	mockEventConn.EXPECT().ReadPacketWithRXTimestampBuf(gomock.Any(), gomock.Any()).Return(0, &unix.SockaddrInet6{}, time.Time{}, fmt.Errorf("some error")).AnyTimes()
-	mockGenConn := NewMockUDPConn(ctrl)
-	mockGenConn.EXPECT().ReadFromUDP(gomock.Any()).Return(2, nil, fmt.Errorf("some error")).AnyTimes()
+	mockGenConn := NewMockUDPConnNoTS(ctrl)
+	mockGenConn.EXPECT().ReadPacketBuf(gomock.Any()).Return(2, "", fmt.Errorf("some error")).AnyTimes()
 	mockClock := NewMockClock(ctrl)
 	mockServo := NewMockServo(ctrl)
 	mockStatsServer := NewMockStatsServer(ctrl)
@@ -480,28 +480,24 @@ func TestRunListenerGood(t *testing.T) {
 		return len(syncBytes), &unix.SockaddrInet4{Addr: addrBytes, Port: 319}, time.Now(), nil
 	}).AnyTimes()
 
-	mockGenConn := NewMockUDPConn(ctrl)
-	mockGenConn.EXPECT().ReadFromUDP(gomock.Any()).DoAndReturn(func(b []byte) (int, *net.UDPAddr, error) {
+	mockGenConn := NewMockUDPConnNoTS(ctrl)
+	mockGenConn.EXPECT().ReadPacketBuf(gomock.Any()).DoAndReturn(func(b []byte) (int, string, error) {
 		// limit how many we send, so we don't overwhelm the client. packets from uknown IPs will be discarded
 		if sentGen > 10 {
-			return 0, &net.UDPAddr{}, nil
+			return 0, "whatever", nil
 		}
 		sentGen++
 		addr := "192.168.0.11"
 		if sentGen%2 == 0 {
 			addr = "192.168.0.10"
 		}
-		udpAddr := &net.UDPAddr{
-			IP:   net.ParseIP(addr),
-			Port: 320,
-		}
 		clear(b)
 		b = announceBytes
 		b[0] = 9
 		b[1] = 3
 		b[2] = 7
 		b[3] = 4
-		return len(announceBytes), udpAddr, nil
+		return len(announceBytes), addr, nil
 	}).AnyTimes()
 
 	mockClock := NewMockClock(ctrl)