deimos-ruby 1.6.2 → 1.8.0.pre.beta2

data/lib/deimos/utils/db_producer.rb

@@ -2,7 +2,8 @@
 
 module Deimos
   module Utils
-    # Class which continually polls the database and sends Kafka messages.
+    # Class which continually polls the kafka_messages table
+    # in the database and sends Kafka messages.
     class DbProducer
       include Phobos::Producer
       attr_accessor :id, :current_topic
@@ -86,9 +87,9 @@ module Deimos
           begin
             produce_messages(compacted_messages.map(&:phobos_message))
           rescue Kafka::BufferOverflow, Kafka::MessageSizeTooLarge, Kafka::RecordListTooLarge
-            Deimos::KafkaMessage.where(id: messages.map(&:id)).delete_all
             @logger.error('Message batch too large, deleting...')
             @logger.error(Deimos::KafkaMessage.decoded(messages))
+            Deimos::KafkaMessage.where(id: messages.map(&:id)).delete_all
             raise
           end
         end
@@ -132,7 +133,11 @@ module Deimos
           metrics.gauge('pending_db_messages_max_wait', 0)
         end
         messages.each do |record|
-          time_diff = Time.zone.now - record.earliest
+          earliest = record.earliest
+          # SQLite gives a string here
+          earliest = Time.zone.parse(earliest) if earliest.is_a?(String)
+
+          time_diff = Time.zone.now - earliest
           metrics.gauge('pending_db_messages_max_wait', time_diff,
                         tags: ["topic:#{record.topic}"])
         end

data/lib/deimos/utils/deadlock_retry.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Utility class to retry a given block if a a deadlock is encountered.
+    # Supports Postgres and MySQL deadlocks and lock wait timeouts.
+    class DeadlockRetry
+      class << self
+        # Maximum number of times to retry the block after encountering a deadlock
+        RETRY_COUNT = 2
+
+        # Need to match on error messages to support older Rails versions
+        DEADLOCK_MESSAGES = [
+          # MySQL
+          'Deadlock found when trying to get lock',
+          'Lock wait timeout exceeded',
+
+          # Postgres
+          'deadlock detected'
+        ].freeze
+
+        # Retry the given block when encountering a deadlock. For any other
+        # exceptions, they are reraised. This is used to handle cases where
+        # the database may be busy but the transaction would succeed if
+        # retried later. Note that your block should be idempotent and it will
+        # be wrapped in a transaction.
+        # Sleeps for a random number of seconds to prevent multiple transactions
+        # from retrying at the same time.
+        # @param tags [Array] Tags to attach when logging and reporting metrics.
+        # @yield Yields to the block that may deadlock.
+        def wrap(tags=[])
+          count = RETRY_COUNT
+
+          begin
+            ActiveRecord::Base.transaction do
+              yield
+            end
+          rescue ActiveRecord::StatementInvalid => e
+            # Reraise if not a known deadlock
+            raise if DEADLOCK_MESSAGES.none? { |m| e.message.include?(m) }
+
+            # Reraise if all retries exhausted
+            raise if count <= 0
+
+            Deimos.config.logger.warn(
+              message: 'Deadlock encountered when trying to execute query. '\
+                "Retrying. #{count} attempt(s) remaining",
+              tags: tags
+            )
+
+            Deimos.config.metrics&.increment(
+              'deadlock',
+              tags: tags
+            )
+
+            count -= 1
+
+            # Sleep for a random amount so that if there are multiple
+            # transactions deadlocking, they don't all retry at the same time
+            sleep(Random.rand(5.0) + 0.5)
+
+            retry
+          end
+        end
+      end
+    end
+  end
+end

data/lib/deimos/utils/lag_reporter.rb

@@ -29,30 +29,21 @@ module Deimos
           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)
+        def assign_current_offset(topic, partition, offset)
           self.topics[topic.to_s] ||= Topic.new(topic, self)
-          self.topics[topic.to_s].compute_lag(partition, offset)
+          self.topics[topic.to_s].assign_current_offset(partition, offset)
         end
       end
 
-      # Topic which has a hash of partition => last known offset lag
+      # Topic which has a hash of partition => last known current offsets
       class Topic
         # @return [String]
         attr_accessor :topic_name
         # @return [Hash<Integer, Integer>]
