Fix N+1s in grouped notifications (#31638)

app/controllers/api/v2_alpha/notifications_controller.rb

@@ -13,7 +13,6 @@ class Api::V2Alpha::NotificationsController < Api::BaseController
   def index
     with_read_replica do
       @notifications = load_notifications
-      @group_metadata = load_group_metadata
       @grouped_notifications = load_grouped_notifications
       @relationships = StatusRelationshipsPresenter.new(target_statuses_from_notifications, current_user&.account_id)
       @presenter = GroupedNotificationsPresenter.new(@grouped_notifications, expand_accounts: expand_accounts_param)
@@ -34,7 +33,7 @@ class Api::V2Alpha::NotificationsController < Api::BaseController
         'app.notification_grouping.expand_accounts_param' => expand_accounts_param
       )
 
-      render json: @presenter, serializer: REST::DedupNotificationGroupSerializer, relationships: @relationships, group_metadata: @group_metadata, expand_accounts: expand_accounts_param
+      render json: @presenter, serializer: REST::DedupNotificationGroupSerializer, relationships: @relationships, expand_accounts: expand_accounts_param
     end
   end
 
@@ -48,7 +47,7 @@ class Api::V2Alpha::NotificationsController < Api::BaseController
 
   def show
     @notification = current_account.notifications.without_suspended.find_by!(group_key: params[:id])
-    presenter = GroupedNotificationsPresenter.new([NotificationGroup.from_notification(@notification)])
+    presenter = GroupedNotificationsPresenter.new(NotificationGroup.from_notifications([@notification]))
     render json: presenter, serializer: REST::DedupNotificationGroupSerializer
   end
 
@@ -77,22 +76,9 @@ class Api::V2Alpha::NotificationsController < Api::BaseController
     end
   end
 
-  def load_group_metadata
-    return {} if @notifications.empty?
-
-    MastodonOTELTracer.in_span('Api::V2Alpha::NotificationsController#load_group_metadata') do
-      browserable_account_notifications
-        .where(group_key: @notifications.filter_map(&:group_key))
-        .where(id: (@notifications.last.id)..(@notifications.first.id))
-        .group(:group_key)
-        .pluck(:group_key, 'min(notifications.id) as min_id', 'max(notifications.id) as max_id', 'max(notifications.created_at) as latest_notification_at')
-        .to_h { |group_key, min_id, max_id, latest_notification_at| [group_key, { min_id: min_id, max_id: max_id, latest_notification_at: latest_notification_at }] }
-    end
-  end
-
   def load_grouped_notifications
     MastodonOTELTracer.in_span('Api::V2Alpha::NotificationsController#load_grouped_notifications') do
-      @notifications.map { |notification| NotificationGroup.from_notification(notification, max_id: @group_metadata.dig(notification.group_key, :max_id), grouped_types: params[:grouped_types]) }
+      NotificationGroup.from_notifications(@notifications, pagination_range: (@notifications.last.id)..(@notifications.first.id), grouped_types: params[:grouped_types])
     end
   end
 

app/models/notification_group.rb

@@ -1,38 +1,49 @@
 # frozen_string_literal: true
 
 class NotificationGroup < ActiveModelSerializers::Model
-  attributes :group_key, :sample_accounts, :notifications_count, :notification, :most_recent_notification_id
+  attributes :group_key, :sample_accounts, :notifications_count, :notification, :most_recent_notification_id, :pagination_data
 
   # Try to keep this consistent with `app/javascript/mastodon/models/notification_group.ts`
   SAMPLE_ACCOUNTS_SIZE = 8
 
-  def self.from_notification(notification, max_id: nil, grouped_types: nil)
+  def self.from_notifications(notifications, pagination_range: nil, grouped_types: nil)
+    return [] if notifications.empty?
+
     grouped_types = grouped_types.presence&.map(&:to_sym) || Notification::GROUPABLE_NOTIFICATION_TYPES
-    groupable = notification.group_key.present? && grouped_types.include?(notification.type)
 
-    if groupable
-      # TODO: caching, and, if caching, preloading
-      scope = notification.account.notifications.where(group_key: notification.group_key)
-      scope = scope.where(id: ..max_id) if max_id.present?
+    grouped_notifications = notifications.filter { |notification| notification.group_key.present? && grouped_types.include?(notification.type) }
+    group_keys = grouped_notifications.pluck(:group_key)
 
