On Thu, Jul 11, 2024 at 05:00:54PM +0200, Peter Zijlstra wrote:

quoted

Let me ponder that a little, I *can* make it work, but all 'solutions'
I've come up with so far are really rather vile.

This is the least horrible solution I could come up with...

---

--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h

@@ -83,6 +83,7 @@ struct uprobe_task {
 
 	struct timer_list		ri_timer;
 	struct callback_head		ri_task_work;
+	bool				ri_work_pending;
 	struct task_struct		*task;
 };
 
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c

@@ -1797,9 +1797,8 @@ void uprobe_free_utask(struct task_struc
 	t->utask = NULL;
 }
 
-static void return_instance_task_work(struct callback_head *head)
+static void __return_instance_put(struct uprobe_task *utask)
 {
-	struct uprobe_task *utask = container_of(head, struct uprobe_task, ri_task_work);
 	struct return_instance *ri;
 
 	for (ri = utask->return_instances; ri; ri = ri->next) {

@@ -1815,9 +1814,43 @@ static void return_instance_task_work(st
 	}
 }
 
+static void return_instance_task_work(struct callback_head *head)
+{
+	struct uprobe_task *utask = container_of(head, struct uprobe_task, ri_task_work);
+	utask->ri_work_pending = false;
+	__return_instance_put(utask);
+}
+
+static int return_instance_blocked(struct task_struct *p, void *arg)
+{
+	unsigned int state = READ_ONCE(p->__state);
+
+	if (state == TASK_RUNNING || state == TASK_WAKING)
+		return 0;
+
+	smp_rmb();
+	if (p->on_rq)
+		return 0;
+
+	/*
+	 * Per __task_needs_rq_locked() we now have: !p->on_cpu and only hold
+	 * p->pi_lock, and can consiter the task fully blocked.
+	 */
+
+	__return_instance_put(p->utask);