Combine partition creation and detaching

By creating and detaching partitions using the same lease, we can
easily skip the partition sync code if another node is currently
executing it. This is safe because we have six months of headroom for
future partitions and the job to sync partitions runs every six hours

app/workers/all_queues.yml

@@ -247,8 +247,8 @@
   :idempotent: true
   :tags:
   - :exclude_from_kubernetes
-- :name: cronjob:database_partition_detach
-  :worker_name: Database::PartitionDetachWorker
+- :name: cronjob:database_partition_management
+  :worker_name: Database::PartitionManagementWorker
   :feature_category: :database
   :has_external_dependencies:
   :urgency: :low

app/workers/database/partition_detach_worker.rb → app/workers/database/partition_management_worker.rb

 # frozen_string_literal: true
 
 module Database
-  class PartitionDetachWorker
+  class PartitionManagementWorker
     include ApplicationWorker
 
     sidekiq_options retry: 3
@@ -11,7 +11,7 @@ module Database
     idempotent!
 
     def perform
-      Gitlab::Database::Partitioning::PartitionManager.new.detach_partitions
+      Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions
     ensure
       Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
     end

app/workers/partition_creation_worker.rb

@@ -10,8 +10,6 @@ class PartitionCreationWorker
   idempotent!
 
   def perform
-    Gitlab::Database::Partitioning::PartitionManager.new.create_partitions
-  ensure
-    Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
+    # Removed in favor of Database::PartitionManagementWorker
   end
 end

config/initializers/1_settings.rb

@@ -540,12 +540,9 @@ Settings.cron_jobs['authorized_project_update_periodic_recalculate_worker']['job
 Settings.cron_jobs['update_container_registry_info_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['update_container_registry_info_worker']['cron'] ||= '0 0 * * *'
 Settings.cron_jobs['update_container_registry_info_worker']['job_class'] = 'UpdateContainerRegistryInfoWorker'
-Settings.cron_jobs['postgres_dynamic_partitions_creator'] ||= Settingslogic.new({})
-Settings.cron_jobs['postgres_dynamic_partitions_creator']['cron'] ||= '21 */6 * * *'
-Settings.cron_jobs['postgres_dynamic_partitions_creator']['job_class'] ||= 'PartitionCreationWorker'
-Settings.cron_jobs['postgres_dynamic_partitions_detacher'] ||= Settingslogic.new({})
-Settings.cron_jobs['postgres_dynamic_partitions_detacher']['cron'] ||= '25 */6 * * *'
-Settings.cron_jobs['postgres_dynamic_partitions_detacher']['job_class'] ||= 'Database::PartitionDetachWorker'
+Settings.cron_jobs['postgres_dynamic_partitions_manager'] ||= Settingslogic.new({})
+Settings.cron_jobs['postgres_dynamic_partitions_manager']['cron'] ||= '21 */6 * * *'
+Settings.cron_jobs['postgres_dynamic_partitions_manager']['job_class'] ||= 'Database::PartitionManagementWorker'
 Settings.cron_jobs['ci_platform_metrics_update_cron_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['ci_platform_metrics_update_cron_worker']['cron'] ||= '47 9 * * *'
 Settings.cron_jobs['ci_platform_metrics_update_cron_worker']['job_class'] = 'CiPlatformMetricsUpdateCronWorker'

config/initializers/postgres_partitioning.rb

@@ -11,7 +11,7 @@ if Gitlab.ee?
 end
 
 begin
-  Gitlab::Database::Partitioning::PartitionManager.new.create_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
+  Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
 rescue ActiveRecord::ActiveRecordError, PG::Error
   # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
 end

lib/gitlab/database/partitioning/partition_manager.rb

@@ -4,8 +4,6 @@ module Gitlab
   module Database
     module Partitioning
       class PartitionManager
-        include Gitlab::ExclusiveLeaseHelpers
-
         def self.register(model)
           raise ArgumentError, "Only models with a #partitioning_strategy can be registered." unless model.respond_to?(:partitioning_strategy)
 
@@ -17,7 +15,7 @@ module Gitlab
         end
 
         LEASE_TIMEOUT = 1.minute
-        LEASE_KEY = 'database_partition_management_%s'
+        MANAGEMENT_LEASE_KEY = 'database_partition_management_%s'
 
         attr_reader :models
 
@@ -25,41 +23,25 @@ module Gitlab
           @models = models
         end
 
-        def create_partitions
+        def sync_partitions
           Gitlab::AppLogger.info("Checking state of dynamic postgres partitions")
 
           models.each do |model|
             # Double-checking before getting the lease:
-            # The prevailing situation is no missing partitions
-            next if missing_partitions(model).empty?
+            # The prevailing situation is no missing partitions and no extra partitions
+            next if missing_partitions(model).empty? && extra_partitions(model).empty?
 
-            only_with_lease_lock(model) do
+            only_with_exclusive_lease(model, lease_key: MANAGEMENT_LEASE_KEY) do
               partitions_to_create = missing_partitions(model)
+              create(partitions_to_create) unless partitions_to_create.empty?
 
-              next if partitions_to_create.empty?
-
-              create(model, partitions_to_create)
-            end
-          rescue StandardError => e
-            Gitlab::AppLogger.error("Failed to create partition(s) for #{model.table_name}: #{e.class}: #{e.message}")
-          end
-        end
-
-        def detach_partitions
-          return unless Feature.enabled?(:partition_pruning_dry_run)
-
-          models.each do |model|
-            next if extra_partitions(model).empty?
-
-            only_with_lease_lock(model) do
-              partitions_to_detach = extra_partitions(model)
-
-              next if partitions_to_detach.empty?
-
-              detach(partitions_to_detach)
+              if Feature.enabled?(:partition_pruning_dry_run)
+                partitions_to_detach = extra_partitions(model)
+                detach(partitions_to_detach) unless partitions_to_detach.empty?
+              end
             end
           rescue StandardError => e
-            Gitlab::AppLogger.error("Failed to remove partition(s) for #{model.table_name}: #{e.class}: #{e.message}")
+            Gitlab::AppLogger.error("Failed to create / detach partition(s) for #{model.table_name}: #{e.class}: #{e.message}")
           end
         end
 
@@ -72,18 +54,21 @@ module Gitlab
         end
 
         def extra_partitions(model)
+          return [] unless Feature.enabled?(:partition_pruning_dry_run)
           return [] unless connection.table_exists?(model.table_name)
 
           model.partitioning_strategy.extra_partitions
         end
 
-        def only_with_lease_lock(model)
-          in_lock(LEASE_KEY % model.table_name, ttl: LEASE_TIMEOUT, sleep_sec: 3.seconds) do
-            yield
-          end
+        def only_with_exclusive_lease(model, lease_key:)
+          lease = Gitlab::ExclusiveLease.new(lease_key % model.table_name, timeout: LEASE_TIMEOUT)
+
+          yield if lease.try_obtain
+        ensure
+          lease&.cancel
         end
 
-        def create(model, partitions)
+        def create(partitions)
           connection.transaction do
             with_lock_retries do
               partitions.each do |partition|

lib/tasks/gitlab/db.rake

@@ -118,7 +118,7 @@ namespace :gitlab do
 
     desc 'Create missing dynamic database partitions'
     task create_dynamic_partitions: :environment do
-      Gitlab::Database::Partitioning::PartitionManager.new.create_partitions
+      Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions
     end
 
     # This is targeted towards deploys and upgrades of GitLab.

spec/lib/gitlab/database/partitioning/partition_manager_spec.rb

@@ -15,12 +15,12 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
   end
 
-  describe '#create_partitions (mocked)' do
-    subject { described_class.new(models).create_partitions }
+  context 'creating partitions (mocked)' do
+    subject { described_class.new(models).sync_partitions }
 
     let(:models) { [model] }
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
-    let(:partitioning_strategy) { double(missing_partitions: partitions) }
+    let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: []) }
     let(:table) { "some_table" }
 
     before do
@@ -28,7 +28,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
       allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
 
-      stub_exclusive_lease(described_class::LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+      stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
 
     let(:partitions) do
@@ -53,8 +53,8 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
         ]
       end
 
-      let(:strategy1) { double('strategy1', missing_partitions: nil) }
-      let(:strategy2) { double('strategy2', missing_partitions: partitions) }
+      let(:strategy1) { double('strategy1', missing_partitions: nil, extra_partitions: []) }
+      let(:strategy2) { double('strategy2', missing_partitions: partitions, extra_partitions: []) }
 
       it 'still creates partitions for the second table' do
         expect(strategy1).to receive(:missing_partitions).and_raise('this should never happen (tm)')
@@ -66,8 +66,8 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
   end
 
-  describe '#create_partitions' do
-    subject { described_class.new([my_model]).create_partitions }
+  context 'creating partitions' do
+    subject { described_class.new([my_model]).sync_partitions }
 
     let(:connection) { ActiveRecord::Base.connection }
     let(:my_model) do
@@ -95,20 +95,20 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
   end
 
-  describe '#detach_partitions (mocked)' do
-    subject { manager.detach_partitions }
+  context 'detaching partitions (mocked)' do
+    subject { manager.sync_partitions }
 
     let(:manager) { described_class.new(models) }
     let(:models) { [model] }
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table)}
-    let(:partitioning_strategy) { double(extra_partitions: extra_partitions) }
+    let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: []) }
     let(:table) { "foo" }
 
     before do
       allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
       allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
 
