work_shaper 0.1.3.1rc1 → 0.1.3.1rc3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ffaaf1badd8ced4836dd5e1ae9f3b60f6f811e249fa617d5278b86cf667d3ba3
-  data.tar.gz: 065b69e44362a8c1bea9aedcbfaa70cd7b36ae6ee310cd780f33905b9036903b
+  metadata.gz: 19828916c7e0f7d9006c2b508566311381edcb8e6a50522b4bffe1ddc79e5a61
+  data.tar.gz: 28d670fd0d3fad8e6705b2392ca00c0c720ac29649fddc04702cdd924910fdc5
 SHA512:
-  metadata.gz: 22e06e944e8f6fa7b062ac8842abad05414d653af9d7d8d68050469a40aa0b2e6bfa1c9dcc490f55e561da56afadb5c336093d7301af62fac8a68d299372bd06
-  data.tar.gz: 8c2187099bac137a99ed48e53a57c5235ea0b63227fc5238d780ec5e3adc668c57c9419020103e22cf0ad786cd10f8efb4e360c035008e18373aa9dae43c1ebb
+  metadata.gz: ca66783166d22a048cdfbf02e1a9a6b36246c60caa42e6612f3008bb8af8029621df2d85ce152315288f8c40ffb1ca28efe4024e3582af7cd5c77957e8a0722e
+  data.tar.gz: 73cdaa9db8bcf924983fc0a6bbfec46bd257fa7c9938a7ed795c77ebbcdb34a1f95e762949876c57618276bf0f00dec9f38daacd2ecb708ea701e3322eb6cc8c

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,14 +28,14 @@ module WorkShaper
       @completed_offsets = {}
       @max_in_queue = max_in_queue
       @semaphore = Mutex.new
-      @shutdown = false
+      @shutting_down = false
 
       @total_enqueued = 0
       @total_acked = 0
 
       @heartbeat = Thread.new do
         while true
-          report(detailed: false)
+          report(detailed: true)
           sleep heartbeat_period_sec
         end
       rescue => e
@@ -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)
+            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.1rc1"
+  VERSION = "0.1.3.1rc3"
 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.1rc1
+  version: 0.1.3.1rc3
 platform: ruby
 authors:
 - Jerry Fernholz
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-03-05 00:00:00.000000000 Z
+date: 2024-03-14 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