From e935f013583f1df76c88bbb3fc6ae408ee5ba740 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Tue, 15 Nov 2022 17:34:59 -0600
Subject: [PATCH] uploads: fix temp files not being cleaned up quickly enough.

Fix temp files generated during the upload process not being cleaned up quickly enough. This included
downloaded files, generated preview images, and Ugoira video conversions.

Before we relied on `Tempfile` cleaning up files automatically. But this only happened when the
Tempfile object was garbage collected, which could take a long time. In the meantime we could have
hundreds of megabytes of temp files hanging around.

The fix is to explicitly close temp files when we're done with them. But the standard `Tempfile`
class doesn't immediately delete the file when it's closed. So we also have to introduce a
Danbooru::Tempfile wrapper that deletes the tempfile as soon as it's closed.
---
 app/logical/bigquery_export_service.rb        | 11 +++----
 app/logical/danbooru/http.rb                  |  4 +--
 app/logical/danbooru/tempfile.rb              | 14 +++++++++
 .../discord_slash_command/tagme_command.rb    |  2 ++
 app/logical/ffmpeg.rb                         |  2 +-
 app/logical/iqdb_client.rb                    |  2 ++
 app/logical/media_file/image.rb               | 29 ++++++++++++-------
 app/logical/media_file/ugoira.rb              | 11 +++----
 app/logical/media_file/video.rb               |  9 ++++--
 app/logical/post_replacement_processor.rb     |  2 ++
 app/logical/storage_manager/rclone.rb         |  2 +-
 app/models/media_asset.rb                     |  2 ++
 app/models/upload_media_asset.rb              |  2 ++
 test/unit/storage_manager_test.rb             |  2 +-
 14 files changed, 67 insertions(+), 27 deletions(-)
 create mode 100644 app/logical/danbooru/tempfile.rb

diff --git a/app/logical/bigquery_export_service.rb b/app/logical/bigquery_export_service.rb
index 6a0801bd6..6d208a1c1 100644
--- a/app/logical/bigquery_export_service.rb
+++ b/app/logical/bigquery_export_service.rb
@@ -55,20 +55,21 @@ class BigqueryExportService
 
     file = dump_records!
     upload_to_bigquery!(file)
+  ensure
+    file&.close
   end
 
   # Dump the table's records to a gzipped, newline-delimited JSON tempfile.
-  def dump_records!
-    file = Tempfile.new("danbooru-export-dump-", binmode: true)
-    file = Zlib::GzipWriter.new(file)
+  def dump_records!(file = Danbooru::Tempfile.new("danbooru-export-dump-#{model.name}-", binmode: true))
+    gzip = Zlib::GzipWriter.new(file)
 
     CurrentUser.scoped(User.anonymous) do
       records.find_each(batch_size: 5_000) do |record|
-        file.puts(record.to_json)
+        gzip.puts(record.to_json)
       end
     end
 
-    file.close # flush zlib footer
+    gzip.finish
     file
   end
 
diff --git a/app/logical/danbooru/http.rb b/app/logical/danbooru/http.rb
index 849cc38f7..7f7fbb59c 100644
--- a/app/logical/danbooru/http.rb
+++ b/app/logical/danbooru/http.rb
@@ -171,11 +171,11 @@ module Danbooru
       # Download a file from `url` and return a {MediaFile}.
       #
       # @param url [String] the URL to download
-      # @param file [Tempfile] the file to download the URL to
+      # @param file [Danbooru::Tempfile] the file to download the URL to
       # @raise [DownloadError] if the server returns a non-200 OK response
       # @raise [FileTooLargeError] if the file exceeds Danbooru's maximum download size.
       # @return [Array<(HTTP::Response, MediaFile)>] the HTTP response and the downloaded file
-      def download_media(url, file: Tempfile.new("danbooru-download-", binmode: true))
+      def download_media(url, file: Danbooru::Tempfile.new("danbooru-download-#{url.parameterize.truncate(96)}-", binmode: true))
         response = get(url)
 
         raise DownloadError, "#{url} failed with code #{response.status}" if response.status != 200
diff --git a/app/logical/danbooru/tempfile.rb b/app/logical/danbooru/tempfile.rb
new file mode 100644
index 000000000..75891a418
--- /dev/null
+++ b/app/logical/danbooru/tempfile.rb
@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+# Like Tempfile, but delete the tempfile when it's closed.
+#
+# The Tempfile class in the standard library doesn't delete the file immediately when you call `file.close`. Instead you
+# have to call `file.close!` or `file.unlink` to delete the file, or wait until the object gets garbage collected, which
+# can take a long time. This makes it so that Tempfiles are cleaned up immediately on close.
+module Danbooru
+  class Tempfile < ::Tempfile
+    def close(unlink = true)
+      super
+    end
+  end
+end
diff --git a/app/logical/discord_slash_command/tagme_command.rb b/app/logical/discord_slash_command/tagme_command.rb
index 2f610ed9c..b2839a135 100644
--- a/app/logical/discord_slash_command/tagme_command.rb
+++ b/app/logical/discord_slash_command/tagme_command.rb
@@ -52,6 +52,8 @@ class DiscordSlashCommand
           },
         }]
       }
