fourkites-sqspoller-v2 1.0.0

data/lib/sqs-ruby-example-create-queue.rb

@@ -0,0 +1,107 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX - License - Identifier: Apache - 2.0
+
+# Purpose
+
+# This code example demonstrates how to create a queue in Amazon Simple Queue Service (Amazon SQS).
+
+require 'aws-sdk-sqs'
+require "net/http"
+require "rest-client"
+require "json"
+require 'concurrent'
+
+# @param sqs_client [Aws::SQS::Client] An initialized Amazon SQS client.
+# @param queue_name [String] The name of the queue.
+# @return [Boolean] true if the queue was created; otherwise, false.
+# @example
+#   exit 1 unless queue_created?(
+#     Aws::SQS::Client.new(region: 'us-west-2'),
+#     'my-queue'
+#   )
+HEADERS = {
+  'Content-Type' => 'application/json',
+  'Accept' => 'application/json'
+}
+
+def queue_created?(sqs_client, queue_name)
+  sqs_client.create_queue(queue_name: queue_name)
+  true
+rescue StandardError => e
+  puts "Error creating queue: #{e.message}"
+  false
+end
+
+def messages_sent?(sqs_client, queue_url, entries)
+  sqs_client.send_message_batch(
+    queue_url: queue_url,
+    entries: entries
+  )
+  true
+rescue StandardError => e
+  puts "Error sending messages: #{e.message}"
+  false
+end
+
+def submit_task(connection_pool, messages, queue_url, sqs_client)
+  connection_pool.post do
+    messages_sent?(sqs_client, queue_url, messages)
+    puts "Messages Sent"
+  end
+end
+
+# Full example call:
+# Replace us-west-2 with the AWS Region you're using for Amazon SQS.
+def run_me
+  region = 'us-west-2'
+  queue_name = 'my-queue'
+  sqs_client = Aws::SQS::Client.new(region: region)
+
+  puts "Creating the queue named '#{queue_name}'..."
+
+  if queue_created?(sqs_client, queue_name)
+    puts 'Queue created.'
+  else
+    puts 'Queue not created.'
+  end
+  queue_url = sqs_client.get_queue_url(queue_name: queue_name).queue_url
+  messages = Array.new
+  connection_pool = Concurrent::RubyThreadPoolExecutor.new(max_threads: 20, min_threads: 1, max_queue: 100000)
+  10000.times do |index|
+    puts "hello - #{index} - #{index % 10}"
+    body = {
+      'MessageType' => 'PROCESS_TRUCK_LOCATION',
+      'Source' => "#{index} SQS updates",
+      'Content' => {
+        'TimeZone' => 'UTC'
+      }
+    }
+    msg = {
+      id: SecureRandom.uuid,
+      message_body: body.to_json
+    }
+=begin
+    RestClient::Request.execute(:method => :post, :url => "http://localhost:5001/process", :payload => body.to_json, :headers => HEADERS,  :timeout => 10, :open_timeout => 5) do |response, request, result|
+      process_http_response response
+    end
+=end
+    messages.push(msg)
+    if index % 10 == 9
+      submit_task(connection_pool, messages.clone, queue_url, sqs_client)
+      messages = Array.new
+    end
+  end
+  connection_pool.shutdown
+  connection_pool.wait_for_termination
+end
+
+def process_http_response(response)
+  case response.code
+  when 200
+    return "OK"
+  else
+    raise "Service did not return 200 OK response. #{response.code}"
+  end
+end
+
+run_me if $PROGRAM_NAME == __FILE__

data/lib/sqspoller/common/ring_buffer.rb

@@ -0,0 +1,74 @@
+module SqsPoller
+  module Common
+    class RingBuffer
+
+      def initialize(size)
+        @size = size
+        @start = 0
+        @count = 0
+        @buffer = Array.new(size)
+        @mutex = Mutex.new
+      end
+
+      def full?
+        @count == @size
+      end
+
+      def count
+        @count
+      end
+
+      def empty?
+        @count == 0
+      end
+
+      def push(value)
+        @mutex.synchronize do
+          stop = (@start + @count) % @size
+          @buffer[stop] = value
+          if full?
+            @start = (@start + 1) % @size
+          else
+            @count += 1
+          end
+          value
+        end
+      end
+
+      alias :<< :push
+
+      def shift
+        @mutex.synchronize do
+          remove_element
+        end
+      end
+
+      def flush
+        values = []
+        @mutex.synchronize do
+          until empty?
+            values << remove_element
+          end
+        end
+        values
+      end
+
+      def clear
+        @buffer = Array.new(@size)
+        @start = 0
+        @count = 0
+      end
+
+      private
+
+      def remove_element
+        return nil if empty?
+        value, @buffer[@start] = @buffer[@start], nil
+        @start = (@start + 1) % @size
+        @count -= 1
+        value
+      end
+
+    end
+  end
+end

