ES: API to rollout namespaces

Async worker to handle potential large number of namespaces.

changelogs/unreleased/33681-api.yml

+---
+title: Add API for rollout Elasticsearch per plan level
+merge_request: 22240
+author:
+type: added

config/sidekiq_queues.yml

@@ -116,6 +116,7 @@
   - [elastic_full_index, 1]
   - [elastic_commit_indexer, 1]
   - [elastic_namespace_indexer, 1]
+  - [elastic_namespace_rollout, 1]
   - [export_csv, 1]
   - [incident_management, 2]
   - [jira_connect, 1]

ee/app/models/elasticsearch_indexed_namespace.rb

@@ -10,10 +10,38 @@ class ElasticsearchIndexedNamespace < ApplicationRecord
 
   validates :namespace_id, presence: true, uniqueness: true
 
+  scope :namespace_in, -> (namespaces) { where(namespace_id: namespaces) }
+
   def self.target_attr_name
     :namespace_id
   end
 
+  # rubocop: disable Naming/UncommunicativeMethodParamName
+  def self.index_first_n_namespaces_of_plan(plan, n)
+    indexed_namespaces = self.select(:namespace_id)
+
+    GitlabSubscription
+      .with_hosted_plan(plan)
+      .where.not(namespace_id: indexed_namespaces)
+      .order(namespace_id: :asc)
+      .limit(n)
+      .pluck(:namespace_id)
+      .each { |id| create!(namespace_id: id) }
+  end
+
+  def self.unindex_last_n_namespaces_of_plan(plan, n)
+    namespaces_under_plan = GitlabSubscription.with_hosted_plan(plan).select(:namespace_id)
+
+    # rubocop: disable Cop/DestroyAll
+    # destroy_all is used in order to trigger `delete_from_index` callback
+    where(namespace: namespaces_under_plan)
+      .order(created_at: :desc)
+      .limit(n)
+      .destroy_all
+    # rubocop: enable Cop/DestroyAll
+  end
+  # rubocop: enable Naming/UncommunicativeMethodParamName
+
   private
 
   def index

ee/app/workers/all_queues.yml

@@ -70,6 +70,7 @@
 - elastic_commit_indexer
 - elastic_indexer
 - elastic_full_index
+- elastic_namespace_rollout
 - export_csv
 - ldap_group_sync
 - new_epic

ee/app/workers/elastic_namespace_rollout_worker.rb

+# frozen_string_literal: true
+
+class ElasticNamespaceRolloutWorker
+  include ApplicationWorker
+
+  feature_category :search
+  sidekiq_options retry: 2
+
+  ROLLOUT = 'rollout'
+  ROLLBACK = 'rollback'
+
+  # @param plan [String] which plan the rollout is scoped to
+  # @param percentage [Integer]
+  # @param mode [ROLLOUT, ROLLBACK] determine whether to rollout or rollback
+  def perform(plan, percentage, mode)
+    total_with_plan = GitlabSubscription.with_hosted_plan(plan).count
+
+    expected_count = total_with_plan * (percentage / 100.0)
+
+    current_count = ElasticsearchIndexedNamespace
+      .namespace_in(GitlabSubscription.with_hosted_plan(plan).select(:namespace_id))
+      .count
+
+    case mode
+    when ROLLOUT
+      rollout(plan, expected_count, current_count)
+    when ROLLBACK
+      rollback(plan, expected_count, current_count)
+    end
+  end
+
+  private
+
+  def rollout(plan, expected_count, current_count)
+    required_count_changes = [expected_count - current_count, 0].max
+
+    logger.info(message: 'rollout_elasticsearch_indexed_namespaces', changes: required_count_changes, expected_count: expected_count, current_count: current_count, plan: plan)
+
+    if required_count_changes > 0
+      ElasticsearchIndexedNamespace.index_first_n_namespaces_of_plan(plan, required_count_changes)
+    end
+  end
+
+  def rollback(plan, expected_count, current_count)
+    required_count_changes = [current_count - expected_count, 0].max
+
+    logger.info(message: 'rollback_elasticsearch_indexed_namespaces', changes: required_count_changes, expected_count: expected_count, current_count: current_count, plan: plan)
+
+    if required_count_changes > 0
+      ElasticsearchIndexedNamespace.unindex_last_n_namespaces_of_plan(plan, required_count_changes)
+    end
+  end
+
+  def logger
+    @logger ||= ::Gitlab::Elasticsearch::Logger.build
+  end
+end

