rails-pipeline 1.1.1

data/lib/rails-pipeline/redis_forwarder.rb

@@ -0,0 +1,207 @@
+
+require 'redis'
+require 'active_support/core_ext'
+require 'rails-pipeline/protobuf/encrypted_message.pb'
+
+# Pipeline forwarder base class that
+#   - reads from redis queue using BRPOPLPUSH for reliable queue pattern
+#   - keeps track of failed tasks in the in_progress queue
+#   - designed to be used with e.g. IronmqPublisher
+
+$redis = ENV["REDISCLOUD_URL"] || ENV["REDISTOGO_URL"] || "localhost:6379"
+
+module RailsPipeline
+  class RedisForwarder
+    if RailsPipeline::HAS_NEWRELIC
+      include ::NewRelic::Agent::MethodTracer
+    end
+
+    def initialize(key)
+      _trap_signals
+      @redis = nil
+      @stop = false
+      @queue = key
+      @in_progress_queue = _in_progress_queue
+
+      @processed = 0
+      @blocking_timeout = 2
+      @failure_check_interval = 30
+      @message_processing_limit = 10 # number of seconds before a message is considered failed
+      @failure_last_checked = Time.now - @failure_check_interval.seconds  # TODO: randomize start time?
+    end
+
+    def _trap_signals
+      trap('SIGTERM') do
+        puts 'Exiting (SIGTERM)'
+        stop
+      end
+      trap('SIGINT') do
+        puts 'Exiting (SIGINT)'
+        stop
+      end
+    end
+
+
+    # Blocking right pop from the queue
+    #    - use BRPOPLPUSH to tenporarily mark the message as "in progress"
+    #    - delete from the in_prgress queue on success
+    #    - restore to the main queue on failure
+    def process_queue
+      # pop from the queue and push onto the in_progress queue
+      data = _redis.brpoplpush(@queue, @in_progress_queue, timeout: @blocking_timeout)
+      if data.nil?  # Timed-out with nothing to process
+        return
+      end
+
+      begin
+        encrypted_data = RailsPipeline::EncryptedMessage.parse(data)
+        RailsPipeline.logger.debug "Processing #{encrypted_data.uuid}"
+
+        # re-publish to wherever (e.g. IronMQ)
+        topic_name = encrypted_data.topic
+        if topic_name.nil?
+          RailsPipeline.logger.error "Damaged message, no topic name"
+          return
+        end
+
+        publish(topic_name, data)
+        @processed += 1
+
+        # Now remove this message from the in_progress queue
+        removed = _redis.lrem(@in_progress_queue, 1, data)
+        if removed != 1
+          RailsPipeline.logger.warn "OHNO! Didn't remove the data I was expecting to: #{data}"
+        end
+      rescue Exception => e
+        RailsPipeline.logger.info e
+        RailsPipeline.logger.info e.backtrace.join("\n")
+        if !data.nil?
+          RailsPipeline.logger.info "Putting message #{encrypted_data.uuid} back on main queue"
+          _put_back_on_queue(data)
+        end
+      end
+    end
+    add_method_tracer :process_queue, "Pipeline/RedisForwarder/process_queue" if RailsPipeline::HAS_NEWRELIC
+
+    # note in redis that we are processing this message
+    def report(uuid)
+      _redis.setex(_report_key(uuid), @message_processing_limit, _client_id)
+    end
+
+    # Search the in-progress queue for messages that are likely to be abandoned
+    # and re-queue them on the main queue
+    def check_for_failures
+      # Lock in_progress queue or return
+      num_in_progress = _redis.llen(@in_progress_queue)
+      if num_in_progress == 0
+        RailsPipeline.logger.debug "No messages in progress, skipping check for failures"
+        return
+      end
+
+      RailsPipeline.logger.debug "Locking '#{@in_progress_queue}' for #{num_in_progress} seconds"
+
+      # Attempt to lock this queue for the next num_in_progress seconds
+      lock_key = "#{@in_progress_queue}__lock"
+      locked = _redis.set(lock_key, _client_id, ex: num_in_progress, nx: true)
+      if !locked
+        RailsPipeline.logger.debug "in progress queue is locked"
+        return
+      end
+
+      # Go through each message, see if there's a 'report' entry. If not,
+      # requeue!
+      in_progress = _redis.lrange(@in_progress_queue, 0, num_in_progress)
+      in_progress.each do |message|
+        enc_message = EncryptedMessage.parse(message)
+        owner = _redis.get(_report_key(enc_message.uuid))
+        if owner.nil?
+          RailsPipeline.logger.info "Putting timed-out message #{enc_message.uuid} back on main queue"
+          _put_back_on_queue(message)
+        else
+          RailsPipeline.logger.debug "Message #{uuid} is owned by #{owner}"
+        end
+      end
+    end
+    add_method_tracer :check_for_failures, "Pipeline/RedisForwarder/check_for_failures" if RailsPipeline::HAS_NEWRELIC
+
+    # Function that runs in the loop
+    def run
+      process_queue
+      RailsPipeline.logger.info "Queue: '#{@queue}'. Processed: #{@processed}"
+      if Time.now - @failure_last_checked > @failure_check_interval
+        @failure_last_checked = Time.now
+        check_for_failures
+      end
+    end
+
+    # Main loop
+    def start
+      while true
+        begin
+          if @stop
+            RailsPipeline.logger.info "Finished"
+            if RailsPipeline::HAS_NEWRELIC
+              RailsPipeline.logger.info "Shutting down NewRelic"
+              ::NewRelic::Agent.shutdown
+            end
+            break
+          end
+          run
+        rescue Exception => e
+          RailsPipeline.logger.info e
+          RailsPipeline.logger.info e.backtrace.join("\n")
+        end
+      end
+    end
+
+    def stop
+      puts "stopping..."
+      @stop = true
+    end
+
+    def _redis
+      if !@redis.nil?
+        return @redis
+      end
+      if $redis.start_with?("redis://")
+        @redis = Redis.new(url: $redis)
+      else
+        host, port = $redis.split(":")
+        @redis = Redis.new(host: host, port: port)
+      end
+      return @redis
+    end
+
+    def _processed
+      return @processed
+    end
+
+    def _in_progress_queue
+      "#{@queue}_in_progress"
+    end
+
+    # The redis key at which we 'claim' the message when we start processing it.
+    def _report_key(uuid)
+      "#{@queue}__#{uuid}"
+    end
+
+    def _client_id
+      self.class.name
+    end
+
+    # Atomically remove a message from the in_progress queue and put it back on
+    # the main queue
+    def _put_back_on_queue(message)
+      future = nil
+      _redis.multi do
+        _redis.rpush(@queue, message)
+        future = _redis.lrem(@in_progress_queue, 1, message)
+      end
+      removed = future.value
+      if removed !=1
+        RailsPipeline.logger.error "ERROR: Didn't remove message from in_progress queue?!!!"
+      end
+    end
+
+  end
+end

