sidekiq-hierarchy 0.1.1

data/lib/sidekiq/hierarchy/notifications.rb

@@ -0,0 +1,8 @@
+module Sidekiq
+  module Hierarchy
+    module Notifications
+      JOB_UPDATE = :job_update
+      WORKFLOW_UPDATE = :workflow_update
+    end
+  end
+end

data/lib/sidekiq/hierarchy/observers.rb

@@ -0,0 +1,9 @@
+require 'sidekiq/hierarchy/observers/job_update'
+require 'sidekiq/hierarchy/observers/workflow_update'
+
+module Sidekiq
+  module Hierarchy
+    module Observers
+    end
+  end
+end

data/lib/sidekiq/hierarchy/observers/job_update.rb

@@ -0,0 +1,15 @@
+module Sidekiq
+  module Hierarchy
+    module Observers
+      class JobUpdate
+        def register(callback_registry)
+          callback_registry.subscribe(Notifications::JOB_UPDATE, self)
+        end
+
+        def call(job, status, old_status)
+          job.workflow.update_status(status)
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/hierarchy/observers/workflow_update.rb

@@ -0,0 +1,18 @@
+module Sidekiq
+  module Hierarchy
+    module Observers
+      class WorkflowUpdate
+        def register(callback_registry)
+          callback_registry.subscribe(Notifications::WORKFLOW_UPDATE, self)
+        end
+
+        def call(workflow, status, old_status)
+          from_set = WorkflowSet.for_status(old_status)
+          to_set = WorkflowSet.for_status(status)
+
+          to_set.move(workflow, from_set)  # Move/add to the new status set
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/hierarchy/rack/middleware.rb

@@ -0,0 +1,27 @@
+require 'rack'
+require 'sidekiq/hierarchy/http'
+
+module Sidekiq
+  module Hierarchy
+    module Rack
+      class Middleware
+        # transform from http header to rack names
+        JID_HEADER_KEY = "HTTP_#{Sidekiq::Hierarchy::Http::JID_HEADER.upcase.gsub('-','_')}".freeze
+        WORKFLOW_HEADER_KEY = "HTTP_#{Sidekiq::Hierarchy::Http::WORKFLOW_HEADER.upcase.gsub('-','_')}".freeze
+
+        def initialize(app)
+          @app = app
+        end
+
+        def call(env)
+          Sidekiq::Hierarchy.current_jid = env[JID_HEADER_KEY]
+          Sidekiq::Hierarchy.current_workflow = Workflow.find_by_jid(env[WORKFLOW_HEADER_KEY]) if env[WORKFLOW_HEADER_KEY]
+          @app.call(env)
+        ensure
+          Sidekiq::Hierarchy.current_workflow = nil
+          Sidekiq::Hierarchy.current_jid = nil
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/hierarchy/server/middleware.rb

@@ -0,0 +1,62 @@
+module Sidekiq
+  module Hierarchy
+    module Server
+      class Middleware
+        def initialize(options={})
+        end
+
+        # Wraps around the actual execution of a job. Takes params:
+        #   worker - the instance of the worker to be used for execution
+        #   msg - the hash of job info, something like {'class' => 'HardWorker', 'args' => [1, 2, 'foo'], 'retry' => true}
+        #   queue - the named queue to use
+        # Must propagate return value upwards.
+        # Since jobs raise errors for signalling, those must be propagated as well.
+        def call(worker, msg, queue)
+          if msg['workflow'] == true  # root job -- start of a new workflow
+            Sidekiq::Hierarchy.current_workflow = Workflow.find_by_jid(worker.jid)
+          elsif msg['workflow'].is_a?(String)  # child job -- inherit parent's workflow
+            Sidekiq::Hierarchy.current_workflow = Workflow.find_by_jid(msg['workflow'])
+          end
+          Sidekiq::Hierarchy.current_jid = worker.jid
+
+          Sidekiq::Hierarchy.record_job_running
+          ret = yield
+          Sidekiq::Hierarchy.record_job_complete
+
+          ret
+        rescue Exception => e
+          if exception_caused_by_shutdown?(e) || retries_remaining?(msg)
+            # job will be pushed back onto queue during hard_shutdown or if retries are permitted
+            Sidekiq::Hierarchy.record_job_requeued
+          else
+            Sidekiq::Hierarchy.record_job_failed
+          end
+
+          raise
+        end
+
+        def retries_remaining?(msg)
+          return false unless msg['retry']
+
+          retry_count = msg['retry_count'] || 0
+          max_retries = if msg['retry'].is_a?(Fixnum)
+                          msg['retry']
+                        else
+                          Sidekiq::Middleware::Server::RetryJobs::DEFAULT_MAX_RETRY_ATTEMPTS
+                        end
+
+          # this check requires prepending the middleware before sidekiq's builtin retry
+          retry_count < max_retries
+        end
+        private :retries_remaining?
+
+        def exception_caused_by_shutdown?(e)
+          e.instance_of?(Sidekiq::Shutdown) ||
+            # In Ruby 2.1+, check if original exception was Shutdown
+            (defined?(e.cause) && exception_caused_by_shutdown?(e.cause))
+        end
+        private :exception_caused_by_shutdown?
+      end
+    end
+  end
+end

