Compress on the fly

Project.toml

@@ -1,10 +1,11 @@
 name = "PProf"
 uuid = "e4faabce-9ead-11e9-39d9-4379958e3056"
 authors = ["Valentin Churavy <[email protected]>", "Nathan Daly <[email protected]>"]
-version = "2.2.2"
+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]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/Allocs.jl

@@ -17,6 +17,7 @@ using Base.StackTraces: StackFrame
 
 using PProf.ProtoBuf
 using PProf.OrderedCollections
+using CodecZlib
 
 using ProgressMeter
 
@@ -201,8 +202,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

src/PProf.jl

@@ -5,6 +5,7 @@ export pprof, @pprof
 using Profile
 using ProtoBuf
 using OrderedCollections
+using CodecZlib
 import pprof_jll
 
 using Profile: clear
@@ -266,8 +267,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

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

test/PProf.jl

@@ -5,6 +5,7 @@ using PProf
 using Test
 using Profile
 using ProtoBuf
+using CodecZlib
 
 const out = tempname() * ".pb.gz"
 
@@ -39,7 +40,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

test/flamegraphs.jl

@@ -5,6 +5,7 @@ using PProf
 using Test
 using Profile
 using ProtoBuf
+using CodecZlib
 
 # Test interactions with FlameGraphs package
 using FlameGraphs
@@ -42,7 +43,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
@@ -54,7 +60,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
@@ -85,4 +96,4 @@ end
     @test load_prof_proto(pprof(fg, out=tempname(), web=false, keep_frames = "foo")).keep_frames != 0
 end
 
-end # module
+end # module