deimos-temp-fork 0.0.1

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,29 @@
+# frozen_string_literal: true
+
+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

data/lib/deimos/utils/db_poller.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+require 'deimos/poll_info'
+require 'sigurd/executor'
+require 'sigurd/signal_handler'
+
+module Deimos
+  module Utils
+    # Class which continually polls the database and sends Kafka messages.
+    class DbPoller
+      BATCH_SIZE = 1000
+
+      # Needed for Executor so it can identify the worker
+      attr_reader :id
+
+      # Begin the DB Poller process.
+      def self.start!
+        if Deimos.config.db_poller_objects.empty?
+          raise('No pollers configured!')
+        end
+
+        pollers = Deimos.config.db_poller_objects.map do |poller_config|
+          self.new(poller_config)
+        end
+        executor = Sigurd::Executor.new(pollers,
+                                        sleep_seconds: 5,
+                                        logger: Deimos.config.logger)
+        signal_handler = Sigurd::SignalHandler.new(executor)
+        signal_handler.run!
+      end
+
+      # @param config [Deimos::Configuration::ConfigStruct]
+      def initialize(config)
+        @config = config
+        @id = SecureRandom.hex
+        begin
+          @producer = @config.producer_class.constantize
+        rescue NameError
+          raise "Class #{@config.producer_class} not found!"
+        end
+        unless @producer < Deimos::ActiveRecordProducer
+          raise "Class #{@producer.class.name} is not an ActiveRecordProducer!"
+        end
+      end
+
+      # Start the poll:
+      # 1) Grab the current PollInfo from the database indicating the last
+      # time we ran
+      # 2) On a loop, process all the recent updates between the last time
+      # we ran and now.
+      def start
+        # Don't send asynchronously
+        if Deimos.config.producers.backend == :kafka_async
+          Deimos.config.producers.backend = :kafka
+        end
+        Deimos.config.logger.info('Starting...')
+        @signal_to_stop = false
+        retrieve_poll_info
+        loop do
+          if @signal_to_stop
+            Deimos.config.logger.info('Shutting down')
+            break
+          end
+          process_updates
+          sleep 0.1
+        end
+      end
+
+      # Grab the PollInfo or create if it doesn't exist.
+      def retrieve_poll_info
+        ActiveRecord::Base.connection.reconnect! unless ActiveRecord::Base.connection.open_transactions.positive?
+        new_time = @config.start_from_beginning ? Time.new(0) : Time.zone.now
+        @info = Deimos::PollInfo.find_by_producer(@config.producer_class) ||
+                Deimos::PollInfo.create!(producer: @config.producer_class,
+                                         last_sent: new_time,
+                                         last_sent_id: 0)
+      end
+
+      # Stop the poll.
+      def stop
+        Deimos.config.logger.info('Received signal to stop')
+        @signal_to_stop = true
+      end
+
+      # Indicate whether this current loop should process updates. Most loops
+      # will busy-wait (sleeping 0.1 seconds) until it's ready.
+      # @return [Boolean]
+      def should_run?
+        Time.zone.now - @info.last_sent - @config.delay_time >= @config.run_every
+      end
+
+      # @param record [ActiveRecord::Base]
+      # @return [ActiveSupport::TimeWithZone]
+      def last_updated(record)
+        record.public_send(@config.timestamp_column)
+      end
+
+      # Send messages for updated data.
+      def process_updates
+        return unless should_run?
+
+        time_from = @config.full_table ? Time.new(0) : @info.last_sent.in_time_zone
+        time_to = Time.zone.now - @config.delay_time
+        Deimos.config.logger.info("Polling #{@producer.topic} from #{time_from} to #{time_to}")
+        message_count = 0
+        batch_count = 0
+
+        # poll_query gets all the relevant data from the database, as defined
+        # by the producer itself.
+        loop do
+          Deimos.config.logger.debug("Polling #{@producer.topic}, batch #{batch_count + 1}")
+          batch = fetch_results(time_from, time_to).to_a
+          break if batch.empty?
+
+          batch_count += 1
+          process_batch(batch)
+          message_count += batch.size
+          time_from = last_updated(batch.last)
+        end
+        Deimos.config.logger.info("Poll #{@producer.topic} complete at #{time_to} (#{message_count} messages, #{batch_count} batches}")
+      end
+
+      # @param time_from [ActiveSupport::TimeWithZone]
+      # @param time_to [ActiveSupport::TimeWithZone]
+      # @return [ActiveRecord::Relation]
+      def fetch_results(time_from, time_to)
+        id = @producer.config[:record_class].primary_key
+        quoted_timestamp = ActiveRecord::Base.connection.quote_column_name(@config.timestamp_column)
+        quoted_id = ActiveRecord::Base.connection.quote_column_name(id)
+        @producer.poll_query(time_from: time_from,
+                             time_to: time_to,
+                             column_name: @config.timestamp_column,
+                             min_id: @info.last_sent_id).
+          limit(BATCH_SIZE).
+          order("#{quoted_timestamp}, #{quoted_id}")
+      end
+
+      # @param batch [Array<ActiveRecord::Base>]
+      def process_batch(batch)
+        record = batch.last
+        id_method = record.class.primary_key
+        last_id = record.public_send(id_method)
+        last_updated_at = last_updated(record)
+        @producer.send_events(batch)
+        @info.attributes = { last_sent: last_updated_at, last_sent_id: last_id }
+        @info.save!
+      end
+    end
+  end
+end

