work_shaper 0.1.3.1rc2 → 0.1.3.1rc4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 6ee5120c590964d5c74c75ef38a3c739d5849213d7c36006050a26d5868378ca
-  data.tar.gz: 1e02c5b221518cbb9e5c96891c831d783cfc14895cffa7e4a6a5e8a24b60f1f7
+  metadata.gz: 883a72259f3ae672a2456139e8b309da4cb72be8ae36e85293a4dfc5a21ad502
+  data.tar.gz: 4fa2fc7b182c6865fd4d761548676a7ada9ae903264f76056d317d90bd7607e2
 SHA512:
-  metadata.gz: dc444bf26f6bc591acf11d6bc283c19fa1a30fbfe33c215e6b102397de2847fabc4358e87a5c32632a9f770323239b87d6d35b7bc468aa4b3dbf767eb775eb27
-  data.tar.gz: efdea370299cbe549ffea400c69f3a386a7b14f2295567fa8844b4e5748a67e6878824eafad4763da583420c7657792935c0da5f0583d007dc8beea5b7fee050
+  metadata.gz: b4caa694ff8cf25082c05bea23b37254ff0d57bd6bbaa1b0a2eb081ceedca2c935dd395c3bfa01640d4ab635ab72535fc81d40d725cfe16ab933e1a8c367a0a5
+  data.tar.gz: 7106024e8f782582836661ad809dd3d033464cb6c2daa68c1e8e10764ea74eeffcf6b688cbdc9c43f93203e750bb3486bbbd24ac89cbbaffeec98a9808d9b5f1

data/Gemfile

@@ -14,3 +14,5 @@ gem "rubocop", "~> 1.21"
 gem "logger", "~> 1.4"
 
 gem "concurrent-ruby", "~> 1.2"
+
+gem 'simplecov', require: false, group: :test

data/lib/work_shaper/manager.rb

@@ -3,7 +3,7 @@ module WorkShaper
   # for each offset in monotonically increasing order (independent of the execution order), and gracefully
   # cleaning up when `#shutdown` is called.
   class Manager
-    attr_reader :total_acked, :total_enqueued
+    attr_reader :total_acked, :total_enqueued, :shutting_down
 
     # Several of the parameters here are Lambdas (not Proc). Note you can pass a method using
     # `method(:some_method)` or a lambda directly `->{ puts 'Hello'}`.
@@ -28,7 +28,7 @@ module WorkShaper
       @completed_offsets = {}
       @max_in_queue = max_in_queue
       @semaphore = Mutex.new
-      @shutdown = false
+      @shutting_down = false
 
       @total_enqueued = 0
       @total_acked = 0
@@ -58,14 +58,16 @@ module WorkShaper
 
     # Enqueue a message to be worked on the given `sub_key`, `partition`, and `offset`.
     def enqueue(sub_key, message, partition, offset)
-      raise StandardError, 'Shutting down' if @shutdown
+      raise StandardError, 'Shutting down' if @shutting_down
       pause_on_overrun
-      WorkShaper.logger.debug "Enqueue: #{sub_key}:#{partition}:#{offset}"
+
+      offset_holder = OffsetHolder.new(partition, offset)
+      WorkShaper.logger.debug "Enqueue: #{sub_key}/#{offset_holder}"
 
       worker = nil
       @semaphore.synchronize do
         @total_enqueued += 1
-        (@received_offsets[partition] ||= Array.new) << offset
+        (@received_offsets[partition] ||= Array.new) << offset_holder
 
         worker =
           @workers[sub_key] ||=
@@ -81,7 +83,7 @@ module WorkShaper
             )
       end
 
-      worker.enqueue(message, partition, offset)
+      worker.enqueue(message, offset_holder)
     end
 
     # Flush any offsets for which work has been completed. Only lowest continuous run of
@@ -108,7 +110,7 @@ module WorkShaper
         if detailed
           WorkShaper.logger.info(
             {
-              messaage: 'Reporting - Extra Detail',
+              message: 'Reporting - Extra Detail',
               pending_ack: @completed_offsets,
               received_offsets: @received_offsets
             })
@@ -118,10 +120,11 @@ module WorkShaper
 
     # Stop the underlying threads
     def shutdown
-      @shutdown = true
-      report(detailed: true)
+      @shutting_down = true
+      WorkShaper.logger.warn({ message: 'Shutting Down' })
       Thread.kill(@heartbeat)
       Thread.kill(@offset_manager)
+      report(detailed: true)
       @workers.each_value(&:shutdown)
     end
 
@@ -137,21 +140,35 @@ module WorkShaper
       completed = @completed_offsets[partition].sort!
       received = @received_offsets[partition].sort!
 
-      offset = completed.first
-      while received.any? && received.first == offset
-        # We observed Kafka sending the same message twice, even after
-        # having committed the offset. Here we skip this offset if we
-        # know it has already been committed.
-        last_offset = @last_ack[partition]
-        if last_offset && offset <= last_offset
-          WorkShaper.logger.warn(
-            { message: 'Received Duplicate Offset',
-              offset: "#{partition}:#{offset}",
-              last_acked: last_offset,
-            })
-        else
-          result = @ack.call(partition, offset)
-          if result.is_a? Exception
+      begin
+        offset = completed.first
+        while received.any? && received.first == offset
+          # We observed Kafka sending the same message twice, even after
+          # having committed the offset. Here we skip this offset if we
+          # know it has already been committed.
+          last_offset = @last_ack[partition]
+          if last_offset && offset <= last_offset
+            WorkShaper.logger.warn(
+              { message: 'Received Duplicate Offset',
+                offset: "#{partition}:#{offset}",
+                last_acked: last_offset,
+              })
+          end
+
+          result =
+            begin
+              @ack.call(partition, offset.to_i)
+            rescue => e
+              # We expect @ack to handle it's own errors and return the error or false if it
+              # is safe to continue. Otherwise @ack should raise an error and we will
+              # shutdown.
+              WorkShaper.logger.error({ message: 'Error in ack', error: e })
+              WorkShaper.logger.error(e.backtrace.join("\n"))
+              shutdown
+              break
+            end
+
+          if result.is_a? Exception || !result
             WorkShaper.logger.warn(
               { message: 'Failed to Ack Offset, likely re-balance',
                 offset: "#{partition}:#{offset}",
@@ -159,18 +176,21 @@ module WorkShaper
                 received: @received_offsets[partition].to_a[0..10].join(',')
               })
           else
