Merge pull request #87 from snprajwal/filter-pid

feat: filter by PID in process tree

checkpointctl.go

@@ -17,6 +17,7 @@ var (
 	format  string
 	stats   bool
 	mounts  bool
+	pID     uint32
 	psTree  bool
 	files   bool
 	showAll bool
@@ -89,6 +90,13 @@ func setupInspect() *cobra.Command {
 		false,
 		"Display an overview of mounts used in the container checkpoint",
 	)
+	flags.Uint32VarP(
+		&pID,
+		"pid",
+		"p",
+		0,
+		"Display the process tree of a specific PID",
+	)
 	flags.BoolVar(
 		&psTree,
 		"ps-tree",
@@ -131,6 +139,11 @@ func inspect(cmd *cobra.Command, args []string) error {
 		requiredFiles = append(requiredFiles, "stats-dump")
 	}
 
+	if pID != 0 {
+		// Enable displaying process tree if the PID filter is passed.
+		psTree = true
+	}
+
 	if files {
 		// Enable displaying process tree, even if it is not passed.
 		// This is necessary to attach the files to the processes

test/checkpointctl.bats

@@ -315,6 +315,31 @@ function teardown() {
 	[[ ${lines[0]} == *"failed to get file descriptors"* ]]
 }
 
+@test "Run checkpointctl inspect with tar file and --ps-tree and valid PID" {
+	cp data/config.dump \
+		data/spec.dump "$TEST_TMP_DIR1"
+	mkdir "$TEST_TMP_DIR1"/checkpoint
+	cp test-imgs/pstree.img \
+		test-imgs/core-*.img "$TEST_TMP_DIR1"/checkpoint
+	( cd "$TEST_TMP_DIR1" && tar cf "$TEST_TMP_DIR2"/test.tar . )
+	checkpointctl inspect "$TEST_TMP_DIR2"/test.tar --ps-tree --pid 1
+	[ "$status" -eq 0 ]
+	[[ ${lines[8]} == *"Process tree"* ]]
+	[[ ${lines[9]} == *"piggie"* ]]
+}
+
+@test "Run checkpointctl inspect with tar file and --ps-tree and invalid PID" {
+	cp data/config.dump \
+		data/spec.dump "$TEST_TMP_DIR1"
+	mkdir "$TEST_TMP_DIR1"/checkpoint
+	cp test-imgs/pstree.img \
+		test-imgs/core-*.img "$TEST_TMP_DIR1"/checkpoint
+	( cd "$TEST_TMP_DIR1" && tar cf "$TEST_TMP_DIR2"/test.tar . )
+	checkpointctl inspect "$TEST_TMP_DIR2"/test.tar --ps-tree --pid 99999
+	[ "$status" -eq 1 ]
+	[[ ${lines[0]} == *"no process with PID 99999"* ]]
+}
+
 @test "Run checkpointctl inspect with tar file and --all and valid spec.dump and valid stats-dump" {
 	cp data/config.dump "$TEST_TMP_DIR1"
 	cp data/spec.dump "$TEST_TMP_DIR1"

tree.go

@@ -37,7 +37,9 @@ func renderTreeView(tasks []task) error {
 				return fmt.Errorf("failed to get process tree: %w", err)
 			}
 
-			addPsTreeToTree(tree, psTree)
+			if err = addPsTreeToTree(tree, psTree); err != nil {
+				return fmt.Errorf("failed to get process tree: %w", err)
+			}
 		}
 
 		if files {
@@ -110,7 +112,30 @@ func addDumpStatsToTree(tree treeprint.Tree, dumpStats *stats_pb.DumpStatsEntry)
 	statsTree.AddBranch(fmt.Sprintf("Pages written: %d us", dumpStats.GetPagesWritten()))
 }
 
-func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree) {
+func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree) error {
+	psRoot := psTree
+	if pID != 0 {
+		// dfs performs a short-circuiting depth-first search.
+		var dfs func(*crit.PsTree) *crit.PsTree
+		dfs = func(root *crit.PsTree) *crit.PsTree {
+			if root.PID == pID {
+				return root
+			}
+			for _, child := range root.Children {
+				if ps := dfs(child); ps != nil {
+					return ps
+				}
+			}
+			return nil
+		}
+
+		ps := dfs(psTree)
+		if ps == nil {
+			return fmt.Errorf("no process with PID %d (use `inspect --ps-tree` to view all PIDs)", pID)
+		}
+		psRoot = ps
+	}
+
 	// processNodes is a recursive function to create
 	// a new branch for each process and add its child
 	// processes as child nodes of the branch.
@@ -122,13 +147,21 @@ func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree) {
 		}
 	}
 	psTreeNode := tree.AddBranch("Process tree")
-	processNodes(psTreeNode, psTree)
+	processNodes(psTreeNode, psRoot)
+
+	return nil
 }
 
 func addFdsToTree(tree treeprint.Tree, fds []*crit.Fd) {
 	var node treeprint.Tree
 	for _, fd := range fds {
 		node = tree.FindByMeta(fd.PId)
+		// If FindByMeta returns nil, then the node with
+		// the PID has been pruned while building the tree.
+		// Hence, skip all associated file descriptors.
+		if node == nil {
+			continue
+		}
 		for _, file := range fd.Files {
 			node.AddMetaBranch(strings.TrimSpace(file.Type+" "+file.Fd), file.Path)
 		}