Amazon SQS を使ったアプリケーションを本番で運用する際に考慮すべき基本的な 5 つのこと

require 'active_record'
require 'aws-sdk-s3'
require 'aws-sdk-sqs'

ActiveRecord::Base.establish_connection(
  adapter:  'mysql2',
  host:     ENV['MYSQL_HOST'],
  port:     ENV['MYSQL_PORT'],
  username: ENV['MYSQL_USERNAME'],
  password: ENV['MYSQL_PASSWORD'],
  database: ENV['MYSQL_DATABASE'],
)

# This table is created by the following DSL:
#   create_table(:s3_objects) do |t|
#     t.string :bucket
#     t.string :key
#     t.boolean :processed
#     t.index [:bucket, :key], name: "ux_bucket_key", unique: true
#   end
class S3Object < ActiveRecord::Base; end

class S3ObjectProcessor
  QUEUE_URL = ENV['QUEUE_URL']

  def run
    @stop = false
    trap(:SIGTERM) { @stop = true }

    poller = Aws::SQS::QueuePoller.new(QUEUE_URL)
    poller.before_request do |stats|
      throw :stop_polling if @stop
    end
    poller.poll do |msg, stats|
      JSON.parse(msg.body)['Records'].each do |record|
        bucket = record.dig('s3', 'bucket', 'name')
        key = record.dig('s3', 'object', 'key')
        s3_object = S3Object.create_or_find_by!(bucket: bucket, key: key)
        unless s3_object.processed?
          process(s3_object)
          s3_object.update!(processed: true)
        end
      rescue
      end
    end
  end

  def process(s3_object)
    Tempfile.open('s3_object') do |f|
      Aws::S3::Client.new.get_object(bucket: s3_object.bucket, key: s3_object.key) do |chunk, _|
        f.write(chunk)
      end
      f.flush
      f.rewind

      f.each do |line|
        # Do something
      end
    end
  end
end

S3ObjectProcessor.new.run

def run
  @stop = false
  trap(:SIGTERM) { @stop = true }

  poller = Aws::SQS::QueuePoller.new(QUEUE)
  poller.before_request do |stats|
    throw :stop_polling if @stop
  end
  poller.poll do |msg, stats|
    JSON.parse(msg.body)['Records'].each do |record|
      bucket = record.dig('s3', 'bucket', 'name')
      key = record.dig('s3', 'object', 'key')
      s3_object = S3Object.create_or_find_by!(bucket: bucket, key: key)
      unless s3_object.processed?
        process(s3_object)
        s3_object.update!(processed: true)
      end
    rescue
    end
  end
end

def run
  logger = Logger.new($stdout)

  @stop = false
  signals = Queue.new
  trap(:SIGTERM) { signals << :SIGTERM }
  Thread.new do
    sig = signals.pop
    logger.info "Received #{sig}, shutting down gracefully"
    @stop = true
  end

  poller = Aws::SQS::QueuePoller.new(QUEUE)
  poller.before_request do |stats|
    throw :stop_polling if @stop
  end
  poller.poll do |msg, stats|
    JSON.parse(msg.body)['Records'].each do |record|
      bucket = record.dig('s3', 'bucket', 'name')
      key = record.dig('s3', 'object', 'key')
      s3_object = S3Object.create_or_find_by!(bucket: bucket, key: key)
      if s3_object.processed?
        logger.info "Skip processing s3://#{bucket}/#{key} because it has already been processed"
        next
      end

      logger.info "Start processing s3://#{bucket}/#{key}"
      process(s3_object)
      s3_object.update!(processed: true)
      logger.info "Finished processing s3://#{bucket}/#{key}"
    rescue => e
      logger.error "Failed to process s3://#{bucket}/#{key}: #{e} (#{e.class})"
      logger.error e.backtrace.join("\n")
    end
  end
end

poller.poll do |msg, stats|
  JSON.parse(msg.body)['Records'].each do |record|
    bucket = record.dig('s3', 'bucket', 'name')
    key = record.dig('s3', 'object', 'key')
    s3_object = S3Object.create_or_find_by!(bucket: bucket, key: key)
    unless s3_object.processed?
      process(s3_object)
      s3_object.update!(processed: true)
    end
  rescue
  end
end

