From 6420885bce8839f8b3a15991d41c74181ebbf468 Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash@gmail.com>
Date: Fri, 22 Oct 2021 01:24:53 -0400
Subject: [PATCH] task: schedule sticky tasks correctly with _wait2 (#42750)

The `_wait2` function is similar to calling `schedule + wait` from
another `@async` task, but optimized, so we want to observe the same
side-effects of making the task sticky to the correct thread (and not
accidentally making it sticky to the later task that handles the event).

Refs #41334 (75858d73322)

Co-authored-by: Takafumi Arakaki <aka.tkf@gmail.com>
---
 base/condition.jl | 11 +++++++++++
 base/task.jl      | 13 +++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/base/condition.jl b/base/condition.jl
index 71db551498170..d5f2277eac863 100644
--- a/base/condition.jl
+++ b/base/condition.jl
@@ -82,6 +82,17 @@ function _wait2(c::GenericCondition, waiter::Task)
     ct = current_task()
     assert_havelock(c)
     push!(c.waitq, waiter)
+    # since _wait2 is similar to schedule, we should observe the sticky bit now
+    if waiter.sticky && Threads.threadid(waiter) == 0
+        # Issue #41324
+        # t.sticky && tid == 0 is a task that needs to be co-scheduled with
+        # the parent task. If the parent (current_task) is not sticky we must
+        # set it to be sticky.
+        # XXX: Ideally we would be able to unset this
+        ct.sticky = true
+        tid = Threads.threadid()
+        ccall(:jl_set_task_tid, Cvoid, (Any, Cint), waiter, tid-1)
+    end
     return
 end
 
diff --git a/base/task.jl b/base/task.jl
index 0d4e5da4ccfd4..b25197e0aadcc 100644
--- a/base/task.jl
+++ b/base/task.jl
@@ -306,6 +306,18 @@ function _wait2(t::Task, waiter::Task)
         if !istaskdone(t)
             push!(t.donenotify.waitq, waiter)
             unlock(t.donenotify)
+            # since _wait2 is similar to schedule, we should observe the sticky
+            # bit, even if we aren't calling `schedule` due to this early-return
+            if waiter.sticky && Threads.threadid(waiter) == 0
+                # Issue #41324
+                # t.sticky && tid == 0 is a task that needs to be co-scheduled with
+                # the parent task. If the parent (current_task) is not sticky we must
+                # set it to be sticky.
+                # XXX: Ideally we would be able to unset this
+                current_task().sticky = true
+                tid = Threads.threadid()
+                ccall(:jl_set_task_tid, Cvoid, (Any, Cint), waiter, tid-1)
+            end
             return nothing
         else
             unlock(t.donenotify)
@@ -455,6 +467,7 @@ function errormonitor(t::Task)
         end
         nothing
     end
+    t2.sticky = false
     _wait2(t, t2)
     return t
 end