deimos-ruby 1.6.1 → 1.8.0.pre.beta1

data/lib/deimos/kafka_message.rb

@@ -42,7 +42,7 @@ module Deimos
       messages.map do |m|
         {
           key: m.key.present? ? decoder&.decode_key(m.key) || m.key : nil,
-          payload: decoder&.decoder&.decode(self.message) || self.message
+          payload: decoder&.decoder&.decode(m.message) || m.message
         }
       end
     end

data/lib/deimos/kafka_source.rb

@@ -93,31 +93,37 @@ module Deimos
 
         # This will contain an array of hashes, where each hash is the actual
         # attribute hash that created the object.
-        ids = if results.is_a?(Array)
-                results[1]
-              elsif results.respond_to?(:ids)
-                results.ids
-              else
-                []
-              end
-        if ids.blank?
-          # re-fill IDs based on what was just entered into the DB.
-          if self.connection.adapter_name.downcase =~ /sqlite/
-            last_id = self.connection.select_value('select last_insert_rowid()')
-            ids = ((last_id - array_of_attributes.size + 1)..last_id).to_a
-          else # mysql
-            last_id = self.connection.select_value('select LAST_INSERT_ID()')
-            ids = (last_id..(last_id + array_of_attributes.size)).to_a
-          end
-        end
         array_of_hashes = []
-        array_of_attributes.each_with_index do |array, i|
-          hash = column_names.zip(array).to_h.with_indifferent_access
-          hash[self.primary_key] = ids[i] if hash[self.primary_key].blank?
-          array_of_hashes << hash
+        array_of_attributes.each do |array|
+          array_of_hashes << column_names.zip(array).to_h.with_indifferent_access
+        end
+        hashes_with_id, hashes_without_id = array_of_hashes.partition { |arr| arr[:id].present? }
+
+        self.kafka_producers.each { |p| p.send_events(hashes_with_id) }
+
+        if hashes_without_id.any?
+          if options[:on_duplicate_key_update].present? &&
+             options[:on_duplicate_key_update] != [:updated_at]
+            unique_columns = column_names.map(&:to_s) -
+                             options[:on_duplicate_key_update].map(&:to_s) - %w(id created_at)
+            records = hashes_without_id.map do |hash|
+              self.where(unique_columns.map { |c| [c, hash[c]] }.to_h).first
+            end
+            self.kafka_producers.each { |p| p.send_events(records) }
+          else
+            # re-fill IDs based on what was just entered into the DB.
+            last_id = if self.connection.adapter_name.downcase =~ /sqlite/
+                        self.connection.select_value('select last_insert_rowid()') -
+                          hashes_without_id.size + 1
+                      else # mysql
+                        self.connection.select_value('select LAST_INSERT_ID()')
+                      end
+            hashes_without_id.each_with_index do |attrs, i|
+              attrs[:id] = last_id + i
+            end
+            self.kafka_producers.each { |p| p.send_events(hashes_without_id) }
+          end
         end
-
-        self.kafka_producers.each { |p| p.send_events(array_of_hashes) }
         results
       end
     end

data/lib/deimos/kafka_topic_info.rb

@@ -14,7 +14,7 @@ module Deimos
         # Try to create it - it's fine if it already exists
         begin
           self.create(topic: topic)
-        rescue ActiveRecord::RecordNotUnique # rubocop:disable Lint/SuppressedException
+        rescue ActiveRecord::RecordNotUnique
           # continue on
         end
 

data/lib/deimos/message.rb

@@ -10,7 +10,7 @@ module Deimos
     # @param producer [Class]
     def initialize(payload, producer, topic: nil, key: nil, partition_key: nil)
       @payload = payload&.with_indifferent_access
-      @producer_name = producer.name
+      @producer_name = producer&.name
       @topic = topic
       @key = key
       @partition_key = partition_key
@@ -70,5 +70,10 @@ module Deimos
     def ==(other)
       self.to_h == other.to_h
     end
+
+    # @return [Boolean] True if this message is a tombstone
+    def tombstone?
+      payload.nil?
+    end
   end
 end

data/lib/deimos/metrics/provider.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-# rubocop:disable Lint/UnusedMethodArgument
 module Deimos
   module Metrics
     # Base class for all metrics providers.
@@ -35,4 +34,3 @@ module Deimos
     end
   end
 end
-# rubocop:enable Lint/UnusedMethodArgument

data/lib/deimos/poll_info.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Deimos
+  # ActiveRecord class to record the last time we polled the database.
+  # For use with DbPoller.
+  class PollInfo < ActiveRecord::Base
+    self.table_name = 'deimos_poll_info'
+  end
+end

data/lib/deimos/tracing/provider.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-# rubocop:disable Lint/UnusedMethodArgument
 module Deimos
   module Tracing
     # Base class for all tracing providers.
@@ -28,4 +27,3 @@ module Deimos
     end
   end
 end
-# rubocop:enable Lint/UnusedMethodArgument

