PERF: Enqueue Job::BackfillBadge in Jobs::BadgeGrant (#30945)

This commit updates the `Jobs::BadgeGrant` scheduled job to enqueue on
`Job::BackfillBadge` regular job for each enabled badge on the site.

The rationale for this change is that we started seeing the
`Jobs::BadgeGrant` job taking hours on sites with lots of enabled badges
as well as users because the job was backfilling all enabled badges
serially within the job. This is bad as it means that a `mini_scheduler`
thread is tied up
by this job thus reducing the overall capacity of `mini_scheduler` for
hours.

app/controllers/user_badges_controller.rb

@@ -144,7 +144,7 @@ class UserBadgesController < ApplicationController
     UserBadge.where(user_id: user_badge.user_id, badge_id: user_badge.badge_id).update_all(
       is_favorite: !user_badge.is_favorite,
     )
-    UserBadge.update_featured_ranks!(user_badge.user_id)
+    UserBadge.update_featured_ranks!([user_badge.user_id])
   end
 
   private

app/jobs/regular/backfill_badge.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Jobs
+  class BackfillBadge < ::Jobs::Base
+    sidekiq_options queue: "low"
+
+    def execute(args)
+      return unless SiteSetting.enable_badges
+
+      badge = Badge.enabled.find_by(id: args[:badge_id])
+      return unless badge
+
+      revoked_user_ids = Set.new
+      granted_user_ids = Set.new
+
+      BadgeGranter.backfill(
+        badge,
+        revoked_callback: ->(user_ids) { revoked_user_ids.merge(user_ids) },
+        granted_callback: ->(user_ids) { granted_user_ids.merge(user_ids) },
+      )
+
+      affected_user_ids = (revoked_user_ids | granted_user_ids).to_a
+
+      BadgeGranter.revoke_ungranted_titles!(revoked_user_ids.to_a)
+      UserBadge.ensure_consistency!(affected_user_ids)
+      UserStat.update_distinct_badge_count(affected_user_ids)
+    end
+  end
+end

app/jobs/scheduled/badge_grant.rb

@@ -10,22 +10,7 @@ module Jobs
 
     def execute(args)
       return unless SiteSetting.enable_badges
-
-      Badge.enabled.find_each do |b|
-        begin
-          BadgeGranter.backfill(b)
-        rescue => ex
-          # TODO - expose errors in UI
-          Discourse.handle_job_exception(
-            ex,
-            error_context({}, code_desc: "Exception granting badges", extra: { badge_id: b.id }),
-          )
-        end
-      end
-
-      BadgeGranter.revoke_ungranted_titles!
-      UserBadge.ensure_consistency! # Badge granter sometimes uses raw SQL, so hooks do not run. Clean up data
-      UserStat.update_distinct_badge_count
+      Badge.enabled.pluck(:id).each { |badge_id| Jobs.enqueue(:backfill_badge, badge_id: badge_id) }
     end
   end
 end

app/jobs/scheduled/ensure_badges_consistency.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Jobs
+  class EnsureBadgesConsistency < ::Jobs::Scheduled
+    every 1.day
+
+    def execute(args)
+      return unless SiteSetting.enable_badges
+
+      BadgeGranter.revoke_ungranted_titles!
+      UserBadge.ensure_consistency!
+      UserStat.update_distinct_badge_count
+    end
+  end
+end

app/models/user_badge.rb

@@ -56,32 +56,36 @@ class UserBadge < ActiveRecord::Base
 
   after_create do
     Badge.increment_counter "grant_count", self.badge_id
-    UserStat.update_distinct_badge_count self.user_id
-    UserBadge.update_featured_ranks! self.user_id
+    user_ids = [self.user_id]
+    UserStat.update_distinct_badge_count(user_ids)
+    UserBadge.update_featured_ranks!(user_ids)
     self.trigger_user_badge_granted_event
   end
 
   after_destroy do
     Badge.decrement_counter "grant_count", self.badge_id
-    UserStat.update_distinct_badge_count self.user_id
-    UserBadge.update_featured_ranks! self.user_id
+    user_ids = [self.user_id]
+    UserStat.update_distinct_badge_count(user_ids)
+    UserBadge.update_featured_ranks!(user_ids)
 
     DiscourseEvent.trigger(:user_badge_removed, self.badge_id, self.user_id)
     DiscourseEvent.trigger(:user_badge_revoked, user_badge: self)
   end
 
-  def self.ensure_consistency!
-    self.update_featured_ranks!
+  def self.ensure_consistency!(user_ids = nil)
+    self.update_featured_ranks!(user_ids)
   end
 
-  def self.update_featured_ranks!(user_id = nil)
+  def self.update_featured_ranks!(user_ids = nil)
+    user_ids = user_ids.join(", ") if user_ids
+
     query = <<~SQL
       WITH featured_tl_badge AS -- Find the best trust level badge for each user
       (
         SELECT user_id, max(badge_id) as badge_id
         FROM user_badges
         WHERE badge_id IN (1,2,3,4)
-        #{"AND user_id = #{user_id.to_i}" if user_id}
+        #{"AND user_id IN (#{user_ids})" if user_ids}
         GROUP BY user_id
       ),
       ranks AS ( -- Take all user badges, group by user_id and badge_id, and calculate a rank for each one
@@ -101,7 +105,7 @@ class UserBadge < ActiveRecord::Base
         FROM user_badges
         INNER JOIN badges ON badges.id = user_badges.badge_id
         LEFT JOIN featured_tl_badge ON featured_tl_badge.user_id = user_badges.user_id AND featured_tl_badge.badge_id = user_badges.badge_id
-        #{"WHERE user_badges.user_id = #{user_id.to_i}" if user_id}
+        #{"WHERE user_badges.user_id IN (#{user_ids})" if user_ids}
         GROUP BY user_badges.user_id, user_badges.badge_id
       )
       -- Now use that data to update the featured_rank column

app/models/user_stat.rb

@@ -192,7 +192,9 @@ class UserStat < ActiveRecord::Base
     SQL
   end
 
-  def self.update_distinct_badge_count(user_id = nil)
+  def self.update_distinct_badge_count(user_ids = nil)
+    user_ids = user_ids.join(", ") if user_ids
+
     sql = <<~SQL
       UPDATE user_stats
       SET distinct_badge_count = x.distinct_badge_count
@@ -204,15 +206,14 @@ class UserStat < ActiveRecord::Base
         GROUP BY users.id
       ) x
       WHERE user_stats.user_id = x.user_id AND user_stats.distinct_badge_count <> x.distinct_badge_count
