compel/infect: Use waitpid with specific PIDs in compel_stop_tasks_on_syscall

When CRIU traces a large number of tasks, calling wait4(-1, ...) can
be slow because the kernel has to iterate over many tasks to find one
that has changed state. This overhead becomes significant when
thousands of tasks are involved. The time complexity of the old approach
was O(n^2) because we called wait4(-1, ...) N times, and each call
took O(N) time in the kernel.

This patch introduces compel_stop_tasks_on_syscall which takes an
array of PIDs and waits for each one specifically. This avoids the
O(N) search in the kernel's do_wait implementation and significantly
improves performance during the restore phase when many tasks are
being resumed and stopped at the rt_sigreturn syscall.

Some parallelism still exists in this approach. The loop waits for each
process to stop on a syscall, verifies it, and then lets it continue
running. While we move to wait for the next process in the array, the
previously processed ones continue their execution in parallel.

The performance profile of the bottleneck showed:
   - 70.45% entry_SYSCALL_64
      - do_syscall_64
         - 68.29% __do_sys_wait4
            - kernel_wait4
               - 68.28% do_wait
                  - 68.26% __do_wait
                     - 65.69% wait_consider_task
                          8.71% _raw_spin_lock_irq
                          2.39% _raw_spin_unlock_irq

Signed-off-by: Andrei Vagin <avagin@google.com>

compel/include/uapi/infect.h

@@ -93,6 +93,7 @@ enum trace_flags {
 };
 
 extern int __must_check compel_stop_on_syscall(int tasks, int sys_nr, int sys_nr_compat);
+extern int compel_stop_tasks_on_syscall(int tasks, pid_t *pids, const int sys_nr, const int sys_nr_compat);
 
 extern int __must_check compel_stop_pie(pid_t pid, void *addr, bool no_bp);
 

compel/src/lib/infect.c

@@ -1853,6 +1853,96 @@ int compel_stop_on_syscall(int tasks, const int sys_nr, const int sys_nr_compat)
 	return 0;
 }
 
+enum {
+	SYS_TRAP_UNSPEC = 0, /* the target syscall isn't found yet */
+	SYS_TRAP_ENTER  = 1, /* enter to the target syscall */
+	SYS_TRAP_EXIT   = 2, /* exit from the target syscall */
+};
+
+/*
+ * Trap multiple tasks on the exit from the specified syscall.
+ *
+ * wait4() with specific PIDs is used instead of wait4(-1, ...) to avoid
+ * the performance overhead of the kernel iterating over many tasks to
+ * find one that has changed state.
+ *
+ * nr_tasks - number of processes, which should be trapped
+ * pids - an array of process IDs
+ * sys_nr - the required syscall number
+ * sys_nr_compat - the required compatible syscall number
+ */
+int compel_stop_tasks_on_syscall(int nr_tasks, pid_t *pids, const int sys_nr, const int sys_nr_compat)
+{
+	user_regs_struct_t regs;
+	int status, ret, exit_code = -1;
+	int cont = 1, i;
+	uint8_t *done;
+	pid_t pid;
+
+	done = xzalloc(sizeof(done[0]) * nr_tasks);
+	if (!done)
+		return -1;
+
+	/* Stop all threads on the enter point in sys_rt_sigreturn */
+	while (cont) {
+		cont = 0;
+
+		for (i = 0; i < nr_tasks; i++) {
+			if (done[i] == SYS_TRAP_EXIT)
+				continue;
+			cont = 1;
+			pid = pids[i];
+			pid = wait4(pid, &status, __WALL, NULL);
+			if (pid == -1) {
+				pr_perror("wait4 failed");
+				goto err;
+			}
+
+			if (!task_is_trapped(status, pid))
+				goto err;
+
+			pr_debug("%d was trapped\n", pid);
+
+			if ((WSTOPSIG(status) & PTRACE_SYSCALL_TRAP) == 0) {
+				/*
+				 * On some platforms such as ARM64, it is impossible to
+				 * pass through a breakpoint, so let's clear it right
+				 * after it has been triggered.
+				*/
+				if (ptrace_flush_breakpoints(pid)) {
+					pr_err("Unable to clear breakpoints\n");
+					goto err;
+				}
+				goto goon;
+			}
+			if (done[i] == SYS_TRAP_ENTER) {
+				done[i] = SYS_TRAP_EXIT;
+				continue;
+			}
+
+			ret = ptrace_get_regs(pid, &regs);
+			if (ret) {
+				pr_perror("ptrace");
+				goto err;
+			}
+
+			if (is_required_syscall(&regs, pid, sys_nr, sys_nr_compat))
+				done[i] = SYS_TRAP_ENTER;
+		goon:
+			/* Let this task run while updating the others. */
+			ret = ptrace(PTRACE_SYSCALL, pid, NULL, NULL);
+			if (ret) {
+				pr_perror("ptrace");
+				goto err;
+			}
+		}
+	}
+	exit_code = 0;
+err:
+	xfree(done);
+	return exit_code;
+}
+
 int compel_mode_native(struct parasite_ctl *ctl)
 {
 	return user_regs_native(&ctl->orig.regs);

criu/cr-restore.c

@@ -1828,10 +1828,11 @@ static int restore_rseq_cs(void)
 	return 0;
 }
 