poller.poll do |msg, stats|
  JSON.parse(msg.body)['Records'].each do |record|
    bucket = record.dig('s3', 'bucket', 'name')
    key = record.dig('s3', 'object', 'key')
    begin
      S3Object.create!(bucket: bucket, key: key)
    rescue ActiveRecord::RecordNotUnique
    end

    S3Object.transaction do
      s3_object = S3Object.lock.find_by!(bucket: bucket, key: key)
      unless s3_object.processed?
        process(s3_object)
        s3_object.update!(processed: true)
      end
    end
  rescue
  end
end

poller.poll(skip_delete: true) do |msg, stats|
  skip_deletion = false
  JSON.parse(msg.body)['Records'].each do |record|
    bucket = record.dig('s3', 'bucket', 'name')
    key = record.dig('s3', 'object', 'key')
    begin
      S3Object.create!(bucket: bucket, key: key)
    rescue ActiveRecord::RecordNotUnique
    end

    skip_process = false
    s3_object = nil
    S3Object.transaction do
      s3_object = S3Object.lock.find_by!(bucket: bucket, key: key)
      if s3_object.processed?
        skip_process = true
      elsif s3_object.processing?
        skip_process = true
        skip_deletion = true
      else
        s3_object.processing!
      end
    end

    unless skip_process
      process(s3_object)
      s3_object.processed!
    end
  rescue
  end
  poller.delete_message(msg) unless skip_deletion
end

# This table is created by the following DSL:
#   create_table(:s3_objects) do |t|
#     t.string :bucket
#     t.string :key
#     t.integer :status, limit: 2, default: 0
#     t.index [:bucket, :key], name: "idx_key", unique: true
#   end
class S3Object < ActiveRecord::Base
  enum status: { unprocessed: 0, processing: 1, processed: 2 }
end

poller.poll(skip_delete: true) do |msg, stats|
  skip_deletion = false
  JSON.parse(msg.body)['Records'].each do |record|
    # snip

    skip_process = false
    s3_object = nil
    S3Object.transaction do
      s3_object = S3Object.lock.find_by!(bucket: bucket, key: key)
      if s3_object.processed?
        skip_process = true
      elsif s3_object.processing? && Time.now - s3_object.start_processing_at < PROCESS_TIMEOUT
        skip_process = true
        skip_deletion = true
      else
        s3_object.update!(status: :processing, start_processing_at: Time.now)
      end
    end

    # snip
  rescue
  end
  poller.delete_message(msg) unless skip_deletion
end

def process(s3_object)
  Tempfile.open('s3_object') do |f|
    Aws::S3::Client.new.get_object(bucket: s3_object.bucket, key: s3_object.key) do |chunk, _|
      f.write(chunk)
    end
    f.flush
    f.rewind

    f.each do |line|
      if @stop
        s3_object.unprocessed!
        break
      end
      # Do something
    end
  end
end

poller.poll(skip_delete: true) do |msg, stats|
  skip_deletion = false
  JSON.parse(msg.body)['Records'].each do |record|
    # snip
  rescue
  end
  poller.delete_message(msg) unless skip_deletion
end

poller.poll(skip_delete: true) do |msg, stats|
  skip_deletion = false
  JSON.parse(msg.body)['Records'].each do |record|
    # snip
  rescue
    skip_deletion = true
  end
  poller.delete_message(msg) unless skip_deletion
end

def process
  Retriable.retriable(on: RETRIABLE_ERRORS, on_retry: method(:log_retries)) do
    # Do something
  end
end

private

def log_retries(exception, try, elapsed_time, interval)
  return unless @logger
  @logger.warn("Failed to ... (#{try}th try): #{exception} (#{exception.class})")
end

require 'active_record'
require 'aws-sdk-s3'
require 'aws-sdk-sqs'
require 'rollbar'

ActiveRecord::Base.establish_connection(
  adapter:  'mysql2',
  host:     ENV['MYSQL_HOST'],
  port:     ENV['MYSQL_PORT'],
  username: ENV['MYSQL_USERNAME'],
  password: ENV['MYSQL_PASSWORD'],
  database: ENV['MYSQL_DATABASE'],
)

Rollbar.configure do |config|
  config.access_token = ENV['ROLLBAR_ACCESS_TOKEN']
  config.environment = ENV['ROLLBAR_ENV'] || 'development'
end

