allow passing a Accept-Encoding header to PkgServer to fetch e.g. zstd compressed files from the storage server

for backwards compatibility, return gzipped files when no `Accept-Encoding` is provided

Author: KristofferC

src/PkgServer.jl

@@ -24,6 +24,49 @@ include("meta.jl")
 include("admin.jl")
 include("dynamic.jl")
 
+# Detect compression format by reading file magic bytes
+function detect_file_compression(filepath::AbstractString)::Union{String,Nothing}
+    magic = open(filepath, "r") do io
+        read(io, min(4, filesize(io)))
+    end
+
+    # Zstd
+    if length(magic) >= 4 && magic[1:4] == [0x28, 0xB5, 0x2F, 0xFD]
+        return "zstd"
+    end
+    # Gzip
+    if length(magic) >= 2 && magic[1:2] == [0x1F, 0x8B]
+        return "gzip"
+    end
+
+    # Not enough bytes to detect or unknown format
+    return nothing
+end
+
+# Check if client accepts a given compression format
+function client_accepts_format(http::HTTP.Stream, format::String)::Bool
+    accept_encoding = HTTP.header(http, "Accept-Encoding", "")
+    # Empty header means only gzip is accepted (backwards compatibility)
+    if isempty(accept_encoding)
+        return format == "gzip"
+    end
+    return occursin(format, accept_encoding)
+end
+
+# Get content type based on file format
+function get_content_type(filepath::AbstractString)
+    file_format = detect_file_compression(filepath)
+
+    if file_format == "zstd"
+        return "application/x-zstd"
+    elseif file_format == "gzip"
+        return "application/x-gzip"
+    else
+        # Unknown or insufficient bytes - return generic octet-stream
+        return "application/octet-stream"
+    end
+end
+
 mutable struct RegistryMeta
     # Upstream registry URL (e.g. "https://github.com/JuliaRegistries/General")
     upstream_url::String
@@ -294,20 +337,30 @@ function handle_request(http::HTTP.Stream)
 
     # If the user asked for something that is an actual resource, send it directly
     if occursin(resource_re, resource)
-        # If the resource already exists locally, yay!  Serve it and quit.
+        # If the resource already exists locally, check if client can accept the format
         resource_path = resource_filepath(resource)
         io = try_open(resource_path)
         if io !== nothing
-            hit!(config.cache, resource[2:end])
-            serve_file(http, io, "application/x-gzip")
-            close(io)
-            return
+            cached_format = detect_file_compression(resource_path)
+
+            # Serve from cache if format is unknown or client accepts it
+            if cached_format === nothing || client_accepts_format(http, cached_format)
+                hit!(config.cache, resource[2:end])
+                content_type = get_content_type(resource_path)
+                serve_file(http, io, content_type)
+                close(io)
+                return
+            else
+                # Cached format not acceptable, close and re-fetch in desired format
+                close(io)
+                # Fall through to fetch_resource below
+            end
         end
 
         # If it doesn't exist locally, let's request a fetch on that resource.
         # This will return either `nothing` (e.g. resource does not exist) or
         # a `DownloadState` that represents a partial download.
-        dl_state = fetch_resource(resource, request_id)
+        dl_state = fetch_resource(resource, request_id, http)
         if dl_state !== nothing
             HTTP.setheader(http, "X-Cache-Miss" => "miss")
             stream_path = temp_resource_filepath(resource)
@@ -318,7 +371,8 @@ function handle_request(http::HTTP.Stream)
             # Try to serve `stream_path` file
             stream_io = try_open(stream_path)
             if stream_io !== nothing
