Add alloc profiling start/stop endpoints (#10)

* Add initial alloc profiling start/stop endpoints

* Add tests

Author: nickrobinson251

src/PerformanceProfilingHttpEndpoints.jl

@@ -123,27 +123,37 @@ end
 
 @static if !(isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs))
 
-function allocations_profile_endpoint(::HTTP.Request)
-    return HTTP.Response(501, "You must use a build of Julia (1.8+) and PProf that support Allocations profiling.")
+for f in (:allocations_profile_endpoint, :allocations_start_endpoint, :allocations_stop_endpoint)
+    @eval function $f(::HTTP.Request)
+        return HTTP.Response(501, "You must use a build of Julia (1.8+) and PProf that support Allocations profiling.")
+    end
 end
 
 else
 
 function allocations_profile_endpoint(req::HTTP.Request)
-
     uri = HTTP.URI(req.target)
     qp = HTTP.queryparams(uri)
     if isempty(qp)
         @info "TODO: interactive HTML input page"
         return HTTP.Response(400, allocs_profile_error_message())
     end
-
     sample_rate = convert(Float64, parse(Float64, get(qp, "sample_rate", default_alloc_sample_rate())))
     duration = parse(Float64, get(qp, "duration", default_duration()))
-
     return _do_alloc_profile(duration, sample_rate)
 end
 
+function allocations_start_endpoint(req::HTTP.Request)
+    uri = HTTP.URI(req.target)
+    qp = HTTP.queryparams(uri)
+    sample_rate = convert(Float64, parse(Float64, get(qp, "sample_rate", default_alloc_sample_rate())))
+    return _start_alloc_profile(sample_rate)
+end
+
+function allocations_stop_endpoint(req::HTTP.Request)
+    return _stop_alloc_profile()
+end
+
 function _do_alloc_profile(duration, sample_rate)
     @info "Starting allocation Profiling from PerformanceProfilingHttpEndpoints with configuration:" duration sample_rate
 
@@ -158,13 +168,30 @@ function _do_alloc_profile(duration, sample_rate)
             "allocs_profile-duration=$duration&sample_rate=$sample_rate.pb.gz")
 end
 
+function _start_alloc_profile(sample_rate)
+    @info "Starting allocation Profiling from PerformanceProfilingHttpEndpoints with configuration:" sample_rate
+    Profile.Allocs.clear()
+    Profile.Allocs.start(; sample_rate)
+    return HTTP.Response(200, "Allocation profiling started.")
+end
+
+function _stop_alloc_profile()
+    Profile.Allocs.stop()
+    prof_name = tempname()
+    PProf.Allocs.pprof(out=prof_name, web=false)
+    prof_name = "$prof_name.pb.gz"
+    return _http_response(read(prof_name), "allocs_profile.pb.gz")
+end
+
 end  # if isdefined
 
 function serve_profiling_server(;addr="127.0.0.1", port=16825, verbose=false, kw...)
     verbose >= 0 && @info "Starting HTTP profiling server on port $port"
     router = HTTP.Router()
     HTTP.register!(router, "/profile", cpu_profile_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)
     # HTTP.serve! returns listening/serving server object
     return HTTP.serve!(router, addr, port; verbose, kw...)
 end
@@ -176,8 +203,12 @@ function __init__()
     precompile(cpu_profile_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
     precompile(_do_cpu_profile, (Int,Float64,Float64,Bool)) || 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")
     if isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs)
         precompile(_do_alloc_profile, (Float64,Float64,)) || error("precompilation of package functions is not supposed to fail")
+        precompile(_start_alloc_profile, (Float64,)) || error("precompilation of package functions is not supposed to fail")
+        precompile(_stop_alloc_profile, ()) || error("precompilation of package functions is not supposed to fail")
     end
 end
 

test/runtests.jl

@@ -43,32 +43,64 @@ const url = "http://127.0.0.1:$port"
     @testset "Allocation profiling" begin
         done = Threads.Atomic{Bool}(false)
         # Schedule some work that's known to be expensive, to profile it
-        t = @async begin
+        workload() = @async begin
             for _ in 1:200
                 if done[] return end
                 global a = [[] for i in 1:1000]
                 yield()  # yield to allow the tests to run
             end
         end
 
-        req = HTTP.get("$url/allocs_profile?duration=3", retry=false, status_exception=false)
-        if !(isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs))
-            # We should be tesing the Allocs profiling if we're on julia nightly.
-            @assert VERSION < v"1.8.0-DEV.1346"
-            @test req.status == 501  # not implemented
-        else
-            @test req.status == 200
-            @test length(req.body) > 0
+        @testset "allocs_profile endpoint" begin
+            done[] = false
+            t = workload()
+            req = HTTP.get("$url/allocs_profile?duration=3", retry=false, status_exception=false)
+            if !(isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs))
+                @assert VERSION < v"1.8.0-DEV.1346"
+                @test req.status == 501  # not implemented
+            else
+                @test req.status == 200
+                @test length(req.body) > 0
+
+                data = read(IOBuffer(req.body), String)
+                # Test that there's something here
+                # TODO: actually parse the profile
+                @test length(data) > 100
+            end
+            @info "Finished `allocs_profile` tests, waiting for workload to finish."
+            done[] = true
+            wait(t)  # handle errors
+        end
 
-            data = read(IOBuffer(req.body), String)
-            # Test that there's something here
-            # TODO: actually parse the profile
-            @test length(data) > 100
+        @testset "allocs_profile_start/stop endpoints" begin
+            done[] = false
+            t = workload()
+            req = HTTP.get("$url/allocs_profile_start", retry=false, status_exception=false)
+            if !(isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs))
+                @assert VERSION < v"1.8.0-DEV.1346"
+                @test req.status == 501  # not implemented
+            else
+                @test req.status == 200
+                @test String(req.body) == "Allocation profiling started."
+            end
 
+            sleep(3)  # Allow workload to run a while before we stop profiling.
+
+            req = HTTP.get("$url/allocs_profile_stop", retry=false, status_exception=false)
+            if !(isdefined(Profile, :Allocs) && isdefined(PProf, :Allocs))
+                @assert VERSION < v"1.8.0-DEV.1346"
+                @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
+            @info "Finished `allocs_profile_stop` tests, waiting for workload to finish."
+            done[] = true
+            wait(t)  # handle errors
         end
-        @info "Finished tests, waiting for workload to finish."
-        done[] = true
-        wait(t)  # handle errors
     end
 
     @testset "error handling" begin