deimos-kafka 1.0.0.pre.beta15

data/lib/deimos/consumer.rb

@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+require 'deimos/avro_data_decoder'
+require 'deimos/shared_config'
+require 'phobos/handler'
+require 'active_support/all'
+require 'ddtrace'
+
+# Class to consume messages coming from the pipeline topic
+# Note: According to the docs, instances of your handler will be created
+# for every incoming message. This class should be lightweight.
+module Deimos
+  # Parent consumer class.
+  class Consumer
+    include Phobos::Handler
+    include SharedConfig
+
+    class << self
+      # @return [AvroDataEncoder]
+      def decoder
+        @decoder ||= AvroDataDecoder.new(schema: config[:schema],
+                                         namespace: config[:namespace])
+      end
+
+      # @return [AvroDataEncoder]
+      def key_decoder
+        @key_decoder ||= AvroDataDecoder.new(schema: config[:key_schema],
+                                             namespace: config[:namespace])
+      end
+    end
+
+    # :nodoc:
+    def around_consume(payload, metadata)
+      _received_message(payload, metadata)
+      benchmark = Benchmark.measure do
+        _with_error_span(payload, metadata) { yield }
+      end
+      _handle_success(benchmark.real, payload, metadata)
+    end
+
+    # :nodoc:
+    def before_consume(payload, metadata)
+      _with_error_span(payload, metadata) do
+        if self.class.config[:key_schema] || self.class.config[:key_field]
+          metadata[:key] = decode_key(metadata[:key])
+        end
+        self.class.decoder.decode(payload) if payload.present?
+      end
+    end
+
+    # Helper method to decode an Avro-encoded key.
+    # @param key [String]
+    # @return [Object] the decoded key.
+    def decode_key(key)
+      return nil if key.nil?
+
+      config = self.class.config
+      if config[:encode_key] && config[:key_field].nil? &&
+         config[:key_schema].nil?
+        raise 'No key config given - if you are not decoding keys, please use `key_config plain: true`'
+      end
+
+      if config[:key_field]
+        self.class.decoder.decode_key(key, config[:key_field])
+      elsif config[:key_schema]
+        self.class.key_decoder.decode(key, schema: config[:key_schema])
+      else # no encoding
+        key
+      end
+    end
+
+    # Consume incoming messages.
+    # @param _payload [String]
+    # @param _metadata [Hash]
+    def consume(_payload, _metadata)
+      raise NotImplementedError
+    end
+
+  private
+
+    # @param payload [Hash|String]
+    # @param metadata [Hash]
+    def _with_error_span(payload, metadata)
+      @span = Deimos.config.tracer&.start(
+        'deimos-consumer',
+        resource: self.class.name.gsub('::', '-')
+      )
+      yield
+    rescue StandardError => e
+      _handle_error(e, payload, metadata)
+    ensure
+      Deimos.config.tracer&.finish(@span)
+    end
+
+    def _received_message(payload, metadata)
+      Deimos.config.logger.info(
+        message: 'Got Kafka event',
+        payload: payload,
+        metadata: metadata
+      )
+      Deimos.config.metrics&.increment('handler', tags: %W(
+                                         status:received
+                                         topic:#{metadata[:topic]}
+                                       ))
+    end
+
+    # @param exception [Throwable]
+    # @param payload [Hash]
+    # @param metadata [Hash]
+    def _handle_error(exception, payload, metadata)
+      Deimos.config.tracer&.set_error(@span, exception)
+      Deimos.config.metrics&.increment(
+        'handler',
+        tags: %W(
+          status:error
+          topic:#{metadata[:topic]}
+        )
+      )
+      Deimos.config.logger.warn(
+        message: 'Error consuming message',
+        handler: self.class.name,
+        metadata: metadata,
+        data: payload,
+        error_message: exception.message,
+        error: exception.backtrace
+      )
+      raise if Deimos.config.reraise_consumer_errors
+    end
+
+    # @param time_taken [Float]
+    # @param payload [Hash]
+    # @param metadata [Hash]
+    def _handle_success(time_taken, payload, metadata)
+      Deimos.config.metrics&.histogram('handler', time_taken, tags: %W(
+                                         time:consume
+                                         topic:#{metadata[:topic]}
+                                       ))
+      Deimos.config.metrics&.increment('handler', tags: %W(
+                                         status:success
+                                         topic:#{metadata[:topic]}
+                                       ))
+      Deimos.config.logger.info(
+        message: 'Finished processing Kafka event',
+        payload: payload,
+        time_elapsed: time_taken,
+        metadata: metadata
+      )
+      return if payload.nil? || payload['timestamp'].blank?
+
+      begin
+        time_delayed = Time.now.in_time_zone - payload['timestamp'].to_datetime
+      rescue ArgumentError
+        Deimos.config.logger.info(
+          message: "Error parsing timestamp! #{payload['timestamp']}"
+        )
+        return
+      end
+      Deimos.config.metrics&.histogram('handler', time_delayed, tags: %W(
+                                         time:time_delayed
+                                         topic:#{metadata[:topic]}
+                                       ))
+    end
+  end
+end

