rst: Helper for restore stage barrier

When finishing a stage we have to report this (decrement the
number of tasks in stage) and wait while stage switch. Write
a helper that does both.
Signed-off-by: Pavel Emelyanov <xemul@parallels.com>

cr-restore.c

@@ -671,13 +671,9 @@ static int restore_one_zombie(int pid, int exit_code)
 	pr_info("Restoring zombie with %d code\n", exit_code);
 
 	if (task_entries != NULL) {
-		futex_dec_and_wake(&task_entries->nr_in_progress);
-		futex_wait_while(&task_entries->start, CR_STATE_RESTORE);
-
+		restore_finish_stage(CR_STATE_RESTORE);
 		zombie_prepare_signals();
-
-		futex_dec_and_wake(&task_entries->nr_in_progress);
-		futex_wait_while(&task_entries->start, CR_STATE_RESTORE_SIGCHLD);
+		restore_finish_stage(CR_STATE_RESTORE_SIGCHLD);
 	}
 
 	if (exit_code & 0x7f) {
@@ -1085,16 +1081,13 @@ static int restore_task_with_children(void *_arg)
 	if (current->pgid == current->pid.virt)
 		restore_pgid();
 
-	futex_dec_and_wake(&task_entries->nr_in_progress);
-	futex_wait_while(&task_entries->start, CR_STATE_FORKING);
+	restore_finish_stage(CR_STATE_FORKING);
 
 	if (current->pgid != current->pid.virt)
 		restore_pgid();
 
-	if (current->state != TASK_HELPER) {
-		futex_dec_and_wake(&task_entries->nr_in_progress);
-		futex_wait_while(&task_entries->start, CR_STATE_RESTORE_PGID);
-	}
+	if (current->state != TASK_HELPER)
+		restore_finish_stage(CR_STATE_RESTORE_PGID);
 
 	return restore_one_task(current->pid.virt);
 }

include/restorer.h

@@ -272,6 +272,12 @@ find_shmem(struct shmems *shmems, unsigned long shmid)
 	return NULL;
 }
 
+#define restore_finish_stage(__stage) do {				\
+		futex_dec_and_wake(&task_entries->nr_in_progress);	\
+		futex_wait_while(&task_entries->start, __stage);	\
+	} while (0)
+
+
 /* the restorer_blob_offset__ prefix is added by gen_offsets.sh */
 #define restorer_sym(rblob, name)	((void *)(rblob) + restorer_blob_offset__##name)
 

pie/restorer.c

@@ -232,13 +232,12 @@ long __export_restore_thread(struct thread_restore_args *args)
 
 	restore_creds(&args->ta->creds);
 
-	futex_dec_and_wake(&task_entries->nr_in_progress);
 
 	pr_info("%ld: Restored\n", sys_gettid());
 
-	futex_wait_while(&task_entries->start, CR_STATE_RESTORE);
-	futex_dec_and_wake(&task_entries->nr_in_progress);
-	futex_wait_while(&task_entries->start, CR_STATE_RESTORE_SIGCHLD);
+	restore_finish_stage(CR_STATE_RESTORE);
+	restore_finish_stage(CR_STATE_RESTORE_SIGCHLD);
+
 	futex_dec_and_wake(&thread_inprogress);
 
 	new_sp = (long)rt_sigframe + 8;
@@ -732,19 +731,15 @@ long __export_restore_task(struct task_restore_core_args *args)
 
 	restore_creds(&args->creds);
 
-	futex_dec_and_wake(&args->task_entries->nr_in_progress);
-
 	pr_info("%ld: Restored\n", sys_getpid());
 
-	futex_wait_while(&args->task_entries->start, CR_STATE_RESTORE);
+	restore_finish_stage(CR_STATE_RESTORE);
 
 	sys_sigaction(SIGCHLD, &args->sigchld_act, NULL, sizeof(rt_sigset_t));
 
 	futex_set_and_wake(&thread_inprogress, args->nr_threads);
 
-	futex_dec_and_wake(&args->task_entries->nr_in_progress);
-
-	futex_wait_while(&args->task_entries->start, CR_STATE_RESTORE_SIGCHLD);
+	restore_finish_stage(CR_STATE_RESTORE_SIGCHLD);
 
 	/* Wait until children stop to use args->task_entries */
 	futex_wait_while_gt(&thread_inprogress, 1);