add lock events

app/controllers/api/submissions_controller.rb

@@ -43,7 +43,7 @@ module Api
       end
 
       if @submission.audit_trail_attachment.blank? && submitters.all?(&:completed_at?)
-        @submission.audit_trail_attachment = Submissions::GenerateAuditTrail.call(@submission)
+        @submission.audit_trail_attachment = Submissions::EnsureAuditGenerated.call(@submission)
       end
 
       render json: Submissions::SerializeForApi.call(@submission, submitters, params)

app/controllers/submissions_download_controller.rb

@@ -70,7 +70,7 @@ class SubmissionsDownloadController < ApplicationController
     return if submitter.submission.submitters.order(:completed_at).last != submitter
 
     attachment = submitter.submission.combined_document_attachment
-    attachment ||= Submissions::GenerateCombinedAttachment.call(submitter)
+    attachment ||= Submissions::EnsureCombinedGenerated.call(submitter)
 
     filename_format = AccountConfig.find_or_initialize_by(account_id: submitter.account_id,
                                                           key: AccountConfig::DOCUMENT_FILENAME_FORMAT_KEY)&.value

app/jobs/process_submitter_completion_job.rb

@@ -14,10 +14,10 @@ class ProcessSubmitterCompletionJob
 
     if is_all_completed && submitter.completed_at == submitter.submission.submitters.maximum(:completed_at)
       if submitter.submission.account.account_configs.exists?(key: AccountConfig::COMBINE_PDF_RESULT_KEY, value: true)
-        Submissions::GenerateCombinedAttachment.call(submitter)
+        Submissions::EnsureCombinedGenerated.call(submitter)
       end
 
-      Submissions::GenerateAuditTrail.call(submitter.submission)
+      Submissions::EnsureAuditGenerated.call(submitter.submission)
 
       enqueue_completed_emails(submitter)
     end

app/models/lock_event.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+# == Schema Information
+#
+# Table name: lock_events
+#
+#  id         :bigint           not null, primary key
+#  event_name :string           not null
+#  key        :string           not null
+#  created_at :datetime         not null
+#  updated_at :datetime         not null
+#
+# Indexes
+#
+#  index_lock_events_on_event_name_and_key  (event_name,key) UNIQUE WHERE ((event_name)::text = ANY ((ARRAY['start'::character varying, 'complete'::character varying])::text[]))
+#  index_lock_events_on_key                 (key)
+#
+class LockEvent < ApplicationRecord
+  enum :event_name, {
+    complete: 'complete',
+    fail: 'fail',
+    start: 'start',
+    retry: 'retry'
+  }, scope: false
+end

db/migrate/20250915060548_create_lock_events.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+class CreateLockEvents < ActiveRecord::Migration[8.0]
+  def change
+    create_table :lock_events do |t|
+      t.string :key, index: true, null: false
+      t.string :event_name, null: false
+
+      t.index %i[event_name key], unique: true, where: "event_name IN ('start', 'complete')"
+
+      t.timestamps
+    end
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.0].define(version: 2025_09_12_090605) do
+ActiveRecord::Schema[8.0].define(version: 2025_09_15_060548) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "btree_gin"
   enable_extension "plpgsql"
@@ -217,6 +217,15 @@ ActiveRecord::Schema[8.0].define(version: 2025_09_12_090605) do
     t.index ["user_id"], name: "index_encrypted_user_configs_on_user_id"
   end
 
+  create_table "lock_events", force: :cascade do |t|
+    t.string "key", null: false
+    t.string "event_name", null: false
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["event_name", "key"], name: "index_lock_events_on_event_name_and_key", unique: true, where: "((event_name)::text = ANY ((ARRAY['start'::character varying, 'complete'::character varying])::text[]))"
+    t.index ["key"], name: "index_lock_events_on_key"
+  end
+
   create_table "oauth_access_grants", force: :cascade do |t|
     t.bigint "resource_owner_id", null: false
     t.bigint "application_id", null: false

lib/submissions/ensure_audit_generated.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module Submissions
+  module EnsureAuditGenerated
+    WAIT_FOR_RETRY = 2.seconds
+    CHECK_EVENT_INTERVAL = 1.second
+    CHECK_COMPLETE_TIMEOUT = 90.seconds
+    KEY_PREFIX = 'audit_trail'
+
+    WaitForCompleteTimeout = Class.new(StandardError)
+    NotCompletedYet = Class.new(StandardError)
+
+    module_function
+
+    def call(submission)
+      return nil unless submission
+
+      raise NotCompletedYet unless submission.submitters.all?(&:completed_at?)
+
+      key = [KEY_PREFIX, submission.id].join(':')
+
+      if ApplicationRecord.uncached { LockEvent.exists?(key:, event_name: :complete) }
+        return submission.audit_trail_attachment
+      end
+
+      events = ApplicationRecord.uncached { LockEvent.where(key:).order(:id).to_a }
+
+      if events.present? && events.last.event_name.in?(%w[start retry])
+        wait_for_complete_or_fail(submission)
+      else
+        LockEvent.create!(key:, event_name: events.present? ? :retry : :start)
+
+        result = Submissions::GenerateAuditTrail.call(submission)
+
+        LockEvent.create!(key:, event_name: :complete)
+
+        result
+      end
+    rescue ActiveRecord::RecordNotUnique
+      sleep WAIT_FOR_RETRY
+
+      retry
+    rescue StandardError => e
+      Rollbar.error(e) if defined?(Rollbar)
+      Rails.logger.error(e)
+
+      LockEvent.create!(key:, event_name: :fail)
+
+      raise
+    end
+
+    def wait_for_complete_or_fail(submission)
+      total_wait_time = 0
+
+      loop do
+        sleep CHECK_EVENT_INTERVAL
+        total_wait_time += CHECK_EVENT_INTERVAL
+
+        last_event =
+          ApplicationRecord.uncached do
+            LockEvent.where(key: [KEY_PREFIX, submission.id].join(':')).order(:id).last
+          end
+
+        if last_event.event_name.in?(%w[complete fail])
+          break ApplicationRecord.uncached do
+            ActiveStorage::Attachment.find_by(record: submission, name: 'audit_trail')
+          end
+        end
+
+        raise WaitForCompleteTimeout if total_wait_time > CHECK_COMPLETE_TIMEOUT
+      end
+    end
+  end
+end

