precompilation: sample precompile workers with the runtime profiler

Most of a package's precompilation happens in a worker process that the driver
spawns, so profiling the session that called `Pkg.precompile()` shows the driver
waiting on its children and nothing about where the time went. Give the workers a
way to sample themselves.

`precompilepkgs` gains a `profile` keyword, taking `true` for every package or a
name or a list of names, and `profile_dir` for where the dumps land. The same is
available through `JULIA_PRECOMPILE_PROFILE` for precompilation started some
other way, such as by loading a package for the first time.

A selected worker starts the runtime's sampling profiler before it loads its
dependencies and dumps once its image has been written, so the profile covers
dependency loading, lowering, inference, code generation and image generation. It
writes the raw sample buffer alongside the stack frames its instruction pointers
resolve to, since a pointer can only be resolved by the process that recorded it.
`contrib/read_precompile_profile.jl` turns the pair back into the arguments that
`Profile.print` and the profile viewer packages take.

The worker side uses only `ccall` and `Base.StackTraces` on purpose: loading
Profile there would put it in the output image's dependency list.

Assisted-by: Claude Code (Fable 5.1)
Claude-Session: https://claude.ai/code/session_019BZkT3Ky7B3kq9xPiXLYVU

Author: IanButterworth

base/loading.jl

@@ -3419,9 +3419,91 @@ end
 const newly_inferred = []
 
 # this is called in the external process that generates precompiled package files
+# Sampling profiler for a precompile worker, enabled by pointing
+# `JULIA_PRECOMPILE_PROFILE` at a directory. Every worker then writes
+# `<pkg>-<pid>.profdata`, the raw sample buffer, and `<pkg>-<pid>.profsyms`, its
+# instruction pointers resolved to stack frames. The dump is deliberately split
+# that way because an instruction pointer can only be resolved by the process
+# that recorded it, while everything else is better done afterwards in a session
+# that can load Profile. See the "Profiling package precompilation" devdocs.
+#
+# This runs inside the worker, so it only uses `ccall` and `Base.StackTraces`:
+# loading a package here would place it in the output image's dependency list.
+const _precompile_profile_nmeta = 4 # threadid, taskid, cpu_cycle_clock, sleepstate
+
+function _precompile_profile_start(pkg::PkgId)
+    dir = get(ENV, "JULIA_PRECOMPILE_PROFILE", "")
+    isempty(dir) && return nothing
+    delay = something(tryparse(Float64, get(ENV, "JULIA_PRECOMPILE_PROFILE_DELAY", "")), 0.001)
+    nsamples = something(tryparse(Int, get(ENV, "JULIA_PRECOMPILE_PROFILE_NSAMPLES", "")), 10_000_000)
+    status = ccall(:jl_profile_init, Cint, (Csize_t, UInt64), nsamples, round(UInt64, 1e9 * delay))
+    if status != 0
+        @warn "Could not allocate the precompile profile buffer" pkg status
+        return nothing
+    end
+    status = ccall(:jl_profile_start_timer, Cint, (Bool,), false)
+    if status != 0
+        @warn "Could not start the precompile profiler" pkg status
+        return nothing
+    end
+    prefix = joinpath(abspath(dir), string(pkg.name, "-", getpid()))
+    # `_postoutput` runs after the package image has been written, so dumping
+    # there covers image generation as well as `include`. It only runs when
+    # native code is emitted; without it, fall back to `atexit`, which runs
+    # before the `.ji` is written and so covers `include` alone.
+    if JLOptions().outputo != C_NULL
+        postoutput(() -> _precompile_profile_dump(prefix))
+    else
+        atexit(() -> _precompile_profile_dump(prefix))
+    end
+    return nothing
+end
+
+function _precompile_profile_dump(prefix::String)
+    ccall(:jl_profile_stop_timer, Cvoid, ())
+    len = Int(ccall(:jl_profile_len_data, Csize_t, ()))
+    ptr = convert(Ptr{UInt}, ccall(:jl_profile_get_data, Ptr{UInt8}, ()))
+    data = Vector{UInt}(undef, len)
+    unsafe_copyto!(pointer(data), ptr, len)
+    try
+        mkpath(dirname(prefix))
+        open(prefix * ".profdata", "w") do io
+            write(io, data)
+        end
+        # Mark the metadata trailer of each block so it is not mistaken for an
+        # instruction pointer: a block ends with the four metadata fields
+        # followed by two null entries.
+        meta = falses(len)
+        for i in (_precompile_profile_nmeta + 2):len
+            if data[i] == 0 && data[i-1] == 0 && data[i-2] != 0
+                meta[(i-_precompile_profile_nmeta-1):i] .= true
+            end
+        end
+        ips = Set{UInt}()
+        for i in 1:len
+            (meta[i] || data[i] == 0) && continue
+            push!(ips, data[i])
+        end
+        open(prefix * ".profsyms", "w") do io
+            println(io, "# ip\tinlined\tline\tfrom_c\tfunc\tfile")
+            for ip in ips, sf in StackTraces.lookup(convert(Ptr{Cvoid}, ip))
+                println(io, ip, '\t', sf.inlined ? 1 : 0, '\t', sf.line, '\t',
+                        sf.from_c ? 1 : 0, '\t', sf.func, '\t', sf.file)
+            end
+        end
+        if ccall(:jl_profile_is_buffer_full, Cint, ()) != 0
+            @warn "The precompile profile buffer filled up; raise JULIA_PRECOMPILE_PROFILE_NSAMPLES or JULIA_PRECOMPILE_PROFILE_DELAY" prefix
+        end
+    catch ex
+        @warn "Could not write the precompile profile" prefix exception=(ex, catch_backtrace())
+    end
+    return nothing
+end
+
 function include_package_for_output(pkg::PkgId, input::String, syntax_version::VersionNumber, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String},
                                     concrete_deps::typeof(_concrete_dependencies), source::Union{Nothing,String},
                                     preresolved::Vector{Pair{PkgId,String}}=Pair{PkgId,String}[])