data/lib/sidekiq/hierarchy/version.rb

@@ -0,0 +1,5 @@
+module Sidekiq
+  module Hierarchy
+    VERSION = '0.1.1'
+  end
+end

data/lib/sidekiq/hierarchy/web.rb

@@ -0,0 +1,149 @@
+# Web interface to Sidekiq-hierarchy
+# Optimised for ease-of-use, not efficiency; it's probably best
+# not to leave this open in a tab forever.
+# Sidekiq seems to use Bootstrap 3.0.0 currently; find docs at
+# http://bootstrapdocs.com/v3.0.0/docs/
+module Sidekiq
+  module Hierarchy
+    module Web
+      module Helpers
+        def sidekiq_hierarchy_template(template_name)
+          File.read(File.join(VIEW_PATH, "#{template_name}.erb"))
+        end
+
+        def job_url(job=nil)
+          "#{root_path}hierarchy/jobs/#{job.jid if job}"
+        end
+
+        def workflow_url(workflow=nil)
+          "#{root_path}hierarchy/workflows/#{workflow.jid if workflow}"
+        end
+
+        def workflow_set_url(status)
+          "#{root_path}hierarchy/workflow_sets/#{status}"
+        end
+
+        def status_updated_at(job)
+          case job.status
+          when :enqueued
+            job.enqueued_at
+          when :running, :requeued
+            job.run_at
+          when :complete
+            job.complete_at
+          when :failed
+            job.failed_at
+          end
+        end
+
+        def bootstrap_status(status)
+          case status
+          when :enqueued, :requeued
+            'warning'
+          when :running
+            'info'
+          when :complete
+            'success'
+          when :failed
+            'danger'
+          end
+        end
+      end
+
+      VIEW_PATH = File.expand_path('../../../../web/views', __FILE__)
+      PER_PAGE = 20
+
+      def self.registered(app)
+        app.helpers Helpers
+
+        app.not_found do
+          erb sidekiq_hierarchy_template(:not_found)
+        end
+
+        app.get '/hierarchy/?' do
+          running_set = RunningSet.new
+          complete_set = CompleteSet.new
+          failed_set = FailedSet.new
+
+          @running = running_set.each.take(PER_PAGE)
+          @complete = complete_set.each.take(PER_PAGE)
+          @failed = failed_set.each.take(PER_PAGE)
+
+          erb sidekiq_hierarchy_template(:status)
+        end
+
+        app.delete '/hierarchy/?' do
+          [RunningSet.new, CompleteSet.new, FailedSet.new].each(&:remove_all)
+          redirect back
+        end
+
+        app.get '/hierarchy/workflow_sets/:status' do |status|
+          @status = status.to_sym
+          if workflow_set = WorkflowSet.for_status(@status)
+            @workflows = workflow_set.each.take(PER_PAGE)
+            erb sidekiq_hierarchy_template(:workflow_set)
+          else
+            halt 404
+          end
+        end
+
+        app.delete '/hierarchy/workflow_sets/:status' do |status|
+          @status = status.to_sym
+          if workflow_set = WorkflowSet.for_status(@status)
+            workflow_set.each(&:delete)
+            redirect back
+          else
+            halt 404
+          end
+        end
+
+        app.get '/hierarchy/workflows/?' do
+          if params['workflow_jid'] =~ /\A\h{24}\z/
+            redirect to("/hierarchy/workflows/#{params['workflow_jid']}")
+          else
+            redirect to(:hierarchy)
+          end
+        end
+
+        app.get %r{\A/hierarchy/workflows/(\h{24})\z} do |workflow_jid|
+          @workflow = Workflow.find_by_jid(workflow_jid)
+          if @workflow.exists?
+            erb sidekiq_hierarchy_template(:workflow)
+          else
+            halt 404
+          end
+        end
+
+        app.delete %r{\A/hierarchy/workflows/(\h{24})\z} do |workflow_jid|
+          workflow = Workflow.find_by_jid(workflow_jid)
+          redirect_url = "/hierarchy/workflow_sets/#{workflow.status}"
+          workflow.delete
+
+          redirect to(redirect_url)
+        end
+
+        app.get '/hierarchy/jobs/?' do
+          if params['jid'] =~ /\A\h{24}\z/
+            redirect to("/hierarchy/jobs/#{params['jid']}")
+          else
+            redirect back
+          end
+        end
+
+        app.get %r{\A/hierarchy/jobs/(\h{24})\z} do |jid|
+          @job = Job.find(jid)
+          @workflow = @job.workflow
+          if @job.exists? && @workflow.exists?
+            erb sidekiq_hierarchy_template(:job)
+          else
+            halt 404
+          end
+        end
+      end
+    end
+  end
+end
+
+require 'sidekiq/web' unless defined?(Sidekiq::Web)
+Sidekiq::Web.register(Sidekiq::Hierarchy::Web)
+Sidekiq::Web.tabs['Hierarchy'] = 'hierarchy'

