critical section

CHANGELOG.md

@@ -1,6 +1,10 @@
 OhMyThreads.jl Changelog
 =========================
 
+Version 0.5.1
+-------------
+- ![Feature][badge-feature] Within a `@tasks` block one can now mark a section as "critical" via `@section :critical begin ... end`.
+
 Version 0.5.0
 -------------
 

docs/src/refs/api.md

@@ -11,6 +11,7 @@ CollapsedDocStrings = true
 @tasks
 @set
 @local
+@section
 ```
 
 ### Functions

src/OhMyThreads.jl

@@ -20,7 +20,7 @@ using .Schedulers: Scheduler, DynamicScheduler, StaticScheduler, GreedyScheduler
                    SerialScheduler
 include("implementation.jl")
 
-export @tasks, @set, @local
+export @tasks, @set, @local, @section
 export treduce, tmapreduce, treducemap, tmap, tmap!, tforeach, tcollect
 export Scheduler, DynamicScheduler, StaticScheduler, GreedyScheduler, SerialScheduler
 

src/macro_impl.jl

@@ -20,6 +20,7 @@ function tasks_macro(forex)
     locals_before, locals_names = _maybe_handle_atlocal_block!(forbody.args)
     tls_names = isnothing(locals_before) ? [] : map(x -> x.args[1], locals_before)
     _maybe_handle_atset_block!(settings, forbody.args)
+    setup_sections = _maybe_handle_atsection_blocks!(forbody.args)
 
     forbody = esc(forbody)
     itrng = esc(itrng)
@@ -39,19 +40,22 @@ function tasks_macro(forex)
     end
     q = if isgiven(settings.reducer)
         quote
+            $setup_sections
             $make_mapping_function
             tmapreduce(mapping_function, $(settings.reducer),
                 $(itrng))
         end
     elseif isgiven(settings.collect)
         maybe_warn_useless_init(settings)
         quote
+            $setup_sections
             $make_mapping_function
             tmap(mapping_function, $(itrng))
         end
     else
         maybe_warn_useless_init(settings)
         quote
+            $setup_sections
             $make_mapping_function
             tforeach(mapping_function, $(itrng))
         end
@@ -68,7 +72,7 @@ function tasks_macro(forex)
     for (k, v) in settings.kwargs
         push!(kwexpr.args, Expr(:kw, k, v))
     end
-    insert!(q.args[4].args, 2, kwexpr)
+    insert!(q.args[6].args, 2, kwexpr)
 
     # wrap everything in a let ... end block
     # and, potentially, define the `TaskLocalValue`s.
@@ -151,16 +155,15 @@ function _atlocal_assign_to_exprs(ex)
         tls_type = esc(left_ex.args[2])
         local_before = :($(tl_storage) = TaskLocalValue{$tls_type}(() -> $(tls_def)))
     else
-        tls_sym  = esc(left_ex)
+        tls_sym = esc(left_ex)
         local_before = :($(tl_storage) = let f = () -> $(tls_def)
-                             TaskLocalValue{Core.Compiler.return_type(f, Tuple{})}(f)
-                         end)
+            TaskLocalValue{Core.Compiler.return_type(f, Tuple{})}(f)
+        end)
     end
     local_name = :($(tls_sym))
     return local_before, local_name
 end
 
-
 function _maybe_handle_atset_block!(settings, args)
     idcs = findall(args) do arg
         arg isa Expr && arg.head == :macrocall && arg.args[1] == Symbol("@set")
@@ -201,3 +204,32 @@ function _handle_atset_single_assign!(settings, ex)
         push!(settings.kwargs, sym => esc(def))
     end
 end
+
+function _maybe_handle_atsection_blocks!(args)
+    idcs = findall(args) do arg
+        arg isa Expr && arg.head == :macrocall && arg.args[1] == Symbol("@section")
+    end
+    isnothing(idcs) && return # no @section blocks
+    setup_sections = quote end
+    for i in idcs
+        kind = args[i].args[3]
+        body = args[i].args[4]
+        if kind isa QuoteNode
+            if kind.value == :critical
+                @gensym critical_lock
+                init_lock_ex = esc(:($(critical_lock) = $(ReentrantLock())))
+                push!(setup_sections.args, init_lock_ex)
+                args[i] = quote
+                    lock($(critical_lock)) do
+                        $(body)
+                    end
+                end
+            else
+                throw(ErrorException("Unknown section kind :$(kind.value)."))
+            end
+        else
+            throw(ErrorException("Wrong usage of @section. Must be followed by a symbol, indicating the kind of section."))
+        end
+    end
+    return setup_sections
+end

src/macros.jl

@@ -151,3 +151,35 @@ end
         error("The @local macro may only be used inside of a @tasks block.")
     end
 end
+
+"""
+    @section kind begin ... end
+
+This can be used inside a `@tasks for ... end` block to identify a section of code that
+must be executed according to a specific synchronization policy
+(as indicated by the symbol `kind`, see below).
+
+Multiple `@section` blocks are supported.
+
+## Kinds
+
+* `:critical`: Section of code that must be executed by a single task at a time (arbitrary order).
+
+## Example
+
+```julia
+@tasks for i in 1:10
+    @set ntasks = 10
+
+    println(i, ": before")
+    @section :critical begin
+        println(i, ": critical")
+        sleep(1)
+    end
+    println(i, ": after")
+end
+```
+"""
+macro section(args...)
+    error("The @section macro may only be used inside of a @tasks block.")
+end