-        attr_accessor :partition_offset_lags
+        attr_accessor :partition_current_offsets
         # @return [ConsumerGroup]
         attr_accessor :consumer_group
 
@@ -61,35 +52,33 @@ module Deimos
         def initialize(topic_name, group)
           self.topic_name = topic_name
           self.consumer_group = group
-          self.partition_offset_lags = {}
+          self.partition_current_offsets = {}
         end
 
         # @param partition [Integer]
-        # @param lag [Integer]
-        def assign_lag(partition, lag)
-          self.partition_offset_lags[partition.to_i] = lag
+        def assign_current_offset(partition, offset)
+          self.partition_current_offsets[partition.to_i] = offset
         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)
+            lag = last_offset - offset
           rescue StandardError # don't do anything, just wait
             Deimos.config.logger.
               debug("Error computing lag for #{self.topic_name}, will retry")
           end
+          lag || 0
         end
 
         # @param partition [Integer]
         def report_lag(partition)
-          lag = self.partition_offset_lags[partition.to_i]
-          return unless lag
+          current_offset = self.partition_current_offsets[partition.to_i]
+          return unless current_offset
 
+          lag = compute_lag(partition, current_offset)
           group = self.consumer_group.id
           Deimos.config.logger.
             debug("Sending lag: #{group}/#{partition}: #{lag}")
@@ -109,16 +98,20 @@ module Deimos
           @groups = {}
         end
 
+        # offset_lag = event.payload.fetch(:offset_lag)
+        # group_id = event.payload.fetch(:group_id)
+        # topic = event.payload.fetch(:topic)
+        # partition = event.payload.fetch(:partition)
         # @param payload [Hash]
         def message_processed(payload)
-          lag = payload[:offset_lag]
+          offset = payload[:offset] || payload[:last_offset]
           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)
+            @groups[group.to_s].assign_current_offset(topic, partition, offset)
           end
         end
 
@@ -131,7 +124,7 @@ module Deimos
 
           synchronize do
             @groups[group.to_s] ||= ConsumerGroup.new(group)
-            @groups[group.to_s].compute_lag(topic, partition, offset)
+            @groups[group.to_s].assign_current_offset(topic, partition, offset)
           end
         end
 

data/lib/deimos/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Deimos
-  VERSION = '1.6.2'
+  VERSION = '1.8.0-beta2'
 end

data/lib/generators/deimos/db_poller/templates/migration

@@ -0,0 +1,11 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :deimos_poll_info, force: true do |t|
+      t.string :producer, null: false
+      t.datetime :last_sent
+      t.bigint :last_sent_id
+    end
+
+    add_index :deimos_poll_info, [:producer]
+  end
+end

data/lib/generators/deimos/db_poller/templates/rails3_migration

@@ -0,0 +1,16 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def self.up
+    create_table :deimos_poll_info, force: true do |t|
+      t.string :producer, null: false
+      t.datetime :last_sent
+      t.bigint :last_sent_id
+    end
+
+    add_index :deimos_poll_info, [:producer]
+  end
+
+  def self.down
+    drop_table :deimos_poll_info
+  end
+
+end

data/lib/generators/deimos/db_poller_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 DbPollerGenerator < 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_poller/templates', __dir__)
+      desc 'Add migrations for the database poller'
+
+      # @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_poller.rb")
+        else
+          migration_template('migration',
+                             "#{db_migrate_path}/create_db_poller.rb")
+        end
+      end
+    end
+  end
+end

data/lib/tasks/deimos.rake

@@ -24,4 +24,11 @@ namespace :deimos do
     Deimos.start_db_backend!(thread_count: thread_count)
   end
 
+  task db_poller: :environment do
+    ENV['DEIMOS_RAKE_TASK'] = 'true'
+    STDOUT.sync = true
+    Rails.logger.info('Running deimos:db_poller rake task.')
+    Deimos::Utils::DbPoller.start!
+  end
+
 end

data/spec/active_record_batch_consumer_spec.rb

