mperham--sidekiq/lib/sidekiq/server.rb

require 'celluloid'
require 'redis'
require 'multi_json'

require 'sidekiq/util'
require 'sidekiq/worker'

module Sidekiq

  ##
  # This is the main router in the system.  This
  # manages the worker state and fetches messages
  # from Redis to be dispatched to ready workers.
  #
  class Server
    include Util
    include Celluloid

    trap_exit :worker_died

    def initialize(location, options={})
      log "Booting sidekiq #{Sidekiq::VERSION} with Redis at #{location}"
      verbose options.inspect
      @count = options[:worker_count]
      @queues = options[:queues]
      @queue_idx = 0
      @queues_size = @queues.size
      @redis = Redis.connect(:url => location)

      @done = false
      @busy = []
      @ready = []
      @count.times do
        @ready << Worker.new_link(current_actor)
      end
    end

    def stop
      @done = true
      @ready.each(&:terminate)
      @ready.clear

      after(5) do
        signal(:shutdown)
      end
    end

    def start
      dispatch(true)
    end

    def worker_done(worker)
      @busy.delete(worker)
      if stopped?
        worker.terminate
      else
        @ready << worker
      end
      dispatch
    end

    def worker_died(worker, reason)
      @busy.delete(worker)

      if reason
        log "Worker death: #{reason}"
        log reason.backtrace.join("\n")
      end

      unless stopped?
        @ready << Worker.new_link(current_actor)
        dispatch
      end
    end

    private

    def find_work(queue_idx)
      current_queue = @queues[queue_idx]
      msg = @redis.lpop("queue:#{current_queue}")
      if msg
        worker = @ready.pop
        @busy << worker
        worker.process! MultiJson.decode(msg)
      end
      msg
    end

    def dispatch(schedule = false)
      watchdog("Fatal error in sidekiq, dispatch loop died") do
        return if stopped?

        # Our dispatch loop
        # Loop through the queues, looking for a message in each.
        queue_idx = 0
        found = false
        loop do
          # return so that we don't dispatch again until worker_done
          break verbose('no workers') if @ready.size == 0

          found ||= find_work(queue_idx)
          queue_idx += 1

          # if we find no messages in any of the queues, we can break
          # out of the loop.  Otherwise we loop again.
          lastq = (queue_idx % @queues.size == 0)
          if lastq && !found
            verbose('nothing to process'); break
          elsif lastq
            queue_idx = 0
            found = false
          end
        end

        after(1) { verbose('ping'); dispatch(schedule) } if schedule
      end
    end

    def stopped?
      @done
    end
  end
end