Clean up temp files after diff generation

lib/diff/tmp_dir.ex

@@ -32,6 +32,17 @@ defmodule Diff.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

lib/diff_web/live/diff_live_view.ex

@@ -144,40 +144,44 @@ defmodule DiffWeb.DiffLiveView do
   def handle_info({:generate_diff, package, from, to}, socket) do
     hex_impl = Application.get_env(:diff, :hex_impl, Diff.Hex)
 
-    case hex_impl.diff(package, from, to) do
-      {:ok, stream} ->
-        case process_stream_to_diffs(package, from, to, stream) do
-          {:ok, metadata, diff_ids} ->
-            initial_batch_size = 5
-            {initial_diffs, _remaining} = Enum.split(diff_ids, initial_batch_size)
-
-            socket =
-              socket
-              |> assign(
-                metadata: metadata,
-                all_diff_ids: diff_ids,
-                loaded_diffs: [],
-                loaded_diff_content: %{},
-                remaining_diffs: diff_ids,
-                generating: false,
-                has_more_diffs: length(diff_ids) > 0
-              )
-
-            send(self(), {:load_diffs, initial_diffs})
-
-            {:noreply, socket}
-
-          {:error, reason} ->
-            Logger.error("Failed to generate diff: #{inspect(reason)}")
-            {:noreply, assign(socket, error: "Failed to generate diff", generating: false)}
-        end
-
-      :error ->
-        {:noreply, assign(socket, error: "Failed to generate diff", generating: false)}
+    try do
+      case hex_impl.diff(package, from, to) do
+        {:ok, stream} ->
+          case process_stream_to_diffs(package, from, to, stream) do
+            {:ok, metadata, diff_ids} ->
+              initial_batch_size = 5
+              {initial_diffs, _remaining} = Enum.split(diff_ids, initial_batch_size)
+
+              socket =
+                socket
+                |> assign(
+                  metadata: metadata,
+                  all_diff_ids: diff_ids,
+                  loaded_diffs: [],
+                  loaded_diff_content: %{},
+                  remaining_diffs: diff_ids,
+                  generating: false,
+                  has_more_diffs: length(diff_ids) > 0
+                )
+
+              send(self(), {:load_diffs, initial_diffs})
+
+              {:noreply, socket}
+
+            {:error, reason} ->
+              Logger.error("Failed to generate diff: #{inspect(reason)}")
+              {:noreply, assign(socket, error: "Failed to generate diff", generating: false)}
+          end
+
+        :error ->
+          {:noreply, assign(socket, error: "Failed to generate diff", generating: false)}
+      end
+    catch
+      :throw, {:diff, :invalid_diff} ->
+        {:noreply, assign(socket, error: "Invalid diff", generating: false)}
+    after
+      Diff.TmpDir.cleanup()
     end
-  catch
-    :throw, {:diff, :invalid_diff} ->
-      {:noreply, assign(socket, error: "Invalid diff", generating: false)}
   end
 
   def handle_info({:load_diffs_and_update, diff_ids}, socket) do

test/diff/tmp_dir_test.exs

@@ -72,6 +72,37 @@ defmodule Diff.TmpDirTest do
     end
   end
 
+  test "cleanup removes paths for calling process" do
+    file = Diff.TmpDir.tmp_file("test")
+    dir = Diff.TmpDir.tmp_dir("test")
+
+    assert File.exists?(file)
+    assert File.dir?(dir)
+
+    Diff.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 = Diff.TmpDir.tmp_file("other")
+      send(test_pid, {:other_path, other_file})
+      Process.sleep(:infinity)
+    end)
+
+    assert_receive {:other_path, other_file}
+
+    file = Diff.TmpDir.tmp_file("test")
+    Diff.TmpDir.cleanup()
+
+    refute File.exists?(file)
+    assert File.exists?(other_file)
+  end
+
   test "paths persist while process is alive" do
     file = Diff.TmpDir.tmp_file("test")
     dir = Diff.TmpDir.tmp_dir("test")