Fix slow uncached VSA label stage query

This change alters the way we query data for label based stages. Instead
of doing a window aggregation, we use a LATERAL JOIN.

db/migrate/20200730044810_add_indexes_to_resource_label_events_to_support_vsa.rb

+# frozen_string_literal: true
+
+class AddIndexesToResourceLabelEventsToSupportVsa < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  NEW_INDEX_NAME_ON_ISSUE_ID = 'index_resource_label_events_issue_id_label_id_action'
+  OLD_INDEX_NAME_ON_ISSUE_ID = 'index_resource_label_events_on_issue_id'
+
+  NEW_INDEX_NAME_ON_MERGE_REQUEST_ID = 'index_resource_label_events_on_merge_request_id_label_id_action'
+  OLD_INDEX_NAME_ON_MERGE_REQUEST_ID = 'index_resource_label_events_on_merge_request_id'
+
+  def up
+    add_concurrent_index :resource_label_events, [:issue_id, :label_id, :action], name: NEW_INDEX_NAME_ON_ISSUE_ID
+    remove_concurrent_index_by_name :resource_label_events, OLD_INDEX_NAME_ON_ISSUE_ID
+
+    add_concurrent_index :resource_label_events, [:merge_request_id, :label_id, :action], name: NEW_INDEX_NAME_ON_MERGE_REQUEST_ID
+    remove_concurrent_index_by_name :resource_label_events, OLD_INDEX_NAME_ON_MERGE_REQUEST_ID
+  end
+
+  def down
+    add_concurrent_index :resource_label_events, :issue_id, name: OLD_INDEX_NAME_ON_ISSUE_ID
+    remove_concurrent_index_by_name(:resource_label_events, NEW_INDEX_NAME_ON_ISSUE_ID)
+
+    add_concurrent_index :resource_label_events, :merge_request_id, name: OLD_INDEX_NAME_ON_MERGE_REQUEST_ID
+    remove_concurrent_index_by_name(:resource_label_events, NEW_INDEX_NAME_ON_MERGE_REQUEST_ID)
+  end
+end

db/schema_migrations/20200730044810

+406594bc48558c3ad50680c7e1fa795f38abb92696acbb94ae2dfb13d8dcaf1a
\ No newline at end of file

db/structure.sql

@@ -20432,13 +20432,13 @@ CREATE INDEX index_requirements_on_title_trigram ON public.requirements USING gi
 
 CREATE INDEX index_requirements_on_updated_at ON public.requirements USING btree (updated_at);
 
-CREATE INDEX index_resource_label_events_on_epic_id ON public.resource_label_events USING btree (epic_id);
+CREATE INDEX index_resource_label_events_issue_id_label_id_action ON public.resource_label_events USING btree (issue_id, label_id, action);
 
-CREATE INDEX index_resource_label_events_on_issue_id ON public.resource_label_events USING btree (issue_id);
+CREATE INDEX index_resource_label_events_on_epic_id ON public.resource_label_events USING btree (epic_id);
 
 CREATE INDEX index_resource_label_events_on_label_id_and_action ON public.resource_label_events USING btree (label_id, action);
 
-CREATE INDEX index_resource_label_events_on_merge_request_id ON public.resource_label_events USING btree (merge_request_id);
+CREATE INDEX index_resource_label_events_on_merge_request_id_label_id_action ON public.resource_label_events USING btree (merge_request_id, label_id, action);
 
 CREATE INDEX index_resource_label_events_on_user_id ON public.resource_label_events USING btree (user_id);
 

ee/changelogs/unreleased/fix-slow-uncached-query-for-vsa-label-stages.yml

+---
+title: Fix slow Value Stream label stage query
+merge_request: 38252
+author:
+type: performance

ee/lib/gitlab/analytics/cycle_analytics/stage_events/label_based_stage_event.rb

@@ -32,8 +32,8 @@ module Gitlab
           # rubocop: disable CodeReuse/ActiveRecord
           def apply_query_customization(query)
             query
