[PATCH 2/2] rv/nomiss: close enable-time task exit race

@@ -181,6 +181,8 @@ static void deadline_storage_init(void)
 	}
 }
 
+static int deadline_create_task_storage(struct task_struct *task);
+
 /*
  * Initialise monitors for all tasks and pre-allocate the storage for servers.
  * This is necessary since we don't have access to the servers here and

@@ -190,9 +192,7 @@ static void deadline_storage_init(void)
 static inline int init_storage(bool skip_tasks)
 {
 	struct task_struct *g, *p;
-	int cpu;
-
-	deadline_storage_init();
+	int cpu, ret;
 
 	for_each_possible_cpu(cpu) {
 		if (!da_create_empty_storage(fair_server_id(cpu)))

@@ -208,25 +208,25 @@ static inline int init_storage(bool skip_tasks)
 	read_lock(&tasklist_lock);
 	for_each_process_thread(g, p) {
 		if (p->policy == SCHED_DEADLINE) {
-			if (!da_create_storage(EXPAND_ID_TASK(p), NULL)) {
+			ret = deadline_create_task_storage(p);
+			if (ret && ret != -ESRCH) {
 				read_unlock(&tasklist_lock);
-				goto fail;
+				return ret;
 			}
-			atomic_inc(&deadline_storage_shard(p->pid)->users);
 		}
 	}
 	read_unlock(&tasklist_lock);
 	return 0;
 
 fail:
-	da_monitor_destroy();
 	return -ENOMEM;
 }
 
-static void deadline_create_task_storage(struct task_struct *task)
+static int deadline_create_task_storage(struct task_struct *task)
 {
 	struct deadline_storage_shard *shard = deadline_storage_shard(task->pid);
 	bool created = false;
+	int ret = -ESRCH;
 
 	/*
 	 * The temporary reference makes an in-flight creator visible to exit.

@@ -239,14 +239,21 @@ static void deadline_create_task_storage(struct task_struct *task)
 	raw_spin_lock(&shard->lock);
 	if (!(READ_ONCE(task->flags) & PF_EXITING)) {
 		guard(rcu)();
-		if (!da_get_monitor(EXPAND_ID_TASK(task)) &&
-		    da_create_storage(EXPAND_ID_TASK(task), NULL))
+		if (da_get_monitor(EXPAND_ID_TASK(task))) {
+			ret = 0;
+		} else if (da_create_storage(EXPAND_ID_TASK(task), NULL)) {
 			created = true;
+			ret = 0;
+		} else {
+			ret = -ENOMEM;
+		}
 	}
 	raw_spin_unlock(&shard->lock);
 
 	if (!created)
 		atomic_dec(&shard->users);
+
+	return ret;
 }
 
 static void deadline_destroy_task_storage(struct task_struct *task)

@@ -231,18 +231,26 @@ static int enable_nomiss(void)
 	if (retval)
 		return retval;
 
+	deadline_storage_init();
+	rv_attach_trace_probe("nomiss", task_newtask, handle_newtask);
+	rv_attach_trace_probe("nomiss", sched_process_exit, handle_exit);
+	if (!should_skip_syscall_handle())
+		rv_attach_trace_probe("nomiss", sys_enter, handle_sys_enter);
+
 	retval = init_storage(false);
-	if (retval)
+	if (retval) {
+		rv_detach_trace_probe("nomiss", task_newtask, handle_newtask);
+		rv_detach_trace_probe("nomiss", sched_process_exit, handle_exit);
+		if (!should_skip_syscall_handle())
+			rv_detach_trace_probe("nomiss", sys_enter, handle_sys_enter);
+		ha_monitor_destroy();
 		return retval;
+	}
 	rv_attach_trace_probe("nomiss", sched_dl_replenish_tp, handle_dl_replenish);
 	rv_attach_trace_probe("nomiss", sched_dl_throttle_tp, handle_dl_throttle);
 	rv_attach_trace_probe("nomiss", sched_dl_server_stop_tp, handle_dl_server_stop);
 	rv_attach_trace_probe("nomiss", sched_switch, handle_sched_switch);
 	rv_attach_trace_probe("nomiss", sched_wakeup, handle_sched_wakeup);
-	if (!should_skip_syscall_handle())
-		rv_attach_trace_probe("nomiss", sys_enter, handle_sys_enter);
-	rv_attach_trace_probe("nomiss", task_newtask, handle_newtask);
-	rv_attach_trace_probe("nomiss", sched_process_exit, handle_exit);
 
 	return 0;
 }