workerholic 0.0.2

data/lib/workerholic/manager.rb

@@ -0,0 +1,40 @@
+module Workerholic
+  # Handles polling from Redis and hands job to worker
+  class Manager
+    attr_reader :workers, :scheduler, :worker_balancer
+
+    def initialize(opts = {})
+      @workers = []
+      Workerholic.workers_count.times { @workers << Worker.new }
+
+      @scheduler = JobScheduler.new
+      @worker_balancer = WorkerBalancer.new(workers: workers, auto_balance: opts[:auto_balance])
+    end
+
+    def start
+      worker_balancer.start
+      workers.each(&:work)
+      scheduler.start
+      sleep
+    rescue SystemExit, Interrupt
+      puts "\nWorkerholic is now shutting down. We are letting the workers finish their current jobs..."
+      shutdown
+      exit
+    end
+
+    def shutdown
+      workers.each(&:kill)
+      worker_balancer.kill
+      scheduler.kill
+    end
+
+=begin
+    def regenerate_workers
+      inactive_workers = WORKERS_COUNT - workers.size
+      if inactive_workers > 0
+        inactive_workers.times { @workers << Worker.new }
+      end
+    end
+=end
+  end
+end

data/lib/workerholic/queue.rb

@@ -0,0 +1,30 @@
+module Workerholic
+  # Handles background job enqueueing/dequeuing functionality
+  class Queue
+    attr_reader :storage, :name
+
+    def initialize(name = 'workerholic:queue:main')
+      @storage = Storage::RedisWrapper.new
+      @name = name
+      @logger = LogManager.new
+    end
+
+    def enqueue(serialized_job)
+      storage.push(name, serialized_job)
+      @logger.log('info', "Your job was placed in the #{name} queue on #{Time.now}.")
+    end
+
+    def dequeue
+      job_info = storage.pop(name)
+      job_info.last if job_info
+    end
+
+    def empty?
+      storage.list_length(name).zero?
+    end
+
+    def size
+      storage.list_length(name)
+    end
+  end
+end

data/lib/workerholic/sorted_set.rb

@@ -0,0 +1,26 @@
+module Workerholic
+  class SortedSet
+    attr_reader :storage, :name
+
+    def initialize(name = 'workerholic:scheduled_jobs')
+      @storage = Storage::RedisWrapper.new
+      @name = name
+    end
+
+    def add(serialized_job, score)
+      storage.add_to_set(name, score, serialized_job)
+    end
+
+    def remove(score)
+      storage.remove_from_set(name, score)
+    end
+
+    def peek
+      storage.peek(name)
+    end
+
+    def empty?
+      storage.set_empty?(name) == 0
+    end
+  end
+end

data/lib/workerholic/statistics.rb

@@ -0,0 +1,21 @@
+module Workerholic
+  class Statistics
+    attr_accessor :enqueued_at, :retry_count, :errors, :started_at, :completed_at
+
+    def initialize(options={})
+      @enqueued_at = options[:enqueued_at]
+      @errors = options[:errors] || []
+      @started_at = options[:started_at]
+      @completed_at = options[:completed_at]
+    end
+
+    def to_hash
+      {
+        enqueued_at: enqueued_at,
+        errors: errors,
+        started_at: started_at,
+        completed_at: completed_at
+      }
+    end
+  end
+end

data/lib/workerholic/storage.rb

@@ -0,0 +1,80 @@
+module Workerholic
+  class Storage
+    # Wraps redis-rb gem methods for enqueueing/dequeuing purposes
+    class RedisWrapper
+      attr_reader :redis, :retries
+
+      def initialize
+        @retries = 0
+        @redis = Workerholic.redis_pool
+
+        redis.with do |conn|
+          conn.ping
+        end
+      end
+
+      def list_length(key)
+        execute { |conn| conn.llen(key) }
+      end
+
+      def push(key, value)
+        execute { |conn| conn.rpush(key, value) }
+      end
+
+      # blocking pop from Redis queue
+      def pop(key, timeout = 1)
+        execute { |conn| conn.blpop(key, timeout) }
+      end
+
+      def add_to_set(key, score, value)
+        execute { |conn| conn.zadd(key, score, value) }
+      end
+
+      def peek(key)
+        execute { |conn| conn.zrange(key, 0, 0, with_scores: true).first }
+      end
+
+      def remove_from_set(key, score)
+        execute { |conn| conn.zremrangebyscore(key, score, score) }
+      end
+
+      def set_empty?(key)
+        execute { |conn| conn.zcount(key, 0, '+inf') }
+      end
+
+      def fetch_queue_names
+        execute { |conn| conn.scan(0, match: 'workerholic:queue*').last }
+      end
+
+      class RedisCannotRecover < Redis::CannotConnectError; end
+
+      private
+
+      def execute
+        begin
+          result = redis.with { |conn| yield conn }
+          reset_retries
+        rescue Redis::CannotConnectError
+          # LogManager might want to output our retries to the user
+          @retries += 1
+          if retries_exhausted?
+            raise RedisCannotRecover, 'Redis reconnect retries exhausted. Main Workerholic thread will be terminated now.'
+          end
+
+          sleep(5)
+          retry
+        end
+
+        result
+      end
+
+      def retries_exhausted?
+        retries == 5
+      end
+
+      def reset_retries
+        @retries = 0
+      end
+    end
+  end
+end

