[CRIU] [PATCH 2/4] restore: TASK_HELPERs live until RESTORE_FS stage

[Tycho Andersen]

In order to use TASK_HELPERS to open files from dead processes, they should
persist until criu is done restoring the filesystem (i.e. until after
restore_fs is called). To accomplish this, we add a new stage called
CR_STATE_RESTORE_FS and require all helpers to wait until that stage is done
before exiting, then we wait() on the helpers after that 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

Signed-off-by: Tycho Andersen <tycho.andersen at canonical.com>
---
 cr-restore.c       | 23 ++++++++++++++++++-----
 include/restorer.h |  1 +
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/cr-restore.c b/cr-restore.c
index 4d5ccd5..7bcdf6e 100644
--- a/cr-restore.c
+++ b/cr-restore.c
@@ -770,9 +770,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;
 
@@ -859,6 +856,7 @@ static int restore_one_zombie(int pid, CoreEntry *core)
 	sys_prctl(PR_SET_NAME, (long)(void *)core->tc->comm, 0, 0, 0);
 
 	if (task_entries != NULL) {
+		restore_finish_stage(CR_STATE_RESTORE_FS);
 		restore_finish_stage(CR_STATE_RESTORE);
 		zombie_prepare_signals();
 		mutex_lock(&task_entries->zombie_lock);
@@ -931,9 +929,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_FS);
 		ret = 0;
-	else {
+	} else {
 		pr_err("Unknown state in code %d\n", (int)core->tc->task_state);
 		ret = -1;
 	}
@@ -1489,6 +1488,8 @@ static inline int stage_participants(int next_stage)
 		return 1;
 	case CR_STATE_FORKING:
 		return task_entries->nr_tasks + task_entries->nr_helpers;
+	case CR_STATE_RESTORE_FS:
+		return task_entries->nr_threads + task_entries->nr_helpers;
 	case CR_STATE_RESTORE:
 	case CR_STATE_RESTORE_SIGCHLD:
 		return task_entries->nr_threads;
@@ -1703,6 +1704,10 @@ static int restore_root_task(struct pstree_item *init)
 
 	timing_stop(TIME_FORK);
 
+	ret = restore_switch_stage(CR_STATE_RESTORE_FS);
+	if (ret < 0)
+		goto out_kill;
+
 	ret = restore_switch_stage(CR_STATE_RESTORE);
 	if (ret < 0)
 		goto out_kill;
@@ -2844,6 +2849,14 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	if (restore_fs(current))
 		goto err;
 
+	/* restore_finish_stage has task_entries hardcoded. Since we moved it
+	 * above, the current pointer is invalid so we need to reset it. */
+	task_entries = task_args->task_entries;
+	restore_finish_stage(CR_STATE_RESTORE_FS);
+
+	if (pstree_wait_helpers() < 0)
+		goto err;
+
 	close_image_dir();
 	close_proc();
 	close_service_fd(ROOT_FD_OFF);
diff --git a/include/restorer.h b/include/restorer.h
index a690341..dbb7109 100644
--- a/include/restorer.h
+++ b/include/restorer.h
@@ -169,6 +169,7 @@ enum {
 	CR_STATE_FAIL		= -1,
 	CR_STATE_RESTORE_NS	= 0, /* is used for executing "setup-namespace" scripts */
 	CR_STATE_FORKING,
+	CR_STATE_RESTORE_FS,
 	CR_STATE_RESTORE,
 	CR_STATE_RESTORE_SIGCHLD,
 	/*
-- 
1.9.1