-            @last_ack[partition] = offset
+            @last_ack[partition] = [@last_ack[partition] || offset, offset].max
           end
-        end
 
-        @total_acked += 1
-        WorkShaper.logger.debug "@total_acked: #{@total_acked}"
-        WorkShaper.logger.debug "completed: [#{completed.join(', ')}]"
-        WorkShaper.logger.debug "received: [#{received.join(', ')}]"
-        completed.shift
-        received.shift
+          @total_acked += 1
+          WorkShaper.logger.debug "@total_acked: #{@total_acked}"
+          WorkShaper.logger.debug "completed: [#{completed.join(', ')}]"
+          WorkShaper.logger.debug "received: [#{received.join(', ')}]"
+          completed.delete(offset)
+          received.delete(offset)
 
-        offset = completed.first
+          offset = completed.first
+        end
+      rescue => e
+        WorkShaper.logger.error({ message: 'Error in offset_ack', error: e })
+        WorkShaper.logger.error(e.backtrace.join("\n"))
       end
     end
 
@@ -179,11 +199,19 @@ module WorkShaper
         @total_enqueued.to_i - @total_acked.to_i > @max_in_queue
       end
 
+      pause_cycles = 0
       # We have to be careful here to avoid a deadlock. Another thread may be waiting
       # for the mutex to ack and remove offsets. If we wrap enqueue in a synchronize
       # block, that would lead to a deadlock. Here the sleep allows other threads
       # to wrap up.
-      sleep 0.005 while @semaphore.synchronize { overrun.call }
+      while @semaphore.synchronize { overrun.call } do
+        if pause_cycles % 12000 == 0
+          WorkShaper.logger.warn 'Paused on Overrun'
+          report(detailed: true)
+        end
+        pause_cycles += 1
+        sleep 0.005
+      end
     end
   end
 end

data/lib/work_shaper/offset_holder.rb

@@ -0,0 +1,54 @@
+module WorkShaper
+  class OffsetHolder
+    attr_reader :partition, :offset, :state
+
+    STATES = {enqueued: 2, acked: 1, completed: 0}
+    def initialize(partition, offset, at: Time.now.to_f)
+      @partition = partition
+      @offset = offset
+      @at = at
+
+      @state = :enqueued
+    end
+
+    def <=(other)
+      self.<=>(other) <= 0
+    end
+
+    def <(other)
+      self.<=>(other) == -1
+    end
+
+    def <=>(other)
+      r = offset <=> other.offset
+
+      if r == 0
+        r = STATES[state] <=> STATES[other.state]
+        puts "States: #{r} | #{STATES[state]} #{STATES[other.state]}"
+      end
+
+      if r == 0
+        r = @at <=> other.instance_variable_get(:@at)
+        puts "At: #{r}"
+      end
+      puts "Final: #{r}"
+      r
+    end
+
+    def ack!
+      @state = :acked
+    end
+
+    def complete!
+      @state = :completed
+    end
+
+    def to_i
+      offset
+    end
+
+    def to_s
+      "#{partition}/#{offset}:#{STATES[state]}"
+    end
+  end
+end

data/lib/work_shaper/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module WorkShaper
-  VERSION = "0.1.3.1rc2"
+  VERSION = "0.1.3.1rc4"
 end

data/lib/work_shaper/worker.rb

@@ -22,7 +22,10 @@ module WorkShaper
     # rubocop:enable Metrics/ParameterLists
     # rubocop:enable Layout/LineLength
 
-    def enqueue(message, partition, offset)
+    def enqueue(message, offset_holder)
+      partition = offset_holder.partition
+      offset = offset_holder.offset
+
       # rubocop:disable Style/RescueStandardError
       @thread_pool.post do
         @work.call(message, partition, offset)
@@ -34,7 +37,8 @@ module WorkShaper
       ensure
         @semaphore.synchronize do
           WorkShaper.logger.debug "Completed: #{partition}:#{offset}"
-          (@completed_offsets[partition] ||= Array.new) << offset
+          offset_holder.complete!
+          (@completed_offsets[partition] ||= Array.new) << offset_holder
         end
       end
       # rubocop:enable Style/RescueStandardError

data/lib/work_shaper.rb

@@ -3,6 +3,7 @@
 require 'logger'
 require 'concurrent-ruby'
 require_relative "work_shaper/version"
+require_relative "work_shaper/offset_holder"
 require_relative "work_shaper/manager"
 require_relative "work_shaper/worker"
 require 'json'

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: work_shaper
 version: !ruby/object:Gem::Version
-  version: 0.1.3.1rc2
+  version: 0.1.3.1rc4
 platform: ruby
 authors:
 - Jerry Fernholz
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-03-05 00:00:00.000000000 Z
+date: 2024-03-21 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: concurrent-ruby
@@ -40,6 +40,7 @@ files:
 - Rakefile
 - lib/work_shaper.rb
 - lib/work_shaper/manager.rb
+- lib/work_shaper/offset_holder.rb
 - lib/work_shaper/version.rb
 - lib/work_shaper/worker.rb
 - work_shaper.gemspec