work_shaper 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 5f008d410f5dc96fc854d8673b38873eacb66d176768d8c8ab8b78107a4b4837
+  data.tar.gz: 621935ffcad5cd2db1b138a2fafeab86eeb5c86138437756a56768cae698a1ff
+SHA512:
+  metadata.gz: 2119feac5a3f1d5e56d69b572a8b764558625d68da7c63d7b7fa4e2747e92a204013fea4247cd5502317fed446ebb7f8826b130f9f15437895de2547a4c927d2
+  data.tar.gz: 591f2a06d61f872e1d194b73dfac057f31b34f903512ba48b6864e9280d71b27ed83004e95274d4dd029ce978853f0dd81b126c5b5eb8e28014a3e9c169f89be

data/CHANGELOG.md

@@ -0,0 +1,5 @@
+## [Unreleased]
+
+## [0.1.0] - 2023-12-01
+
+- Initial release

data/Gemfile

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+source "https://rubygems.org"
+
+# Specify your gem's dependencies in work_shaper.gemspec
+gemspec
+
+gem "rake", "~> 13.0"
+
+gem "rspec", "~> 3.0"
+
+gem "rubocop", "~> 1.21"

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Broadvoice
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Jerry Fernholz
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,92 @@
+# WorkShaper
+
+WorkShaper is inspired by Kafka partitions and offsets, but could be used to organize and
+parallelize other forms of work. The original goal was to parallelize processing offsets in
+a given partition while maintaining order for a subset of the messages based on Sub Keys.
+
+The key concepts include Sub Key, Partition, and Offset. Work on a given Sub Key must be
+executed in the order in which it is enqueued. However, work on different Sub Keys can run
+in parallel. All Work (offsets) on a given Partition must be Acknowledged in continuous
+monotonically increasing order. If a higher offset's work is completed before a lower offset,
+the Manager will hold the acknowledgement until all lower offsets are acknowledged. Remember,
+work (offsets) for a given sub key are still processed in order.
+
+## Installation
+
+TODO: Replace `UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG` with your gem name right after releasing it to RubyGems.org. Please do not do it earlier due to security reasons. Alternatively, replace this section with instructions to install your gem from git if you don't plan to release to RubyGems.org.
+
+Install the gem and add to the application's Gemfile by executing:
+
+    $ bundle add UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG
+
+If bundler is not being used to manage dependencies, install the gem by executing:
+
+    $ gem install UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG
+
+## Usage
+
+### Example
+
+```ruby
+consumer = MyRdKafka.consumer(...)
+
+# Called for each message
+work = ->(message, _p, _o) do
+  MsgProcessor.process(message)
+end
+
+# Called each time `work` completes
+done = ->(_m, _p, _o) {}
+
+# Called periodically after work is complete to acknowledge the
+# completed work. Completed offsets are queued and processed every
+# 5 ms by the OffsetManager.
+ack = ->(p, o) do
+  consumer.store_offset(ENV.fetch('TOPIC_NAME'), p, o)
+rescue InvalidOffset => e
+  # On rebalance, RdKafka sets the offset to _INVALID for the consumer
+  # losing that offset. In this scenario InvalidOffset is expected
+  # and we should move on.
+  # TODO: This can probably be more elegantly handled.
+end
+
+# Call if an exception in encountered in `done`. It is important to
+# understand `work` is being called in a sub thread, so the exception
+# will not bubble up.
+error = ->(e, m, p, o) do
+  logger.error "#{e} on #{p} #{o}"
+  @fatal_error = e
+end
+
+max_in_queue = ENV.fetch('MAX_THREAD_QUEUE_SIZE', 25)
+
+work_shaper = WorkShaper::Manager.new(work, done, ack, error, max_in_queue)
+
+@value_to_subkey = {}
+max_sub_keys = ENV.fetch('MAX_SUB_KEYS', 100)
+consumer.each_message do |message|
+  break if @fatal_error
+  
+  sub_key = @value_to_subkey[message.payload['some attribute']] ||=
+    MurmurHash3::V32.str_hash(message.payload['some attribute']) % max_sub_keys
+
+  work_shaper.enqueue(
+    sub_key,
+    message,
+    message.partition,
+    message.offset
+  )
+end
+```
+
+
+## Development
+
+
+## Contributing
+
+Bug reports and pull requests are welcome on GitHub at https://github.com/broadvoice/work-shaper.
+
+## License
+
+The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).

