deimos-ruby 1.0.0.pre.beta22

data/lib/deimos/utils/executor.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+# rubocop:disable Lint/RescueException
+module Deimos
+  module Utils
+    # Mostly copied from Phobos::Executor. We should DRY this up by putting in a
+    # PR to make it more generic. Might even make sense to move to a separate
+    # gem.
+    class Executor
+      # @return [Array<#start, #stop, #id>]
+      attr_accessor :runners
+
+      # @param runners [Array<#start, #stop, #id>] A list of objects that can be
+      # started or stopped.
+      # @param logger [Logger]
+      def initialize(runners, logger=Logger.new(STDOUT))
+        @threads = Concurrent::Array.new
+        @runners = runners
+        @logger = logger
+      end
+
+      # Start the executor.
+      def start
+        @logger.info('Starting executor')
+        @signal_to_stop = false
+        @threads.clear
+        @thread_pool = Concurrent::FixedThreadPool.new(@runners.size)
+
+        @runners.each do |runner|
+          @thread_pool.post do
+            thread = Thread.current
+            thread.abort_on_exception = true
+            @threads << thread
+            run_object(runner)
+          end
+        end
+
+        true
+      end
+
+      # Stop the executor.
+      def stop
+        return if @signal_to_stop
+
+        @logger.info('Stopping executor')
+        @signal_to_stop = true
+        @runners.each(&:stop)
+        @threads.select(&:alive?).each do |thread|
+          begin
+            thread.wakeup
+          rescue StandardError
+            nil
+          end
+        end
+        @thread_pool&.shutdown
+        @thread_pool&.wait_for_termination
+        @logger.info('Executor stopped')
+      end
+
+    private
+
+      # @param exception [Throwable]
+      # @return [Hash]
+      def error_metadata(exception)
+        {
+          exception_class: exception.class.name,
+          exception_message: exception.message,
+          backtrace: exception.backtrace
+        }
+      end
+
+      def run_object(runner)
+        retry_count = 0
+
+        begin
+          @logger.info("Running #{runner.id}")
+          runner.start
+          retry_count = 0 # success - reset retry count
+        rescue Exception => e
+          handle_crashed_runner(runner, e, retry_count)
+          retry_count += 1
+          retry unless @signal_to_stop
+        end
+      rescue Exception => e
+        @logger.error("Failed to run listener (#{e.message}) #{error_metadata(e)}")
+        raise e
+      end
+
+      # @return [ExponentialBackoff]
+      def create_exponential_backoff
+        min = 1
+        max = 60
+        ExponentialBackoff.new(min, max).tap do |backoff|
+          backoff.randomize_factor = rand
+        end
+      end
+
+      # When "runner#start" is interrupted / crashes we assume it's
+      # safe to be called again
+      def handle_crashed_runner(runner, error, retry_count)
+        backoff = create_exponential_backoff
+        interval = backoff.interval_at(retry_count).round(2)
+
+        metadata = {
+          listener_id: runner.id,
+          retry_count: retry_count,
+          waiting_time: interval
+        }.merge(error_metadata(error))
+
+        @logger.error("Runner crashed, waiting #{interval}s (#{error.message}) #{metadata}")
+        sleep(interval)
+      end
+    end
+  end
+end
+
+# rubocop:enable Lint/RescueException