data/lib/deimos/utils/db_producer.rb

@@ -0,0 +1,243 @@
+# frozen_string_literal: true
+
+module Deimos
+  module Utils
+    # Class which continually polls the kafka_messages table
+    # in the database and sends Kafka messages.
+    class DbProducer
+      include Phobos::Producer
+      attr_accessor :id, :current_topic
+
+      BATCH_SIZE = 1000
+      DELETE_BATCH_SIZE = 10
+      MAX_DELETE_ATTEMPTS = 3
+
+      # @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
+
+      # @return [Deimos::DbProducerConfig]
+      def config
+        Deimos.config.db_producer
+      end
+
+      # Start the poll.
+      def start
+        @logger.info('Starting...')
+        @signal_to_stop = false
+        ActiveRecord::Base.connection.reconnect!
+        loop do
+          if @signal_to_stop
+            @logger.info('Shutting down')
+            break
+          end
+          send_pending_metrics
+          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))
+        KafkaTopicInfo.ping_empty_topics(topics)
+        sleep(0.5)
+      end
+
+      # @return [Array<String>]
+      def retrieve_topics
+        KafkaMessage.select('distinct topic').map(&:topic).uniq
+      end
+
+      # @param topic [String]
+      # @return [String, nil] 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
+
+        loop { break unless process_topic_batch }
+
+        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)
+        shutdown_producer
+      end
+
+      # Process a single batch in a topic.
+      def process_topic_batch
+        messages = retrieve_messages
+        return false if messages.empty?
+
+        batch_size = messages.size
+        compacted_messages = compact_messages(messages)
+        log_messages(compacted_messages)
+        Deimos.instrument('db_producer.produce', topic: @current_topic, messages: compacted_messages) do
+          begin
+            produce_messages(compacted_messages.map(&:phobos_message))
+          rescue Kafka::BufferOverflow, Kafka::MessageSizeTooLarge, Kafka::RecordListTooLarge
+            delete_messages(messages)
+            @logger.error('Message batch too large, deleting...')
+            @logger.error(Deimos::KafkaMessage.decoded(messages))
+            raise
+          end
+        end
+        delete_messages(messages)
+        Deimos.config.metrics&.increment(
+          'db_producer.process',
+          tags: %W(topic:#{@current_topic}),
+          by: messages.size
+        )
+        return false if batch_size < BATCH_SIZE
+
+        KafkaTopicInfo.heartbeat(@current_topic, @id) # keep alive
+        send_pending_metrics
+        true
+      end
+
+      # @param messages [Array<Deimos::KafkaMessage>]
+      def delete_messages(messages)
+        attempts = 1
+        begin
+          messages.in_groups_of(DELETE_BATCH_SIZE, false).each do |batch|
+            Deimos::KafkaMessage.where(topic: batch.first.topic,
+                                       id: batch.map(&:id)).
+              delete_all
+          end
+        rescue StandardError => e
+          if (e.message =~ /Lock wait/i || e.message =~ /Lost connection/i) &&
+             attempts <= MAX_DELETE_ATTEMPTS
+            attempts += 1
+            ActiveRecord::Base.connection.verify!
+            sleep(1)
+            retry
+          end
+          raise
+        end
+      end
+
+      # @return [Array<Deimos::KafkaMessage>]
+      def retrieve_messages
+        KafkaMessage.where(topic: @current_topic).order(:id).limit(BATCH_SIZE)
+      end
+
+      # @param messages [Array<Deimos::KafkaMessage>]
+      def log_messages(messages)
+        return if config.log_topics != :all && !config.log_topics.include?(@current_topic)
+
+        @logger.debug do
+          decoded_messages = Deimos::KafkaMessage.decoded(messages)
+          "DB producer: Topic #{@current_topic} Producing messages: #{decoded_messages}}"
+        end
+      end
+
+      # Send metrics to Datadog.
+      def send_pending_metrics
+        metrics = Deimos.config.metrics
+        return unless metrics
+
+        topics = KafkaTopicInfo.select(%w(topic last_processed_at))
+        messages = Deimos::KafkaMessage.
+          select('count(*) as num_messages, min(created_at) as earliest, topic').
+          group(:topic).
+          index_by(&:topic)
+        topics.each do |record|
+          message_record = messages[record.topic]
+          # We want to record the last time we saw any activity, meaning either
+          # the oldest message, or the last time we processed, whichever comes
+          # last.
+          if message_record
+            record_earliest = message_record.earliest
+            # SQLite gives a string here
+            if record_earliest.is_a?(String)
+              record_earliest = Time.zone.parse(record_earliest)
+            end
+
+            earliest = [record.last_processed_at, record_earliest].max
+            time_diff = Time.zone.now - earliest
+            metrics.gauge('pending_db_messages_max_wait', time_diff,
+                          tags: ["topic:#{record.topic}"])
+          else
+            # no messages waiting
+            metrics.gauge('pending_db_messages_max_wait', 0,
+                          tags: ["topic:#{record.topic}"])
+          end
+          metrics.gauge('pending_db_messages_count', message_record&.num_messages || 0,
+                        tags: ["topic:#{record.topic}"])
+        end
+      end
+
+      # Shut down the sync producer if we have to. Phobos will automatically
+      # create a new one. We should call this if the producer can be in a bad
+      # state and e.g. we need to clear the buffer.
+      def shutdown_producer
+        if self.class.producer.respond_to?(:sync_producer_shutdown) # Phobos 1.8.3
+          self.class.producer.sync_producer_shutdown
+        end
+      end
+
+      # Produce messages in batches, reducing the size 1/10 if the batch is too
+      # large. Does not retry batches of messages that have already been sent.
+      # @param batch [Array<Hash>]
+      def produce_messages(batch)
+        batch_size = batch.size
+        current_index = 0
+        begin
+          batch[current_index..-1].in_groups_of(batch_size, false).each do |group|
+            @logger.debug("Publishing #{group.size} messages to #{@current_topic}")
+            producer.publish_list(group)
+            Deimos.config.metrics&.increment(
+              'publish',
+              tags: %W(status:success topic:#{@current_topic}),
+              by: group.size
+            )
+            current_index += group.size
+            @logger.info("Sent #{group.size} messages to #{@current_topic}")
+          end
+        rescue Kafka::BufferOverflow, Kafka::MessageSizeTooLarge,
+               Kafka::RecordListTooLarge => e
+          if batch_size == 1
+            shutdown_producer
+            raise
+          end
+
+          @logger.error("Got error #{e.class.name} when publishing #{batch.size} in groups of #{batch_size}, retrying...")
+          batch_size = if batch_size < 10
+                         1
+                       else
+                         (batch_size / 10).to_i
+                       end
+          shutdown_producer
+          retry
+        end
+      end
+
+      # @param batch [Array<Deimos::KafkaMessage>]
+      # @return [Array<Deimos::KafkaMessage>]
+      def compact_messages(batch)
+        return batch if batch.first&.key.blank?
+
+        topic = batch.first.topic
+        return batch if config.compact_topics != :all &&
+                        !config.compact_topics.include?(topic)
+
+        batch.reverse.uniq(&:key).reverse!
+      end
+    end
+  end
+end