data/lib/sidekiq/hierarchy/workflow.rb

@@ -0,0 +1,130 @@
+module Sidekiq
+  module Hierarchy
+    class Workflow
+      extend Forwardable
+
+      attr_reader :root
+
+      def initialize(root)
+        @root = root
+      end
+
+      class << self
+        alias_method :find, :new
+
+        def find_by_jid(root_jid)
+          find(Job.find(root_jid))
+        end
+      end
+
+      delegate [:jid, :[], :[]=, :exists?] => :@root
+
+      def ==(other_workflow)
+        other_workflow.instance_of?(self.class) &&
+          self.jid == other_workflow.jid
+      end
+
+      def workflow_set
+        WorkflowSet.for_status(status)
+      end
+
+      def delete
+        wset = workflow_set  # save it for later
+        root.delete  # deleting nodes is more important than a dangling reference
+        wset.remove(self) if wset  # now we can clear out from the set
+      end
+
+      # Walks the tree in DFS order (for optimal completion checking)
+      # Returns an Enumerator; use #to_a to get an array instead
+      def jobs
+        to_visit = [root]
+        Enumerator.new do |y|
+          while node = to_visit.pop
+            y << node  # sugar for yielding a value
+            to_visit += node.children
+          end
+        end
+      end
+
+
+      ### Status
+
+      def status
+        case self[Job::WORKFLOW_STATUS_FIELD]
+        when Job::STATUS_RUNNING
+          :running
+        when Job::STATUS_COMPLETE
+          :complete
+        when Job::STATUS_FAILED
+          :failed
+        else
+          :unknown
+        end
+      end
+
+      def update_status(from_job_status)
+        old_status = status
+        return if [:failed, :complete].include?(old_status)  # these states are final
+
+        if [:enqueued, :running, :requeued].include?(from_job_status)
+          new_status, s_val = :running, Job::STATUS_RUNNING
+        elsif from_job_status == :failed
+          new_status, s_val = :failed, Job::STATUS_FAILED
+        elsif from_job_status == :complete && jobs.all?(&:complete?)
+          new_status, s_val = :complete, Job::STATUS_COMPLETE
+        end
+
+        return if !new_status || new_status == old_status  # don't publish null updates
+        self[Job::WORKFLOW_STATUS_FIELD] = s_val
+
+        Sidekiq::Hierarchy.publish(Notifications::WORKFLOW_UPDATE, self, new_status, old_status)
+      end
+
+      def running?
+        status == :running
+      end
+
+      def complete?
+        status == :complete
+      end
+
+      def failed?
+        status == :failed
+      end
+
+
+      ### Calculated metrics
+
+      def enqueued_at
+        root.enqueued_at
+      end
+
+      def run_at
+        root.run_at
+      end
+
+      # Returns the time at which all jobs were complete;
+      # nil if any jobs are still incomplete
+      def complete_at
+        jobs.map(&:complete_at).max if complete?
+      end
+
+      # Returns the earliest time at which a job failed;
+      # nil if none did
+      def failed_at
+        jobs.map(&:failed_at).compact.min if failed?
+      end
+
+
+      ### Serialisation
+
+      def as_json(options={})
+        root.as_json(options)
+      end
+
+      def to_s
+        Sidekiq.dump_json(self.as_json)
+      end
+    end
+  end
+end