+    ensure
+      preview&.close
     end
 
     def get_last_message_with_url(limit: 10)
diff --git a/app/logical/ffmpeg.rb b/app/logical/ffmpeg.rb
index 0fe3dc3ea..7a3c9c3ab 100644
--- a/app/logical/ffmpeg.rb
+++ b/app/logical/ffmpeg.rb
@@ -22,7 +22,7 @@ class FFmpeg
   #
   # @return [MediaFile] the preview image
   def smart_video_preview
-    vp = Tempfile.new(["video-preview", ".png"], binmode: true)
+    vp = Danbooru::Tempfile.new(["danbooru-video-preview-#{file.md5}-", ".png"], binmode: true)
 
     # https://ffmpeg.org/ffmpeg.html#Main-options
     # https://ffmpeg.org/ffmpeg-filters.html#thumbnail
diff --git a/app/logical/iqdb_client.rb b/app/logical/iqdb_client.rb
index e17adc780..fd007dc2a 100644
--- a/app/logical/iqdb_client.rb
+++ b/app/logical/iqdb_client.rb
@@ -101,6 +101,8 @@ class IqdbClient
       preview = media_file.preview!(Danbooru.config.small_image_width, Danbooru.config.small_image_width)
       file = HTTP::FormData::File.new(preview)
       request(:post, "query", form: { file: file }, params: { limit: limit })
+    ensure
+      preview&.close
     end
 
     # Add a post to IQDB.
diff --git a/app/logical/media_file/image.rb b/app/logical/media_file/image.rb
index 1b41b0f90..9c9cc46c8 100644
--- a/app/logical/media_file/image.rb
+++ b/app/logical/media_file/image.rb
@@ -7,6 +7,11 @@
 class MediaFile::Image < MediaFile
   delegate :thumbnail_image, to: :image
 
+  def close
+    super
+    @preview_frame&.close unless @preview_frame == self
+  end
+
   def dimensions
     image.size
   rescue Vips::Error
@@ -107,7 +112,7 @@ class MediaFile::Image < MediaFile
       resized_image = resized_image.flatten(background: 255)
     end
 
-    output_file = Tempfile.new(["image-preview-#{md5}", ".#{format.to_s}"])
+    output_file = Danbooru::Tempfile.new(["danbooru-image-preview-#{md5}-", ".#{format.to_s}"])
     case format.to_sym
     when :jpeg
       # https://www.libvips.org/API/current/VipsForeignSave.html#vips-jpegsave
@@ -130,14 +135,6 @@ class MediaFile::Image < MediaFile
     preview_frame.resize!(w, h, size: :force, **options)
   end
 
-  def preview_frame
-    if is_animated?
-      FFmpeg.new(self).smart_video_preview
-    else
-      self
-    end
-  end
-
   def is_animated?
     frame_count.to_i > 1
   end
@@ -166,6 +163,8 @@ class MediaFile::Image < MediaFile
     false
   end
 
+  private
+
   # @return [Vips::Image] the Vips image object for the file
   def image
     Vips::Image.new_from_file(file.path, fail: false).autorot
@@ -175,5 +174,15 @@ class MediaFile::Image < MediaFile
     FFmpeg.new(self)
   end
 
-  memoize :image, :video, :preview_frame, :dimensions, :error, :metadata, :is_corrupt?, :is_animated_gif?, :is_animated_png?
+  def preview_frame
+    @preview_frame ||= begin
+      if is_animated?
+        FFmpeg.new(self).smart_video_preview
+      else
+        self
+      end
+    end
+  end
+
+  memoize :image, :video, :dimensions, :error, :metadata, :is_corrupt?, :is_animated_gif?, :is_animated_png?
 end
diff --git a/app/logical/media_file/ugoira.rb b/app/logical/media_file/ugoira.rb
index bf7f10c05..dc68df42e 100644
--- a/app/logical/media_file/ugoira.rb
+++ b/app/logical/media_file/ugoira.rb
@@ -17,8 +17,9 @@ class MediaFile::Ugoira < MediaFile
   end
 
   def close
