pass the environment stack to precompile workers via serialization

Author: KristofferC

base/codeloading2.jl

@@ -305,7 +305,19 @@ struct EnvironmentStack
     envs::Vector{Union{ImplicitEnv, ExplicitEnv}}
 end
 
-function EnvironmentStack(environments = load_path())
+
+# Caching
+
+const CACHED_ENV_STACK = Ref{Union{EnvironmentStack, Nothing}}(nothing)
+
+function EnvironmentStack(environments = nothing)
+    if CACHED_ENV_STACK[] !== nothing
+        return CACHED_ENV_STACK[]
+    end
+    # Avoid looking up `load_path` until it is really needed.
+    if environments === nothing
+        environments = load_path()
+    end
     envs = Union{ImplicitEnv, ExplicitEnv}[]
     for env in environments
         if isfile(env)

base/loading.jl

@@ -159,7 +159,6 @@ end
 const ns_dummy_uuid = UUID("fe0723d6-3a44-4c41-8065-ee0f42c8ceab")
 
 function dummy_uuid(project_file::String)
-    @lock require_lock begin
     project_path = try
         realpath(project_file)
     catch ex
@@ -168,7 +167,6 @@ function dummy_uuid(project_file::String)
     end
     uuid = uuid5(ns_dummy_uuid, project_path)
     return uuid
-    end
 end
 
 ## package path slugs: turning UUID + SHA1 into a pair of 4-byte "slugs" ##
