[CRIU] [PATCH v3 34/55] pstree: Extract __pstree_item_by_virt() to act on any pid_ns

[Kirill Tkhai]

Extrack a helper to find a task by its pid value in any pid_ns of hierarhy
by its value in this pid_ns.

Signed-off-by: Kirill Tkhai <ktkhai at virtuozzo.com>
---
 criu/include/pstree.h |   14 ++++++++++++--
 criu/pstree.c         |   19 ++++++++++++-------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/criu/include/pstree.h b/criu/include/pstree.h
index b0b440229..b8b8dd3f5 100644
--- a/criu/include/pstree.h
+++ b/criu/include/pstree.h
@@ -93,7 +93,12 @@ extern void init_pstree_helper(struct pstree_item *ret);
 
 extern struct pstree_item *lookup_create_item(pid_t *pid, int level, uint32_t ns_id);
 extern void pstree_insert_pid(struct pid *pid_node, uint32_t ns_id);
-extern struct pid *pstree_pid_by_virt(pid_t pid);
+extern struct pid *__pstree_pid_by_virt(struct ns_id *ns, pid_t pid);
+static inline struct pid *pstree_pid_by_virt(pid_t pid)
+{
+	extern struct ns_id *top_pid_ns;
+	return __pstree_pid_by_virt(top_pid_ns, pid);
+}
 
 extern struct pstree_item *root_item;
 extern struct pstree_item *pstree_item_next(struct pstree_item *item);
@@ -107,7 +112,12 @@ extern int prepare_dummy_pstree(void);
 extern int dump_pstree(struct pstree_item *root_item);
 
 struct pstree_item *pstree_item_by_real(pid_t virt);
-struct pstree_item *pstree_item_by_virt(pid_t virt);
+extern struct pstree_item *__pstree_item_by_virt(struct ns_id *ns, pid_t virt);
+static inline struct pstree_item *pstree_item_by_virt(pid_t virt)
+{
+	extern struct ns_id *top_pid_ns;
+	return __pstree_item_by_virt(top_pid_ns, virt);
+}
 
 extern int pid_to_virt(pid_t pid);
 
diff --git a/criu/pstree.c b/criu/pstree.c
index 6b7daba67..ecc46761e 100644
--- a/criu/pstree.c
+++ b/criu/pstree.c
@@ -563,16 +563,21 @@ struct pstree_item *lookup_create_item(pid_t *pid, int level, uint32_t ns_id)
 	return node->item;
 }
 
-struct pid *pstree_pid_by_virt(pid_t pid)
+struct pid *__pstree_pid_by_virt(struct ns_id *ns, pid_t pid)
 {
-	struct rb_node *node = top_pid_ns->pid.rb_root.rb_node;
+	struct rb_node *node = ns->pid.rb_root.rb_node;
+	struct ns_id *i = ns;
+	int level = 0;
+
+	while ((i = i->parent) != NULL)
+		level++;
 
 	while (node) {
-		struct pid *this = rb_entry(node, struct pid, ns[0].node);
+		struct pid *this = rb_entry(node, struct pid, ns[level].node);
 
-		if (pid < this->ns[0].virt)
+		if (pid < this->ns[level].virt)
 			node = node->rb_left;
-		else if (pid > this->ns[0].virt)
+		else if (pid > this->ns[level].virt)
 			node = node->rb_right;
 		else
 			return this;
@@ -1216,11 +1221,11 @@ bool restore_before_setsid(struct pstree_item *child)
 	return false;
 }
 
-struct pstree_item *pstree_item_by_virt(pid_t virt)
+struct pstree_item *__pstree_item_by_virt(struct ns_id *ns, pid_t virt)
 {
 	struct pid *pid;
 
-	pid = pstree_pid_by_virt(virt);
+	pid = __pstree_pid_by_virt(ns, virt);
 	if (pid == NULL)
 		return NULL;
 	BUG_ON(pid->state == TASK_THREAD);