Merge pull request #263 from IanButterworth/ib/cache_download_tweaks

Symbol caches: Larger download batching. Don't ask for non-General pkgs

Author: pfitzseb

src/SymbolServer.jl

@@ -27,6 +27,65 @@ mutable struct SymbolServerInstance
     end
 end
 
+function get_general_pkgs()
+    dp_before = copy(Base.DEPOT_PATH)
+    try
+        # because the env var JULIA_DEPOT_PATH is overritten this is probably the best
+        # guess depot location
+        push!(empty!(Base.DEPOT_PATH), joinpath(homedir(), ".julia"))
+        @static if VERSION >= v"1.7-"
+            regs = Pkg.Types.Context().registries
+            i = findfirst(r -> r.name == "General" && r.uuid == UUID("23338594-aafe-5451-b93e-139f81909106"), regs)
+            i === nothing && return Dict{UUID, PkgEntry}()
+            return regs[i].pkgs
+        else
+            for r in Pkg.Types.collect_registries()
+                (r.name == "General" && r.uuid == UUID("23338594-aafe-5451-b93e-139f81909106")) || continue
+                reg = Pkg.Types.read_registry(joinpath(r.path, "Registry.toml"))
+                return reg["packages"]
+            end
+            return Dict{UUID, PkgEntry}()
+        end
+    finally
+        append!(empty!(Base.DEPOT_PATH), dp_before)
+    end
+end
+
+"""
+    remove_non_general_pkgs!(pkgs)
+
+Removes packages that aren't going to be on the symbol cache server because they aren't in the General registry.
+This avoids leaking private package name & uuid pairs via the url requests to the symbol server.
+
+If the General registry cannot be found packages cannot be checked, so all packages will be removed.
+"""
+function remove_non_general_pkgs!(pkgs)
+    general_pkgs = get_general_pkgs()
+    if isempty(general_pkgs)
+        @warn """
+        Could not find the General registry when checking for whether packages are public.
+        All package symbol caches will be generated locally"""
+        return empty!(pkgs)
+    end
+    filter!(pkgs) do pkg
+        packageuuid(pkg) === nothing && return false
+        packagename(pkg) === nothing && return false
+        tree_hash(pkg) === nothing && return false # stdlibs and dev-ed packages don't have tree_hash and aren't cached
+        @static if VERSION >= v"1.7-"
+            uuid_match = get(general_pkgs, packageuuid(pkg), nothing)
+            uuid_match === nothing && return false
+            uuid_match.name != packagename(pkg) && return false
+            return true
+        else
+            uuid_match = get(general_pkgs, string(packageuuid(pkg)), nothing)
+            uuid_match === nothing && return false
+            uuid_match["name"] != packagename(pkg) && return false
+            return true
+        end
+    end
+    return pkgs
+end
+
 function getstore(ssi::SymbolServerInstance, environment_path::AbstractString, progress_callback=nothing, error_handler=nothing; download = false)
     !ispath(environment_path) && return :success, recursive_copy(stdlibs)
 