@@ -789,6 +787,9 @@ For more details regarding code loading, see the manual sections on [modules](@r
 """
 function require(into::Module, mod::Symbol)
     @lock require_lock begin
+    no_initial_cache = CACHED_ENV_STACK[] === nothing
+    try
+    no_initial_cache && (CACHED_ENV_STACK[] = EnvironmentStack())
     uuidkey = identify_package(into, String(mod))
     # Core.println("require($(PkgId(into)), $mod) -> $uuidkey from env \"$env\"")
     if uuidkey === nothing
@@ -824,6 +825,9 @@ function require(into::Module, mod::Symbol)
         push!(_require_dependencies, (into, binpack(uuidkey), 0.0))
     end
     return _require_prelocked(uuidkey)
+    finally
+        no_initial_cache && (CACHED_ENV_STACK[] = nothing)
+    end # try
     end # @lock
 end
 
@@ -1209,6 +1213,7 @@ function include_package_for_output(pkg::PkgId, input::String, depot_path::Vecto
 end
 
 const PRECOMPILE_TRACE_COMPILE = Ref{String}()
+const ENV_SERIALIZATION_FILE = Ref{Union{Nothing,String}}(nothing)
 function create_expr_cache(pkg::PkgId, input::String, output::String, concrete_deps::typeof(_concrete_dependencies), internal_stderr::IO = stderr, internal_stdout::IO = stdout)
     @nospecialize internal_stderr internal_stdout
     rm(output, force=true)   # Remove file if it exists
@@ -1241,6 +1246,22 @@ function create_expr_cache(pkg::PkgId, input::String, output::String, concrete_d
                        -`, stderr = internal_stderr, stdout = internal_stdout),
               "w", stdout)
     # write data over stdin to avoid the (unlikely) case of exceeding max command line size
+    serialization_pkgid = PkgId(UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b"), "Serialization")
+    if haskey(loaded_modules, serialization_pkgid)
+        local ftmp
+        if ENV_SERIALIZATION_FILE[] === nothing || !isfile(ENV_SERIALIZATION_FILE[])
+            Serialization = loaded_modules[serialization_pkgid]
+            ftmp, iotmp = mktemp()
+            Serialization.serialize(iotmp, EnvironmentStack())
+            close(iotmp)
+        else
+            ftmp = ENV_SERIALIZATION_FILE[]
+        end
+        write(io.in, """
+            Base.ENV_SERIALIZATION_FILE[] = $(repr(ftmp))
+            Base.CACHED_ENV_STACK[] = (Base.loaded_modules[Base.PkgId(Base.UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b"), "Serialization")].deserialize(IOBuffer(read($(repr(ftmp))))))
+        """)
+    end
     write(io.in, """
         Base.include_package_for_output($(pkg_str(pkg)), $(repr(abspath(input))), $(repr(depot_path)), $(repr(dl_load_path)),
             $(repr(load_path)), $deps, $(repr(source_path(nothing))))

contrib/generate_precompile.jl

@@ -186,6 +186,31 @@ if Libdl !== nothing
     """
 end
 
+Serialization = get(Base.loaded_modules,
+          Base.PkgId(Base.UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b"), "Serialization"),
+          nothing)
+if Serialization !== nothing
+    # Deserializing environment
+    hardcoded_precompile_statements *= """
+        precompile(Tuple{typeof(Serialization.deserialize), String})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.EnvironmentStack}})
+        precompile(Tuple{typeof(Serialization.deserialize_fillarray!), Array{String, 1}, Serialization.Serializer{Base.IOStream}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Union}})
+        precompile(Tuple{typeof(Serialization.deserialize_fillarray!), Array{Union{Base.ExplicitEnv, Base.ImplicitEnv}, 1}, Serialization.Serializer{Base.IOStream}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.ExplicitEnv}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.Dict{String, Base.UUID}}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.UUID}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.Dict{Base.UUID, Base.Dict{String, Base.UUID}}}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.Dict{Base.UUID, Union{Base.Missing, Nothing, Base.SHA1, String}}}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.SHA1}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.ImplicitEnv}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.Dict{String, Base.ImplicitEnvPkg}}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.ImplicitEnvPkg}})
+        precompile(Tuple{typeof(Serialization.deserialize_fillarray!), Array{Base.PkgId, 1}, Serialization.Serializer{Base.IOStream}})
+        precompile(Tuple{typeof(Serialization.deserialize), Serialization.Serializer{Base.IOStream}, Type{Base.PkgId}})
+    """
+end
+
 Test = get(Base.loaded_modules,
           Base.PkgId(Base.UUID("8dfed614-e22c-5e08-85e1-65c5234f0b40"), "Test"),
           nothing)

deps/checksums/Pkg-3cbbd860afd4c2a50a80a04fa229fe5cd5bddc76.tar.gz/md5

@@ -0,0 +1 @@
+da299f922804dae87cefd024e9c96866

deps/checksums/Pkg-3cbbd860afd4c2a50a80a04fa229fe5cd5bddc76.tar.gz/sha512

@@ -0,0 +1 @@
+fa9edd15a7e4e754ece0ae7f7a5b3dda744340211e0857e4f59d5050c86fe73ff938cb6afe6588caaa0945dedcd0f0373a92e38fc47c2673dbec02df4ce40251

deps/checksums/Pkg-45003c2e00a50d0594b927782f8342739860e3da.tar.gz/md5

@@ -1,4 +1,4 @@
 PKG_BRANCH = master
-PKG_SHA1 = 3cbbd860afd4c2a50a80a04fa229fe5cd5bddc76
+PKG_SHA1 = 45003c2e00a50d0594b927782f8342739860e3da
 PKG_GIT_URL := https://github.com/JuliaLang/Pkg.jl.git
 PKG_TAR_URL = https://api.github.com/repos/JuliaLang/Pkg.jl/tarball/$1

deps/checksums/Pkg-45003c2e00a50d0594b927782f8342739860e3da.tar.gz/sha512

@@ -1465,13 +1465,13 @@ function deserialize_string(s::AbstractSerializer, len::Int)
 end
 
 # default DataType deserializer
-function deserialize(s::AbstractSerializer, t::DataType)
-    nf = length(t.types)
+function deserialize(s::AbstractSerializer, ::Type{T}) where {T}
+    nf = length(T.types)
     if nf == 0 && t.size > 0
         # bits type
-        return read(s.io, t)
-    elseif ismutabletype(t)
-        x = ccall(:jl_new_struct_uninit, Any, (Any,), t)
+        return read(s.io, T)
+    elseif ismutabletype(T)
+        x = ccall(:jl_new_struct_uninit, Any, (Any,), T)
         deserialize_cycle(s, x)
         for i in 1:nf
             tag = Int32(read(s.io, UInt8)::UInt8)
@@ -1481,7 +1481,7 @@ function deserialize(s::AbstractSerializer, t::DataType)
         end
         return x
     elseif nf == 0
-        return ccall(:jl_new_struct_uninit, Any, (Any,), t)
+        return ccall(:jl_new_struct_uninit, Any, (Any,), T)
     else
         na = nf
         vflds = Vector{Any}(undef, nf)
@@ -1494,7 +1494,7 @@ function deserialize(s::AbstractSerializer, t::DataType)
                 na >= i && (na = i - 1) # rest of tail must be undefined values
             end
         end
-        return ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), t, vflds, na)
+        return ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), T, vflds, na)
     end
 end