task-tls: Refactor into DTaskTLS struct

Author: jpsamaroo

src/Dagger.jl

@@ -22,6 +22,8 @@ else
     import Base.ScopedValues: ScopedValue, with
 end
 
+import TaskLocalValues: TaskLocalValue
+
 if !isdefined(Base, :get_extension)
     import Requires: @require
 end

src/array/indexing.jl

@@ -1,5 +1,3 @@
-import TaskLocalValues: TaskLocalValue
-
 ### getindex
 
 struct GetIndex{T,N} <: ArrayOp{T,N}

src/sch/Sch.jl

@@ -1202,7 +1202,7 @@ function proc_states(f::Base.Callable, uid::UInt64)
     end
 end
 proc_states(f::Base.Callable) =
-    proc_states(f, task_local_storage(:_dagger_sch_uid)::UInt64)
+    proc_states(f, Dagger.get_tls().sch_uid)
 
 task_tid_for_processor(::Processor) = nothing
 task_tid_for_processor(proc::Dagger.ThreadProc) = proc.tid

src/sch/dynamic.jl

@@ -17,7 +17,7 @@ struct SchedulerHandle
 end
 
 "Gets the scheduler handle for the currently-executing thunk."
-sch_handle() = task_local_storage(:_dagger_sch_handle)::SchedulerHandle
+sch_handle() = Dagger.get_tls().sch_handle::SchedulerHandle
 
 "Thrown when the scheduler halts before finishing processing the DAG."
 struct SchedulerHaltedException <: Exception end

src/task-tls.jl

@@ -1,41 +1,42 @@
 # In-Thunk Helpers
 
+struct DTaskTLS
+    processor::Processor
+    sch_uid::UInt
+    sch_handle::Any # FIXME: SchedulerHandle
+    task_spec::Vector{Any} # FIXME: TaskSpec
+end
+
+const DTASK_TLS = TaskLocalValue{Union{DTaskTLS,Nothing}}(()->nothing)
+
 """
-    task_processor()
+    get_tls() -> DTaskTLS
 
-Get the current processor executing the current Dagger task.
+Gets all Dagger TLS variable as a `DTaskTLS`.
 """
-task_processor() = task_local_storage(:_dagger_processor)::Processor
-@deprecate thunk_processor() task_processor()
+get_tls() = DTASK_TLS[]::DTaskTLS
 
 """
-    in_task()
+    set_tls!(tls)
 
-Returns `true` if currently executing in a [`DTask`](@ref), else `false`.
+Sets all Dagger TLS variables from `tls`, which may be a `DTaskTLS` or a `NamedTuple`.
 """
-in_task() = haskey(task_local_storage(), :_dagger_sch_uid)
-@deprecate in_thunk() in_task()
+function set_tls!(tls)
+    DTASK_TLS[] = DTaskTLS(tls.processor, tls.sch_uid, tls.sch_handle, tls.task_spec)
+end
 
 """
-    get_tls()
+    in_task() -> Bool
 
-Gets all Dagger TLS variable as a `NamedTuple`.
+Returns `true` if currently executing in a [`DTask`](@ref), else `false`.
 """
-get_tls() = (
-    sch_uid=task_local_storage(:_dagger_sch_uid),
-    sch_handle=task_local_storage(:_dagger_sch_handle),
-    processor=task_processor(),
-    task_spec=task_local_storage(:_dagger_task_spec),
-)
+in_task() = DTASK_TLS[] !== nothing
+@deprecate in_thunk() in_task()
 
 """
-    set_tls!(tls)
+    task_processor() -> Processor
 
-Sets all Dagger TLS variables from the `NamedTuple` `tls`.
+Get the current processor executing the current [`DTask`](@ref).
 """
-function set_tls!(tls)
-    task_local_storage(:_dagger_sch_uid, tls.sch_uid)
-    task_local_storage(:_dagger_sch_handle, tls.sch_handle)
-    task_local_storage(:_dagger_processor, tls.processor)
-    task_local_storage(:_dagger_task_spec, tls.task_spec)
-end
+task_processor() = get_tls().processor
+@deprecate thunk_processor() task_processor()