ee/lib/api/elasticsearch_indexed_namespaces.rb

+# frozen_string_literal: true
+
+module API
+  class ElasticsearchIndexedNamespaces < Grape::API
+    before { authenticated_as_admin! }
+
+    resource :elasticsearch_indexed_namespaces do
+      desc 'Rollout namespaces to be indexed up to n%' do
+        detail <<~END
+          This feature was introduced in GitLab 12.7.
+
+          This will only ever increase the number of indexed namespaces. Providing a value lower than the current rolled out percentage will have no effect.
+
+          This percentage is never persisted but is used to calculate the number of new namespaces to rollout.
+
+          If the same percentage is applied again at a later time, due to possible new namespaces being created during the period, some of them will also be indexed. Therefore you may expect that setting this to 10%, then waiting a month and setting to 10% again will trigger new namespaces to be added (i.e. 10% of the number of newly created namespaces in the last month within the given plan).
+        END
+      end
+      params do
+        requires :percentage, type: Integer, values: 0..100
+        requires :plan, type: String, values: Plan::ALL_HOSTED_PLANS
+      end
+      put 'rollout' do
+        ElasticNamespaceRolloutWorker.perform_async(params[:plan], params[:percentage], ElasticNamespaceRolloutWorker::ROLLOUT)
+      end
+
+      desc 'Rollback namespaces to be indexed down to n%' do
+        detail <<~END
+          This feature was introduced in GitLab 12.7.
+
+          This will only ever decrease the number of indexed namespaces. Providing a value higher than the current rolled out percentage will have no effect.
+
+          This percentage is never persisted but is used to calculate the number of namespaces to rollback.
+        END
+      end
+      params do
+        requires :percentage, type: Integer, values: 0..100
+        requires :plan, type: String, values: Plan::ALL_HOSTED_PLANS
+      end
+      put 'rollback' do
+        ElasticNamespaceRolloutWorker.perform_async(params[:plan], params[:percentage], ElasticNamespaceRolloutWorker::ROLLBACK)
+      end
+    end
+  end
+end

ee/lib/ee/api/api.rb

@@ -18,6 +18,7 @@ module EE
         mount ::API::EpicIssues
         mount ::API::EpicLinks
         mount ::API::Epics
+        mount ::API::ElasticsearchIndexedNamespaces
         mount ::API::FeatureFlags
         mount ::API::FeatureFlagScopes
         mount ::API::ContainerRegistryEvent

ee/spec/models/elasticsearch_indexed_namespace_spec.rb

@@ -7,6 +7,16 @@ describe ElasticsearchIndexedNamespace do
     stub_ee_application_setting(elasticsearch_indexing: true)
   end
 
+  describe 'scope' do
+    describe '.namespace_in' do
+      let(:records) { create_list(:elasticsearch_indexed_namespace, 3) }
+
+      it 'returns records of the ids' do
+        expect(described_class.namespace_in(records.last(2).map(&:id))).to eq(records.last(2))
+      end
+    end
+  end
+
   it_behaves_like 'an elasticsearch indexed container' do
     let(:container) { :elasticsearch_indexed_namespace }
     let(:attribute) { :namespace_id }
@@ -17,4 +27,73 @@ describe ElasticsearchIndexedNamespace do
       expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(subject.namespace_id, :delete)
     end
   end
