[system-probe] Add packet counting to existing UDP byte counters

pkg/network/ebpf/c/co-re/tracer-fentry.c

@@ -201,7 +201,7 @@ int BPF_PROG(udp_sendpage_exit, struct sock *sk, struct page *page, int offset,
         return 0;
     }
 
-    return handle_message(&t, sent, 0, CONN_DIRECTION_UNKNOWN, 0, 0, PACKET_COUNT_NONE, sk);
+    return handle_message(&t, sent, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_INCREMENT, sk);
 }
 
 SEC("fexit/tcp_recvmsg")
@@ -262,7 +262,7 @@ static __always_inline int handle_udp_send(struct sock *sk, int sent) {
 
     if (sent > 0) {
         log_debug("udp_sendmsg: sent: %d", sent);
-        handle_message(t, sent, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_NONE, sk);
+        handle_message(t, sent, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_INCREMENT, sk);
     }
 
     bpf_map_delete_elem(&udp_send_skb_args, &pid_tgid);

pkg/network/ebpf/c/tracer.c

@@ -187,7 +187,7 @@ int BPF_BYPASSABLE_KRETPROBE(kretprobe__udp_sendpage, int sent) {
         return 0;
     }
 
-    return handle_message(&t, sent, 0, CONN_DIRECTION_UNKNOWN, 0, 0, PACKET_COUNT_NONE, skp);
+    return handle_message(&t, sent, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_INCREMENT, skp);
 }
 
 SEC("kprobe/tcp_done")
@@ -420,7 +420,7 @@ static __always_inline int handle_ip6_skb(struct sock *sk, size_t size, struct f
     }
 
     log_debug("kprobe/ip6_make_skb: pid_tgid: %llu, size: %zu", pid_tgid, size);
-    handle_message(&t, size, 0, CONN_DIRECTION_UNKNOWN, 0, 0, PACKET_COUNT_NONE, sk);
+    handle_message(&t, size, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_INCREMENT, sk);
     increment_telemetry_count(udp_send_processed);
 
     return 0;
@@ -597,9 +597,7 @@ static __always_inline int handle_ip_skb(struct sock *sk, size_t size, struct fl
 
     log_debug("kprobe/ip_make_skb: pid_tgid: %llu, size: %zu", pid_tgid, size);
 
-    // segment count is not currently enabled on prebuilt.
-    // to enable, change PACKET_COUNT_NONE => PACKET_COUNT_INCREMENT
-    handle_message(&t, size, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_NONE, sk);
+    handle_message(&t, size, 0, CONN_DIRECTION_UNKNOWN, 1, 0, PACKET_COUNT_INCREMENT, sk);
     increment_telemetry_count(udp_send_processed);
 
     return 0;
@@ -767,9 +765,7 @@ static __always_inline int handle_ret_udp_recvmsg_pre_4_7_0(int copied, void *ud
     bpf_map_delete_elem(udp_sock_map, &pid_tgid);
 
     log_debug("kretprobe/udp_recvmsg: pid_tgid: %llu, return: %d", pid_tgid, copied);
-    // segment count is not currently enabled on prebuilt.
-    // to enable, change PACKET_COUNT_NONE => PACKET_COUNT_INCREMENT
-    handle_message(&t, 0, copied, CONN_DIRECTION_UNKNOWN, 0, 1, PACKET_COUNT_NONE, st->sk);
+    handle_message(&t, 0, copied, CONN_DIRECTION_UNKNOWN, 0, 1, PACKET_COUNT_INCREMENT, st->sk);
 
     return 0;
 }

pkg/network/tracer/tracer_linux_test.go

@@ -1126,31 +1126,61 @@ func (s *TracerSuite) TestSelfConnect() {
 	}, 5*time.Second, 100*time.Millisecond, "could not find expected number of tcp connections, expected: 2")
 }
 