-                serve_file(http, stream_io, "application/x-gzip";
+                content_type = get_content_type(stream_path)
+                serve_file(http, stream_io, content_type;
                            content_length=dl_state.content_length,
                            dl_task=dl_state.dl_task)
                 close(stream_io)
@@ -329,7 +383,8 @@ function handle_request(http::HTTP.Stream)
             # downloading since we last checked 20 lines ago.  Check again.
             io = try_open(resource_path)
             if io !== nothing
-                serve_file(http, io, "application/x-gzip")
+                content_type = get_content_type(resource_path)
+                serve_file(http, io, content_type)
                 close(io)
                 return
             end
@@ -377,7 +432,7 @@ function handle_request(http::HTTP.Stream)
     return
 end
 
-# precompilation                                                                                                                    
+# precompilation
 include(joinpath(dirname(@__DIR__), "deps", "precompile.jl"))
 if get(ENV, "PKGSERVER_GENERATING_PRECOMPILE", nothing) === nothing
     _precompile_()

src/resource.jl

@@ -257,20 +257,21 @@ Given a resource and a list of storage servers, return the storage server that r
 most quickly to a HEAD request for that storage server, as well as the HEAD response
 itself so that metadata such as the content-length of the resource can be inspected.
 """
-function select_server(resource::AbstractString, servers::Vector{<:AbstractString}; timeout = 5, retries = 2)
-    function head_req(server, resource)
+function select_server(resource::AbstractString, servers::Vector{<:AbstractString}, accept_encoding::AbstractString="gzip"; timeout = 5, retries = 2)
+    function head_req(server, resource, accept_encoding)
         @try_printerror begin
             response = HTTP.head(
                 string(server, resource);
                 status_exception = false,
                 timeout = timeout,
                 retries = retries,
+                headers = ["Accept-Encoding" => accept_encoding],
             )
             return server, response
         end
     end
     # Launch one task per server, performing a HEAD request
-    tasks = [@spawn head_req(server, resource) for server in servers]
+    tasks = [@spawn head_req(server, resource, accept_encoding) for server in servers]
 
     # Wait for the first Task that gives us an HTTP 200 OK, returning that server.
     # If none have it, we return `nothing`. :(
@@ -314,12 +315,26 @@ are recorded and future downloads of that same resource will be skipped, until
 `forget_failures()` is called.  The `DownloadState` object contains within it enough
 information to still serve a resource as it is being downloaded in the background task.
 """
-function fetch_resource(resource::AbstractString, request_id::AbstractString; servers::Vector{String}=config.storage_servers)
+function fetch_resource(resource::AbstractString, request_id::AbstractString, http::Union{HTTP.Stream,Nothing}=nothing; servers::Vector{String}=config.storage_servers)
     if isempty(servers)
         @error("fetch called with no servers", resource)
         error("fetch called with no servers")
     end
 
+    # Determine what encoding to request from storage servers based on client preferences
+    # For backwards compatibility, if no Accept-Encoding header, default to gzip
+    if http !== nothing
+        client_accept = HTTP.header(http, "Accept-Encoding", "")
+        if isempty(client_accept)
+            accept_encoding = "gzip" # no header means gzip
+        else
+            accept_encoding = client_accept
+        end
+    else
+        # No client context, prefer zstd for storage efficiency
+        accept_encoding = "zstd, gzip"
+    end
+
     # with_fetch_state() will wait for a lock
     with_fetch_state(resource) do state
         # check if this has failed to download recently
@@ -335,15 +350,15 @@ function fetch_resource(resource::AbstractString, request_id::AbstractString; se
         end
 
         # If not, let's figure out which storage server we're going to download from:
-        server, response = select_server(resource, servers)
+        server, response = select_server(resource, servers, accept_encoding)
         if response === nothing
             @debug("no upstream server", resource, servers)
             return nothing
         end
 
         # Launch download process in a separate task:
         dl_task = @async begin
-            success = download(server, resource, content_length(response), request_id)
+            success = download(server, resource, content_length(response), request_id, accept_encoding)
             lock(state.lock) do
                 if success
                     global fetch_hits += 1
@@ -422,8 +437,8 @@ function stream_file(io_in::IO, start_byte::Int, length::Int, dl_task::Task, io_
     return transmitted
 end
 
-function download(server::AbstractString, resource::AbstractString, content_length::Int, request_id::AbstractString)
-    @info("downloading resource", server, resource, request_id)
+function download(server::AbstractString, resource::AbstractString, content_length::Int, request_id::AbstractString, accept_encoding::AbstractString="gzip")
+    @info("downloading resource", server, resource, request_id, accept_encoding)
     t_start = time()
     hash = basename(resource)
 
@@ -465,6 +480,8 @@ function download(server::AbstractString, resource::AbstractString, content_leng
             req = HTTP.get(server * resource,
                 status_exception = false,
                 response_stream = file_io,
+                headers = ["Accept-Encoding" => accept_encoding],
+                decompress = false,
             )
             close(file_io)
             return req

test/tests.jl

@@ -71,13 +71,32 @@ end
     @test endswith(eager_response.request.target, ".eager")
 
     # Test asking for that registry directly, unpacking it and verifying the treehash
+    # Clients without Accept-Encoding header should get gzip
     mktemp() do tarball_path, tarball_io
         response = HTTP.get("$(server_url)/registry/$(registry_uuid)/$(registry_treehash)"; response_stream=tarball_io)
         close(tarball_io)
         @test response.status == 200
+
+        format = PkgServer.detect_file_compression(tarball_path)
+        @test format == "gzip"
+
         @test registry_treehash == Tar.tree_hash(open(pipeline(`cat $(tarball_path)`, `gzip -d`), read=true))
     end
 
+    # Clients that only accept zstd should get zstd
+    mktemp() do tarball_path, tarball_io
+        response = HTTP.get("$(server_url)/registry/$(registry_uuid)/$(registry_treehash)",
+                           ["Accept-Encoding" => "zstd"];
+                           response_stream=tarball_io)
+        close(tarball_io)
+        @test response.status == 200
+
+        format = PkgServer.detect_file_compression(tarball_path)
+        @test format == "zstd"
+
+        @test registry_treehash == Tar.tree_hash(open(pipeline(`zstd -d -c $(tarball_path)`), read=true))
+    end
+
     # Verify that these files exist within the cache
     @test isfile(joinpath(cache_dir, "..", "static", "registries.eager"))
     @test isfile(joinpath(cache_dir, "..", "static", "registries.conservative"))
@@ -278,6 +297,20 @@ end
         # Also test that it's available at its nskip hash:
         @test HTTP.head("$(server_url)/artifact/$(art_yskip_hash)").status == 200
     end
+
+    # Test artifact download with zstd compression
+    mktemp() do tarball_path, tarball_io
+        response = HTTP.get("$(server_url)/artifact/$(art_yskip_hash)",
+                           ["Accept-Encoding" => "zstd"];
+                           response_stream=tarball_io)
+        close(tarball_io)
+        @test response.status == 200
+
+        format = PkgServer.detect_file_compression(tarball_path)
+        @test format == "zstd"
+
+        @test art_nskip_hash == Tar.tree_hash(open(pipeline(`zstd -d -c $(tarball_path)`), read=true))
+    end
 end
 
 @testset "Partial Content" begin