data/lib/sqspoller/common/utils.rb

@@ -0,0 +1,47 @@
+require "logger"
+require "concurrent"
+require "net/http"
+require "aws-sdk"
+
+module SqsPoller
+  module Common
+    class Utils
+      class << self
+
+        def get_current_time_in_millis
+          current_time = Time.now
+          (current_time.to_f * 1000).to_i
+        end
+
+        def start_timer
+          SqsPoller::Common::Utils::Timer.new
+        end
+
+      end
+
+      class Timer
+        def initialize
+          start
+        end
+
+        def start_time
+          @start_time
+        end
+
+        def stop_time
+          @stop_time
+        end
+
+        def start
+          @start_time = SqsPoller::Common::Utils.get_current_time_in_millis
+        end
+
+        def stop
+          @stop_time = SqsPoller::Common::Utils.get_current_time_in_millis
+          @stop_time - @start_time
+        end
+
+      end
+    end
+  end
+end

data/lib/sqspoller/logger/logger.rb

@@ -0,0 +1,40 @@
+require 'logger'
+
+module SqsPoller
+  module Logger
+    class << self
+
+      @logger_file = nil
+      @logger = nil
+      @log_level
+
+      def get_new_logger (prog_name = nil, log_level = nil)
+        if @logger_file.nil?
+          logger = ::Logger.new(STDOUT)
+        else
+          logger = ::Logger.new(@logger_file)
+        end
+        @log_level = @log_level || ::Logger::ERROR
+        logger.level = log_level || @log_level
+        logger.progname = prog_name
+        logger
+      end
+
+      def set_log_level (log_level)
+        @log_level = log_level
+      end
+
+      def set_logger_file (logger_file)
+        @logger_file = logger_file
+      end
+
+      def get_logger
+        if @logger.nil?
+          @logger = get_new_logger
+        end
+        @logger
+      end
+
+    end
+  end
+end

data/lib/sqspoller/metrics/log_reporter.rb

@@ -0,0 +1,19 @@
+require 'ruby-metrics'
+require 'sqspoller/logger/logger'
+
+# Report local metrics in logs.
+module SqsPoller
+  module Metrics
+    class LogReporter
+
+      def initialize(options = {})
+        @logger = SqsPoller::Logger.get_new_logger(self.class.name)
+      end
+
+      def report(agent)
+        @logger.info(JSON.pretty_generate(agent.to_json))
+      end
+
+    end
+  end
+end

data/lib/sqspoller/metrics/queue_stats_reporter.rb

@@ -0,0 +1,42 @@
+require 'sqspoller/logger/logger'
+
+# This reports the Aws::SQS::QueuePoller stats for each queue and each worker in the log.
+module SqsPoller
+  module Metrics
+    class SqsStatsReporter
+      # Default reporting delay is 60 seconds
+      DEFAULT_REPORTING_DELAY = 60
+
+      def stop
+        @running = false
+      end
+
+      def initialize(options = {})
+        @running = true
+        @logger = SqsPoller::Logger.get_new_logger(self.class.name)
+
+        if options[:queue_controller] == nil
+          raise "Need an agent to report data from"
+        end
+
+        delay = options[:delay] || DEFAULT_REPORTING_DELAY
+        queue_controller = options[:queue_controller]
+
+        Thread.new {
+          while @running
+            sleep delay
+            next if queue_controller.pollers.size == 0
+            queue_controller.pollers.each do |queue_name, pollers|
+              pollers.each.with_index(1) do |poller, index|
+                stats = poller.get_poller_stats
+                next if stats == nil
+                @logger.info("Queue: #{queue_name}, Worker: #{index} started: #{stats.polling_started_at}, requests: #{stats.request_count}, messages: #{stats.received_message_count}, last-timestamp: #{stats.last_message_received_at}")
+              end
+            end
+          end
+        }
+      end
+
+    end
+  end
+end

