<commit_before>require 'monitor'

module Que
  class Worker
    # Each worker has a thread that does the actual work of running jobs.
    # Since both the worker's thread and whatever thread is managing the
    # worker are capable of affecting the worker's state, we need to
    # synchronize access to it.
    include MonitorMixin

    # A custom exception to immediately kill a worker and its current job.
    class Stop < Interrupt; end

    attr_reader :thread, :state

    def initialize
      super # For MonitorMixin.
      @state  = :working
      @thread = Thread.new { work_loop }
    end

    def alive?
      !!@thread.status
    end

    def sleeping?
      synchronize do
        if @state == :sleeping
          # There's a very small period of time between when the Worker marks
          # itself as sleeping and when it actually goes to sleep. Only report
          # true when we're certain the thread is sleeping.
          wait until @thread.status == 'sleep'
          true
        end
      end
    end

    def working?
      synchronize { @state == :working }
    end

    def wake!
      synchronize do
        if sleeping?
          # Have to set the state here so that another thread checking
          # immediately after this won't see the worker as asleep.
          @state = :working
          @thread.wakeup
          true
        end
      end
    end

    # #stop informs the worker that it should shut down after its next job,
    # while #stop! kills the job and worker immediately. #stop! is bad news
    # because its results are unpredictable (it can leave the DB connection
    # in an unusable state), so it should only be used when we're shutting
    # down the whole process anyway and side effects aren't a big deal.
    def stop
      synchronize do
        @stop = true
        @thread.wakeup if sleeping?
      end
    end

    def stop!
      @thread.raise Stop
    end

    def wait_until_stopped
      wait while alive?
    end

    private

    # Sleep very briefly while waiting for a thread to get somewhere.
    def wait
      sleep 0.0001
    end

    def work_loop
      loop do
        job = Job.work

        # Grab the lock and figure out what we should do next.
        synchronize { @state = :sleeping unless @stop || job }

        sleep if @state == :sleeping
        break if @stop
      end
    rescue Stop
      # This process is shutting down; let it.
    ensure
      @state = :stopped
    end

    # Defaults for the Worker pool.
    @sleep_period = 5

    class << self
      attr_reader :mode, :sleep_period

      def mode=(mode)
        return if @mode == mode
        set_mode(mode)

        case mode
        when :async
          wrangler # Make sure the wrangler thread is initialized.
          set_worker_count 4 if worker_count.zero?
        else
          set_worker_count 0
        end
      end

      def workers
        @workers ||= []
      end

      def worker_count=(count)
        set_mode(count > 0 ? :async : :off)
        set_worker_count(count)
      end

      def worker_count
        workers.count
      end

      def sleep_period=(period)
        @sleep_period = period
        wrangler.wakeup if period
      end

      def stop!
        # The behavior of Worker#stop! is unpredictable - what it does is
        # dependent on what the Worker is currently doing. Sometimes it won't
        # work the first time, so we need to try again, but sometimes it'll
        # never do anything, so we can't repeat indefinitely. So, compromise.
        5.times do
          break if workers.select(&:alive?).each(&:stop!).none?
          sleep 0.001
        end

        workers.clear
      end

      def wake!
        workers.find &:wake!
      end

      def wake_all!
        workers.each &:wake!
      end

      private

      def wrangler
        @wrangler ||= Thread.new { loop { sleep(*sleep_period); wake! } }
      end

      def set_mode(mode)
        return if mode == @mode
        @mode = mode
        Que.log :info, "Set mode to #{mode.inspect}"
      end

      def set_worker_count(count)
        return if count == worker_count
        Que.log :info, "Set worker_count to #{count.inspect}"

        if count > worker_count
          (count - worker_count).times { workers << new }
        elsif count < worker_count
          workers.pop(worker_count - count).each(&:stop).each(&:wait_until_stopped)
        end
      end
    end
  end
end
<commit_msg>Let the wrangler thread run all the time - don't bother with initialization.<commit_after>require 'monitor'

