feat: filter by PID in process tree

Currently, the entire process tree is rendered in the tree view. This allows for the tree to be pruned for a specific process using the `--pid` or `-p` flag. Signed-off-by: Prajwal S N <[email protected]>
checkpoint-restore · Jul 12, 2023 · 40d4073 · 40d4073
1 parent 29d4681
commit 40d4073
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 3 deletions.
diff --git a/checkpointctl.go b/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,
+		"Specify a particular process in the container checkpoint",
+	)
 	flags.BoolVar(
 		&psTree,
 		"ps-tree",

diff --git a/test/checkpointctl.bats b/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"

diff --git a/tree.go b/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", 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,7 +147,9 @@ 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) {