Compress on the fly (#82)

* Compress on the fly * Fixed tests * Bump breaking version change in Project.toml This is a breaking change, since the format of the files we write has changed, and the process to open them has differed. * Fix tests that were loading proto to decompress * Fix test to do decompressing * Update flamegraphs.jl * Update Allocs.jl --------- Co-authored-by: Nathan Daly <NHDaly@gmail.com>
JuliaPerf · Oct 6, 2023 · 6e9f56c · 6e9f56c · vchuravy · Oct 17, 2023
1 parent 45399a8
commit 6e9f56c
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 10 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,10 +1,11 @@
 name = "PProf"
 uuid = "e4faabce-9ead-11e9-39d9-4379958e3056"
 authors = ["Valentin Churavy <v.churavy@gmail.com>", "Nathan Daly <nhdaly@gmail.com>"]
-version = "2.3.0"
+version = "3.0.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FlameGraphs = "08572546-2f56-4bcf-ba4e-bab62c3a3f89"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -23,6 +24,7 @@ ProgressMeter = "1.7"
 ProtoBuf = "1"
 julia = "1.6"
 pprof_jll = "0.1, 1"
+CodecZlib = "0.7"
 
 [extras]
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"

diff --git a/src/Allocs.jl b/src/Allocs.jl
@@ -17,6 +17,7 @@ using Base.StackTraces: StackFrame
 
 using PProf.ProtoBuf
 using PProf.OrderedCollections
+using CodecZlib
 
 using ProgressMeter
 
@@ -206,8 +207,11 @@ function pprof(alloc_profile::Profile.Allocs.AllocResults = Profile.Allocs.fetch
     )
 
     # Write to disk
-    open(out, "w") do io
+    io = GzipCompressorStream(open(out, "w"))
+    try
         ProtoBuf.encode(ProtoBuf.ProtoEncoder(io), prof)
+    finally
+        close(io)
     end
 
     if web

diff --git a/src/PProf.jl b/src/PProf.jl
@@ -5,6 +5,7 @@ export pprof, @pprof
 using Profile
 using ProtoBuf
 using OrderedCollections
+using CodecZlib
 import pprof_jll
 
 using Profile: clear
@@ -260,8 +261,11 @@ function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
     )
 
     # Write to disk
-    open(out, "w") do io
+    io = GzipCompressorStream(open(out, "w"))
+    try
         ProtoBuf.encode(ProtoBuf.ProtoEncoder(io), prof)
+    finally
+        close(io)
     end
 
     if web

diff --git a/src/flamegraphs.jl b/src/flamegraphs.jl
@@ -199,8 +199,11 @@ function pprof(fg::Node{NodeData},
     )
 
     # Write to disk
-    open(out, "w") do io
+    io = GzipCompressorStream(open(out, "w"))
+    try
         ProtoBuf.encode(ProtoBuf.ProtoEncoder(io), prof)
+    finally
+        close(io)
     end
 
     if web

diff --git a/test/Allocs.jl b/test/Allocs.jl
@@ -3,6 +3,7 @@ module PProfAllocsTest
 import PProf
 import Profile
 using ProtoBuf
+using CodecZlib
 
 using Test
 
@@ -20,7 +21,12 @@ const out = tempname()
     outf = PProf.Allocs.pprof(out=out, web=false)
 
     # Read the exported profile
-    prof = open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), outf, "r")
+    io = GzipDecompressorStream(open(outf, "r"))
+    prof = try
+        decode(ProtoDecoder(io), PProf.perftools.profiles.Profile)
+    finally
+        close(io)
+    end
 
     # Verify that we exported stack trace samples:
     @test length(prof.sample) > 0
@@ -43,7 +49,7 @@ end
     outf = PProf.Allocs.pprof(out=out, web=false)
 
     # Read the exported profile
-    prof = open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), outf, "r")
+    prof = open(io->decode(ProtoDecoder(GzipDecompressorStream(io)), PProf.perftools.profiles.Profile), outf, "r")
 
     # Test for both functions:
     @test in("foo(::Float64)", prof.string_table)

diff --git a/test/PProf.jl b/test/PProf.jl
@@ -5,6 +5,7 @@ using PProf
 using Test
 using Profile
 using ProtoBuf
+using CodecZlib
 
 const out = tempname() * ".pb.gz"
 
@@ -42,7 +43,12 @@ end
     outf = pprof(out=out, web=false)
 
     # Read the exported profile
-    prof = open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), outf, "r")
+    io = GzipDecompressorStream(open(outf, "r"))
+    prof = try
+        decode(ProtoDecoder(io), PProf.perftools.profiles.Profile)
+    finally
+        close(io)
+    end
 
     # Verify that we exported stack trace samples:
     @test length(prof.sample) > 0
@@ -56,7 +62,7 @@ end
 
 function load_prof_proto(file)
     @show file
-    open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), file, "r")
+    open(io->decode(ProtoDecoder(GzipDecompressorStream(io)), PProf.perftools.profiles.Profile), file, "r")
 end
 
 @testset "with_c" begin

diff --git a/test/flamegraphs.jl b/test/flamegraphs.jl
@@ -5,6 +5,7 @@ using PProf
 using Test
 using Profile
 using ProtoBuf
+using CodecZlib
 
 # Test interactions with FlameGraphs package
 using FlameGraphs
@@ -43,7 +44,12 @@ end
     outf = pprof(fg, out=out, web=false)
 
     # Read the exported profile
-    fg_prof = open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), outf, "r")
+    io = GzipDecompressorStream(open(outf, "r"))
+    fg_prof = try
+        decode(ProtoDecoder(io), PProf.perftools.profiles.Profile)
+    finally
+        close(io)
+    end
 
     # Verify that we exported stack trace samples:
     @test length(fg_prof.sample) > 0
@@ -55,7 +61,12 @@ end
 
 function load_prof_proto(file)
     @show file
-    open(io->decode(ProtoDecoder(io), PProf.perftools.profiles.Profile), file, "r")
+    io = GzipDecompressorStream(open(file, "r"))
+    fg_prof = try
+        decode(ProtoDecoder(io), PProf.perftools.profiles.Profile)
+    finally
+        close(io)
+    end
 end
 
 @testset "with_c" begin