-static int catch_tasks(bool root_seized)
+static int catch_tasks(bool root_seized, pid_t *pids, int nr_tasks)
 {
 	struct pstree_item *item;
 	bool nobp = fault_injected(FI_NO_BREAKPOINTS) || !kdat.has_breakpoints;
+	int npids = 0;
 
 	for_each_pstree_item(item) {
 		int status, i, ret;
@@ -1849,6 +1850,12 @@ static int catch_tasks(bool root_seized)
 		for (i = 0; i < item->nr_threads; i++) {
 			pid_t pid = item->threads[i].real;
 
+			if (nr_tasks <= npids) {
+				pr_err("Too many threads discovered while catching tasks\n");
+				return -1;
+			}
+			pids[npids] = pid;
+			npids++;
 			if (ptrace(PTRACE_INTERRUPT, pid, 0, 0)) {
 				pr_perror("Can't interrupt the %d task", pid);
 				return -1;
@@ -2011,6 +2018,7 @@ static int restore_root_task(struct pstree_item *init)
 	int ret, fd, mnt_ns_fd = -1;
 	int root_seized = 0;
 	struct pstree_item *item;
+	pid_t *pids = NULL;
 
 	ret = run_scripts(ACT_PRE_RESTORE);
 	if (ret != 0) {
@@ -2229,7 +2237,10 @@ skip_ns_bouncing:
 
 	timing_stop(TIME_RESTORE);
 
-	if (catch_tasks(root_seized)) {
+	pids = xzalloc(sizeof(pid_t) * task_entries->nr_threads);
+	if (!pids)
+		goto out_kill_network_unlocked;
+	if (catch_tasks(root_seized, pids, task_entries->nr_threads)) {
 		pr_err("Can't catch all tasks\n");
 		goto out_kill_network_unlocked;
 	}
@@ -2239,11 +2250,13 @@ skip_ns_bouncing:
 
 	__restore_switch_stage(CR_STATE_COMPLETE);
 
-	ret = compel_stop_on_syscall(task_entries->nr_threads, __NR(rt_sigreturn, 0), __NR(rt_sigreturn, 1));
+	ret = compel_stop_tasks_on_syscall(task_entries->nr_threads, pids, __NR(rt_sigreturn, 0), __NR(rt_sigreturn, 1));
 	if (ret) {
 		pr_err("Can't stop all tasks on rt_sigreturn\n");
 		goto out_kill_network_unlocked;
 	}
+	xfree(pids);
+	pids = NULL;
 
 	finalize_restore();
 
@@ -2329,6 +2342,7 @@ out_kill:
 	}
 
 out:
+	xfree(pids);
 	depopulate_roots_yard(mnt_ns_fd, true);
 	stop_usernsd();
 	__restore_switch_stage(CR_STATE_FAIL);