deimos-ruby 1.7.0.pre.beta1 → 1.8.0.pre.beta1

data/lib/deimos/utils/db_producer.rb

@@ -87,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
@@ -133,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.7.0-beta1'
+  VERSION = '1.8.0-beta1'
 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