@@ -0,0 +1,481 @@
+# frozen_string_literal: true
+
+# Wrapped in a module to prevent class leakage
+module ActiveRecordBatchConsumerTest
+  describe Deimos::ActiveRecordConsumer do
+    # Create ActiveRecord table and model
+    before(:all) do
+      ActiveRecord::Base.connection.create_table(:widgets, force: true) do |t|
+        t.string(:test_id)
+        t.string(:part_one)
+        t.string(:part_two)
+        t.integer(:some_int)
+        t.boolean(:deleted, default: false)
+        t.timestamps
+
+        t.index(%i(part_one part_two), unique: true)
+      end
+
+      # Sample model
+      class Widget < ActiveRecord::Base
+        validates :test_id, presence: true
+
+        default_scope -> { where(deleted: false) }
+      end
+
+      Widget.reset_column_information
+    end
+
+    after(:all) do
+      ActiveRecord::Base.connection.drop_table(:widgets)
+    end
+
+    prepend_before(:each) do
+      stub_const('MyBatchConsumer', consumer_class)
+    end
+
+    around(:each) do |ex|
+      # Set and freeze example time
+      travel_to start do
+        ex.run
+      end
+    end
+
+    # Default starting time
+    let(:start) { Time.zone.local(2019, 1, 1, 10, 30, 0) }
+
+    # Basic uncompacted consumer
+    let(:consumer_class) do
+      Class.new(described_class) do
+        schema 'MySchema'
+        namespace 'com.my-namespace'
+        key_config plain: true
+        record_class Widget
+        compacted false
+      end
+    end
+
+    # Helper to get all instances, ignoring default scopes
+    def all_widgets
+      Widget.unscoped.all
+    end
+
+    # Helper to publish a list of messages and call the consumer
+    def publish_batch(messages)
+      keys = messages.map { |m| m[:key] }
+      payloads = messages.map { |m| m[:payload] }
+
+      test_consume_batch(MyBatchConsumer, payloads, keys: keys, call_original: true)
+    end
+
+    it 'should handle an empty batch' do
+      expect { publish_batch([]) }.not_to raise_error
+    end
+
+    it 'should create records from a batch' do
+      publish_batch(
+        [
+          { key: 1,
+            payload: { test_id: 'abc', some_int: 3 } },
+          { key: 2,
+            payload: { test_id: 'def', some_int: 4 } }
+        ]
+      )
+
+      expect(all_widgets).
+        to match_array(
+          [
+            have_attributes(id: 1, test_id: 'abc', some_int: 3, updated_at: start, created_at: start),
+            have_attributes(id: 2, test_id: 'def', some_int: 4, updated_at: start, created_at: start)
+          ]
+        )
+    end
+
+    it 'should handle deleting a record that doesn\'t exist' do
+      publish_batch(
+        [
+          { key: 1,
+            payload: nil }
+        ]
+      )
+
+      expect(all_widgets).to be_empty
+    end
+
+    it 'should handle an update, followed by a delete in the correct order' do
+      Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+      publish_batch(
+        [
+          { key: 1,
+            payload: { test_id: 'abc', some_int: 3 } },
+          { key: 1,
+            payload: nil }
+        ]
+      )
+
+      expect(all_widgets).to be_empty
+    end
+
+    it 'should handle a delete, followed by an update in the correct order' do
+      Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+      travel 1.day
+
+      publish_batch(
+        [
+          { key: 1,
+            payload: nil },
+          { key: 1,
+            payload: { test_id: 'abc', some_int: 3 } }
+        ]
+      )
+
+      expect(all_widgets).
+        to match_array(
+          [
+            have_attributes(id: 1, test_id: 'abc', some_int: 3, updated_at: Time.zone.now, created_at: Time.zone.now)
+          ]
+        )
+    end
+
+    it 'should handle a double update' do
+      Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+      travel 1.day
+
+      publish_batch(
+        [
+          { key: 1,
+            payload: { test_id: 'def', some_int: 3 } },
+          { key: 1,
+            payload: { test_id: 'ghi', some_int: 4 } }
+        ]
+      )
+
+      expect(all_widgets).
+        to match_array(
+          [
+            have_attributes(id: 1, test_id: 'ghi', some_int: 4, updated_at: Time.zone.now, created_at: start)
+          ]
+        )
+    end
+
+    it 'should handle a double deletion' do
+      Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+      publish_batch(
+        [
+          { key: 1,
+            payload: nil },
+          { key: 1,
+            payload: nil }
+        ]
+      )
+
+      expect(all_widgets).to be_empty
+    end
+
+    it 'should ignore default scopes' do
+      Widget.create!(id: 1, test_id: 'abc', some_int: 2, deleted: true)
+      Widget.create!(id: 2, test_id: 'def', some_int: 3, deleted: true)
+
+      publish_batch(
+        [
+          { key: 1,
+            payload: nil },
+          { key: 2,
+            payload: { test_id: 'def', some_int: 5 } }
+        ]
+      )
+
+      expect(all_widgets).
+        to match_array(
+          [
+            have_attributes(id: 2, test_id: 'def', some_int: 5)
+          ]
+        )
+    end
+
+    describe 'compacted mode' do
+      # Create a compacted consumer
+      let(:consumer_class) do
+        Class.new(described_class) do
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config plain: true
+          record_class Widget
+
+          # :no-doc:
+          def deleted_query(_records)
+            raise 'Should not have anything to delete!'
+          end
+        end
+      end
+
+      it 'should allow for compacted batches' do
+        expect(Widget).to receive(:import!).once.and_call_original
+
+        publish_batch(
+          [
+            { key: 1,
+              payload: nil },
+            { key: 2,
+              payload: { test_id: 'xyz', some_int: 5 } },
+            { key: 1,
+              payload: { test_id: 'abc', some_int: 3 } },
+            { key: 2,
+              payload: { test_id: 'def', some_int: 4 } },
+            { key: 3,
+              payload: { test_id: 'hij', some_int: 9 } }
+          ]
+        )
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(id: 1, test_id: 'abc', some_int: 3),
+              have_attributes(id: 2, test_id: 'def', some_int: 4),
+              have_attributes(id: 3, test_id: 'hij', some_int: 9)
+            ]
+          )
+      end
+    end
+
+    describe 'batch atomicity' do
+      it 'should roll back if there was an exception while deleting' do
+        Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+        travel 1.day
+
+        expect(Widget.connection).to receive(:delete).and_raise('Some error')
+
+        expect {
+          publish_batch(
+            [
+              { key: 1,
+                payload: { test_id: 'def', some_int: 3 } },
+              { key: 1,
+                payload: nil }
+            ]
+          )
+        }.to raise_error('Some error')
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(id: 1, test_id: 'abc', some_int: 2, updated_at: start, created_at: start)
+            ]
+          )
+      end
+
+      it 'should roll back if there was an invalid instance while upserting' do
+        Widget.create!(id: 1, test_id: 'abc', some_int: 2) # Updated but rolled back
+        Widget.create!(id: 3, test_id: 'ghi', some_int: 3) # Removed but rolled back
+
+        travel 1.day
+
+        expect {
+          publish_batch(
+            [
+              { key: 1,
+                payload: { test_id: 'def', some_int: 3 } },
+              { key: 2,
+                payload: nil },
+              { key: 2,
+                payload: { test_id: '', some_int: 4 } }, # Empty string is not valid for test_id
+              { key: 3,
+                payload: nil }
+            ]
+          )
+        }.to raise_error(ActiveRecord::RecordInvalid)
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(id: 1, test_id: 'abc', some_int: 2, updated_at: start, created_at: start),
+              have_attributes(id: 3, test_id: 'ghi', some_int: 3, updated_at: start, created_at: start)
+            ]
+          )
+      end
+    end
+
+    describe 'compound keys' do
+      let(:consumer_class) do
+        Class.new(described_class) do
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config schema: 'MySchemaCompound-key'
+          record_class Widget
+          compacted false
+
+          # :no-doc:
+          def deleted_query(records)
+            keys = records.
+              map { |m| record_key(m.key) }.
+              reject(&:empty?)
+
+            # Only supported on Rails 5+
+            keys.reduce(@klass.none) do |query, key|
+              query.or(@klass.unscoped.where(key))
+            end
+          end
+        end
+      end
+
+      it 'should consume with compound keys' do
+        Widget.create!(test_id: 'xxx', some_int: 2, part_one: 'ghi', part_two: 'jkl')
+        Widget.create!(test_id: 'yyy', some_int: 7, part_one: 'mno', part_two: 'pqr')
+
+        publish_batch(
+          [
+            { key: { part_one: 'abc', part_two: 'def' }, # To be created
+              payload: { test_id: 'aaa', some_int: 3 } },
+            { key: { part_one: 'ghi', part_two: 'jkl' }, # To be updated
+              payload: { test_id: 'bbb', some_int: 4 } },
+            { key: { part_one: 'mno', part_two: 'pqr' }, # To be deleted
+              payload: nil }
+          ]
+        )
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(test_id: 'aaa', some_int: 3, part_one: 'abc', part_two: 'def'),
+              have_attributes(test_id: 'bbb', some_int: 4, part_one: 'ghi', part_two: 'jkl')
+            ]
+          )
+      end
+    end
+
+    describe 'no keys' do
+      let(:consumer_class) do
+        Class.new(described_class) do
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config none: true
+          record_class Widget
+        end
+      end
+
+      it 'should handle unkeyed topics' do
+        Widget.create!(test_id: 'xxx', some_int: 2)
+
+        publish_batch(
+          [
+            { payload: { test_id: 'aaa', some_int: 3 } },
+            { payload: { test_id: 'bbb', some_int: 4 } },
+            { payload: nil } # Should be ignored. Can't delete with no key
+          ]
+        )
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(test_id: 'xxx', some_int: 2),
+              have_attributes(test_id: 'aaa', some_int: 3),
+              have_attributes(test_id: 'bbb', some_int: 4)
+            ]
+          )
+      end
+    end
+
+    describe 'soft deletion' do
+      let(:consumer_class) do
+        Class.new(described_class) do
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config plain: true
+          record_class Widget
+          compacted false
+
+          # Sample customization: Soft delete
+          def remove_records(records)
+            deleted = deleted_query(records)
+
+            deleted.update_all(
+              deleted: true,
+              updated_at: Time.zone.now
+            )
+          end
+
+          # Sample customization: Undelete records
+          def record_attributes(payload, key)
+            super.merge(deleted: false)
+          end
+        end
+      end
+
+      it 'should mark records deleted' do
+        Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+        Widget.create!(id: 3, test_id: 'xyz', some_int: 4)
+        Widget.create!(id: 4, test_id: 'uvw', some_int: 5, deleted: true)
+
+        travel 1.day
+
+        publish_batch(
+          [
+            { key: 1,
+              payload: nil },
+            { key: 1, # Double delete for key 1
+              payload: nil },
+            { key: 2, # Create 2
+              payload: { test_id: 'def', some_int: 3 } },
+            { key: 2, # Delete 2
+              payload: nil },
+            { key: 3, # Update non-deleted
+              payload: { test_id: 'ghi', some_int: 4 } },
+            { key: 4, # Revive
+              payload: { test_id: 'uvw', some_int: 5 } }
+          ]
+        )
+
+        expect(all_widgets).
+          to match_array(
+            [
+              have_attributes(id: 1, test_id: 'abc', some_int: 2, deleted: true,
+                              created_at: start, updated_at: Time.zone.now),
+              have_attributes(id: 2, test_id: 'def', some_int: 3, deleted: true,
+                              created_at: Time.zone.now, updated_at: Time.zone.now),
+              have_attributes(id: 3, test_id: 'ghi', some_int: 4, deleted: false,
+                              created_at: start, updated_at: Time.zone.now),
+              have_attributes(id: 4, test_id: 'uvw', some_int: 5, deleted: false,
+                              created_at: start, updated_at: Time.zone.now)
+            ]
+          )
+      end
+    end
+
+    describe 'skipping records' do
+      let(:consumer_class) do
+        Class.new(described_class) do
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config plain: true
+          record_class Widget
+
+          # Sample customization: Skipping records
+          def record_attributes(payload, key)
+            return nil if payload[:test_id] == 'skipme'
+
+            super
+          end
+        end
+      end
+
+      it 'should allow overriding to skip any unwanted records' do
+        publish_batch(
+          [
+            { key: 1, # Record that consumer can decide to skip
+              payload: { test_id: 'skipme' } },
+            { key: 2,
+              payload: { test_id: 'abc123' } }
+          ]
+        )
+
+        expect(all_widgets).
+          to match_array([have_attributes(id: 2, test_id: 'abc123')])
+      end
+    end
+  end
+end