deimos-kafka 1.0.0.pre.beta15

data/lib/deimos/tracing/datadog.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require 'deimos/tracing/provider'
+
+module Deimos
+  module Tracing
+    # Tracing wrapper class for Datadog.
+    class Datadog < Tracing::Provider
+      # :nodoc:
+      def initialize(config)
+        raise 'Tracing config must specify service_name' if config[:service_name].nil?
+
+        @service = config[:service_name]
+      end
+
+      # :nodoc:
+      def start(span_name, options={})
+        span = ::Datadog.tracer.trace(span_name)
+        span.service = @service
+        span.resource = options[:resource]
+        span
+      end
+
+      # :nodoc:
+      def finish(span)
+        span.finish
+      end
+
+      # :nodoc:
+      def set_error(span, exception)
+        span.set_error(exception)
+      end
+    end
+  end
+end

data/lib/deimos/tracing/mock.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require 'deimos/tracing/provider'
+
+module Deimos
+  module Tracing
+    # Class that mocks out tracing functionality
+    class Mock < Tracing::Provider
+      # :nodoc:
+      def initialize(logger=nil)
+        @logger = logger || Logger.new(STDOUT)
+        @logger.info('MockTracingProvider initialized')
+      end
+
+      # :nodoc:
+      def start(span_name, _options={})
+        @logger.info("Mock span '#{span_name}' started")
+        {
+          name: span_name,
+          started_at: Time.zone.now
+        }
+      end
+
+      # :nodoc:
+      def finish(span)
+        name = span[:name]
+        start = span[:started_at]
+        finish = Time.zone.now
+        @logger.info("Mock span '#{name}' finished: #{start} to #{finish}")
+      end
+
+      # :nodoc:
+      def set_error(span, exception)
+        span[:exception] = exception
+        name = span[:name]
+        @logger.info("Mock span '#{name}' set an error: #{exception}")
+      end
+    end
+  end
+end

data/lib/deimos/tracing/provider.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+# rubocop:disable Lint/UnusedMethodArgument
+module Deimos
+  module Tracing
+    # Base class for all tracing providers.
+    class Provider
+      # Returns a span object and starts the trace.
+      # @param span_name [String] The name of the span/trace
+      # @param options [Hash] Options for the span
+      # @return [Object] The span object
+      def start(span_name, options={})
+        raise NotImplementedError
+      end
+
+      # Finishes the trace on the span object.
+      # @param span [Object] The span to finish trace on
+      def finish(span)
+        raise NotImplementedError
+      end
+
+      # Set an error on the span.
+      # @param span [Object] The span to set error on
+      # @param exception [Exception] The exception that occurred
+      def set_error(span, exception)
+        raise NotImplementedError
+      end
+    end
+  end
+end
+# rubocop:enable Lint/UnusedMethodArgument

data/lib/deimos/utils/db_producer.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Class which continually polls the database and sends Kafka messages.
+    class DbProducer
+      include Phobos::Producer
+      attr_accessor :id, :current_topic
+
+      BATCH_SIZE = 1000
+
+      # @param logger [Logger]
+      def initialize(logger=Logger.new(STDOUT))
+        @id = SecureRandom.uuid
+        @logger = logger
+        @logger.push_tags("DbProducer #{@id}") if @logger.respond_to?(:push_tags)
+      end
+
+      # Start the poll.
+      def start
+        @logger.info('Starting...')
+        @signal_to_stop = false
+        loop do
+          if @signal_to_stop
+            @logger.info('Shutting down')
+            break
+          end
+          process_next_messages
+        end
+      end
+
+      # Stop the poll.
+      def stop
+        @logger.info('Received signal to stop')
+        @signal_to_stop = true
+      end
+
+      # Complete one loop of processing all messages in the DB.
+      def process_next_messages
+        topics = retrieve_topics
+        @logger.info("Found topics: #{topics}")
+        topics.each(&method(:process_topic))
+        sleep(0.5)
+      end
+
+      # @return [Array<String>]
+      def retrieve_topics
+        KafkaMessage.select('distinct topic').map(&:topic).uniq
+      end
+
+      # @param topic [String]
+      # @return [String] the topic that was locked, or nil if none were.
+      def process_topic(topic)
+        # If the topic is already locked, another producer is currently
+        # working on it. Move on to the next one.
+        unless KafkaTopicInfo.lock(topic, @id)
+          @logger.debug("Could not lock topic #{topic} - continuing")
+          return
+        end
+        @current_topic = topic
+        messages = retrieve_messages
+
+        while messages.any?
+          produce_messages(messages.map(&:phobos_message))
+          messages.first.class.where(id: messages.map(&:id)).delete_all
+          break if messages.size < BATCH_SIZE
+
+          KafkaTopicInfo.heartbeat(@current_topic, @id) # keep alive
+          messages = retrieve_messages
+        end
+        KafkaTopicInfo.clear_lock(@current_topic, @id)
+      rescue StandardError => e
+        @logger.error("Error processing messages for topic #{@current_topic}: #{e.class.name}: #{e.message} #{e.backtrace.join("\n")}")
+        KafkaTopicInfo.register_error(@current_topic, @id)
+      end
+
+      # @return [Array<KafkaMessage>]
+      def retrieve_messages
+        KafkaMessage.where(topic: @current_topic).order(:id).limit(BATCH_SIZE)
+      end
+
+      # @param batch [Array<Hash>]
+      def produce_messages(batch)
+        @logger.debug("Publishing #{batch.size} messages to #{@current_topic}")
+        producer.publish_list(batch)
+        Deimos.metrics&.increment(
+          'publish',
+          tags: %W(status:success topic:#{@current_topic}),
+          by: batch.size
+        )
+        @logger.info("Sent #{batch.size} messages to #{@current_topic}")
+      end
+    end
+  end
+end

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