online_migrations 0.26.0 → 0.27.1

data/lib/online_migrations/background_data_migrations/migration_job.rb

@@ -0,0 +1,160 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # Sidekiq job responsible for running background data migrations.
+    class MigrationJob
+      include Sidekiq::IterableJob
+
+      sidekiq_retry_in do |count, _exception, jobhash|
+        migration_id = jobhash["args"].fetch(0)
+        migration = Migration.find(migration_id)
+
+        if count + 1 >= migration.max_attempts
+          :kill
+        end
+      end
+
+      sidekiq_retries_exhausted do |jobhash, exception|
+        migration_id = jobhash["args"].fetch(0)
+        migration = Migration.find(migration_id)
+        migration.persist_error(exception)
+
+        OnlineMigrations.config.background_data_migrations.error_handler.call(exception, migration)
+      end
+
+      TICKER_INTERVAL = 5 # seconds
+
+      def initialize
+        super
+
+        @migration = nil
+        @data_migration = nil
+
+        @ticker = Ticker.new(TICKER_INTERVAL) do |ticks, duration|
+          # TODO: use 'cursor' accessor from sidekiq in the future.
+          # https://github.com/sidekiq/sidekiq/pull/6606
+          @migration.persist_progress(@_cursor, ticks, duration)
+          @migration.reload
+        end
+
+        @throttle_checked_at = current_time
+      end
+
+      def on_start
+        @migration.start
+      end
+
+      def on_resume
+        @data_migration.after_resume
+      end
+
+      def on_stop
+        @ticker.persist
+        @migration.stop
+      end
+
+      def on_complete
+        # Job was manually cancelled.
+        @migration.cancel if cancelled?
+
+        @migration.complete
+      end
+
+      def build_enumerator(migration_id, cursor:)
+        @migration = BackgroundDataMigrations::Migration.find(migration_id)
+        cursor ||= @migration.cursor
+
+        @migration.on_shard_if_present do
+          @data_migration = @migration.data_migration
+          collection_enum = @data_migration.build_enumerator(cursor: cursor)
+
+          if collection_enum
+            if !collection_enum.is_a?(Enumerator)
+              raise ArgumentError, <<~MSG.squish
+                #{@data_migration.class.name}#build_enumerator must return an Enumerator,
+                got #{collection_enum.class.name}.
+              MSG
+            end
+
+            collection_enum
+          else
+            collection = @data_migration.collection
+
+            case collection
+            when ActiveRecord::Relation
+              options = {
+                cursor: cursor,
+                batch_size: @data_migration.class.active_record_enumerator_batch_size || 100,
+              }
+              active_record_records_enumerator(collection, **options)
+            when ActiveRecord::Batches::BatchEnumerator
+              if collection.start || collection.finish
+                raise ArgumentError, <<~MSG.squish
+                  #{@data_migration.class.name}#collection does not support
+                  a batch enumerator with the "start" or "finish" options.
+                MSG
+              end
+
+              active_record_relations_enumerator(
+                collection.relation,
+                batch_size: collection.batch_size,
+                cursor: cursor,
+                use_ranges: collection.use_ranges
+              )
+            when Array
+              array_enumerator(collection, cursor: cursor)
+            else
+              raise ArgumentError, <<~MSG.squish
+                #{@data_migration.class.name}#collection must be either an ActiveRecord::Relation,
+                ActiveRecord::Batches::BatchEnumerator, or Array.
+              MSG
+            end
+          end
+        end
+      end
+
+      def each_iteration(item, _migration_id)
+        if @migration.cancelling? || @migration.pausing? || @migration.paused?
+          # Finish this exact sidekiq job. When the migration is paused
+          # and will be resumed, a new job will be enqueued.
+          finished = true
+          throw :abort, finished
+        elsif should_throttle?
+          ActiveSupport::Notifications.instrument("throttled.background_data_migrations", migration: @migration)
+          finished = false
+          throw :abort, finished
+        else
+          @data_migration.around_process do
+            @migration.data_migration.process(item)
+
+            pause = OnlineMigrations.config.background_data_migrations.iteration_pause
+            sleep(pause) if pause > 0
+          end
+          @ticker.tick
+        end
+      end
+
+      private
+        # It would be better for sidekiq to have a callback like `around_perform`,
+        # but currently this is the way to make job iteration shard aware.
+        def iterate_with_enumerator(enumerator, arguments)
+          @migration.on_shard_if_present { super }
+        end
+
+        THROTTLE_CHECK_INTERVAL = 5 # seconds
+        private_constant :THROTTLE_CHECK_INTERVAL
+
+        def should_throttle?
+          if current_time - @throttle_checked_at >= THROTTLE_CHECK_INTERVAL
+            @throttle_checked_at = current_time
+            OnlineMigrations.config.throttler.call
+          end
+        end
+
+        def current_time
+          ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+        end
+    end
+  end
+end