# This table is created by the following DSL:
#   create_table(:s3_objects) do |t|
#     t.string :bucket
#     t.string :key
#     t.integer :status, limit: 2, default: 0
#     t.timestamps
#     t.index [:bucket, :key], name: "ux_bucket_key", unique: true
#   end
class S3Object < ActiveRecord::Base
  enum status: { unprocessed: 0, processing: 1, processed: 2 }

  def self.uri(bucket, key)
    "s3://#{bucket}/#{key}"
  end

  def uri
    self.class.uri(bucket, key)
  end
end

class S3ObjectProcessor
  QUEUE_URL = ENV['QUEUE_URL']
  PROCESS_TIMEOUT = 60
  VISIBILITY_TIMEOUT = PROCESS_TIMEOUT * 2
  LINES_PER_HEARTBEAT = 10_000

  def initialize
    @logger = Logger.new($stdout)
  end

  def run
    @stop = false
    signals = Queue.new
    trap(:SIGTERM) { signals << :SIGTERM }
    Thread.new do
      sig = signals.pop
      @logger.info "Received #{sig}, shutting down gracefully"
      @stop = true
    end

    poller.poll(skip_delete: true, visibility_timeout: VISIBILITY_TIMEOUT) do |msg, stats|
      skip_deletion = false
      s3_object_uris = []
      records = JSON.parse(msg.body)['Records']
      if records.nil?
        @logger.warn 'Received a message without the field Records'
      else
        records.each do |record|
          bucket = record.dig('s3', 'bucket', 'name')
          key = record.dig('s3', 'object', 'key')
          if process_record(msg, bucket, key)
            s3_object_uris << S3Object.uri(bucket, key)
          else
            skip_deletion = true
          end
        rescue => e
          skip_deletion = true
          @logger.error "Failed to process #{S3Object.uri(bucket, key)}: #{e} (#{e.class})"
          @logger.error e.backtrace.join("\n")
          Rollbar.error(e, bucket: bucket, key: key)
        end
      end

      unless skip_deletion
        poller.delete_message(msg)
        @logger.info "Deleted the message for #{s3_object_uris.join(', ')}"
      end
    rescue => e
      @logger.error "Unknown error: #{e} (#{e.class})"
      @logger.error e.backtrace.join("\n")
      Rollbar.error(e)
    end
  end

  private

  def poller
    @poller ||= Aws::SQS::QueuePoller.new(QUEUE_URL).tap do |p|
      p.before_request do |stats|
        throw :stop_polling if @stop
      end
    end
  end

  def process_record(msg, bucket, key)
    begin
      S3Object.create!(bucket: bucket, key: key)
    rescue ActiveRecord::RecordNotUnique
    end

    skip_process = false
    delete_message = true
    s3_object = nil
    S3Object.transaction do
      s3_object = S3Object.lock.find_by!(bucket: bucket, key: key)
      if s3_object.processed?
        skip_process = true
      elsif s3_object.processing? && Time.now - s3_object.updated_at < PROCESS_TIMEOUT
        skip_process = true
        delete_message = false
      else
        s3_object.update!(status: :processing, updated_at: Time.now)
      end
    end

    if skip_process
      @logger.info "Skip processing #{s3_object.uri} because it has already been processed"
      return delete_message
    end

    @logger.info "Start processing #{s3_object.uri}"
    process_s3_object(s3_object) do
      if @stop
        s3_object.unprocessed!
        @logger.info "Changed the status to 'unprocessed' for #{s3_object.uri}"
        return false
      end

      s3_object.touch
      poller.change_message_visibility_timeout(msg, VISIBILITY_TIMEOUT)
      @logger.info "Updated the message visibility timeout for #{s3_object.uri}"
    end

    s3_object.processed!
    @logger.info "Finished processing #{s3_object.uri}"
  end

  def process_s3_object(s3_object)
    Tempfile.open('s3_object') do |f|
      Aws::S3::Client.new.get_object(bucket: s3_object.bucket, key: s3_object.key) do |chunk, _|
        f.write(chunk)
      end
      f.flush
      f.rewind

      f.each.with_index(1) do |line, i|
        yield if (i % LINES_PER_HEARTBEAT).zero?

        # Do something
      rescue ArgumentError => e
        if e.message != 'invalid byte sequence in UTF-8'
          raise
        end
        @logger.warn "Skip the line at #{i} due to invalid byte sequence"
      end
    end
  end
end

S3ObjectProcessor.new.run