[CRIU] [PATCH 6/6] pstree: try to find a free pid between busy pids (v2)

[Andrey Vagin]

From: Andrew Vagin <avagin at virtuozzo.com>

Currently our pid allocator returns max_pid++ and it can return a pid
which is bigger than kernel.max_pid.

(00.821430)   5506: Error (cr-restore.c:1540): Pid 300 do not match expected 32768

https://jira.sw.ru/browse/PSBM-42202

v2: handle error code from insert_item()
Signed-off-by: Andrew Vagin <avagin at virtuozzo.com>
---
 criu/pstree.c | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 54 insertions(+), 9 deletions(-)

diff --git a/criu/pstree.c b/criu/pstree.c
index 3a278fa..096abec 100644
--- a/criu/pstree.c
+++ b/criu/pstree.c
@@ -325,8 +325,6 @@ err:
 	return ret;
 }
 
-static int max_pid = 0;
-
 static int prepare_pstree_for_shell_job(void)
 {
 	pid_t current_sid = getsid(getpid());
@@ -373,12 +371,31 @@ static int prepare_pstree_for_shell_job(void)
 			pi->sid = current_sid;
 	}
 
-	max_pid = max((int)current_sid, max_pid);
-	max_pid = max((int)current_gid, max_pid);
+	if (lookup_create_item(current_sid))
+		return -1;
+	if (lookup_create_item(current_gid))
+		return -1;
 
 	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
@@ -465,11 +482,8 @@ static int read_pstree_image(void)
 			break;
 
 		pi->pid.virt = e->pid;
-		max_pid = max((int)e->pid, max_pid);
 		pi->pgid = e->pgid;
-		max_pid = max((int)e->pgid, max_pid);
 		pi->sid = e->sid;
-		max_pid = max((int)e->sid, max_pid);
 		pi->pid.state = TASK_ALIVE;
 
 		if (e->ppid == 0) {
@@ -558,11 +572,34 @@ static int read_pstree_image(void)
 				goto err;
 		}
 	}
+
 err:
 	close_image(img);
 	return ret;
 }
 
+#define RESERVED_PIDS           300
+static int get_free_pid()
+{
+	static struct pstree_item *prev, *next;
+
+	if (prev == NULL)
+		prev = rb_entry(rb_first(&pid_root_rb), struct pstree_item, pid.node);
+
+	while (1) {
+		pid_t pid;
+		pid = prev->pid.virt + 1;
+		pid = pid < RESERVED_PIDS ? RESERVED_PIDS + 1 : pid;
+
+		next = rb_entry(rb_next(&prev->pid.node), struct pstree_item, pid.node);
+		if (&next->pid.node == NULL || next->pid.virt > pid)
+			return pid;
+		prev = next;
+	}
+
+	return -1;
+}
+
 static int prepare_pstree_ids(void)
 {
 	struct pstree_item *item, *child, *helper, *tmp;
@@ -578,6 +615,7 @@ static int prepare_pstree_ids(void)
 	 */
 	list_for_each_entry(item, &root_item->children, sibling) {
 		struct pstree_item *leader;
+		struct pid *node;
 
 		/*
 		 * If a child belongs to the root task's session or it's
@@ -587,9 +625,16 @@ static int prepare_pstree_ids(void)
 		if (item->sid == root_item->sid || item->sid == item->pid.virt)
 			continue;
 
-		leader = insert_item(item->sid);
+		node = pstree_pid_by_virt(item->sid);
+		BUG_ON(node == NULL);
+		leader = container_of(node, struct pstree_item, pid);
 		if (leader->pid.state != TASK_UNDEF) {
-			helper = insert_item(++max_pid);
+			pid_t pid;
+
+			pid = get_free_pid();
+			if (pid < 0)
+				break;
+			helper = lookup_create_item(pid);
 			if (helper == NULL)
 				return -1;
 
-- 
2.5.0