data/lib/online_migrations/background_data_migrations/migration_status_validator.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # @private
+    class MigrationStatusValidator < ActiveModel::Validator
+      # Valid status transitions a Migration can make.
+      VALID_STATUS_TRANSITIONS = {
+        # enqueued -> running occurs when the migration starts performing.
+        # enqueued -> paused occurs when the migration is paused before starting.
+        # enqueued -> cancelled occurs when the migration is cancelled before starting.
+        # enqueued -> failed occurs when the migration job fails to be enqueued, or
+        #   if the migration is deleted before is starts running.
+        "enqueued" => ["running", "paused", "cancelled", "failed"],
+        # running -> succeeded occurs when the migration completes successfully.
+        # running -> pausing occurs when a user pauses the migration as it's performing.
+        # running -> cancelling occurs when a user cancels the migration as it's performing.
+        # running -> failed occurs when the job raises an exception when running.
+        "running" => [
+          "succeeded",
+          "pausing",
+          "cancelling",
+          "failed",
+        ],
+        # pausing -> paused occurs when the migration actually halts performing and
+        #   occupies a status of paused.
+        # pausing -> cancelling occurs when the user cancels a migration immediately
+        #   after it was paused, such that the migration had not actually halted yet.
+        # pausing -> succeeded occurs when the migration completes immediately after
+        #   being paused. This can happen if the migration is on its last iteration
+        #   when it is paused, or if the migration is paused after enqueue but has
+        #   nothing in its collection to process.
+        # pausing -> failed occurs when the job raises an exception after the
+        #   user has paused it.
+        "pausing" => ["paused", "cancelling", "succeeded", "failed"],
+        # paused -> enqueued occurs when the migration is resumed after being paused.
+        # paused -> cancelled when the user cancels the migration after it is paused.
+        "paused" => ["enqueued", "cancelled"],
+        # failed -> enqueued occurs when the migration is retried after encounting an error.
+        "failed" => ["enqueued"],
+        # cancelling -> cancelled occurs when the migration actually halts performing
+        #   and occupies a status of cancelled.
+        # cancelling -> succeeded occurs when the migration completes immediately after
+        #   being cancelled. See description for pausing -> succeeded.
+        # cancelling -> failed occurs when the job raises an exception after the
+        #   user has cancelled it.
+        "cancelling" => ["cancelled", "succeeded", "failed"],
+      }
+
+      def validate(record)
+        return if !record.status_changed?
+
+        previous_status, new_status = record.status_change
+        valid_new_statuses = VALID_STATUS_TRANSITIONS.fetch(previous_status, [])
+
+        if !valid_new_statuses.include?(new_status)
+          record.errors.add(
+            :status,
+            "cannot transition data migration from status '#{previous_status}' to '#{new_status}'"
+          )
+        end
+      end
+    end
+  end
+end

data/lib/online_migrations/{background_migrations → background_data_migrations}/perform_action_on_relation.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     # @private
-    class PerformActionOnRelation < BackgroundMigration
+    class PerformActionOnRelation < DataMigration
       attr_reader :model, :conditions, :action, :options
 
       def initialize(model_name, conditions, action, options = {})
@@ -13,11 +13,11 @@ module OnlineMigrations
         @options = options.symbolize_keys
       end
 
-      def relation
-        model.unscoped.where(conditions)
+      def collection
+        model.unscoped.where(conditions).in_batches(of: 100)
       end
 
-      def process_batch(relation)
+      def process(relation)
         case action
         when :update_all
           updates = options.fetch(:updates)

data/lib/online_migrations/{background_migrations → background_data_migrations}/reset_counters.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     # @private
-    class ResetCounters < BackgroundMigration
+    class ResetCounters < DataMigration
       attr_reader :model, :counters, :touch
 
       def initialize(model_name, counters, options = {})
@@ -12,11 +12,11 @@ module OnlineMigrations
         @touch = options[:touch]
       end
 
-      def relation
-        model.unscoped
+      def collection
+        model.unscoped.in_batches(of: 100)
       end
 
-      def process_batch(relation)
+      def process(relation)
         updates = counters.map do |counter_association|
           has_many_association = has_many_association(counter_association)
 