data/lib/deimos/utils/inline_consumer.rb

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+# Class to consume messages. Can be used with integration testing frameworks.
+# Assumes that you have a topic with only one partition.
+module Deimos
+  module Utils
+    # Listener that can seek to get the last X messages in a topic.
+    class SeekListener < Phobos::Listener
+      attr_accessor :num_messages
+
+      # :nodoc:
+      def start_listener
+        @num_messages ||= 10
+        @consumer = create_kafka_consumer
+        @consumer.subscribe(topic, @subscribe_opts)
+
+        begin
+          last_offset = @kafka_client.last_offset_for(topic, 0)
+          offset = last_offset - num_messages
+          if offset.positive?
+            Deimos.config.logger.info("Seeking to #{offset}")
+            @consumer.seek(topic, 0, offset)
+          end
+        rescue StandardError => e
+          "Could not seek to offset: #{e.message}"
+        end
+
+        instrument('listener.start_handler', listener_metadata) do
+          @handler_class.start(@kafka_client)
+        end
+        log_info('Listener started', listener_metadata)
+      end
+    end
+
+    # Class to return the messages consumed.
+    class MessageBankHandler < Deimos::Consumer
+      include Phobos::Handler
+
+      cattr_accessor :total_messages
+
+      # @param klass [Class < Deimos::Consumer]
+      def self.config_class=(klass)
+        self.config.merge!(klass.config)
+      end
+
+      # :nodoc:
+      def self.start(_kafka_client)
+        self.total_messages = []
+      end
+
+      # :nodoc:
+      def consume(payload, metadata)
+        puts "Got #{payload}"
+        self.class.total_messages << {
+          key: metadata[:key],
+          payload: payload
+        }
+      end
+    end
+
+    # Class which can process/consume messages inline.
+    class InlineConsumer
+      MAX_MESSAGE_WAIT_TIME = 1.second
+      MAX_TOPIC_WAIT_TIME = 10.seconds
+
+      # Get the last X messages from a topic. You can specify a subclass of
+      # Deimos::Consumer or Deimos::Producer, or provide the
+      # schema, namespace and key_config directly.
+      # @param topic [String]
+      # @param config_class [Class < Deimos::Consumer|Deimos::Producer>]
+      # @param schema [String]
+      # @param namespace [String]
+      # @param key_config [Hash]
+      # @param num_messages [Number]
+      # @return [Array<Hash>]
+      def self.get_messages_for(topic:, schema: nil, namespace: nil, key_config: nil,
+                                config_class: nil, num_messages: 10)
+        if config_class
+          MessageBankHandler.config_class = config_class
+        elsif schema.nil? || key_config.nil?
+          raise 'You must specify either a config_class or a schema, namespace and key_config!'
+        else
+          MessageBankHandler.class_eval do
+            schema schema
+            namespace namespace
+            key_config key_config
+            @decoder = nil
+            @key_decoder = nil
+          end
+        end
+        self.consume(topic: topic,
+                     frk_consumer: MessageBankHandler,
+                     num_messages: num_messages)
+        messages = MessageBankHandler.total_messages
+        messages.size <= num_messages ? messages : messages[-num_messages..-1]
+      end
+
+      # Consume the last X messages from a topic.
+      # @param topic [String]
+      # @param frk_consumer [Class]
+      # @param num_messages [Integer] If this number is >= the number
+      #   of messages in the topic, all messages will be consumed.
+      def self.consume(topic:, frk_consumer:, num_messages: 10)
+        listener = SeekListener.new(
+          handler: frk_consumer,
+          group_id: SecureRandom.hex,
+          topic: topic,
+          heartbeat_interval: 1
+        )
+        listener.num_messages = num_messages
+
+        # Add the start_time and last_message_time attributes to the
+        # consumer class so we can kill it if it's gone on too long
+        class << frk_consumer
+          attr_accessor :start_time, :last_message_time
+        end
+
+        subscribers = []
+        subscribers << ActiveSupport::Notifications.
+          subscribe('phobos.listener.process_message') do
+            frk_consumer.last_message_time = Time.zone.now
+          end
+        subscribers << ActiveSupport::Notifications.
+          subscribe('phobos.listener.start_handler') do
+            frk_consumer.start_time = Time.zone.now
+            frk_consumer.last_message_time = nil
+          end
+        subscribers << ActiveSupport::Notifications.
+          subscribe('heartbeat.consumer.kafka') do
+            if frk_consumer.last_message_time
+              if Time.zone.now - frk_consumer.last_message_time > MAX_MESSAGE_WAIT_TIME
+                raise Phobos::AbortError
+              end
+            elsif Time.zone.now - frk_consumer.start_time > MAX_TOPIC_WAIT_TIME
+              Deimos.config.logger.error('Aborting - initial wait too long')
+              raise Phobos::AbortError
+            end
+          end
+        listener.start
+        subscribers.each { |s| ActiveSupport::Notifications.unsubscribe(s) }
+      end
+    end
+  end
+end