-func (s *TracerSuite) TestUDPPeekCount() {
-	t := s.T()
-	config := testConfig()
-	tr := setupTracer(t, config)
+// sets up two udp sockets talking to each other locally.
+// returns (listener, dialer)
+func setupUdpSockets(t *testing.T, udpnet, ip string) (*net.UDPConn, *net.UDPConn) {
+	serverAddr := fmt.Sprintf("%s:%d", ip, 0)
 
-	ln, err := net.ListenPacket("udp", "127.0.0.1:0")
+	laddr, err := net.ResolveUDPAddr(udpnet, serverAddr)
 	require.NoError(t, err)
-	defer ln.Close()
 
-	saddr := ln.LocalAddr().String()
+	var ln, c *net.UDPConn = nil, nil
+	t.Cleanup(func() {
+		if ln != nil {
+			ln.Close()
+		}
+		if c != nil {
+			c.Close()
+		}
+	})
 
-	laddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
+	ln, err = net.ListenUDP(udpnet, laddr)
 	require.NoError(t, err)
-	raddr, err := net.ResolveUDPAddr("udp", saddr)
+
+	saddr := ln.LocalAddr().String()
+
+	raddr, err := net.ResolveUDPAddr(udpnet, saddr)
 	require.NoError(t, err)
 
-	c, err := net.DialUDP("udp", laddr, raddr)
+	c, err = net.DialUDP(udpnet, laddr, raddr)
 	require.NoError(t, err)
-	defer c.Close()
+
+	return ln, c
+}
+
+func (s *TracerSuite) TestUDPPeekCount() {
+	t := s.T()
+	t.Run("v4", func(t *testing.T) {
+		testUDPPeekCount(t, "udp4", "127.0.0.1")
+	})
+	t.Run("v6", func(t *testing.T) {
+		if !testConfig().CollectUDPv6Conns {
+			t.Skip("UDPv6 disabled")
+		}
+		testUDPPeekCount(t, "udp6", "[::1]")
+	})
+}
+func testUDPPeekCount(t *testing.T, udpnet, ip string) {
+	config := testConfig()
+	tr := setupTracer(t, config)
+
+	ln, c := setupUdpSockets(t, udpnet, ip)
 
 	msg := []byte("asdf")
-	_, err = c.Write(msg)
+	_, err := c.Write(msg)
 	require.NoError(t, err)
 
-	rawConn, err := ln.(*net.UDPConn).SyscallConn()
+	rawConn, err := ln.SyscallConn()
 	require.NoError(t, err)
 	err = rawConn.Control(func(fd uintptr) {
 		buf := make([]byte, 1024)
@@ -1203,12 +1233,82 @@ func (s *TracerSuite) TestUDPPeekCount() {
 	m := outgoing.Monotonic
 	require.Equal(t, len(msg), int(m.SentBytes))
 	require.Equal(t, 0, int(m.RecvBytes))
+	require.Equal(t, 1, int(m.SentPackets))
+	require.Equal(t, 0, int(m.RecvPackets))
 	require.True(t, outgoing.IntraHost)
 
 	// make sure the inverse values are seen for the other message
 	m = incoming.Monotonic
 	require.Equal(t, 0, int(m.SentBytes))
 	require.Equal(t, len(msg), int(m.RecvBytes))
+	require.Equal(t, 0, int(m.SentPackets))
+	require.Equal(t, 1, int(m.RecvPackets))
+	require.True(t, incoming.IntraHost)
+}
+
+func (s *TracerSuite) TestUDPPacketSumming() {
+	t := s.T()
+	t.Run("v4", func(t *testing.T) {
+		testUDPPacketSumming(t, "udp4", "127.0.0.1")
+	})
+	t.Run("v6", func(t *testing.T) {
+		if !testConfig().CollectUDPv6Conns {
+			t.Skip("UDPv6 disabled")
+		}
+		testUDPPacketSumming(t, "udp6", "[::1]")
+	})
+}
+func testUDPPacketSumming(t *testing.T, udpnet, ip string) {
+	config := testConfig()
+	tr := setupTracer(t, config)
+
+	ln, c := setupUdpSockets(t, udpnet, ip)
+
+	msg := []byte("asdf")
+	// send UDP packets of increasing length
+	for i := range msg {
+		_, err := c.Write(msg[:i+1])
+		require.NoError(t, err)
+	}
+	expectedBytes := 1 + 2 + 3 + 4
+
+	buf := make([]byte, 256)
+	recvBytes := 0
+	for range msg {
+		n, _, err := ln.ReadFrom(buf)
+		require.NoError(t, err)
+		recvBytes += n
+	}
+	// sanity check: did userspace get all four expected packets?
+	require.Equal(t, recvBytes, expectedBytes)
+
+	var incoming *network.ConnectionStats
+	var outgoing *network.ConnectionStats
+	require.Eventuallyf(t, func() bool {
+		conns := getConnections(t, tr)
+		if outgoing == nil {
+			outgoing, _ = findConnection(c.LocalAddr(), c.RemoteAddr(), conns)
+		}
+		if incoming == nil {
+			incoming, _ = findConnection(c.RemoteAddr(), c.LocalAddr(), conns)
+		}
+
+		return outgoing != nil && incoming != nil
+	}, 3*time.Second, 100*time.Millisecond, "couldn't find incoming and outgoing connections matching")
+
+	m := outgoing.Monotonic
+	require.Equal(t, expectedBytes, int(m.SentBytes))
+	require.Equal(t, 0, int(m.RecvBytes))
+	require.Equal(t, int(len(msg)), int(m.SentPackets))
+	require.Equal(t, 0, int(m.RecvPackets))
+	require.True(t, outgoing.IntraHost)
+
+	// make sure the inverse values are seen for the other message
+	m = incoming.Monotonic
+	require.Equal(t, 0, int(m.SentBytes))
+	require.Equal(t, expectedBytes, int(m.RecvBytes))
+	require.Equal(t, 0, int(m.SentPackets))
+	require.Equal(t, int(len(msg)), int(m.RecvPackets))
 	require.True(t, incoming.IntraHost)
 }
 
@@ -1506,10 +1606,18 @@ func (s *TracerSuite) TestSendfileRegression() {
 		}, 3*time.Second, 100*time.Millisecond, "couldn't find connections used by sendfile(2)")
 
 		if assert.NotNil(t, outConn, "couldn't find outgoing connection used by sendfile(2)") {
-			assert.Equalf(t, int64(clientMessageSize), int64(outConn.Monotonic.SentBytes), "sendfile send data wasn't properly traced")
+			assert.Equalf(t, int64(clientMessageSize), int64(outConn.Monotonic.SentBytes), "sendfile sent bytes wasn't properly traced")
+			if connType == network.UDP {
+				assert.Equalf(t, int64(1), int64(outConn.Monotonic.SentPackets), "sendfile UDP should send exactly 1 packet")
+				assert.Equalf(t, int64(0), int64(outConn.Monotonic.RecvPackets), "sendfile outConn shouldn't have any RecvPackets")
+			}
 		}
 		if assert.NotNil(t, inConn, "couldn't find incoming connection used by sendfile(2)") {
-			assert.Equalf(t, int64(clientMessageSize), int64(inConn.Monotonic.RecvBytes), "sendfile recv data wasn't properly traced")
+			assert.Equalf(t, int64(clientMessageSize), int64(inConn.Monotonic.RecvBytes), "sendfile recv bytes wasn't properly traced")
+			if connType == network.UDP {
+				assert.Equalf(t, int64(1), int64(inConn.Monotonic.RecvPackets), "sendfile UDP should recv exactly 1 packet")
+				assert.Equalf(t, int64(0), int64(inConn.Monotonic.SentPackets), "sendfile inConn shouldn't have any SentPackets")
+			}
 		}
 	}
 
@@ -1540,7 +1648,7 @@ func (s *TracerSuite) TestSendfileRegression() {
 					t.Skip("UDP will fail with prebuilt tracer")
 				}
 
-				// Start TCP server
+				// Start UDP server
 				var rcvd int64
 				server := &UDPServer{
 					network: "udp" + strings.TrimPrefix(family.String(), "v"),

releasenotes/notes/npm-udp-packet-counting-ebce8760bec29f50.yaml

@@ -0,0 +1,11 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+features:
+  - |
+    NPM - adds UDP "Packets Sent" and "Packets Received" to the network telemetry in Linux.