Clean up temp files after each queue message (#91)

TmpDir only cleans up on process exit, but Broadway processors are
long-lived. Accumulated temp files cause page cache and kernel slab
memory to grow, counted against the container's cgroup memory limit,
eventually triggering OOM kills.

Add TmpDir.cleanup/0 and call it after each message is processed.

Author: ericmj

lib/hexdocs/queue.ex

@@ -58,26 +58,30 @@ defmodule Hexdocs.Queue do
 
     case key_components(key) do
       {:ok, repository, package, version} ->
-        tarball_path = Hexdocs.TmpDir.tmp_file("docs-tarball")
-
-        case Hexdocs.Store.get_to_file(:repo_bucket, key, tarball_path) do
-          :ok ->
-            case Hexdocs.Tar.unpack_to_dir({:file, tarball_path},
-                   repository: repository,
-                   package: package,
-                   version: version
-                 ) do
-              {:ok, _dir, files} ->
-                update_index_sitemap(repository, key)
-                update_package_sitemap(repository, key, package, files)
-                Logger.info("#{key}: done")
-
-              {:error, reason} ->
-                Logger.error("Failed unpack #{repository}/#{package} #{version}: #{reason}")
-            end
-
-          nil ->
-            Logger.error("#{key}: package not found in store")
+        try do
+          tarball_path = Hexdocs.TmpDir.tmp_file("docs-tarball")
+
+          case Hexdocs.Store.get_to_file(:repo_bucket, key, tarball_path) do
+            :ok ->
+              case Hexdocs.Tar.unpack_to_dir({:file, tarball_path},
+                     repository: repository,
+                     package: package,
+                     version: version
+                   ) do
+                {:ok, _dir, files} ->
+                  update_index_sitemap(repository, key)
+                  update_package_sitemap(repository, key, package, files)
+                  Logger.info("#{key}: done")
+
+                {:error, reason} ->
+                  Logger.error("Failed unpack #{repository}/#{package} #{version}: #{reason}")
+              end
+
+            nil ->
+              Logger.error("#{key}: package not found in store")
+          end
+        after
+          Hexdocs.TmpDir.cleanup()
         end
 
       :error ->
@@ -114,20 +118,24 @@ defmodule Hexdocs.Queue do
           version: version
         })
 
-        tarball_path = Hexdocs.TmpDir.tmp_file("docs-tarball")
+        try do
+          tarball_path = Hexdocs.TmpDir.tmp_file("docs-tarball")
 
-        case Hexdocs.Store.get_to_file(:repo_bucket, key, tarball_path) do
-          :ok ->
-            case type do
-              :upload ->
-                process_upload(key, repository, package, version, {:file, tarball_path}, start)
+          case Hexdocs.Store.get_to_file(:repo_bucket, key, tarball_path) do
+            :ok ->
+              case type do
+                :upload ->
+                  process_upload(key, repository, package, version, {:file, tarball_path}, start)
 
-              :search ->
-                process_search(key, repository, package, version, {:file, tarball_path}, start)
-            end
+                :search ->
+                  process_search(key, repository, package, version, {:file, tarball_path}, start)
+              end
 
-          nil ->
-            Logger.error("#{log_prefix} #{key}: package not found in store")
+            nil ->
+              Logger.error("#{log_prefix} #{key}: package not found in store")
+          end
+        after
+          Hexdocs.TmpDir.cleanup()
         end
 
       :error ->
@@ -254,18 +262,22 @@ defmodule Hexdocs.Queue do
           version: version
         })
 
-        version = Version.parse!(version)
-        all_versions = all_versions(repository, package)
-        Hexdocs.Bucket.delete(repository, package, version, all_versions)
-        update_index_sitemap(repository, key)
+        try do
+          version = Version.parse!(version)
+          all_versions = all_versions(repository, package)
+          Hexdocs.Bucket.delete(repository, package, version, all_versions)
+          update_index_sitemap(repository, key)
 
-        if repository == "hexpm" do
-          Hexdocs.Search.delete(package, version)
-        end
+          if repository == "hexpm" do
+            Hexdocs.Search.delete(package, version)
+          end
 
-        elapsed = System.os_time(:millisecond) - start
-        Logger.info("FINISHED DELETING DOCS #{key} #{elapsed}ms")
-        :ok
+          elapsed = System.os_time(:millisecond) - start
+          Logger.info("FINISHED DELETING DOCS #{key} #{elapsed}ms")
+          :ok
+        after
+          Hexdocs.TmpDir.cleanup()
+        end
 
       {:ok, _repository, _package, _version} ->
         :skip

lib/hexdocs/tmp_dir.ex

@@ -32,6 +32,17 @@ defmodule Hexdocs.TmpDir do
     dir
   end
 
+  def cleanup do
+    pid = self()
+    entries = :ets.lookup(@table, pid)
+
+    Enum.each(entries, fn {_pid, path} ->
+      File.rm_rf(path)
+    end)
+
+    :ets.delete(@table, pid)
+  end
+
   def await_cleanup(pid) do
     GenServer.call(__MODULE__, {:await_cleanup, pid}, 5000)
   end

test/hexdocs/tmp_dir_test.exs

@@ -72,6 +72,37 @@ defmodule Hexdocs.TmpDirTest do
     end
   end
 
+  test "cleanup removes paths for calling process" do
+    file = Hexdocs.TmpDir.tmp_file("test")
+    dir = Hexdocs.TmpDir.tmp_dir("test")
+
+    assert File.exists?(file)
+    assert File.dir?(dir)
+
+    Hexdocs.TmpDir.cleanup()
+
+    refute File.exists?(file)
+    refute File.exists?(dir)
+  end
+
+  test "cleanup only removes paths for calling process" do
+    test_pid = self()
+
+    Task.start(fn ->
+      other_file = Hexdocs.TmpDir.tmp_file("other")
+      send(test_pid, {:other_path, other_file})
+      Process.sleep(:infinity)
+    end)
+
+    assert_receive {:other_path, other_file}
+
+    file = Hexdocs.TmpDir.tmp_file("test")
+    Hexdocs.TmpDir.cleanup()
+
+    refute File.exists?(file)
+    assert File.exists?(other_file)
+  end
+
   test "paths persist while process is alive" do
     file = Hexdocs.TmpDir.tmp_file("test")
     dir = Hexdocs.TmpDir.tmp_dir("test")