From 40dc5293c5dd1c5db157db914f2092d1d09c2187 Mon Sep 17 00:00:00 2001
From: Jingxian He <hejingxian@huawei.com>
Date: Wed, 19 May 2021 21:59:24 +0800
Subject: [PATCH 18/54] task_exit_notify: add task exit notify mask method for
 criu

Add task exit notify mask method for criu during kernel module upgrade.

Signed-off-by: Jingxian He <hejingxian@huawei.com>
---
 criu/config.c             |  1 +
 criu/cr-restore.c         |  8 ++++++++
 criu/crtools.c            |  1 +
 criu/include/cr_options.h |  1 +
 criu/include/util.h       |  5 +++++
 criu/seize.c              | 33 ++++++++++++++++++++++++++++++++-
 6 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/criu/config.c b/criu/config.c
index 0ccd2b5..9c4d8ce 100644
--- a/criu/config.c
+++ b/criu/config.c
@@ -547,6 +547,7 @@ int parse_options(int argc, char **argv, bool *usage_error,
 		BOOL_OPT("with-notifier", &opts.with_notifier_kup),
 		BOOL_OPT("with-fd-cred", &opts.with_fd_cred),
 		BOOL_OPT("dump-char-dev", &opts.dump_char_dev),
+		BOOL_OPT("mask-exit-notify", &opts.mask_exit_notify),
 		{ },
 	};
 
diff --git a/criu/cr-restore.c b/criu/cr-restore.c
index 7c198ce..ecebdfe 100644
--- a/criu/cr-restore.c
+++ b/criu/cr-restore.c
@@ -1583,6 +1583,14 @@ static inline int fork_with_pid(struct pstree_item *item)
 				item->pid->real, vpid(item));
 	}
 
+	if (opts.mask_exit_notify) {
+		int mask_pid = ret;
+		pr_info("start unmask for %d\n", mask_pid);
+		ret = mask_task_exit_notify(mask_pid, false);
+		if (ret)
+			pr_err("unmask exit notify fail for: %d\n", mask_pid);
+	}
+
 err_unlock:
 	if (!(clone_flags & CLONE_NEWPID))
 		unlock_last_pid();
diff --git a/criu/crtools.c b/criu/crtools.c
index 26010b5..8694ed0 100644
--- a/criu/crtools.c
+++ b/criu/crtools.c
@@ -451,6 +451,7 @@ usage:
 "  --with-fd-cred      Allow to make the restored process has the same cred\n"
 "                      as checkout assisted by kernel.\n"
 "  --dump-char-dev     Dump char dev files as normal file with repair cmd\n"
+"  --mask-exit-notify     Mask task exit notify during dump and restore\n"
 "\n"
 "Check options:\n"
 "  Without options, \"criu check\" checks availability of absolutely required\n"
diff --git a/criu/include/cr_options.h b/criu/include/cr_options.h
index 5ca177a..5b3ff86 100644
--- a/criu/include/cr_options.h
+++ b/criu/include/cr_options.h
@@ -181,6 +181,7 @@ struct cr_options {
 	int			with_notifier_kup;
 	int			with_fd_cred;
 	int			dump_char_dev;
+	int			mask_exit_notify;
 };
 
 extern struct cr_options opts;
diff --git a/criu/include/util.h b/criu/include/util.h
index cf9a8f4..3a4b8f9 100644
--- a/criu/include/util.h
+++ b/criu/include/util.h
@@ -425,4 +425,9 @@ void do_notifier_rollback(bool, enum notifier_state);
 int parse_devname(void);
 bool find_devname(const char *name);
 
+#define PID_BUF_SIZE  32
+#define MASK_EXIT_NOTIFY_DIR  "/sys/kernel/mask_exit_notify"
+#define UNMASK_EXIT_NOTIFY_DIR "/sys/kernel/unmask_exit_notify"
+int mask_task_exit_notify(int pid, bool mask);
+
 #endif /* __CR_UTIL_H__ */
diff --git a/criu/seize.c b/criu/seize.c
index a661097..387681c 100644
--- a/criu/seize.c
+++ b/criu/seize.c
@@ -649,9 +649,35 @@ free:
 	return ret < 0 ? ret : nr_inprogress;
 }
 
+int mask_task_exit_notify(int pid, bool mask)
+{
+	int fd, retval;
+	char buf[PID_BUF_SIZE] = {0};
+
+	if (pid <= 0)
+		return -1;
+
+	snprintf(buf, PID_BUF_SIZE - 1, "%d", pid);
+	if (mask)
+		fd = open(MASK_EXIT_NOTIFY_DIR, O_WRONLY, 0);
+	else
+		fd = open(UNMASK_EXIT_NOTIFY_DIR, O_WRONLY, 0);
+	if (fd < 0) {
+		pr_err("open mask exit notify file fail\n");
+		return fd;
+	}
+
+	retval = write(fd, buf, PID_BUF_SIZE);
+	if (retval < 0)
+		pr_err("Write mask exit pid: %s fail\n", buf);
+	close(fd);
+
+	return retval < 0 ? -1 : 0;
+}
+
 static void unseize_task_and_threads(const struct pstree_item *item, int st)
 {
-	int i;
+	int i, ret;
 
 	if (item->pid->state == TASK_DEAD)
 		return;
@@ -660,6 +686,11 @@ static void unseize_task_and_threads(const struct pstree_item *item, int st)
         	for (i = 0; i < item->nr_threads; i++)
                 	dump_task_special_pages(item->threads[i].real);
 	}
+	if (opts.mask_exit_notify) {
+		ret = mask_task_exit_notify(item->threads[0].real, true);
+		if (ret)
+			pr_err("mask exit notify for %d fail.\n", item->threads[0].real);
+	}
 
 	/*
 	 * The st is the state we want to switch tasks into,
-- 
2.33.1