data/lib/deimos/utils/db_poller.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require 'deimos/poll_info'
+require 'sigurd/executor'
+require 'sigurd/signal_handler'
+
+module Deimos
+  module Utils
+    # Class which continually polls the database and sends Kafka messages.
+    class DbPoller
+      BATCH_SIZE = 1000
+
+      # Needed for Executor so it can identify the worker
+      attr_reader :id
+
+      # Begin the DB Poller process.
+      def self.start!
+        if Deimos.config.db_poller_objects.empty?
+          raise('No pollers configured!')
+        end
+
+        pollers = Deimos.config.db_poller_objects.map do |poller_config|
+          self.new(poller_config)
+        end
+        executor = Sigurd::Executor.new(pollers,
+                                        sleep_seconds: 5,
+                                        logger: Deimos.config.logger)
+        signal_handler = Sigurd::SignalHandler.new(executor)
+        signal_handler.run!
+      end
+
+      # @param config [Deimos::Configuration::ConfigStruct]
+      def initialize(config)
+        @config = config
+        @id = SecureRandom.hex
+        begin
+          @producer = @config.producer_class.constantize
+        rescue NameError
+          raise "Class #{@config.producer_class} not found!"
+        end
+        unless @producer < Deimos::ActiveRecordProducer
+          raise "Class #{@producer.class.name} is not an ActiveRecordProducer!"
+        end
+      end
+
+      # Start the poll:
+      # 1) Grab the current PollInfo from the database indicating the last
+      # time we ran
+      # 2) On a loop, process all the recent updates between the last time
+      # we ran and now.
+      def start
+        # Don't send asynchronously
+        if Deimos.config.producers.backend == :kafka_async
+          Deimos.config.producers.backend = :kafka
+        end
+        Deimos.config.logger.info('Starting...')
+        @signal_to_stop = false
+        retrieve_poll_info
+        loop do
+          if @signal_to_stop
+            Deimos.config.logger.info('Shutting down')
+            break
+          end
+          process_updates
+          sleep 0.1
+        end
+      end
+
+      # Grab the PollInfo or create if it doesn't exist.
+      def retrieve_poll_info
+        ActiveRecord::Base.connection.reconnect!
+        new_time = @config.start_from_beginning ? Time.new(0) : Time.zone.now
+        @info = Deimos::PollInfo.find_by_producer(@config.producer_class) ||
+                Deimos::PollInfo.create!(producer: @config.producer_class,
+                                         last_sent: new_time,
+                                         last_sent_id: 0)
+      end
+
+      # Stop the poll.
+      def stop
+        Deimos.config.logger.info('Received signal to stop')
+        @signal_to_stop = true
+      end
+
+      # Indicate whether this current loop should process updates. Most loops
+      # will busy-wait (sleeping 0.1 seconds) until it's ready.
+      # @return [Boolean]
+      def should_run?
+        Time.zone.now - @info.last_sent - @config.delay_time >= @config.run_every
+      end
+
+      # @param record [ActiveRecord::Base]
+      # @return [ActiveSupport::TimeWithZone]
+      def last_updated(record)
+        record.public_send(@config.timestamp_column)
+      end
+
+      # Send messages for updated data.
+      def process_updates
+        return unless should_run?
+
+        time_from = @config.full_table ? Time.new(0) : @info.last_sent.in_time_zone
+        time_to = Time.zone.now - @config.delay_time
+        Deimos.config.logger.info("Polling #{@producer.topic} from #{time_from} to #{time_to}")
+        message_count = 0
+        batch_count = 0
+
+        # poll_query gets all the relevant data from the database, as defined
+        # by the producer itself.
+        loop do
+          Deimos.config.logger.debug("Polling #{@producer.topic}, batch #{batch_count + 1}")
+          batch = fetch_results(time_from, time_to).to_a
+          break if batch.empty?
+
+          batch_count += 1
+          process_batch(batch)
+          message_count += batch.size
+          time_from = last_updated(batch.last)
+        end
+        Deimos.config.logger.info("Poll #{@producer.topic} complete at #{time_to} (#{message_count} messages, #{batch_count} batches}")
+      end
+
+      # @param time_from [ActiveSupport::TimeWithZone]
+      # @param time_to [ActiveSupport::TimeWithZone]
+      # @return [ActiveRecord::Relation]
+      def fetch_results(time_from, time_to)
+        id = @producer.config[:record_class].primary_key
+        quoted_timestamp = ActiveRecord::Base.connection.quote_column_name(@config.timestamp_column)
+        quoted_id = ActiveRecord::Base.connection.quote_column_name(id)
+        @producer.poll_query(time_from: time_from,
+                             time_to: time_to,
+                             column_name: @config.timestamp_column,
+                             min_id: @info.last_sent_id).
+          limit(BATCH_SIZE).
+          order("#{quoted_timestamp}, #{quoted_id}")
+      end
+
+      # @param batch [Array<ActiveRecord::Base>]
+      def process_batch(batch)
+        record = batch.last
+        id_method = record.class.primary_key
+        last_id = record.public_send(id_method)
+        last_updated_at = last_updated(record)
+        @producer.send_events(batch)
+        @info.update_attributes!(last_sent: last_updated_at, last_sent_id: last_id)
+      end
+    end
+  end
+end

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