Refactoring Environment#stop_actions for CI DB

This commit refactors the code not to use CTE joins
between deployments/environments and ci_builds/pipelines.

app/models/environment.rb

@@ -77,6 +77,7 @@ class Environment < ApplicationRecord
   scope :in_review_folder, -> { where(environment_type: "review") }
   scope :for_name, -> (name) { where(name: name) }
   scope :preload_cluster, -> { preload(last_deployment: :cluster) }
+  scope :preload_project, -> { preload(:project) }
   scope :auto_stoppable, -> (limit) { available.where('auto_stop_at < ?', Time.zone.now).limit(limit) }
   scope :auto_deletable, -> (limit) { stopped.where('auto_delete_at < ?', Time.zone.now).limit(limit) }
 
@@ -132,6 +133,10 @@ class Environment < ApplicationRecord
     state :available
     state :stopped
 
+    before_transition any => :stopped do |environment|
+      environment.auto_stop_at = nil
+    end
+
     after_transition do |environment|
       environment.expire_etag_cache
     end
@@ -168,33 +173,6 @@ class Environment < ApplicationRecord
   end
 
   class << self
-    ##
-    # This method returns stop actions (jobs) for multiple environments within one
-    # query. It's useful to avoid N+1 problem.
-    #
-    # NOTE: The count of environments should be small~medium (e.g. < 5000)
-    def stop_actions
-      cte = cte_for_deployments_with_stop_action
-      ci_builds = Ci::Build.arel_table
-
-      inner_join_stop_actions = ci_builds.join(cte.table).on(
-        ci_builds[:project_id].eq(cte.table[:project_id])
-          .and(ci_builds[:ref].eq(cte.table[:ref]))
-          .and(ci_builds[:name].eq(cte.table[:on_stop]))
-      ).join_sources
-
-      pipeline_ids = ci_builds.join(cte.table).on(
-        ci_builds[:id].eq(cte.table[:deployable_id])
-      ).project(:commit_id)
-
-      Ci::Build.joins(inner_join_stop_actions)
-               .with(cte.to_arel)
-               .where(ci_builds[:commit_id].in(pipeline_ids))
-               .where(status: Ci::HasStatus::BLOCKED_STATUS)
-               .preload_project_and_pipeline_project
-               .preload(:user, :metadata, :deployment)
-    end
-
     def count_by_state
       environments_count_by_state = group(:state).count
 
@@ -202,15 +180,6 @@ class Environment < ApplicationRecord
         count_hash[state] = environments_count_by_state[state.to_s] || 0
       end
     end
-
-    private
-
-    def cte_for_deployments_with_stop_action
-      Gitlab::SQL::CTE.new(:deployments_with_stop_action,
-        Deployment.where(environment_id: select(:id))
-          .distinct_on_environment
-          .stoppable)
-    end
   end
 
   def clear_prometheus_reactive_cache!(query_name)

app/services/environments/auto_stop_service.rb

@@ -28,11 +28,17 @@ module Environments
     private
 
     def stop_in_batch
-      environments = Environment.auto_stoppable(BATCH_SIZE)
+      environments = Environment.preload_project.select(:id, :project_id).auto_stoppable(BATCH_SIZE)
 
-      return false unless environments.exists?
+      return false if environments.empty?
 
-      Environments::StopService.execute_in_batch(environments)
+      Environments::AutoStopWorker.bulk_perform_async_with_contexts(
+        environments,
+        arguments_proc: -> (environment) { environment.id },
+        context_proc: -> (environment) { { project: environment.project } }
+      )
+
+      true
     end
   end
 end

app/services/environments/stop_service.rb

@@ -22,22 +22,6 @@ module Environments
       merge_request.environments.each { |environment| execute(environment) }
     end
 
-    ##
-    # This method is for stopping multiple environments in a batch style.
-    # The maximum acceptable count of environments is roughly 5000. Please
-    # apply acceptable `LIMIT` clause to the `environments` relation.
-    def self.execute_in_batch(environments)
-      stop_actions = environments.stop_actions.load
-
-      environments.update_all(auto_stop_at: nil, state: 'stopped')
-
-      stop_actions.each do |stop_action|
-        stop_action.play(stop_action.user)
-      rescue StandardError => e
-        Gitlab::ErrorTracking.track_exception(e, deployable_id: stop_action.id)
-      end
-    end
-
     private
 
     def environments

