diff --git a/checkpointctl.go b/checkpointctl.go index f58409aa..9723e0cb 100644 --- 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 index 72d4e2f1..5d8fd675 100644 --- 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 index 581bc12a..20de8955 100644 --- 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) {