From 00c1aa4ce00404d4e0af02b616cd3924e9ade554 Mon Sep 17 00:00:00 2001
From: Ian Butterworth <i.r.butterworth@gmail.com>
Date: Sat, 23 Mar 2024 21:37:18 -0400
Subject: [PATCH] tidy up streamed files when not requesting streaming mode

---
 stdlib/Profile/src/Profile.jl                 |  1 +
 stdlib/Profile/src/heapsnapshot_reassemble.jl | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/stdlib/Profile/src/Profile.jl b/stdlib/Profile/src/Profile.jl
index 17e903ed373b1..062b608b25c59 100644
--- a/stdlib/Profile/src/Profile.jl
+++ b/stdlib/Profile/src/Profile.jl
@@ -1286,6 +1286,7 @@ function take_heap_snapshot(filepath::AbstractString, all_one::Bool=false; strea
         prefix = filepath
         _stream_heap_snapshot(prefix, all_one)
         Profile.HeapSnapshot.assemble_snapshot(prefix, filepath)
+        Profile.HeapSnapshot.cleanup_streamed_files(prefix)
     end
     return filepath
 end
diff --git a/stdlib/Profile/src/heapsnapshot_reassemble.jl b/stdlib/Profile/src/heapsnapshot_reassemble.jl
index d6d6e1cade4cd..27d7d12318ed1 100644
--- a/stdlib/Profile/src/heapsnapshot_reassemble.jl
+++ b/stdlib/Profile/src/heapsnapshot_reassemble.jl
@@ -204,6 +204,19 @@ function assemble_snapshot(in_prefix, io::IO)
     return nothing
 end
 
+"""
+    cleanup_streamed_files(prefix::AbstractString)
+
+Remove files streamed during `take_heap_snapshot` in streaming mode.
+"""
+function cleanup_streamed_files(prefix::AbstractString)
+    rm(string(prefix, ".metadata.json"))
+    rm(string(prefix, ".nodes"))
+    rm(string(prefix, ".edges"))
+    rm(string(prefix, ".strings"))
+    return nothing
+end
+
 function print_str_escape_json(stream::IO, s::AbstractString)
     print(stream, '"')
     for c in s