From 3d5f710ae6849d3977f178ba0aeeda681eedae37 Mon Sep 17 00:00:00 2001
From: Dmitry Gruzd <dgruzd@gitlab.com>
Date: Mon, 19 Feb 2024 15:39:46 +0000
Subject: [PATCH] Add Zoekt dot_com_rollout scheduling task
---
.../search/zoekt/scheduling_service.rb | 55 +++++++++++++++++++
.../beta/zoekt_dot_com_rollout.yml | 9 +++
.../search/zoekt/scheduling_service_spec.rb | 36 ++++++++++++
3 files changed, 100 insertions(+)
create mode 100644 ee/config/feature_flags/beta/zoekt_dot_com_rollout.yml
diff --git a/ee/app/services/search/zoekt/scheduling_service.rb b/ee/app/services/search/zoekt/scheduling_service.rb
index a5141cb67693c..2a54977b40b2f 100644
--- a/ee/app/services/search/zoekt/scheduling_service.rb
+++ b/ee/app/services/search/zoekt/scheduling_service.rb
@@ -6,6 +6,7 @@ class SchedulingService
include Gitlab::Loggable
TASKS = %i[
+ dot_com_rollout
remove_expired_subscriptions
node_assignment
].freeze
@@ -13,6 +14,9 @@ class SchedulingService
BUFFER_FACTOR = 3
WATERMARK_LIMIT = 0.8
+ DOT_COM_ROLLOUT_TARGET_BYTES = 100.gigabytes
+ DOT_COM_ROLLOUT_LIMIT = 2000
+
attr_reader :task
def initialize(task)
@@ -36,6 +40,57 @@ def logger
@logger ||= ::Zoekt::Logger.build
end
+ def execute_every(period, cache_key:)
+ Rails.cache.fetch([self.class.name, :execute_every, cache_key], expires_in: period) do
+ yield
+ end
+ end
+
+ # A temporary task to simplify the .com Zoekt rollout
+ # rubocop:disable CodeReuse/ActiveRecord -- this is a temporary task, which will be removed after the rollout
+ def dot_com_rollout
+ return false unless ::Gitlab::Saas.feature_available?(:exact_code_search)
+ return false if Feature.disabled?(:zoekt_dot_com_rollout)
+ return false if EnabledNamespace.with_missing_indices.exists?
+
+ execute_every 6.hours, cache_key: :dot_com_rollout do
+ size = 0
+ sizes = {}
+
+ indexed_namespaces_ids = Search::Zoekt::EnabledNamespace.find_each.map(&:root_namespace_id).to_set
+ namespaces_to_add = GitlabSubscription.with_a_paid_hosted_plan
+ .where('end_date > ? OR end_date IS NULL', Date.today)
+ scope = Group.includes(:root_storage_statistics)
+ .where(parent_id: nil)
+ .where(id: namespaces_to_add.select(:namespace_id))
+ scope.find_each do |n|
+ next if indexed_namespaces_ids.include?(n.id)
+
+ sizes[n.id] = n.root_storage_statistics.repository_size if n.root_storage_statistics
+ end
+
+ sorted = sizes.to_a.sort_by { |_k, v| v }
+
+ count = 0
+ sorted.take(DOT_COM_ROLLOUT_LIMIT).each do |id, s|
+ size += s
+ break count if size > DOT_COM_ROLLOUT_TARGET_BYTES
+
+ Search::Zoekt::EnabledNamespace.create!(root_namespace_id: id, search: false)
+ count += 1
+ end
+
+ logger.info(build_structured_payload(
+ task: :dot_com_rollout,
+ message: 'Rollout has been completed',
+ namespace_count: count
+ ))
+
+ count
+ end
+ end
+ # rubocop:enable CodeReuse/ActiveRecord
+
def remove_expired_subscriptions
return false unless ::Gitlab::Saas.feature_available?(:exact_code_search)
diff --git a/ee/config/feature_flags/beta/zoekt_dot_com_rollout.yml b/ee/config/feature_flags/beta/zoekt_dot_com_rollout.yml
new file mode 100644
index 0000000000000..140275a581c5d
--- /dev/null
+++ b/ee/config/feature_flags/beta/zoekt_dot_com_rollout.yml
@@ -0,0 +1,9 @@
+---
+name: zoekt_dot_com_rollout
+feature_issue_url:
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/144988
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/442203
+milestone: '16.10'
+group: group::global search
+type: beta
+default_enabled: false
diff --git a/ee/spec/services/search/zoekt/scheduling_service_spec.rb b/ee/spec/services/search/zoekt/scheduling_service_spec.rb
index 402b746c4d6ae..b951a0c2a5695 100644
--- a/ee/spec/services/search/zoekt/scheduling_service_spec.rb
+++ b/ee/spec/services/search/zoekt/scheduling_service_spec.rb
@@ -42,6 +42,42 @@
end
end
+ describe '#dot_com_rollout' do
+ let(:task) { :dot_com_rollout }
+
+ it 'returns false unless saas' do
+ expect(execute_task).to eq(false)
+ end
+
+ it 'returns false if there are unassigned namespaces' do
+ create(:zoekt_enabled_namespace)
+
+ expect(execute_task).to eq(false)
+ end
+
+ context 'when on .com', :saas do
+ let_it_be(:group) { create(:group) }
+ let_it_be(:subscription) { create(:gitlab_subscription, namespace: group) }
+ let_it_be(:root_storage_statistics) { create(:namespace_root_storage_statistics, namespace: group) }
+
+ context 'when feature flag is disabled' do
+ before do
+ stub_feature_flags(zoekt_dot_com_rollout: false)
+ end
+
+ it 'returns false' do
+ expect(execute_task).to eq(false)
+ end
+ end
+
+ it 'assigns namespaces to a node' do
+ expect { execute_task }.to change { ::Search::Zoekt::EnabledNamespace.count }.by(1)
+
+ expect(::Search::Zoekt::EnabledNamespace.pluck(:root_namespace_id)).to contain_exactly(group.id)
+ end
+ end
+ end
+
describe '#remove_expired_subscriptions' do
let(:task) { :remove_expired_subscriptions }
--
GitLab