data/lib/workerholic/worker.rb

@@ -0,0 +1,43 @@
+module Workerholic
+  # handles job execution in threads
+  class Worker
+    attr_reader :thread
+    attr_accessor :alive, :queue
+
+    def initialize(queue=nil)
+      @queue = queue
+      @alive = true
+      @logger = LogManager.new
+    end
+
+    def work
+      @thread = Thread.new do
+        while alive
+          serialized_job = poll
+          JobProcessor.new(serialized_job).process if serialized_job
+        end
+
+        #puts "DONE!"
+      end
+    rescue ThreadError => e
+      @logger.log('info', e.message)
+      raise Interrupt
+    end
+
+    def kill
+      self.alive = false
+      thread.join if thread
+    end
+
+    private
+
+    def poll
+      if queue
+        queue.dequeue
+      else
+        sleep 0.1
+        nil
+      end
+    end
+  end
+end

data/lib/workerholic/worker_balancer.rb

@@ -0,0 +1,128 @@
+module Workerholic
+  class WorkerBalancer
+    attr_reader :storage, :workers, :thread, :alive, :auto
+    attr_accessor :queues
+
+    def initialize(opts = {})
+      @storage = Storage::RedisWrapper.new
+      @queues = fetch_queues
+      @workers = opts[:workers] || []
+      @alive = true
+      @logger = LogManager.new
+      @auto = opts[:auto_balance]
+    end
+
+    def start
+      if auto
+        auto_balance_workers
+      else
+        evenly_balance_workers
+      end
+    end
+
+    def kill
+      thread.kill
+      thread.join
+    end
+
+    private
+
+    def auto_balance_workers
+      @thread = Thread.new do
+        while alive
+          self.queues = fetch_queues
+
+          total_workers_count = assign_one_worker_per_queue
+
+          remaining_workers_count = workers.size - (total_workers_count + 1)
+          average_job_count_per_worker = total_jobs / remaining_workers_count.to_f
+
+          queues.each do |queue|
+            workers_count = queue.size / average_job_count_per_worker
+
+            if workers_count % 1 == 0.5
+              workers_count = workers_count.floor
+            else
+              workers_count = workers_count.round
+            end
+
+            assign_workers_to_queue(queue, workers_count, total_workers_count)
+
+            total_workers_count += workers_count
+          end
+
+          distribute_unassigned_worker(total_workers_count)
+          output_balancer_stats
+
+          sleep 2
+        end
+      end
+    end
+
+    def evenly_balance_workers
+      @thread = Thread.new do
+        while alive
+          self.queues = fetch_queues
+
+          total_workers_count = assign_one_worker_per_queue
+
+          remaining_workers_count = workers.size - (total_workers_count + 1)
+
+          queues.each do |queue|
+            workers_count = remaining_workers_count / queues.size
+            assign_workers_to_queue(queue, workers_count, total_workers_count)
+            total_workers_count += workers_count
+          end
+
+          distribute_unassigned_worker(total_workers_count)
+          output_balancer_stats
+
+          sleep 2
+        end
+      end
+    end
+
+    def distribute_unassigned_worker(total_workers_count)
+      workers[workers.size - 1].queue = queues.sample if workers.size - total_workers_count == 1
+    end
+
+    def output_balancer_stats
+      @logger.log('info', queues.map { |q| { name: q.name, size: q.size } })
+      @logger.log('info', current_workers_count_per_queue)
+    end
+
+    def assign_one_worker_per_queue
+      index = 0
+      while index < queues.size && index < workers.size
+        workers[index].queue = queues[index]
+        index += 1
+      end
+
+      index
+    end
+
+    def fetch_queues
+      storage.fetch_queue_names.map { |queue_name| Queue.new(queue_name) }
+    end
+
+    def total_jobs
+      @queues.map(&:size).reduce(:+) || 0
+    end
+
+    def assign_workers_to_queue(queue, workers_count, total_workers_count)
+      total_workers_count.upto(total_workers_count + workers_count - 1) do |i|
+        workers.to_a[i].queue = Queue.new(queue.name)
+      end
+    end
+
+    def current_workers_count_per_queue
+      workers.reduce({}) do |result, worker|
+        if worker.queue
+          result[worker.queue.name] = result[worker.queue.name] ? result[worker.queue.name] + 1 : 1
+        end
+
+        result
+      end
+    end
+  end
+end

