Read TLS directly if possible (#38)

Author: kabiroberai

Sources/NodeAPI/NodeActor.swift

@@ -39,47 +39,45 @@ private final class NodeExecutor: SerialExecutor {
     }
 
     func enqueue(_ job: UnownedJob) {
-        schedulerQueue.async {
-            // We want to access `job`'s task-local storage. To do so,
-            // this temporarily swaps ResumeTask for our own function.
-            // Then, swift_job_run is called, which sets the active task to
-            // the receiver and invokes its ResumeTask. We then execute the
-            // given closure, allowing us to grab task-local values. Finally,
-            // we "suspend" the task and return ResumeTask to its old value.
-            //
-            // on Darwin we can instead replace the "current task" thread-local
-            // (key 103) temporarily, but that isn't portable.
-            //
-            // This is sort of like inserting a "work(); await Task.yield()" block
-            // at the top of the task, since when a Task awaits it similarly changes
-            // the Resume function and suspends. Note that we can assume that this
-            // is a Task and not a basic Job, because Executor.enqueue is only
-            // called from swift_task_enqueue.
-            //
-            // Regarding `schedulerQueue.async`:
-            // Pre Swift 6.0 we didn't need a scheduler queue as enqueue would always
-            // run on the global queue. However, Swift 6 introduces optimizations in
-            // Task dispatch that allow tasks to be enqueued more efficiently, including
-            // that Task.init avoids a hop when possible. This, however, interferes with
-            // our `job.asCurrent` code because `asCurrent` relies on there being no already-
-            // running task (swift_job_run doesn't play well with nesting, it's possible but
-            // requires more private APIs, cf [swift_task_startOnMainActor][1]). The simplest
-            // solution is to hop onto our own queue for scheduling.
-            //
-            // [1]: https://github.com/apple/swift/blob/876c056153554f93b89dfd134794a05426ee789a/stdlib/public/Concurrency/Task.cpp#L1739
-            let target = job.asCurrent { NodeActor.target }
-
-            guard let q = target?.queue else {
-                nodeFatalError("There is no target NodeAsyncQueue associated with this Task")
-            }
-
-            let ref = self.asUnownedSerialExecutor()
-
-            do {
-                try q.run { job.runSynchronously(on: ref) }
-            } catch {
-                nodeFatalError("Could not execute job on NodeActor: \(error)")
-            }
+        // We want to enqueue the job on the "current" NodeActor, which is
+        // stored in task-local storage. In the Swift 6 compiler,
+        // NodeExecutor.enqueue is invoked with the same isolation as the caller,
+        // which means we can simply read out the TaskLocal value to obtain
+        // this.
+        let target: NodeAsyncQueue.Handle?
+        #if compiler(>=6.0)
+        target = NodeActor.target
+        #else
+        // It's a bit trickier in Swift <= 5.10 as Swift first makes a hop to the
+        // global executor before invoking this method. So instead we have to use
+        // some runtime spelunking to read the TaskLocal value from `job`.
+        // To do so, we temporarily swap ResumeTask for our own function.
+        // Then, swift_job_run is called, which sets the active task to
+        // the receiver and invokes its ResumeTask. We then execute the
+        // given closure, allowing us to grab task-local values. Finally,
+        // we "suspend" the task and return ResumeTask to its old value.
+        //
+        // on Darwin we can instead replace the "current task" thread-local
+        // (key 103) temporarily, but that isn't portable.
+        //
+        // This is sort of like inserting a "work(); await Task.yield()" block
+        // at the top of the task, since when a Task awaits it similarly changes
+        // the Resume function and suspends. Note that we can assume that this
+        // is a Task and not a basic Job, because Executor.enqueue is only
+        // called from swift_task_enqueue.
+        target = job.asCurrent { NodeActor.target }
+        #endif
+
+        guard let q = target?.queue else {
+            nodeFatalError("There is no target NodeAsyncQueue associated with this Task")
+        }
+
+        let ref = self.asUnownedSerialExecutor()
+
+        do {
+            try q.run { job.runSynchronously(on: ref) }
+        } catch {
+            nodeFatalError("Could not execute job on NodeActor: \(error)")
         }
     }