@@ -42,20 +101,33 @@ function getstore(ssi::SymbolServerInstance, environment_path::AbstractString, p
                     if manifest !== nothing
                         @debug "Downloading cache files for manifest at $(manifest_filename)."
                         to_download = collect(validate_disc_store(ssi.store_path, manifest))
-                        batches = Iterators.partition(to_download, max(1, floor(Int, length(to_download)÷50)))
-                        for (i, batch) in enumerate(batches)
-                            percentage = round(Int, 100*(i - 1)/length(batches))
-                            progress_callback !== nothing && progress_callback("Downloading caches...", percentage)
-                            @sync for pkg in batch
-                                @async begin
-                                    yield()
-                                    uuid = packageuuid(pkg)
-                                    get_file_from_cloud(manifest, uuid, environment_path, ssi.depot_path, ssi.store_path, download_dir, ssi.symbolcache_upstream)
-                                    yield()
+                        try
+                            remove_non_general_pkgs!(to_download)
+                        catch err
+                            # if any errors, err on the side of caution and mark all as private, and continue
+                            @error """
+                            Symbol cache downloading: Failed to identify which packages to omit based on the General registry.
+                            All packages will be processsed locally""" err
+                            empty!(to_download)
+                        end
+                        if !isempty(to_download)
+                            n_done = 0
+                            n_total = length(to_download)
+                            for batch in Iterators.partition(to_download, 100) # 100 connections at a time
+                                @sync for pkg in batch
+                                    @async begin
+                                        yield()
+                                        uuid = packageuuid(pkg)
+                                        get_file_from_cloud(manifest, uuid, environment_path, ssi.depot_path, ssi.store_path, download_dir, ssi.symbolcache_upstream)
+                                        yield()
+                                        n_done += 1
+                                        percentage = round(Int, 100*(n_done/n_total))
+                                        progress_callback !== nothing && progress_callback("Downloading caches...", percentage)
+                                    end
                                 end
                             end
+                            progress_callback !== nothing && progress_callback("All cache files downloaded.", 100)
                         end
-                        progress_callback !== nothing && progress_callback("All cache files downloaded.", 100)
                     end
                 end
             end

src/utils.jl

@@ -143,7 +143,7 @@ else
     packageuuid(pkg::Pair{String,UUID}) = last(pkg)
     packageuuid(pkg::Pair{UUID,PackageEntry}) = first(pkg)
 
-    packagename(pkg::Pair{UUID,PackageEntry})::String = last(pkg).name
+    packagename(pkg::Pair{UUID,PackageEntry})::Union{Nothing,String} = last(pkg).name
     packagename(c::Pkg.Types.Context, uuid::UUID) = manifest(c)[uuid].name
     packagename(manifest::Dict{UUID,PackageEntry}, uuid::UUID) = manifest[uuid].name
 
@@ -162,7 +162,8 @@ else
     version(pe::Pair{UUID,PackageEntry}) = last(pe).version
     frommanifest(c::Pkg.Types.Context, uuid) = manifest(c)[uuid]
     frommanifest(manifest::Dict{UUID,PackageEntry}, uuid) = manifest[uuid]
-    tree_hash(pe::PackageEntry) = VERSION >= v"1.3" ? pe.tree_hash : get(pe.other, "git-tree-sha1", nothing)
+    tree_hash(pkg::Pair{UUID,PackageEntry}) = tree_hash(last(pkg))
+    tree_hash(pe::PackageEntry) = VERSION >= v"1.3" ? pe.tree_hash : (pe.other === nothing ? nothing : get(pe.other, "git-tree-sha1", nothing))
 
     is_package_deved(manifest, uuid) = manifest[uuid].path !== nothing
 end

test/runtests.jl

@@ -146,6 +146,34 @@ end
         @test !isempty(SymbolServer.stdlibs[:Base][:VecOrMat].doc)     # Union
         @test occursin("Cint", SymbolServer.stdlibs[:Base][:Cint].doc) # Alias
     end
+
+    if VERSION >= v"1.1-"
+        @testset "Excluding private packages from cache download requests" begin
+            pkgs = Dict{Base.UUID, Pkg.Types.PackageEntry}()
+            if VERSION < v"1.3-"
+                pkgs[UUID("7876af07-990d-54b4-ab0e-23690620f79a")] = Pkg.Types.PackageEntry(name="Example", other=Dict("git-tree-sha1" => Base.SHA1("0"^40)))
+                pkgs[UUID("3e13f8c9-a9aa-412e-8b2a-fda000b375e2")] = Pkg.Types.PackageEntry(name="NotInGeneral", other=Dict("git-tree-sha1" => Base.SHA1("0"^40)))
+                pkgs[UUID("eb4ab7d2-1172-48bd-a954-ae6825f2e6e3")] = Pkg.Types.PackageEntry(other=Dict("git-tree-sha1" => Base.SHA1("0"^40))) # no name
+                pkgs[UUID("1fec9e91-426f-45f4-a317-da8b2730f864")] = Pkg.Types.PackageEntry(name="NoTreeHash") # no tree_hash, like stdlibs
+            else
+                pkgs[UUID("7876af07-990d-54b4-ab0e-23690620f79a")] = Pkg.Types.PackageEntry(name="Example", tree_hash=Base.SHA1("0"^40))
+                pkgs[UUID("3e13f8c9-a9aa-412e-8b2a-fda000b375e2")] = Pkg.Types.PackageEntry(name="NotInGeneral", tree_hash=Base.SHA1("1"^40))
+                pkgs[UUID("eb4ab7d2-1172-48bd-a954-ae6825f2e6e3")] = Pkg.Types.PackageEntry(tree_hash=Base.SHA1("2"^40)) # no name
+                pkgs[UUID("1fec9e91-426f-45f4-a317-da8b2730f864")] = Pkg.Types.PackageEntry(name="NoTreeHash") # no tree_hash, like stdlibs
+            end
+
+            SymbolServer.remove_non_general_pkgs!(pkgs)
+
+            @test length(pkgs) == 1
+            @test haskey(pkgs, UUID("7876af07-990d-54b4-ab0e-23690620f79a"))
+            @test pkgs[UUID("7876af07-990d-54b4-ab0e-23690620f79a")].name == "Example"
+            if VERSION < v"1.3"
+                @test pkgs[UUID("7876af07-990d-54b4-ab0e-23690620f79a")].other["git-tree-sha1"] == Base.SHA1("0"^40)
+            else
+                @test pkgs[UUID("7876af07-990d-54b4-ab0e-23690620f79a")].tree_hash == Base.SHA1("0"^40)
+            end
+        end
+    end
 end
 
 using SymbolServer: FakeTypeName