restore: TASK_HELPERs live until RESTORE stage ends

In order to use TASK_HELPERS to open files from dead processes, they should
persist until criu is done restoring the filesystem, which happens in the
RESTORE stage. To do this, we need to pass each helper's PIDs to the restorer
blob, so that it can wait() on them when the restore stage is done.

This commit is in preparation for the remap_dead_pid commits.

v2: wait() on helpers after restore stage is over
v3: add CR_STATE_RESTORE_FS stage
v4: CR_STATE_RESTORE_FS waits for nr_tasks + nr_helpers, not nr_threads
v5: ditch CR_STATE_RESTORE_FS in favor of passing helpers to restorer blob
Signed-off-by: Tycho Andersen <tycho.andersen@canonical.com>
Signed-off-by: Pavel Emelyanov <xemul@parallels.com>

cr-restore.c

@@ -95,6 +95,9 @@ static int prepare_posix_timers(int pid, CoreEntry *core);
 static int prepare_signals(int pid, CoreEntry *core);
 
 static int root_as_sibling;
+static pid_t *helpers = NULL;
+static unsigned long helpers_pos = 0;
+static int n_helpers = 0;
 
 static int crtools_prepare_shared(void)
 {
@@ -698,30 +701,29 @@ err:
 	return ret;
 }
 
