<commit_before>module Leveret
  class Worker
    extend Forwardable

    attr_accessor :queues, :consumers

    def_delegators :Leveret, :log, :channel

    def initialize(*queue_names)
      queue_names = [Leveret.configuration.default_queue_name] if queue_names.empty?

      self.queues = queue_names.map { |name| Leveret::Queue.new(name) }
      self.consumers = []
      @time_to_die = false
    end

    def do_work
      log.info "Worker starting for #{queues.map(&:name).join(', ')}"
      setup_traps
      subscribe_to_queues

      loop do
        if @time_to_die
          log.info "Interrupt received, preparing to exit"
          consumers.each do |consumer|
            log.debug "Cancelling consumer on #{consumer.queue.name}"
            consumer.cancel
          end
          break
        end
        sleep 1
      end
      log.info "Exiting master process"
    end

    private

    def setup_traps
      trap('INT') do
        @time_to_die = true
      end
    end

    def subscribe_to_queues
      queues.map do |queue|
        consumers << queue.subscribe do |delivery_tag, payload|
          fork_and_run(delivery_tag, payload)
        end
      end
    end

    def fork_and_run(delivery_tag, payload)
      pid = fork do
        Leveret.configuration.after_fork.call

        log.info "[#{delivery_tag}] Forked to child process #{Process.pid} to run #{payload[:job]}"

        job_klass = Object.const_get(payload[:job])
        result = job_klass.perform(Leveret::Parameters.new(payload[:params]))

        log.info "[#{delivery_tag}] Job returned #{result}"

        ack_message(delivery_tag, result)

        log.info "[#{delivery_tag}] Exiting child process #{Process.pid}"
        exit!(0)
      end

      Process.wait(pid)
    end

    def ack_message(delivery_tag, result)
      if result == :reject
        log.debug "[#{delivery_tag}] Rejecting message"
        channel.reject(delivery_tag)
      elsif result == :requeue<commit_msg>light refactoring for readability<commit_after>module Leveret
  # Subscribes to one or more queues and forks workers to perform jobs as they arrive
  #
  # Call #do_work to subscribe to all queues and block the main thread.
  class Worker
    extend Forwardable

    attr_accessor :queues, :consumers

    def_delegators :Leveret, :log, :channel

    def initialize(*queue_names)
      queue_names = [Leveret.configuration.default_queue_name] if queue_names.empty?

      self.queues = queue_names.map { |name| Leveret::Queue.new(name) }
      self.consumers = []
      @time_to_die = false
    end

    def do_work
      log.info "Starting master process for #{queues.map(&:name).join(', ')}"
      prepare_for_work

      loop do
        if @time_to_die
          cancel_subscriptions
          break
        end
        sleep 1
      end
      log.info "Exiting master process"
    end

    private

    def prepare_for_work
      setup_traps
      start_subscriptions
    end

    def setup_traps
      trap('INT') do
        @time_to_die = true
      end
      trap('TERM') do
        @time_to_die = true
      end
    end

    def start_subscriptions
      queues.map do |queue|
        consumers << queue.subscribe do |delivery_tag, payload|
          fork_and_run(delivery_tag, payload)
        end
      end
    end

    def cancel_subscriptions
      log.info "Interrupt received, preparing to exit"
      consumers.each do |consumer|
        log.debug "Cancelling consumer on #{consumer.queue.name}"
        consumer.cancel
      end
    end

    def fork_and_run(delivery_tag, payload)
      pid = fork do
        log.info "[#{delivery_tag}] Forked to child process #{Process.pid} to run #{payload[:job]}"

        Leveret.configuration.after_fork.call

        result = perform_job(payload)
        log.info "[#{delivery_tag}] Job returned #{result}"
        ack_message(delivery_tag, result)

        log.info "[#{delivery_tag}] Exiting child process #{Process.pid}"
        exit!(0)
      end

      Process.wait(pid)
    end

    def perform_job(payload)
      job_klass = Object.const_get(payload[:job])
      job_klass.perform(Leveret::Parameters.new(payload[:params]))
    end

    def ack_message(delivery_tag, result)
      if result == :reject
        log.debug "[#{delivery_tag}] Rejecting message"
        channel.reject(delivery_tag)
      elsif result == :requeue
{'fraction_non_alphanumeric': 0.09057555606512267, 'fraction_numerical': 0.0009172208209126347, 'mean_word_length': 2.4265514532600156, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1739340', 'n_tokens_mistral': 1361, 'n_tokens_neox': 1330, 'n_words': 331}