Using concurrent-ruby for a 1.5x real-world speedup

class ChannelPipeline
  # Ruby isn't typed by default, so we can use a simple string as a marker for when processing is done.
  # All we need to do is ensure that it does not look like any data coming in the system
  CHANNEL_FINISH_TOKEN = '__finish__'

  def load_metrics
	# Create the channel which we are going to use to download files from s3 and make them available
    # on the local disk.
    # You want this to be more than 1. The 5 here is somewhat arbitrary. 
	# The ideal number depends on the relative speed of the processes
    s3_channel = Concurrent::Channel.new(capacity: 5)
    # Create the channel which is going to be doing the db insert.
    # In my case, this has to be single threaded
    db_channel = Concurrent::Channel.new(capacity: 1)
    # A single threaded channel will be used to communicate that the process is done, either due to
    # an error or with successful completion
    finish_channel = Concurrent::Channel.new(capacity: 1)

	# I'm going to admit that I found the naming in the library a little confusing. 
    # `Concurrent::Channel.new` creates something you use for communicating, but 
    # `Concurrent::Channel.go` creates a worker thread you can use for processing.
    # So I'm starting a channel, but also sending a channel as arguments. 
    start_s3_channel(s3_channel, db_channel, finish_channel)
    start_db_channel(db_channel, finish_channel)

    # On the main thread, list the files in s3 
    start_aws_file_list(s3_channel)

    Logger.info('Workers started, waiting on the finish token')
    success = ~finish_channel
	s3_channel.close
	db_channel.close
	finish_channel.close
    raise 'Unable to complete pipeline' unless success

    Logger.info('Pipeline complete')
  end

  def start_s3_channel(s3_channel, db_channel, finish_channel)
    # Start the worker that will perform the downloads
    Concurrent::Channel.go do
      # Loop on the channel. `~` is an alias for `take`. If there is nothing in the channel, 
      # this is a blocking action
      while (s3_channel_message = ~s3_channel)
        if s3_channel_message == CHANNEL_FINISH_TOKEN
          # We received the finish token, pass it forward then break the loop
  		  Logger.info('Completed downloads of files from s3')
          db_channel << CHANNEL_FINISH_TOKEN
          break
        else
		  # Pull the s3 object out of the message, download it, then publish a new message
          # to the db thread that it is available for loading
          s3_channel_message => { s3_object: }
		  local_file = S3Helper.download(s3_object)
          db_channel << { local_file: }
        end
      end
    rescue StandardError => e
      # Ruby threads don't bubble errors up to the main thread. I want to halt on failure, 
      # so catch it here then tell the main thread to stop execution. 
      Logger.error(e)
      finish_channel << false
    end
  end

  def start_db_channel(db_channel, finish_channel)
    # The structure of this is very similar to the structure above. Limiting the size of db_channel
    # is what gets the throttling behavior I wanted.
    Concurrent::Channel.go do
      while (db_channel_message = ~db_channel)
        if db_channel_message == CHANNEL_FINISH_TOKEN
          finish_channel << true
          break
        else
          db_channel_message => { local_file: }
		  Db::Loader.load(local_file)
        end
      end
    rescue StandardError => e
      Logger.error(e)
      finish_channel << false
    end
  end

  def start_aws_file_list(s3_channel)
    AwsUtil.list_objects(bucket:, prefix: "/") do |s3_object|
      s3_channel << { s3_object: }
    end

    # Push the finish token the let the channels know we've reached the end of the input
    s3_channel << CHANNEL_FINISH_TOKEN
  end
end

class ChannelPipeline
  def load_metrics
    AwsUtil.list_objects(bucket:, prefix: "/") do |s3_object|
      local_file = S3Helper.download(s3_object)
      Db::Loader.load(local_file)
    end
  end
end