conn: split read buffer on GSO

Signed-off-by: Alex Valiushko <alexvaliushko@tailscale.com>
Change-Id: If61dbfce4174bd07f12996c47f495e426a6a6964

Author: illotum

conn/bind_std.go

@@ -208,7 +208,7 @@ func (s *StdNetBind) putMessages(msgs *[]ipv6.Message) {
 		// Non coalesced write paths access only batch.msgs[i].Buffers[0],
 		// but we append more during [coalesceMessages].
 		// Leave index zero accessible:
-		(*msgs)[i] = ipv6.Message{Buffers: (*msgs)[i].Buffers[:1], OOB: (*msgs)[i].OOB}
+		(*msgs)[i] = ipv6.Message{Buffers: (*msgs)[i].Buffers[:1], OOB: (*msgs)[i].OOB[:cap((*msgs)[i].OOB)]}
 	}
 	s.msgsPool.Put(msgs)
 }
@@ -236,53 +236,92 @@ func (s *StdNetBind) receiveIP(
 	rxOffload bool,
 	bufs []iobuf.View,
 	eps []Endpoint,
-) (n int, err error) {
+) (int, error) {
 	msgs := s.getMessages()
-	// TODO: placeholder until bind implements right-sized buffers.
-	iobuf.EnsureAllocated(bufs)
-	for i := range bufs {
-		(*msgs)[i].Buffers[0] = bufs[i].Bytes
-		(*msgs)[i].OOB = (*msgs)[i].OOB[:cap((*msgs)[i].OOB)]
-	}
 	defer s.putMessages(msgs)
-	var numMsgs int
-	if runtime.GOOS == "linux" {
-		if rxOffload {
-			readAt := len(*msgs) - 2
-			numMsgs, err = br.ReadBatch((*msgs)[readAt:], 0)
-			if err != nil {
-				return 0, err
+	readData := make([]*iobuf.Shared, len(*msgs)) // TODO: optimize this allocation
+	defer func() {
+		for i := range readData {
+			if readData[i] != nil {
+				// release backing arrays for any unclaimed buffers
+				readData[i].Release()
+				readData[i] = nil
+			}
+		}
+	}()
+
+	switch {
+	case runtime.GOOS == "linux" && rxOffload:
+		const readBatchSize = 2
+		for i := range readBatchSize {
+			readData[i] = iobuf.SharedBufPool.Get()
+			(*msgs)[i].Buffers[0] = readData[i].Bytes[:]
+		}
+		msgsN, err := br.ReadBatch((*msgs)[:readBatchSize], 0)
+		if err != nil {
+			return 0, err // expect atomic reads
+		}
+		var n int
+		for i, msg := range (*msgs)[:msgsN] {
+			if msg.N == 0 {
+				continue
 			}
-			numMsgs, err = splitCoalescedMessages(*msgs, readAt, getGSOSize)
+			gsoSize, err := getGSOSize(msg.OOB[:msg.NN])
 			if err != nil {
-				return 0, err
+				return n, err
 			}
-		} else {
-			numMsgs, err = br.ReadBatch(*msgs, 0)
+			if gsoSize == 0 {
+				gsoSize = msg.N
+			}
+			split, err := readData[i].SplitCoalesced(bufs[n:], gsoSize, msg.N)
 			if err != nil {
-				return 0, err
+				return n, err
+			}
+			addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
+			ep := &StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation
+			getSrcFromControl(msg.OOB[:msg.NN], ep)
+			for j := n; j < n+split; j++ {
+				eps[j] = ep
 			}
+			n += split
 		}
-	} else {
+		return n, nil
+	case runtime.GOOS == "linux":
+		readBatchSize := min(len(bufs), len(*msgs))
+		for i := range readBatchSize {
+			readData[i] = iobuf.SharedBufPool.Get()
+			(*msgs)[i].Buffers[0] = readData[i].Bytes[:]
+		}
+		msgsN, err := br.ReadBatch((*msgs)[:readBatchSize], 0)
+		if err != nil {
+			return 0, err // expect atomic reads
+		}
+		var n int
+		for i, msg := range (*msgs)[:msgsN] {
+			if msg.N == 0 {
+				continue
+			}
+			readData[i].Refer(&bufs[n], 0, msg.N)
+			addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
+			ep := &StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation
+			getSrcFromControl(msg.OOB[:msg.NN], ep)
+			eps[n] = ep
+			n++
+		}
+		return n, nil
+	default:
 		msg := &(*msgs)[0]
-		msg.N, msg.NN, _, msg.Addr, err = conn.ReadMsgUDP(msg.Buffers[0], msg.OOB)
+		readData[0] = iobuf.SharedBufPool.Get()
+		n, _, _, addr, err := conn.ReadMsgUDP(readData[0].Bytes[:], msg.OOB)
 		if err != nil {
 			return 0, err
 		}
-		numMsgs = 1
-	}
-	for i := 0; i < numMsgs; i++ {
-		msg := &(*msgs)[i]
-		bufs[i].Bytes = bufs[i].Bytes[:msg.N]
-		if len(bufs[i].Bytes) == 0 {
-			continue
-		}
-		addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
-		ep := &StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation
+		readData[0].Refer(&bufs[0], 0, n)
+		ep := &StdNetEndpoint{AddrPort: addr.AddrPort()}
 		getSrcFromControl(msg.OOB[:msg.NN], ep)
-		eps[i] = ep
+		eps[0] = ep
+		return 1, nil
 	}
-	return numMsgs, nil
 }
 
 func (s *StdNetBind) makeReceiveIPv4(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {
@@ -523,49 +562,3 @@ func coalesceMessages(addr *net.UDPAddr, ep *StdNetEndpoint, bufs [][]byte, offs
 	}
 	return base + 1
 }
-
-type getGSOFunc func(control []byte) (int, error)
-
-func splitCoalescedMessages(msgs []ipv6.Message, firstMsgAt int, getGSO getGSOFunc) (n int, err error) {
-	for i := firstMsgAt; i < len(msgs); i++ {
-		msg := &msgs[i]
-		if msg.N == 0 {
-			return n, err
-		}
-		var (
-			gsoSize    int
-			start      int
-			end        = msg.N
-			numToSplit = 1
-		)
-		gsoSize, err = getGSO(msg.OOB[:msg.NN])
-		if err != nil {
-			return n, err
-		}
-		if gsoSize > 0 {
-			numToSplit = (msg.N + gsoSize - 1) / gsoSize
-			end = gsoSize
-		}
-		for j := 0; j < numToSplit; j++ {
-			if n > i {
-				return n, errors.New("splitting coalesced packet resulted in overflow")
-			}
-			copied := copy(msgs[n].Buffers[0], msg.Buffers[0][start:end])
-			msgs[n].N = copied
-			msgs[n].Addr = msg.Addr
-			start = end
-			end += gsoSize
-			if end > msg.N {
-				end = msg.N
-			}
-			n++
-		}
-		if i != n-1 {
-			// It is legal for bytes to move within msg.Buffers[0] as a result
-			// of splitting, so we only zero the source msg len when it is not
-			// the destination of the last split operation above.
-			msg.N = 0
-		}
-	}
-	return n, nil
-}

conn/bind_std_test.go

@@ -135,123 +135,3 @@ func mockGetGSOSize(control []byte) (int, error) {
 	}
 	return int(binary.LittleEndian.Uint16(control)), nil
 }
-
-func Test_splitCoalescedMessages(t *testing.T) {
-	newMsg := func(n, gso int) ipv6.Message {
-		msg := ipv6.Message{
-			Buffers: [][]byte{make([]byte, 1<<16-1)},
-			N:       n,
-			OOB:     make([]byte, 2),
-		}
-		binary.LittleEndian.PutUint16(msg.OOB, uint16(gso))
-		if gso > 0 {
-			msg.NN = 2
-		}
-		return msg
-	}
-
-	cases := []struct {
-		name        string
-		msgs        []ipv6.Message
-		firstMsgAt  int
-		wantNumEval int
-		wantMsgLens []int
-		wantErr     bool
-	}{
-		{
-			name: "second last split last empty",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(3, 1),
-				newMsg(0, 0),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 3,
-			wantMsgLens: []int{1, 1, 1, 0},
-			wantErr:     false,
-		},
-		{
-			name: "second last no split last empty",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(1, 0),
-				newMsg(0, 0),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 1,
-			wantMsgLens: []int{1, 0, 0, 0},
-			wantErr:     false,
-		},
-		{
-			name: "second last no split last no split",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(1, 0),
-				newMsg(1, 0),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 2,
-			wantMsgLens: []int{1, 1, 0, 0},
-			wantErr:     false,
-		},
-		{
-			name: "second last no split last split",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(1, 0),
-				newMsg(3, 1),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 4,
-			wantMsgLens: []int{1, 1, 1, 1},
-			wantErr:     false,
-		},
-		{
-			name: "second last split last split",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(2, 1),
-				newMsg(2, 1),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 4,
-			wantMsgLens: []int{1, 1, 1, 1},
-			wantErr:     false,
-		},
-		{
-			name: "second last no split last split overflow",
-			msgs: []ipv6.Message{
-				newMsg(0, 0),
-				newMsg(0, 0),
-				newMsg(1, 0),
-				newMsg(4, 1),
-			},
-			firstMsgAt:  2,
-			wantNumEval: 4,
-			wantMsgLens: []int{1, 1, 1, 1},
-			wantErr:     true,
-		},
-	}
-
-	for _, tt := range cases {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := splitCoalescedMessages(tt.msgs, 2, mockGetGSOSize)
-			if err != nil && !tt.wantErr {
-				t.Fatalf("err: %v", err)
-			}
-			if got != tt.wantNumEval {
-				t.Fatalf("got to eval: %d want: %d", got, tt.wantNumEval)
-			}
-			for i, msg := range tt.msgs {
-				if msg.N != tt.wantMsgLens[i] {
-					t.Fatalf("msg[%d].N: %d want: %d", i, msg.N, tt.wantMsgLens[i])
-				}
-			}
-		})
-	}
-}