data/lib/online_migrations/background_data_migrations/scheduler.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # Class responsible for scheduling background data migrations.
+    #
+    # Scheduler should be configured to run periodically, for example, via cron.
+    #
+    # @example Run via whenever
+    #   # add this to schedule.rb
+    #   every 1.minute do
+    #     runner "OnlineMigrations.run_background_data_migrations"
+    #   end
+    #
+    # @example Specific shard
+    #   every 1.minute do
+    #     runner "OnlineMigrations.run_background_data_migrations(shard: :shard_two)"
+    #   end
+    #
+    # @example Custom concurrency
+    #   every 1.minute do
+    #     # Allow to run 2 data migrations in parallel.
+    #     runner "OnlineMigrations.run_background_data_migrations(concurrency: 2)"
+    #   end
+    #
+    class Scheduler
+      def self.run(**options)
+        new.run(**options)
+      end
+
+      # Runs Scheduler
+      def run(shard: nil, concurrency: 1)
+        relation = Migration.queue_order
+        relation = relation.where(shard: shard) if shard
+
+        with_lock do
+          running = relation.running
+          enqueued = relation.enqueued
+
+          # Ensure no more than 'concurrency' migrations are running at the same time.
+          remaining_to_enqueue = concurrency - running.count
+          if remaining_to_enqueue > 0
+            migrations_to_enqueue = enqueued.limit(remaining_to_enqueue)
+            migrations_to_enqueue.each do |migration|
+              enqueue_migration(migration)
+            end
+          end
+        end
+
+        true
+      end
+
+      private
+        def with_lock(&block)
+          # Don't lock the whole table if we can lock only a single record (must be always the same).
+          first_record = Migration.queue_order.first
+          if first_record
+            first_record.with_lock(&block)
+          else
+            Migration.transaction do
+              Migration.connection.execute("LOCK #{Migration.table_name} IN ACCESS EXCLUSIVE MODE")
+              yield
+            end
+          end
+        end
+
+        def enqueue_migration(migration)
+          job = OnlineMigrations.config.background_data_migrations.job
+          job_class = job.constantize
+
+          jid = job_class.perform_async(migration.id)
+          if jid
+            migration.update!(status: :running, jid: jid)
+          end
+        end
+    end
+  end
+end

data/lib/online_migrations/background_data_migrations/ticker.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # This class encapsulates the logic behind updating the tick counter.
+    #
+    # It's initialized with a duration for the throttle, and a block to persist
+    # the number of ticks to increment.
+    #
+    # When +tick+ is called, the block will be called with the increment,
+    # provided the duration since the last update (or initialization) has been
+    # long enough.
+    #
+    # To not lose any increments, +persist+ should be used, which may call the
+    # block with any leftover ticks.
+    #
+    # @private
+    class Ticker
+      # Creates a Ticker that will call the block each time +tick+ is called,
+      # unless the tick is being throttled.
+      #
+      # @param interval [ActiveSupport::Duration, Numeric] Duration
+      #   since initialization or last call that will cause a throttle.
+      # @yieldparam ticks [Integer] the increment in ticks to be persisted.
+      #
+      def initialize(interval, &block)
+        @interval = interval
+        @block = block
+        @last_persisted_at = Time.current
+        @ticks_recorded = 0
+      end
+
+      # Increments the tick count by one, and may persist the new value if the
+      # threshold duration has passed since initialization or the tick count was
+      # last persisted.
+      #
+      def tick
+        @ticks_recorded += 1
+        persist if persist?
+      end
+
+      # Persists the tick increments by calling the block passed to the
+      # initializer. This is idempotent in the sense that calling it twice in a
+      # row will call the block at most once (if it had been throttled).
+      #
+      def persist
+        return if @ticks_recorded == 0
+
+        now = Time.current
+        duration = now - @last_persisted_at
+        @last_persisted_at = now
+        @block.call(@ticks_recorded, duration)
+        @ticks_recorded = 0
+      end
+
+      private
+        def persist?
+          Time.now - @last_persisted_at >= @interval
+        end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/config.rb

@@ -4,9 +4,9 @@ module OnlineMigrations
   module BackgroundSchemaMigrations
     # Class representing configuration options for background schema migrations.
     class Config
-      # Maximum number of run attempts
+      # Maximum number of run attempts.
       #
-      # When attempts are exhausted, the migration is marked as failed.
+      # When attempts are exhausted, the schema migration is marked as failed.
       # @return [Integer] defaults to 5
       #
       attr_accessor :max_attempts