only try download artifacts from pkg server if the corresponding package is in a registry from it (#4297)

Co-authored-by: KristofferC <[email protected]>

Author: KristofferC

src/Artifacts.jl

@@ -434,6 +434,7 @@ Ensures an artifact is installed, downloading it via the download information st
 function ensure_artifact_installed(name::String, artifacts_toml::String;
                                    platform::AbstractPlatform = HostPlatform(),
                                    pkg_uuid::Union{Base.UUID,Nothing}=nothing,
+                                   pkg_server_eligible::Bool=true,
                                    verbose::Bool = false,
                                    quiet_download::Bool = false,
                                    progress::Union{Function,Nothing} = nothing,
@@ -444,23 +445,23 @@ function ensure_artifact_installed(name::String, artifacts_toml::String;
     end
 
     return ensure_artifact_installed(name, meta, artifacts_toml;
-                                    platform, verbose, quiet_download, progress, io)
+                                     pkg_server_eligible, platform, verbose, quiet_download, progress, io)
 end
 
 function ensure_artifact_installed(name::String, meta::Dict, artifacts_toml::String;
+                                   pkg_server_eligible::Bool=true,
                                    platform::AbstractPlatform = HostPlatform(),
                                    verbose::Bool = false,
                                    quiet_download::Bool = false,
                                    progress::Union{Function,Nothing} = nothing,
                                    io::IO=stderr_f())
-
     hash = SHA1(meta["git-tree-sha1"])
     if !artifact_exists(hash)
         if isnothing(progress) || verbose == true
-            return try_artifact_download_sources(name, hash, meta, artifacts_toml; platform, verbose, quiet_download, io)
+            return try_artifact_download_sources(name, hash, meta, artifacts_toml; pkg_server_eligible, platform, verbose, quiet_download, io)
         else
             # if a custom progress handler is given it is taken to mean the caller wants to handle the download scheduling
-            return () -> try_artifact_download_sources(name, hash, meta, artifacts_toml; platform, quiet_download=true, io, progress)
+            return () -> try_artifact_download_sources(name, hash, meta, artifacts_toml; pkg_server_eligible, platform, quiet_download=true, io, progress)
         end
     else
         return artifact_path(hash)
@@ -469,16 +470,16 @@ end
 
 function try_artifact_download_sources(
             name::String, hash::SHA1, meta::Dict, artifacts_toml::String;
+            pkg_server_eligible::Bool=true,
             platform::AbstractPlatform=HostPlatform(),
             verbose::Bool=false,
             quiet_download::Bool=false,
             io::IO=stderr_f(),
             progress::Union{Function,Nothing}=nothing)
 
     errors = Any[]
-    # first try downloading from Pkg server
-    # TODO: only do this if Pkg server knows about this package
-    if (server = pkg_server()) !== nothing
+    # first try downloading from Pkg server if the Pkg server knows about this package
+    if pkg_server_eligible && (server = pkg_server()) !== nothing
         url = "$server/artifact/$hash"
         download_success = let url = url
             @debug "Downloading artifact from Pkg server" name artifacts_toml platform url

src/Operations.jl

@@ -936,15 +936,23 @@ function download_artifacts(ctx::Context;
     env = ctx.env
     io = ctx.io
     fancyprint = can_fancyprint(io)
-    pkg_roots = String[]
+    pkg_info = Tuple{String, Union{Base.UUID, Nothing}}[]
     for (uuid, pkg) in env.manifest
         pkg = manifest_info(env.manifest, uuid)
         pkg_root = source_path(env.manifest_file, pkg, julia_version)
-        pkg_root === nothing || push!(pkg_roots, pkg_root)
+        pkg_root === nothing || push!(pkg_info, (pkg_root, uuid))
     end
-    push!(pkg_roots, dirname(env.project_file))
+    push!(pkg_info, (dirname(env.project_file), env.pkg !== nothing ? env.pkg.uuid : nothing))
     download_jobs = Dict{SHA1, Function}()
 
+    # Check what registries the current pkg server tracks
+    # Disable if precompiling to not access internet
+    server_registry_info = if Base.JLOptions().incremental == 0
+        Registry.pkg_server_registry_info()
+    else
+        nothing
+    end
+
     print_lock = Base.ReentrantLock() # for non-fancyprint printing
 
     download_states = Dict{SHA1, DownloadState}()
@@ -958,12 +966,17 @@ function download_artifacts(ctx::Context;
     ansi_enablecursor = "\e[?25h"
     ansi_disablecursor = "\e[?25l"
 
-    all_collected_artifacts = reduce(vcat, map(pkg_root -> collect_artifacts(pkg_root; platform, include_lazy), pkg_roots))
-    used_artifact_tomls = Set{String}(map(first, all_collected_artifacts))
-    longest_name_length = maximum(all_collected_artifacts; init=0) do (artifacts_toml, artifacts)
-        maximum(textwidth, keys(artifacts); init=0)
+    all_collected_artifacts = reduce(
+        vcat, map(
+            ((pkg_root, pkg_uuid),) ->
+            map(ca -> (ca[1], ca[2], pkg_uuid), collect_artifacts(pkg_root; platform, include_lazy)), pkg_info
+        )
+    )
+    used_artifact_tomls = Set{String}(map(ca -> ca[1], all_collected_artifacts))
+    longest_name_length = maximum(all_collected_artifacts; init = 0) do (artifacts_toml, artifacts, pkg_uuid)
+        maximum(textwidth, keys(artifacts); init = 0)
     end
-    for (artifacts_toml, artifacts) in all_collected_artifacts
+    for (artifacts_toml, artifacts, pkg_uuid) in all_collected_artifacts
         # For each Artifacts.toml, install each artifact we've collected from it
         for name in keys(artifacts)
             local rname = rpad(name, longest_name_length)
@@ -981,9 +994,12 @@ function download_artifacts(ctx::Context;
                     dstate.status_update_time = t
                 end
             end
+            # Check if the current package is eligible for PkgServer artifact downloads
+            local pkg_server_eligible = pkg_uuid !== nothing && Registry.is_pkg_in_pkgserver_registry(pkg_uuid, server_registry_info, ctx.registries)
+
             # returns a string if exists, or function that downloads the artifact if not
             local ret = ensure_artifact_installed(name, artifacts[name], artifacts_toml;
-                                        verbose, quiet_download=!(usable_io(io)), io, progress)
+                                        pkg_server_eligible, verbose, quiet_download=!(usable_io(io)), io, progress)
             if ret isa Function
                 download_states[hash] = dstate
                 download_jobs[hash] =
@@ -1175,17 +1191,10 @@ function download_source(ctx::Context, pkgs; readonly=true)
                         archive_urls = Pair{String,Bool}[]
                         # Check if the current package is available in one of the registries being tracked by the pkg server
                         # In that case, download from the package server
-                        if server_registry_info !== nothing
+                        if Registry.is_pkg_in_pkgserver_registry(pkg.uuid, server_registry_info, ctx.registries)
                             server, registry_info = server_registry_info
-                            for reg in ctx.registries
-                                if reg.uuid in keys(registry_info)
-                                    if haskey(reg, pkg.uuid)
-                                        url = "$server/package/$(pkg.uuid)/$(pkg.tree_hash)"
-                                        push!(archive_urls, url => true)
-                                        break
-                                    end
-                                end
-                            end
+                            url = "$server/package/$(pkg.uuid)/$(pkg.tree_hash)"
+                            push!(archive_urls, url => true)
                         end
                         for repo_url in urls
                             url = get_archive_url_for_version(repo_url, pkg.tree_hash)

src/Registry/Registry.jl

@@ -103,6 +103,27 @@ end
 
 pkg_server_url_hash(url::String) = Base.SHA1(split(url, '/')[end])
 
+"""
+    is_pkg_in_pkgserver_registry(pkg_uuid::Base.UUID, server_registry_info, registries)
+
+Check if a package UUID is tracked by the PkgServer by verifying it exists in
+a registry that is known to the PkgServer.
+"""
+function is_pkg_in_pkgserver_registry(pkg_uuid::Base.UUID, server_registry_info, registries)
+    server_registry_info === nothing && return false
+    registries === nothing && return false
+
+    server, registry_info = server_registry_info
+    for reg in registries
+        if reg.uuid in keys(registry_info)
+            if haskey(reg, pkg_uuid)
+                return true
+            end
+        end
+    end
+    return false
+end
+
 function download_default_registries(io::IO; only_if_empty::Bool = true, depots::Union{String, Vector{String}}=depots())
     # Check the specified depots for installed registries
     installed_registries = reachable_registries(; depots)