app/workers/all_queues.yml

@@ -2043,6 +2043,15 @@
   :weight: 2
   :idempotent:
   :tags: []
+- :name: environments_auto_stop
+  :worker_name: Environments::AutoStopWorker
+  :feature_category: :continuous_delivery
+  :has_external_dependencies:
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: environments_canary_ingress_update
   :worker_name: Environments::CanaryIngress::UpdateWorker
   :feature_category: :continuous_delivery

app/workers/environments/auto_stop_worker.rb

+# frozen_string_literal: true
+
+module Environments
+  class AutoStopWorker
+    include ApplicationWorker
+
+    data_consistency :always
+    idempotent!
+    feature_category :continuous_delivery
+
+    def perform(environment_id, params = {})
+      Environment.find_by_id(environment_id).try do |environment|
+        user = environment.stop_action&.user
+        environment.stop_with_action!(user)
+      end
+    end
+  end
+end

config/sidekiq_queues.yml

@@ -126,6 +126,8 @@
   - 2
 - - emails_on_push
   - 2
+- - environments_auto_stop
+  - 1
 - - environments_canary_ingress_update
   - 1
 - - epics

spec/models/environment_spec.rb

@@ -135,6 +135,20 @@ RSpec.describe Environment, :use_clean_rails_memory_store_caching do
 
       environment.stop
     end
+
+    context 'when environment has auto stop period' do
+      let!(:environment) { create(:environment, :available, :auto_stoppable, project: project) }
+
+      it 'clears auto stop period when the environment has stopped' do
+        environment.stop!
+
+        expect(environment.auto_stop_at).to be_nil
+      end
+
+      it 'does not clear auto stop period when the environment has not stopped' do
+        expect(environment.auto_stop_at).to be_present
+      end
+    end
   end
 
   describe '.for_name_like' do
@@ -233,55 +247,6 @@ RSpec.describe Environment, :use_clean_rails_memory_store_caching do
     end
   end
 
-  describe '.stop_actions' do
-    subject { environments.stop_actions }
-
-    let_it_be(:project) { create(:project, :repository) }
-    let_it_be(:user) { create(:user) }
-
-    let(:environments) { Environment.all }
-
-    before_all do
-      project.add_developer(user)
-      project.repository.add_branch(user, 'review/feature-1', 'master')
-      project.repository.add_branch(user, 'review/feature-2', 'master')
-    end
-
-    shared_examples_for 'correct filtering' do
-      it 'returns stop actions for available environments only' do
-        expect(subject.count).to eq(1)
-        expect(subject.first.name).to eq('stop_review_app')
-        expect(subject.first.ref).to eq('review/feature-1')
-      end
-    end
-
-    before do
-      create_review_app(user, project, 'review/feature-1')
-      create_review_app(user, project, 'review/feature-2')
-    end
-
-    it 'returns stop actions for environments' do
-      expect(subject.count).to eq(2)
-      expect(subject).to match_array(Ci::Build.where(name: 'stop_review_app'))
-    end
-
-    context 'when one of the stop actions has already been executed' do
-      before do
-        Ci::Build.where(ref: 'review/feature-2').find_by_name('stop_review_app').enqueue!
-      end
-
-      it_behaves_like 'correct filtering'
-    end
-
-    context 'when one of the deployments does not have stop action' do
-      before do
-        Deployment.where(ref: 'review/feature-2').update_all(on_stop: nil)
-      end
-
-      it_behaves_like 'correct filtering'
-    end
-  end
-
   describe '.pluck_names' do
     subject { described_class.pluck_names }
 

spec/services/environments/auto_stop_service_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Environments::AutoStopService, :clean_gitlab_redis_shared_state do
+RSpec.describe Environments::AutoStopService, :clean_gitlab_redis_shared_state, :sidekiq_inline do
   include CreateEnvironmentsHelpers
   include ExclusiveLeaseHelpers
 
@@ -42,6 +42,15 @@ RSpec.describe Environments::AutoStopService, :clean_gitlab_redis_shared_state d
       expect(Ci::Build.where(name: 'stop_review_app').map(&:status).uniq).to eq(['pending'])
     end
 
