[CRIU] [PATCH 6/6] pstree: use rbtree to find a specified pid (v2)

[Andrey Vagin]

From: Andrew Vagin <avagin at virtuozzo.com>

v2: rename lookup_pid to pstree_pid_by_virt
    use pstree_pid_by_virt in pstree_item_by_virt

Signed-off-by: Andrew Vagin <avagin at virtuozzo.com>
---
 criu/pstree.c | 49 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 17 deletions(-)

diff --git a/criu/pstree.c b/criu/pstree.c
index 00dc174..5729266 100644
--- a/criu/pstree.c
+++ b/criu/pstree.c
@@ -388,6 +388,23 @@ static int prepare_pstree_for_shell_job(void)
 	return 0;
 }
 
+static struct pid *pstree_pid_by_virt(pid_t pid)
+{
+	struct rb_node *node = pid_root_rb.rb_node;
+
+	while (node) {
+		struct pid *this = rb_entry(node, struct pid, node);
+
+		if (pid < this->virt)
+			node = node->rb_left;
+		else if (pid > this->virt)
+			node = node->rb_right;
+		else
+			return this;
+	}
+	return NULL;
+}
+
 /*
  * Try to find a pid node in the tree and insert a new one,
  * it is not there yet. If pid_node isn't set, pstree_item
@@ -712,18 +729,15 @@ static int prepare_pstree_ids(void)
 
 	/* Add a process group leader if it is absent  */
 	for_each_pstree_item(item) {
-		struct pstree_item *gleader;
+		struct pid *pid;
 
 		if (!item->pgid || item->pid.virt == item->pgid)
 			continue;
 
-		for_each_pstree_item(gleader) {
-			if (gleader->pid.virt == item->pgid)
-				break;
-		}
-
-		if (gleader) {
-			rsti(item)->pgrp_leader = gleader;
+		pid = pstree_pid_by_virt(item->pgid);
+		if (pid->state != TASK_UNDEF) {
+			BUG_ON(pid->state == TASK_THREAD);
+			rsti(item)->pgrp_leader = container_of(pid, struct pstree_item, pid);
 			continue;
 		}
 
@@ -735,9 +749,9 @@ static int prepare_pstree_ids(void)
 		if (current_pgid == item->pgid)
 			continue;
 
-		helper = alloc_pstree_helper();
-		if (helper == NULL)
-			return -1;
+		helper = container_of(pid, struct pstree_item, pid);
+		init_pstree_helper(helper);
+
 		helper->sid = item->sid;
 		helper->pgid = item->pgid;
 		helper->pid.virt = item->pgid;
@@ -982,13 +996,14 @@ bool restore_before_setsid(struct pstree_item *child)
 
 struct pstree_item *pstree_item_by_virt(pid_t virt)
 {
-	struct pstree_item *item;
+	struct pid *pid;
 
-	for_each_pstree_item(item) {
-		if (item->pid.virt == virt)
-			return item;
-	}
-	return NULL;
+	pid = pstree_pid_by_virt(virt);
+	if (pid == NULL)
+		return NULL;
+	BUG_ON(pid->state == TASK_THREAD);
+
+	return container_of(pid, struct pstree_item, pid);
 }
 
 struct pstree_item *pstree_item_by_real(pid_t real)
-- 
2.5.0