plugin: Add DUMP_DEVICE_LATE callback

Previously, amdgpu plugin was determining when to call its
UNPAUSE ioctl by counting the files that have been restored.

This was not reliable; there may be more or fewer device files
than expected and there may be other processes still checkpointing
when unpause was called.

Add a new plugin callback DUMP_DEVICE_LATE which is called after
files are finished checkpointing for all processes.

Signed-off-by: David Francis <David.Francis@amd.com>

Author: fdavid-amd

criu/cr-dump.c

@@ -2225,6 +2225,9 @@ int cr_dump_tasks(pid_t pid)
 			goto err;
 	}
 
+	if(run_plugins(DUMP_DEVICE_LATE, pid))
+		goto err;
+
 	if (parent_ie) {
 		inventory_entry__free_unpacked(parent_ie, NULL);
 		parent_ie = NULL;

criu/include/criu-plugin.h

@@ -62,6 +62,8 @@ enum {
 
 	CR_PLUGIN_HOOK__DMABUF_FD = 12,
 
+	CR_PLUGIN_HOOK__DUMP_DEVICE_LATE = 13,
+
 	CR_PLUGIN_HOOK__MAX
 };
 
@@ -81,6 +83,7 @@ DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__RESUME_DEVICES_LATE, int pid);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__PAUSE_DEVICES, int pid);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__CHECKPOINT_DEVICES, int pid);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__DMABUF_FD, int handle, int fd);
+DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__DUMP_DEVICE_LATE, int id);
 
 enum {
 	CR_PLUGIN_STAGE__DUMP,

criu/plugin.c

@@ -60,6 +60,7 @@ static cr_plugin_desc_t *cr_gen_plugin_desc(void *h, char *path)
 	__assign_hook(PAUSE_DEVICES, "cr_plugin_pause_devices");
 	__assign_hook(CHECKPOINT_DEVICES, "cr_plugin_checkpoint_devices");
 	__assign_hook(DMABUF_FD, "cr_plugin_dmabuf_fd");
+	__assign_hook(DUMP_DEVICE_LATE, "cr_plugin_dump_device_late");
 
 #undef __assign_hook
 

plugins/amdgpu/amdgpu_plugin.c

@@ -49,13 +49,6 @@ struct vma_metadata {
 
 /************************************ Global Variables ********************************************/
 
-/**
- * FD of KFD device used to checkpoint. On a multi-process
- * tree the order of checkpointing goes from parent to child
- * and so on - so saving the FD will not be overwritten
- */
-static int kfd_checkpoint_fd;
-
 static LIST_HEAD(update_vma_info_list);
 
 size_t kfd_max_buffer_size;
@@ -1007,28 +1000,39 @@ int restore_hsakmt_shared_mem(const uint64_t shared_mem_size, const uint32_t sha
 	return 0;
 }
 
-static int unpause_process(int fd)
+int amdgpu_unpause_processes(int pid)
 {
 	int ret = 0;
 	struct kfd_ioctl_criu_args args = { 0 };
+	struct list_head *l = get_dumped_fds();
+	struct dumped_fd *st;
+
+	list_for_each_entry(st, l, l) {
+		if (st->is_drm) {
+			ret = amdgpu_plugin_drm_unpause_file(st->fd);
+			if (ret) {
+				pr_perror("Failed to unpause drm device file");
+				goto exit;
+			}
+			close(st->fd);
+		} else {
+			args.op = KFD_CRIU_OP_UNPAUSE;
 
-	args.op = KFD_CRIU_OP_UNPAUSE;
-
-	ret = kmtIoctl(fd, AMDKFD_IOC_CRIU_OP, &args);
-	if (ret) {
-		pr_perror("Failed to unpause process");
-		goto exit;
+			ret = kmtIoctl(st->fd, AMDKFD_IOC_CRIU_OP, &args);
+			if (ret) {
+				pr_perror("Failed to unpause process");
+				goto exit;
+			}
+		}
 	}
 
-	// Reset the KFD FD
-	kfd_checkpoint_fd = -1;
-	sys_close_drm_render_devices(&src_topology);
-
 exit:
 	pr_info("Process unpaused %s (ret:%d)\n", ret ? "Failed" : "Ok", ret);
+	clear_dumped_fds();
 
 	return ret;
 }
+CR_PLUGIN_REGISTER_HOOK(CR_PLUGIN_HOOK__DUMP_DEVICE_LATE, amdgpu_unpause_processes)
 
 static int save_devices(int fd, struct kfd_ioctl_criu_args *args, struct kfd_criu_device_bucket *device_buckets,
 			CriuKfd *e)
@@ -1230,9 +1234,6 @@ int amdgpu_plugin_dump_file(int fd, int id)
 		return -1;
 	}
 
-	/* Initialize number of device files that will be checkpointed */
-	init_gpu_count(&src_topology);
-
 	/* Check whether this plugin was called for kfd or render nodes */
 	if (major(st.st_rdev) != major(st_kfd.st_rdev) || minor(st.st_rdev) != 0) {
 
@@ -1244,11 +1245,9 @@ int amdgpu_plugin_dump_file(int fd, int id)
 		if (ret)
 			return ret;
 
-		/* Invoke unpause process if needed */
-		decrement_checkpoint_count();
-		if (checkpoint_is_complete()) {
-			ret = unpause_process(kfd_checkpoint_fd);
-		}
+		ret = record_dumped_fd(fd, true);
+		if (ret)
+			return ret;
 
 		/* Need to return success here so that criu can call plugins for renderD nodes */
 		return ret;
@@ -1346,14 +1345,11 @@ int amdgpu_plugin_dump_file(int fd, int id)
 
 	xfree(buf);
 
-exit:
-	/* Restore all queues if conditions permit */
-	kfd_checkpoint_fd = fd;
-	decrement_checkpoint_count();
-	if (checkpoint_is_complete()) {
-		ret = unpause_process(fd);
-	}
+	ret = record_dumped_fd(fd, false);
+	if (ret)
+		goto exit;
 
+exit:
 	xfree((void *)args.devices);
 	xfree((void *)args.bos);
 	xfree((void *)args.priv_data);

plugins/amdgpu/amdgpu_plugin_drm.c

@@ -349,6 +349,22 @@ int amdgpu_plugin_drm_dump_file(int fd, int id, struct stat *drm)
 	return ret;
 }
 
+int amdgpu_plugin_drm_unpause_file(int fd)
+{
+	struct drm_amdgpu_criu_args args = {0};
+	int ret = 0;
+
+	args.op = AMDGPU_CRIU_OP_UNPAUSE;
+	if (drmIoctl(fd, DRM_IOCTL_AMDGPU_CRIU_OP, &args) == -1) {
+		pr_perror("Failed to call unpause ioctl");
+		ret = -1;
+		goto exit;
+	}
+
+	exit:
+	return ret;
+}
+
 int amdgpu_plugin_drm_restore_file(int fd, CriuRenderNode *rd)
 {
 	struct drm_amdgpu_criu_args args = {0};

plugins/amdgpu/amdgpu_plugin_drm.h

@@ -26,6 +26,8 @@ int amdgpu_plugin_drm_dump_file(int fd, int id, struct stat *drm);
 
 int amdgpu_plugin_drm_restore_file(int fd, CriuRenderNode *rd);
 
+int amdgpu_plugin_drm_unpause_file(int fd);
+
 int get_gem_handle(amdgpu_device_handle h_dev, int dmabuf_fd);
 
 int save_vma_updates(uint64_t offset, uint64_t addr, uint64_t restored_offset, int gpu_id);

plugins/amdgpu/amdgpu_plugin_util.c

@@ -38,9 +38,7 @@
 #include "amdgpu_plugin_util.h"
 #include "amdgpu_plugin_topology.h"
 
-/* Tracks number of device files that need to be checkpointed */
-static int dev_file_cnt = 0;
-
+static LIST_HEAD(dumped_fds);
 static LIST_HEAD(shared_bos);
 static LIST_HEAD(shared_dmabuf_fds);
 static LIST_HEAD(completed_work);
@@ -53,23 +51,23 @@ struct tp_system dest_topology;
 struct device_maps checkpoint_maps;
 struct device_maps restore_maps;
 
-bool checkpoint_is_complete()
-{
-	return (dev_file_cnt == 0);
-}
+int record_dumped_fd(int fd, bool is_drm) {
+	int newfd = dup(fd);
 
-void decrement_checkpoint_count()
-{
-	dev_file_cnt--;
-}
+	if (newfd < 0)
+		return newfd;
+	struct dumped_fd *st = malloc(sizeof(struct dumped_fd));
+	if (!st)
+		return -1;
+	st->fd = newfd;
+	st->is_drm = is_drm;
+	list_add(&st->l, &dumped_fds);
 
-void init_gpu_count(struct tp_system *topo)
-{
-	if (dev_file_cnt != 0)
-		return;
+	return 0;
+}
 
-	/* We add ONE to include checkpointing of KFD device */
-	dev_file_cnt = 1 + topology_gpu_count(topo);
+struct list_head *get_dumped_fds() {
+	return &dumped_fds;
 }
 
 bool shared_bo_has_exporter(int handle) {
@@ -174,6 +172,15 @@ void clear_completed_work_and_dmabuf_fds() {
 	}
 }
 
+void clear_dumped_fds() {
+	while (!list_empty(&dumped_fds)) {
+		struct dumped_fd *st = list_first_entry(&dumped_fds, struct dumped_fd, l);
+		list_del(&st->l);
+		close(st->fd);
+		free(st);
+	}
+}
+
 int read_fp(FILE *fp, void *buf, const size_t buf_len)
 {
        size_t len_read;

plugins/amdgpu/amdgpu_plugin_util.h

@@ -123,9 +123,9 @@ int read_file(const char *file_path, void *buf, const size_t buf_len);
 int write_img_file(char *path, const void *buf, const size_t buf_len);
 FILE *open_img_file(char *path, bool write, size_t *size);
 
-bool checkpoint_is_complete();
-void decrement_checkpoint_count();
-void init_gpu_count(struct tp_system *topology);
+int record_dumped_fd(int fd, bool is_drm);
+struct list_head *get_dumped_fds();
+void clear_dumped_fds();
 
 bool shared_bo_has_exporter(int handle);
 int record_shared_bo(int handle, bool is_imported);