+      #{"AND user_stats.user_id IN (#{user_ids})" if user_ids}
     SQL
 
-    sql = sql + " AND user_stats.user_id = #{user_id.to_i}" if user_id
-
     DB.exec sql
   end
 
   def update_distinct_badge_count
-    self.class.update_distinct_badge_count(self.user_id)
+    self.class.update_distinct_badge_count([self.user_id])
   end
 
   def self.update_draft_count(user_id = nil)

app/services/badge_granter.rb

@@ -110,7 +110,7 @@ class BadgeGranter
         WHERE notification_id IS NULL AND user_id = :user_id AND badge_id = :badge_id
       SQL
 
-      UserBadge.update_featured_ranks!(user.id)
+      UserBadge.update_featured_ranks!([user.id])
     end
   end
 
@@ -404,8 +404,9 @@ class BadgeGranter
     post_clause = badge.target_posts ? "AND (q.post_id = ub.post_id OR NOT :multiple_grant)" : ""
     post_id_field = badge.target_posts ? "q.post_id" : "NULL"
 
-    sql = <<~SQL
-      DELETE FROM user_badges
+    if badge.auto_revoke && full_backfill
+      sql = <<~SQL
+        DELETE FROM user_badges
         WHERE id IN (
           SELECT ub.id
           FROM user_badges ub
@@ -415,17 +416,22 @@ class BadgeGranter
           #{post_clause}
           WHERE ub.badge_id = :id AND q.user_id IS NULL
         )
-    SQL
+        RETURNING user_badges.user_id
+      SQL
 
-    if badge.auto_revoke && full_backfill
-      DB.exec(
-        sql,
-        id: badge.id,
-        post_ids: [-1],
-        user_ids: [-2],
-        backfill: true,
-        multiple_grant: true, # cheat here, cause we only run on backfill and are deleting
-      )
+      rows =
+        DB.query(
+          sql,
+          id: badge.id,
+          post_ids: [-1],
+          user_ids: [-2],
+          backfill: true,
+          multiple_grant: true, # cheat here, cause we only run on backfill and are deleting
+        )
+
+      if (revoked_callback = opts&.dig(:revoked_callback)) && rows.size > 0
+        revoked_callback.call(rows.map { |r| r.user_id })
+      end
     end
 
     sql = <<~SQL
@@ -465,34 +471,41 @@ class BadgeGranter
       return
     end
 