-      stub_exclusive_lease(described_class::LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+      stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
 
     let(:extra_partitions) do
@@ -136,7 +136,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
           ]
         end
 
-        let(:error_strategy) { double }
+        let(:error_strategy) { double(extra_partitions: nil, missing_partitions: []) }
 
         it 'still drops partitions for the other model' do
           expect(error_strategy).to receive(:extra_partitions).and_raise('injected error!')

spec/workers/database/partition_detach_worker_spec.rb → spec/workers/database/partition_management_worker_spec.rb

@@ -2,11 +2,11 @@
 
 require "spec_helper"
 
-RSpec.describe Database::PartitionDetachWorker do
+RSpec.describe Database::PartitionManagementWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:manager) {  instance_double('PartitionManager', detach_partitions: nil) }
+    let(:manager) {  instance_double('PartitionManager', sync_partitions: nil) }
     let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
 
     before do
@@ -15,7 +15,7 @@ RSpec.describe Database::PartitionDetachWorker do
     end
 
     it 'delegates to PartitionManager' do
-      expect(manager).to receive(:detach_partitions)
+      expect(manager).to receive(:sync_partitions)
 
       subject
     end

spec/workers/partition_creation_worker_spec.rb

-# frozen_string_literal: true
-
-require "spec_helper"
-
-RSpec.describe PartitionCreationWorker do
-  describe '#perform' do
-    subject { described_class.new.perform }
-
-    let(:manager) {  instance_double('PartitionManager', create_partitions: nil) }
-    let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
-
-    before do
-      allow(Gitlab::Database::Partitioning::PartitionManager).to receive(:new).and_return(manager)
-      allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
-    end
-
-    it 'delegates to PartitionCreator' do
-      expect(manager).to receive(:create_partitions)
-
-      subject
-    end
-
-    it 'reports partition metrics' do
-      expect(monitoring).to receive(:report_metrics)
-
-      subject
-    end
-  end
-end