feat: Make net_(tcp|udp) read limit configurable

Defaults to 4GiB otherwise, as earlier. Fixes: prometheus#364 Signed-off-by: Pranshu Srivastava <[email protected]>
rexagod · Mar 22, 2024 · 2b570ae · 2b570ae
1 parent f7c2619
commit 2b570ae
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 11 deletions.
diff --git a/go.mod b/go.mod
@@ -6,4 +6,5 @@ require (
 	github.com/google/go-cmp v0.6.0
 	golang.org/x/sync v0.6.0
 	golang.org/x/sys v0.18.0
+	k8s.io/utils v0.0.0-20240310230437-4693a0247e57
 )
diff --git a/go.sum b/go.sum
@@ -4,3 +4,5 @@ golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ=
 golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
 golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+k8s.io/utils v0.0.0-20240310230437-4693a0247e57 h1:gbqbevonBh57eILzModw6mrkbwM0gQBEuevE/AaBsHY=
+k8s.io/utils v0.0.0-20240310230437-4693a0247e57/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
diff --git a/net_ip_socket.go b/net_ip_socket.go
@@ -22,16 +22,18 @@ import (
 	"os"
 	"strconv"
 	"strings"
+
+	"k8s.io/utils/ptr"
 )
 
-const (
-	// readLimit is used by io.LimitReader while reading the content of the
+var (
+	// defaultReadLimitPtr is used by io.LimitReader while reading the content of the
 	// /proc/net/udp{,6} files. The number of lines inside such a file is dynamic
 	// as each line represents a single used socket.
 	// In theory, the number of available sockets is 65535 (2^16 - 1) per IP.
 	// With e.g. 150 Byte per line and the maximum number of 65535,
 	// the reader needs to handle 150 Byte * 65535 =~ 10 MB for a single IP.
-	readLimit = 4294967296 // Byte -> 4 GiB
+	defaultReadLimit = ptr.To(4294967296) // Byte -> 4 GiB
 )
 
 // This contains generic data structures for both udp and tcp sockets.
@@ -73,7 +75,7 @@ type (
 	}
 )
 
-func newNetIPSocket(file string) (NetIPSocket, error) {
+func newNetIPSocket(file string, readLimit *int) (NetIPSocket, error) {
 	f, err := os.Open(file)
 	if err != nil {
 		return nil, err
@@ -83,7 +85,10 @@ func newNetIPSocket(file string) (NetIPSocket, error) {
 	var netIPSocket NetIPSocket
 	isUDP := strings.Contains(file, "udp")
 
-	lr := io.LimitReader(f, readLimit)
+	if readLimit == nil {
+		readLimit = defaultReadLimit
+	}
+	lr := io.LimitReader(f, int64(*readLimit))
 	s := bufio.NewScanner(lr)
 	s.Scan() // skip first line with headers
 	for s.Scan() {
@@ -101,7 +106,7 @@ func newNetIPSocket(file string) (NetIPSocket, error) {
 }
 
 // newNetIPSocketSummary creates a new NetIPSocket{,6} from the contents of the given file.
-func newNetIPSocketSummary(file string) (*NetIPSocketSummary, error) {
+func newNetIPSocketSummary(file string, readLimit *int) (*NetIPSocketSummary, error) {
 	f, err := os.Open(file)
 	if err != nil {
 		return nil, err
@@ -112,7 +117,10 @@ func newNetIPSocketSummary(file string) (*NetIPSocketSummary, error) {
 	var udpPacketDrops uint64
 	isUDP := strings.Contains(file, "udp")
 
-	lr := io.LimitReader(f, readLimit)
+	if readLimit == nil {
+		readLimit = defaultReadLimit
+	}
+	lr := io.LimitReader(f, int64(*readLimit))
 	s := bufio.NewScanner(lr)
 	s.Scan() // skip first line with headers
 	for s.Scan() {

diff --git a/net_tcp.go b/net_tcp.go
@@ -49,13 +49,13 @@ func (fs FS) NetTCP6Summary() (*NetTCPSummary, error) {
 
 // newNetTCP creates a new NetTCP{,6} from the contents of the given file.
 func newNetTCP(file string) (NetTCP, error) {
-	n, err := newNetIPSocket(file)
+	n, err := newNetIPSocket(file, nil)
 	n1 := NetTCP(n)
 	return n1, err
 }
 
 func newNetTCPSummary(file string) (*NetTCPSummary, error) {
-	n, err := newNetIPSocketSummary(file)
+	n, err := newNetIPSocketSummary(file, nil)
 	if n == nil {
 		return nil, err
 	}

diff --git a/net_udp.go b/net_udp.go
@@ -49,13 +49,13 @@ func (fs FS) NetUDP6Summary() (*NetUDPSummary, error) {
 
 // newNetUDP creates a new NetUDP{,6} from the contents of the given file.
 func newNetUDP(file string) (NetUDP, error) {
-	n, err := newNetIPSocket(file)
+	n, err := newNetIPSocket(file, nil)
 	n1 := NetUDP(n)
 	return n1, err
 }
 
 func newNetUDPSummary(file string) (*NetUDPSummary, error) {
-	n, err := newNetIPSocketSummary(file)
+	n, err := newNetIPSocketSummary(file, nil)
 	if n == nil {
 		return nil, err
 	}