restore: use less convoluted method to calculate the number of zombies

Instead of doing this crazy thing where we pass the number of zombies back
from the restorer blob to the main restore thread, let's just calculate the
number in the main restore thread when we need it.

Also, note that this patch shifts the location of this calculation: we used
to do a decrement of nr_in_progress when we waited() on a zombie, to sort
of lie about the fact that they did go through CR_STATE_RESTORE_SIGCHLD,
but since we're waiting on them anyway explicitly, there's no reason to do
all this hackery. Let's just calculate the number once in the right place
and be done with it.
Signed-off-by: Tycho Andersen <tycho.andersen@canonical.com>
Signed-off-by: Pavel Emelyanov <xemul@virtuozzo.com>

criu/cr-restore.c

@@ -1590,6 +1590,7 @@ static int restore_root_task(struct pstree_item *init)
 	enum trace_flags flag = TRACE_ALL;
 	int ret, fd, mnt_ns_fd = -1;
 	int clean_remaps = 1, root_seized = 0;
+	struct pstree_item *item;
 
 	ret = run_scripts(ACT_PRE_RESTORE);
 	if (ret != 0) {
@@ -1709,16 +1710,16 @@ static int restore_root_task(struct pstree_item *init)
 	if (ret < 0)
 		goto out_kill;
 
+	/* Zombies die after CR_STATE_RESTORE */
+	for_each_pstree_item(item) {
+		if (item->pid.state == TASK_DEAD)
+			task_entries->nr_threads--;
+	}
+
 	ret = restore_switch_stage(CR_STATE_RESTORE_SIGCHLD);
 	if (ret < 0)
 		goto out_kill;
 
-	/*
-	 * The task_entries->nr_zombies is updated in the
-	 * CR_STATE_RESTORE_SIGCHLD in pie code.
-	 */
-	task_entries->nr_threads -= atomic_read(&task_entries->nr_zombies);
-
 	/*
 	 * There is no need to call try_clean_remaps() after this point,
 	 * as restore went OK and all ghosts were removed by the openers.
@@ -1850,7 +1851,6 @@ static int prepare_task_entries(void)
 	task_entries->nr_threads = 0;
 	task_entries->nr_tasks = 0;
 	task_entries->nr_helpers = 0;
-	atomic_set(&task_entries->nr_zombies, 0);
 	futex_set(&task_entries->start, CR_STATE_RESTORE_NS);
 	mutex_init(&task_entries->userns_sync_lock);
 

criu/include/rst_info.h

@@ -7,7 +7,6 @@
 
 struct task_entries {
 	int nr_threads, nr_tasks, nr_helpers;
-	atomic_t nr_zombies;
 	futex_t nr_in_progress;
 	futex_t start;
 	atomic_t cr_err;

criu/pie/restorer.c

@@ -978,8 +978,6 @@ static int wait_zombies(struct task_restore_args *task_args)
 {
 	int i;
 
-	atomic_add(task_args->zombies_n, &task_entries->nr_zombies);
-
 	for (i = 0; i < task_args->zombies_n; i++) {
 		if (sys_waitid(P_PID, task_args->zombies[i], NULL, WNOWAIT | WEXITED, NULL) < 0) {
 			pr_err("Wait on %d zombie failed\n", task_args->zombies[i]);
@@ -987,7 +985,6 @@ static int wait_zombies(struct task_restore_args *task_args)
 		}
 		pr_debug("%ld: Collect a zombie with pid %d\n",
 			sys_getpid(), task_args->zombies[i]);
-		futex_dec_and_wake(&task_entries->nr_in_progress);
 	}
 
 	return 0;