-      # Ideally, we would not load accounts for each notification group
-      most_recent_notifications = scope.order(id: :desc).includes(:from_account).take(SAMPLE_ACCOUNTS_SIZE)
-      most_recent_id = most_recent_notifications.first.id
-      sample_accounts = most_recent_notifications.map(&:from_account)
-      notifications_count = scope.count
-    else
-      most_recent_id = notification.id
-      sample_accounts = [notification.from_account]
-      notifications_count = 1
-    end
+    groups_data = load_groups_data(notifications.first.account_id, group_keys, pagination_range: pagination_range)
+    accounts_map = Account.where(id: groups_data.values.pluck(1).flatten).index_by(&:id)
+
+    notifications.map do |notification|
+      if notification.group_key.present? && grouped_types.include?(notification.type)
+        most_recent_notification_id, sample_account_ids, count, *raw_pagination_data = groups_data[notification.group_key]
+
+        pagination_data = raw_pagination_data.empty? ? nil : { min_id: raw_pagination_data[0], latest_notification_at: raw_pagination_data[1] }
 
         NotificationGroup.new(
           notification: notification,
-      group_key: groupable ? notification.group_key : "ungrouped-#{notification.id}",
-      sample_accounts: sample_accounts,
-      notifications_count: notifications_count,
-      most_recent_notification_id: most_recent_id
+          group_key: notification.group_key,
+          sample_accounts: sample_account_ids.map { |id| accounts_map[id] },
+          notifications_count: count,
+          most_recent_notification_id: most_recent_notification_id,
+          pagination_data: pagination_data
         )
+      else
+        pagination_data = pagination_range.blank? ? nil : { min_id: notification.id, latest_notification_at: notification.created_at }
+
+        NotificationGroup.new(
+          notification: notification,
+          group_key: "ungrouped-#{notification.id}",
+          sample_accounts: [notification.from_account],
+          notifications_count: 1,
+          most_recent_notification_id: notification.id,
+          pagination_data: pagination_data
+        )
+      end
+    end
   end
 
   delegate :type,
@@ -41,4 +52,50 @@ class NotificationGroup < ActiveModelSerializers::Model
            :account_relationship_severance_event,
            :account_warning,
            to: :notification, prefix: false
+
+  class << self
+    private
+
+    def load_groups_data(account_id, group_keys, pagination_range: nil)
+      return {} if group_keys.empty?
+
+      if pagination_range.present?
+        binds = [
+          account_id,
+          SAMPLE_ACCOUNTS_SIZE,
+          pagination_range.begin,
+          pagination_range.end,
+          ActiveRecord::Relation::QueryAttribute.new('group_keys', group_keys, ActiveRecord::ConnectionAdapters::PostgreSQL::OID::Array.new(ActiveModel::Type::String.new)),
+        ]
+
+        ActiveRecord::Base.connection.select_all(<<~SQL.squish, 'grouped_notifications', binds).cast_values.to_h { |k, *values| [k, values] }
+          SELECT
+            groups.group_key,
+            (SELECT id FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key AND id <= $4 ORDER BY id DESC LIMIT 1),
+            array(SELECT from_account_id FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key AND id <= $4 ORDER BY id DESC LIMIT $2),
+            (SELECT count(*) FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key AND id <= $4) AS notifications_count,
+            (SELECT id FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key AND id >= $3 ORDER BY id ASC LIMIT 1) AS min_id,
+            (SELECT created_at FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key AND id <= $4 ORDER BY id DESC LIMIT 1)
+          FROM
+            unnest($5::text[]) AS groups(group_key);
+        SQL
+      else
+        binds = [
+          account_id,
+          SAMPLE_ACCOUNTS_SIZE,
+          ActiveRecord::Relation::QueryAttribute.new('group_keys', group_keys, ActiveRecord::ConnectionAdapters::PostgreSQL::OID::Array.new(ActiveModel::Type::String.new)),
+        ]
+
+        ActiveRecord::Base.connection.select_all(<<~SQL.squish, 'grouped_notifications', binds).cast_values.to_h { |k, *values| [k, values] }
+          SELECT
+            groups.group_key,
+            (SELECT id FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key ORDER BY id DESC LIMIT 1),
+            array(SELECT from_account_id FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key ORDER BY id DESC LIMIT $2),
+            (SELECT count(*) FROM notifications WHERE notifications.account_id = $1 AND notifications.group_key = groups.group_key) AS notifications_count
+          FROM
+            unnest($3::text[]) AS groups(group_key);
+        SQL
+      end
+    end
+  end
 end

app/serializers/rest/notification_group_serializer.rb

@@ -39,21 +39,18 @@ class REST::NotificationGroupSerializer < ActiveModel::Serializer
   end
 
   def page_min_id
-    range = instance_options[:group_metadata][object.group_key]
-    range.present? ? range[:min_id].to_s : object.notification.id.to_s
+    object.pagination_data[:min_id].to_s
   end
 
   def page_max_id
-    range = instance_options[:group_metadata][object.group_key]
-    range.present? ? range[:max_id].to_s : object.notification.id.to_s
+    object.most_recent_notification_id.to_s
   end
 
   def latest_page_notification_at
-    range = instance_options[:group_metadata][object.group_key]
-    range.present? ? range[:latest_notification_at] : object.notification.created_at
+    object.pagination_data[:latest_notification_at]
   end
 
   def paginated?
-    !instance_options[:group_metadata].nil?
+    object.pagination_data.present?
   end
 end