data/lib/sqspoller/metrics/sqs_poller_metrics.rb

@@ -0,0 +1,109 @@
+require 'ruby-metrics'
+require 'new_relic/agent'
+require 'sqspoller/metrics/log_reporter'
+
+# collect the metrics around queue, message_type and message count.
+module SqsPoller
+  module Metrics
+
+    @new_relic_metrics_enabled = true
+    @duration_units = :milliseconds
+    @agent = nil
+
+    UNKNOWN_MESSAGE_TYPE = "UNKNOWN"
+    TASK_SUCCESS_STATUS = "success"
+    TASK_FAILURE_STATUS = "failed"
+    NEW_RELIC_CUSTOM_METRICS_PATH = "Custom/sqspoller/"
+
+    # @param [Object] task
+    def self.get_message_type(task)
+      message_type = UNKNOWN_MESSAGE_TYPE
+      begin
+        parsed_message = JSON.parse(task[:message].body, symbolize_names: true)
+        if parsed_message[:MessageType] != nil
+          message_type = parsed_message[:MessageType]
+        end
+      rescue Exception => e
+        # Ignoring When message Unknown. This happens when json is not parsable or message does not contains MessageType.
+      end
+      return message_type
+    end
+
+    def self.get_message_status(task_status)
+      task_status ? TASK_SUCCESS_STATUS : TASK_FAILURE_STATUS
+    end
+
+    def self.get_queue_name(task)
+      return task[:queue_name]
+    end
+
+    def self.get_metrics_recorder
+      return @agent
+    end
+
+    def self.get_qps
+      return @agent.timer(:total_messages)
+    end
+
+    def self.get_queue_wait_time(task, timer)
+      queue_time = task[:queue_time]
+      start_time = timer.start_time
+      return queue_time - start_time
+    end
+
+    def self.record(task, task_status, timer, elapsed_time)
+      return if @agent.nil?
+      queue_name = get_queue_name(task)
+      message_type = get_message_type(task)
+      message_status = get_message_status(task_status)
+      queue_wait_time = get_queue_wait_time(task, timer)
+      local_metrics(elapsed_time, message_status, message_type, queue_name, queue_wait_time)
+
+      new_relic_metrics(elapsed_time, message_status, message_type, queue_name, queue_wait_time) if @new_relic_metrics_enabled
+
+    end
+
+    def self.enable_new_relic_metrics
+      @new_relic_metrics_enabled = true
+    end
+
+    def self.disable_new_relic_metrics
+      @new_relic_metrics_enabled = false
+    end
+
+    private
+
+    def self.local_metrics(elapsed_time, message_status, message_type, queue_name, queue_wait_time)
+      @agent.timer(:total_messages, @duration_units).update(elapsed_time, @duration_units)
+      @agent.timer(queue_name.to_sym, @duration_units).update(elapsed_time, @duration_units)
+      @agent.timer("#{queue_name}_#{message_type}".to_sym, @duration_units).update(elapsed_time, @duration_units)
+      @agent.timer("#{queue_name}_#{message_type}_wait_time".to_sym, @duration_units).update(queue_wait_time, @duration_units)
+
+      @agent.counter("#{message_status}".to_sym).inc
+      @agent.counter("#{queue_name}_count".to_sym).inc
+      @agent.counter("#{queue_name}_#{message_status}".to_sym).inc
+      @agent.counter("#{queue_name}_#{message_type}_#{message_status}".to_sym).inc
+    end
+
+    def self.new_relic_metrics(elapsed_time, message_status, message_type, queue_name, queue_wait_time)
+      ::NewRelic::Agent.record_metric("#{NEW_RELIC_CUSTOM_METRICS_PATH}all", elapsed_time)
+      ::NewRelic::Agent.record_metric("#{NEW_RELIC_CUSTOM_METRICS_PATH}#{queue_name}", elapsed_time)
+      ::NewRelic::Agent.record_metric("#{NEW_RELIC_CUSTOM_METRICS_PATH}#{queue_name}_#{message_type}", elapsed_time)
+      ::NewRelic::Agent.record_metric("#{NEW_RELIC_CUSTOM_METRICS_PATH}#{queue_name}_#{message_type}_wait_time", queue_wait_time)
+
+      ::NewRelic::Agent.increment_metric("#{NEW_RELIC_CUSTOM_METRICS_PATH}count")
+      ::NewRelic::Agent.increment_metric("#{queue_name}_count")
+      ::NewRelic::Agent.increment_metric("#{queue_name}_#{message_status}_count")
+      ::NewRelic::Agent.increment_metric("#{queue_name}_#{message_type}_#{message_status}_count")
+    end
+
+    def self.start_metrics_agent
+      @agent = Metrics::Agent.new
+      @agent.report_to("log_reporter", SqsPoller::Metrics::LogReporter.new)
+      @agent.report_periodically
+    end
+
+  end
+
+end
+

