pstree: collect pid-s in rb-tree

criu/files-reg.c

@@ -357,18 +357,17 @@ static int open_remap_dead_process(struct reg_file_info *rfi,
 {
 	struct pstree_item *helper;
 
-	for_each_pstree_item(helper) {
-		/* don't need to add multiple tasks */
-		if (helper->pid.virt == rfe->remap_id) {
-			pr_info("Skipping helper for restoring /proc/%d; pid exists\n", rfe->remap_id);
-			return 0;
-		}
-	}
-
-	helper = alloc_pstree_helper();
+	helper = insert_item(rfe->remap_id);
 	if (!helper)
 		return -1;
 
+	if (helper->pid.state != TASK_UNDEF) {
+		pr_info("Skipping helper for restoring /proc/%d; pid exists\n", rfe->remap_id);
+		return 0;
+	}
+
+	init_pstree_helper(helper);
+
 	helper->sid = root_item->sid;
 	helper->pgid = root_item->pgid;
 	helper->pid.virt = rfe->remap_id;

criu/include/image.h

@@ -112,10 +112,12 @@
 
 #define TASK_COMM_LEN 16
 
+#define TASK_UNDEF		0x0
 #define TASK_ALIVE		0x1
 #define TASK_DEAD		0x2
 #define TASK_STOPPED		0x3
 #define TASK_HELPER		0x4
+#define TASK_THREAD		0x5
 
 #define CR_PARENT_LINK "parent"
 

criu/pstree.c

@@ -331,8 +331,6 @@ int dump_pstree(struct pstree_item *root_item)
 	return ret;
 }
 
-static int max_pid = 0;
-
 static int prepare_pstree_for_shell_job(void)
 {
 	pid_t current_sid = getsid(getpid());
@@ -385,6 +383,23 @@ static int prepare_pstree_for_shell_job(void)
 	return 0;
 }
 
+static struct pid *lookup_pid(pid_t pid)
+{
+	struct rb_node *node = 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;
+}
+
 static struct pid *insert_pid(pid_t pid, struct pid *pid_node)
 {
 	struct rb_node *node = root_rb.rb_node;
@@ -448,18 +463,19 @@ static int read_pstree_image(void)
 			break;
 
 		ret = -1;
-		pi = alloc_pstree_item_with_rst();
+		pi = insert_item(e->pid);
 		if (pi == NULL)
 			break;
 
-		pi->pid.virt = e->pid;
-		max_pid = max((int)e->pid, max_pid);
+		if (insert_item(e->pgid) == NULL)
+			break;
+		if (insert_item(e->sid) == NULL)
+			break;
 
+		pi->pid.virt = e->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) {
 			if (root_item) {
@@ -509,7 +525,8 @@ static int read_pstree_image(void)
 		for (i = 0; i < e->n_threads; i++) {
 			pi->threads[i].real = -1;
 			pi->threads[i].virt = e->threads[i];
-			max_pid = max((int)e->threads[i], max_pid);
+			pi->threads[i].state = TASK_THREAD;
+			insert_pid(pi->threads[i].virt, &pi->threads[i]);
 		}
 
 		task_entries->nr_threads += e->n_threads;
@@ -537,11 +554,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(&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;
@@ -556,6 +596,7 @@ static int prepare_pstree_ids(void)
 	 * reparented to init.
 	 */
 	list_for_each_entry(item, &root_item->children, sibling) {
+		struct pstree_item *leader;
 
 		/*
 		 * If a child belongs to the root task's session or it's
@@ -565,15 +606,36 @@ static int prepare_pstree_ids(void)
 		if (item->sid == root_item->sid || item->sid == item->pid.virt)
 			continue;
 
-		helper = alloc_pstree_helper();
-		if (helper == NULL)
-			return -1;
-		helper->sid = item->sid;
-		helper->pgid = item->sid;
-		helper->pid.virt = item->sid;
-		helper->parent = root_item;
-		helper->ids = root_item->ids;
-		list_add_tail(&helper->sibling, &helpers);
+		leader = insert_item(item->sid);
+		if (leader->pid.state != TASK_UNDEF) {
+			pid_t pid;
+
+			pid = get_free_pid();
+			if (pid < 0)
+				break;
+			helper = insert_item(pid);
+			if (helper == NULL)
+				return -1;
+
+			pr_info("Session leader %d\n", item->sid);
+
+			helper->sid = item->sid;
+			helper->pgid = leader->pgid;
+			helper->ids = leader->ids;
+			helper->parent = leader;
+			list_add(&helper->sibling, &leader->children);
+
+			pr_info("Attach %d to the task %d\n",
+					helper->pid.virt, leader->pid.virt);
+		} else {
+			helper = leader;
+			helper->sid = item->sid;
+			helper->pgid = item->sid;
+			helper->parent = root_item;
+			helper->ids = root_item->ids;
+			list_add_tail(&helper->sibling, &helpers);
+		}
+		init_pstree_helper(helper);
 
 		pr_info("Add a helper %d for restoring SID %d\n",
 				helper->pid.virt, helper->sid);
@@ -585,6 +647,8 @@ static int prepare_pstree_ids(void)
 		 * Stack on helper task all children with target sid.
 		 */
 		list_for_each_entry_safe_continue(child, tmp, &root_item->children, sibling) {
+			if (child == helper)
+				continue;
 			if (child->sid != helper->sid)
 				continue;
 			if (child->sid == child->pid.virt)
@@ -633,46 +697,22 @@ static int prepare_pstree_ids(void)
 
 			continue;
 		}
-
-		pr_info("Session leader %d\n", item->sid);
-
-		/* Try to find helpers, who should be connected to the leader */
-		list_for_each_entry(child, &helpers, sibling) {
-			if (child->pid.state != TASK_HELPER)
-				continue;
-
-			if (child->sid != item->sid)
-				continue;
-
-			child->pgid = item->pgid;
-			child->pid.virt = ++max_pid;
-			child->parent = item;
-			list_move(&child->sibling, &item->children);
-
-			pr_info("Attach %d to the task %d\n",
-					child->pid.virt, item->pid.virt);
-
-			break;
-		}
 	}
 
 	/* All other helpers are session leaders for own sessions */
 	list_splice(&helpers, &root_item->children);
 
 	/* 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 = lookup_pid(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;
 		}
 
@@ -684,9 +724,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;