From ba1cf14c7ec9290518bf5bd3dead5d446d526c79 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Sun, 20 Nov 2022 23:34:18 -0600
Subject: [PATCH] uploads: mark uploads as failed if they're stuck processing
 for more than 4 hours.

---
 app/jobs/prune_uploads_job.rb       | 11 +++++++++++
 app/logical/danbooru_maintenance.rb |  1 +
 app/models/media_asset.rb           |  6 ++++++
 app/models/upload.rb                |  7 +++++++
 app/models/upload_media_asset.rb    |  5 +++++
 test/jobs/prune_uploads_job_test.rb | 17 +++++++++++++++++
 6 files changed, 47 insertions(+)
 create mode 100644 app/jobs/prune_uploads_job.rb
 create mode 100644 test/jobs/prune_uploads_job_test.rb

diff --git a/app/jobs/prune_uploads_job.rb b/app/jobs/prune_uploads_job.rb
new file mode 100644
index 000000000..f899a9903
--- /dev/null
+++ b/app/jobs/prune_uploads_job.rb
@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+# A job that runs hourly to mark as failed all uploads that got stuck in the 'processing' state for more than 4 hours.
+# Spawned by {DanbooruMaintenance}.
+class PruneUploadsJob < ApplicationJob
+  def perform
+    MediaAsset.prune!
+    UploadMediaAsset.prune!
+    Upload.prune!
+  end
+end
diff --git a/app/logical/danbooru_maintenance.rb b/app/logical/danbooru_maintenance.rb
index 45427a0d3..ba6a1e655 100644
--- a/app/logical/danbooru_maintenance.rb
+++ b/app/logical/danbooru_maintenance.rb
@@ -7,6 +7,7 @@ module DanbooruMaintenance
     queue PrunePostsJob
     queue PruneRateLimitsJob
     queue RegeneratePostCountsJob
+    queue PruneUploadsJob
     #queue AmcheckDatabaseJob
   end
 
diff --git a/app/models/media_asset.rb b/app/models/media_asset.rb
index 0e6682515..8d4d29d07 100644
--- a/app/models/media_asset.rb
+++ b/app/models/media_asset.rb
@@ -51,6 +51,12 @@ class MediaAsset < ApplicationRecord
 
   before_create :initialize_file_key
 
+  scope :expired, -> { processing.where(created_at: ..4.hours.ago) }
+
+  def self.prune!
+    expired.update_all(status: :failed)
+  end
+
   class Variant
     extend Memoist
     include ActiveModel::Serializers::JSON
diff --git a/app/models/upload.rb b/app/models/upload.rb
index 876f22a7b..f835ebdbf 100644
--- a/app/models/upload.rb
+++ b/app/models/upload.rb
@@ -24,6 +24,7 @@ class Upload < ApplicationRecord
 
   validates :source, format: { with: %r{\Ahttps?://}i, message: "is not a valid URL" }, if: -> { source.present? }
   validates :referer_url, format: { with: %r{\Ahttps?://}i, message: "is not a valid URL" }, if: -> { referer_url.present? }
+  validates :status, inclusion: { in: %w[pending processing completed error] }
   validate :validate_file_and_source, on: :create
   validate :validate_archive_files, on: :create
   validate :validate_uploader_is_not_limited, on: :create
@@ -31,8 +32,10 @@ class Upload < ApplicationRecord
   after_create :async_process_upload!
 
   scope :pending, -> { where(status: "pending") }
+  scope :processing, -> { where(status: "processing") }
   scope :completed, -> { where(status: "completed") }
   scope :failed, -> { where(status: "error") }
+  scope :expired, -> { processing.where(created_at: ..4.hours.ago) }
 
   def self.visible(user)
     if user.is_admin?
@@ -42,6 +45,10 @@ class Upload < ApplicationRecord
     end
   end
 
+  def self.prune!
+    expired.update_all(status: "error", error: "Stuck processing for more than 4 hours")
+  end
+
   concerning :StatusMethods do
     def is_pending?
       status == "pending"
diff --git a/app/models/upload_media_asset.rb b/app/models/upload_media_asset.rb
index 1a59a4093..dd09fccdb 100644
--- a/app/models/upload_media_asset.rb
+++ b/app/models/upload_media_asset.rb
@@ -25,6 +25,7 @@ class UploadMediaAsset < ApplicationRecord
 
   scope :unfinished, -> { where(status: %w[pending processing]) }
   scope :finished, -> { where(status: %w[active failed]) }
+  scope :expired, -> { unfinished.where(created_at: ..4.hours.ago) }
 
   def self.visible(user)
     if user.is_admin?
@@ -36,6 +37,10 @@ class UploadMediaAsset < ApplicationRecord
     end
   end
 
+  def self.prune!
+    expired.update_all(status: :failed, error: "Stuck processing for more than 4 hours")
+  end
+
   def self.search(params, current_user)
     q = search_attributes(params, [:id, :created_at, :updated_at, :status, :source_url, :page_url, :error, :upload, :media_asset, :post], current_user: current_user)
 
diff --git a/test/jobs/prune_uploads_job_test.rb b/test/jobs/prune_uploads_job_test.rb
new file mode 100644
index 000000000..0cc0c575a
--- /dev/null
+++ b/test/jobs/prune_uploads_job_test.rb
@@ -0,0 +1,17 @@
+require 'test_helper'
+
+class PruneUploadsJobTest < ActiveJob::TestCase
+  context "PruneUploadsJob" do
+    should "prune expired uploads and media assets" do
+      upload = create(:upload, created_at: 6.hours.ago, status: "processing")
+      media_asset = create(:media_asset, created_at: 6.hours.ago, status: "processing")
+      upload_media_asset = create(:upload_media_asset, created_at: 6.hours.ago, media_asset: media_asset, upload: upload, status: "processing")
+
+      PruneUploadsJob.perform_now
+
+      assert_equal("error", upload.reload.status)
+      assert_equal("failed", upload_media_asset.reload.status)
+      assert_equal("failed", media_asset.reload.status)
+    end
+  end
+end