data/lib/sqspoller/poll/queue_controller.rb

@@ -0,0 +1,141 @@
+require "concurrent"
+require "net/http"
+require "aws-sdk"
+require 'sqspoller/logger/logger'
+require "sqspoller/poll/queue_poller"
+require "sqspoller/metrics/queue_stats_reporter"
+
+# QueueController Responsible for controlling ThreadPools per sqs queue.
+# Each queue can configure max number of threads to fetch messages from SQS and queue it.
+# For each queue there will be threadpool created and managed :workers_pool
+# Each worker thread will create Aws::SQS::QueuePoller and poll messages from SQS.
+module SqsPoller
+  module Poller
+    class QueueController
+
+      private_class_method :new
+
+      def initialize(queues_config, task_queue, aws_config)
+        @queues_config = queues_config
+        @aws_config = aws_config
+        @task_queue = task_queue
+        @logger = SqsPoller::Logger.get_new_logger(self.class.name)
+        @counter = Concurrent::MutexAtomicFixnum.new(0)
+        @sqs_client = Aws::SQS::Client.new(:access_key_id => @aws_config[:access_key_id], :secret_access_key => @aws_config[:secret_access_key], :region => @aws_config[:region])
+        @started = true
+        # hash of threadpools by queue name
+        @pollers_thread_pool = {}
+        # hash of QueuePollers by queue name
+        @pollers = {}
+        @queue_urls = {}
+      end
+
+      def self.get
+        return @instance if @instance
+        raise "QueueController not yet started"
+      end
+
+      def self.delete_messages(queue, messages)
+        self.get.delete_messages(queue, messages)
+      end
+
+      def self.delete_message(queue, message)
+        self.get.delete_message(queue, message)
+      end
+
+      def delete_messages(queue, messages)
+        @sqs_client.delete_message_batch(
+          queue_url: get_queue_url(queue),
+          entries: messages.map { |msg|
+            { id: msg.message_id, receipt_handle: msg.receipt_handle }
+          }
+        )
+      end
+
+      def get_queue_url(queue)
+        @queue_urls[queue]
+      end
+
+      def delete_message(queue, message)
+        delete_msg = {
+          queue_url: get_queue_url(queue),
+          receipt_handle: message.receipt_handle,
+        }
+        @client.delete_message(delete_msg)
+      end
+
+      def started?
+        @started
+      end
+
+      def pollers
+        @pollers
+      end
+
+      def self.start (queues_config, task_queue, aws_config)
+        return @instance if @instance
+        @instance = new(queues_config, task_queue, aws_config)
+        @instance.start_queue_controller
+        @instance.start_poller_stats_reporter
+        @instance
+      end
+
+      def start_poller_stats_reporter
+        SqsPoller::Metrics::SqsStatsReporter.new(
+          {
+            :queue_controller => self
+          }
+        )
+      end
+
+      def start_queue_controller
+        @queues_config.keys.each { |queue|
+          queue_config = @queues_config[queue]
+          polling_threads = queue_config[:polling_threads]
+          if polling_threads == 0
+            @logger.info "Polling disabled for queue: #{queue}"
+            next
+          end
+          started = false
+          begin
+            sqs_queue_config = @sqs_client.get_queue_url(queue_name: queue)
+            started = start_pollers(polling_threads, queue, sqs_queue_config.queue_url, queue_config)
+            @queue_urls[queue] = sqs_queue_config.queue_url
+          rescue Exception => e
+            @logger.error "Failed to start Queue Pollers. Caught error: #{e.message}\n #{e.backtrace.join("\n")}"
+          end
+          unless started
+            @started = false
+            @logger.error "Failed to start Queue Pollers for the queue #{queue}"
+          end
+        }
+      end
+
+      private
+
+      def start_pollers(polling_threads, queue, queue_url, queue_config)
+        workers = Concurrent::RubyThreadPoolExecutor.new(max_threads: polling_threads, min_threads: 1, max_queue: polling_threads)
+        @pollers_thread_pool[queue] = workers
+        @pollers[queue] = []
+        polling_threads.times do |index|
+          worker = QueuePoller.new "QueuePoller-#{index}", queue, queue_config, @task_queue, @sqs_client, queue_url, @counter
+          @pollers[queue] << worker
+          workers.post do
+            @logger.info "Starting QueuePoller-#{queue}-#{index} object for queue: #{queue}"
+            loop do
+              begin
+                worker.run
+              rescue Exception => e
+                @logger.error "Poller killed for queue: #{e} and restarted"
+              end
+            end
+          end
+          sleep 0.01
+        end
+        all_workers_started = workers.length == polling_threads
+        all_workers_started
+      end
+
+    end
+  end
+end