+    it 'schedules stop processes in bulk' do
+      args = [[Environment.find_by_name('review/feature-1').id], [Environment.find_by_name('review/feature-2').id]]
+
+      expect(Environments::AutoStopWorker)
+        .to receive(:bulk_perform_async).with(args).once.and_call_original
+
+      subject
+    end
+
     context 'when the other sidekiq worker has already been running' do
       before do
         stub_exclusive_lease_taken(described_class::EXCLUSIVE_LOCK_KEY)

spec/services/environments/stop_service_spec.rb

@@ -237,60 +237,6 @@ RSpec.describe Environments::StopService do
     end
   end
 
-  describe '.execute_in_batch' do
-    subject { described_class.execute_in_batch(environments) }
-
-    let_it_be(:project) { create(:project, :repository) }
-    let_it_be(:user) { create(:user) }
-
-    let(:environments) { Environment.available }
-
-    before_all do
-      project.add_developer(user)
-      project.repository.add_branch(user, 'review/feature-1', 'master')
-      project.repository.add_branch(user, 'review/feature-2', 'master')
-    end
-
-    before do
-      create_review_app(user, project, 'review/feature-1')
-      create_review_app(user, project, 'review/feature-2')
-    end
-
-    it 'stops environments' do
-      expect { subject }
-        .to change { project.environments.all.map(&:state).uniq }
-        .from(['available']).to(['stopped'])
-
-      expect(project.environments.all.map(&:auto_stop_at).uniq).to eq([nil])
-    end
-
-    it 'plays stop actions' do
-      expect { subject }
-        .to change { Ci::Build.where(name: 'stop_review_app').map(&:status).uniq }
-        .from(['manual']).to(['pending'])
-    end
-
-    context 'when user does not have a permission to play the stop action' do
-      before do
-        project.team.truncate
-      end
-
-      it 'tracks the exception' do
-        expect(Gitlab::ErrorTracking)
-          .to receive(:track_exception)
-          .with(Gitlab::Access::AccessDeniedError, anything)
-          .twice
-          .and_call_original
-
-        subject
-      end
-
-      after do
-        project.add_developer(user)
-      end
-    end
-  end
-
   def expect_environment_stopped_on(branch)
     expect { service.execute_for_branch(branch) }
       .to change { Environment.last.state }.from('available').to('stopped')

spec/workers/environments/auto_stop_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Environments::AutoStopWorker do
+  include CreateEnvironmentsHelpers
+
+  subject { worker.perform(environment_id) }
+
+  let_it_be(:project) { create(:project, :repository) }
+  let_it_be(:developer) { create(:user).tap { |u| project.add_developer(u) } }
+  let_it_be(:reporter) { create(:user).tap { |u| project.add_reporter(u) } }
+
+  before_all do
+    project.repository.add_branch(developer, 'review/feature', 'master')
+  end
+
+  let!(:environment) { create_review_app(user, project, 'review/feature').environment }
+  let(:environment_id) { environment.id }
+  let(:worker) { described_class.new }
+  let(:user) { developer }
+
+  it 'stops the environment' do
+    expect { subject }
+      .to change { Environment.find_by_name('review/feature').state }
+      .from('available').to('stopped')
+  end
+
+  it 'executes the stop action' do
+    expect { subject }
+      .to change { Ci::Build.find_by_name('stop_review_app').status }
+      .from('manual').to('pending')
+  end
+
+  context 'when user does not have a permission to play the stop action' do
+    let(:user) { reporter }
+
+    it 'raises an error' do
+      expect { subject }.to raise_error(Gitlab::Access::AccessDeniedError)
+    end
+  end
+
+  context 'when the environment has already been stopped' do
+    before do
+      environment.stop!
+    end
+
+    it 'does not execute the stop action' do
+      expect { subject }
+        .not_to change { Ci::Build.find_by_name('stop_review_app').status }
+    end
+  end
+
+  context 'when there are no deployments and associted stop actions' do
+    let!(:environment) { create(:environment) }
+
+    it 'stops the environment' do
+      subject
+
+      expect(environment.reload).to be_stopped
+    end
+  end
+
+  context 'when there are no corresponding environment record' do
+    let!(:environment) { double(:environment, id: non_existing_record_id) }
+
+    it 'ignores the invalid record' do
+      expect { subject }.not_to raise_error
+    end
+  end
+end