+
+  context 'with plans' do
+    Plan::PAID_HOSTED_PLANS.each do |plan|
+      plan_factory = "#{plan}_plan"
+      let_it_be(plan_factory) { create(plan_factory) }
+    end
+
+    let_it_be(:namespaces) { create_list(:namespace, 3) }
+    let_it_be(:subscription1) { create(:gitlab_subscription, namespace: namespaces[2]) }
+    let_it_be(:subscription2) { create(:gitlab_subscription, namespace: namespaces[0]) }
+    let_it_be(:subscription3) { create(:gitlab_subscription, :silver, namespace: namespaces[1]) }
+
+    before do
+      stub_ee_application_setting(elasticsearch_indexing: false)
+    end
+
+    def get_indexed_namespaces
+      described_class.order(:created_at).pluck(:namespace_id)
+    end
+
+    describe '.index_first_n_namespaces_of_plan' do
+      it 'creates records, scoped by plan and ordered by namespace id' do
+        ids = namespaces.map(&:id)
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[0], :index)
+
+        described_class.index_first_n_namespaces_of_plan('gold', 1)
+
+        expect(get_indexed_namespaces).to eq([ids[0]])
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[2], :index)
+
+        described_class.index_first_n_namespaces_of_plan('gold', 2)
+
+        expect(get_indexed_namespaces).to eq([ids[0], ids[2]])
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[1], :index)
+
+        described_class.index_first_n_namespaces_of_plan('silver', 1)
+
+        expect(get_indexed_namespaces).to eq([ids[0], ids[2], ids[1]])
+      end
+    end
+
+    describe '.unindex_last_n_namespaces_of_plan' do
+      before do
+        described_class.index_first_n_namespaces_of_plan('gold', 2)
+        described_class.index_first_n_namespaces_of_plan('silver', 1)
+      end
+
+      it 'creates records, scoped by plan and ordered by namespace id' do
+        ids = namespaces.map(&:id)
+
+        expect(get_indexed_namespaces).to eq([ids[0], ids[2], ids[1]])
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[2], :delete)
+
+        described_class.unindex_last_n_namespaces_of_plan('gold', 1)
+
+        expect(get_indexed_namespaces).to eq([ids[0], ids[1]])
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[1], :delete)
+
+        described_class.unindex_last_n_namespaces_of_plan('silver', 1)
+
+        expect(get_indexed_namespaces).to eq([ids[0]])
+        expect(ElasticNamespaceIndexerWorker).to receive(:perform_async).with(ids[0], :delete)
+
+        described_class.unindex_last_n_namespaces_of_plan('gold', 1)
+
+        expect(get_indexed_namespaces).to be_empty
+      end
+    end
+  end
 end

ee/spec/requests/api/elasticsearch_indexed_namespaces_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe API::ElasticsearchIndexedNamespaces do
+  let_it_be(:admin) { create(:admin) }
+  let_it_be(:non_admin) { create(:user) }
+
+  shared_examples 'rollout related' do
+    context 'when parameters are incorrect' do
+      using RSpec::Parameterized::TableSyntax
+
+      where(:percentage, :plan) do
+        -1    | 'gold'
+        101   | 'gold'
+        nil   | 'gold'
+        1     | nil
+        1     | 'foobar'
+      end
+
+      with_them do
+        it 'errs' do
+          put api(path, admin), params: { plan: plan, percentage: percentage }
+          expect(response).to have_gitlab_http_status(400)
+        end
+      end
+    end
+
+    it 'prohibits non-admin' do
+      put api(path, non_admin), params: { plan: 'gold', percentage: 50 }
+      expect(response).to have_gitlab_http_status(403)
+    end
+  end
+
+  describe 'PUT /elasticsearch_indexed_namespaces/rollout' do
+    let(:path) { "/elasticsearch_indexed_namespaces/rollout" }
+
+    include_context 'rollout related'
+
+    it 'invokes ElasticNamespaceRolloutWorker rollout' do
+      expect(ElasticNamespaceRolloutWorker).to receive(:perform_async).with('gold', 50, ElasticNamespaceRolloutWorker::ROLLOUT)
+
+      put api(path, admin), params: { plan: 'gold', percentage: 50 }
+
+      expect(response).to have_gitlab_http_status(200)
+    end
+  end
+
+  describe 'PUT /elasticsearch_indexed_namespaces/rollback' do
+    let(:path) { "/elasticsearch_indexed_namespaces/rollback" }
+
+    include_context 'rollout related'
+
+    it 'invokes ElasticNamespaceRolloutWorker rollback' do
+      expect(ElasticNamespaceRolloutWorker).to receive(:perform_async).with('gold', 50, ElasticNamespaceRolloutWorker::ROLLBACK)
+
+      put api(path, admin), params: { plan: 'gold', percentage: 50 }
+
+      expect(response).to have_gitlab_http_status(200)
+    end
+  end
+end