lib/submissions/ensure_combined_generated.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module Submissions
+  module EnsureCombinedGenerated
+    WAIT_FOR_RETRY = 2.seconds
+    CHECK_EVENT_INTERVAL = 1.second
+    CHECK_COMPLETE_TIMEOUT = 90.seconds
+    KEY_PREFIX = 'combined_document'
+
+    WaitForCompleteTimeout = Class.new(StandardError)
+    NotCompletedYet = Class.new(StandardError)
+
+    module_function
+
+    def call(submitter)
+      return nil unless submitter
+
+      raise NotCompletedYet unless submitter.completed_at?
+
+      key = [KEY_PREFIX, submitter.id].join(':')
+
+      if ApplicationRecord.uncached { LockEvent.exists?(key:, event_name: :complete) }
+        return submitter.submission.combined_document_attachment
+      end
+
+      events = ApplicationRecord.uncached { LockEvent.where(key:).order(:id).to_a }
+
+      if events.present? && events.last.event_name.in?(%w[start retry])
+        wait_for_complete_or_fail(submitter)
+      else
+        LockEvent.create!(key:, event_name: events.present? ? :retry : :start)
+
+        result = Submissions::GenerateCombinedAttachment.call(submitter)
+
+        LockEvent.create!(key:, event_name: :complete)
+
+        result
+      end
+    rescue ActiveRecord::RecordNotUnique
+      sleep WAIT_FOR_RETRY
+
+      retry
+    rescue StandardError => e
+      Rollbar.error(e) if defined?(Rollbar)
+      Rails.logger.error(e)
+
+      LockEvent.create!(key:, event_name: :fail)
+
+      raise
+    end
+
+    def wait_for_complete_or_fail(submitter)
+      total_wait_time = 0
+
+      loop do
+        sleep CHECK_EVENT_INTERVAL
+        total_wait_time += CHECK_EVENT_INTERVAL
+
+        last_event =
+          ApplicationRecord.uncached do
+            LockEvent.where(key: [KEY_PREFIX, submitter.id].join(':')).order(:id).last
+          end
+
+        if last_event.event_name.in?(%w[complete fail])
+          break ApplicationRecord.uncached do
+            ActiveStorage::Attachment.find_by(record: submitter.submission, name: 'combined_document')
+          end
+        end
+
+        raise WaitForCompleteTimeout if total_wait_time > CHECK_COMPLETE_TIMEOUT
+      end
+    end
+  end
+end

lib/submissions/ensure_result_generated.rb

@@ -16,21 +16,20 @@ module Submissions
 
       raise NotCompletedYet unless submitter.completed_at?
 
-      return submitter.documents if ApplicationRecord.uncached { submitter.document_generation_events.complete.exists? }
+      key = ['result_attachments', submitter.id].join(':')
 
-      events =
-        ApplicationRecord.uncached do
-          DocumentGenerationEvent.where(submitter:).order(:created_at).to_a
-        end
+      return submitter.documents if ApplicationRecord.uncached { LockEvent.exists?(key:, event_name: :complete) }
+
+      events = ApplicationRecord.uncached { LockEvent.where(key:).order(:id).to_a }
 
       if events.present? && events.last.event_name.in?(%w[start retry])
         wait_for_complete_or_fail(submitter)
       else
-        submitter.document_generation_events.create!(event_name: events.present? ? :retry : :start)
+        LockEvent.create!(key:, event_name: events.present? ? :retry : :start)
 
         documents = GenerateResultAttachments.call(submitter)
 
-        submitter.document_generation_events.create!(event_name: :complete)
+        LockEvent.create!(key:, event_name: :complete)
 
         documents
       end
@@ -42,7 +41,7 @@ module Submissions
       Rollbar.error(e) if defined?(Rollbar)
       Rails.logger.error(e)
 
-      submitter.document_generation_events.create!(event_name: :fail)
+      LockEvent.create!(key:, event_name: :fail)
 
       raise
     end
@@ -56,7 +55,7 @@ module Submissions
 
         last_event =
           ApplicationRecord.uncached do
-            DocumentGenerationEvent.where(submitter:).order(:created_at).last
+            LockEvent.where(key: ['result_attachments', submitter.id].join(':')).order(:id).last
           end
 
         break submitter.documents.reload if last_event.event_name.in?(%w[complete fail])

lib/submissions/serialize_for_api.rb

@@ -74,7 +74,7 @@ module Submissions
       if !attachment && params[:include].to_s.include?('combined_document_url')
         submitter = submitters.max_by(&:completed_at)
 
-        attachment = Submissions::GenerateCombinedAttachment.call(submitter)
+        attachment = Submissions::EnsureCombinedGenerated.call(submitter)
       end
 
       ActiveStorage::Blob.proxy_url(attachment.blob, expires_at:) if attachment