-static int pstree_wait_helpers()
+static int collect_helper_pids()
 {
 	struct pstree_item *pi;
 
 	list_for_each_entry(pi, &current->children, sibling) {
-		int status, ret;
 
 		if (pi->state != TASK_HELPER)
 			continue;
 
-		/* Check, that a helper completed. */
-		ret = waitpid(pi->pid.virt, &status, 0);
-		if (ret == -1) {
-			if (errno == ECHILD)
-				continue; /* It has been waited in sigchld_handler */
-			pr_err("waitpid(%d) failed\n", pi->pid.virt);
+		if (helpers) {
+			void *m;
+			m = rst_mem_alloc(sizeof(*helpers) * ++n_helpers, RM_PRIVATE);
+			if (!m)
 				return -1;
-		}
-		if (!WIFEXITED(status) || WEXITSTATUS(status)) {
-			pr_err("%d exited with non-zero code (%d,%d)\n", pi->pid.virt,
-				WEXITSTATUS(status), WTERMSIG(status));
+		} else {
+			helpers_pos = rst_mem_cpos(RM_PRIVATE);
+			helpers = rst_mem_alloc(sizeof(*helpers), RM_PRIVATE);
+			if (!helpers)
 				return -1;
+			n_helpers = 1;
 		}
 
+		helpers[n_helpers - 1] = pi->pid.virt;
 	}
 
 	return 0;
@@ -770,9 +772,6 @@ static int restore_one_alive_task(int pid, CoreEntry *core)
 
 	rst_mem_switch_to_private();
 
-	if (pstree_wait_helpers())
-		return -1;
-
 	if (prepare_fds(current))
 		return -1;
 
@@ -794,6 +793,9 @@ static int restore_one_alive_task(int pid, CoreEntry *core)
 	if (prepare_rlimits(pid, core) < 0)
 		return -1;
 
+	if (collect_helper_pids() < 0)
+		return -1;
+
 	return sigreturn_restore(pid, core);
 }
 
@@ -931,9 +933,10 @@ static int restore_one_task(int pid, CoreEntry *core)
 		ret = restore_one_alive_task(pid, core);
 	else if (current->state == TASK_DEAD)
 		ret = restore_one_zombie(pid, core);
-	else if (current->state == TASK_HELPER)
+	else if (current->state == TASK_HELPER) {
+		restore_finish_stage(CR_STATE_RESTORE);
 		ret = 0;
-	else {
+	} else {
 		pr_err("Unknown state in code %d\n", (int)core->tc->task_state);
 		ret = -1;
 	}
@@ -1496,6 +1499,7 @@ static inline int stage_participants(int next_stage)
 	case CR_STATE_FORKING:
 		return task_entries->nr_tasks + task_entries->nr_helpers;
 	case CR_STATE_RESTORE:
+		return task_entries->nr_threads + task_entries->nr_helpers;
 	case CR_STATE_RESTORE_SIGCHLD:
 		return task_entries->nr_threads;
 	case CR_STATE_RESTORE_CREDS:
@@ -2712,6 +2716,12 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	task_args->tcp_socks_nr = rst_tcp_socks_nr;
 	task_args->tcp_socks = rst_mem_remap_ptr(tcp_socks, RM_PRIVATE);
 
+	task_args->n_helpers = n_helpers;
+	if (n_helpers > 0)
+		task_args->helpers = rst_mem_remap_ptr(helpers_pos, RM_PRIVATE);
+	else
+		task_args->helpers = NULL;
+
 	/*
 	 * Arguments for task restoration.
 	 */

include/restorer.h

@@ -150,6 +150,9 @@ struct task_restore_args {
 
 	int				fd_last_pid; /* sys.ns_last_pid for threads rst */
 
+	pid_t				*helpers /* the TASK_HELPERS to wait on at the end of restore */;
+	int				n_helpers;
+
 #ifdef CONFIG_VDSO
 	unsigned long			vdso_rt_size;
 	struct vdso_symtable		vdso_sym_rt;		/* runtime vdso symbols */

pie/restorer.c

@@ -51,6 +51,8 @@ static struct task_entries *task_entries;
 static futex_t thread_inprogress;
 static futex_t zombies_inprogress;
 static int cap_last_cap;
+static pid_t *helpers;
+static int n_helpers;
 
 extern void cr_restore_rt (void) asm ("__cr_restore_rt")
 			__attribute__ ((visibility ("hidden")));
@@ -58,6 +60,13 @@ extern void cr_restore_rt (void) asm ("__cr_restore_rt")
 static void sigchld_handler(int signal, siginfo_t *siginfo, void *data)
 {
 	char *r;
+	int i;
+
+	/* We can ignore helpers that die, we expect them to after
+	 * CR_STATE_RESTORE is finished. */
+	for (i = 0; i < n_helpers; i++)
+		if (siginfo->si_pid == helpers[i])
+			return;
 
 	if (futex_get(&task_entries->start) == CR_STATE_RESTORE_SIGCHLD) {
 		pr_debug("%ld: Collect a zombie with (pid %d, %d)\n",
@@ -702,6 +711,29 @@ static int unmap_old_vmas(void *premmapped_addr, unsigned long premmapped_len,
 	return 0;
 }
 
+static int wait_helpers(struct task_restore_args *task_args)
+{
+	int i;
+
+	for (i = 0; i < task_args->n_helpers; i++) {
+		int status;
+		pid_t pid = task_args->helpers[i];
+
+		/* Check that a helper completed. */
+		if (sys_waitpid(pid, &status, 0, NULL) == -1) {
+			/* It has been waited in sigchld_handler */
+			continue;
+		}
+		if (!WIFEXITED(status) || WEXITSTATUS(status)) {
+			pr_err("%d exited with non-zero code (%d,%d)\n", pid,
+				WEXITSTATUS(status), WTERMSIG(status));
+			return -1;
+		}
+	}
+
+	return 0;
+}
+
 /*
  * The main routine to restore task via sigreturn.
  * This one is very special, we never return there
@@ -730,6 +762,8 @@ long __export_restore_task(struct task_restore_args *args)
 #endif
 
 	task_entries = args->task_entries;
+	helpers = args->helpers;
+	n_helpers = args->n_helpers;
 
 	ksigfillset(&act.rt_sa_mask);
 	act.rt_sa_handler = sigchld_handler;
@@ -1020,6 +1054,9 @@ long __export_restore_task(struct task_restore_args *args)
 
 	futex_wait_while_gt(&zombies_inprogress, 0);
 
+	if (wait_helpers(args) < 0)
+		goto core_restore_end;
+
 	ksigfillset(&to_block);
 	ret = sys_sigprocmask(SIG_SETMASK, &to_block, NULL, sizeof(k_rtsigset_t));
 	if (ret) {