Add heap snapshot endpoint

.github/workflows/CI.yml

@@ -13,8 +13,9 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1.6"
-          - "1"
+          - "1.6"  # LTS
+          - "1"    # Latest
+          - "~1.9.0-0"  # To test heap snapshot; remove when v1.9 is latest Julia release.
           - "nightly"
         os:
           - ubuntu-latest

README.md

@@ -9,6 +9,8 @@ Currently provides:
     - `/profile?n=1e8&delay=0.01&duration=10&pprof=true`
 - `/allocs_profile` endpoint, with default query params:
     - `/allocs_profile?sample_rate=0.0001&duration=10`
+- `/heap_snapshot` endpoint, with default query params:
+    - `/heap_snapshot?all_one=false`
 
 
 ## Example

src/PerformanceProfilingHttpEndpoints.jl

@@ -133,10 +133,29 @@ end
 ### Allocs
 ###
 
+# If `all_one=true`, then every object is given size 1 so they can be easily counted.
+# Otherwise, if `false`, every object reports its actual size on the heap.
+default_heap_all_one() = "false"
+
+@static if !isdefined(Profile, :take_heap_snapshot)
+
+function heap_snapshot_endpoint(::HTTP.Request)
+    return HTTP.Response(501, "You must use a build of Julia (1.9+) that supports heap snapshots.")
+end
+
+else
+
 function heap_snapshot_endpoint(req::HTTP.Request)
-    # TODO: implement this once https://github.com/JuliaLang/julia/pull/42286 is merged
+    uri = HTTP.URI(req.target)
+    qp = HTTP.queryparams(uri)
+    all_one = parse(Bool, get(qp, "all_one", default_heap_all_one()))
+    filename = Profile.take_heap_snapshot(all_one)
+    @info "Taking heap snapshot from PerformanceProfilingHttpEndpoints" all_one filename
+    return _http_response(read(filename), filename)
 end
 
+end  # if isdefined
+
 default_alloc_sample_rate() = "0.0001"
 
 allocs_profile_error_message() = """Need to provide query params:
@@ -233,6 +252,7 @@ function serve_profiling_server(;addr="127.0.0.1", port=16825, verbose=false, kw
     HTTP.register!(router, "/profile", cpu_profile_endpoint)
     HTTP.register!(router, "/profile_start", cpu_profile_start_endpoint)
     HTTP.register!(router, "/profile_stop", cpu_profile_stop_endpoint)
+    HTTP.register!(router, "/heap_snapshot", heap_snapshot_endpoint)
     HTTP.register!(router, "/allocs_profile", allocations_profile_endpoint)
     HTTP.register!(router, "/allocs_profile_start", allocations_start_endpoint)
     HTTP.register!(router, "/allocs_profile_stop", allocations_stop_endpoint)
@@ -251,6 +271,8 @@ function __init__()
     precompile(_do_cpu_profile, (Int,Float64,Float64,Bool)) || error("precompilation of package functions is not supposed to fail")
     precompile(_start_cpu_profile, (Int,Float64,)) || error("precompilation of package functions is not supposed to fail")
 
+    precompile(heap_snapshot_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
+
     precompile(allocations_profile_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
     precompile(allocations_start_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
     precompile(allocations_stop_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")

test/runtests.jl

@@ -75,6 +75,22 @@ const url = "http://127.0.0.1:$port"
         end
     end
 
+    @testset "Heap snapshot $query" for query in ("", "?all_one=true")
+        req = HTTP.get("$url/heap_snapshot$query", retry=false, status_exception=false)
+        if !isdefined(Profile, :take_heap_snapshot)
+            # Assert the version is before https://github.com/JuliaLang/julia/pull/46862
+            # Although we actually also need https://github.com/JuliaLang/julia/pull/47300
+            @assert VERSION < v"1.9.0-DEV.1643"
+            @test req.status == 501  # not implemented
+        else
+            @test req.status == 200
+            data = read(IOBuffer(req.body), String)
+            # Test that there's something here
+            # TODO: actually parse the profile
+            @test length(data) > 100
+        end
+    end
+
     @testset "Allocation profiling" begin
         done = Threads.Atomic{Bool}(false)
         # Schedule some work that's known to be expensive, to profile it