ee/spec/workers/elastic_namespace_rollout_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe ElasticNamespaceRolloutWorker do
+  ROLLOUT = described_class::ROLLOUT
+  ROLLBACK = described_class::ROLLBACK
+
+  Plan::PAID_HOSTED_PLANS.each do |plan|
+    plan_factory = "#{plan}_plan"
+    let_it_be(plan_factory) { create(plan_factory) }
+  end
+
+  before_all do
+    Plan::PAID_HOSTED_PLANS.each do |plan|
+      4.times do
+        create(:gitlab_subscription, hosted_plan: public_send("#{plan}_plan"))
+      end
+    end
+  end
+
+  def expect_percentage_to_result_in_records(percentage, record_count, mode)
+    subject.perform('gold', percentage, mode)
+
+    namespace_ids = GitlabSubscription
+      .with_hosted_plan('gold')
+      .order(id: :asc)
+      .pluck(:namespace_id)
+
+    expect(
+      ElasticsearchIndexedNamespace.pluck(:namespace_id)
+    ).to contain_exactly(*namespace_ids.first(record_count))
+  end
+
+  def get_namespace_ids(plan, count)
+    GitlabSubscription
+      .with_hosted_plan(plan)
+      .order(id: :asc)
+      .pluck(:namespace_id)
+      .first(count)
+  end
+
+  it 'rolls out and back' do
+    # Rollout
+    expect_percentage_to_result_in_records(0, 0, ROLLOUT)
+    expect_percentage_to_result_in_records(50, 2, ROLLOUT)
+    expect_percentage_to_result_in_records(25, 2, ROLLOUT) # no op
+    expect_percentage_to_result_in_records(100, 4, ROLLOUT)
+
+    # Rollback
+    expect_percentage_to_result_in_records(50, 2, ROLLBACK)
+    expect_percentage_to_result_in_records(75, 2, ROLLBACK) # no op
+    expect_percentage_to_result_in_records(0, 0, ROLLBACK)
+  end
+
+  it 'distinguishes different plans' do
+    # Rollout
+    subject.perform('gold', 50, ROLLOUT)
+    subject.perform('silver', 25, ROLLOUT)
+
+    expect(
+      ElasticsearchIndexedNamespace.pluck(:namespace_id)
+    ).to contain_exactly(
+      *get_namespace_ids(:gold, 2),
+      *get_namespace_ids(:silver, 1)
+    )
+
+    # Rollback
+    subject.perform('gold', 25, ROLLBACK)
+    subject.perform('silver', 0, ROLLBACK)
+
+    expect(
+      ElasticsearchIndexedNamespace.pluck(:namespace_id)
+    ).to contain_exactly(
+      *get_namespace_ids(:gold, 1)
+    )
+  end
+
+  it 'logs' do
+    logger = subject.send(:logger)
+
+    expect(logger).to receive(:info).with(
+      hash_including(
+        message: "rollout_elasticsearch_indexed_namespaces",
+        changes: 3,
+        expected_count: 3,
+        current_count: 0,
+        plan: 'gold'
+      )
+    ).and_call_original
+
+    subject.perform('gold', 75, ROLLOUT)
+
+    expect(logger).to receive(:info).with(
+      hash_including(
+        message: "rollback_elasticsearch_indexed_namespaces",
+        changes: 2,
+        expected_count: 1,
+        current_count: 3,
+        plan: 'gold'
+      )
+    ).and_call_original
+
+    subject.perform('gold', 25, ROLLBACK)
+  end
+end