data/lib/deimos/instrumentation.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require 'active_support/notifications'
+require 'active_support/concern'
+
+# :nodoc:
+module Deimos
+  # Copied from Phobos instrumentation.
+  module Instrumentation
+    extend ActiveSupport::Concern
+    NAMESPACE = 'Deimos'
+
+    # :nodoc:
+    module ClassMethods
+      # :nodoc:
+      def subscribe(event)
+        ActiveSupport::Notifications.subscribe("#{NAMESPACE}.#{event}") do |*args|
+          yield(ActiveSupport::Notifications::Event.new(*args)) if block_given?
+        end
+      end
+
+      # :nodoc:
+      def unsubscribe(subscriber)
+        ActiveSupport::Notifications.unsubscribe(subscriber)
+      end
+
+      # :nodoc:
+      def instrument(event, extra={})
+        ActiveSupport::Notifications.instrument("#{NAMESPACE}.#{event}", extra) do |extra2|
+          yield(extra2) if block_given?
+        end
+      end
+    end
+  end
+
+  include Instrumentation
+
+  # This module listens to events published by RubyKafka.
+  module KafkaListener
+    # Listens for any exceptions that happen during publishing and re-publishes
+    # as a Deimos event.
+    # @param event [ActiveSupport::Notification]
+    def self.send_produce_error(event)
+      exception = event.payload[:exception_object]
+      return if !exception || !exception.respond_to?(:failed_messages)
+
+      messages = exception.failed_messages
+      messages.group_by(&:topic).each do |topic, batch|
+        next if batch.empty?
+
+        producer = batch.first.metadata[:producer_name]
+        payloads = batch.map { |m| m.metadata[:decoded_payload] }
+
+        Deimos.metrics&.count('publish_error', payloads.size,
+                              tags: %W(topic:#{topic}))
+        Deimos.instrument(
+          'produce_error',
+          producer: producer,
+          topic: topic,
+          exception_object: exception,
+          payloads: payloads
+        )
+      end
+    end
+  end
+
+  ActiveSupport::Notifications.subscribe('deliver_messages.producer.kafka') do |*args|
+    event = ActiveSupport::Notifications::Event.new(*args)
+    KafkaListener.send_produce_error(event)
+  end
+end

data/lib/deimos/kafka_message.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Deimos
+  # Store Kafka messages into the database.
+  class KafkaMessage < ActiveRecord::Base
+    self.table_name = 'kafka_messages'
+
+    validates_presence_of :message, :topic
+
+    # Ensure it gets turned into a string, e.g. for testing purposes. It
+    # should already be a string.
+    # @param mess [Object]
+    def message=(mess)
+      write_attribute(:message, mess.to_s)
+    end
+
+    # @return [Hash]
+    def phobos_message
+      {
+        payload: self.message,
+        partition_key: self.partition_key,
+        key: self.key,
+        topic: self.topic
+      }
+    end
+  end
+end

data/lib/deimos/kafka_source.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module Deimos
+  # Represents an object which needs to inform Kafka when it is saved or
+  # bulk imported.
+  module KafkaSource
+    extend ActiveSupport::Concern
+
+    included do
+      after_create(:send_kafka_event_on_create)
+      after_update(:send_kafka_event_on_update)
+      after_destroy(:send_kafka_event_on_destroy)
+    end
+
+    # Send the newly created model to Kafka.
+    def send_kafka_event_on_create
+      return unless self.persisted?
+      return unless self.class.kafka_config[:create]
+
+      self.class.kafka_producers.each { |p| p.send_event(self) }
+    end
+
+    # Send the newly updated model to Kafka.
+    def send_kafka_event_on_update
+      return unless self.class.kafka_config[:update]
+
+      producers = self.class.kafka_producers
+      fields = producers.flat_map(&:watched_attributes).uniq
+      fields -= ['updated_at']
+      # Only send an event if a field we care about was changed.
+      any_changes = fields.any? do |field|
+        field_change = self.previous_changes[field]
+        field_change.present? && field_change[0] != field_change[1]
+      end
+      return unless any_changes
+
+      producers.each { |p| p.send_event(self) }
+    end
+
+    # Send a deletion (null payload) event to Kafka.
+    def send_kafka_event_on_destroy
+      return unless self.class.kafka_config[:delete]
+
+      self.class.kafka_producers.each { |p| p.send_event(self.deletion_payload) }
+    end
+
+    # Payload to send after we are destroyed.
+    # @return [Hash]
+    def deletion_payload
+      { payload_key: self[self.class.primary_key] }
+    end
+
+    # :nodoc:
+    module ClassMethods
+      # @return [Hash]
+      def kafka_config
+        {
+          update: true,
+          delete: true,
+          import: true,
+          create: true
+        }
+      end
+
+      # @return [Array<Deimos::ActiveRecordProducer>] the producers to run.
+      def kafka_producers
+        raise NotImplementedError if self.method(:kafka_producer).
+          owner == Deimos::KafkaSource
+
+        [self.kafka_producer]
+      end
+
+      # Deprecated - use #kafka_producers instead.
+      # @return [Deimos::ActiveRecordProducer] the producer to use.
+      def kafka_producer
+        raise NotImplementedError if self.method(:kafka_producers).
+          owner == Deimos::KafkaSource
+
+        self.kafka_producers.first
+      end
+
+      # This is an internal method, part of the activerecord_import gem. It's
+      # the one that actually does the importing, having already normalized
+      # the inputs (arrays, hashes, records etc.)
+      # Basically we want to first do the import, then reload the records
+      # and send them to Kafka.
+      def import_without_validations_or_callbacks(column_names,
+                                                  array_of_attributes,
+                                                  options={})
+        results = super
+        return unless self.kafka_config[:import]
+        return if array_of_attributes.empty?
+
+        # This will contain an array of hashes, where each hash is the actual
+        # attribute hash that created the object.
+        ids =
+          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
+        end
+
+        self.kafka_producers.each { |p| p.send_events(array_of_hashes) }
+        results
+      end
+    end
+  end
+end

data/lib/deimos/kafka_topic_info.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Deimos
+  # Record that keeps track of which topics are being worked on by DbProducers.
+  class KafkaTopicInfo < ActiveRecord::Base
+    self.table_name = 'kafka_topic_info'
+
+    class << self
+      # Lock a topic for the given ID. Returns whether the lock was successful.
+      # @param topic [String]
+      # @param lock_id [String]
+      # @return [Boolean]
+      def lock(topic, lock_id)
+        # Try to create it - it's fine if it already exists
+        begin
+          self.create(topic: topic)
+        rescue ActiveRecord::RecordNotUnique # rubocop:disable Lint/HandleExceptions
+          # continue on
+        end
+
+        # Lock the record
+        qtopic = self.connection.quote(topic)
+        qlock_id = self.connection.quote(lock_id)
+        qtable = self.connection.quote_table_name('kafka_topic_info')
+        qnow = self.connection.quote(Time.zone.now.to_s(:db))
+        qfalse = self.connection.quoted_false
+        qtime = self.connection.quote(1.minute.ago.to_s(:db))
+
+        # If a record is marked as error and less than 1 minute old,
+        # we don't want to pick it up even if not currently locked because
+        # we worry we'll run into the same problem again.
+        # Once it's more than 1 minute old, we figure it's OK to try again
+        # so we can pick up any topic that's that old, even if it was
+        # locked by someone, because it's the job of the producer to keep
+        # updating the locked_at timestamp as they work on messages in that
+        # topic. If the locked_at timestamp is that old, chances are that
+        # the producer crashed.
+        sql = <<~SQL
+          UPDATE #{qtable}
+          SET locked_by=#{qlock_id}, locked_at=#{qnow}, error=#{qfalse}
+          WHERE topic=#{qtopic} AND
+           ((locked_by IS NULL AND error=#{qfalse}) OR locked_at < #{qtime})
+        SQL
+        self.connection.update(sql)
+        self.where(locked_by: lock_id, topic: topic).any?
+      end
+
+      # This is called once a producer is finished working on a topic, i.e.
+      # there are no more messages to fetch. It unlocks the topic and
+      # moves on to the next one.
+      # @param topic [String]
+      # @param lock_id [String]
+      def clear_lock(topic, lock_id)
+        self.where(topic: topic, locked_by: lock_id).
+          update_all(locked_by: nil, locked_at: nil, error: false, retries: 0)
+      end
+
+      # The producer calls this if it gets an error sending messages. This
+      # essentially locks down this topic for 1 minute (for all producers)
+      # and allows the caller to continue to the next topic.
+      # @param topic [String]
+      # @param lock_id [String]
+      def register_error(topic, lock_id)
+        record = self.where(topic: topic, locked_by: lock_id).last
+        record.update(locked_by: nil, locked_at: Time.zone.now, error: true,
+                      retries: record.retries + 1)
+      end
+
+      # Update the locked_at timestamp to indicate that the producer is still
+      # working on those messages and to continue.
+      # @param topic [String]
+      # @param lock_id [String]
+      def heartbeat(topic, lock_id)
+        self.where(topic: topic, locked_by: lock_id).
+          update_all(locked_at: Time.zone.now)
+      end
+    end
+  end
+end