mktemp/dir: delete temp files on exit by default (#32851)

StefanKarpinski · web-flow · commit 57a022f74ca7 · 2019-08-14T14:22:28.000-04:00
The main change here is that the non-higher-order versions of mktemp/dir will
now by default put the temp files/dirs that they create into a temp cleanup list
which will be deleted upon process exit. Instead of putting lots of atexit
handlers in the list, this uses a single handler and keeps a list of temporary
paths to cleanup. If the list gets long (&gt; 1024) it will be scanned for
already-deleted paths when adding a new path to the cleanup list. To avoid
running too frequently, it doubles the "too many paths" threshold after this.

If cleanup fails for higher order `mktemp() do ... end` forms, they will also
put their temp files/dirs in the cleanup list but marked for "asap" deletion.
Every time the above mechanism for cleaning out the temp cleanup list scans for
already-deleted temp files, this will try again to delete these temp files.
diff --git a/NEWS.md b/NEWS.md
@@ -63,6 +63,7 @@ Standard library changes
 * `nothing` can now be `print`ed, and interplated into strings etc. as the string `"nothing"`. It is still not permitted to be interplated into Cmds (i.e. ``echo `$(nothing)` `` will still error without running anything.) ([#32148])
 * When `open` is called with a function, command, and keyword argument (e.g. ```open(`ls`, read=true) do f ...```)
   it now correctly throws a `ProcessFailedException` like other similar calls ([#32193]).
+* `mktemp` and `mktempdir` now try, by default, to remove temporary paths they create before the process exits ([#32851]).
 
 #### Libdl
 
diff --git a/base/file.jl b/base/file.jl
@@ -448,6 +448,34 @@ function tempdir()
     end
 end
 
+const TEMP_CLEANUP_MIN = Ref(1024)
+const TEMP_CLEANUP_MAX = Ref(1024)
+const TEMP_CLEANUP = Dict{String,Bool}()
+const TEMP_CLEANUP_LOCK = ReentrantLock()
+
+function temp_cleanup_later(path::AbstractString; asap::Bool=false)
+    lock(TEMP_CLEANUP_LOCK)
+    TEMP_CLEANUP[path] = asap
+    if length(TEMP_CLEANUP) > TEMP_CLEANUP_MAX[]
+        temp_cleanup_purge(false)
+        TEMP_CLEANUP_MAX[] = max(TEMP_CLEANUP_MIN[], 2*length(TEMP_CLEANUP))
+    end
+    unlock(TEMP_CLEANUP_LOCK)
+    return nothing
+end
+
+function temp_cleanup_purge(all::Bool=true)
+    need_gc = Sys.iswindows()
+    for (path, asap) in TEMP_CLEANUP
+        if (all || asap) && ispath(path)
+            need_gc && GC.gc(true)
+            need_gc = false
+            rm(path, recursive=true, force=true)
+        end
+        !ispath(path) && delete!(TEMP_CLEANUP, path)
+    end
+end
+
 const temp_prefix = "jl_"
 
 if Sys.iswindows()
@@ -466,8 +494,9 @@ function _win_tempname(temppath::AbstractString, uunique::UInt32)
     return transcode(String, tname)
 end
 
-function mktemp(parent=tempdir())
+function mktemp(parent::AbstractString=tempdir(); cleanup::Bool=true)
     filename = _win_tempname(parent, UInt32(0))
+    cleanup && temp_cleanup_later(filename)
     return (filename, Base.open(filename, "r+"))
 end
 
@@ -498,10 +527,11 @@ function tempname()
 end
 
 # Create and return the name of a temporary file along with an IOStream
-function mktemp(parent=tempdir())
+function mktemp(parent::AbstractString=tempdir(); cleanup::Bool=true)
     b = joinpath(parent, temp_prefix * "XXXXXX")
     p = ccall(:mkstemp, Int32, (Cstring,), b) # modifies b
     systemerror(:mktemp, p == -1)
+    cleanup && temp_cleanup_later(b)
     return (b, fdio(p, true))
 end
 
@@ -524,22 +554,25 @@ created. The path is likely to be unique, but this cannot be guaranteed.
 tempname()
 
 """
-    mktemp(parent=tempdir())
+    mktemp(parent=tempdir(); cleanup=true) -> (path, io)
 
-Return `(path, io)`, where `path` is the path of a new temporary file in `parent` and `io`
-is an open file object for this path.
+Return `(path, io)`, where `path` is the path of a new temporary file in `parent`
+and `io` is an open file object for this path. The `cleanup` option controls whether
+the temporary file is automatically deleted when the process exits.
 """
 mktemp(parent)
 
 """
-    mktempdir(parent=tempdir(); prefix=$(repr(temp_prefix)))
+    mktempdir(parent=tempdir(); prefix=$(repr(temp_prefix)), cleanup=true) -> path
 
 Create a temporary directory in the `parent` directory with a name
 constructed from the given prefix and a random suffix, and return its path.
 Additionally, any trailing `X` characters may be replaced with random characters.
-If `parent` does not exist, throw an error.
+If `parent` does not exist, throw an error. The `cleanup` option controls whether
+the temporary directory is automatically deleted when the process exits.
 """
-function mktempdir(parent=tempdir(); prefix=temp_prefix)
+function mktempdir(parent::AbstractString=tempdir();
+    prefix::AbstractString=temp_prefix, cleanup::Bool=true)
     if isempty(parent) || occursin(path_separator_re, parent[end:end])
         # append a path_separator only if parent didn't already have one
         tpath = "$(parent)$(prefix)XXXXXX"
@@ -558,6 +591,7 @@ function mktempdir(parent=tempdir(); prefix=temp_prefix)
         end
         path = unsafe_string(ccall(:jl_uv_fs_t_path, Cstring, (Ptr{Cvoid},), req))
         ccall(:uv_fs_req_cleanup, Cvoid, (Ptr{Cvoid},), req)
+        cleanup && temp_cleanup_later(path)
         return path
     finally
         Libc.free(req)
@@ -571,17 +605,18 @@ end
 Apply the function `f` to the result of [`mktemp(parent)`](@ref) and remove the
 temporary file upon completion.
 """
-function mktemp(fn::Function, parent=tempdir())
-    (tmp_path, tmp_io) = mktemp(parent)
+function mktemp(fn::Function, parent::AbstractString=tempdir())
+    (tmp_path, tmp_io) = mktemp(parent, cleanup=false)
     try
         fn(tmp_path, tmp_io)
     finally
-        # TODO: should we call GC.gc() first on error, to make it much more likely that `rm` succeeds?
         try
             close(tmp_io)
             rm(tmp_path)
         catch ex
             @error "mktemp cleanup" _group=:file exception=(ex, catch_backtrace())
+            # might be possible to remove later
+            temp_cleanup_later(tmp_path, asap=true)
         end
     end
 end
@@ -592,16 +627,18 @@ end
 Apply the function `f` to the result of [`mktempdir(parent; prefix)`](@ref) and remove the
 temporary directory all of its contents upon completion.
 """
-function mktempdir(fn::Function, parent=tempdir(); prefix=temp_prefix)
-    tmpdir = mktempdir(parent; prefix=prefix)
+function mktempdir(fn::Function, parent::AbstractString=tempdir();
+    prefix::AbstractString=temp_prefix)
+    tmpdir = mktempdir(parent; prefix=prefix, cleanup=false)
     try
         fn(tmpdir)
     finally
-        # TODO: should we call GC.gc() first on error, to make it much more likely that `rm` succeeds?
         try
             rm(tmpdir, recursive=true)
         catch ex
             @error "mktempdir cleanup" _group=:file exception=(ex, catch_backtrace())
+            # might be possible to remove later
+            temp_cleanup_later(tmpdir, asap=true)
         end
     end
 end
diff --git a/base/initdefs.jl b/base/initdefs.jl
@@ -293,7 +293,7 @@ end
 
 ## atexit: register exit hooks ##
 
-const atexit_hooks = []
+const atexit_hooks = Callable[Filesystem.temp_cleanup_purge]
 
 """
     atexit(f)
diff --git a/doc/src/base/file.md b/doc/src/base/file.md
@@ -29,10 +29,10 @@ Base.Filesystem.rm
 Base.Filesystem.touch
 Base.Filesystem.tempname
 Base.Filesystem.tempdir
-Base.Filesystem.mktemp(::Any)
-Base.Filesystem.mktemp(::Function, ::Any)
-Base.Filesystem.mktempdir(::Any)
-Base.Filesystem.mktempdir(::Function, ::Any)
+Base.Filesystem.mktemp(::AbstractString)
+Base.Filesystem.mktemp(::Function, ::AbstractString)
+Base.Filesystem.mktempdir(::AbstractString)
+Base.Filesystem.mktempdir(::Function, ::AbstractString)
 Base.Filesystem.isblockdev
 Base.Filesystem.ischardev
 Base.Filesystem.isdir
diff --git a/test/file.jl b/test/file.jl
@@ -48,6 +48,193 @@ if !Sys.iswindows()
     cd(pwd_)
 end
 
+child_eval(code::String) = eval(Meta.parse(readchomp(`$(Base.julia_cmd()) -E $code`)))
+
+@testset "mktemp/dir basic cleanup" begin
+    # mktemp without cleanup
+    t = child_eval("t = mktemp(cleanup=false)[1]; @assert isfile(t); t")
+    @test isfile(t)
+    rm(t, force=true)
+    @test !ispath(t)
+    # mktemp with cleanup
+    t = child_eval("t = mktemp()[1]; @assert isfile(t); t")
+    @test !ispath(t)
+    # mktempdir without cleanup
+    t = child_eval("t = mktempdir(cleanup=false); touch(joinpath(t, \"file.txt\")); t")
+    @test isfile("$t/file.txt")
+    rm(t, recursive=true, force=true)
+    @test !ispath(t)
+    # mktempdir with cleanup
+    t = child_eval("t = mktempdir(); touch(joinpath(t, \"file.txt\")); t")
+    @test !ispath(t)
+end
+
+import Base.Filesystem: TEMP_CLEANUP_MIN, TEMP_CLEANUP_MAX, TEMP_CLEANUP
+
+function with_temp_cleanup(f::Function, n::Int)
+    SAVE_TEMP_CLEANUP_MIN = TEMP_CLEANUP_MIN[]
+    SAVE_TEMP_CLEANUP_MAX = TEMP_CLEANUP_MAX[]
+    SAVE_TEMP_CLEANUP = copy(TEMP_CLEANUP)
+    empty!(TEMP_CLEANUP)
+    TEMP_CLEANUP_MIN[] = n
+    TEMP_CLEANUP_MAX[] = n
+    try f()
+    finally
+        Sys.iswindows() && GC.gc(true)
+        for t in keys(TEMP_CLEANUP)
+            rm(t, recursive=true, force=true)
+        end
+        copy!(TEMP_CLEANUP, SAVE_TEMP_CLEANUP)
+        TEMP_CLEANUP_MAX[] = SAVE_TEMP_CLEANUP_MAX
+        TEMP_CLEANUP_MIN[] = SAVE_TEMP_CLEANUP_MIN
+    end
+end
+
+function mktempfile(; cleanup=true)
+    (file, io) = mktemp(cleanup=cleanup)
+    Sys.iswindows() && close(io)
+    return file
+end
+
+@testset "mktemp/dir cleanup list purging" begin
+    n = 12 # cleanup min & max
+    @assert n % 2 == n % 3 == 0 # otherwise tests won't work
+    with_temp_cleanup(n) do
+        @test length(TEMP_CLEANUP) == 0
+        @test TEMP_CLEANUP_MAX[] == n
+        # for n mktemps, no purging is triggered
+        temps = String[]
+        for i = 1:n
+            t = i % 2 == 0 ? mktempfile() : mktempdir()
+            push!(temps, t)
+            @test ispath(t)
+            @test length(TEMP_CLEANUP) == i 
+            @test TEMP_CLEANUP_MAX[] == n
+            # delete 1/3 of the temp paths
+            i % 3 == 0 && rm(t, recursive=true, force=true)
+        end
+        # without cleanup no purge is triggered
+        t = mktempdir(cleanup=false)
+        @test isdir(t)
+        @test length(TEMP_CLEANUP) == n
+        @test TEMP_CLEANUP_MAX[] == n
+        rm(t, recursive=true, force=true)
+        # purge triggered by next mktemp with cleanup
+        t = mktempfile()
+        push!(temps, t)
+        n′ = 2n÷3 + 1
+        @test 2n′ > n
+        @test isfile(t)
+        @test length(TEMP_CLEANUP) == n′
+        @test TEMP_CLEANUP_MAX[] == 2n′
+        # remove all temp files
+        for t in temps
+            rm(t, recursive=true, force=true)
+        end
+        # for n′ mktemps, no purging is triggered
+        for i = 1:n′
+            t = i % 2 == 0 ? mktempfile() : mktempdir()
+            push!(temps, t)
+            @test ispath(t)
+            @test length(TEMP_CLEANUP) == n′ + i
+            @test TEMP_CLEANUP_MAX[] == 2n′
+            # delete 2/3 of the temp paths
+            i % 3 != 0 && rm(t, recursive=true, force=true)
+        end
+        # without cleanup no purge is triggered
+        t = mktempfile(cleanup=false)
+        @test isfile(t)
+        @test length(TEMP_CLEANUP) == 2n′
+        @test TEMP_CLEANUP_MAX[] == 2n′
+        rm(t, force=true)
+        # purge triggered by next mktemp
+        t = mktempdir()
+        push!(temps, t)
+        n′′ = n′÷3 + 1
+        @test 2n′′ < n
+        @test isdir(t)
+        @test length(TEMP_CLEANUP) == n′′
+        @test TEMP_CLEANUP_MAX[] == n
+    end
+end
+
+no_error_logging(f::Function) =
+    Base.CoreLogging.with_logger(f, Base.CoreLogging.NullLogger())
+
+@testset "hof mktemp/dir when cleanup is prevented" begin
+    d = mktempdir()
+    with_temp_cleanup(3) do
+        @test length(TEMP_CLEANUP) == 0
+        @test TEMP_CLEANUP_MAX[] == 3
+        local t, f
+        temps = String[]
+        # mktemp is normally cleaned up on completion
+        mktemp(d) do path, _
+            @test isfile(path)
+            t = path
+        end
+        @test !ispath(t)
+        @test length(TEMP_CLEANUP) == 0
+        @test TEMP_CLEANUP_MAX[] == 3
+        # mktemp when cleanup is prevented
+        no_error_logging() do
+            mktemp(d) do path, _
+                @test isfile(path)
+                f = open(path) # make undeletable on Windows
+                chmod(d, 0o400) # make undeletable on UNIX
+                t = path
+            end
+        end
+        chmod(d, 0o700)
+        close(f)
+        @test isfile(t)
+        @test length(TEMP_CLEANUP) == 1
+        @test TEMP_CLEANUP_MAX[] == 3
+        push!(temps, t)
+        # mktempdir is normally cleaned up on completion
+        mktempdir(d) do path
+            @test isdir(path)
+            t = path
+        end
+        @test !ispath(t)
+        @test length(TEMP_CLEANUP) == 1
+        @test TEMP_CLEANUP_MAX[] == 3
+        # mktempdir when cleanup is prevented
+        no_error_logging() do
+            mktempdir(d) do path
+                @test isdir(path)
+                # make undeletable on Windows:
+                f = open(joinpath(path, "file.txt"), "w+")
+                chmod(d, 0o400) # make undeletable on UNIX
+                t = path
+            end
+        end
+        chmod(d, 0o700)
+        close(f)
+        @test isdir(t)
+        @test length(TEMP_CLEANUP) == 2
+        @test TEMP_CLEANUP_MAX[] == 3
+        push!(temps, t)
+        # make one more temp file
+        t = mktemp()[1]
+        @test isfile(t)
+        @test length(TEMP_CLEANUP) == 3
+        @test TEMP_CLEANUP_MAX[] == 3
+        # nothing has been deleted yet
+        for t in temps
+            @test ispath(t)
+        end
+        # another temp file triggers purge
+        t = mktempdir()
+        @test isdir(t)
+        @test length(TEMP_CLEANUP) == 2
+        @test TEMP_CLEANUP_MAX[] == 4
+        # now all the temps are gone
+        for t in temps
+            @test !ispath(t)
+        end
+    end
+end
 
 #######################################################################
 # This section tests some of the features of the stat-based file info #
