Merge pull request #268 from DataDog/conor/connections-max

Add ConnectionsMax function that limits connections per pid.

net/net_fallback.go

@@ -19,3 +19,7 @@ func ProtoCounters(protocols []string) ([]ProtoCountersStat, error) {
 func Connections(kind string) ([]ConnectionStat, error) {
 	return []ConnectionStat{}, common.ErrNotImplementedError
 }
+
+func ConnectionsMax(kind string, max int) ([]ConnectionStat, error) {
+	return []ConnectionStat{}, common.ErrNotImplementedError
+}

net/net_linux.go

@@ -292,6 +292,12 @@ func Connections(kind string) ([]ConnectionStat, error) {
 	return ConnectionsPid(kind, 0)
 }
 
+// Return a list of network connections opened returning at most `max`
+// connections for each running process.
+func ConnectionsMax(kind string, max int) ([]ConnectionStat, error) {
+	return ConnectionsPidMax(kind, 0, max)
+}
+
 // Return a list of network connections opened by a process.
 func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 	tmap, ok := netConnectionKindMap[kind]
@@ -302,9 +308,9 @@ func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 	var err error
 	var inodes map[string][]inodeMap
 	if pid == 0 {
-		inodes, err = getProcInodesAll(root)
+		inodes, err = getProcInodesAll(root, 0)
 	} else {
-		inodes, err = getProcInodes(root, pid)
+		inodes, err = getProcInodes(root, pid, 0)
 		if len(inodes) == 0 {
 			// no connection for the pid
 			return []ConnectionStat{}, nil
@@ -313,10 +319,38 @@ func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 	if err != nil {
 		return nil, fmt.Errorf("cound not get pid(s), %d", pid)
 	}
+	return statsFromInodes(root, pid, tmap, inodes)
+}
 
+// Return up to `max` network connections opened by a process.
+func ConnectionsPidMax(kind string, pid int32, max int) ([]ConnectionStat, error) {
+	tmap, ok := netConnectionKindMap[kind]
+	if !ok {
+		return nil, fmt.Errorf("invalid kind, %s", kind)
+	}
+	root := common.HostProc()
+	var err error
+	var inodes map[string][]inodeMap
+	if pid == 0 {
+		inodes, err = getProcInodesAll(root, max)
+	} else {
+		inodes, err = getProcInodes(root, pid, max)
+		if len(inodes) == 0 {
+			// no connection for the pid
+			return []ConnectionStat{}, nil
+		}
+	}
+	if err != nil {
+		return nil, fmt.Errorf("cound not get pid(s), %d", pid)
+	}
+	return statsFromInodes(root, pid, tmap, inodes)
+}
+
+func statsFromInodes(root string, pid int32, tmap []netConnectionKindType, inodes map[string][]inodeMap) ([]ConnectionStat, error) {
 	dupCheckMap := make(map[string]bool)
 	var ret []ConnectionStat
 
+	var err error
 	for _, t := range tmap {
 		var path string
 		var ls []connTmp
@@ -368,11 +402,15 @@ func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 }
 
 // getProcInodes returnes fd of the pid.
-func getProcInodes(root string, pid int32) (map[string][]inodeMap, error) {
+func getProcInodes(root string, pid int32, max int) (map[string][]inodeMap, error) {
 	ret := make(map[string][]inodeMap)
 
 	dir := fmt.Sprintf("%s/%d/fd", root, pid)
-	files, err := ioutil.ReadDir(dir)
+	f, err := os.Open(dir)
+	if err != nil {
+		return ret, nil
+	}
+	files, err := f.Readdir(max)
 	if err != nil {
 		return ret, nil
 	}
@@ -484,7 +522,7 @@ func (p *process) fillFromStatus() error {
 	return nil
 }
 
-func getProcInodesAll(root string) (map[string][]inodeMap, error) {
+func getProcInodesAll(root string, max int) (map[string][]inodeMap, error) {
 	pids, err := Pids()
 	if err != nil {
 		return nil, err
@@ -492,7 +530,7 @@ func getProcInodesAll(root string) (map[string][]inodeMap, error) {
 	ret := make(map[string][]inodeMap)
 
 	for _, pid := range pids {
-		t, err := getProcInodes(root, pid)
+		t, err := getProcInodes(root, pid, max)
 		if err != nil {
 			return ret, err
 		}

net/net_linux_test.go

@@ -15,11 +15,32 @@ func TestGetProcInodesAll(t *testing.T) {
 	}
 
 	root := common.HostProc("")
-	v, err := getProcInodesAll(root)
+	v, err := getProcInodesAll(root, 0)
 	assert.Nil(t, err)
 	assert.NotEmpty(t, v)
 }
 
+func TestConnectionsMax(t *testing.T) {
+	if os.Getenv("CIRCLECI") == "true" {
+		t.Skip("Skip CI")
+	}
+
+	max := 10
+	v, err := ConnectionsMax("tcp", max)
+	assert.Nil(t, err)
+	assert.NotEmpty(t, v)
+
+	cxByPid := map[int32]int{}
+	for _, cx := range v {
+		if cx.Pid > 0 {
+			cxByPid[cx.Pid]++
+		}
+	}
+	for _, c := range cxByPid {
+		assert.True(t, c <= max)
+	}
+}
+
 type AddrTest struct {
 	IP    string
 	Port  int

net/net_unix.go

@@ -13,6 +13,12 @@ func Connections(kind string) ([]ConnectionStat, error) {
 	return ConnectionsPid(kind, 0)
 }
 
+// Return a list of network connections opened returning at most `max`
+// connections for each running process.
+func ConnectionsMax(kind string, max int) ([]ConnectionStat, error) {
+	return []ConnectionStat{}, common.ErrNotImplementedError
+}
+
 // Return a list of network connections opened by a process.
 func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 	var ret []ConnectionStat
@@ -66,3 +72,8 @@ func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
 
 	return ret, nil
 }
+
+// Return up to `max` network connections opened by a process.
+func ConnectionsPidMax(kind string, pid int32, max int) ([]ConnectionStat, error) {
+	return []ConnectionStat{}, common.ErrNotImplementedError
+}

net/net_windows.go

@@ -76,6 +76,12 @@ func Connections(kind string) ([]ConnectionStat, error) {
 	return ret, common.ErrNotImplementedError
 }
 
+// Return a list of network connections opened returning at most `max`
+// connections for each running process.
+func ConnectionsMax(kind string, max int) ([]ConnectionStat, error) {
+	return []ConnectionStat{}, common.ErrNotImplementedError
+}
+
 func FilterCounters() ([]FilterStat, error) {
 	return nil, errors.New("NetFilterCounters not implemented for windows")
 }