-              .joins("INNER JOIN (#{subquery.to_sql}) #{join_expression_name} on #{join_expression_name}.model_id = #{quote_table_name(object_type.table_name)}.id")
-              .where("#{join_expression_name}.label_assignment_order = 1")
+              .from(Arel::Nodes::Grouping.new(Arel.sql(object_type.all.to_sql)).as(object_type.table_name)) # This is needed for the LATERAL JOIN: FROM (SELECT * FROM table) as table
+              .joins("INNER JOIN LATERAL (#{subquery.to_sql}) #{join_expression_name} ON TRUE")
           end
           # rubocop: enable CodeReuse/ActiveRecord
 
@@ -58,11 +58,6 @@ module Gitlab
           # - IssueLabelAdded event: find the first assignment (add, id = 1)
           # - IssueLabelRemoved event: find the latest unassignment (remove, id = 4)
           #
-          # This can be achieved with the PARTITION window function.
-          #
-          #  - IssueLabelAdded: order by `created_at` ASC and take the row number 1
-          #  - IssueLabelRemoved: order by `created_at` DESC and take the row number 1
-          #
           #  Arguments:
           #    foreign_key: :issue_id or :merge_request_id (based on resource_label_events table)
           #    label: label model,
@@ -72,9 +67,12 @@ module Gitlab
           # rubocop: disable CodeReuse/ActiveRecord
           def resource_label_events_with_subquery(foreign_key, label, action, order)
             ResourceLabelEvent
-              .select(:created_at, resource_label_events_table[foreign_key].as('model_id'), partition_select(foreign_key, order).as('label_assignment_order'))
+              .select(:created_at)
               .where(action: action)
               .where(label_id: label.id)
+              .where(ResourceLabelEvent.arel_table[foreign_key].eq(object_type.arel_table[:id]))
+              .order(order_expression(order))
+              .limit(1)
           end
           # rubocop: enable CodeReuse/ActiveRecord
 
@@ -83,23 +81,16 @@ module Gitlab
             @join_expression_name ||= quote_table_name("#{self.class.to_s.demodulize.underscore}_#{SecureRandom.hex(5)}")
           end
 
-          # rubocop: disable CodeReuse/ActiveRecord
-          def partition_select(foreign_key, order)
-            order_expression = case order
-                               when :asc
-                                 resource_label_events_table[:created_at].asc
-                               when :desc
-                                 resource_label_events_table[:created_at].desc
-                               else
-                                 raise "unsupported order option: #{order}"
-                               end
-
-            Arel::Nodes::Over.new(
-              Arel::Nodes::NamedFunction.new('row_number', []),
-              Arel::Nodes::Window.new.partition(resource_label_events_table[foreign_key]).order(order_expression)
-            )
+          def order_expression(order)
+            case order
+            when :asc
+              resource_label_events_table[:created_at].asc
+            when :desc
+              resource_label_events_table[:created_at].desc
+            else
+              raise "unsupported order option: #{order}"
+            end
           end
-          # rubocop: enable CodeReuse/ActiveRecord
         end
       end
     end

lib/gitlab/analytics/cycle_analytics/records_fetcher.rb

@@ -90,7 +90,7 @@ module Gitlab
         end
 
         def ordered_and_limited_query
-          order_by_end_event(query).limit(MAX_RECORDS)
+          order_by_end_event(query, columns).limit(MAX_RECORDS)
         end
 
         def records

lib/gitlab/analytics/cycle_analytics/stage_query_helpers.rb

@@ -24,13 +24,13 @@ module Gitlab
         end
 
         # rubocop: disable CodeReuse/ActiveRecord
-        def order_by_end_event(query)
+        def order_by_end_event(query, extra_columns_to_select = [:id])
           ordered_query = query.reorder(stage.end_event.timestamp_projection.desc)
 
           # When filtering for more than one label, postgres requires the columns in ORDER BY to be present in the GROUP BY clause
           if requires_grouping?
             column_list = [
-              ordered_query.arel_table[:id],
+              *extra_columns_to_select,
               *stage.end_event.column_list,
               *stage.start_event.column_list
             ]