data/lib/deimos/utils/lag_reporter.rb

@@ -0,0 +1,182 @@
+# frozen_string_literal: true
+
+require 'mutex_m'
+
+# :nodoc:
+module Deimos
+  module Utils
+    # Class that manages reporting lag.
+    class LagReporter
+      extend Mutex_m
+
+      # Class that has a list of topics
+      class ConsumerGroup
+        # @return [Hash<String, Topic>]
+        attr_accessor :topics
+        # @return [String]
+        attr_accessor :id
+
+        # @param id [String]
+        def initialize(id)
+          self.id = id
+          self.topics = {}
+        end
+
+        # @param topic [String]
+        # @param partition [Integer]
+        def report_lag(topic, partition)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          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)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          self.topics[topic.to_s].compute_lag(partition, offset)
+        end
+      end
+
+      # Topic which has a hash of partition => last known offset lag
+      class Topic
+        # @return [String]
+        attr_accessor :topic_name
+        # @return [Hash<Integer, Integer>]
+        attr_accessor :partition_offset_lags
+        # @return [ConsumerGroup]
+        attr_accessor :consumer_group
+
+        # @param topic_name [String]
+        # @param group [ConsumerGroup]
+        def initialize(topic_name, group)
+          self.topic_name = topic_name
+          self.consumer_group = group
+          self.partition_offset_lags = {}
+        end
+
+        # @param partition [Integer]
+        # @param lag [Integer]
+        def assign_lag(partition, lag)
+          self.partition_offset_lags[partition.to_i] = lag
+        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)
+          rescue StandardError # don't do anything, just wait
+            Deimos.config.logger.
+              debug("Error computing lag for #{self.topic_name}, will retry")
+          end
+        end
+
+        # @param partition [Integer]
+        def report_lag(partition)
+          lag = self.partition_offset_lags[partition.to_i]
+          return unless lag
+
+          group = self.consumer_group.id
+          Deimos.config.logger.
+            debug("Sending lag: #{group}/#{partition}: #{lag}")
+          Deimos.config.metrics&.gauge('consumer_lag', lag, tags: %W(
+                                         consumer_group:#{group}
+                                         partition:#{partition}
+                                         topic:#{self.topic_name}
+                                       ))
+        end
+      end
+
+      @groups = {}
+
+      class << self
+        # Reset all group information.
+        def reset
+          @groups = {}
+        end
+
+        # @param payload [Hash]
+        def message_processed(payload)
+          lag = payload[:offset_lag]
+          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)
+          end
+        end
+
+        # @param payload [Hash]
+        def offset_seek(payload)
+          offset = payload[:offset]
+          topic = payload[:topic]
+          group = payload[:group_id]
+          partition = payload[:partition]
+
+          synchronize do
+            @groups[group.to_s] ||= ConsumerGroup.new(group)
+            @groups[group.to_s].compute_lag(topic, partition, offset)
+          end
+        end
+
+        # @param payload [Hash]
+        def heartbeat(payload)
+          group = payload[:group_id]
+          synchronize do
+            @groups[group.to_s] ||= ConsumerGroup.new(group)
+            consumer_group = @groups[group.to_s]
+            payload[:topic_partitions].each do |topic, partitions|
+              partitions.each do |partition|
+                consumer_group.report_lag(topic, partition)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+
+  ActiveSupport::Notifications.subscribe('start_process_message.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.message_processed(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('start_process_batch.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.message_processed(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('seek.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.offset_seek(event.payload)
+  end
+
+  ActiveSupport::Notifications.subscribe('heartbeat.consumer.kafka') do |*args|
+    next unless Deimos.config.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.heartbeat(event.payload)
+  end
+end

data/lib/deimos/utils/signal_handler.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Mostly copied free-form from Phobos::Cli::Runner. We should add a PR to
+    # basically replace that implementation with this one to make it more generic.
+    class SignalHandler
+      SIGNALS = %i(INT TERM QUIT).freeze
+
+      # Takes any object that responds to the `start` and `stop` methods.
+      # @param runner[#start, #stop]
+      def initialize(runner)
+        @signal_queue = []
+        @reader, @writer = IO.pipe
+        @runner = runner
+      end
+
+      # Run the runner.
+      def run!
+        setup_signals
+        @runner.start
+
+        loop do
+          case signal_queue.pop
+          when *SIGNALS
+            @runner.stop
+            break
+          else
+            ready = IO.select([reader, writer])
+
+            # drain the self-pipe so it won't be returned again next time
+            reader.read_nonblock(1) if ready[0].include?(reader)
+          end
+        end
+      end
+
+    private
+
+      attr_reader :reader, :writer, :signal_queue, :executor
+
+      # https://stackoverflow.com/questions/29568298/run-code-when-signal-is-sent-but-do-not-trap-the-signal-in-ruby
+      def prepend_handler(signal)
+        previous = Signal.trap(signal) do
+          previous = -> { raise SignalException, signal } unless previous.respond_to?(:call)
+          yield
+          previous.call
+        end
+      end
+
+      # Trap signals using the self-pipe trick.
+      def setup_signals
+        at_exit { @runner&.stop }
+        SIGNALS.each do |signal|
+          prepend_handler(signal) do
+            unblock(signal)
+          end
+        end
+      end
+
+      # Save the signal to the queue and continue on.
+      # @param signal [Symbol]
+      def unblock(signal)
+        writer.write_nonblock('.')
+        signal_queue << signal
+      end
+    end
+  end
+end

data/lib/deimos/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Deimos
+  VERSION = '1.0.0-beta22'
+end

data/lib/deimos.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require 'avro-patches'
+require 'avro_turf'
+require 'phobos'
+require 'deimos/version'
+require 'deimos/avro_data_encoder'
+require 'deimos/avro_data_decoder'
+require 'deimos/producer'
+require 'deimos/active_record_producer'
+require 'deimos/active_record_consumer'
+require 'deimos/consumer'
+require 'deimos/configuration'
+require 'deimos/instrumentation'
+require 'deimos/utils/lag_reporter'
+
+require 'deimos/publish_backend'
+require 'deimos/backends/kafka'
+require 'deimos/backends/kafka_async'
+
+require 'deimos/monkey_patches/ruby_kafka_heartbeat'
+require 'deimos/monkey_patches/schema_store'
+require 'deimos/monkey_patches/phobos_producer'
+require 'deimos/monkey_patches/phobos_cli'
+
+require 'deimos/railtie' if defined?(Rails)
+if defined?(ActiveRecord)
+  require 'deimos/kafka_source'
+  require 'deimos/kafka_topic_info'
+  require 'deimos/backends/db'
+  require 'deimos/utils/signal_handler.rb'
+  require 'deimos/utils/executor.rb'
+  require 'deimos/utils/db_producer.rb'
+end
+require 'deimos/utils/inline_consumer'
+require 'yaml'
+require 'erb'
+
+# Parent module.
+module Deimos
+  class << self
+    attr_accessor :config
+
+    # Configure Deimos.
+    def configure
+      first_time_config = self.config.nil?
+      self.config ||= Configuration.new
+      old_config = self.config.dup
+      yield(config)
+
+      # Don't re-configure Phobos every time
+      if first_time_config || config.phobos_config_changed?(old_config)
+
+        file = config.phobos_config_file
+        phobos_config = YAML.load(ERB.new(File.read(File.expand_path(file))).result)
+
+        configure_kafka_for_phobos(phobos_config)
+        configure_loggers(phobos_config)
+
+        Phobos.configure(phobos_config)
+      end
+
+      validate_db_backend if self.config.publish_backend == :db
+    end
+
+    # Ensure everything is set up correctly for the DB backend.
+    def validate_db_backend
+      begin
+        require 'activerecord-import'
+      rescue LoadError
+        raise 'Cannot set publish_backend to :db without activerecord-import! Please add it to your Gemfile.'
+      end
+      if Phobos.config.producer_hash[:required_acks] != :all
+        raise 'Cannot set publish_backend to :db unless required_acks is set to ":all" in phobos.yml!'
+      end
+    end
+
+    # Start the DB producers to send Kafka messages.
+    # @param thread_count [Integer] the number of threads to start.
+    def start_db_backend!(thread_count: 1)
+      if self.config.publish_backend != :db # rubocop:disable Style/IfUnlessModifier
+        raise('Publish backend is not set to :db, exiting')
+      end
+
+      if thread_count.nil? || thread_count.zero?
+        raise('Thread count is not given or set to zero, exiting')
+      end
+
+      producers = (1..thread_count).map do
+        Deimos::Utils::DbProducer.new(self.config.logger)
+      end
+      executor = Deimos::Utils::Executor.new(producers,
+                                             self.config.logger)
+      signal_handler = Deimos::Utils::SignalHandler.new(executor)
+      signal_handler.run!
+    end
+
+    # @param phobos_config [Hash]
+    def configure_kafka_for_phobos(phobos_config)
+      if config.ssl_enabled
+        %w(ssl_ca_cert ssl_client_cert ssl_client_cert_key).each do |key|
+          next if config.send(key).blank?
+
+          phobos_config['kafka'][key] = ssl_var_contents(config.send(key))
+        end
+      end
+      phobos_config['kafka']['seed_brokers'] = config.seed_broker if config.seed_broker
+    end
+
+    # @param phobos_config [Hash]
+    def configure_loggers(phobos_config)
+      phobos_config['custom_logger'] = config.phobos_logger
+      phobos_config['custom_kafka_logger'] = config.kafka_logger
+    end
+
+    # @param filename [String] a file to read, or the contents of the SSL var
+    # @return [String] the contents of the file
+    def ssl_var_contents(filename)
+      File.exist?(filename) ? File.read(filename) : filename
+    end
+  end
+end
+
+at_exit do
+  begin
+    Deimos::Backends::KafkaAsync.producer.async_producer_shutdown
+    Deimos::Backends::KafkaAsync.producer.kafka_client&.close
+  rescue StandardError => e
+    Deimos.config.logger.error(
+      "Error closing async producer on shutdown: #{e.message} #{e.backtrace.join("\n")}"
+    )
+  end
+end

data/lib/generators/deimos/db_backend/templates/migration

@@ -0,0 +1,24 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :kafka_messages, force: true do |t|
+      t.string :topic, null: false
+      t.binary :message, limit: 10.megabytes
+      t.binary :key
+      t.string :partition_key
+      t.timestamps
+    end
+
+    add_index :kafka_messages, [:topic, :id]
+
+    create_table :kafka_topic_info, force: true do |t| # rubocop:disable Rails/CreateTableWithTimestamps
+      t.string :topic, null: false
+      t.string :locked_by
+      t.datetime :locked_at
+      t.boolean :error, null: false, default: false
+      t.integer :retries, null: false, default: 0
+    end
+    add_index :kafka_topic_info, :topic, unique: true
+    add_index :kafka_topic_info, [:locked_by, :error]
+    add_index :kafka_topic_info, :locked_at
+  end
+end