module Que
  class Worker
    # Each worker has a thread that does the actual work of running jobs.
    # Since both the worker's thread and whatever thread is managing the
    # worker are capable of affecting the worker's state, we need to
    # synchronize access to it.
    include MonitorMixin

    # A custom exception to immediately kill a worker and its current job.
    class Stop < Interrupt; end

    attr_reader :thread, :state

    def initialize
      super # For MonitorMixin.
      @state  = :working
      @thread = Thread.new { work_loop }
    end

    def alive?
      !!@thread.status
    end

    def sleeping?
      synchronize do
        if @state == :sleeping
          # There's a very small period of time between when the Worker marks
          # itself as sleeping and when it actually goes to sleep. Only report
          # true when we're certain the thread is sleeping.
          wait until @thread.status == 'sleep'
          true
        end
      end
    end

    def working?
      synchronize { @state == :working }
    end

    def wake!
      synchronize do
        if sleeping?
          # Have to set the state here so that another thread checking
          # immediately after this won't see the worker as asleep.
          @state = :working
          @thread.wakeup
          true
        end
      end
    end

    # #stop informs the worker that it should shut down after its next job,
    # while #stop! kills the job and worker immediately. #stop! is bad news
    # because its results are unpredictable (it can leave the DB connection
    # in an unusable state), so it should only be used when we're shutting
    # down the whole process anyway and side effects aren't a big deal.
    def stop
      synchronize do
        @stop = true
        @thread.wakeup if sleeping?
      end
    end

    def stop!
      @thread.raise Stop
    end

    def wait_until_stopped
      wait while alive?
    end

    private

    # Sleep very briefly while waiting for a thread to get somewhere.
    def wait
      sleep 0.0001
    end

    def work_loop
      loop do
        job = Job.work

        # Grab the lock and figure out what we should do next.
        synchronize { @state = :sleeping unless @stop || job }

        sleep if @state == :sleeping
        break if @stop
      end
    rescue Stop
      # This process is shutting down; let it.
    ensure
      @state = :stopped
    end

    @sleep_period = 5
    @wrangler = Thread.new { loop { sleep(*@sleep_period); wake! } }

    class << self
      attr_reader :mode, :sleep_period

      def mode=(mode)
        return if @mode == mode
        set_mode(mode)

        case mode
        when :async
          set_worker_count 4 if worker_count.zero?
        else
          set_worker_count 0
        end
      end

      def workers
        @workers ||= []
      end

      def worker_count=(count)
        set_mode(count > 0 ? :async : :off)
        set_worker_count(count)
      end

      def worker_count
        workers.count
      end

      def sleep_period=(period)
        @sleep_period = period
        @wrangler.wakeup if period
      end

      def stop!
        # The behavior of Worker#stop! is unpredictable - what it does is
        # dependent on what the Worker is currently doing. Sometimes it won't
        # work the first time, so we need to try again, but sometimes it'll
        # never do anything, so we can't repeat indefinitely. So, compromise.
        5.times do
          break if workers.select(&:alive?).each(&:stop!).none?
          sleep 0.001
        end

        workers.clear
      end

      def wake!
        workers.find &:wake!
      end

      def wake_all!
        workers.each &:wake!
      end

      private

      def set_mode(mode)
        return if mode == @mode
        @mode = mode
        Que.log :info, "Set mode to #{mode.inspect}"
      end

      def set_worker_count(count)
        return if count == worker_count
        Que.log :info, "Set worker_count to #{count.inspect}"

        if count > worker_count
          (count - worker_count).times { workers << new }
        elsif count < worker_count
          workers.pop(worker_count - count).each(&:stop).each(&:wait_until_stopped)
        end
      end
    end
  end
end

{'fraction_non_alphanumeric': 0.06403885249604699, 'fraction_numerical': 0.003162412468940592, 'mean_word_length': 1.8656957928802589, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 7, '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': '3258956', 'n_tokens_mistral': 2491, 'n_tokens_neox': 2382, 'n_words': 983}