From 187283775e76e0c0fe3cd02cd223e2157c515d31 Mon Sep 17 00:00:00 2001
From: Nathan Daly <NHDaly@gmail.com>
Date: Thu, 23 Sep 2021 07:29:19 -0400
Subject: [PATCH] Fix bug in `@time` compilation time measurement, using
 tryfinally macro. (#41923)

Use a tryfinally macro to ensure compile timer is closed, while maintaining scope.

This fixes a bug where an exception thrown _anywhere_, on any Task on
any thread, completely disables the compilation time measurement for any
concurrently executing `@time` blocks, without warning the user about
it.

Here are some examples of the current, broken behavior:
- Throwing (and catching) an exception disables compilation time
  measurement, even though the exception didn't escape, without warning
  the user - note 0 compilation time reported:
    ```julia
    julia> @time begin
              # The following exception disables compilation time measurement
              try throw("ha HAH!") catch end
              @eval module M ; f(x) = 2*x ; end
              @eval M.f(2)
          end
      0.003950 seconds (2.17 k allocations: 170.093 KiB)
    4
    ```
- Throwing an exception while the first task is sleeping disables
  compilation time measurement, so this ends up incorrectly reporting 0
  compilation time:
    ```julia
    julia> f(n) = begin
              # while we're sleeping, someone throws an exception somewhere else :'(
              sleep(n)
              @eval module M ; f(x) = 2*x ; end
              @eval M.f(2)
          end
    f (generic function with 1 method)

    julia> f(0) # warmup
    WARNING: replacing module M.
    4

    julia> @async @time f(5)
    Task (runnable) @0x000000010f2c0780

    julia> throw("oh no sad")
    ERROR: "oh no sad"
    Stacktrace:
    [1] top-level scope
      @ REPL[9]:1

    WARNING: replacing module M.
      5.008401 seconds (1.92 k allocations: 120.515 KiB)
    ```

After this commit, both of those examples correctly report their
compilation time.
---
 base/timing.jl | 24 ++++++++++++++++++------
 src/task.c     |  6 ------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/base/timing.jl b/base/timing.jl
index e229fbeb328e9..157fe288b9708 100644
--- a/base/timing.jl
+++ b/base/timing.jl
@@ -164,6 +164,16 @@ function timev_print(elapsedtime, diff::GC_Diff, compile_time)
     padded_nonzero_print(diff.full_sweep,   "full collections")
 end
 
+# Like a try-finally block, except without introducing the try scope
+# NOTE: This is deprecated and should not be used from user logic. A proper solution to
+# this problem will be introduced in https://github.com/JuliaLang/julia/pull/39217
+macro __tryfinally(ex, fin)
+    Expr(:tryfinally,
+       :($(esc(ex))),
+       :($(esc(fin)))
+       )
+end
+
 """
     @time
 
@@ -207,9 +217,10 @@ macro time(ex)
         local stats = gc_num()
         local elapsedtime = time_ns()
         local compile_elapsedtime = cumulative_compile_time_ns_before()
-        local val = $(esc(ex))
-        compile_elapsedtime = cumulative_compile_time_ns_after() - compile_elapsedtime
-        elapsedtime = time_ns() - elapsedtime
+        local val = @__tryfinally($(esc(ex)),
+            (elapsedtime = time_ns() - elapsedtime;
+            compile_elapsedtime = cumulative_compile_time_ns_after() - compile_elapsedtime)
+        )
         local diff = GC_Diff(gc_num(), stats)
         time_print(elapsedtime, diff.allocd, diff.total_time, gc_alloc_count(diff), compile_elapsedtime, true)
         val
@@ -253,9 +264,10 @@ macro timev(ex)
         local stats = gc_num()
         local elapsedtime = time_ns()
         local compile_elapsedtime = cumulative_compile_time_ns_before()
-        local val = $(esc(ex))
-        compile_elapsedtime = cumulative_compile_time_ns_after() - compile_elapsedtime
-        elapsedtime = time_ns() - elapsedtime
+        local val = @__tryfinally($(esc(ex)),
+            (elapsedtime = time_ns() - elapsedtime;
+            compile_elapsedtime = cumulative_compile_time_ns_after() - compile_elapsedtime)
+        )
         local diff = GC_Diff(gc_num(), stats)
         timev_print(elapsedtime, diff, compile_elapsedtime)
         val
diff --git a/src/task.c b/src/task.c
index 52cd2dfc5bcba..d7dda74f3ac99 100644
--- a/src/task.c
+++ b/src/task.c
@@ -560,12 +560,6 @@ static void JL_NORETURN throw_internal(jl_task_t *ct, jl_value_t *exception JL_M
     assert(!jl_get_safe_restore());
     jl_ptls_t ptls = ct->ptls;
     ptls->io_wait = 0;
-    // @time needs its compile timer disabled on error,
-    // and cannot use a try-finally as it would break scope for assignments
-    // We blindly disable compilation time tracking here, for all running Tasks, even though
-    // it may cause some incorrect measurements. This is a known bug, and is being tracked
-    // here: https://github.com/JuliaLang/julia/pull/39138
-    jl_atomic_store_relaxed(&jl_measure_compile_time_enabled, 0);
     JL_GC_PUSH1(&exception);
     jl_gc_unsafe_enter(ptls);
     if (exception) {