+    _precompile_profile_start(pkg)
 
     @lock require_lock begin
     m = start_loading(pkg, UInt128(0), false)
@@ -3517,7 +3599,7 @@ const PRECOMPILE_VERBOSE_TIMING_MARKER = "__JL_PRECOMP_VERBOSE_TIMING__"
 function create_expr_cache(pkg::PkgId, input::PkgLoadSpec, output::String, output_o::Union{Nothing, String},
                            concrete_deps::typeof(_concrete_dependencies), flags::Cmd=``, cacheflags::CacheFlags=CacheFlags(),
                            internal_stderr::IO = stderr, internal_stdout::IO = stdout, loadable_exts::Union{Vector{PkgId},Nothing}=nothing;
-                           report_timing::Bool=false,
+                           report_timing::Bool=false, profile_dir::Union{Nothing,String}=nothing,
                            preresolved::Vector{Pair{PkgId,String}} = @lock(require_lock, collect(preresolved_cachefiles)))
     @nospecialize internal_stderr internal_stdout
     depot_path = String[abspath(x) for x in DEPOT_PATH]
@@ -3574,6 +3656,8 @@ function create_expr_cache(pkg::PkgId, input::PkgLoadSpec, output::String, outpu
     # Only request per-package timing reports when explicitly asked for (e.g. by
     # precompilepkgs), so that the marker lines don't leak into normal load logs.
     report_timing && (cmd = addenv(cmd, "JULIA_PRECOMP_REPORT_TIMING" => 1))
+    # Ask this worker to profile itself; see `_precompile_profile_start`.
+    profile_dir === nothing || (cmd = addenv(cmd, "JULIA_PRECOMPILE_PROFILE" => profile_dir))
     io = open(pipeline(cmd, stderr = internal_stderr, stdout = internal_stdout),
               "w", stdout)
     # write data over stdin to avoid the (unlikely) case of exceeding max command line size
@@ -3637,11 +3721,11 @@ This can be used to reduce package load times. Cache files are stored in
 `DEPOT_PATH[1]/compiled`. See [Module initialization and precompilation](@ref)
 for important notes.
 """
-function compilecache(pkg::PkgId, internal_stderr::IO = stderr, internal_stdout::IO = stdout; flags::Cmd=``, cacheflags::CacheFlags=CacheFlags(), loadable_exts::Union{Vector{PkgId},Nothing}=nothing, signal_channel::Union{Channel{Int32},Nothing}=nothing, report_timing::Bool=false)
+function compilecache(pkg::PkgId, internal_stderr::IO = stderr, internal_stdout::IO = stdout; flags::Cmd=``, cacheflags::CacheFlags=CacheFlags(), loadable_exts::Union{Vector{PkgId},Nothing}=nothing, signal_channel::Union{Channel{Int32},Nothing}=nothing, report_timing::Bool=false, profile_dir::Union{Nothing,String}=nothing)
     @nospecialize internal_stderr internal_stdout
     spec = locate_package_load_spec(pkg)
     spec === nothing && throw(ArgumentError("$(repr("text/plain", pkg)) not found during precompilation"))
-    return compilecache(pkg, spec, internal_stderr, internal_stdout; flags, cacheflags, loadable_exts, signal_channel, report_timing)
+    return compilecache(pkg, spec, internal_stderr, internal_stdout; flags, cacheflags, loadable_exts, signal_channel, report_timing, profile_dir)
 end
 
 const MAX_NUM_PRECOMPILE_FILES = Ref(10)
@@ -3650,6 +3734,7 @@ function compilecache(pkg::PkgId, spec::PkgLoadSpec, internal_stderr::IO = stder
                       keep_loaded_modules::Bool = true; flags::Cmd=``, cacheflags::CacheFlags=CacheFlags(),
                       loadable_exts::Union{Vector{PkgId},Nothing}=nothing, signal_channel::Union{Channel{Int32},Nothing}=nothing,
                       pid_channel::Union{Channel{Int32},Nothing}=nothing, report_timing::Bool=false,
+                      profile_dir::Union{Nothing,String}=nothing,
                       preresolved::Vector{Pair{PkgId,String}} = @lock(require_lock, collect(preresolved_cachefiles)))
 
     @nospecialize internal_stderr internal_stdout
@@ -3688,7 +3773,7 @@ function compilecache(pkg::PkgId, spec::PkgLoadSpec, internal_stderr::IO = stder
             close(tmpio_o)
             close(tmpio_so)
         end
-        p = create_expr_cache(pkg, spec, tmppath, tmppath_o, concrete_deps, flags, cacheflags, internal_stderr, internal_stdout, loadable_exts; report_timing, preresolved)
+        p = create_expr_cache(pkg, spec, tmppath, tmppath_o, concrete_deps, flags, cacheflags, internal_stderr, internal_stdout, loadable_exts; report_timing, profile_dir, preresolved)
 
         # Report the PID of the compilation subprocess
         if pid_channel !== nothing

base/precompilation.jl

@@ -279,12 +279,14 @@ mutable struct BackgroundPrecompileState
     confirm_deadline::Float64  # time() deadline for confirmation
     info_requested::Bool  # whether SIGINFO/SIGUSR1 has been broadcast at least once
     key_listening::Bool  # whether a key listener task is currently consuming stdin
+    profile_dir::Union{Nothing, String}  # directory for worker profile dumps, or nothing to not profile
+    profile_pkgs::Union{Nothing, Set{String}}  # packages to profile, or nothing for all of them
 end
 Base.lock(f, bg::BackgroundPrecompileState) = lock(f, bg.lock)
 Base.lock(bg::BackgroundPrecompileState) = lock(bg.lock)
 Base.unlock(bg::BackgroundPrecompileState) = unlock(bg.lock)
 
-const BG = BackgroundPrecompileState(nothing, false, false, false, nothing, nothing, nothing, nothing, ReentrantLock(), Threads.Condition(), Channel{Int32}[], Dict{PkgId, Int}(), Set{PkgId}(), Threads.Condition(), Channel{PrecompileRequest}(Inf), false, false, :none, 0.0, false, false)
+const BG = BackgroundPrecompileState(nothing, false, false, false, nothing, nothing, nothing, nothing, ReentrantLock(), Threads.Condition(), Channel{Int32}[], Dict{PkgId, Int}(), Set{PkgId}(), Threads.Condition(), Channel{PrecompileRequest}(Inf), false, false, :none, 0.0, false, false, nothing, nothing)
 
 # Serializes the inject-vs-launch decision in `_precompilepkgs` with the launch
 # itself. Lock ordering: acquired before (outside) BG.lock, never while holding it.
@@ -1105,6 +1107,14 @@ precompiles only the given packages and their dependencies (unless
                   samples may be missed. Linux/macOS only, `-` elsewhere.
   Values under 5 ms and zero counts are dimmed for readability.
 
+- `profile::Union{Bool,AbstractString,AbstractVector{<:AbstractString}}`: When not `false`,
+  each selected package's worker samples itself and writes its profile to a file. Pass `true`
+  for every package, or a package name or list of names to profile only those. The dumps are
+  read back in a separate session; see the "Profiling package precompilation" devdocs.
+
+- `profile_dir::Union{Nothing,AbstractString}`: Where to write those dumps. Defaults to a
+  fresh temporary directory, whose path is printed when profiling starts.
+
 - `_from_loading::Bool`: Internal flag indicating the call originated from the
   package loading system. When `true` (not default): returns early instead of
   throwing when packages are not found; suppresses progress messages when not
@@ -1175,6 +1185,48 @@ function preresolved_snapshot(s::PrecompileSession)
     @lock s.cache_lock Pair{Base.PkgId,String}[k => first(v) for (k, v) in s.cachepath_cache if !isempty(v)]
 end
 
+# Turn the `profile` / `profile_dir` keywords of `precompilepkgs` into the
+# settings that `spawn_precompile_tasks!` hands to each worker. Kept on `BG`
+# rather than on the session so a request merged into a running background run
+# picks it up, the same way `verbose` does.
+function setup_worker_profiling!(profile, profile_dir, io)
+    if profile === false
+        if profile_dir !== nothing
+            throw(ArgumentError("`profile_dir` was given but `profile` is false; pass `profile=true` to enable profiling"))
+        end
+        # clear any setting left by an earlier profiling run
+        @lock BG begin
+            BG.profile_dir = nothing
+            BG.profile_pkgs = nothing
+        end
+        return nothing
+    end
+    pkgs = if profile === true
+        nothing
+    elseif profile isa AbstractString
+        Set{String}((profile,))
+    else
+        Set{String}(profile)
+    end
+    dir = profile_dir === nothing ? mktempdir(; prefix="jl_precompile_profile_", cleanup=false) : abspath(String(profile_dir))
+    mkpath(dir)
+    @lock BG begin
+        BG.profile_dir = dir
+        BG.profile_pkgs = pkgs
+    end
+    which = pkgs === nothing ? "all packages" : join(sort!(collect(pkgs)), ", ")
+    printpkgstyle(io, :Profiling, "$which into $dir", color = Base.info_color())
+    return nothing
+end
+
+# The directory this package's worker should dump its profile into, or `nothing`.
+function worker_profile_dir(pkg::PkgId)
+    dir, pkgs = @lock BG (BG.profile_dir, BG.profile_pkgs)
+    dir === nothing && return nothing
+    (pkgs === nothing || pkg.name in pkgs) || return nothing
+    return dir
+end
+
 function precompilepkgs(pkgs::Union{Vector{String}, Vector{PkgId}}=String[];
                         internal_call::Bool=false,
                         strict::Bool = false,
@@ -1189,10 +1241,13 @@ function precompilepkgs(pkgs::Union{Vector{String}, Vector{PkgId}}=String[];
                         fancyprint::Bool = can_fancyprint(io) && !timing && !verbose,
                         manifest::Bool=false,
                         ignore_loaded::Bool=true,
-                        detachable::Bool=false)
+                        detachable::Bool=false,
+                        profile::Union{Bool,AbstractString,AbstractVector{<:AbstractString}}=false,
+                        profile_dir::Union{Nothing,AbstractString}=nothing)
     # verbose timing mode requires timing to be enabled (per-package breakdown
     # is only shown alongside timing lines in non-fancy mode)
     verbose && (timing = true)
+    setup_worker_profiling!(profile, profile_dir, io)
     @debug "precompilepkgs called with" pkgs internal_call strict warn_loaded timing verbose _from_loading configs fancyprint manifest ignore_loaded detachable
     # monomorphize this to avoid latency problems
     _precompilepkgs(pkgs, internal_call, strict, warn_loaded, timing, verbose, _from_loading,
@@ -2255,6 +2310,7 @@ function spawn_precompile_tasks!(s::PrecompileSession;
                                 Base.compilecache(pkg, sourcespec, std_pipe, std_pipe, !s.ignore_loaded;
                                                   flags=flags_, cacheflags, loadable_exts, signal_channel=make_signal_channel(),
                                                   pid_channel=pid_ch, report_timing=true,
+                                                  profile_dir=worker_profile_dir(pkg),
                                                   preresolved=preresolved_snapshot(s))
                             end
                         else
@@ -2281,6 +2337,7 @@ function spawn_precompile_tasks!(s::PrecompileSession;
                                 Base.compilecache(pkg, sourcespec, std_pipe, std_pipe, !s.ignore_loaded;
                                                   flags=flags_, cacheflags, loadable_exts, signal_channel=make_signal_channel(),
                                                   pid_channel=pid_ch, report_timing=true,
+                                                  profile_dir=worker_profile_dir(pkg),
                                                   preresolved=preresolved_snapshot(s))
                             end
                         end

contrib/read_precompile_profile.jl

@@ -0,0 +1,54 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+# Read a profile written by a precompilation worker, as enabled by
+# `precompilepkgs(profile=...)` or the `JULIA_PRECOMPILE_PROFILE` environment
+# variable, and hand it to `Profile` for analysis.
+#
+#     include("contrib/read_precompile_profile.jl")
+#     data, lidict = read_precompile_profile("/tmp/jl_precompile_profile_XXXX/Makie-12345")
+#     Profile.print(; C=true)  # see the docstring below for the exact call
+#
+# The worker writes two files because instruction pointers can only be resolved
+# in the process that recorded them: `<prefix>.profdata` holds the raw sample
+# buffer and `<prefix>.profsyms` the stack frames those pointers resolve to.
+
+using Profile
+using Base.StackTraces: StackFrame
+
+"""
+    read_precompile_profile(prefix) -> (data, lidict)
+
+Load the precompile-worker profile written to `\$prefix.profdata` and
+`\$prefix.profsyms`. `prefix` may also be either of those paths.
+
+The result is accepted directly by the `Profile` reporting functions:
+
+```julia
+data, lidict = read_precompile_profile(prefix)
+Profile.print(stdout, data, lidict; C=true, format=:tree)
+```
+"""
+function read_precompile_profile(prefix::AbstractString)
+    prefix = replace(String(prefix), r"\.(profdata|profsyms)$" => "")
+    data = collect(reinterpret(UInt, read(prefix * ".profdata")))
+    lidict = Dict{UInt64,Vector{StackFrame}}()
+    for line in eachline(prefix * ".profsyms")
+        (isempty(line) || startswith(line, '#')) && continue
+        ip_str, inlined, lineno, from_c, func, file = split(line, '\t'; limit=6)
+        ip = parse(UInt64, ip_str)
+        frame = StackFrame(Symbol(func), Symbol(file), parse(Int, lineno), nothing,
+                           from_c == "1", inlined == "1", ip)
+        push!(get!(Vector{StackFrame}, lidict, ip), frame)
+    end
+    return data, lidict
+end
+
+"""
+    print_precompile_profile(prefix; kwargs...)
+
+Read the profile at `prefix` and print it, passing `kwargs` on to `Profile.print`.
+"""
+function print_precompile_profile(prefix::AbstractString; io::IO=stdout, kwargs...)
+    data, lidict = read_precompile_profile(prefix)
+    Profile.print(io, data, lidict; kwargs...)
+end

doc/make.jl

@@ -361,6 +361,7 @@ DevDocs = [
         "devdocs/valgrind.md",
         "devdocs/gc-debug.md",
         "devdocs/external_profilers.md",
+        "devdocs/precompile_profiling.md",
         "devdocs/sanitizers.md",
         "devdocs/probes.md",
     ],

doc/src/devdocs/external_profilers.md

@@ -89,6 +89,8 @@ end
 
 Here, we use a custom port for tracy which makes it easier to find the correct client in the Tracy UI to connect to.
 
+For sampled stack traces of the same worker processes, rather than the zones described here, see [Profiling package precompilation](@ref).
+
 ### Adding metadata to zones
 
 The various `jl_timing_show_*` and `jl_timing_printf` functions can be used to attach a string (or strings) to a zone. For example, the trace zone for inference shows the method instance that is being inferred.