deimos-temp-fork 0.0.1

data/lib/deimos/instrumentation.rb

@@ -0,0 +1,76 @@
+# 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|
+        producer = Deimos::Producer.descendants.find { |c| c.topic == topic }
+        next if batch.empty? || !producer
+
+        decoder = Deimos.schema_backend(schema: producer.config[:schema],
+                                        namespace: producer.config[:namespace])
+        payloads = batch.map { |m| decoder.decode(m.value) }
+
+        Deimos.config.metrics&.increment(
+          'publish_error',
+          tags: %W(topic:#{topic}),
+          by: payloads.size
+        )
+        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,60 @@
+# frozen_string_literal: true
+
+module Deimos
+  # Store Kafka messages into the database.
+  class KafkaMessage < ActiveRecord::Base
+    self.table_name = 'kafka_messages'
+
+    validates_presence_of :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 ? mess.to_s : nil)
+    end
+
+    # Decoded payload for this message.
+    # @return [Hash]
+    def decoded_message
+      self.class.decoded([self]).first
+    end
+
+    # Get a decoder to decode a set of messages on the given topic.
+    # @param topic [String]
+    # @return [Deimos::Consumer]
+    def self.decoder(topic)
+      producer = Deimos::Producer.descendants.find { |c| c.topic == topic }
+      return nil unless producer
+
+      consumer = Class.new(Deimos::Consumer)
+      consumer.config.merge!(producer.config)
+      consumer
+    end
+
+    # Decoded payloads for a list of messages.
+    # @param messages [Array<Deimos::KafkaMessage>]
+    # @return [Array<Hash>]
+    def self.decoded(messages=[])
+      return [] if messages.empty?
+
+      decoder = self.decoder(messages.first.topic)&.new
+      messages.map do |m|
+        {
+          key: m.key.present? ? decoder&.decode_key(m.key) || m.key : nil,
+          payload: decoder&.decoder&.decode(m.message) || m.message
+        }
+      end
+    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,128 @@
+# 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
+
+    DEPRECATION_WARNING = 'The kafka_producer interface will be deprecated ' \
+        'in future releases. Please use kafka_producers instead.'
+
+    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.publish_list([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
+        if self.respond_to?(:kafka_producer)
+          Deimos.config.logger.warn(message: DEPRECATION_WARNING)
+          return [self.kafka_producer]
+        end
+
+        raise NotImplementedError
+      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
+        if !self.kafka_config[:import] || array_of_attributes.empty?
+          return results
+        end
+
+        # This will contain an array of hashes, where each hash is the actual
+        # attribute hash that created the object.
+        array_of_hashes = []
+        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
+        results
+      end
+    end
+  end
+end

data/lib/deimos/kafka_topic_info.rb

@@ -0,0 +1,102 @@
+# 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, last_processed_at: Time.zone.now)
+        rescue ActiveRecord::RecordNotUnique
+          # 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,
+                     last_processed_at: Time.zone.now)
+      end
+
+      # Update all topics that aren't currently locked and have no messages
+      # waiting. It's OK if some messages get inserted in the middle of this
+      # because the point is that at least within a few milliseconds of each
+      # other, it wasn't locked and had no messages, meaning the topic
+      # was in a good state.
+      # @param except_topics [Array<String>] the list of topics we've just
+      # realized had messages in them, meaning all other topics were empty.
+      def ping_empty_topics(except_topics)
+        records = KafkaTopicInfo.where(locked_by: nil).
+          where('topic not in(?)', except_topics)
+        records.each do |info|
+          info.update_attribute(:last_processed_at, Time.zone.now)
+        end
+      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
+        attr_hash = { locked_by: nil,
+                      locked_at: Time.zone.now,
+                      error: true,
+                      retries: record.retries + 1 }
+        record.attributes = attr_hash
+        record.save!
+      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

data/lib/deimos/message.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Deimos
+  # Basically a struct to hold the message as it's processed.
+  class Message
+    attr_accessor :payload, :key, :partition_key, :encoded_key,
+                  :encoded_payload, :topic, :producer_name
+
+    # @param payload [Hash]
+    # @param producer [Class]
+    def initialize(payload, producer, topic: nil, key: nil, partition_key: nil)
+      @payload = payload&.with_indifferent_access
+      @producer_name = producer&.name
+      @topic = topic
+      @key = key
+      @partition_key = partition_key
+    end
+
+    # Add message_id and timestamp default values if they are in the
+    # schema and don't already have values.
+    # @param fields [Array<String>] existing name fields in the schema.
+    def add_fields(fields)
+      return if @payload.except(:payload_key, :partition_key).blank?
+
+      if fields.include?('message_id')
+        @payload['message_id'] ||= SecureRandom.uuid
+      end
+      if fields.include?('timestamp')
+        @payload['timestamp'] ||= Time.now.in_time_zone.to_s
+      end
+    end
+
+    # @param encoder [Deimos::SchemaBackends::Base]
+    def coerce_fields(encoder)
+      return if payload.nil?
+
+      @payload = encoder.coerce(@payload)
+    end
+
+    # @return [Hash]
+    def encoded_hash
+      {
+        topic: @topic,
+        key: @encoded_key,
+        partition_key: @partition_key || @encoded_key,
+        payload: @encoded_payload,
+        metadata: {
+          decoded_payload: @payload,
+          producer_name: @producer_name
+        }
+      }
+    end
+
+    # @return [Hash]
+    def to_h
+      {
+        topic: @topic,
+        key: @key,
+        partition_key: @partition_key || @key,
+        payload: @payload,
+        metadata: {
+          decoded_payload: @payload,
+          producer_name: @producer_name
+        }
+      }
+    end
+
+    # @param other [Message]
+    # @return [Boolean]
+    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