data/Rakefile

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "rspec/core/rake_task"
+
+RSpec::Core::RakeTask.new(:spec)
+
+require "rubocop/rake_task"
+
+RuboCop::RakeTask.new
+
+task default: %i[spec rubocop]

data/lib/work_shaper/manager.rb

@@ -0,0 +1,182 @@
+module WorkShaper
+  # The Manager is responsible for organizing the work to be done, triggering calls to acknowledge work done
+  # for each offset in monotonically increasing order (independent of the execution order), and gracefully
+  # cleaning up when `#shutdown` is called.
+  class Manager
+    # 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'}`.
+    #
+    # @param work [#call(message, partition, offset)] Lambda that we will call to execute work.
+    # @param on_done [#call(message, partition, offset)] Lambda that we call when work is done.
+    # @param ack [] Lambda we will call when it is safe to commit an offset. This is not the
+    #   same as Done.
+    # @param on_error [#call(exception, message, partition, offset)] Lambda that we call if an
+    #   error is encountered.
+    # @param max_in_queue [Integer] The maximum in flight jobs per Sub Key. This affects how many
+    #   message could get replayed if your process crashes before the offsets are committed.
+    def initialize(work:, on_done:, ack:, on_error:, max_in_queue: 3,
+                   heartbeat_period_sec: 60, offset_commit_period_ms: 5)
+      @work = work
+      @on_done = on_done
+      @ack = ack
+      @on_error = on_error
+      @workers = {}
+      @last_ack = {}
+      @received_offsets = {}
+      @completed_offsets = {}
+      @max_in_queue = max_in_queue
+      @semaphore = Mutex.new
+      @shutdown = false
+
+      @total_enqueued = 0
+
+      @heartbeat = Thread.new do
+        while true
+          report
+          sleep heartbeat_period_sec
+        end
+      rescue => e
+        WorkShaper.logger.warn({ message: 'Shutdown from Heartbeat', error: e })
+        shutdown
+      end
+
+      @offset_manager = Thread.new do
+        while true
+          @completed_offsets.each_key do |partition|
+            offset_ack(partition)
+          end
+          sleep offset_commit_period_ms / 1000.0
+        end
+      rescue => e
+        WorkShaper.logger.warn({ message: 'Shutdown from Offset Manager', error: e })
+        shutdown
+      end
+    end
+
+    # 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
+      pause_on_overrun
+
+      worker = nil
+      @semaphore.synchronize do
+        @total_enqueued += 1
+        (@received_offsets[partition] ||= SortedSet.new) << offset
+
+        worker =
+          @workers[sub_key] ||=
+            Worker.new(
+              @work,
+              @on_done,
+              method(:offset_ack),
+              @on_error,
+              @last_ack,
+              @completed_offsets,
+              @semaphore,
+              @max_in_queue
+            )
+      end
+
+      worker.enqueue(message, partition, offset)
+    end
+
+    # Flush any offsets for which work has been completed. Only lowest continuous run of
+    # offsets will be acknowledged. Any offset after a discontinuity will be replayed when
+    # the consumer restarts.
+    def flush(safe: true)
+      sleep 5
+      @completed_offsets.each_key do |k|
+        safe ? offset_ack(k) : offset_ack_unsafe(k)
+      end
+    end
+
+    # Output state of Last Acked and Pending Offset Ack's.
+    def report(detailed: false)
+      @semaphore.synchronize do
+        WorkShaper.logger.info(
+          { message: 'Reporting', total_enqueued: @total_enqueued,
+            total_acked: @total_acked,
+            in_flight: (@total_enqueued.to_i - @total_acked.to_i),
+            last_acked_offsets: @last_ack,
+            worker_count: @workers.keys.count
+          })
+        if detailed
+          WorkShaper.logger.info(
+            {
+              messaage: 'Reporting - Extra Detail',
+              pending_ack: @completed_offsets,
+              received_offsets: @received_offsets
+            })
+        end
+      end
+    end
+
+    # Stop the underlying threads
+    def shutdown
+      @shutdown = true
+      report(detailed: true)
+      Thread.kill(@heartbeat)
+      Thread.kill(@offset_manager)
+      @workers.each_value(&:shutdown)
+    end
+
+    private
+
+    def offset_ack(partition)
+      @semaphore.synchronize do
+        offset_ack_unsafe(partition)
+      end
+    end
+
+    def offset_ack_unsafe(partition)
+      @total_acked ||= 0
+
+      completed = @completed_offsets[partition]
+      received = @received_offsets[partition]
+
+      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 Dupilcate Offset',
+              offset: "#{partition}:#{offset}"
+            })
+        else
+          result = @ack.call(partition, offset)
+          if result.is_a? Exception
+            WorkShaper.logger.warn(
+              { message: 'Failed to Ack Offset, likely re-balance',
+                offset: "#{partition}:#{offset}",
+                completed: @completed_offsets[partition].to_a[0..10].join(','),
+                received: @received_offsets[partition].to_a[0..10].join(',')
+              })
+          else
+            @total_acked += 1
+            @last_ack[partition] = offset
+          end
+        end
+
+        completed.delete(offset)
+        received.delete(offset)
+
+        offset = completed.first
+      end
+    end
+
+    def pause_on_overrun
+      overrun = lambda do
+        @total_enqueued.to_i - @total_acked.to_i > @max_in_queue
+      end
+
+      # 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 }
+    end
+  end
+end

