Safer and more efficient email sends (forem#22784)

Author: benhalpern

app/models/email.rb

@@ -3,7 +3,7 @@ class Email < ApplicationRecord
   belongs_to :user_query, optional: true
   has_many :email_messages
 
-  after_save :deliver_to_users
+  after_commit :deliver_to_users, on: [:create, :update]
 
   validates :subject, presence: true
   validates :body, presence: true
@@ -91,7 +91,7 @@ def deliver_to_test_emails(addresses_string)
 
   def deliver_to_users
     return if type_of == "onboarding_drip"
-    return if status != "active"
+    return unless saved_change_to_status? && active?
 
     Emails::EnqueueCustomBatchSendWorker.perform_async(id)
     update_columns(status: "delivered")

app/services/user_query_executor.rb

@@ -49,6 +49,41 @@ def execute
     execute_safe_query
   end
 
+  def each_id_batch(batch_size: 1000)
+    return unless block_given?
+    return [] unless valid?
+
+    ReadOnlyDatabaseService.with_connection do |conn|
+      setup_execution_environment(conn)
+      final_query = user_query.substitute_variables(variables)
+      safe_query = build_safe_query(final_query)
+
+      begin
+        result = execute_with_timeout(conn, safe_query)
+        return unless result.is_a?(PG::Result)
+
+        current_batch = []
+        result.each do |row|
+          user_id = row["id"] || row["user_id"] || row["users.id"]
+          next unless user_id
+
+          current_batch << user_id.to_i
+
+          if current_batch.size >= batch_size
+            yield current_batch
+            current_batch = []
+          end
+        end
+
+        yield current_batch if current_batch.any?
+
+        update_execution_tracking
+      rescue StandardError => e
+        handle_execution_error(e)
+      end
+    end
+  end
+
   def test_execute(limit: MAX_TEST_USER_LIMIT)
     @limit = limit
     execute

app/workers/emails/enqueue_custom_batch_send_worker.rb

@@ -5,124 +5,111 @@ class EnqueueCustomBatchSendWorker
 
     sidekiq_options queue: :medium_priority, retry: 15
 
-    BATCH_SIZE = Rails.env.production? ? 500 : 10
+    BATCH_SIZE = Rails.env.production? ? 1000 : 10 # Increased batch size is safe with pluck
 
-    def perform(email_id)
-      # 1) Remember the old ENV timeout (milliseconds)
-      original_timeout = ENV.fetch("STATEMENT_TIMEOUT", 10_000).to_i
+    def perform(email_id, min_id = nil, max_id = nil)
+      # 1) Define timeout locally (do not touch ENV)
+      custom_timeout = 10_000 
 
-      # 2) Check out exactly one connection for the entire block
-      # Use read-only database if available, otherwise fall back to main database
+      # 2) Open the connection block
       ReadOnlyDatabaseService.with_connection do |conn|
-        # 3) Ensure subsequent PG SETs inherit that same original timeout
-        #    (in case some middleware or initializer reads ENV["STATEMENT_TIMEOUT"])
-        ENV["STATEMENT_TIMEOUT"] = original_timeout.to_s
-        conn.execute("SET statement_timeout TO #{original_timeout}")
+        
+        # SAFE: Set timeout only on this current connection thread
+        conn.execute("SET statement_timeout TO #{custom_timeout}")
 
         email = Email.find_by(id: email_id)
         return unless email
 
-        user_scope = if email.user_query.present?
-                       # Use custom user query for targeting
-                       begin
-                         # Extract variables from email if available
-                         variables = extract_email_variables(email)
-                         executor = UserQueryExecutor.new(email.user_query, variables: variables)
-                         # Apply the same filtering as other paths to ensure proper notification settings and roles
-                         executor.execute
-                           .registered
-                           .joins(:notification_setting)
-                           .without_role(:suspended)
-                           .without_role(:spam)
-                           .where(notification_setting: { email_newsletter: true })
-                           .where.not(email: "")
-                       rescue StandardError => e
-                         Rails.logger.error("UserQuery execution failed for email #{email.id}: #{e.message}")
-                         User.none
-                       end
-                     elsif email.audience_segment
-                       email.audience_segment.users
-                         .registered
-                         .joins(:notification_setting)
-                         .without_role(:suspended)
-                         .without_role(:spam)
-                         .where(notification_setting: { email_newsletter: true })
-                         .where.not(email: "")
-                     else
-                       User.registered
-                         .joins(:notification_setting)
-                         .without_role(:suspended)
-                         .without_role(:spam)
-                         .where(notification_setting: { email_newsletter: true })
-                         .where.not(email: "")
-                     end
-
-        # 4) Now open a transaction. Everything inside here uses "SET LOCAL statement_timeout = 0"
-        conn.transaction do
-          # This sets statement_timeout=0 for *this* transaction only.
-          # No matter how many queries ActiveRecord fires inside this block,
-          # they all see infinite timeout.
-          conn.execute("SET LOCAL statement_timeout TO 0")
-
-          # 5) Run your batches inside the same transaction, so every SELECT is "no timeout"
-          batch_count = 0
-          total_users = 0
-
-          user_scope.find_in_batches(batch_size: BATCH_SIZE) do |users_batch|
-            batch_count += 1
-
-            # Skip empty batches
-            next if users_batch.empty?
-
-            # Validate we have valid user IDs
-            user_ids = users_batch.map(&:id).compact
-            next if user_ids.empty?
-
-            total_users += user_ids.size
-
-            Rails.logger.info("Processing email batch #{batch_count} for email #{email.id}: #{user_ids.size} users (first ID: #{user_ids.first})")
-
-            begin
-              Emails::BatchCustomSendWorker.perform_async(
-                user_ids,
-                email.subject,
-                email.body,
-                email.type_of,
-                email.id,
-              )
-            rescue StandardError => e
-              Rails.logger.error("Failed to enqueue batch #{batch_count} for email #{email.id}: #{e.message}")
-              # Continue processing other batches even if one fails
-              next
-            end
-          end
-
-          Rails.logger.info("Completed email processing for email #{email.id}: #{batch_count} batches, #{total_users} total users")
+        if email.user_query.present?
+          process_custom_query(email, min_id, max_id)
+        else
+          process_standard_scope(email, min_id, max_id)
         end
-      # As soon as this transaction block ends, Postgres automatically reverts
-      # statement_timeout back to the previous session value (original_timeout).
-      ensure
-        # 6) Just to be safe, set ENV back and explicitly restore on the connection
-        ENV["STATEMENT_TIMEOUT"] = original_timeout.to_s
-        conn.execute("SET statement_timeout TO #{ENV.fetch('STATEMENT_TIMEOUT', 10_000)}")
+        
+        # Connection is returned to pool automatically; 
+        # Ideally, the pool creates new connections or resets them, 
+        # but resetting the timeout here is good practice if connections are reused immediately.
+        conn.execute("RESET statement_timeout") 
       end
     end
-  end
-
-  private
 
-  def extract_email_variables(email)
-    variables = {}
+    private
+
+    def process_custom_query(email, min_id = nil, max_id = nil)
+      variables = extract_email_variables(email)
+      executor = UserQueryExecutor.new(email.user_query, variables: variables)
+      
+      executor.each_id_batch(batch_size: 1000) do |id_batch|
+        # Optional: Apply ID range filtering in Ruby for custom query results
+        filtered_ids = id_batch
+        filtered_ids = filtered_ids.select { |id| id >= min_id.to_i } if min_id
+        filtered_ids = filtered_ids.select { |id| id <= max_id.to_i } if max_id
+
+        next if filtered_ids.empty?
+
+        # Load ONLY ids, don't instantiate User objects if possible
+        filtered_user_ids = User.where(id: filtered_ids)
+                                .registered
+                                .joins(:notification_setting)
+                                .without_role(:suspended)
+                                .without_role(:spam)
+                                .where(notification_setting: { email_newsletter: true })
+                                .where.not(email: "")
+                                .pluck(:id)
+
+        enqueue_batch(email, filtered_user_ids, "Custom Query")
+      end
+    end
 
-    # Extract variables from email's variables field if it exists
-    if email.respond_to?(:variables) && email.variables.present?
-      begin
-        variables.merge!(JSON.parse(email.variables))
-      rescue JSON::ParserError
-        Rails.logger.warn("Invalid variables JSON in email #{email.id}")
+    def process_standard_scope(email, min_id = nil, max_id = nil)
+      # Build the relation
+      base_scope = User.registered
+                       .joins(:notification_setting)
+                       .without_role(:suspended)
+                       .without_role(:spam)
+                       .where(notification_setting: { email_newsletter: true })
+                       .where.not(email: "")
+
+      # Apply ID range filtering at the DB level for standard scopes
+      base_scope = base_scope.where("users.id >= ?", min_id) if min_id
+      base_scope = base_scope.where("users.id <= ?", max_id) if max_id
+
+      user_scope = if email.audience_segment
+                     base_scope.merge(email.audience_segment.users)
+                   else
+                     base_scope
+                   end
+
+      # OPTIMIZED: Use in_batches + pluck to avoid loading User models
+      user_scope.in_batches(of: BATCH_SIZE) do |relation|
+        # relation is an ActiveRecord::Relation for the batch (e.g. "WHERE id BETWEEN 1 and 1000")
+        # .pluck(:id) executes "SELECT id FROM ..." directly
+        user_ids = relation.pluck(:id)
+        
+        enqueue_batch(email, user_ids, "Segment/Default")
       end
     end
 
-    variables
+    def enqueue_batch(email, user_ids, source_label)
+      return if user_ids.empty?
+
+      Rails.logger.info("Processing #{source_label} batch for email #{email.id}: #{user_ids.size} users")
+
+      Emails::BatchCustomSendWorker.perform_async(
+        user_ids,
+        email.subject,
+        email.body,
+        email.type_of,
+        email.id,
+      )
+    end
+
+    def extract_email_variables(email)
+      return {} unless email.respond_to?(:variables) && email.variables.present?
+      JSON.parse(email.variables)
+    rescue JSON::ParserError
+      Rails.logger.warn("Invalid variables JSON in email #{email.id}")
+      {}
+    end
   end
-end
+end

spec/models/email_spec.rb

@@ -7,10 +7,11 @@
   end
 
   describe "Callbacks" do
-    it "calls #deliver_to_users after create" do
-      email = build(:email)
-      expect(email).to receive(:deliver_to_users)
-      email.save
+    it "registers #deliver_to_users as an after_commit callback" do
+      # Verify the callback is registered in the chain
+      # Note: checking private internal Rails structure is brittle but confirms configuration
+      callback_names = Email._commit_callbacks.select { |cb| cb.kind == :after }.map(&:filter)
+      expect(callback_names).to include(:deliver_to_users)
     end
   end
 
@@ -27,7 +28,8 @@
 
       it "does not enqueue any jobs to EnqueueCustomBatchSendWorker" do
         expect(Emails::EnqueueCustomBatchSendWorker).not_to receive(:perform_async)
-        email.send(:deliver_to_users)
+        # Manually trigger since after_commit doesn't run in transactional tests
+        email.deliver_to_users
       end
     end
 
@@ -36,7 +38,7 @@
 
       it "does not enqueue any jobs to EnqueueCustomBatchSendWorker" do
         expect(Emails::EnqueueCustomBatchSendWorker).not_to receive(:perform_async)
-        email.send(:deliver_to_users)
+        email.deliver_to_users
       end
     end
 
@@ -45,18 +47,28 @@
 
       it "enqueues jobs to EnqueueCustomBatchSendWorker" do
         email.update(status: "active")
+        email.deliver_to_users
         expect(Emails::EnqueueCustomBatchSendWorker).to have_received(:perform_async).with(email.id)
       end
 
       it "only enqueues once even if re-saved" do
         email.update(status: "active")
+        email.deliver_to_users
+        expect(Emails::EnqueueCustomBatchSendWorker).to have_received(:perform_async).with(email.id)
+        
+        # Clear expectations
+        RSpec::Mocks.space.proxy_for(Emails::EnqueueCustomBatchSendWorker).reset
+        allow(Emails::EnqueueCustomBatchSendWorker).to receive(:perform_async)
+
         email.reload.save
-        expect(Emails::EnqueueCustomBatchSendWorker).to have_received(:perform_async).with(email.id).once
+        email.deliver_to_users
+        expect(Emails::EnqueueCustomBatchSendWorker).not_to have_received(:perform_async)
       end
     end
 
     it "updates the email status to 'delivered'" do
-      email = create(:email, status: "active")
+      email = create(:email, status: "draft") # Start as draft
+      email.update(status: "active") # Make it active/dirty
       email.deliver_to_users
       expect(email.reload.status).to eq("delivered")
     end