data/lib/rails-pipeline/redis_ironmq_forwarder.rb

@@ -0,0 +1,12 @@
+
+require 'rails-pipeline/redis_forwarder'
+require 'rails-pipeline/ironmq_publisher'
+
+# Mix-in the IronMQ publisher into a RedisForwarder to create a
+# class that will forward redis messages onto IronMQ
+
+module RailsPipeline
+  class RedisIronmqForwarder < RedisForwarder
+    include IronmqPublisher
+  end
+end

data/lib/rails-pipeline/redis_publisher.rb

@@ -0,0 +1,71 @@
+
+require 'redis'
+
+# Backend for data pipeline that publishes to redis queue
+# (typically for consumption by a log sender)
+#
+# Typically initialized in rails initialzer e.g.
+# RailsPipeline::RedisPublisher.redis = Redis.new(ENV["REDIS_URL"])
+# RailsPipeline::RedisPublisher.namespace = "my-app-pipeline"
+
+module RailsPipeline
+  module RedisPublisher
+    class << self
+      # Allow configuration via initializer
+      @@redis = nil
+      @@namespace = "pipeline" # default redis queue name
+      attr_accessor :namespace
+      def _redis
+        if @@redis.nil?
+          if $redis.start_with?("redis://")
+            @@redis = Redis.new(url: $redis)
+          else
+            host, port = $redis.split(":")
+            @@redis = Redis.new(host: host, port: port)
+          end
+        end
+        @@redis
+      end
+      def redis=(redis)
+        @@redis = redis
+      end
+      def namespace=(namespace)
+        @@namespace = namespace
+      end
+      def namespace
+        @@namespace
+      end
+
+
+    end
+
+    def self.included(base)
+      base.extend ClassMethods
+      base.send :include, InstanceMethods
+    end
+
+    module InstanceMethods
+      def publish(topic_name, data)
+        t0 = Time.now
+        _redis.lpush(_key, data)
+        t1 = Time.now
+        ::NewRelic::Agent.record_metric('Pipeline/Redis/publish', t1-t0) if RailsPipeline::HAS_NEWRELIC
+        RailsPipeline.logger.debug "Publishing to redis '#{topic_name}' took #{t1-t0}s"
+      end
+
+      def _redis
+        RedisPublisher._redis
+      end
+
+      def _key
+        RedisPublisher.namespace
+      end
+
+    end
+
+    module ClassMethods
+
+    end
+  end
+
+end

