From 91d30ef9eec56103f00ee7fb7c9eb4a8e7ed18dd Mon Sep 17 00:00:00 2001
From: Nathan Daly <nathan.daly@relational.ai>
Date: Fri, 20 Oct 2023 16:54:34 -0600
Subject: [PATCH] Fix for v1.6

---
 src/PProf.jl  | 20 +++++++++++++++++---
 test/PProf.jl | 18 ++++++++++++++++--
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/PProf.jl b/src/PProf.jl
index 1b08603..a97f73f 100644
--- a/src/PProf.jl
+++ b/src/PProf.jl
@@ -48,6 +48,7 @@ using Base.StackTraces: StackFrame
             web = true, webhost = "localhost", webport = 57599,
             out = "profile.pb.gz", from_c = true, full_signatures = true, drop_frames = "",
             keep_frames = "", ui_relative_percentages = true, sampling_delay = nothing,
+            tagroot = "taskid,threadid"
          )
     pprof(FlameGraphs.flamegraph(); kwargs...)
 
@@ -78,6 +79,14 @@ You can also use `PProf.refresh(file="...")` to open a new file in the server.
 - `from_c::Bool`: If `false`, exclude frames that come from from_c. Defaults to `true`.
 - `full_signatures::Bool`: If `true`, methods are printed as signatures with full argument
                            types. If `false`, as only names. E.g. `eval(::Module, ::Any)` vs `eval`.
+- `tagroot`: Set which metadata tags you want to turn into root frames for the profile. This
+             is used to view the metadata tags in the Flamegraph view. This should be a
+             comma-separated string, chosing from the following metadata options:
+             - `taskid`
+             - `threadid`
+             - `thread_sleeping`
+             - `cycle_clock`
+             Defaults to `"taskid,threadid"`, grouping by taskid then threadid.
 - `drop_frames`: frames with function_name fully matching regexp string will be dropped from the samples,
                  along with their successors.
 - `keep_frames`: frames with function_name fully matching regexp string will be kept, even if it matches drop_functions.
@@ -96,6 +105,7 @@ function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
                drop_frames::Union{Nothing, AbstractString} = nothing,
                keep_frames::Union{Nothing, AbstractString} = nothing,
                ui_relative_percentages::Bool = true,
+               tagroot::Union{Nothing, AbstractString} = "taskid,threadid",
             )
     has_meta = false
     if data === nothing
@@ -319,8 +329,7 @@ function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
     end
 
     if web
-        refresh(webhost = webhost, webport = webport, file = out,
-            ui_relative_percentages = ui_relative_percentages)
+        refresh(; webhost, webport, file = out, ui_relative_percentages, tagroot)
     end
 
     out
@@ -361,6 +370,7 @@ function refresh(; webhost::AbstractString = "localhost",
                    webport::Integer = 57599,
                    file::AbstractString = "profile.pb.gz",
                    ui_relative_percentages::Bool = true,
+                   tagroot::Union{AbstractString,Nothing} = "taskid,threadid",
                 )
 
     if proc[] === nothing
@@ -374,7 +384,11 @@ function refresh(; webhost::AbstractString = "localhost",
     relative_percentages_flag = ui_relative_percentages ? "-relative_percentages" : ""
 
     proc[] = pprof_jll.pprof() do pprof_path
-        open(pipeline(`$pprof_path -http=$webhost:$webport $relative_percentages_flag $file`))
+        if tagroot !== nothing && !isempty(tagroot)
+            open(pipeline(`$pprof_path -tagroot $tagroot -http=$webhost:$webport $relative_percentages_flag $file`))
+        else
+            open(pipeline(`$pprof_path -http=$webhost:$webport $relative_percentages_flag $file`))
+        end
     end
 end
 
diff --git a/test/PProf.jl b/test/PProf.jl
index 2d81d0d..b5a86bd 100644
--- a/test/PProf.jl
+++ b/test/PProf.jl
@@ -110,18 +110,32 @@ end
     end
     @testset for i in 1:4
         if i == 1
+            if !Profile.has_meta
+                continue
+            end
             data = Profile.fetch(include_meta = true)
             args = (data,)
         elseif i == 2
+            if !Profile.has_meta
+                continue
+            end
             data,lidict = Profile.retrieve(include_meta = true)
             args = (data, lidict)
         elseif i == 3
             # Ensure we are backwards compatible with older, non-meta profiles
-            data = Profile.fetch(include_meta = false)
+            if Profile.has_meta
+                data = Profile.fetch(include_meta = false)
+            else
+                data = Profile.fetch()
+            end
             args = (data,)
         else
             # Ensure we are backwards compatible with older, non-meta profiles
-            data,lidict = Profile.retrieve(include_meta = false)
+            if Profile.has_meta
+                data,lidict = Profile.retrieve(include_meta = false)
+            else
+                data,lidict = Profile.retrieve()
+            end
             args = (data, lidict)
         end