-    file.close
-    preview_frame.close
+    super
+    @preview_frame&.close
+    # XXX should clean up `convert` too
   end
 
   def metadata
@@ -52,7 +53,7 @@ class MediaFile::Ugoira < MediaFile
     raise RuntimeError, "can't convert ugoira to webm: no ugoira frame data was provided" unless frame_delays.present?
 
     Danbooru::Archive.extract!(file) do |tmpdir, filenames|
-      output_file = Tempfile.new(["ugoira-conversion", ".webm"], binmode: true)
+      output_file = Danbooru::Tempfile.new(["danbooru-ugoira-conversion-#{md5}-", ".webm"], binmode: true)
 
       # Duplicate last frame to avoid it being displayed only for a very short amount of time.
       last_file_name = File.basename(filenames.last)
@@ -90,8 +91,8 @@ class MediaFile::Ugoira < MediaFile
   private
 
   def preview_frame
-    FFmpeg.new(convert).smart_video_preview
+    @preview_frame ||= FFmpeg.new(convert).smart_video_preview
   end
 
-  memoize :preview_frame, :dimensions, :convert, :metadata
+  memoize :dimensions, :convert, :metadata
 end
diff --git a/app/logical/media_file/video.rb b/app/logical/media_file/video.rb
index 615c88302..997ad8897 100644
--- a/app/logical/media_file/video.rb
+++ b/app/logical/media_file/video.rb
@@ -10,6 +10,11 @@ class MediaFile::Video < MediaFile
     :audio_stream, :audio_streams, :silence_duration, :silence_percentage, :average_loudness,
     :peak_loudness, :loudness_range, :error, to: :video
 
+  def close
+    super
+    @preview_frame&.close
+  end
+
   def dimensions
     [video.width, video.height]
   end
@@ -65,8 +70,8 @@ class MediaFile::Video < MediaFile
   end
 
   def preview_frame
-    video.smart_video_preview
+    @preview_frame ||= video.smart_video_preview
   end
 
-  memoize :video, :preview_frame, :dimensions, :metadata, :duration, :has_audio?
+  memoize :video, :dimensions, :metadata, :duration, :has_audio?
 end
diff --git a/app/logical/post_replacement_processor.rb b/app/logical/post_replacement_processor.rb
index 3891f2256..3816b2e8d 100644
--- a/app/logical/post_replacement_processor.rb
+++ b/app/logical/post_replacement_processor.rb
@@ -53,6 +53,8 @@ class PostReplacementProcessor
   rescue Exception => exception
     replacement.errors.add(:base, exception.message)
     raise ActiveRecord::Rollback
+  ensure
+    media_file&.close
   end
 
   def rescale_notes(post)
diff --git a/app/logical/storage_manager/rclone.rb b/app/logical/storage_manager/rclone.rb
index de2ffc249..c9eb40d4f 100644
--- a/app/logical/storage_manager/rclone.rb
+++ b/app/logical/storage_manager/rclone.rb
@@ -27,7 +27,7 @@ class StorageManager::Rclone < StorageManager
   end
 
   def open(path)
-    file = Tempfile.new(binmode: true)
+    file = Danbooru::Tempfile.new(binmode: true)
     rclone "copyto", key(path), file.path
     file
   end
diff --git a/app/models/media_asset.rb b/app/models/media_asset.rb
index a09576efd..0e6682515 100644
--- a/app/models/media_asset.rb
+++ b/app/models/media_asset.rb
@@ -68,6 +68,8 @@ class MediaAsset < ApplicationRecord
       file = convert_file(original_file)
       storage_service.store(file, file_path)
       backup_storage_service.store(file, file_path)
+    ensure
+      file&.close
     end
 
     def trash_file!
diff --git a/app/models/upload_media_asset.rb b/app/models/upload_media_asset.rb
index a2e6dfb91..6da34154b 100644
--- a/app/models/upload_media_asset.rb
+++ b/app/models/upload_media_asset.rb
@@ -112,6 +112,8 @@ class UploadMediaAsset < ApplicationRecord
     update!(status: :active)
   rescue Exception => e
     update!(status: :failed, error: e.message)
+  ensure
+    media_file&.close
   end
 
   def update_upload_status
diff --git a/test/unit/storage_manager_test.rb b/test/unit/storage_manager_test.rb
index 1dfce4e9e..794cbaf82 100644
--- a/test/unit/storage_manager_test.rb
+++ b/test/unit/storage_manager_test.rb
@@ -2,7 +2,7 @@ require 'test_helper'
 
 class StorageManagerTest < ActiveSupport::TestCase
   def tempfile(data)
-    file = Tempfile.new
+    file = Danbooru::Tempfile.new
     file.write(data)
     file.flush
     file