data/lib/rails-pipeline/sns_publisher.rb

@@ -0,0 +1,62 @@
+
+
+require 'aws-sdk'
+
+# Backend for data pipeline that publishes to Amazon Simple Notification
+# Service (SNS).
+#
+# Configure via an initializer like:
+# PipelineSnsEmitter.account_id = "6982739827398"
+
+module RailsPipeline::SnsPublisher
+  class << self
+    # Allow configuration via initializer
+    @@account_id = nil
+    def account_id
+      @@account_id
+    end
+    def account_id=(account_id)
+      @@account_id = account_id
+    end
+  end
+
+  def self.included(base)
+    base.send :include, InstanceMethods
+    base.extend ClassMethods
+  end
+
+  module InstanceMethods
+    def publish(topic_name, data)
+      t0 = Time.now
+      topic = _sns.topics[_topic_arn(topic_name)]
+      topic.publish(data, subject: _subject, sqs: data)
+      t1 = Time.now
+      RailsPipeline.logger.debug "Published to SNS '#{topic_name}' in #{t1-t0}s"
+    end
+
+    def _sns
+      @sns = AWS::SNS.new if @sns.nil?
+      return @sns
+    end
+
+    def _topic_arn(topic_name, region="us-east-1")
+      "arn:aws:sns:#{region}:#{_account_id}:#{topic_name}"
+    end
+
+    # Subject of SNS message is ClassName-id
+    def _subject
+      "#{self.class.name}-#{self.id}"
+    end
+
+    def _account_id
+      if ENV.has_key?("AWS_ACCOUNT_ID")
+        return ENV["AWS_ACCOUNT_ID"]
+      end
+      return RailsPipeline::SnsPublisher.account_id
+    end
+  end
+
+  module ClassMethods
+  end
+
+end

data/lib/rails-pipeline/subscriber.rb

