Repository: JuliaFolds2/OhMyThreads.jl (julia)

Issue:
`@only_one` block is executed by a single iteration (not a single task)
From my understanding of the docstring
```
  @only_one begin ... end

  This can be used inside a @tasks for ... end block to mark a region of code to be executed by only one
  of the parallel tasks (all other tasks skip over this region).
```
I was expecting
```julia
@tasks for i in 1:8
    @set ntasks = 2

    println(i, ": before")
    @only_one begin
        println(i, ": after")
    end
end
```
to print after 4 times (since a single task has to run 4 iterations of the loop).  Instead, I get the output
```
5: before
1: before
5: after
6: before
2: before
7: before
8: before
3: before
4: before
```
Is this expected? Should the docstring be changed or I just got it wrong? It seems to me that this macro should be called `@only_once` instead.

Pull request description:
Closes #100

Patch:
diff --git a/src/tools.jl b/src/tools.jl
index b0d95fd8..8029afe0 100644
--- a/src/tools.jl
+++ b/src/tools.jl
@@ -31,8 +31,8 @@ May be used to mark a region in parallel code to be executed by a single task on
 See [`try_enter!`](@ref) and [`reset!`](@ref).
 """
 mutable struct OnlyOneRegion
-    @atomic latch::Bool
-    OnlyOneRegion() = new(false)
+    @atomic task::Union{Task, Nothing}
+    OnlyOneRegion() = new(nothing)
 end
 
 """
@@ -62,15 +62,14 @@ end
 ```
 """
 function try_enter!(f, s::OnlyOneRegion)
-    latch = @atomic :monotonic s.latch
-    if latch
+    ct = current_task()
+    t = @atomic :monotonic s.task
+    if !isnothing(t) && ct != t
         return
     end
-    (_, success) = @atomicreplace s.latch false=>true
-    if !success
-        return
+    if ct == t || (@atomicreplace s.task nothing=>ct).success
+        f()
     end
-    f()
     return
 end
 
@@ -78,8 +77,8 @@ end
 Reset the `OnlyOneRegion` (so that it can be used again).
 """
 function reset!(s::OnlyOneRegion)
-    @atomicreplace s.latch true=>false
-    nothing
+    @atomic s.task = nothing
+    return
 end
 
 """


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index c436ab05..5681665f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -454,6 +454,22 @@ end
         catch ErrorException
             @test false
         end
+
+        x = 0
+        y = 0
+        try
+            @tasks for i in 1:10
+                @set ntasks = 2
+
+                y += 1 # not safe (race condition)
+                @only_one begin
+                    x += 1 # parallel-safe because only a single task will execute this
+                end
+            end
+            @test x == 5 # a single task should have incremented x 5 times
+        catch ErrorException
+            @test false
+        end
     end
 
     @testset "@only_one + @one_by_one" begin