# This file is a part of Julia. License is MIT: https://julialang.org/license

using Test
using Distributed
using Dates
using Printf: @sprintf
using Base: Experimental
using Base.ScopedValues

include("choosetests.jl")
include("testenv.jl")
include("buildkitetestjson.jl")

const longrunning_delay = parse(Int, get(ENV, "JULIA_TEST_LONGRUNNING_DELAY", "45")) * 60 # minutes
const longrunning_interval = parse(Int, get(ENV, "JULIA_TEST_LONGRUNNING_INTERVAL", "15")) * 60 # minutes

(; tests, net_on, exit_on_error, use_revise, buildroot, seed) = choosetests(ARGS)
tests = unique(tests)

if Sys.islinux()
    const SYS_rrcall_check_presence = 1008
    global running_under_rr() = 0 == ccall(:syscall, Int,
        (Int, Int, Int, Int, Int, Int, Int),
        SYS_rrcall_check_presence, 0, 0, 0, 0, 0, 0)
else
    global running_under_rr() = false
end

const rmwait_timeout = running_under_rr() ? 300 : 30

ENV["JULIA_TEST_BUILDROOT"] = buildroot
if use_revise
    # First put this at the top of the DEPOT PATH to install revise if necessary.
    # Once it's loaded, we swizzle it to the end, to avoid confusing any tests.
    pushfirst!(DEPOT_PATH, joinpath(buildroot, "deps", "jlutilities", "depot"))
    using Pkg
    Pkg.activate(joinpath(@__DIR__, "..", "deps", "jlutilities", "revise"))
    Pkg.instantiate()
    using Revise
    union!(Revise.stdlib_names, Symbol.(STDLIBS))
    push!(DEPOT_PATH, popfirst!(DEPOT_PATH))
    # Remote-eval the following to initialize Revise in workers
    const revise_init_expr = quote
        ENV["JULIA_REVISE_WORKER_ONLY"] = "1"
        using Revise
        const STDLIBS = $STDLIBS
        union!(Revise.stdlib_names, Symbol.(STDLIBS))
        revise_trackall()
    end
end

if isempty(tests)
    println("No tests selected. Exiting.")
    exit()
end

const max_worker_rss = if haskey(ENV, "JULIA_TEST_MAXRSS_MB")
    parse(Int, ENV["JULIA_TEST_MAXRSS_MB"]) * 2^20
else
    typemax(Csize_t)
end
limited_worker_rss = max_worker_rss != typemax(Csize_t)

# Check all test files exist
isfiles = isfile.(test_path.(tests) .* ".jl")
if !all(isfiles)
    error("did not find test files for the following tests: ",
          join(tests[.!(isfiles)], ", "))
end

const node1_tests = String[]
function move_to_node1(t)
    if t in tests
        splice!(tests, findfirst(isequal(t), tests))
        push!(node1_tests, t)
    end
    nothing
end

# Base.compilecache only works from node 1, so precompile test is handled specially
move_to_node1("ccall")
move_to_node1("precompile")
move_to_node1("SharedArrays")
move_to_node1("threads")
move_to_node1("Distributed")
move_to_node1("gc")
# Ensure things like consuming all kernel pipe memory doesn't interfere with other tests
move_to_node1("stress")

# In a constrained memory environment, run the "distributed" test after all other tests
# since it starts a lot of workers and can easily exceed the maximum memory
limited_worker_rss && move_to_node1("Distributed")

# Move LinearAlgebra and Pkg tests to the front, because they take a while, so we might
# as well get them all started early. JuliaLowering_stdlibs both takes a while and
# uses a lot of memory at the beginning so try to run it early to keep total memory
# use flatter.
for prependme in ["LinearAlgebra", "Pkg", "JuliaLowering_stdlibs"]
    prependme_test_ids = findall(x->occursin(prependme, x), tests)
    prependme_tests = tests[prependme_test_ids]
    deleteat!(tests, prependme_test_ids)
    prepend!(tests, prependme_tests)
end

