deimos-kafka 1.0.0.pre.beta15

data/lib/deimos/utils/lag_reporter.rb

@@ -0,0 +1,182 @@
+# frozen_string_literal: true
+
+require 'mutex_m'
+
+# :nodoc:
+module Deimos
+  module Utils
+    # Class that manages reporting lag.
+    class LagReporter
+      extend Mutex_m
+
+      # Class that has a list of topics
+      class ConsumerGroup
+        # @return [Hash<String, Topic>]
+        attr_accessor :topics
+        # @return [String]
+        attr_accessor :id
+
+        # @param id [String]
+        def initialize(id)
+          self.id = id
+          self.topics = {}
+        end
+
+        # @param topic [String]
+        # @param partition [Integer]
+        def report_lag(topic, partition)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          self.topics[topic.to_s].report_lag(partition)
+        end
+
+        # @param topic [String]
+        # @param partition [Integer]
+        # @param lag [Integer]
+        def assign_lag(topic, partition, lag)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          self.topics[topic.to_s].assign_lag(partition, lag)
+        end
+
+        # Figure out the current lag by asking Kafka based on the current offset.
+        # @param topic [String]
+        # @param partition [Integer]
+        # @param offset [Integer]
+        def compute_lag(topic, partition, offset)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          self.topics[topic.to_s].compute_lag(partition, offset)
+        end
+      end
+
+      # Topic which has a hash of partition => last known offset lag
+      class Topic
+        # @return [String]
+        attr_accessor :topic_name
+        # @return [Hash<Integer, Integer>]
+        attr_accessor :partition_offset_lags
+        # @return [ConsumerGroup]
+        attr_accessor :consumer_group
+
+        # @param topic_name [String]
+        # @param group [ConsumerGroup]
+        def initialize(topic_name, group)
+          self.topic_name = topic_name
+          self.consumer_group = group
+          self.partition_offset_lags = {}
+        end
+
+        # @param partition [Integer]
+        # @param lag [Integer]
+        def assign_lag(partition, lag)
+          self.partition_offset_lags[partition.to_i] = lag
+        end
+
+        # @param partition [Integer]
+        # @param offset [Integer]
+        def compute_lag(partition, offset)
+          return if self.partition_offset_lags[partition.to_i]
+
+          begin
+            client = Phobos.create_kafka_client
+            last_offset = client.last_offset_for(self.topic_name, partition)
+            assign_lag(partition, [last_offset - offset, 0].max)
+          rescue StandardError # don't do anything, just wait
+            Deimos.config.logger.
+              debug("Error computing lag for #{self.topic_name}, will retry")
+          end
+        end
+
+        # @param partition [Integer]
+        def report_lag(partition)
+          lag = self.partition_offset_lags[partition.to_i]
+          return unless lag
+
+          group = self.consumer_group.id
+          Deimos.config.logger.
+            debug("Sending lag: #{group}/#{partition}: #{lag}")
+          Deimos.config.metrics&.gauge('consumer_lag', lag, tags: %W(
+                                         consumer_group:#{group}
+                                         partition:#{partition}
+                                         topic:#{self.topic_name}
+                                       ))
+        end
+      end
+
+      @groups = {}
+
+      class << self
+        # Reset all group information.
+        def reset
+          @groups = {}
+        end
+
+        # @param payload [Hash]
+        def message_processed(payload)
+          lag = payload[:offset_lag]
+          topic = payload[:topic]
+          group = payload[:group_id]
+          partition = payload[:partition]
+
+          synchronize do
+            @groups[group.to_s] ||= ConsumerGroup.new(group)
+            @groups[group.to_s].assign_lag(topic, partition, lag)
+          end
+        end
+
+        # @param payload [Hash]
+        def offset_seek(payload)
+          offset = payload[:offset]
+          topic = payload[:topic]
+          group = payload[:group_id]
+          partition = payload[:partition]
+
+          synchronize do
+            @groups[group.to_s] ||= ConsumerGroup.new(group)
+            @groups[group.to_s].compute_lag(topic, partition, offset)
+          end
+        end
+
+        # @param payload [Hash]
+        def heartbeat(payload)
+          group = payload[:group_id]
+          synchronize do
+            @groups[group.to_s] ||= ConsumerGroup.new(group)
+            consumer_group = @groups[group.to_s]
+            payload[:topic_partitions].each do |topic, partitions|
+              partitions.each do |partition|
+                consumer_group.report_lag(topic, partition)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+
+  ActiveSupport::Notifications.subscribe('start_process_message.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.message_processed(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('start_process_batch.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.message_processed(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('seek.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.offset_seek(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('heartbeat.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.heartbeat(event.payload)
+  end
+end

data/lib/deimos/utils/signal_handler.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Mostly copied free-form from Phobos::Cli::Runner. We should add a PR to
+    # basically replace that implementation with this one to make it more generic.
+    class SignalHandler
+      SIGNALS = %i(INT TERM QUIT).freeze
+
+      # Takes any object that responds to the `start` and `stop` methods.
+      # @param runner[#start, #stop]
+      def initialize(runner)
+        @signal_queue = []
+        @reader, @writer = IO.pipe
+        @runner = runner
+      end
+
+      # Run the runner.
+      def run!
+        setup_signals
+        @runner.start
+
+        loop do
+          case signal_queue.pop
+          when *SIGNALS
+            @runner.stop
+            break
+          else
+            ready = IO.select([reader, writer])
+
+            # drain the self-pipe so it won't be returned again next time
+            reader.read_nonblock(1) if ready[0].include?(reader)
+          end
+        end
+      end
+
+    private
+
+      attr_reader :reader, :writer, :signal_queue, :executor
+
+      # https://stackoverflow.com/questions/29568298/run-code-when-signal-is-sent-but-do-not-trap-the-signal-in-ruby
+      def prepend_handler(signal)
+        previous = Signal.trap(signal) do
+          previous = -> { raise SignalException, signal } unless previous.respond_to?(:call)
+          yield
+          previous.call
+        end
+      end
+
+      # Trap signals using the self-pipe trick.
+      def setup_signals
+        at_exit { @runner&.stop }
+        SIGNALS.each do |signal|
+          prepend_handler(signal) do
+            unblock(signal)
+          end
+        end
+      end
+
+      # Save the signal to the queue and continue on.
+      # @param signal [Symbol]
+      def unblock(signal)
+        writer.write_nonblock('.')
+        signal_queue << signal
+      end
+    end
+  end
+end

data/lib/deimos/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Deimos
+  VERSION = '1.0.0-beta15'
+end

data/lib/generators/deimos/db_backend/templates/migration

@@ -0,0 +1,24 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :kafka_messages, force: true do |t|
+      t.string :topic, null: false
+      t.binary :message, limit: 10.megabytes
+      t.binary :key
+      t.string :partition_key
+      t.timestamps
+    end
+
+    add_index :kafka_messages, [:topic, :id]
+
+    create_table :kafka_topic_info, force: true do |t|
+      t.string :topic, null: false
+      t.string :locked_by
+      t.datetime :locked_at
+      t.boolean :error, null: false, default: false
+      t.integer :retries, null: false, default: 0
+    end
+    add_index :kafka_topic_info, :topic, unique: true
+    add_index :kafka_topic_info, [:locked_by, :error]
+    add_index :kafka_topic_info, :locked_at
+  end
+end

data/lib/generators/deimos/db_backend/templates/rails3_migration

@@ -0,0 +1,30 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def self.up
+    create_table :kafka_messages, force: true do |t|
+      t.string :topic, null: false
+      t.binary :message, limit: 10.megabytes
+      t.binary :key
+      t.string :partition_key
+      t.timestamps
+    end
+
+    add_index :kafka_messages, [:topic, :id]
+
+    create_table :kafka_topic_info, force: true do |t|
+      t.string :topic, null: false
+      t.string :locked_by
+      t.datetime :locked_at
+      t.boolean :error, null: false, default: false
+      t.integer :retries, null: false, default: 0
+    end
+    add_index :kafka_topic_info, :topic, unique: true
+    add_index :kafka_topic_info, [:locked_by, :error]
+    add_index :kafka_topic_info, :locked_at
+  end
+
+  def self.down
+    drop_table :kafka_messages
+    drop_table :kafka_topic_info
+  end
+
+end

data/lib/generators/deimos/db_backend_generator.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'rails/generators'
+require 'rails/generators/active_record/migration'
+
+module Deimos
+  module Generators
+    # Generate the database backend migration.
+    class DbBackendGenerator < Rails::Generators::Base
+      include Rails::Generators::Migration
+      if Rails.version < '4'
+        extend(ActiveRecord::Generators::Migration)
+      else
+        include ActiveRecord::Generators::Migration
+      end
+      source_root File.expand_path('db_backend/templates', __dir__)
+      desc 'Add migrations for the database backend'
+
+      # @return [String]
+      def migration_version
+        "[#{ActiveRecord::Migration.current_version}]"
+      rescue StandardError
+        ''
+      end
+
+      # @return [String]
+      def db_migrate_path
+        if defined?(Rails.application) && Rails.application
+          paths = Rails.application.config.paths['db/migrate']
+          paths.respond_to?(:to_ary) ? paths.to_ary.first : paths.to_a.first
+        else
+          'db/migrate'
+        end
+      end
+
+      # Main method to create all the necessary files
+      def generate
+        if Rails.version < '4'
+          migration_template('rails3_migration',
+                             "#{db_migrate_path}/create_db_backend.rb")
+        else
+          migration_template('migration',
+                             "#{db_migrate_path}/create_db_backend.rb")
+        end
+      end
+    end
+  end
+end

data/lib/tasks/deimos.rake

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+require 'phobos'
+require 'phobos/cli'
+
+namespace :deimos do
+  desc 'Starts Deimos in the rails environment'
+  task start: :environment do
+    Deimos.configure do |config|
+      config.publish_backend = :kafka_sync if config.publish_backend == :kafka_async
+    end
+    ENV['DEIMOS_RAKE_TASK'] = 'true'
+    STDOUT.sync = true
+    Rails.logger.info('Running deimos:start rake task.')
+    Phobos::CLI::Commands.start(%w(start --skip_config))
+  end
+end

data/spec/active_record_consumer_spec.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require 'date'
+
+# :nodoc:
+module ActiveRecordProducerTest
+  describe Deimos::ActiveRecordConsumer do
+
+    before(:all) do
+      ActiveRecord::Base.connection.create_table(:widgets, force: true) do |t|
+        t.string(:test_id)
+        t.integer(:some_int)
+        t.boolean(:some_bool)
+        t.datetime(:some_datetime_int)
+        t.timestamps
+      end
+
+      # :nodoc:
+      class Widget < ActiveRecord::Base
+      end
+      Widget.reset_column_information
+    end
+
+    after(:all) do
+      ActiveRecord::Base.connection.drop_table(:widgets)
+    end
+
+    prepend_before(:each) do
+
+      consumer_class = Class.new(Deimos::ActiveRecordConsumer) do
+        schema 'MySchemaWithDateTimes'
+        namespace 'com.my-namespace'
+        key_config plain: true
+        record_class Widget
+      end
+      stub_const('MyConsumer', consumer_class)
+
+      Time.zone = 'Eastern Time (US & Canada)'
+    end
+
+    it 'should receive events correctly' do
+      travel 1.day do
+        expect(Widget.count).to eq(0)
+        test_consume_message(MyConsumer, {
+                               test_id: 'abc',
+                               some_int: 3,
+                               updated_at: 1.day.ago.to_i,
+                               some_datetime_int: Time.zone.now.to_i,
+                               timestamp: 2.minutes.ago.to_s
+                             }, { call_original: true, key: 5 })
+
+        expect(Widget.count).to eq(1)
+        widget = Widget.last
+        expect(widget.id).to eq(5)
+        expect(widget.test_id).to eq('abc')
+        expect(widget.some_int).to eq(3)
+        expect(widget.some_datetime_int).to eq(Time.zone.now)
+        expect(widget.updated_at).to eq(Time.zone.now)
+
+        # test update
+        test_consume_message(MyConsumer, {
+                               test_id: 'abcd',
+                               some_int: 3,
+                               some_datetime_int: Time.zone.now.to_i,
+                               timestamp: 2.minutes.ago.to_s
+                             }, { call_original: true, key: 5 })
+        widget = Widget.last
+        expect(widget.id).to eq(5)
+        expect(widget.test_id).to eq('abcd')
+        expect(widget.some_int).to eq(3)
+
+        # test delete
+        test_consume_message(MyConsumer, nil, call_original: true, key: 5)
+        expect(Widget.count).to eq(0)
+
+      end
+
+    end
+
+  end
+end