data/lib/work_shaper/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module WorkShaper
+  VERSION = "0.1.0"
+end

data/lib/work_shaper/worker.rb

@@ -0,0 +1,54 @@
+module WorkShaper
+  # The worker that runs the stuff
+  class Worker
+    # rubocop:disable Metrics/ParameterLists
+    # rubocop:disable Layout/LineLength
+    # @param work [Lambda] Lambda that we will #call(message) to execute work.
+    # @param on_done [Lambda] Lambda that we #call(partition, offset) when work is done.
+    # @param on_error [Lambda] Lambda that we #call(exception) if an error is encountered.
+    def initialize(work, on_done, ack_handler, on_error, last_ack, offset_stack, semaphore, max_in_queue)
+      @jobs = []
+      @work = work
+      @on_done = on_done
+      @ack_handler = ack_handler
+      @on_error = on_error
+      @last_ack = last_ack
+      @completed_offsets = offset_stack
+      @semaphore = semaphore
+      @max_in_queue = max_in_queue
+      @thread_pool = Concurrent::FixedThreadPool.new(1, auto_terminate: false)
+    end
+
+    # rubocop:enable Metrics/ParameterLists
+    # rubocop:enable Layout/LineLength
+
+    def enqueue(message, partition, offset)
+      # rubocop:disable Style/RescueStandardError
+      @thread_pool.post do
+        @work.call(message, partition, offset)
+        @on_done.call(message, partition, offset)
+        @semaphore.synchronize do
+          (@completed_offsets[partition] ||= SortedSet.new) << offset
+        end
+        # @ack_handler.call(partition, offset)
+      rescue => e
+        puts("Error processing #{partition}:#{offset} #{e}")
+        puts(e.backtrace.join("\n"))
+        # logger.error("Acking it anyways, why not?")
+        @on_error.call(e, message, partition, offset)
+        # @ack_handler.call(partition, offset)
+      end
+      # rubocop:enable Style/RescueStandardError
+    end
+
+    def shutdown
+      # Cannot call logger from trap{}
+      WorkShaper.logger.info({message: 'Shutting down worker'})
+      @thread_pool.shutdown
+      @thread_pool.wait_for_termination
+      sleep 0.05 while @thread_pool.queue_length.positive?
+    end
+
+    private
+  end
+end