import LinearAlgebra
cd(@__DIR__) do
    # `net_on` implies that we have access to the loopback interface which is
    # necessary for Distributed multi-processing. There are some test
    # environments that do not allow access to loopback, so we must disable
    # addprocs when `net_on` is false. Note that there exist build environments,
    # including Nix, where `net_on` is false but we still have access to the
    # loopback interface. It would be great to make this check more specific to
    # identify those situations somehow. See
    #   * https://github.com/JuliaLang/julia/issues/6722
    #   * https://github.com/JuliaLang/julia/pull/29384
    #   * https://github.com/JuliaLang/julia/pull/40348
    n = 1
    JULIA_TEST_USE_MULTIPLE_WORKERS = Base.get_bool_env("JULIA_TEST_USE_MULTIPLE_WORKERS", false)
    # If the `JULIA_TEST_USE_MULTIPLE_WORKERS` environment variable is set to `true`, we use
    # multiple worker processes regardless of the value of `net_on`.
    # Otherwise, we use multiple worker processes if and only if `net_on` is true.
    if net_on || JULIA_TEST_USE_MULTIPLE_WORKERS
        n = min(Sys.EFFECTIVE_CPU_THREADS, length(tests))
        n > 1 && addprocs_with_testenv(n)
        LinearAlgebra.BLAS.set_num_threads(1)
    end
    skipped = 0

    @everywhere include("testdefs.jl")

    if use_revise
        @invokelatest revise_trackall()
        Distributed.remotecall_eval(Main, workers(), revise_init_expr)
    end

    println("""
        Running parallel tests with:
          getpid() = $(getpid())
          nworkers() = $(nworkers())
          nthreads(:interactive) = $(Threads.threadpoolsize(:interactive))
          nthreads(:default) = $(Threads.threadpoolsize(:default))
          Sys.CPU_THREADS = $(Sys.CPU_THREADS)
          Sys.total_memory() = $(Base.format_bytes(Sys.total_memory()))
          Sys.free_memory() = $(Base.format_bytes(Sys.free_memory()))
          Sys.uptime() = $(Sys.uptime()) ($(round(Sys.uptime() / (60 * 60), digits=1)) hours)
        """)

    #pretty print the information about gc and mem usage
    testgroupheader = "Test"
    workerheader = "(Worker)"
    name_align    = maximum([textwidth(testgroupheader) + textwidth(" ") + textwidth(workerheader); map(x -> textwidth(x) + 3 + ndigits(nworkers()), tests)])
    elapsed_align = textwidth("Time (s)")
    gc_align      = textwidth("GC (s)")
    percent_align = textwidth("GC %")
    alloc_align   = textwidth("Alloc (MB)")
    rss_align     = textwidth("RSS (MB)")
    printstyled(testgroupheader, color=:white)
    printstyled(lpad(workerheader, name_align - textwidth(testgroupheader) + 1), " | ", color=:white)
    printstyled("Time (s) | GC (s) | GC % | Alloc (MB) | RSS (MB)\n", color=:white)
    results = []
    print_lock = stdout isa Base.LibuvStream ? stdout.lock : ReentrantLock()
    if stderr isa Base.LibuvStream
        stderr.lock = print_lock
    end

    function print_testworker_stats(test, wrkr, resp)
        @nospecialize resp
        lock(print_lock)
        try
            printstyled(test, color=:white)
            printstyled(lpad("($wrkr)", name_align - textwidth(test) + 1, " "), " | ", color=:white)
            time_str = @sprintf("%7.2f",resp[2])
            printstyled(lpad(time_str, elapsed_align, " "), " | ", color=:white)
            gc_str = @sprintf("%5.2f", resp[5].total_time / 10^9)
            printstyled(lpad(gc_str, gc_align, " "), " | ", color=:white)

            # since there may be quite a few digits in the percentage,
            # the left-padding here is less to make sure everything fits
            percent_str = @sprintf("%4.1f", 100 * resp[5].total_time / (10^9 * resp[2]))
            printstyled(lpad(percent_str, percent_align, " "), " | ", color=:white)
            alloc_str = @sprintf("%5.2f", resp[3] / 2^20)
            printstyled(lpad(alloc_str, alloc_align, " "), " | ", color=:white)
            rss_str = @sprintf("%5.2f", resp[6] / 2^20)
            printstyled(lpad(rss_str, rss_align, " "), "\n", color=:white)
        finally
            unlock(print_lock)
        end
        nothing
    end

    global print_testworker_started = (name, wrkr)->begin
        pid = running_under_rr() ? remotecall_fetch(getpid, wrkr) : 0
        at = lpad("($wrkr)", name_align - textwidth(name) + 1, " ")
        lock(print_lock)
        try
            printstyled(name, at, " |", " "^elapsed_align, color=:white)
            printstyled("started at $(now())",
                    (pid > 0 ? " on pid $pid" : ""),
                    "\n", color=:light_black)
        finally
            unlock(print_lock)
        end
        nothing
    end

    function print_testworker_errored(name, wrkr, @nospecialize(e))
        lock(print_lock)
        try
            printstyled(name, color=:red)
            printstyled(lpad("($wrkr)", name_align - textwidth(name) + 1, " "), " |",
                " "^elapsed_align, " failed at $(now())\n", color=:red)
            if isa(e, Test.TestSetException)
                for t in e.errors_and_fails
                    show(t)
                    println()
                end
            elseif e !== nothing
                Base.showerror(stdout, e)
            end
            println()
        finally
            unlock(print_lock)
        end
    end


    all_tests = [tests; node1_tests]

    local stdin_monitor
    all_tasks = Task[]
    o_ts_duration = 0.0
    try
        # Monitor stdin and kill this task on ^C
        # but don't do this on Windows, because it may deadlock in the kernel
        running_tests = Dict{String, DateTime}()

        # Track timeout timers for each test
        test_timers = Dict{String, Timer}()

        if !Sys.iswindows() && isa(stdin, Base.TTY)
            t = current_task()
            stdin_monitor = @async begin
                trylock(stdin.raw_lock) || return
                term = Base.Terminals.TTYTerminal("xterm", stdin, stdout, stderr)
                try
                    Base.Terminals.raw!(term, true)
                    while true
                        c = read(term, Char)
                        if c == '\x3'
                            Base.throwto(t, InterruptException())
                            break
                        elseif c == '?'
                            println("Currently running: ")
                            tests = sort(collect(running_tests), by=x->x[2])
                            foreach(tests) do (test, date)
                                println(test, " (running for ", round(now()-date, Minute), ")")
                            end
                        end
                    end
                catch e
                    isa(e, InterruptException) || rethrow()
                finally
                    Base.Terminals.raw!(term, false)
                    unlock(stdin.raw_lock)
                end
            end
            Base.errormonitor(stdin_monitor)
        end
        o_ts_duration = @elapsed Experimental.@sync begin
            for p in workers()
                @async begin
                    push!(all_tasks, current_task())
                    while length(tests) > 0
                        test = popfirst!(tests)
                        running_tests[test] = now()
                        wrkr = p

                        # Create a timer for this test to report long-running status
                        test_timers[test] = Timer(longrunning_delay, interval=longrunning_interval) do timer
                            if haskey(running_tests, test)  # Check test is still running
                                start_time = running_tests[test]
                                elapsed = now() - start_time
                                elapsed_minutes = elapsed.value ÷ (1000 * 60)

                                elapsed_str = if elapsed_minutes >= 60
                                    hours, mins = divrem(elapsed_minutes, 60)
                                    "$(hours)h $(mins)m"
                                else
                                    "$(elapsed_minutes)m"
                                end

                                @lock print_lock begin
                                    print(test)
                                    print(lpad("($(wrkr))", name_align - textwidth(test) + 1, " "), " | ")
                                    # Calculate total width of data columns: "Time (s) | GC (s) | GC % | Alloc (MB) | RSS (MB)"
                                    # This is: elapsed_align + 3 + gc_align + 3 + percent_align + 3 + alloc_align + 3 + rss_align
                                    data_width = elapsed_align + gc_align + percent_align + alloc_align + rss_align + 12  # 12 = 4 * " | "
                                    message = "has been running for $(elapsed_str)"
                                    centered_message = lpad(rpad(message, (data_width + textwidth(message)) ÷ 2), data_width)
                                    printstyled(centered_message, "\n", color=:light_black)
                                end
                            end
                        end
                        before = time()
                        resp, duration = try
                                r = remotecall_fetch(@Base.world(runtests, ∞), wrkr, test, test_path(test); seed=seed)
                                r, time() - before
                            catch e
                                isa(e, InterruptException) && return
                                Any[CapturedException(e, catch_backtrace())], time() - before
                            end
                        delete!(running_tests, test)
                        if haskey(test_timers, test)
                            close(test_timers[test])
                            delete!(test_timers, test)
                        end
                        push!(results, (test, resp, duration))
                        if length(resp) == 1
                            print_testworker_errored(test, wrkr, exit_on_error ? nothing : resp[1])
                            if exit_on_error
                                skipped = length(tests)
                                empty!(tests)
                            elseif n > 1
                                # the worker encountered some failure, recycle it
                                # so future tests get a fresh environment
                                rmprocs(wrkr, waitfor=rmwait_timeout)
                                p = addprocs_with_testenv(1)[1]
                                remotecall_fetch(include, p, "testdefs.jl")
                                if use_revise
                                    Distributed.remotecall_eval(Main, p, revise_init_expr)
                                end
                            end
                        else
                            print_testworker_stats(test, wrkr, resp)
                            if resp[end] > max_worker_rss
                                # the worker has reached the max-rss limit, recycle it
                                # so future tests start with a smaller working set
                                if n > 1
                                    rmprocs(wrkr, waitfor=rmwait_timeout)
                                    p = addprocs_with_testenv(1)[1]
                                    remotecall_fetch(include, p, "testdefs.jl")
                                    if use_revise
                                        Distributed.remotecall_eval(Main, p, revise_init_expr)
                                    end
                                else # single process testing
                                    error("Halting tests. Memory limit reached : $resp > $max_worker_rss")
                                end
                            end
                       end
                    end
                    if p != 1
                        # Free up memory =)
                        rmprocs(p, waitfor=rmwait_timeout)
                    end
                end
            end
        end

        n > 1 && length(node1_tests) > 1 && print("\nExecuting tests that run on node 1 only:\n")
        for t in node1_tests
            # As above, try to run each test
            # which must run on node 1. If
            # the test fails, catch the error,
            # and either way, append the results
            # to the overall aggregator
            isolate = true
            t == "SharedArrays" && (isolate = false)
            before = time()
            resp, duration = try
                    r = @invokelatest runtests(t, test_path(t), isolate, seed=seed) # runtests is defined by the include above
                    r, time() - before
                catch e
                    isa(e, InterruptException) && rethrow()
                    Any[CapturedException(e, catch_backtrace())], time() - before
                end
            if length(resp) == 1
                print_testworker_errored(t, 1, resp[1])
            else
                print_testworker_stats(t, 1, resp)
            end
            push!(results, (t, resp, duration))
        end
    catch e
        isa(e, InterruptException) || rethrow()
        # If the test suite was merely interrupted, still print the
        # summary, which can be useful to diagnose what's going on
        foreach(task -> begin
                istaskstarted(task) || return
                istaskdone(task) && return
                try
                    schedule(task, InterruptException(); error=true)
                catch ex
                    @error "InterruptException" exception=ex,catch_backtrace()
                end
            end, all_tasks)
        foreach(wait, all_tasks)
    finally
        if @isdefined stdin_monitor
            istaskdone(stdin_monitor) || schedule(stdin_monitor, InterruptException(); error=true)
        end
        if @isdefined test_timers
            foreach(close, values(test_timers))
        end
    end

    #=
    Construct a testset on the master node which will hold results from all the
    test files run on workers and on node1. The loop goes through the results,
    inserting them as children of the overall testset if they are testsets,
    handling errors otherwise.

    Since the workers don't return information about passing/broken tests, only
    errors or failures, those Result types get passed `nothing` for their test
    expressions (and expected/received result in the case of Broken).

    If a test failed, returning a `RemoteException`, the error is displayed and
    the overall testset has a child testset inserted, with the (empty) Passes
    and Brokens from the worker and the full information about all errors and
    failures encountered running the tests. This information will be displayed
    as a summary at the end of the test run.

    If a test failed, returning an `Exception` that is not a `RemoteException`,
    it is likely the julia process running the test has encountered some kind
    of internal error, such as a segfault.  The entire testset is marked as
    Errored, and execution continues until the summary at the end of the test
    run, where the test file is printed out as the "failed expression".
    =#
    @with Test.TESTSET_PRINT_ENABLE=>false begin
        o_ts = Test.DefaultTestSet("Overall")
        @atomic o_ts.time_end = o_ts.time_start + o_ts_duration # manually populate the timing
        BuildkiteTestJSON.write_testset_json_files(@__DIR__, o_ts)
        Test.@with_testset o_ts begin
            completed_tests = Set{String}()
            for (testname, (resp,), duration) in results
                push!(completed_tests, testname)
                if isa(resp, Test.DefaultTestSet)
                    @atomic resp.time_end = resp.time_start + duration
                    Test.@with_testset resp begin
                        Test.record(o_ts, resp)
                    end
                elseif isa(resp, Test.TestSetException)
                    fake = Test.DefaultTestSet(testname)
                    @atomic fake.time_end = fake.time_start + duration
                    for i in 1:resp.pass
                        Test.record(fake, Test.Pass(:test, nothing, nothing, nothing, LineNumberNode(@__LINE__, @__FILE__)))
                    end
                    for i in 1:resp.broken
                        Test.record(fake, Test.Broken(:test, nothing))
                    end
                    for t in resp.errors_and_fails
                        Test.record(fake, t)
                    end
                    Test.@with_testset fake begin
                        Test.record(o_ts, fake)
                    end
                else
                    if !isa(resp, Exception)
                        resp = ErrorException(string("Unknown result type : ", typeof(resp)))
                    end
                    # If this test raised an exception that is not a remote testset exception,
                    # i.e. not a RemoteException capturing a TestSetException that means
                    # the test runner itself had some problem, so we may have hit a segfault,
                    # deserialization errors or something similar.  Record this testset as Errored.
                    fake = Test.DefaultTestSet(testname)
                    @atomic fake.time_end = fake.time_start + duration
                    Test.record(fake, Test.Error(:nontest_error, testname, nothing, Base.ExceptionStack(NamedTuple[(;exception = resp, backtrace = [])]), LineNumberNode(1), nothing))
                    Test.@with_testset fake begin
                        Test.record(o_ts, fake)
                    end
                end
            end
            for test in all_tests
                (test in completed_tests) && continue
                fake = Test.DefaultTestSet(test)
                Test.record(fake, Test.Error(:test_interrupted, test, nothing, Base.ExceptionStack(NamedTuple[(;exception = "skipped", backtrace = [])]), LineNumberNode(1), nothing))
                Test.@with_testset fake begin
                    Test.record(o_ts, fake)
                end
            end
        end
    end

    println()
    # o_ts.verbose = true # set to true to show all timings when successful
    Test.print_test_results(o_ts, 1)
    if !Test.anynonpass(o_ts)
        printstyled("    SUCCESS\n"; bold=true, color=:green)
    else
        printstyled("    FAILURE\n\n"; bold=true, color=:red)
        skipped > 0 &&
            println("$skipped test", skipped > 1 ? "s were" : " was", " skipped due to failure.")
        println("The global RNG seed was 0x$(string(seed, base = 16)).\n")
        Test.print_test_errors(o_ts)
        throw(Test.FallbackTestSetException("Test run finished with errors"))
    end
end
