From 550c01a295e5690dc6b79247205df5c2aa3d2b5f Mon Sep 17 00:00:00 2001
From: Steve Xuereb <sxuereb@gitlab.com>
Date: Wed, 9 Oct 2024 09:10:15 +0000
Subject: [PATCH] Revert "Merge branch 'sc1-session-map' into 'master'"
This reverts merge request !166986
---
app/models/application_record.rb | 2 +-
app/models/ci/runner.rb | 2 +-
app/models/ci/runner_manager.rb | 2 +-
app/models/concerns/packages/downloadable.rb | 2 +-
...se_sql_function_for_primary_key_lookups.rb | 2 +-
app/models/concerns/users/visitable.rb | 3 +-
app/models/internal_id.rb | 2 +-
app/models/project.rb | 3 +-
app/models/sent_notification.rb | 2 +-
.../finished_pipelines_sync_service.rb | 2 +-
app/services/ci/register_job_service.rb | 2 +-
app/services/issuable_base_service.rb | 2 +-
.../last_used_service.rb | 3 +-
app/services/users/activity_service.rb | 3 +-
.../user_refresh_from_replica_worker.rb | 4 +-
.../cleanup_container_repository_worker.rb | 4 +-
.../container_expiration_policy_worker.rb | 4 +-
.../image_ttl_group_policy_worker.rb | 4 +-
.../cleanup_package_registry_worker.rb | 4 +-
.../use_load_balancing_session_map.yml | 9 -
.../ci_finished_builds_sync_service.rb | 2 +-
ee/lib/analytics/group_activity_calculator.rb | 12 +-
.../license_scanning/package_licenses.rb | 4 +-
.../license_scanning/package_licenses_spec.rb | 8 +-
ee/spec/lib/gitlab/llm/ai_message_spec.rb | 3 +-
lib/api/ci/helpers/runner.rb | 5 +-
lib/api/maven_packages.rb | 3 +-
lib/gitlab/auth.rb | 3 +-
lib/gitlab/checks/matching_merge_request.rb | 8 +-
lib/gitlab/database.rb | 8 +-
lib/gitlab/database/consistency.rb | 5 +-
.../indicators/write_ahead_log.rb | 2 +-
lib/gitlab/database/load_balancing.rb | 4 -
.../load_balancing/action_cable_callbacks.rb | 2 +-
.../load_balancing/connection_proxy.rb | 2 +-
.../load_balancing/rack_middleware.rb | 7 +-
lib/gitlab/database/load_balancing/session.rb | 1 -
.../database/load_balancing/session_map.rb | 128 --------
.../sidekiq_client_middleware.rb | 14 +-
.../sidekiq_server_middleware.rb | 6 +-
.../database/load_balancing/sticking.rb | 12 +-
.../load_balancing/wal_tracking_sender.rb | 21 +-
.../partitioning/partition_manager.rb | 2 +-
.../database/postgres_autovacuum_activity.rb | 4 +-
.../action_cable_with_load_balancing.rb | 7 +-
.../json/streaming_serializer.rb | 4 +-
spec/initializers/load_balancing_spec.rb | 2 +-
spec/lib/feature_spec.rb | 5 -
.../auth/token_expiration_banner_spec.rb | 3 +-
spec/lib/gitlab/auth_spec.rb | 3 +-
.../checks/matching_merge_request_spec.rb | 6 +-
spec/lib/gitlab/database/bulk_update_spec.rb | 8 +-
spec/lib/gitlab/database/consistency_spec.rb | 6 +-
.../indicators/write_ahead_log_spec.rb | 3 +-
.../action_cable_callbacks_spec.rb | 4 +-
.../load_balancing/connection_proxy_spec.rb | 20 +-
.../load_balancing/rack_middleware_spec.rb | 11 +-
.../load_balancing/session_map_spec.rb | 306 ------------------
.../sidekiq_client_middleware_spec.rb | 19 +-
.../sidekiq_server_middleware_spec.rb | 12 +-
.../database/load_balancing/sticking_spec.rb | 12 +-
.../gitlab/database/load_balancing_spec.rb | 68 ++--
.../postgres_autovacuum_activity_spec.rb | 3 +-
spec/lib/gitlab/database_spec.rb | 14 +-
.../action_cable_with_load_balancing_spec.rb | 17 +-
.../json/streaming_serializer_spec.rb | 7 +-
spec/models/application_record_spec.rb | 3 +-
.../concerns/packages/downloadable_spec.rb | 3 +-
spec/models/sent_notification_spec.rb | 4 +-
.../api/ci/runner/jobs_artifacts_spec.rb | 17 +-
spec/requests/api/maven_packages_spec.rb | 8 +-
.../last_used_service_spec.rb | 9 +-
spec/services/users/activity_service_spec.rb | 11 +-
spec/support/database_load_balancing.rb | 4 +-
spec/support/db_cleaner.rb | 21 --
.../wal_tracking_shared_context.rb | 12 +-
.../gitlab/db/validate_config_rake_spec.rb | 1 +
.../user_refresh_from_replica_worker_spec.rb | 6 +-
...leanup_container_repository_worker_spec.rb | 7 +-
...container_expiration_policy_worker_spec.rb | 3 +-
.../image_ttl_group_policy_worker_spec.rb | 6 +-
.../cleanup_package_registry_worker_spec.rb | 6 +-
82 files changed, 196 insertions(+), 792 deletions(-)
delete mode 100644 config/feature_flags/gitlab_com_derisk/use_load_balancing_session_map.yml
delete mode 100644 lib/gitlab/database/load_balancing/session_map.rb
delete mode 100644 spec/lib/gitlab/database/load_balancing/session_map_spec.rb
diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 5faa3b2dcc1fd..6a0cbfc220da5 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -69,7 +69,7 @@ def self.safe_find_or_create_by!(*args, &block)
# to allow callers gracefully handling the errors to still complete within
# the 5s target duration of a low urgency request.
def self.with_fast_read_statement_timeout(timeout_ms = 4500)
- ::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer).fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
transaction(requires_new: true) do # rubocop:disable Performance/ActiveRecordSubtransactions
connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
diff --git a/app/models/ci/runner.rb b/app/models/ci/runner.rb
index e6d8c9726bee1..b51bf6f9b3999 100644
--- a/app/models/ci/runner.rb
+++ b/app/models/ci/runner.rb
@@ -469,7 +469,7 @@ def heartbeat
# not want to upgrade database connection proxy to use the primary
# database after heartbeat write happens.
#
- ::Gitlab::Database::LoadBalancing::SessionMap.current(connection.load_balancer).without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
values = { contacted_at: Time.current, creation_state: :finished }
merge_cache_attributes(values)
diff --git a/app/models/ci/runner_manager.rb b/app/models/ci/runner_manager.rb
index 8081d2a7b0770..3a2920172a3f3 100644
--- a/app/models/ci/runner_manager.rb
+++ b/app/models/ci/runner_manager.rb
@@ -137,7 +137,7 @@ def heartbeat(values, update_contacted_at: true)
# not want to upgrade database connection proxy to use the primary
# database after heartbeat write happens.
#
- ::Gitlab::Database::LoadBalancing::SessionMap.current(connection.load_balancer).without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
values = values&.slice(:version, :revision, :platform, :architecture, :ip_address, :config, :executor) || {}
values.merge!(contacted_at: Time.current, creation_state: :finished) if update_contacted_at
diff --git a/app/models/concerns/packages/downloadable.rb b/app/models/concerns/packages/downloadable.rb
index ad9bee6901dfa..011f5ddda9c4d 100644
--- a/app/models/concerns/packages/downloadable.rb
+++ b/app/models/concerns/packages/downloadable.rb
@@ -5,7 +5,7 @@ module Downloadable
extend ActiveSupport::Concern
def touch_last_downloaded_at
- ::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer).without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
update_column(:last_downloaded_at, Time.zone.now)
end
end
diff --git a/app/models/concerns/use_sql_function_for_primary_key_lookups.rb b/app/models/concerns/use_sql_function_for_primary_key_lookups.rb
index 6b7b7432dcc25..1e1dea42225ce 100644
--- a/app/models/concerns/use_sql_function_for_primary_key_lookups.rb
+++ b/app/models/concerns/use_sql_function_for_primary_key_lookups.rb
@@ -58,7 +58,7 @@ def try_replace_with_function_call(arel)
return unless verification_arel.ast == arel.ast
if table_name == "namespaces" && Feature.enabled?(:log_sql_function_namespace_lookups, Feature.current_request)
- using_primary = Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer).use_primary?
+ using_primary = Gitlab::Database::LoadBalancing::Session.current.use_primary?
Gitlab::AppLogger.info(
message: "Namespaces lookup using function",
backtrace: caller,
diff --git a/app/models/concerns/users/visitable.rb b/app/models/concerns/users/visitable.rb
index 854561862bc5c..029d60d61eeda 100644
--- a/app/models/concerns/users/visitable.rb
+++ b/app/models/concerns/users/visitable.rb
@@ -48,8 +48,7 @@ def self.frecent_visits_scores(user_id:, limit:)
LIMIT #{limit}
SQL
- ::Gitlab::Database::LoadBalancing::SessionMap
- .current(connection.load_balancer).fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
connection.execute(sql).to_a
end
end
diff --git a/app/models/internal_id.rb b/app/models/internal_id.rb
index 7ecc2aed4d937..5e78f46862182 100644
--- a/app/models/internal_id.rb
+++ b/app/models/internal_id.rb
@@ -191,7 +191,7 @@ def initial_value(subject, scope)
# `init` computes the maximum based on actual records. We use the
# primary to make sure we have up to date results
- Gitlab::Database::LoadBalancing::SessionMap.current(subject.load_balancer).use_primary do
+ Gitlab::Database::LoadBalancing::Session.current.use_primary do
instance = subject.is_a?(::Class) ? nil : subject
init.call(instance, scope) || 0
diff --git a/app/models/project.rb b/app/models/project.rb
index 52ff49814f4e6..a9428e67382ac 100644
--- a/app/models/project.rb
+++ b/app/models/project.rb
@@ -1491,8 +1491,7 @@ def log_import_activity(job_id, type: :import)
job_type = type.to_s.capitalize
if job_id
- use_primary = ::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer).use_primary?
- Gitlab::AppLogger.info("#{job_type} job scheduled for #{full_path} with job ID #{job_id} (primary: #{use_primary}).")
+ Gitlab::AppLogger.info("#{job_type} job scheduled for #{full_path} with job ID #{job_id} (primary: #{::Gitlab::Database::LoadBalancing::Session.current.use_primary?}).")
else
Gitlab::AppLogger.error("#{job_type} job failed to create for #{full_path}.")
end
diff --git a/app/models/sent_notification.rb b/app/models/sent_notification.rb
index 2cdb013e256e8..10765f826977e 100644
--- a/app/models/sent_notification.rb
+++ b/app/models/sent_notification.rb
@@ -45,7 +45,7 @@ def record(noteable, recipient_id, reply_key = self.reply_key, attrs = {})
# Non-sticky write is used as `.record` is only used in ActionMailer
# where there are no queries to SentNotification.
- ::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer).without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
create(attrs)
end
end
diff --git a/app/services/ci/click_house/data_ingestion/finished_pipelines_sync_service.rb b/app/services/ci/click_house/data_ingestion/finished_pipelines_sync_service.rb
index 4f546224a9acd..6570523a96325 100644
--- a/app/services/ci/click_house/data_ingestion/finished_pipelines_sync_service.rb
+++ b/app/services/ci/click_house/data_ingestion/finished_pipelines_sync_service.rb
@@ -52,7 +52,7 @@ def execute
# Prevent parallel jobs
in_lock("#{self.class.name.underscore}/worker/#{@worker_index}", ttl: MAX_TTL, retries: 0) do
- ::Gitlab::Database::LoadBalancing::SessionMap.without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
report = insert_new_finished_pipelines
ServiceResponse.success(payload: report.merge(service_payload))
diff --git a/app/services/ci/register_job_service.rb b/app/services/ci/register_job_service.rb
index 044b7fcc4fdbe..54b43f369349e 100644
--- a/app/services/ci/register_job_service.rb
+++ b/app/services/ci/register_job_service.rb
@@ -144,7 +144,7 @@ def retrieve_queue(queue_query_proc)
# We want to reset a load balancing session to discard the side
# effects of writes that could have happened prior to this moment.
#
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
@metrics.observe_queue_time(:retrieve, @runner.runner_type) do
queue_query_proc.call
diff --git a/app/services/issuable_base_service.rb b/app/services/issuable_base_service.rb
index f2ec829335109..28abf33ea908d 100644
--- a/app/services/issuable_base_service.rb
+++ b/app/services/issuable_base_service.rb
@@ -274,7 +274,7 @@ def handle_description_updated(issuable)
# rubocop:disable Metrics/AbcSize -- Method is only slightly over the limit due to decomposition method
def update(issuable)
- ::Gitlab::Database::LoadBalancing::SessionMap.current(issuable.load_balancer).use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
old_associations = associations_before_update(issuable)
diff --git a/app/services/personal_access_tokens/last_used_service.rb b/app/services/personal_access_tokens/last_used_service.rb
index 536d7fd90c1a4..4e402a44dec70 100644
--- a/app/services/personal_access_tokens/last_used_service.rb
+++ b/app/services/personal_access_tokens/last_used_service.rb
@@ -18,9 +18,8 @@ def execute
# would be updated when using #touch).
return unless update?
- lb = @personal_access_token.load_balancer
try_obtain_lease do
- ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
@personal_access_token.update_column(:last_used_at, Time.zone.now)
end
end
diff --git a/app/services/users/activity_service.rb b/app/services/users/activity_service.rb
index 5749df9013f64..8f3c2ea313729 100644
--- a/app/services/users/activity_service.rb
+++ b/app/services/users/activity_service.rb
@@ -20,8 +20,7 @@ def execute
return unless user
return if user.last_activity_on == Date.today
- ::Gitlab::Database::LoadBalancing::SessionMap.current(user.load_balancer)
- .without_sticky_writes { record_activity }
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes { record_activity }
end
private
diff --git a/app/workers/authorized_project_update/user_refresh_from_replica_worker.rb b/app/workers/authorized_project_update/user_refresh_from_replica_worker.rb
index e86e39a682507..b5e3332d39d24 100644
--- a/app/workers/authorized_project_update/user_refresh_from_replica_worker.rb
+++ b/app/workers/authorized_project_update/user_refresh_from_replica_worker.rb
@@ -30,9 +30,7 @@ def perform(user_id)
# does not allow us to deduplicate these jobs.
# https://gitlab.com/gitlab-org/gitlab/-/issues/325291
def use_replica_if_available(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord])
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
def project_authorizations_needs_refresh?(user)
diff --git a/app/workers/container_expiration_policies/cleanup_container_repository_worker.rb b/app/workers/container_expiration_policies/cleanup_container_repository_worker.rb
index c129b4ae213c2..44aa8081e1f1f 100644
--- a/app/workers/container_expiration_policies/cleanup_container_repository_worker.rb
+++ b/app/workers/container_expiration_policies/cleanup_container_repository_worker.rb
@@ -185,9 +185,7 @@ def project
end
def use_replica_if_available(&blk)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord])
- .use_replicas_for_read_queries(&blk)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&blk)
end
end
end
diff --git a/app/workers/container_expiration_policy_worker.rb b/app/workers/container_expiration_policy_worker.rb
index 0559a5716e0f1..80499aff43164 100644
--- a/app/workers/container_expiration_policy_worker.rb
+++ b/app/workers/container_expiration_policy_worker.rb
@@ -51,9 +51,7 @@ def log_counts
# not perfomed with a delay
# https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63635#note_603771207
def use_replica_if_available(&blk)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .current(ContainerRepository.load_balancer)
- .use_replicas_for_read_queries(&blk)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&blk)
end
def process_stale_ongoing_cleanups
diff --git a/app/workers/dependency_proxy/image_ttl_group_policy_worker.rb b/app/workers/dependency_proxy/image_ttl_group_policy_worker.rb
index e140d468aff56..5a59c2d23c93e 100644
--- a/app/workers/dependency_proxy/image_ttl_group_policy_worker.rb
+++ b/app/workers/dependency_proxy/image_ttl_group_policy_worker.rb
@@ -43,9 +43,7 @@ def log_counts
end
def use_replica_if_available(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord])
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
end
end
diff --git a/app/workers/packages/cleanup_package_registry_worker.rb b/app/workers/packages/cleanup_package_registry_worker.rb
index 31947c42fdc8e..50036923e943a 100644
--- a/app/workers/packages/cleanup_package_registry_worker.rb
+++ b/app/workers/packages/cleanup_package_registry_worker.rb
@@ -53,9 +53,7 @@ def log_counts
end
def use_replica_if_available(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord])
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
end
end
diff --git a/config/feature_flags/gitlab_com_derisk/use_load_balancing_session_map.yml b/config/feature_flags/gitlab_com_derisk/use_load_balancing_session_map.yml
deleted file mode 100644
index 9f85f9d17c7ef..0000000000000
--- a/config/feature_flags/gitlab_com_derisk/use_load_balancing_session_map.yml
+++ /dev/null
@@ -1,9 +0,0 @@
----
-name: use_load_balancing_session_map
-feature_issue_url: https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/3834
-introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/166986
-rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/496505
-milestone: '17.5'
-group: group::scalability
-type: gitlab_com_derisk
-default_enabled: false
diff --git a/ee/app/services/click_house/data_ingestion/ci_finished_builds_sync_service.rb b/ee/app/services/click_house/data_ingestion/ci_finished_builds_sync_service.rb
index 79c60b8553be7..a2b2be83fbed1 100644
--- a/ee/app/services/click_house/data_ingestion/ci_finished_builds_sync_service.rb
+++ b/ee/app/services/click_house/data_ingestion/ci_finished_builds_sync_service.rb
@@ -35,7 +35,7 @@ def execute
# Prevent parallel jobs
in_lock("#{self.class.name.underscore}/worker/#{@worker_index}", ttl: MAX_TTL, retries: 0) do
- ::Gitlab::Database::LoadBalancing::SessionMap.without_sticky_writes do
+ ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
report = insert_new_finished_builds
ServiceResponse.success(payload: report.merge(service_payload))
diff --git a/ee/lib/analytics/group_activity_calculator.rb b/ee/lib/analytics/group_activity_calculator.rb
index 6b83d916dad8f..3a40d1e0da4d8 100644
--- a/ee/lib/analytics/group_activity_calculator.rb
+++ b/ee/lib/analytics/group_activity_calculator.rb
@@ -12,19 +12,19 @@ def initialize(group, current_user)
end
def issues_count
- @issues_count ||= fetch_cached(:issues, Issue) do
+ @issues_count ||= fetch_cached(:issues) do
IssuesFinder.new(@current_user, issuable_params).execute.limit(RECENT_COUNT_LIMIT).reorder(nil).count # rubocop:disable CodeReuse/ActiveRecord
end
end
def merge_requests_count
- @merge_requests_count ||= fetch_cached(:merge_requests, MergeRequest) do
+ @merge_requests_count ||= fetch_cached(:merge_requests) do
MergeRequestsFinder.new(@current_user, issuable_params).execute.limit(RECENT_COUNT_LIMIT).reorder(nil).count # rubocop:disable CodeReuse/ActiveRecord
end
end
def new_members_count
- @new_members_count ||= fetch_cached(:new_members, Member) do
+ @new_members_count ||= fetch_cached(:new_members) do
GroupMembersFinder.new(
@group,
@current_user,
@@ -44,11 +44,9 @@ def issuable_params
attempt_project_search_optimizations: true }
end
- def fetch_cached(type, klass, &block)
+ def fetch_cached(type, &block)
Rails.cache.fetch(cache_key(type), CACHE_OPTIONS) do
- ::Gitlab::Database::LoadBalancing::SessionMap
- .current(klass.load_balancer)
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end.to_i
end
diff --git a/ee/lib/gitlab/license_scanning/package_licenses.rb b/ee/lib/gitlab/license_scanning/package_licenses.rb
index 094e30f3957ed..d97b1f3526f0b 100644
--- a/ee/lib/gitlab/license_scanning/package_licenses.rb
+++ b/ee/lib/gitlab/license_scanning/package_licenses.rb
@@ -62,9 +62,7 @@ def fetch
attr_reader :components, :component_data, :all_records, :project
def use_replica_if_available(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord])
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
# set the default license of all components to an unknown license.
diff --git a/ee/spec/lib/gitlab/license_scanning/package_licenses_spec.rb b/ee/spec/lib/gitlab/license_scanning/package_licenses_spec.rb
index 0c858086c50c3..109573230aceb 100644
--- a/ee/spec/lib/gitlab/license_scanning/package_licenses_spec.rb
+++ b/ee/spec/lib/gitlab/license_scanning/package_licenses_spec.rb
@@ -242,12 +242,8 @@
context 'when load balancing enabled', :db_load_balancing do
it 'uses the replica' do
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries)
+ .and_call_original
fetch
end
diff --git a/ee/spec/lib/gitlab/llm/ai_message_spec.rb b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
index d2006d4bcaeef..933622679528c 100644
--- a/ee/spec/lib/gitlab/llm/ai_message_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
@@ -41,8 +41,7 @@
end
it 'generates id' do
- # TODO: revert changes to .once when removing use_load_balancing_session_map feature flag
- allow(SecureRandom).to receive(:uuid).at_least(:once).and_return('123')
+ allow(SecureRandom).to receive(:uuid).once.and_return('123')
expect(described_class.new(data.except(:id)).id).to eq('123')
end
diff --git a/lib/api/ci/helpers/runner.rb b/lib/api/ci/helpers/runner.rb
index fb57ce8ef02c8..fa0724061e888 100644
--- a/lib/api/ci/helpers/runner.rb
+++ b/lib/api/ci/helpers/runner.rb
@@ -101,10 +101,7 @@ def authenticate_job!(heartbeat_runner: false)
end
def authenticate_job_via_dependent_job!
- # Use primary for both main and ci database as authenticating in the scope of runners will load
- # Ci::Build model and other standard authn related models like License, Project and User.
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord]).use_primary { authenticate! }
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary { authenticate! }
forbidden! unless current_job
forbidden! unless can?(current_user, :read_build, current_job)
diff --git a/lib/api/maven_packages.rb b/lib/api/maven_packages.rb
index ed76f15ddfa10..e0ffc0f358050 100644
--- a/lib/api/maven_packages.rb
+++ b/lib/api/maven_packages.rb
@@ -266,8 +266,7 @@ def find_and_present_package_file(package, file_name, format, params)
# so we need to skip the second FIPS check here.
file_name, format = extract_format(params[:file_name], skip_fips_check: true)
- lb = ::ApplicationRecord.load_balancer
- ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
result = ::Packages::Maven::FindOrCreatePackageService
.new(user_project, current_user, params.merge(build: current_authenticated_job)).execute
diff --git a/lib/gitlab/auth.rb b/lib/gitlab/auth.rb
index 642486f2f6884..46dd5c38af4aa 100644
--- a/lib/gitlab/auth.rb
+++ b/lib/gitlab/auth.rb
@@ -471,8 +471,7 @@ def non_admin_available_scopes
end
def find_build_by_token(token)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord]).use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
::Ci::AuthJobFinder.new(token: token).execute
end
end
diff --git a/lib/gitlab/checks/matching_merge_request.rb b/lib/gitlab/checks/matching_merge_request.rb
index 27949f0ae7769..d06448561055b 100644
--- a/lib/gitlab/checks/matching_merge_request.rb
+++ b/lib/gitlab/checks/matching_merge_request.rb
@@ -28,7 +28,7 @@ def match?
# the write location to ensure the replica can make this query.
# Adding use_primary_on_empty_location: true for extra precaution in case there happens to be
# no LSN saved for the project then we will use the primary.
- track_session_metrics(::ApplicationRecord.load_balancer) do
+ track_session_metrics do
::ApplicationRecord.sticking.find_caught_up_replica(:project, @project.id, use_primary_on_empty_location: true)
end
@@ -42,12 +42,12 @@ def match?
private
- def track_session_metrics(lb)
- before = ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?
+ def track_session_metrics
+ before = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
yield
- after = ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?
+ after = ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
increment_attempt_count
diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index b01c187b8cc20..e4c694654230f 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -214,10 +214,10 @@ def self.all_uncached(&block)
# Calls to #uncached only disable caching for the current connection. Since the load balancer
# can potentially upgrade from read to read-write mode (using a different connection), we specify
# up-front that we'll explicitly use the primary for the duration of the operation.
- base_models = database_base_models_using_load_balancing.values
- base_models.reduce(block) do |blk, model|
- -> { Gitlab::Database::LoadBalancing::SessionMap.current(model.load_balancer).use_primary { model.uncached(&blk) } }
- end.call
+ Gitlab::Database::LoadBalancing::Session.current.use_primary do
+ base_models = database_base_models_using_load_balancing.values
+ base_models.reduce(block) { |blk, model| -> { model.uncached(&blk) } }.call
+ end
end
def self.allow_cross_joins_across_databases(url:)
diff --git a/lib/gitlab/database/consistency.rb b/lib/gitlab/database/consistency.rb
index d54c3881f5f22..17c16640e4c95 100644
--- a/lib/gitlab/database/consistency.rb
+++ b/lib/gitlab/database/consistency.rb
@@ -13,9 +13,8 @@ class Consistency
# require read consistency after recent writes.
#
def self.with_read_consistency(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions(Gitlab::Database::LoadBalancing.base_models)
- .use_primary(&block)
+ ::Gitlab::Database::LoadBalancing::Session
+ .current.use_primary(&block)
end
end
end
diff --git a/lib/gitlab/database/health_status/indicators/write_ahead_log.rb b/lib/gitlab/database/health_status/indicators/write_ahead_log.rb
index 253af7c9f0dec..1614b17df4808 100644
--- a/lib/gitlab/database/health_status/indicators/write_ahead_log.rb
+++ b/lib/gitlab/database/health_status/indicators/write_ahead_log.rb
@@ -59,7 +59,7 @@ def enabled?
# Returns number of WAL segments pending archival
def pending_wal_count
- Gitlab::Database::LoadBalancing::SessionMap.current(connection.load_balancer).use_primary do
+ Gitlab::Database::LoadBalancing::Session.current.use_primary do
connection.execute(PENDING_WAL_COUNT_SQL).to_a.first&.fetch('pending_wal_count')
end
end
diff --git a/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb
index bc4d33291039f..5f9416fb4db16 100644
--- a/lib/gitlab/database/load_balancing.rb
+++ b/lib/gitlab/database/load_balancing.rb
@@ -34,10 +34,6 @@ def self.primary_only?
each_load_balancer.all?(&:primary_only?)
end
- def self.primary?(name)
- each_load_balancer.find { |c| c.name == name }&.primary_only?
- end
-
def self.release_hosts
each_load_balancer(&:release_host)
end
diff --git a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
index 7758be013ae9c..0e97f49eb4900 100644
--- a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
+++ b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
@@ -13,7 +13,7 @@ def self.wrapper
inner.call
ensure
::Gitlab::Database::LoadBalancing.release_hosts
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
end
end
end
diff --git a/lib/gitlab/database/load_balancing/connection_proxy.rb b/lib/gitlab/database/load_balancing/connection_proxy.rb
index 00a8a05321ef9..67cde3e687d3a 100644
--- a/lib/gitlab/database/load_balancing/connection_proxy.rb
+++ b/lib/gitlab/database/load_balancing/connection_proxy.rb
@@ -131,7 +131,7 @@ def write_using_load_balancer(...)
private
def current_session
- ::Gitlab::Database::LoadBalancing::SessionMap.current(@load_balancer)
+ ::Gitlab::Database::LoadBalancing::Session.current
end
def track_read_only_transaction!
diff --git a/lib/gitlab/database/load_balancing/rack_middleware.rb b/lib/gitlab/database/load_balancing/rack_middleware.rb
index 98ed683ea9731..554205c8dcce9 100644
--- a/lib/gitlab/database/load_balancing/rack_middleware.rb
+++ b/lib/gitlab/database/load_balancing/rack_middleware.rb
@@ -45,17 +45,18 @@ def find_caught_up_replica(env)
# Determine if we need to stick after handling a request.
def stick_if_necessary(env)
+ return unless ::Gitlab::Database::LoadBalancing::Session.current.performed_write?
+
namespaces_and_ids = sticking_namespaces(env)
namespaces_and_ids.each do |sticking, namespace, id|
- lb = sticking.load_balancer
- sticking.stick(namespace, id) if ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).performed_write?
+ sticking.stick(namespace, id)
end
end
def clear
::Gitlab::Database::LoadBalancing.release_hosts
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
end
# Determines the sticking namespace and identifier based on the Rack
diff --git a/lib/gitlab/database/load_balancing/session.rb b/lib/gitlab/database/load_balancing/session.rb
index d6c7cf94fb274..3682c9265c221 100644
--- a/lib/gitlab/database/load_balancing/session.rb
+++ b/lib/gitlab/database/load_balancing/session.rb
@@ -56,7 +56,6 @@ def ignore_writes(&block)
ensure
@ignore_writes = false
end
- alias_method :without_sticky_writes, :ignore_writes
# Indicates that the read SQL statements from anywhere inside this
# blocks should use a replica, regardless of the current primary
diff --git a/lib/gitlab/database/load_balancing/session_map.rb b/lib/gitlab/database/load_balancing/session_map.rb
deleted file mode 100644
index e254da8460a8a..0000000000000
--- a/lib/gitlab/database/load_balancing/session_map.rb
+++ /dev/null
@@ -1,128 +0,0 @@
-# frozen_string_literal: true
-
-module Gitlab
- module Database
- module LoadBalancing
- class SessionMap
- CACHE_KEY = :gitlab_load_balancer_session_map
-
- InvalidLoadBalancerNameError = Class.new(StandardError)
-
- # lb - Gitlab::Database::LoadBalancing::LoadBalancer instance
- def self.current(lb)
- return cached_instance.lookup(lb) if use_session_map?
-
- Session.current
- end
-
- # models - Array<ActiveRecord::Base>
- def self.with_sessions(models)
- dbs = models.map { |m| m.load_balancer.name }.uniq
- dbs.each { |db| cached_instance.validate_db_name(db) }
- ScopedSessions.new(dbs, cached_instance.session_map)
- end
-
- def self.clear_session
- return RequestStore.delete(CACHE_KEY) if use_session_map?
-
- Session.clear_session
- end
-
- def self.without_sticky_writes(&)
- return with_sessions(Gitlab::Database::LoadBalancing.base_models).ignore_writes(&) if use_session_map?
-
- Session.without_sticky_writes(&)
- end
-
- def self.use_session_map?
- ::Feature.enabled?(:use_load_balancing_session_map, :current_request, type: :gitlab_com_derisk)
- rescue ActiveRecord::StatementInvalid,
- Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError
- # If the feature_gates table is missing, we should default to a false.
- # In a migration scope, we also rescue and default to false.
- false
- end
- private_class_method :use_session_map?
-
- def self.cached_instance
- RequestStore[CACHE_KEY] ||= new
- end
- private_class_method :cached_instance
-
- attr_reader :session_map
-
- def initialize
- @session_map = Gitlab::Database.all_database_names.to_h do |k|
- [k.to_sym, Gitlab::Database::LoadBalancing::Session.new]
- end
-
- @session_map[:primary] = Gitlab::Database::LoadBalancing::Session.new
- end
-
- def lookup(lb)
- name = lb.name
- validate_db_name(name)
- session_map[name]
- end
-
- def validate_db_name(db)
- # Allow :primary only for rake task db migrations as ActiveRecord::Tasks::PostgresqlDatabaseTasks calls
- # .establish_connection using a hash which resets the name from :main/:ci to :primary.
- # See
- # https://github.com/rails/rails/blob/v7.0.8.4/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb#L97
- #
- # In the case of derailed test in memory-on-boot job, the runtime is unknown.
- return if db == :primary && (Gitlab::Runtime.rake? || Gitlab::Runtime.safe_identify.nil?)
-
- # Disallow :primary usage outside of rake or unknown runtimes as the db config should be
- # main/ci/embedding/ci/geo.
- return if db != :primary && session_map[db]
-
- raise InvalidLoadBalancerNameError, "Invalid load balancer name #{db} in #{Gitlab::Runtime.safe_identify}."
- end
- end
-
- class ScopedSessions
- attr_reader :scoped_sessions
-
- def initialize(scope, session_map)
- @scope = scope
- @scoped_sessions = session_map.slice(*@scope).values
- end
-
- def use_primary!
- scoped_sessions.each(&:use_primary!)
- end
-
- def ignore_writes(&)
- nest_sessions(scoped_sessions, :without_sticky_writes, &)
- end
-
- def use_primary(&)
- nest_sessions(scoped_sessions, :use_primary, &)
- end
-
- def use_replicas_for_read_queries(&)
- nest_sessions(scoped_sessions, :use_replicas_for_read_queries, &)
- end
-
- def fallback_to_replicas_for_ambiguous_queries(&)
- nest_sessions(scoped_sessions, :fallback_to_replicas_for_ambiguous_queries, &)
- end
-
- private
-
- def nest_sessions(sessions, method, &block)
- if sessions.empty?
- yield if block
- else
- session = sessions.shift
- session.public_send(method) do # rubocop: disable GitlabSecurity/PublicSend -- methods are verified
- nest_sessions(sessions, method, &block)
- end
- end
- end
- end
- end
- end
-end
diff --git a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index e48a92bbb84a6..346d951414a8a 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -34,7 +34,19 @@ def load_balancing_enabled?(worker_class)
def set_data_consistency_locations!(job)
job['wal_locations'] = wal_locations_by_db_name
- job['wal_location_sources'] = wal_location_sources_by_db_name
+ job['wal_location_source'] = wal_location_source
+ end
+
+ def wal_location_source
+ if ::Gitlab::Database::LoadBalancing.primary_only? || uses_primary?
+ ::Gitlab::Database::LoadBalancing::ROLE_PRIMARY
+ else
+ ::Gitlab::Database::LoadBalancing::ROLE_REPLICA
+ end
+ end
+
+ def uses_primary?
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
end
end
end
diff --git a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
index 7babe6422ad14..7c3331b2a0eb4 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
+++ b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
@@ -22,9 +22,7 @@ def call(worker, job, _queue)
job['load_balancing_strategy'] = strategy.to_s
if use_primary?(strategy)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions(Gitlab::Database::LoadBalancing.base_models)
- .use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
elsif strategy == :retry
raise JobReplicaNotUpToDate, "Sidekiq job #{resolved_class} JID-#{job['jid']} couldn't use the replica. "\
"Replica was not up to date."
@@ -41,7 +39,7 @@ def call(worker, job, _queue)
def clear
::Gitlab::Database::LoadBalancing.release_hosts
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
end
def use_primary?(strategy)
diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb
index 744c215d16e14..df07f510c8d2f 100644
--- a/lib/gitlab/database/load_balancing/sticking.rb
+++ b/lib/gitlab/database/load_balancing/sticking.rb
@@ -10,8 +10,6 @@ class Sticking
# the primary.
EXPIRATION = 30
- attr_reader :load_balancer
-
def initialize(load_balancer)
@load_balancer = load_balancer
end
@@ -37,7 +35,7 @@ def find_caught_up_replica(namespace, id, use_primary_on_failure: true, use_prim
!use_primary_on_empty_location
end
- use_primary! if !result && use_primary_on_failure
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary! if !result && use_primary_on_failure
result
end
@@ -48,7 +46,7 @@ def stick(namespace, id)
with_primary_write_location do |location|
set_write_location_for(namespace, id, location)
end
- use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
end
def bulk_stick(namespace, ids)
@@ -58,7 +56,7 @@ def bulk_stick(namespace, ids)
end
end
- use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
end
private
@@ -102,10 +100,6 @@ def redis_key_for(namespace, id)
def with_redis(&block)
Gitlab::Redis::DbLoadBalancing.with(&block)
end
-
- def use_primary!
- ::Gitlab::Database::LoadBalancing::SessionMap.current(@load_balancer).use_primary!
- end
end
end
end
diff --git a/lib/gitlab/database/load_balancing/wal_tracking_sender.rb b/lib/gitlab/database/load_balancing/wal_tracking_sender.rb
index bf8502bad84c6..ed0af58d06441 100644
--- a/lib/gitlab/database/load_balancing/wal_tracking_sender.rb
+++ b/lib/gitlab/database/load_balancing/wal_tracking_sender.rb
@@ -18,31 +18,12 @@ def wal_location_for(load_balancer)
# When only using the primary there's no need for any WAL queries.
return if load_balancer.primary_only?
- if SessionMap.current(load_balancer).use_primary?
+ if Session.current.use_primary?
load_balancer.primary_write_location
else
load_balancer.host&.database_replica_location || load_balancer.primary_write_location
end
end
-
- def wal_location_sources_by_db_name
- {}.tap do |locations|
- ::Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- if (location = wal_location_source(lb))
- locations[lb.name] = location
- end
- end
- end
- end
-
- def wal_location_source(lb)
- if ::Gitlab::Database::LoadBalancing.primary?(lb.name) ||
- ::Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?
- ::Gitlab::Database::LoadBalancing::ROLE_PRIMARY
- else
- ::Gitlab::Database::LoadBalancing::ROLE_REPLICA
- end
- end
end
end
end
diff --git a/lib/gitlab/database/partitioning/partition_manager.rb b/lib/gitlab/database/partitioning/partition_manager.rb
index 7cbb8e608f1b9..a07534b3e46ad 100644
--- a/lib/gitlab/database/partitioning/partition_manager.rb
+++ b/lib/gitlab/database/partitioning/partition_manager.rb
@@ -196,7 +196,7 @@ def analyze_interval
end
def primary_transaction(statement_timeout: nil)
- Gitlab::Database::LoadBalancing::SessionMap.current(connection.load_balancer).use_primary do
+ Gitlab::Database::LoadBalancing::Session.current.use_primary do
connection.transaction(requires_new: false) do
if statement_timeout.present?
connection.execute(
diff --git a/lib/gitlab/database/postgres_autovacuum_activity.rb b/lib/gitlab/database/postgres_autovacuum_activity.rb
index 4523519b8b5ca..6a80c631abbc0 100644
--- a/lib/gitlab/database/postgres_autovacuum_activity.rb
+++ b/lib/gitlab/database/postgres_autovacuum_activity.rb
@@ -9,9 +9,7 @@ class PostgresAutovacuumActivity < SharedModel
scope :wraparound_prevention, -> { where(wraparound_prevention: true) }
def self.for_tables(tables)
- Gitlab::Database::LoadBalancing::SessionMap
- .current(connection.load_balancer)
- .use_primary do
+ Gitlab::Database::LoadBalancing::Session.current.use_primary do
# calling `.to_a` here to execute the query in the primary's scope
# and to avoid having the scope chained and re-executed
#
diff --git a/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing.rb b/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing.rb
index dacfb9a30e5fd..851750163af18 100644
--- a/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing.rb
+++ b/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing.rb
@@ -27,11 +27,8 @@ def execute_update(subscription_id, event, object)
# and others expect the original payload.
return super(subscription_id, event, object) unless wrapped_payload?(object)
- if use_primary?(object[KEY_WAL_LOCATIONS])
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions(Gitlab::Database::LoadBalancing.base_models)
- .use_primary!
- end
+ wal_locations = object[KEY_WAL_LOCATIONS]
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary! if use_primary?(wal_locations)
super(subscription_id, event, object[KEY_PAYLOAD])
end
diff --git a/lib/gitlab/import_export/json/streaming_serializer.rb b/lib/gitlab/import_export/json/streaming_serializer.rb
index c6b47448c8a13..9c684d20e09fb 100644
--- a/lib/gitlab/import_export/json/streaming_serializer.rb
+++ b/lib/gitlab/import_export/json/streaming_serializer.rb
@@ -249,9 +249,7 @@ def batch_size(relation_name)
end
def read_from_replica_if_available(&block)
- ::Gitlab::Database::LoadBalancing::SessionMap
- .with_sessions(Gitlab::Database::LoadBalancing.base_models)
- .use_replicas_for_read_queries(&block)
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
end
def before_read_callback(record)
diff --git a/spec/initializers/load_balancing_spec.rb b/spec/initializers/load_balancing_spec.rb
index 9b7459bd2fcaa..eddedcb2f386f 100644
--- a/spec/initializers/load_balancing_spec.rb
+++ b/spec/initializers/load_balancing_spec.rb
@@ -73,7 +73,7 @@ def simulate_puma_worker
it 'makes a read query successfully' do
# Clear any previous sticky writes
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
initialize_load_balancer
diff --git a/spec/lib/feature_spec.rb b/spec/lib/feature_spec.rb
index 8675c191280b9..b2bc24970a1df 100644
--- a/spec/lib/feature_spec.rb
+++ b/spec/lib/feature_spec.rb
@@ -560,11 +560,6 @@ def wrap_all_methods_with_flag_check(lb, flag)
context 'and feature has been disabled' do
before do
- # TODO: revert changes to .once when removing use_load_balancing_session_map feature flag
- # We stub SessionMap.current since it uses a feature flag. Within this spec, all feature flags
- # definitions are missing due to the stubbed Feature::Definition.definitions method.
- allow(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:current).and_return(Gitlab::Database::LoadBalancing::Session.current)
described_class.disable(:my_feature_flag)
end
diff --git a/spec/lib/gitlab/auth/token_expiration_banner_spec.rb b/spec/lib/gitlab/auth/token_expiration_banner_spec.rb
index b1f40c61df604..d5175e9d168f4 100644
--- a/spec/lib/gitlab/auth/token_expiration_banner_spec.rb
+++ b/spec/lib/gitlab/auth/token_expiration_banner_spec.rb
@@ -13,8 +13,7 @@
# for specs, we need to reset it each time
before do
described_class.instance_variable_set(:@show_token_expiration_banner, nil)
- # TODO: revert changes to .twice when removing use_load_balancing_session_map feature flag
- allow(Gitlab).to receive(:version_info).at_least(:twice).and_return(gitlab_version)
+ allow(Gitlab).to receive(:version_info).twice.and_return(gitlab_version)
end
it { is_expected.to be(false) }
diff --git a/spec/lib/gitlab/auth_spec.rb b/spec/lib/gitlab/auth_spec.rb
index 1300218ff97a7..33e0a69dc5562 100644
--- a/spec/lib/gitlab/auth_spec.rb
+++ b/spec/lib/gitlab/auth_spec.rb
@@ -1047,8 +1047,7 @@ def operation
it 'executes query using primary database' do
expect(Ci::Build).to receive(:find_by_token).with(build.token).and_wrap_original do |m, *args|
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(Ci::Build.load_balancer).use_primary?)
- .to eq(true)
+ expect(::Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
m.call(*args)
end
diff --git a/spec/lib/gitlab/checks/matching_merge_request_spec.rb b/spec/lib/gitlab/checks/matching_merge_request_spec.rb
index cd69ccacff613..85fe669b8cf72 100644
--- a/spec/lib/gitlab/checks/matching_merge_request_spec.rb
+++ b/spec/lib/gitlab/checks/matching_merge_request_spec.rb
@@ -32,7 +32,7 @@
end
context 'with load balancing enabled', :redis do
- let(:session) { ::Gitlab::Database::LoadBalancing::SessionMap.current(project.load_balancer) }
+ let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
before do
# Need to mock as though we actually have replicas
@@ -43,7 +43,7 @@
# Put some sticking position for the primary in Redis
::ApplicationRecord.sticking.stick(:project, project.id)
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
# Mock the load balancer result since we don't actually have real replicas to match against
expect(::ApplicationRecord.load_balancer)
@@ -59,7 +59,7 @@
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
context 'when any secondary is caught up' do
diff --git a/spec/lib/gitlab/database/bulk_update_spec.rb b/spec/lib/gitlab/database/bulk_update_spec.rb
index 7f5c43a571840..2f0859dba743d 100644
--- a/spec/lib/gitlab/database/bulk_update_spec.rb
+++ b/spec/lib/gitlab/database/bulk_update_spec.rb
@@ -100,14 +100,10 @@
end
before do
- db_config = ActiveRecord::Base.connection_db_config
+ configuration_hash = ActiveRecord::Base.connection_db_config.configuration_hash
ActiveRecord::Base.establish_connection( # rubocop: disable Database/EstablishConnection
- ActiveRecord::DatabaseConfigurations::HashConfig.new(
- db_config.env_name,
- db_config.name,
- db_config.configuration_hash.merge(prepared_statements: prepared_statements)
- )
+ configuration_hash.merge(prepared_statements: prepared_statements)
)
end
diff --git a/spec/lib/gitlab/database/consistency_spec.rb b/spec/lib/gitlab/database/consistency_spec.rb
index 2043c5fa18f05..a42d97dc1cbc7 100644
--- a/spec/lib/gitlab/database/consistency_spec.rb
+++ b/spec/lib/gitlab/database/consistency_spec.rb
@@ -4,15 +4,15 @@
RSpec.describe Gitlab::Database::Consistency do
let(:session) do
- Gitlab::Database::LoadBalancing::SessionMap.current(ApplicationRecord.load_balancer)
+ Gitlab::Database::LoadBalancing::Session.current
end
before do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '.with_read_consistency' do
diff --git a/spec/lib/gitlab/database/health_status/indicators/write_ahead_log_spec.rb b/spec/lib/gitlab/database/health_status/indicators/write_ahead_log_spec.rb
index cb849a7662156..aa2aee4f94a8c 100644
--- a/spec/lib/gitlab/database/health_status/indicators/write_ahead_log_spec.rb
+++ b/spec/lib/gitlab/database/health_status/indicators/write_ahead_log_spec.rb
@@ -44,8 +44,7 @@
end
it 'uses primary database' do
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(connection.load_balancer))
- .to receive(:use_primary).and_yield
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary).and_yield
evaluate
end
diff --git a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
index ed1e083657190..2c69f39f131c5 100644
--- a/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb
@@ -6,7 +6,7 @@
describe '.wrapper' do
it 'releases the connection and clears the session' do
expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:clear_session)
+ expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
described_class.wrapper.call(nil, -> {})
end
@@ -14,7 +14,7 @@
context 'with an exception' do
it 'releases the connection and clears the session' do
expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:clear_session)
+ expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
expect do
described_class.wrapper.call(nil, -> { raise 'test_exception' })
diff --git a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
index 4d41fa5d7c613..a2076f5b95037 100644
--- a/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
@@ -39,7 +39,7 @@
arel = double(:arel, locked: true)
session = Gitlab::Database::LoadBalancing::Session.new
- allow(Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).with(load_balancer)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
.and_return(session)
expect(session).to receive(:write!)
@@ -68,8 +68,8 @@
it 'runs the query on the primary and sticks to it' do
session = Gitlab::Database::LoadBalancing::Session.new
- allow(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:current).with(load_balancer).and_return(session)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+ .and_return(session)
expect(session).to receive(:write!)
expect(proxy).to receive(:write_using_load_balancer).with(name, 'foo')
@@ -96,8 +96,8 @@
end
end
- allow(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:current).with(model_class.load_balancer).and_return(session)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+ .and_return(session)
end
after do
@@ -137,8 +137,8 @@
let(:session) { Gitlab::Database::LoadBalancing::Session.new }
before do
- allow(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:current).with(load_balancer).and_return(session)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+ .and_return(session)
end
context 'session fallbacks ambiguous queries to replicas' do
@@ -231,7 +231,7 @@
let(:session) { double(:session) }
before do
- allow(Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).with(load_balancer)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
.and_return(session)
allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries?).and_return(true)
allow(session).to receive(:use_primary?).and_return(false)
@@ -259,7 +259,7 @@
let(:connection) { double(:connection) }
before do
- allow(Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).with(load_balancer)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
.and_return(session)
end
@@ -319,7 +319,7 @@
let(:connection) { double(:connection) }
before do
- allow(Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).with(load_balancer)
+ allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
.and_return(session)
end
diff --git a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
index 4ec32c574bb81..1ba35f74fb251 100644
--- a/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
@@ -16,7 +16,7 @@
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '#call' do
@@ -100,11 +100,8 @@
let(:warden) { warden_user }
before do
- Gitlab::Database::LoadBalancing.base_models.each do |model|
- allow(::Gitlab::Database::LoadBalancing::SessionMap.current(model.load_balancer))
- .to receive(:performed_write?)
- .and_return(write_performed)
- end
+ allow(::Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?)
+ .and_return(write_performed)
end
subject { middleware.stick_if_necessary(env) }
@@ -183,7 +180,7 @@
it 'clears the currently used host and session' do
session = spy(:session)
- stub_const('Gitlab::Database::LoadBalancing::SessionMap', session)
+ stub_const('Gitlab::Database::LoadBalancing::Session', session)
expect(Gitlab::Database::LoadBalancing).to receive(:release_hosts)
diff --git a/spec/lib/gitlab/database/load_balancing/session_map_spec.rb b/spec/lib/gitlab/database/load_balancing/session_map_spec.rb
deleted file mode 100644
index 95813e04fb28e..0000000000000
--- a/spec/lib/gitlab/database/load_balancing/session_map_spec.rb
+++ /dev/null
@@ -1,306 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::LoadBalancing::SessionMap, feature_category: :database do
- let(:lb) { ::ApplicationRecord.load_balancer }
-
- describe '.current' do
- let(:session) { Gitlab::Database::LoadBalancing::Session.new }
-
- before do
- described_class.clear_session
- end
-
- context 'when feature flag is disabled' do
- before do
- described_class.current(lb) # initialise a SessionMap
- stub_feature_flags(use_load_balancing_session_map: false)
- end
-
- it 'returns sessions from Gitlab::Database::LoadBalancing::Session.current' do
- expect(described_class.current(lb)).to eq(Gitlab::Database::LoadBalancing::Session.current)
- end
- end
-
- context 'when already initialised' do
- before do
- described_class.current(lb)
- end
-
- it 're-use memoized SessionMap' do
- expect(described_class).not_to receive(:new)
- described_class.current(lb)
- end
- end
-
- context 'when using a non-rake runtime' do
- before do
- allow_next_instance_of(described_class) do |inst|
- allow(inst).to receive(:lookup).and_return(session)
- end
- end
-
- it 'returns desired Session instance' do
- expect(described_class.current(lb)).to eq(session)
- end
- end
-
- context 'when using a rake runtime' do
- let(:pri_session) { Gitlab::Database::LoadBalancing::Session.new }
- let(:pri_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :primary) }
-
- before do
- allow(Gitlab::Runtime).to receive(:rake?).and_return(true)
- sm = described_class.new
- sm.session_map[:primary] = pri_session
- RequestStore[described_class::CACHE_KEY] = sm
- end
-
- it 'returns desired Session instance' do
- expect(described_class.current(pri_lb)).to eq(pri_session)
- end
- end
-
- context 'when receiving invalid db type' do
- using RSpec::Parameterized::TableSyntax
-
- let(:pri_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :primary) }
- let(:invalid_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :invalid) }
-
- where(:runtime, :db, :error) do
- :rake? | :primary | false
- :rake? | :invalid | true
-
- :puma? | :primary | true
- :puma? | :invalid | true
-
- :sidekiq? | :primary | true
- :sidekiq? | :invalid | true
-
- :test_suite? | :primary | true
- :test_suite? | :invalid | true
- end
-
- with_them do
- let(:lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: db) }
-
- before do
- allow(Gitlab::Runtime).to receive(runtime).and_return(true)
- allow(Gitlab::Runtime).to receive(:safe_identify).and_return(runtime)
- end
-
- subject(:current) { described_class.current(lb) }
-
- it 'handles invalid db' do
- if error
- expect do
- current
- end.to raise_error(instance_of(Gitlab::Database::LoadBalancing::SessionMap::InvalidLoadBalancerNameError))
- else
- expect(current).to be_instance_of(Gitlab::Database::LoadBalancing::Session)
- end
- end
- end
-
- it 'handles unknown runtimes' do
- allow(Gitlab::Runtime).to receive(:rake?).and_return(false)
- allow(Gitlab::Runtime).to receive(:safe_identify).and_return(nil)
-
- expect(described_class.current(pri_lb)).to be_instance_of(Gitlab::Database::LoadBalancing::Session)
- expect do
- described_class.current(invalid_lb)
- end.to raise_error(instance_of(Gitlab::Database::LoadBalancing::SessionMap::InvalidLoadBalancerNameError))
- end
- end
- end
-
- describe '.clear_session' do
- before do
- described_class.current(::ApplicationRecord.load_balancer)
- end
-
- it 'clears instance from RequestStore' do
- described_class.clear_session
-
- expect(RequestStore[described_class::CACHE_KEY]).to eq(nil)
- end
-
- context 'when feature flag is disabled' do
- before do
- described_class.current(lb) # initialise a SessionMap
- stub_feature_flags(use_load_balancing_session_map: false)
- end
-
- it 'clears session from Gitlab::Database::LoadBalancing::Session.current' do
- expect(Gitlab::Database::LoadBalancing::Session.current).not_to eq(nil)
-
- described_class.clear_session
-
- expect(RequestStore[Gitlab::Database::LoadBalancing::Session::CACHE_KEY]).to eq(nil)
- end
- end
- end
-
- describe '.without_sticky_writes' do
- let(:dbs) { [::ApplicationRecord, ::Ci::ApplicationRecord] }
- let(:names) { dbs.map { |m| m.load_balancer.name }.uniq }
-
- let(:scoped_session) { Gitlab::Database::LoadBalancing::ScopedSessions.new(dbs, {}) }
-
- before do
- described_class.clear_session
- # This makes the spec more robust in single-db scenarios
- allow(Gitlab::Database::LoadBalancing).to receive(:names).and_return([:main, :ci])
- described_class.current(::ApplicationRecord.load_balancer)
- end
-
- it 'initialises ScopedSessions with all valid lb names and calls ignore_writes' do
- expect(Gitlab::Database::LoadBalancing::ScopedSessions)
- .to receive(:new).with(names, RequestStore[described_class::CACHE_KEY].session_map).and_return(scoped_session)
-
- expect(scoped_session).to receive(:ignore_writes).and_yield
-
- described_class.without_sticky_writes do
- # exact logic for ignore_writes is tested in `.with_sessions` test suite
- end
- end
-
- context 'when feature flag is disabled' do
- before do
- stub_feature_flags(use_load_balancing_session_map: false)
- end
-
- it 'calls Gitlab::Database::LoadBalancing::Session instead' do
- expect(Gitlab::Database::LoadBalancing::Session).to receive(:without_sticky_writes).and_yield
-
- described_class.without_sticky_writes do
- # exact logic for ignore_writes is tested in `.with_sessions` test suite
- end
- end
- end
- end
-
- describe '.with_sessions' do
- let(:main_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :main) }
- let(:ci_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :ci) }
- let(:sec_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :sec) }
- let(:invalid_lb) { instance_double('Gitlab::Database::LoadBalancing::LoadBalancer', name: :invalid) }
-
- let(:main) { instance_double('ActiveRecord::Base', load_balancer: main_lb) }
- let(:ci) { instance_double('ActiveRecord::Base', load_balancer: ci_lb) }
- let(:sec) { instance_double('ActiveRecord::Base', load_balancer: sec_lb) }
- let(:invalid) { instance_double('ActiveRecord::Base', load_balancer: invalid_lb) }
-
- let(:all_dbs) { [main, ci, sec] }
- let(:scoped_dbs) { [main, ci] }
-
- before do
- described_class.clear_session
-
- # This makes the spec more robust in single-db scenarios
- allow(Gitlab::Database::LoadBalancing).to receive(:names).and_return(all_dbs)
- end
-
- subject(:with_sessions) { described_class.with_sessions(scoped_dbs) }
-
- it 'returns a ScopedSession instance' do
- expect(with_sessions)
- .to be_an_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions)
- end
-
- it 'validates invalid dbs' do
- expect do
- described_class.with_sessions(scoped_dbs + [invalid])
- end.to raise_error(instance_of(described_class::InvalidLoadBalancerNameError))
- end
-
- context 'when calling use_primary!' do
- it 'applies use_primary! to all sessions' do
- with_sessions.use_primary!
-
- scoped_dbs.each do |db|
- expect(described_class.current(db.load_balancer).use_primary?).to eq(true)
- end
-
- (all_dbs - scoped_dbs).each do |db|
- expect(described_class.current(db.load_balancer).use_primary?).to eq(false)
- end
- end
- end
-
- context 'when calling use_primary' do
- it 'applies use_primary to all scoped sessions' do
- with_sessions.use_primary do
- scoped_dbs.each do |db|
- expect(described_class.current(db.load_balancer).use_primary?).to eq(true)
- end
-
- (all_dbs - scoped_dbs).each do |db|
- expect(described_class.current(db.load_balancer).use_primary?).to eq(false)
- end
- end
-
- all_dbs.each do |db|
- expect(described_class.current(db.load_balancer).use_primary?).to eq(false)
- end
- end
- end
-
- context 'when calling ignore_writes' do
- it 'applies ignore_writes to all scoped sessions' do
- with_sessions.ignore_writes do
- all_dbs.each do |db|
- described_class.current(db.load_balancer).write!
- end
-
- scoped_dbs.each do |db|
- expect(described_class.current(db.load_balancer).performed_write?).to eq(true)
- expect(described_class.current(db.load_balancer).use_primary?).to eq(false)
- end
-
- (all_dbs - scoped_dbs).each do |db|
- expect(described_class.current(db.load_balancer).performed_write?).to eq(true)
- expect(described_class.current(db.load_balancer).use_primary?).to eq(true)
- end
- end
- end
- end
-
- context 'when calling use_replicas_for_read_queries' do
- it 'applies use_replicas_for_read_queries to all scoped sessions' do
- with_sessions.use_replicas_for_read_queries do
- scoped_dbs.each do |db|
- expect(described_class.current(db.load_balancer).use_replicas_for_read_queries?).to eq(true)
- end
-
- (all_dbs - scoped_dbs).each do |db|
- expect(described_class.current(db.load_balancer).use_replicas_for_read_queries?).to eq(false)
- end
- end
-
- all_dbs.each do |db|
- expect(described_class.current(db.load_balancer).use_replicas_for_read_queries?).to eq(false)
- end
- end
- end
-
- context 'when calling fallback_to_replicas_for_ambiguous_queries' do
- it 'applies fallback_to_replicas_for_ambiguous_queries to all scoped sessions' do
- with_sessions.fallback_to_replicas_for_ambiguous_queries do
- scoped_dbs.each do |db|
- expect(described_class.current(db.load_balancer).fallback_to_replicas_for_ambiguous_queries?).to eq(true)
- end
-
- (all_dbs - scoped_dbs).each do |db|
- expect(described_class.current(db.load_balancer).fallback_to_replicas_for_ambiguous_queries?).to eq(false)
- end
- end
-
- all_dbs.each do |db|
- expect(described_class.current(db.load_balancer).fallback_to_replicas_for_ambiguous_queries?).to eq(false)
- end
- end
- end
- end
-end
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
index 3accd73eb6729..397db793d9ca4 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -13,7 +13,7 @@
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
def run_middleware
@@ -55,7 +55,7 @@ def perform(*args); end
run_middleware
expect(job['wal_locations']).to be_nil
- expect(job['wal_location_sources']).to be_nil
+ expect(job['wal_location_source']).to be_nil
end
include_examples 'job data consistency'
@@ -80,24 +80,22 @@ def perform(*args); end
context 'when replica hosts are available' do
it 'passes database_replica_location' do
expected_locations = expect_tracked_locations_when_replicas_available
- expected_sources = expected_locations.keys.index_with { |_| :replica }
run_middleware
expect(job['wal_locations']).to eq(expected_locations)
- expect(job['wal_location_sources']).to eq(expected_sources)
+ expect(job['wal_location_source']).to eq(:replica)
end
end
context 'when no replica hosts are available' do
it 'passes primary_write_location' do
expected_locations = expect_tracked_locations_when_no_replicas_available
- expected_sources = expected_locations.keys.index_with { |_| :replica }
run_middleware
expect(job['wal_locations']).to eq(expected_locations)
- expect(job['wal_location_sources']).to eq(expected_sources)
+ expect(job['wal_location_source']).to eq(:replica)
end
end
@@ -111,12 +109,11 @@ def perform(*args); end
it 'passes primary write location', :aggregate_failures do
expected_locations = expect_tracked_locations_from_primary_only
- expected_sources = expected_locations.keys.index_with { |_| :primary }
run_middleware
expect(job['wal_locations']).to eq(expected_locations)
- expect(job['wal_location_sources']).to eq(expected_sources)
+ expect(job['wal_location_source']).to eq(:primary)
end
include_examples 'job data consistency'
@@ -159,7 +156,6 @@ def perform(*args); end
let(:new_location) { 'AB/12345' }
let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => old_location } }
let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
- let(:session) { Gitlab::Database::LoadBalancing::Session.new }
before do
Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
@@ -170,15 +166,14 @@ def perform(*args); end
shared_examples_for 'does not set database location again' do |use_primary|
before do
- allow(Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).and_return(session)
- allow(session).to receive(:use_primary?).and_return(use_primary)
+ allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
end
it 'does not set database locations again' do
run_middleware
expect(job['wal_locations']).to eq(wal_locations)
- expect(job['wal_location_sources']).to be_nil
+ expect(job['wal_location_source']).to be_nil
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
index 4d2ff61a4adf4..793632c81f881 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
@@ -16,14 +16,14 @@
skip_default_enabled_yaml_check
replication_lag!(false)
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
stub_const("#{described_class.name}::REPLICA_WAIT_SLEEP_SECONDS", 0.02)
stub_const("#{described_class.name}::URGENT_REPLICA_WAIT_SLEEP_SECONDS", 0.01)
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '#call' do
@@ -59,9 +59,7 @@ def perform(*args); end
shared_examples_for 'stick to the primary' do |expected_strategy|
it 'sticks to the primary' do
run_middleware do
- Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?).to be_truthy
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to be_truthy
end
end
@@ -76,9 +74,7 @@ def perform(*args); end
.and_return(any_caught_up)
run_middleware do
- Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?).not_to be_truthy
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
end
end
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index e9782f214c745..8c2901c3b8994 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -14,7 +14,6 @@
let(:redis) { instance_double(::Gitlab::Redis::MultiStore) }
before do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
allow(::Gitlab::Redis::DbLoadBalancing).to receive(:with).and_yield(redis)
allow(ActiveRecord::Base.load_balancer)
@@ -27,7 +26,7 @@
end
after do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
describe '#find_caught_up_replica' do
@@ -49,7 +48,7 @@
expect(load_balancer).not_to receive(:select_up_to_date_host)
expect(redis).not_to receive(:del)
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer)).to receive(:use_primary!)
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
expect(sticking.find_caught_up_replica(:user, 42, use_primary_on_empty_location: true)).to eq(false)
end
@@ -88,7 +87,7 @@
it 'returns false, does not unstick and calls use_primary!' do
expect(redis).not_to receive(:del)
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer)).to receive(:use_primary!)
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
expect(sticking.find_caught_up_replica(:user, 42)).to eq(false)
end
@@ -96,7 +95,7 @@
context 'when use_primary_on_failure is false' do
it 'does not call use_primary!' do
expect(redis).not_to receive(:del)
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer)).not_to receive(:use_primary!)
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
expect(sticking.find_caught_up_replica(:user, 42, use_primary_on_failure: false)).to eq(false)
end
@@ -116,7 +115,8 @@
.with("database-load-balancing/write-location/#{load_balancer.name}/user/#{id}", 'the-primary-lsn', ex: 30)
end
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(load_balancer)).to receive(:use_primary!)
+ expect(Gitlab::Database::LoadBalancing::Session.current)
+ .to receive(:use_primary!)
subject
end
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index 8bdbaea422f6a..f65c27688b8e2 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -153,10 +153,6 @@
end
end
- def current_session
- ::Gitlab::Database::LoadBalancing::SessionMap.current(model.load_balancer)
- end
-
where(:queries, :include_transaction, :expected_results) do
[
# Read methods
@@ -239,7 +235,7 @@ def current_session
# use_primary
[
-> {
- current_session.use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
model.first
model.where(name: 'test1').to_a
end
@@ -252,7 +248,7 @@ def current_session
[
-> {
model.first
- current_session.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
model.where(name: 'test1').to_a
},
false, [:replica, :primary]
@@ -261,7 +257,7 @@ def current_session
# use_replicas_for_read_queries does not affect read queries
[
-> {
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.where(name: 'test1').to_a
end
},
@@ -271,7 +267,7 @@ def current_session
# use_replicas_for_read_queries does not affect write queries
[
-> {
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.create!(name: 'test1')
end
},
@@ -281,7 +277,7 @@ def current_session
# use_replicas_for_read_queries does not affect ambiguous queries
[
-> {
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.connection.exec_query("SELECT 1")
end
},
@@ -291,8 +287,8 @@ def current_session
# use_replicas_for_read_queries ignores use_primary! for read queries
[
-> {
- current_session.use_primary!
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.where(name: 'test1').to_a
end
},
@@ -302,8 +298,8 @@ def current_session
# use_replicas_for_read_queries adheres use_primary! for write queries
[
-> {
- current_session.use_primary!
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.create!(name: 'test1')
end
},
@@ -313,8 +309,8 @@ def current_session
# use_replicas_for_read_queries adheres use_primary! for ambiguous queries
[
-> {
- current_session.use_primary!
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.connection.exec_query('SELECT 1')
end
},
@@ -324,8 +320,8 @@ def current_session
# use_replicas_for_read_queries ignores use_primary blocks
[
-> {
- current_session.use_primary do
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.where(name: 'test1').to_a
end
end
@@ -336,8 +332,8 @@ def current_session
# use_replicas_for_read_queries ignores a session already performed write
[
-> {
- current_session.write!
- current_session.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.write!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
model.where(name: 'test1').to_a
end
},
@@ -347,7 +343,7 @@ def current_session
# fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.first
model.where(name: 'test1').to_a
end
@@ -358,7 +354,7 @@ def current_session
# fallback_to_replicas_for_ambiguous_queries for read-only transaction
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.transaction do
model.first
model.where(name: 'test1').to_a
@@ -371,7 +367,7 @@ def current_session
# A custom read query inside fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query("SELECT 1")
end
},
@@ -381,7 +377,7 @@ def current_session
# A custom read query inside a transaction fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.transaction do
model.connection.exec_query("SET LOCAL statement_timeout = 5000")
model.count
@@ -395,7 +391,7 @@ def current_session
[
-> {
model.create!(name: 'Test1')
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query("SELECT 1")
end
},
@@ -405,8 +401,8 @@ def current_session
# fallback_to_replicas_for_ambiguous_queries after use_primary!
[
-> {
- current_session.use_primary!
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query("SELECT 1")
end
},
@@ -416,8 +412,8 @@ def current_session
# fallback_to_replicas_for_ambiguous_queries inside use_primary
[
-> {
- current_session.use_primary do
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query("SELECT 1")
end
end
@@ -428,8 +424,8 @@ def current_session
# use_primary inside fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
- current_session.use_primary do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
model.connection.exec_query("SELECT 1")
end
end
@@ -440,7 +436,7 @@ def current_session
# A write query inside fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query("SELECT 1")
model.delete_all
model.connection.exec_query("SELECT 1")
@@ -452,8 +448,8 @@ def current_session
# use_replicas_for_read_queries incorporates with fallback_to_replicas_for_ambiguous_queries
[
-> {
- current_session.use_replicas_for_read_queries do
- current_session.fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.connection.exec_query('SELECT 1')
model.where(name: 'test1').to_a
end
@@ -504,7 +500,7 @@ def current_session
# replica even when the current session is stuck to the primary.
[
-> {
- current_session.use_primary!
+ ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
model.connection.clear_cache!
model.connection.schema_cache.add('users')
model.connection.pool.release_connection
@@ -565,9 +561,7 @@ def current_session
context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
it 'raises an exception' do
expect do
- ::Gitlab::Database::LoadBalancing::SessionMap
- .current(model.load_balancer)
- .fallback_to_replicas_for_ambiguous_queries do
+ ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
model.transaction do
model.first
model.create!(name: 'hello')
diff --git a/spec/lib/gitlab/database/postgres_autovacuum_activity_spec.rb b/spec/lib/gitlab/database/postgres_autovacuum_activity_spec.rb
index 74806231ed8e0..5367cf1fb9bfd 100644
--- a/spec/lib/gitlab/database/postgres_autovacuum_activity_spec.rb
+++ b/spec/lib/gitlab/database/postgres_autovacuum_activity_spec.rb
@@ -22,8 +22,7 @@
create(:postgres_autovacuum_activity, table: table)
end
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(ApplicationRecord.load_balancer))
- .to receive(:use_primary).and_yield
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary).and_yield
end
it 'returns autovacuum activity for queries tables' do
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 68ab6540ff9bb..c52599baef2a8 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -211,11 +211,11 @@
it 'returns primary db config even if ambiguous queries default to replica' do
Gitlab::Database.database_base_models_using_load_balancing.each_value do |database_base_model|
connection = database_base_model.connection
- Gitlab::Database::LoadBalancing::SessionMap.with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord]).use_primary!
+ Gitlab::Database::LoadBalancing::Session.current.use_primary!
primary_config = described_class.db_config_for_connection(connection)
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
- Gitlab::Database::LoadBalancing::SessionMap.with_sessions([::ApplicationRecord, ::Ci::ApplicationRecord]).fallback_to_replicas_for_ambiguous_queries do
+ Gitlab::Database::LoadBalancing::Session.clear_session
+ Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
expect(described_class.db_config_for_connection(connection)).to eq(primary_config)
end
end
@@ -466,11 +466,6 @@ def self.uncached
yield
end
-
- def self.load_balancer
- lb = Struct.new(:name)
- lb.new(:main)
- end
end
end
@@ -486,8 +481,7 @@ def self.load_balancer
expect(model1.instance_variable_get(:@uncached)).to be_nil
expect(model2.instance_variable_get(:@uncached)).to be_nil
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(::ApplicationRecord.load_balancer))
- .to receive(:use_primary).twice.and_yield
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary).and_yield
expect(model2).to receive(:uncached).and_call_original
expect(model1).to receive(:uncached).and_call_original
diff --git a/spec/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing_spec.rb b/spec/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing_spec.rb
index a461493b7e105..8d8b879a90f6a 100644
--- a/spec/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing_spec.rb
+++ b/spec/lib/gitlab/graphql/subscriptions/action_cable_with_load_balancing_spec.rb
@@ -100,12 +100,7 @@ def handle_event!(wal_locations: nil)
context 'when WAL locations are not present' do
it 'uses the primary' do
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:with_sessions).with(Gitlab::Database::LoadBalancing.base_models).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_primary!).and_call_original
- end
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
handle_event!(wal_locations: {})
end
@@ -124,8 +119,7 @@ def handle_event!(wal_locations: nil)
it 'does not use the primary' do
stub_replica_available!(true)
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .not_to receive(:with_sessions).with(Gitlab::Database::LoadBalancing.base_models).and_call_original
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
handle_event!
end
@@ -135,12 +129,7 @@ def handle_event!(wal_locations: nil)
it 'uses the primary' do
stub_replica_available!(false)
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:with_sessions).with(Gitlab::Database::LoadBalancing.base_models).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_primary!).and_call_original
- end
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary!)
handle_event!
end
diff --git a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
index 7dfd4504cbf5a..98a72ff8e554c 100644
--- a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
+++ b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
@@ -309,12 +309,7 @@ def write_relation_array(_, _, enumerator)
describe 'load balancing' do
it 'reads from replica' do
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:with_sessions).with(Gitlab::Database::LoadBalancing.base_models).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
subject.execute
end
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index e9a880bdcea4b..a2ef469a8459c 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -207,8 +207,7 @@
let(:session) { Gitlab::Database::LoadBalancing::Session.new }
before do
- allow(::Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:current).with(described_class.load_balancer).and_return(session)
+ allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries).and_yield
end
diff --git a/spec/models/concerns/packages/downloadable_spec.rb b/spec/models/concerns/packages/downloadable_spec.rb
index 59c9a9228ed64..79e0d684b7cf7 100644
--- a/spec/models/concerns/packages/downloadable_spec.rb
+++ b/spec/models/concerns/packages/downloadable_spec.rb
@@ -10,8 +10,7 @@
subject { package.touch_last_downloaded_at }
it 'updates the downloaded_at' do
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(package.load_balancer))
- .to receive(:without_sticky_writes).and_call_original
+ expect(::Gitlab::Database::LoadBalancing::Session).to receive(:without_sticky_writes).and_call_original
expect { subject }
.to change { package.last_downloaded_at }.from(nil).to(instance_of(ActiveSupport::TimeWithZone))
end
diff --git a/spec/models/sent_notification_spec.rb b/spec/models/sent_notification_spec.rb
index 8aad3c64a492f..8da87c2349aac 100644
--- a/spec/models/sent_notification_spec.rb
+++ b/spec/models/sent_notification_spec.rb
@@ -56,9 +56,7 @@
it 'writes without sticking to primary' do
subject
- Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(lb).use_primary?).to be false
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to be false
end
end
diff --git a/spec/requests/api/ci/runner/jobs_artifacts_spec.rb b/spec/requests/api/ci/runner/jobs_artifacts_spec.rb
index 428089b0da692..e20cd1f231bef 100644
--- a/spec/requests/api/ci/runner/jobs_artifacts_spec.rb
+++ b/spec/requests/api/ci/runner/jobs_artifacts_spec.rb
@@ -897,12 +897,19 @@ def upload_artifacts(file, headers = {}, params = {})
let(:token) { job.token }
def expect_use_primary
- expect(Gitlab::Database::LoadBalancing::SessionMap)
- .to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
+ lb_session = ::Gitlab::Database::LoadBalancing::Session.current
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_primary).and_call_original
- end
+ expect(lb_session).to receive(:use_primary).and_call_original
+
+ allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(lb_session)
+ end
+
+ def expect_no_use_primary
+ lb_session = ::Gitlab::Database::LoadBalancing::Session.current
+
+ expect(lb_session).not_to receive(:use_primary)
+
+ allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(lb_session)
end
it_behaves_like 'API::CI::Runner application context metadata', 'GET /api/:version/jobs/:id/artifacts' do
diff --git a/spec/requests/api/maven_packages_spec.rb b/spec/requests/api/maven_packages_spec.rb
index 7f6132f68d7b3..a81dd7e6df09b 100644
--- a/spec/requests/api/maven_packages_spec.rb
+++ b/spec/requests/api/maven_packages_spec.rb
@@ -39,10 +39,6 @@
let(:version) { '1.0-SNAPSHOT' }
let(:param_path) { "#{package_name}/#{version}" }
- before do
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
- end
-
shared_examples 'handling groups and subgroups for' do |shared_example_name, shared_example_args = {}, visibilities: { public: :redirect }|
context 'within a group' do
visibilities.each do |visibility, not_found_response|
@@ -1321,11 +1317,11 @@ def authorize_upload_with_token(params = {}, request_headers = headers_with_toke
end
def expect_use_primary
- lb_session = ::Gitlab::Database::LoadBalancing::SessionMap.current(ApplicationRecord.load_balancer)
+ lb_session = ::Gitlab::Database::LoadBalancing::Session.current
expect(lb_session).to receive(:use_primary).and_call_original
- allow(::Gitlab::Database::LoadBalancing::SessionMap).to receive(:current).and_return(lb_session)
+ allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(lb_session)
end
end
diff --git a/spec/services/personal_access_tokens/last_used_service_spec.rb b/spec/services/personal_access_tokens/last_used_service_spec.rb
index 28fe8ecf98edb..b31cc4dd754fe 100644
--- a/spec/services/personal_access_tokens/last_used_service_spec.rb
+++ b/spec/services/personal_access_tokens/last_used_service_spec.rb
@@ -48,15 +48,14 @@
context 'when database load balancing is configured' do
let!(:service) { described_class.new(personal_access_token) }
- let(:lb) { personal_access_token.load_balancer }
it 'does not stick to primary' do
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).not_to be_performed_write
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
expect { service.execute }.to change { personal_access_token.last_used_at }
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).to be_performed_write
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).not_to be_using_primary
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
end
end
end
diff --git a/spec/services/users/activity_service_spec.rb b/spec/services/users/activity_service_spec.rb
index c2688195ac50c..3e0f9b829ce74 100644
--- a/spec/services/users/activity_service_spec.rb
+++ b/spec/services/users/activity_service_spec.rb
@@ -124,29 +124,28 @@
context 'when last activity is in the past' do
let(:user) { create(:user, last_activity_on: Date.today - 1.week) }
- let(:lb) { User.load_balancer }
context 'database load balancing is configured' do
before do
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
end
let(:service) do
service = described_class.new(author: user)
- ::Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ ::Gitlab::Database::LoadBalancing::Session.clear_session
service
end
it 'does not stick to primary' do
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).not_to be_performed_write
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
service.execute
expect(user.last_activity_on).to eq(Date.today)
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).to be_performed_write
- expect(::Gitlab::Database::LoadBalancing::SessionMap.current(lb)).not_to be_using_primary
+ expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
+ expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
end
end
diff --git a/spec/support/database_load_balancing.rb b/spec/support/database_load_balancing.rb
index cc176dce3c944..d2902ddcc7c01 100644
--- a/spec/support/database_load_balancing.rb
+++ b/spec/support/database_load_balancing.rb
@@ -15,12 +15,12 @@
model.connection = Gitlab::Database::LoadBalancing::ConnectionProxy.new(model.load_balancer)
end
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
redis_shared_state_cleanup!
example.run
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
redis_shared_state_cleanup!
old_proxies.each do |model, proxy|
diff --git a/spec/support/db_cleaner.rb b/spec/support/db_cleaner.rb
index 2df82511246ca..b6539d6215b6a 100644
--- a/spec/support/db_cleaner.rb
+++ b/spec/support/db_cleaner.rb
@@ -114,24 +114,3 @@ def force_disconnect_all_connections!
end
DbCleaner.prepend_mod_with('DbCleaner')
-
-# We patch the establish_master_connection so that it establishes a connection
-# using a ActiveRecord::DatabaseConfigurations::HashConfig instead of a hash.
-#
-# Using a HashConfig avoids resetting the name of the connection.
-module PostgreSQLDatabaseTasksPatch
- def establish_master_connection
- establish_connection(
- ActiveRecord::DatabaseConfigurations::HashConfig.new(
- db_config.env_name,
- db_config.name,
- db_config.configuration_hash.merge(
- database: "postgres",
- schema_search_path: "public"
- )
- )
- )
- end
-end
-
-ActiveRecord::Tasks::PostgreSQLDatabaseTasks.prepend(PostgreSQLDatabaseTasksPatch)
diff --git a/spec/support/shared_contexts/lib/gitlab/database/load_balancing/wal_tracking_shared_context.rb b/spec/support/shared_contexts/lib/gitlab/database/load_balancing/wal_tracking_shared_context.rb
index 0d1d86eb96164..46d6a1fbac0cc 100644
--- a/spec/support/shared_contexts/lib/gitlab/database/load_balancing/wal_tracking_shared_context.rb
+++ b/spec/support/shared_contexts/lib/gitlab/database/load_balancing/wal_tracking_shared_context.rb
@@ -4,9 +4,9 @@
let(:current_location) { '0/D525E3A8' }
around do |example|
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
example.run
- Gitlab::Database::LoadBalancing::SessionMap.clear_session
+ Gitlab::Database::LoadBalancing::Session.clear_session
end
def expect_tracked_locations_when_replicas_available
@@ -53,15 +53,11 @@ def stub_load_balancing_enabled!
end
def stub_no_writes_performed!
- Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- allow(Gitlab::Database::LoadBalancing::SessionMap.current(lb)).to receive(:use_primary?).and_return(false)
- end
+ allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(false)
end
def stub_write_performed!
- Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
- allow(Gitlab::Database::LoadBalancing::SessionMap.current(lb)).to receive(:use_primary?).and_return(true)
- end
+ allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(true)
end
def stub_replica_available!(available)
diff --git a/spec/tasks/gitlab/db/validate_config_rake_spec.rb b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
index 8a5d5b2c30046..4a0108f00cebf 100644
--- a/spec/tasks/gitlab/db/validate_config_rake_spec.rb
+++ b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
@@ -170,6 +170,7 @@
context "when there's no main: but something different, as currently we only can share with main:" do
let(:test_config) do
{
+ archive: main_database_config,
ci: main_database_config.merge(database_tasks: false)
}
end
diff --git a/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb b/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb
index 8aa84b248e475..d2ed5dcbbe173 100644
--- a/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb
+++ b/spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb
@@ -46,11 +46,7 @@
context 'with load balancing enabled' do
it 'reads from the replica database' do
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
execute_worker
end
diff --git a/spec/workers/container_expiration_policies/cleanup_container_repository_worker_spec.rb b/spec/workers/container_expiration_policies/cleanup_container_repository_worker_spec.rb
index 381f3433904c3..b957189f0d962 100644
--- a/spec/workers/container_expiration_policies/cleanup_container_repository_worker_spec.rb
+++ b/spec/workers/container_expiration_policies/cleanup_container_repository_worker_spec.rb
@@ -24,12 +24,7 @@
service_response = cleanup_service_response(repository: repository)
expect(ContainerExpirationPolicies::CleanupService)
.to receive(:new).with(repository).and_return(double(execute: service_response))
-
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
-
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
expect_log_extra_metadata(service_response: service_response)
expect_log_info(project_id: project.id, container_repository_id: repository.id)
diff --git a/spec/workers/container_expiration_policy_worker_spec.rb b/spec/workers/container_expiration_policy_worker_spec.rb
index 77e07d009973e..fcb43b8608496 100644
--- a/spec/workers/container_expiration_policy_worker_spec.rb
+++ b/spec/workers/container_expiration_policy_worker_spec.rb
@@ -81,8 +81,7 @@
context 'with load balancing enabled' do
it 'reads the counts from the replica' do
- expect(Gitlab::Database::LoadBalancing::SessionMap.current(ContainerRepository.load_balancer))
- .to receive(:use_replicas_for_read_queries).and_call_original
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
subject
end
diff --git a/spec/workers/dependency_proxy/image_ttl_group_policy_worker_spec.rb b/spec/workers/dependency_proxy/image_ttl_group_policy_worker_spec.rb
index 6ce22ec798cf2..ce5209f1d09aa 100644
--- a/spec/workers/dependency_proxy/image_ttl_group_policy_worker_spec.rb
+++ b/spec/workers/dependency_proxy/image_ttl_group_policy_worker_spec.rb
@@ -48,11 +48,7 @@
context 'with load balancing enabled', :db_load_balancing do
it 'reads the counts from the replica' do
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
subject
end
diff --git a/spec/workers/packages/cleanup_package_registry_worker_spec.rb b/spec/workers/packages/cleanup_package_registry_worker_spec.rb
index f27d089feada3..63ce0943f05d9 100644
--- a/spec/workers/packages/cleanup_package_registry_worker_spec.rb
+++ b/spec/workers/packages/cleanup_package_registry_worker_spec.rb
@@ -116,11 +116,7 @@
context 'with load balancing enabled', :db_load_balancing do
it 'reads the count from the replica' do
- expect(Gitlab::Database::LoadBalancing::SessionMap).to receive(:with_sessions).with([::ApplicationRecord, ::Ci::ApplicationRecord]).and_call_original
-
- expect_next_instance_of(Gitlab::Database::LoadBalancing::ScopedSessions) do |inst|
- expect(inst).to receive(:use_replicas_for_read_queries).and_call_original
- end
+ expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
perform
end
--
GitLab