From 820dd9b5a6e76c74eba7b63f56cb96545a4f6909 Mon Sep 17 00:00:00 2001
From: Jamie Gaskins <jamie@forem.com>
Date: Fri, 15 Oct 2021 11:53:55 -0400
Subject: [PATCH] Allow `Emails::RemoveOldEmailsWorker` to run longer (#15085)

* Allow Emails::RemoveOldEmailsWorker to run longer

We set the SQL statement timeout based on the STATEMENT_TIMEOUT
environment variable (defaulting to 2.5 seconds in production), but not
all SQL queries are created equal. Since some may take longer out of
necessity, this PR introduces a `Model.with_statement_timeout` method
that allows you to change the SQL statement timeout only for a given
block.

* EmailMessage doesn't inherit ApplicationRecord

It inherits from `Ahoy::Message`, which inherits directly from
`ActiveRecord::Base`.

[3] pry(main)> EmailMessage.ancestors
=> [EmailMessage(id: integer, clicked_at: datetime, content: text, feedback_message_id: integer, mailer: string, sent_at: datetime, subject: text, to: text, token: string, user_id: integer, user_type: string, utm_campaign: string, utm_content: string, utm_medium: string, utm_source: string, utm_term: string),
 EmailMessage::GeneratedAssociationMethods,
 EmailMessage::GeneratedAttributeMethods,
 Ahoy::Message(id: integer, clicked_at: datetime, content: text, feedback_message_id: integer, mailer: string, sent_at: datetime, subject: text, to: text, token: string, user_id: integer, user_type: string, utm_campaign: string, utm_content: string, utm_medium: string, utm_source: string, utm_term: string),
 Kaminari::ConfigurationMethods,
 Kaminari::ActiveRecordModelExtension,
 Ahoy::Message::GeneratedAssociationMethods,
 Ahoy::Message::GeneratedAttributeMethods,
 Bullet::SaveWithBulletSupport,
 ActiveRecord::Base,
 ...

* Use milliseconds to match STATEMENT_TIMEOUT units
---
 app/models/application_record.rb              | 31 +++++++++++++++++++
 .../emails/remove_old_emails_worker.rb        |  6 +++-
 spec/models/application_record_spec.rb        | 18 +++++++++++
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index fb38ec9a9..faaf138fd 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -53,6 +53,37 @@ class ApplicationRecord < ActiveRecord::Base
     raise UninferrableDecoratorError, "Could not infer a decorator for #{called_on.class.name}."
   end
 
+  def self.statement_timeout
+    connection.execute("SELECT setting FROM pg_settings WHERE name = 'statement_timeout'")
+      .first["setting"]
+      .to_i
+      .seconds / 1000
+  end
+
+  def self.with_statement_timeout(duration)
+    # Using SET LOCAL only works inside a transaction.
+    transaction do
+      original_timeout = statement_timeout
+      milliseconds = (duration.to_f * 1000).to_i
+      connection.execute "SET LOCAL statement_timeout = #{milliseconds}"
+      yield
+    ensure
+      # Typically this value is reset after a transaction is completed (whether
+      # successfully or unsuccessfully), but a nested transaction will not reset
+      # this value after the inner transaction is complete because they are not
+      # isolated from the outer transaction the way the outer transaction is
+      # from transactions on other connections. Nested transactions are
+      # implemented simply as savepoints inside the top-level transaction. For
+      # this reason, we need to explicitly reset the statement timeout so that
+      # when the block completes the statement timeout is reset to its previous
+      # value.
+      #
+      # See: https://www.postgresql.org/docs/11/sql-savepoint.html
+      milliseconds = original_timeout.to_i * 1000
+      connection.execute "SET LOCAL statement_timeout = #{milliseconds}"
+    end
+  end
+
   def errors_as_sentence
     errors.full_messages.to_sentence
   end
diff --git a/app/workers/emails/remove_old_emails_worker.rb b/app/workers/emails/remove_old_emails_worker.rb
index 04b8e8b36..61aa940ad 100644
--- a/app/workers/emails/remove_old_emails_worker.rb
+++ b/app/workers/emails/remove_old_emails_worker.rb
@@ -4,8 +4,12 @@ module Emails
 
     sidekiq_options queue: :low_priority, retry: 10
 
+    STATEMENT_TIMEOUT = ENV.fetch("STATEMENT_TIMEOUT_REMOVE_OLD_EMAILS", 10_000).to_i.seconds / 1_000.to_f
+
     def perform
-      EmailMessage.fast_destroy_old_retained_email_deliveries
+      ApplicationRecord.with_statement_timeout STATEMENT_TIMEOUT do
+        EmailMessage.fast_destroy_old_retained_email_deliveries
+      end
     end
   end
 end
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index f2e1c4ec4..bc6e2a51b 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -38,4 +38,22 @@ RSpec.describe ApplicationRecord, type: :model do
       expect(decorated_collection.first).to be_a(SponsorshipDecorator)
     end
   end
+
+  describe ".with_statement_timeout" do
+    it "sets the SQL statement timeout to the specified duration" do
+      original_timeout = described_class.statement_timeout
+
+      described_class.with_statement_timeout 1.second do
+        expect(described_class.statement_timeout).to eq 1.second
+
+        described_class.with_statement_timeout 10.seconds do
+          expect(described_class.statement_timeout).to eq 10.seconds
+        end
+
+        expect(described_class.statement_timeout).to eq 1.second
+      end
+
+      expect(described_class.statement_timeout).to eq original_timeout
+    end
+  end
 end