Type inference profiling

Project.toml

@@ -4,10 +4,13 @@ authors = ["Nathan Daly <[email protected]>", "Dana Wilson <[email protected]>
 version = "1.0.0"
 
 [deps]
+FlameGraphs = "08572546-2f56-4bcf-ba4e-bab62c3a3f89"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 PProf = "e4faabce-9ead-11e9-39d9-4379958e3056"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+SnoopCompile = "aa65fe97-06da-5843-b5b1-d5d13cad87d2"
+SnoopCompileCore = "e2b509da-e806-4183-be48-004708413034"
 
 [compat]
 HTTP = "1"

src/ProfileEndpoints.jl

@@ -4,6 +4,10 @@ import HTTP
 import Profile
 import PProf
 
+using FlameGraphs
+using SnoopCompile
+import SnoopCompileCore
+
 using Serialization: serialize
 
 #----------------------------------------------------------
@@ -242,6 +246,42 @@ end
 
 end  # if isdefined
 
+###
+### Type Inference
+###
+
+function typeinf_start_endpoint(req::HTTP.Request)
+    if !isdefined(Core.Compiler.Timings, :clear_and_fetch_timings)
+        # See: https://github.com/JuliaLang/julia/pull/47615.
+        return HTTP.Response(501, "Type inference profiling isn't thread safe without Julia #47615.")
+    end
+    SnoopCompileCore.start_deep_timing()
+    return HTTP.Response(200, "Type inference profiling started.")
+end
+
+function typeinf_stop_endpoint(req::HTTP.Request)
+    if !isdefined(Core.Compiler.Timings, :clear_and_fetch_timings)
+        # See: https://github.com/JuliaLang/julia/pull/47615.
+        return HTTP.Response(501, "Type inference profiling isn't thread safe without Julia #47615.")
+    end
+
+    SnoopCompileCore.stop_deep_timing()
+    timings = SnoopCompileCore.finish_snoopi_deep()
+
+    # Currently, SnoopCompile will throw an error if timings is empty..
+    # Reported, here: https://github.com/timholy/SnoopCompile.jl/pull/212/files#r1062926193
+    if isempty(timings.children)
+        # So just return an empty profile..
+        return _http_response("", "inference_profile.pb.gz")
+    end
+
+    flame_graph = flamegraph(timings)
+    prof_name = tempname()
+    PProf.pprof(flame_graph; out=prof_name, web=false)
+    prof_name = "$prof_name.pb.gz"
+    return _http_response(read(prof_name), "inference_profile.pb.gz")
+end
+
 ###
 ### Server
 ###
@@ -256,6 +296,8 @@ function serve_profiling_server(;addr="127.0.0.1", port=16825, verbose=false, kw
     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.register!(router, "/typeinf_profile_start", typeinf_start_endpoint)
+    HTTP.register!(router, "/typeinf_profile_stop", typeinf_stop_endpoint)
     # HTTP.serve! returns listening/serving server object
     return HTTP.serve!(router, addr, port; verbose, kw...)
 end
@@ -281,6 +323,9 @@ function __init__()
         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
+
+    precompile(typeinf_start_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
+    precompile(typeinf_stop_endpoint, (HTTP.Request,)) || error("precompilation of package functions is not supposed to fail")
 end
 
 end # module ProfileEndpoints

test/runtests.jl

@@ -1,8 +1,8 @@
 module ProfileEndpointsTests
 
 using ProfileEndpoints
-using Test
 using Serialization
+using Test
 
 import InteractiveUtils
 import HTTP
@@ -154,6 +154,30 @@ const url = "http://127.0.0.1:$port"
         end
     end
 
+    @testset "Type inference profiling" begin
+        if !isdefined(Core.Compiler.Timings, :clear_and_fetch_timings)
+            @test HTTP.get("$url/typeinf_profile_start", retry=false, status_exception=false).status == 501
+            @test HTTP.get("$url/typeinf_profile_stop", retry=false, status_exception=false).status == 501
+        else
+            @testset "typeinf start/stop endpoints" begin
+                resp = HTTP.get("$url/typeinf_profile_start", retry=false, status_exception=false)
+                @test resp.status == 200
+                @test String(resp.body) == "Type inference profiling started."
+
+                # workload
+                @eval foo() = 2
+                @eval foo()
+
+                resp = HTTP.get("$url/typeinf_profile_stop", retry=false, status_exception=false)
+                @test resp.status == 200
+                data = read(IOBuffer(resp.body), String)
+                # Test that there's something here
+                # TODO: actually parse the profile
+                @test length(data) > 100
+            end
+        end
+    end
+
     @testset "error handling" begin
         let res = HTTP.get("$url/profile", status_exception=false)
             @test 400 <= res.status < 500