data/lib/sidekiq/hierarchy/workflow_set.rb

@@ -0,0 +1,134 @@
+module Sidekiq
+  module Hierarchy
+
+    ### Implementations
+
+    # A sorted set of Workflows that permits enumeration
+    class WorkflowSet
+      PAGE_SIZE = 100
+
+      def self.for_status(status)
+        case status
+        when :running
+          RunningSet.new
+        when :complete
+          CompleteSet.new
+        when :failed
+          FailedSet.new
+        end
+      end
+
+      def initialize(status)
+        raise ArgumentError, 'status cannot be nil' if status.nil?
+        @status = status
+      end
+
+      def ==(other_workflow_set)
+        other_workflow_set.instance_of?(self.class)
+      end
+
+      def size
+        Sidekiq.redis { |conn| conn.zcard(redis_zkey) }
+      end
+
+      def add(workflow)
+        Sidekiq.redis { |conn| conn.zadd(redis_zkey, Time.now.to_f, workflow.jid) }
+      end
+
+      def contains?(workflow)
+        !!Sidekiq.redis { |conn| conn.zscore(redis_zkey, workflow.jid) }
+      end
+
+      # Remove a workflow from the set if it is present. This operation can
+      # only be executed as cleanup (i.e., on a workflow that has been
+      # unpersisted/deleted); otherwise it will fail in order to avoid
+      # memory leaks.
+      def remove(workflow)
+        raise 'Workflow still exists' if workflow.exists?
+        Sidekiq.redis { |conn| conn.zrem(redis_zkey, workflow.jid) }
+      end
+
+      # Move a workflow to this set from its current one
+      # This should really be done in Lua, but unit testing support is just not there,
+      # so there is a potential race condition in which a workflow could end up in
+      # multiple sets. the effect of this is minimal, so we'll fix it later.
+      def move(workflow, from_set=nil)
+        Sidekiq.redis do |conn|
+          conn.multi do
+            conn.zrem(from_set.redis_zkey, workflow.jid) if from_set
+            conn.zadd(redis_zkey, Time.now.to_f, workflow.jid)
+          end.last
+        end
+      end
+
+      def each
+        return enum_for(:each) unless block_given?
+
+        elements = []
+        last_max_score = Time.now.to_f
+        loop do
+          elements = Sidekiq.redis do |conn|
+            conn.zrevrangebyscore(redis_zkey, last_max_score, '-inf', limit: [0, PAGE_SIZE], with_scores: true)
+              .drop_while { |elt| elements.include?(elt) }
+          end
+          break if elements.empty?
+          elements.each { |jid, _| yield Workflow.find_by_jid(jid) }
+          _, last_max_score = elements.last  # timestamp of last element
+        end
+      end
+
+      def redis_zkey
+        "hierarchy:set:#{@status}"
+      end
+    end
+
+    # An implementation of WorkflowSet that auto-prunes by time & size
+    # to stay within space constraints. Do _not_ use for workflows that
+    # cannot be lost (i.e., are in any state of progress, or require followup)
+    class PruningSet < WorkflowSet
+      def self.timeout
+        Sidekiq.options[:dead_timeout_in_seconds]
+      end
+
+      def self.max_workflows
+        Sidekiq.options[:dead_max_workflows] || Sidekiq.options[:dead_max_jobs]
+      end
+
+      def add(workflow)
+        prune
+        super
+      end
+
+      def prune
+        Sidekiq.redis do |conn|
+          conn.multi do
+            conn.zrangebyscore(redis_zkey, '-inf', Time.now.to_f - self.class.timeout)  # old workflows
+            conn.zrevrange(redis_zkey, self.class.max_workflows, -1)  # excess workflows
+          end.flatten.uniq  # take the union of both pruning strategies
+            .tap { |to_remove| conn.zrem(redis_zkey, to_remove) if to_remove.any? }
+        end.each { |jid| Workflow.find_by_jid(jid).delete }
+      end
+    end
+
+
+    ### Instances
+
+    class RunningSet < WorkflowSet
+      def initialize
+        super 'running'
+      end
+    end
+
+    class CompleteSet < PruningSet
+      def initialize
+        super 'complete'
+      end
+    end
+
+    class FailedSet < PruningSet
+      def initialize
+        super 'failed'
+      end
+    end
+  end
+end