-    builder
-      .query(
+    rows =
+      builder.query(
         id: badge.id,
         multiple_grant: badge.multiple_grant,
         backfill: full_backfill,
         post_ids: post_ids || [-2],
         user_ids: user_ids || [-2],
       )
-      .each do |row|
-        next if suppress_notification?(badge, row.granted_at, row.skip_new_user_tips)
-        next if row.staff && badge.awarded_for_trust_level?
 
-        notification = send_notification(row.user_id, row.username, row.locale, badge)
-        UserBadge.trigger_user_badge_granted_event(badge.id, row.user_id)
+    if (granted_callback = opts&.dig(:granted_callback)) && rows.size > 0
+      granted_callback.call(rows.map { |r| r.user_id })
+    end
 
-        DB.exec(
-          "UPDATE user_badges SET notification_id = :notification_id WHERE id = :id",
-          notification_id: notification.id,
-          id: row.id,
-        )
-      end
+    rows.each do |row|
+      next if suppress_notification?(badge, row.granted_at, row.skip_new_user_tips)
+      next if row.staff && badge.awarded_for_trust_level?
+
+      notification = send_notification(row.user_id, row.username, row.locale, badge)
+      UserBadge.trigger_user_badge_granted_event(badge.id, row.user_id)
+
+      DB.exec(
+        "UPDATE user_badges SET notification_id = :notification_id WHERE id = :id",
+        notification_id: notification.id,
+        id: row.id,
+      )
+    end
 
     badge.reset_grant_count!
   rescue => e
     raise GrantError, "Failed to backfill '#{badge.name}' badge: #{opts}. Reason: #{e.message}"
   end
 
-  def self.revoke_ungranted_titles!
+  def self.revoke_ungranted_titles!(user_ids = nil)
+    user_ids = user_ids.join(", ") if user_ids
+
     DB.exec <<~SQL
       UPDATE users u
       SET title = ''
@@ -509,6 +522,7 @@ class BadgeGranter
             AND b.allow_title
             AND b.enabled
         )
+        #{user_ids.present? ? "AND u.id IN (:user_ids)" : ""}
     SQL
 
     DB.exec <<~SQL
@@ -518,6 +532,7 @@ class BadgeGranter
       WHERE up.user_id = u.id
         AND (u.title IS NULL OR u.title = '')
         AND up.granted_title_badge_id IS NOT NULL
+        #{user_ids.present? ? "AND up.user_id IN (:user_ids)" : ""}
     SQL
   end
 

spec/models/user_badge_spec.rb

@@ -121,14 +121,14 @@ RSpec.describe UserBadge do
     it "can ensure consistency per user" do
       user_badge_tl2.update_column(:featured_rank, 20) # Update without hooks
       expect(user_badge_tl2.reload.featured_rank).to eq(20) # Double check
-      UserBadge.update_featured_ranks! user.id
+      UserBadge.update_featured_ranks!([user.id])
       expect(user_badge_tl2.reload.featured_rank).to eq(1)
     end
 
     it "can ensure consistency for all users" do
       user_badge_tl2.update_column(:featured_rank, 20) # Update without hooks
       expect(user_badge_tl2.reload.featured_rank).to eq(20) # Double check
-      UserBadge.update_featured_ranks!
+      UserBadge.update_featured_ranks!([user.id])
       expect(user_badge_tl2.reload.featured_rank).to eq(1)
     end
   end

spec/services/badge_granter_spec.rb

@@ -115,7 +115,7 @@ RSpec.describe BadgeGranter do
     end
   end
 
-  describe "backfill" do
+  describe ".backfill" do
     it "has no broken badge queries" do
       Badge.all.each { |b| BadgeGranter.backfill(b) }
     end
@@ -211,6 +211,46 @@ RSpec.describe BadgeGranter do
 
       expect(UserBadge.where(user_id: user_id).count).to eq(0)
     end
+
+    it "auto revokes badges from users when badge is set to auto revoke and user no longer satisfy the badge's query" do
+      user.update!(username: "cool_username")
+
+      badge_for_having_cool_username =
+        Fabricate(
+          :badge,
+          query:
+            "SELECT users.id user_id, CURRENT_TIMESTAMP granted_at FROM users WHERE users.username = 'cool_username'",
+          auto_revoke: true,
+        )
+
+      granted_user_ids = []
+
+      BadgeGranter.backfill(
+        badge_for_having_cool_username,
+        granted_callback: ->(user_ids) { granted_user_ids.concat(user_ids) },
+      )
+
+      expect(granted_user_ids).to eq([user.id])
+
+      expect(
+        UserBadge.exists?(user_id: user.id, badge_id: badge_for_having_cool_username.id),
+      ).to eq(true)
+
+      user.update!(username: "not_cool_username")
+
+      revoked_user_ids = []
+
+      BadgeGranter.backfill(
+        badge_for_having_cool_username,
+        revoked_callback: ->(user_ids) { revoked_user_ids.concat(user_ids) },
+      )
+
+      expect(revoked_user_ids).to eq([user.id])
+
+      expect(
+        UserBadge.exists?(user_id: user.id, badge_id: badge_for_having_cool_username.id),
+      ).to eq(false)
+    end
   end
 
   describe "grant" do