From 179efd94677d1d30bdbbbbaafc899c8c908181d2 Mon Sep 17 00:00:00 2001
From: tusooa <tusooa@kazv.moe>
Date: Sat, 24 Dec 2022 00:17:17 -0500
Subject: [PATCH] Make backup parameters configurable

---
 config/config.exs          |  4 +++-
 config/description.exs     | 15 +++++++++++++++
 lib/pleroma/user/backup.ex | 37 ++++++++++++++++++++++++++-----------
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/config/config.exs b/config/config.exs
index e41ec2f91..ecb592b9c 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -874,7 +874,9 @@
 config :pleroma, Pleroma.User.Backup,
   purge_after_days: 30,
   limit_days: 7,
-  dir: nil
+  dir: nil,
+  process_wait_time: 30_000,
+  process_chunk_size: 100
 
 config :pleroma, ConcurrentLimiter, [
   {Pleroma.Web.RichMedia.Helpers, [max_running: 5, max_waiting: 5]},
diff --git a/config/description.exs b/config/description.exs
index bf4734426..996267558 100644
--- a/config/description.exs
+++ b/config/description.exs
@@ -3394,6 +3394,21 @@
         type: :integer,
         description: "Limit user to export not more often than once per N days",
         suggestions: [7]
+      },
+      %{
+        key: :process_wait_time,
+        type: :integer,
+        label: "Process Wait Time",
+        description:
+          "The amount of time to wait for backup to report progress, in milliseconds. If no progress is received from the backup job for that much time, terminate it and deem it failed.",
+        suggestions: [30_000]
+      },
+      %{
+        key: :process_chunk_size,
+        type: :integer,
+        label: "Process Chunk Size",
+        description: "The number of activities to fetch in the backup job for each chunk.",
+        suggestions: [100]
       }
     ]
   },
diff --git a/lib/pleroma/user/backup.ex b/lib/pleroma/user/backup.ex
index cb9a40ba1..447fca2a1 100644
--- a/lib/pleroma/user/backup.ex
+++ b/lib/pleroma/user/backup.ex
@@ -35,8 +35,6 @@ defmodule Pleroma.User.Backup do
     timestamps()
   end
 
-  @report_every 100
-
   def create(user, admin_id \\ nil) do
     with :ok <- validate_limit(user, admin_id),
          {:ok, backup} <- user |> new() |> Repo.insert() do
@@ -160,6 +158,8 @@ def do_process(backup, current_pid) do
   end
 
   defp wait_backup(backup, current_processed, task) do
+    wait_time = Pleroma.Config.get([__MODULE__, :process_wait_time])
+
     receive do
       {:progress, new_processed} ->
         total_processed = current_processed + new_processed
@@ -175,6 +175,8 @@ defp wait_backup(backup, current_processed, task) do
 
           {:ok, backup} = set_state(backup, :failed)
 
+          cleanup(backup)
+
           {:error,
            %{
              backup: backup,
@@ -185,15 +187,17 @@ defp wait_backup(backup, current_processed, task) do
           {:ok, backup}
         end
     after
-      30_000 ->
+      wait_time ->
         Logger.error(
-          "Backup #{backup.id} timed out after no response for 30 seconds, terminating"
+          "Backup #{backup.id} timed out after no response for #{wait_time}ms, terminating"
         )
 
         Task.Supervisor.terminate_child(Pleroma.TaskSupervisor, task.pid)
 
         {:ok, backup} = set_state(backup, :failed)
 
+        cleanup(backup)
+
         {:error,
          %{
            backup: backup,
@@ -205,8 +209,7 @@ defp wait_backup(backup, current_processed, task) do
   @files ['actor.json', 'outbox.json', 'likes.json', 'bookmarks.json']
   def export(%__MODULE__{} = backup, caller_pid) do
     backup = Repo.preload(backup, :user)
-    name = String.trim_trailing(backup.file_name, ".zip")
-    dir = dir(name)
+    dir = backup_tempdir(backup)
 
     with :ok <- File.mkdir(dir),
          :ok <- actor(dir, backup.user, caller_pid),
@@ -264,16 +267,28 @@ defp write_header(file, name) do
     )
   end
 
-  defp should_report?(num), do: rem(num, @report_every) == 0
+  defp should_report?(num, chunk_size), do: rem(num, chunk_size) == 0
+
+  defp backup_tempdir(backup) do
+    name = String.trim_trailing(backup.file_name, ".zip")
+    dir(name)
+  end
+
+  defp cleanup(backup) do
+    dir = backup_tempdir(backup)
+    File.rm_rf(dir)
+  end
 
   defp write(query, dir, name, fun, caller_pid) do
     path = Path.join(dir, "#{name}.json")
 
+    chunk_size = Pleroma.Config.get([__MODULE__, :process_chunk_size])
+
     with {:ok, file} <- File.open(path, [:write, :utf8]),
          :ok <- write_header(file, name) do
       total =
         query
-        |> Pleroma.Repo.chunk_stream(100)
+        |> Pleroma.Repo.chunk_stream(chunk_size, _returns_as = :one, timeout: :infinity)
         |> Enum.reduce(0, fn i, acc ->
           with {:ok, data} <-
                  (try do
@@ -283,8 +298,8 @@ defp write(query, dir, name, fun, caller_pid) do
                   end),
                {:ok, str} <- Jason.encode(data),
                :ok <- IO.write(file, str <> ",\n") do
-            if should_report?(acc + 1) do
-              send(caller_pid, {:progress, @report_every})
+            if should_report?(acc + 1, chunk_size) do
+              send(caller_pid, {:progress, chunk_size})
             end
 
             acc + 1
@@ -301,7 +316,7 @@ defp write(query, dir, name, fun, caller_pid) do
           end
         end)
 
-      send(caller_pid, {:progress, rem(total, @report_every)})
+      send(caller_pid, {:progress, rem(total, chunk_size)})
 
       with :ok <- :file.pwrite(file, {:eof, -2}, "\n],\n  \"totalItems\": #{total}}") do
         File.close(file)