task-hook-timer-demo.jl

function task_hook(code, to)
    from_tls = task_local_storage()
    if code == 0
        to::Task
        # Creation, setup child counters
        @assert to.storage === nothing
        to_tls = to.storage = IdDict{Any, Any}()
        to_tls[:counter] = UInt64(0)
        to_tls[:child_counter] = UInt64(0)

        # Attach finalizer to child so we can collect their counter
        finalizer(to) do _
            from_tls[:child_counter] += to_tls[:counter]
        end

        # Setup task_hook on child
        Base.attach_task_hook!(to, task_hook_ptr)
    elseif code == 1 || code == 3
        # First execution or switch back, start our timer
        from_tls[:last_time] = time_ns()
    elseif code == 2 || code == 4
        # Switch away or finished, stop our timer
        time_now = time_ns()
        last_time = from_tls[:last_time]::UInt64
        from_tls[:last_time] = time_now
        from_tls[:counter] += time_now - last_time
    end
    nothing
end
const task_hook_ptr = @cfunction(task_hook, Cvoid, (UInt8, Any))

function start_timing!()
    Base.attach_task_hook!(task_hook_ptr)
    tls = task_local_storage()
    tls[:last_time] = time_ns()
    tls[:counter] = UInt64(0)
    tls[:child_counter] = UInt64(0)
end

function stop_timing!()
    time_now = time_ns()

    Base.detach_task_hook!(task_hook_ptr)

    tls = task_local_storage()
    last_time = get(tls, :last_time, time_now)::UInt64
    child_time = tls[:child_counter]::UInt64 / 1000^3
    our_time = (tls[:counter]::UInt64 + time_now - last_time) / 1000^3
    return (;child_time, our_time)
end

function test_timings()
    local times
    wait(@async begin
        start_timing!()

        t = @async begin
            for i in 1:5
                println(devnull, "test")
            end
            sleep(1)
        end

        wait(t)
        finalize(t)

        times = stop_timing!()
    end)
    return times
end
test_timings()
times = test_timings()
@show times