data/spec/helpers/helper_methods.rb

@@ -0,0 +1,15 @@
+TEST_QUEUE = 'workerholic:queue:_test_queue'
+TEST_SCHEDULED_SORTED_SET = 'workerholic:test:scheduled_jobs'
+
+def expect_during(duration_in_secs, target)
+  timeout = Time.now.to_f + duration_in_secs
+
+  while Time.now.to_f <= timeout
+    result = yield
+    return if result == target
+
+    sleep(0.001)
+  end
+
+  expect(result).to eq(target)
+end

data/spec/helpers/job_tests.rb

@@ -0,0 +1,17 @@
+class SimpleJobTest
+  include Workerholic::Job
+  job_options queue_name: TEST_QUEUE
+
+  def perform(s)
+    s
+  end
+end
+
+class ComplexJobTest
+  include Workerholic::Job
+  job_options queue_name: TEST_QUEUE
+
+  def perform(arg1, arg2, arg3)
+    [arg1, arg2, arg3]
+  end
+end

data/spec/integration/dequeuing_and_job_processing_spec.rb

@@ -0,0 +1,24 @@
+require_relative '../spec_helper'
+
+describe 'dequeuing and processesing of jobs' do
+  let(:redis) { Redis.new }
+  before { redis.del(TEST_QUEUE) }
+
+  xit 'successfully dequeues and process a simple job' do
+    serialized_job = Workerholic::JobSerializer.serialize(
+      class: SimpleJobTest,
+      arguments: ['test job']
+    )
+    redis.rpush(TEST_QUEUE, serialized_job)
+    manager = Workerholic::Manager.new
+
+    Thread.new { manager.start }
+    expect_during(1, false) { redis.exists(TEST_QUEUE) }
+  end
+
+  it 'successfully dequeues and process a complex job'
+
+  context 'user interrupts process' do
+    it 'finishes executing the current job before gracefully shutting down'
+  end
+end

data/spec/integration/enqueuing_jobs_spec.rb

@@ -0,0 +1,53 @@
+require_relative '../spec_helper'
+
+describe 'enqueuing jobs to Redis' do
+  let(:redis) { Redis.new }
+  before { redis.del(TEST_QUEUE) }
+
+  context 'successfully creates a job and enqueues it in Redis' do
+    it 'enqueues a simple job in redis' do
+      SimpleJobTest.new.perform_async('test job')
+      serialized_job = redis.lpop(TEST_QUEUE)
+      job_from_redis = Workerholic::JobSerializer.deserialize(serialized_job)
+
+      expected_job = Workerholic::JobWrapper.new(class: SimpleJobTest, arguments: ['test job'])
+      expected_job.statistics.enqueued_at = job_from_redis.statistics.enqueued_at
+
+      expect(job_from_redis.to_hash).to eq(expected_job.to_hash)
+    end
+
+    it 'enqueues a complex job in redis' do
+      ComplexJobTest.new.perform_async('test job', { a: 1, b: 2 }, [1, 2, 3])
+      serialized_job = redis.lpop(TEST_QUEUE)
+      job_from_redis = Workerholic::JobSerializer.deserialize(serialized_job)
+
+      expected_job = Workerholic::JobWrapper.new(
+        class: ComplexJobTest,
+        arguments: ['test job', { a: 1, b: 2 }, [1, 2, 3]]
+      )
+      expected_job.statistics.enqueued_at = job_from_redis.statistics.enqueued_at
+
+      expect(job_from_redis.to_hash).to eq(expected_job.to_hash)
+    end
+
+    it 'enqueues a job with the right statistics' do
+      SimpleJobTest.new.perform_async('test_job')
+      serialized_job = redis.lpop(TEST_QUEUE)
+      job_from_redis = Workerholic::JobSerializer.deserialize(serialized_job)
+
+      expect(job_from_redis.statistics.enqueued_at).to be < Time.now.to_f
+    end
+  end
+
+  context 'handles user errors' do
+    it 'raises an error if Redis server is not running' do
+      allow(Workerholic::Storage::RedisWrapper).to receive(:new).and_raise(Redis::CannotConnectError)
+
+      expect { SimpleJobTest.new.perform_async('test job') }.to raise_error(Redis::CannotConnectError)
+    end
+
+    it 'raises an error when wrong number of arguments is specified to perform_async' do
+      expect { SimpleJobTest.new.perform_async(1, 2, 3) }.to raise_error(ArgumentError)
+    end
+  end
+end