data/lib/work_shaper.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require_relative "work_shaper/version"
+require_relative "work_shaper/manager"
+require_relative "work_shaper/worker"
+
+# WorkShaper is inspired by Kafka partitions and offsets, but could be used to organize and
+# parallelize other forms of work. The original goal was to parallelize processing offsets in
+# a given partition while maintaining order for a subset of the messages based on Sub Keys.
+#
+# The key concepts include Sub Key, Partition, and Offset. Work on a given Sub Key must be
+# executed in the order in which it is enqueued. However, work on different Sub Keys can run
+# in parallel. All Work (offset) on a given Partition must be Acknowledged in continuous
+# monotonically increasing order. If a higher offset's work is completed before a lower offset,
+# the Manager will hold the acknowledgement until all lower offsets are acknowledged. Remember,
+# work (offsets) for a given sub key are still processed in order.
+module WorkShaper
+  def self.logger=(logger)
+    @logger = logger
+  end
+  def self.logger
+    @logger ||= Logger.new(
+      $stdout,
+      level: ENV['LOG_LEVEL'] || 'DEBUG',
+      formatter: Ruby::JSONFormatter::Base.new
+    )
+  end
+end
+

data/work_shaper.gemspec

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require_relative "lib/work_shaper/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "work_shaper"
+  spec.version = WorkShaper::VERSION
+  spec.authors = ["Jerry Fernholz"]
+  spec.email = ["jerryf@broadvoice.com"]
+
+  spec.summary = "Parallelize work across many threads."
+  spec.description = "WorkShaper was built to parallelize the work needed to process Kafka messages."
+  spec.homepage = "https://github.com/broadvoice/work-shaper"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 2.6.0"
+
+  spec.metadata["allowed_push_host"] = "https://rubygems.org"
+
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = "https://github.com/broadvoice/work-shaper"
+  spec.metadata["changelog_uri"] = "https://github.com/broadvoice/work-shaper/blob/main/CHANGELOG.md"
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files = Dir.chdir(__dir__) do
+    `git ls-files -z`.split("\x0").reject do |f|
+      (File.expand_path(f) == __FILE__) || f.start_with?(*%w[bin/ test/ spec/ features/ .git .circleci appveyor])
+    end
+  end
+  spec.bindir = "exe"
+  spec.executables = spec.files.grep(%r{\Aexe/}) { |f| File.basename(f) }
+  spec.require_paths = ["lib"]
+
+  # Uncomment to register a new dependency of your gem
+  spec.add_dependency "sorted_set", "~> 1.0"
+
+  # For more information and examples about making a new gem, check out our
+  # guide at: https://bundler.io/guides/creating_gem.html
+end

metadata

@@ -0,0 +1,73 @@
+--- !ruby/object:Gem::Specification
+name: work_shaper
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Jerry Fernholz
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2024-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: sorted_set
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+description: WorkShaper was built to parallelize the work needed to process Kafka
+  messages.
+email:
+- jerryf@broadvoice.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- Gemfile
+- LICENSE
+- LICENSE.txt
+- README.md
+- Rakefile
+- lib/work_shaper.rb
+- lib/work_shaper/manager.rb
+- lib/work_shaper/version.rb
+- lib/work_shaper/worker.rb
+- work_shaper.gemspec
+homepage: https://github.com/broadvoice/work-shaper
+licenses:
+- MIT
+metadata:
+  allowed_push_host: https://rubygems.org
+  homepage_uri: https://github.com/broadvoice/work-shaper
+  source_code_uri: https://github.com/broadvoice/work-shaper
+  changelog_uri: https://github.com/broadvoice/work-shaper/blob/main/CHANGELOG.md
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 2.6.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.4.10
+signing_key:
+specification_version: 4
+summary: Parallelize work across many threads.
+test_files: []