data/lib/sqspoller/poll/queue_poller.rb

@@ -0,0 +1,73 @@
+require "concurrent"
+require "net/http"
+require "aws-sdk"
+require 'sqspoller/common/utils'
+require 'sqspoller/logger/logger'
+
+# QueuePoller responsible for polling batch message[s] using Aws::SQS::QueuePoller based on configured :max_number_of_messages and push to the task queue.
+# If there is no messages available SQS it will keep polling based on :wait_time_seconds until the message arrives.
+module SqsPoller
+  module Poller
+    class QueuePoller
+
+      DEFAULT_SQS_WAIT_TIME_SECONDS = 20
+      DEFAULT_SQS_MAX_NUMBER_OF_MESSAGES = 10
+
+      def initialize(worker_name, queue_name, queue_config, task_queue, sqs_client, queue_url, counter)
+        @worker_name = worker_name
+        @queue_name = queue_name
+        @queue_config = queue_config
+        @task_queue = task_queue
+        @wait_time_seconds = get_wait_time_seconds(@queue_config)
+        @max_number_of_messages = get_max_number_of_messages(@queue_config)
+        @sqs_client = sqs_client
+        @queue_url = queue_url
+        @logger = SqsPoller::Logger.get_new_logger("#{self.class.name}-#{@worker_name}")
+        @counter = counter
+      end
+
+      def get_poller_stats
+        @poller_stats
+      end
+
+      def run
+        poller = Aws::SQS::QueuePoller.new(@queue_url, { client: @sqs_client })
+        poller.before_request do |stats|
+          @poller_stats = stats
+        end
+        poller.poll(skip_delete: true, :wait_time_seconds => @wait_time_seconds, :max_number_of_messages => @max_number_of_messages) do |messages|
+          if @max_number_of_messages == 1
+            messages = [messages]
+          end
+          messages.each { |message|
+            task = {
+              :message => message,
+              :queue_name => @queue_name,
+              :queue_time => SqsPoller::Common::Utils.get_current_time_in_millis,
+              :index => @counter.increment
+            }
+            @task_queue.push task
+          }
+          @logger.info "Queued #{messages.size} messages from #{@queue_name}"
+        end
+      end
+
+      def get_wait_time_seconds(queue_config)
+        wait_time_seconds = queue_config[:wait_time_seconds]
+        unless wait_time_seconds && wait_time_seconds >= 0 && wait_time_seconds <= DEFAULT_SQS_WAIT_TIME_SECONDS
+          wait_time_seconds = DEFAULT_SQS_WAIT_TIME_SECONDS
+        end
+        wait_time_seconds
+      end
+
+      def get_max_number_of_messages(queue_config)
+        max_number_of_messages = queue_config[:max_number_of_messages]
+        unless max_number_of_messages && max_number_of_messages >= 0 && max_number_of_messages <= DEFAULT_SQS_MAX_NUMBER_OF_MESSAGES
+          max_number_of_messages = DEFAULT_SQS_MAX_NUMBER_OF_MESSAGES
+        end
+        max_number_of_messages
+      end
+
+    end
+  end
+end