PERF: Add partial index for alert topics. (#21)

PERF: Add partial index for alert topics. (#21)

Instead of scanning all the opened topics like the previous query did, we only scan for topics inside the configured alert categories. A partial index is also used here so that we don’t have to index other topics unnecessarily.

On one of our own internal instance, this reduces the time spent on this query from ~20ms to ~4ms.

diff --git a/app/jobs/regular/create_alert_topics_index.rb b/app/jobs/regular/create_alert_topics_index.rb
new file mode 100644
index 0000000..bf2a3f1
--- /dev/null
+++ b/app/jobs/regular/create_alert_topics_index.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Jobs
+  class CreateAlertTopicsIndex < ::Jobs::Base
+    ALERT_TOPICS_INDEX_NAME = 'idx_alert_topics'
+
+    def execute(_)
+      DB.exec(<<~SQL)
+      DROP INDEX IF EXISTS #{ALERT_TOPICS_INDEX_NAME};
+      SQL
+
+      if (category_ids = Topic.alerts_category_ids).present?
+        DB.exec(<<~SQL)
+        CREATE INDEX #{Rails.env.test? ? '' : 'CONCURRENTLY'} #{ALERT_TOPICS_INDEX_NAME}
+        ON topics (id) WHERE deleted_at IS NULL AND NOT closed AND category_id IN (#{category_ids.join(",")})
+        SQL
+      end
+    end
+  end
+end
diff --git a/db/migrate/20210608034155_add_alert_topics_index.rb b/db/migrate/20210608034155_add_alert_topics_index.rb
new file mode 100644
index 0000000..eaed873
--- /dev/null
+++ b/db/migrate/20210608034155_add_alert_topics_index.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+class AddAlertTopicsIndex < ActiveRecord::Migration[6.1]
+  disable_ddl_transaction!
+
+  def up
+    category_ids = DB.query_single(<<~SQL)
+    SELECT
+      value::json->'category_id'
+    FROM plugin_store_rows
+    WHERE plugin_name = 'discourse-prometheus-alert-receiver'
+    SQL
+
+    if category_ids.present?
+      DB.exec(<<~SQL)
+      CREATE INDEX CONCURRENTLY idx_alert_topics
+      ON topics (id)
+      WHERE deleted_at IS NULL
+      AND NOT closed
+      AND category_id IN (#{category_ids.join(",")})
+      SQL
+    end
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end
diff --git a/plugin.rb b/plugin.rb
index 9b3420d..1ee284c 100644
--- a/plugin.rb
+++ b/plugin.rb
@@ -30,6 +30,7 @@ after_initialize do
     '../app/jobs/concerns/alert_post_mixin.rb',
     '../app/jobs/regular/process_alert.rb',
     '../app/jobs/regular/process_grouped_alerts.rb',
+    '../app/jobs/regular/create_alert_topics_index.rb',
   ].each { |path| load File.expand_path(path, __FILE__) }
 
   unless Rails.env.test?
@@ -66,15 +67,36 @@ after_initialize do
     Topic.has_many :alert_receiver_alerts, dependent: :delete_all
   end
 
+  self.add_model_callback('PluginStoreRow', :after_commit) do
+    if self.plugin_name == ::DiscoursePrometheusAlertReceiver::PLUGIN_NAME
+      Topic.clear_alerts_category_ids_cache
+    end
+  end
+
+  add_class_method(:topic, :clear_alerts_category_ids_cache) do
+    @alerts_category_ids = nil
+    Jobs.enqueue(:create_alert_topics_index)
+  end
+
+  add_class_method(:topic, :alerts_category_ids_cache) do
+    @alerts_category_ids ||= alerts_category_ids
+  end
+
+  add_class_method(:topic, :alerts_category_ids) do
+    category_ids = PluginStoreRow
+      .where(plugin_name: ::DiscoursePrometheusAlertReceiver::PLUGIN_NAME)
+      .pluck("value::json->'category_id'")
+  end
+
   add_class_method(:topic, :open_alerts) do
     joins(:alert_receiver_alerts)
-      .where("not closed")
+      .where("not topics.closed AND topics.category_id IN (?)", alerts_category_ids_cache)
   end
 
   add_class_method(:topic, :firing_alerts) do
     joins(:alert_receiver_alerts)
       .where("alert_receiver_alerts.status": 'firing')
-      .where("not closed")
+      .where("not topics.closed AND topics.category_id IN (?)", alerts_category_ids_cache)
   end
 
   add_to_class(:user, :include_alert_counts?) do
diff --git a/spec/model/topic_spec.rb b/spec/model/topic_spec.rb
new file mode 100644
index 0000000..4660dff
--- /dev/null
+++ b/spec/model/topic_spec.rb
@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+describe Topic do
+  fab!(:topic) { Fabricate(:topic) }
+  fab!(:category) { topic.category }
+  fab!(:closed_topic) { Fabricate(:topic, category: category, closed: true) }
+
+  fab!(:plugin_store_row) do
+    PluginStore.set(
+      ::DiscoursePrometheusAlertReceiver::PLUGIN_NAME,
+      'sometoken',
+      { category_id: category.id }
+    )
+  end
+
+  fab!(:firing_alert) do
+    AlertReceiverAlert.create!(
+      topic: topic,
+      status: 'firing',
+      identifier: 'someidentifier',
+      starts_at: Time.zone.now,
+      external_url: "someurl"
+    )
+  end
+
+  fab!(:closed_alert) do
+    AlertReceiverAlert.create!(
+      topic: closed_topic,
+      status: 'resolved',
+      identifier: 'someidentifier',
+      starts_at: Time.zone.now,
+      external_url: "someurl"
+    )
+  end
+
+  describe '.open_alerts' do
+    it 'should return the right count' do
+      expect(Topic.open_alerts).to contain_exactly(firing_alert.topic)
+    end
+  end
+
+  describe '.firing_alerts' do
+    it 'should return the right topics' do
+      expect(Topic.firing_alerts).to contain_exactly(firing_alert.topic)
+    end
+  end
+end

GitHub sha: ce5e52ba317381fa516da1a08ea095940224457c

This commit appears in #21 which was approved by SamSaffron. It was merged by tgxworld.