ParallelPrecompilation.jl

module ParallelPrecompilation

using Base: PkgId
using Pkg.Types: is_stdlib
using Pkg: TOML
using ProgressLogging: @logprogress, @withprogress
using UUIDs: UUID

function manifest_path(project)
    manifest_names = ("JuliaManifest.toml", "Manifest.toml")
    if basename(project) in manifest_names
        return project
    else
        if isfile(project)
            project = dirname(project)
        end
        candidates = joinpath.(project, manifest_names)
        i = findfirst(isfile, candidates)
        i === nothing || return candidates[i]
        error("Manifest file does not exist at: ", project)
    end
end

load_manifest(project = Base.active_project()) = TOML.parsefile(manifest_path(project))

function foreach_async(f, c; ntasks = Sys.CPU_THREADS)
    @sync for _ in 1:ntasks
        @async foreach(f, c)
    end
end

precompile(; kwargs...) =
    precompile(keys(TOML.parsefile(Base.active_project())["deps"]); kwargs...)

function precompile(packages; ntasks = Sys.CPU_THREADS)
    packages = [Base.identify_package(p) for p in packages]
    manifest = load_manifest()

    function depsof(pkg::PkgId)
        for spec in manifest[pkg.name]
            if UUID(spec["uuid"]) == pkg.uuid
                deps = (Base.identify_package(pkg, d) for d in get(spec, "deps", Union{}[]))
                deps = Iterators.filter(!isnothing, deps)
                return collect(deps)
            end
        end
        return PkgId[]
    end

    function collectdeps!(deps::Set{PkgId}, pkg::PkgId)
        pkg in deps && return deps
        foldl(collectdeps!, depsof(pkg); init = deps)
        push!(deps, pkg)
        return deps
    end

    total = let deps = Set{PkgId}()
        for pkg in packages
            collectdeps!(deps, pkg)
        end
        length(deps)
    end

    workqueue = Channel{Task}(ntasks) do workqueue
        foreach_async(wait ∘ schedule, workqueue; ntasks = ntasks)
    end

    request = Channel{Tuple{PkgId,Channel{Task}}}() do request
        counter = Ref(0)
        loaded = Dict{UUID,Task}()
        @withprogress name = "Precompile" for (pkg, response) in request
            task = get!(loaded, pkg.uuid) do
                @async try
                    if is_stdlib(pkg.uuid)
                        counter[] += 1
                        return
                    end
                    @sync for dep in depsof(pkg)
                        @debug "$(pkg.name) => $(dep.name)"
                        Base.@sync_add compiled(dep)
                    end

                    work = @task try
                        @info "Compiling $pkg..."
                        seconds = @elapsed was_stale = try
                            compilepkg(pkg)
                        catch err
                            err isa ProcessFailedException || rethrow()
                            @error(
                                "Failed to compile $pkg",
                                exception = (err, catch_backtrace())
                            )
                            return
                        end
                        if was_stale
                            @info "Compiling $pkg... (took $seconds seconds)"
                        else
                            @info "Compiling $pkg... (cache is fresh)"
                        end
                        c = counter[] += 1
                        @logprogress c / total
                        return
                    catch err
                        @error(
                            "Compilation (inner) task failed",
                            pkg,
                            exception = (err, catch_backtrace())
                        )
                        rethrow()
                    end
                    put!(workqueue, work)
                    wait(work)
                catch err
                    @error(
                        "Compilation wrapper task failed",
                        pkg,
                        exception = (err, catch_backtrace())
                    )
                    rethrow()
                end
            end
            put!(response, task)
        end
    end
    function compiled(pkg::PkgId)
        @debug "Start: compiled($pkg)"
        response = Channel{Task}()
        try
            put!(request, (pkg, response))
            task = take!(response)
            @debug "OK: compiled($pkg)"
            return task
        catch err
            @error "Failed: compiled($pkg)" exception = (err, catch_backtrace())
            rethrow()
        finally
            close(response)
        end
    end

    try
        @sync for pkg in packages
            Base.@sync_add compiled(pkg)
        end
    finally
        close(request)
        close(workqueue)
    end
end

#=
# Taken from `Pkg.precompile`
function needcompile(pkg::PkgId)
    paths = Base.find_all_in_cache_path(pkg)
    sourcepath = Base.locate_package(pkg)
    sourcepath === nothing && return false
    # Heuristic for when precompilation is disabled
    occursin(r"\b__precompile__\(\s*false\s*\)", read(sourcepath, String)) && return false
    for path_to_try in paths::Vector{String}
        staledeps = Base.stale_cachefile(sourcepath, path_to_try)
        staledeps !== true && return true
        # This is probably not the right thing to do. (But probably
        # there is no right thing one can do...)
    end
    return false
end
=#

function compilepkg(pkg::PkgId)
    # needcompile(pkg) || return false
    uuid = UInt128(pkg.uuid)
    code = """
    $(Base.load_path_setup_code())
    Base.require(Base.PkgId(Base.UUID($(repr(uuid))), $(repr(pkg.name))))
    """
    #=
    Base.compilecache(Base.PkgId(Base.UUID($(repr(uuid))), $(repr(pkg.name))))
    =#
    cmd = `$(Base.julia_cmd()) --startup-file=no -e $code`
    @debug "Running: $(Base.julia_cmd())" code = Text(code)
    run(cmd)
    return true
end

end # module