@@ -0,0 +1,185 @@
+require "rails-pipeline/symmetric_encryptor"
+
+module RailsPipeline
+
+  module Subscriber
+
+    Error = Class.new(StandardError)
+    NoApiKeyError = Class.new(Error)
+    WrongApiKeyError = Class.new(Error)
+
+    class << self
+
+      @@registered_models = {}
+      @@registered_handlers = {}
+
+      def register(payload_class, target_class, handler = nil)
+        @@registered_models[payload_class] = target_class
+        @@registered_handlers[payload_class] = handler
+      end
+
+      def target_class(payload_class)
+        @@registered_models[payload_class]
+      end
+
+      def target_handler(payload_class)
+        @@registered_handlers[payload_class]
+      end
+
+      def registered_handlers
+        @@registered_handlers
+      end
+    end
+
+
+    def self.included(base)
+      RailsPipeline::SymmetricEncryptor.included(base)
+      base.send :include, InstanceMethods
+      base.extend ClassMethods
+      if RailsPipeline::HAS_NEWRELIC
+        base.send :include, ::NewRelic::Agent::Instrumentation::ControllerInstrumentation
+        base.extend  ::NewRelic::Agent::Instrumentation::ControllerInstrumentation::ClassMethods
+        base.add_transaction_tracer :handle_envelope, category: :task
+        base.add_transaction_tracer :handle_payload, category: :task
+      end
+    end
+
+    module InstanceMethods
+
+      # Take an EncryptedMessage envelope, and decrypt the cipher text, then
+      # get the protobuf object out of it
+      def handle_envelope(envelope)
+        if ENV.has_key?("DISABLE_RAILS_PIPELINE") || ENV.has_key?("DISABLE_RAILS_PIPELINE_PROCESSING")
+          RailsPipeline.logger.debug "Skipping incoming pipeline messages (disabled by env vars)"
+          return
+        end
+        verify_api_key(envelope)
+        payload_str = self.class.decrypt(envelope)
+
+        # Find the registered minor version & its related handler to parse and
+        # process this message.
+        clazz = registered_class_on_same_major_version(envelope.type_info)
+
+        if clazz.nil?
+          # No compatible version of this message is registered for this app.
+          RailsPipeline.logger.info "Dropping unclaimed message #{envelope.type_info} (no compatible version registered)."
+          return
+        end
+
+        # Parse and handle the payload.
+        payload = clazz.parse(payload_str)
+        handle_payload(payload, envelope)
+      end
+
+      # Take a protobuf object (payload) and forward it to the appropriate
+      # handler/method/proc
+      def handle_payload(payload, envelope)
+        version = _version(payload)
+        clazz = target_class(payload)
+        handler_class = target_handler(payload)
+        event_type = envelope.event_type
+        method = most_suitable_handler_method_name(version, clazz)
+
+        if clazz.is_a?(Class)
+          if handler_class
+            # If a built in handler_class is registered, then just use it
+            handler_class.new(payload, target_class: clazz, envelope: envelope).handle_payload
+          elsif method
+            # Target class had a from_pipeline method, so just call it and move on
+            target = clazz.send(method, payload, event_type)
+          else
+            RailsPipeline.logger.info "No handler set, dropping message #{payload.class.name}"
+          end
+          return target
+        elsif clazz.is_a?(Proc)
+          return clazz.call(payload)
+        end
+      end
+
+      def registered_class_on_same_major_version(payload_class_name)
+        if Subscriber.registered_handlers.has_key?(payload_class_name)
+          # The version we've been given has been registered. Return that.
+          return Object.const_get(payload_class_name)
+        end
+
+        # Lops off the minor version from the payload class name so
+        # we can find the registered message type with the same major version.
+        class_name_with_major_version =
+          payload_class_name
+          .split("_", 3)[0, 2]
+          .join("_")
+
+        # Look for message types with the same major version.
+        available_classes = Subscriber.registered_handlers
+          .keys
+          .map(&:to_s)
+          .select { |class_name| class_name.start_with?(class_name_with_major_version) }
+
+        if available_classes.empty?
+          # No message types with the same major version.
+          return nil
+        else
+          # There's a message type with the same major version.
+          return Object.const_get(available_classes.first)
+        end
+      end
+
+      def most_suitable_handler_method_name(version, receiver_class)
+        # Returns the closest lower implemented method in target_class for the given version
+        cached_method = self.class.handler_method_cache[version]
+        if cached_method
+          return cached_method
+        end
+        available_methods = receiver_class.methods.grep(%r{^from_pipeline_#{version.major}})
+          .reject { |method_name| method_name.to_s.split('_').last.to_i > version.minor }
+          .sort
+          .reverse
+
+        # cache handler method for this version
+        self.class.handler_method_cache[version] = available_methods.first
+        return available_methods.first
+      end
+
+      def verify_api_key(envelope)
+        if envelope.api_key.present?
+          if _api_keys.include?(envelope.api_key)
+            return true
+          else
+            raise WrongApiKeyError.new
+          end
+        else
+          raise NoApiKeyError.new
+        end
+      end
+
+      def target_class(payload)
+        RailsPipeline::Subscriber.target_class(payload.class)
+      end
+
+      def target_handler(payload)
+        RailsPipeline::Subscriber.target_handler(payload.class)
+      end
+
+      def _version(payload)
+        _, version = payload.class.name.split('_', 2)
+        return RailsPipeline::PipelineVersion.new(version)
+      end
+
+      def _api_keys
+        return ENV.fetch('PIPELINE_API_KEYS', "").split(',')
+      end
+
+    end
+
+    module ClassMethods
+
+      def handler_method_cache
+        @handler_method_cache ||= {}
+      end
+
+      def handler_method_cache=(cache)
+        @handler_method_cache = cache
+      end
+    end
+  end
+end