[CRIU] [PATCH 3/6] pstree: prepare to store all pid-s in rb-tree

Andrey Vagin avagin at openvz.org
Thu Mar 10 11:24:40 PST 2016


From: Andrew Vagin <avagin at virtuozzo.com>

It will allow us to find a free pid and will speed up
searching an item by pid.

The current algorithm to searching a free pid may return
a value bigger that kernel.max_pid.

v2: add a coment before insert_pid()
    rename root_rb to pid_root_rb and make it "static"

Signed-off-by: Andrew Vagin <avagin at virtuozzo.com>
---
 criu/files-reg.c      | 17 ++++++-----
 criu/include/image.h  |  2 ++
 criu/include/pid.h    |  3 ++
 criu/include/pstree.h |  2 ++
 criu/pstree.c         | 79 ++++++++++++++++++++++++++++++++++++++++++++++++---
 5 files changed, 90 insertions(+), 13 deletions(-)

diff --git a/criu/files-reg.c b/criu/files-reg.c
index a4687c3..53469ed 100644
--- a/criu/files-reg.c
+++ b/criu/files-reg.c
@@ -360,18 +360,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;
diff --git a/criu/include/image.h b/criu/include/image.h
index 305febf..f141915 100644
--- a/criu/include/image.h
+++ b/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"
 
diff --git a/criu/include/pid.h b/criu/include/pid.h
index 2a709a2..533be88 100644
--- a/criu/include/pid.h
+++ b/criu/include/pid.h
@@ -2,6 +2,7 @@
 #define __CR_PID_H__
 
 #include "stdbool.h"
+#include "rbtree.h"
 
 struct pid {
 	/*
@@ -19,6 +20,8 @@ struct pid {
 	pid_t virt;
 
 	int state;	/* TASK_XXX constants */
+
+	struct rb_node node;
 };
 
 /*
diff --git a/criu/include/pstree.h b/criu/include/pstree.h
index 226d99a..4a1d50b 100644
--- a/criu/include/pstree.h
+++ b/criu/include/pstree.h
@@ -68,6 +68,8 @@ extern struct pstree_item *__alloc_pstree_item(bool rst);
 extern struct pstree_item *alloc_pstree_helper(void);
 extern void init_pstree_helper(struct pstree_item *ret);
 
+extern struct pstree_item *insert_item(pid_t pid);
+
 extern struct pstree_item *root_item;
 extern struct pstree_item *pstree_item_next(struct pstree_item *item);
 #define for_each_pstree_item(pi) \
diff --git a/criu/pstree.c b/criu/pstree.c
index bdf8ff6..bbd797e 100644
--- a/criu/pstree.c
+++ b/criu/pstree.c
@@ -20,6 +20,7 @@
 #include "images/pstree.pb-c.h"
 
 struct pstree_item *root_item;
+static struct rb_root pid_root_rb;
 
 #define CLONE_ALLNS     (CLONE_NEWPID | CLONE_NEWNET | CLONE_NEWIPC | CLONE_NEWUTS | CLONE_NEWNS | CLONE_NEWUSER)
 
@@ -387,6 +388,55 @@ static int prepare_pstree_for_shell_job(void)
 	return 0;
 }
 
+/*
+ * 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
+ * is inserted.
+ */
+static struct pid *insert_pid(pid_t pid, struct pid *pid_node)
+{
+	struct rb_node *node = pid_root_rb.rb_node;
+	struct rb_node **new = &pid_root_rb.rb_node;
+	struct rb_node *parent = NULL;
+
+	while (node) {
+		struct pid *this = rb_entry(node, struct pid, node);
+
+		parent = *new;
+		if (pid < this->virt)
+			node = node->rb_left, new = &((*new)->rb_left);
+		else if (pid > this->virt)
+			node = node->rb_right, new = &((*new)->rb_right);
+		else
+			return this;
+	}
+
+	if (!pid_node) {
+		struct pstree_item *item;
+
+		item = alloc_pstree_item_with_rst();
+		if (item == NULL)
+			return NULL;
+
+		item->pid.virt = pid;
+		pid_node = &item->pid;
+	}
+	rb_link_and_balance(&pid_root_rb, &pid_node->node, parent, new);
+	return pid_node;
+}
+
+struct pstree_item *insert_item(pid_t pid)
+{
+	struct pid *node;;
+
+	node = insert_pid(pid, NULL);
+	if (!node)
+		return NULL;
+	BUG_ON(node->state == TASK_THREAD);
+
+	return container_of(node, struct pstree_item, pid);
+}
+
 static int read_pstree_image(void)
 {
 	int ret = 0, i;
@@ -407,18 +457,29 @@ static int read_pstree_image(void)
 			break;
 
 		ret = -1;
-		pi = alloc_pstree_item_with_rst();
+		pi = insert_item(e->pid);
 		if (pi == NULL)
 			break;
+		BUG_ON(pi->pid.state != TASK_UNDEF);
+
+		/*
+		 * All pids should be added in the tree to be able to find
+		 * free pid-s for helpers. pstree_item for these pid-s will
+		 * be initialized when we meet PstreeEntry with this pid or
+		 * we will create helpers for them.
+		 */
+		if (insert_item(e->pgid) == NULL)
+			break;
+		if (insert_item(e->sid) == NULL)
+			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) {
 			if (root_item) {
@@ -466,9 +527,19 @@ static int read_pstree_image(void)
 			break;
 
 		for (i = 0; i < e->n_threads; i++) {
+			struct pid *node;
 			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;
+			if (i == 0)
+				continue; /* A thread leader is in a tree already */
+			node = insert_pid(pi->threads[i].virt, &pi->threads[i]);
+
+			BUG_ON(node == NULL);
+			if (node != &pi->threads[i]) {
+				pr_err("Unexpected task %d in a tree %d\n", e->threads[i], i);
+				return -1;
+			}
 		}
 
 		task_entries->nr_threads += e->n_threads;
-- 
2.5.0



More information about the CRIU mailing list