deimos-temp-fork 0.0.1

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/inline_consumer.rb

@@ -0,0 +1,150 @@
+# 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
+      MAX_SEEK_RETRIES = 3
+      attr_accessor :num_messages
+
+      # :nodoc:
+      def start_listener
+        @num_messages ||= 10
+        @consumer = create_kafka_consumer
+        @consumer.subscribe(topic, @subscribe_opts)
+        attempt = 0
+
+        begin
+          attempt += 1
+          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
+          if attempt < MAX_SEEK_RETRIES
+            sleep(1.seconds * attempt)
+            retry
+          end
+          log_error("Could not seek to offset: #{e.message} after #{MAX_SEEK_RETRIES} retries", listener_metadata)
+        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)
+        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,175 @@
+# 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 offset [Integer]
+        def assign_current_offset(topic, partition, offset)
+          self.topics[topic.to_s] ||= Topic.new(topic, self)
+          self.topics[topic.to_s].assign_current_offset(partition, offset)
+        end
+      end
+
+      # 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_current_offsets
+        # @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_current_offsets = {}
+        end
+
+        # @param partition [Integer]
+        def assign_current_offset(partition, offset)
+          self.partition_current_offsets[partition.to_i] = offset
+        end
+
+        # @param partition [Integer]
+        def compute_lag(partition, offset)
+          begin
+            client = Phobos.create_kafka_client
+            last_offset = client.last_offset_for(self.topic_name, partition)
+            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)
+          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}")
+          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
+
+        # 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)
+          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_current_offset(topic, partition, offset)
+          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].assign_current_offset(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.consumers.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.consumers.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.consumers.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.consumers.report_lag
+
+    event = ActiveSupport::Notifications::Event.new(*args)
+    Deimos::Utils::LagReporter.heartbeat(event.payload)
+  end
+end

data/lib/deimos/utils/schema_controller_mixin.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Mixin to automatically decode schema-encoded payloads when given the correct content type,
+    # and provide the `render_schema` method to encode the payload for responses.
+    module SchemaControllerMixin
+      extend ActiveSupport::Concern
+
+      included do
+        Mime::Type.register('avro/binary', :avro)
+
+        attr_accessor :payload
+
+        if respond_to?(:before_filter)
+          before_filter(:decode_schema, if: :schema_format?)
+        else
+          before_action(:decode_schema, if: :schema_format?)
+        end
+      end
+
+      # :nodoc:
+      module ClassMethods
+        # @return [Hash<String, Hash<Symbol, String>>]
+        def schema_mapping
+          @schema_mapping ||= {}
+        end
+
+        # Indicate which schemas should be assigned to actions.
+        # @param actions [Symbol]
+        # @param kwactions [String]
+        # @param request [String]
+        # @param response [String]
+        def schemas(*actions, request: nil, response: nil, **kwactions)
+          actions.each do |action|
+            request ||= action.to_s.titleize
+            response ||= action.to_s.titleize
+            schema_mapping[action.to_s] = { request: request, response: response }
+          end
+          kwactions.each do |key, val|
+            schema_mapping[key.to_s] = { request: val, response: val }
+          end
+        end
+
+        # @return [Hash<Symbol, String>]
+        def namespaces
+          @namespaces ||= {}
+        end
+
+        # Set the namespace for both requests and responses.
+        # @param name [String]
+        def namespace(name)
+          request_namespace(name)
+          response_namespace(name)
+        end
+
+        # Set the namespace for requests.
+        # @param name [String]
+        def request_namespace(name)
+          namespaces[:request] = name
+        end
+
+        # Set the namespace for repsonses.
+        # @param name [String]
+        def response_namespace(name)
+          namespaces[:response] = name
+        end
+      end
+
+      # @return [Boolean]
+      def schema_format?
+        request.content_type == Deimos.schema_backend_class.content_type
+      end
+
+      # Get the namespace from either an existing instance variable, or tease it out of the schema.
+      # @param type [Symbol] :request or :response
+      # @return [Array<String, String>] the namespace and schema.
+      def parse_namespace(type)
+        namespace = self.class.namespaces[type]
+        schema = self.class.schema_mapping[params['action']][type]
+        if schema.nil?
+          raise "No #{type} schema defined for #{params[:controller]}##{params[:action]}!"
+        end
+
+        if namespace.nil?
+          last_period = schema.rindex('.')
+          namespace, schema = schema.split(last_period)
+        end
+        if namespace.nil? || schema.nil?
+          raise "No request namespace defined for #{params[:controller]}##{params[:action]}!"
+        end
+
+        [namespace, schema]
+      end
+
+      # Decode the payload with the parameters.
+      def decode_schema
+        namespace, schema = parse_namespace(:request)
+        decoder = Deimos.schema_backend(schema: schema, namespace: namespace)
+        @payload = decoder.decode(request.body.read).with_indifferent_access
+        request.body.rewind if request.body.respond_to?(:rewind)
+      end
+
+      # Render a hash into a payload as specified by the configured schema and namespace.
+      # @param payload [Hash]
+      def render_schema(payload, schema: nil, namespace: nil)
+        namespace, schema = parse_namespace(:response) if !schema && !namespace
+        encoder = Deimos.schema_backend(schema: schema, namespace: namespace)
+        encoded = encoder.encode(payload)
+        response.headers['Content-Type'] = encoder.class.content_type
+        send_data(encoded)
+      end
+    end
+  end
+end