dynamo-autoscale 0.1

data/config/environment/common.rb

@@ -0,0 +1,114 @@
+require 'logger'
+require 'time'
+require 'csv'
+require 'tempfile'
+require 'aws-sdk'
+require 'active_support/all'
+require 'rbtree'
+require 'colored'
+
+require_relative '../../lib/dynamo-autoscale/logger'
+require_relative '../../lib/dynamo-autoscale/poller'
+
+module DynamoAutoscale
+  include DynamoAutoscale::Logger
+
+  DEFAULT_AWS_REGION = 'us-east-1'
+
+  def self.root
+    File.expand_path(File.join(File.dirname(__FILE__), '..', '..'))
+  end
+
+  def self.data_dir
+    File.join(self.root, 'data')
+  end
+
+  def self.env
+    ENV['RACK_ENV'] || 'development'
+  end
+
+  def self.with_config name, opts ={}
+    path = nil
+
+    if opts[:absolute]
+      path = name
+    else
+      path = File.join(root, 'config', "#{name}.yml")
+    end
+
+    conf = YAML.load_file(path)[env]
+    yield conf if block_given?
+    conf
+  end
+
+  def self.require_all path
+    Dir[File.join(root, path, '*.rb')].each { |file| require file }
+  end
+
+  def self.dispatcher= new_dispatcher
+    @@dispatcher = new_dispatcher
+  end
+
+  def self.dispatcher
+    @@dispatcher ||= Dispatcher.new
+  end
+
+  def self.poller= new_poller
+    @@poller = new_poller
+  end
+
+  def self.poller
+    @@poller ||= LocalDataPoll.new
+  end
+
+  def self.actioner_class= klass
+    @@actioner_class = klass
+  end
+
+  def self.actioner_class
+    @@actioner_class ||= LocalActioner
+  end
+
+  def self.actioner_opts= new_opts
+    @@actioner_opts = new_opts
+  end
+
+  def self.actioner_opts
+    @@actioner_opts ||= {}
+  end
+
+  def self.actioners
+    @@actioners ||= Hash.new { |h, k| h[k] = actioner_class.new(k, actioner_opts) }
+  end
+
+  def self.tables= new_tables
+    @@tables = new_tables
+  end
+
+  def self.tables
+    @@tables ||= Hash.new { |h, k| h[k] = TableTracker.new(k) }
+  end
+
+  def self.current_table= new_current_table
+    @@current_table = new_current_table
+  end
+
+  def self.current_table
+    @@current_table
+  end
+
+  def self.rules= new_rules
+    @@rules = new_rules
+  end
+
+  def self.rules
+    @@rules ||= RuleSet.new
+  end
+end
+
+DynamoAutoscale.require_all 'lib/dynamo-autoscale'
+DynamoAutoscale.require_all 'lib/dynamo-autoscale/ext/**'
+
+Dir[File.join(DynamoAutoscale.root, 'config', 'services', '*.rb')].each do |path|
+  load path
+end

data/config/environment/console.rb

@@ -0,0 +1,2 @@
+require_relative 'common'
+require 'pp'

data/config/environment/test.rb

@@ -0,0 +1,3 @@
+ENV['RACK_ENV'] = "test"
+require_relative 'common'
+require 'timecop'

data/config/logger.yml

@@ -0,0 +1,11 @@
+test:
+  :level: INFO
+  :sync:  false
+
+development:
+  :level: INFO
+  :sync:  true
+
+production:
+  :level: INFO
+  :sync:  true

data/config/services/aws.rb

@@ -0,0 +1,20 @@
+config_location = nil
+
+if File.exists? './aws.yml'
+  config_location = './aws.yml'
+elsif ENV['AWS_CONFIG'] and File.exists? ENV['AWS_CONFIG']
+  config_location = ENV['AWS_CONFIG']
+elsif File.exists?(File.join(DynamoAutoscale.root, 'config', 'aws.yml'))
+  config_location = File.join(DynamoAutoscale.root, 'config', 'aws.yml')
+end
+
+if config_location.nil?
+  STDERR.puts "Could not load AWS configuration. Searched in: ./aws.yml and " +
+    "ENV['AWS_CONFIG']"
+
+  exit 1
+end
+
+DynamoAutoscale.with_config(config_location, absolute: true) do |config|
+  AWS.config(config)
+end

data/config/services/logger.rb

@@ -0,0 +1,35 @@
+DynamoAutoscale.with_config 'logger' do |config|
+  if config[:sync]
+    STDOUT.sync = true
+    STDERR.sync = true
+  end
+
+  if config[:log_to]
+    STDOUT.reopen(config[:log_to])
+    STDERR.reopen(config[:log_to])
+  end
+
+  DynamoAutoscale::Logger.logger = ::Logger.new(STDOUT)
+
+  if ENV['PRETTY_LOG']
+    DynamoAutoscale::Logger.logger.formatter = DynamoAutoscale::PrettyFormatter.new
+  else
+    DynamoAutoscale::Logger.logger.formatter = Logger::Formatter.new
+  end
+
+  if ENV['DEBUG']
+    DynamoAutoscale::Logger.logger.level = ::Logger::DEBUG
+  elsif config[:level]
+    DynamoAutoscale::Logger.logger.level = ::Logger.const_get(config[:level])
+  end
+
+  if ENV['SILENT'] # or DynamoAutoscale.env == "test"
+    DynamoAutoscale::Logger.logger.level = ::Logger::FATAL
+  end
+end
+
+if ENV['DEBUG']
+  AWS.config({
+    logger: DynamoAutoscale::Logger.logger,
+  })
+end

data/dynamo-autoscale.gemspec

@@ -0,0 +1,29 @@
+require 'date'
+require './lib/dynamo-autoscale/version'
+
+Gem::Specification.new do |gem|
+  gem.name    = 'dynamo-autoscale'
+  gem.version = DynamoAutoscale::VERSION
+  gem.date    = Date.today.to_s
+
+  gem.summary = "Autoscaling for DynamoDB provisioned throughputs."
+  gem.description = "Will automatically monitor DynamoDB tables and scale them based on rules."
+
+  gem.authors  = ['InvisibleHand']
+  gem.email    = 'developers@getinvisiblehand.com'
+  gem.homepage = 'http://github.com/invisiblehand/dynamo-autoscale'
+
+  gem.bindir      = ['bin']
+  gem.executables = ['dynamo-autoscale']
+
+  gem.license  = 'MIT'
+
+  gem.add_dependency 'aws-sdk'
+  gem.add_dependency 'rbtree'
+  gem.add_dependency 'ruby-prof'
+  gem.add_dependency 'colored'
+  gem.add_dependency 'activesupport'
+
+  # ensure the gem is built out of versioned files
+  gem.files = `git ls-files -z`.split("\0")
+end

data/lib/dynamo-autoscale/actioner.rb

@@ -0,0 +1,265 @@
+module DynamoAutoscale
+  class Actioner
+    include DynamoAutoscale::Logger
+
+    def self.minimum_throughput
+      @minimum_throughput ||= 10
+    end
+
+    def self.minimum_throughput= new_minimum_throughput
+      @minimum_throughput = new_minimum_throughput
+    end
+
+    def self.maximum_throughput
+      @maximum_throughput ||= 20000
+    end
+
+    def self.maximum_throughput= new_maximum_throughput
+      @maximum_throughput = new_maximum_throughput
+    end
+
+    attr_accessor :table, :upscales, :downscales
+
+    def initialize table, opts = {}
+      @table            = table
+      @downscales       = 0
+      @upscales         = 0
+      @provisioned      = { reads: RBTree.new, writes: RBTree.new }
+      @pending          = { reads: nil, writes: nil }
+      @last_action      = Time.now.utc
+      @last_scale_check = Time.now.utc
+      @downscale_warn   = false
+      @opts             = opts
+    end
+
+    def provisioned_for metric
+      @provisioned[normalize_metric(metric)]
+    end
+
+    def provisioned_writes
+      @provisioned[:writes]
+    end
+
+    def provisioned_reads
+      @provisioned[:reads]
+    end
+
+    def check_day_reset!
+      now = Time.now.utc
+
+      if now >= (check = (@last_scale_check + 1.day).midnight)
+        logger.info "[scales] New day! Reset scaling counts back to 0."
+        logger.debug "[scales] now: #{now}, comp: #{check}"
+
+        if @downscales < 4
+          logger.warn "[scales] Unused downscales. Used: #{@downscales}"
+        end
+
+        @upscales   = 0
+        @downscales = 0
+        @downscale_warn    = false
+      end
+
+      @last_scale_check = now
+    end
+
+    def upscales
+      check_day_reset!
+      @upscales
+    end
+
+    def downscales new_val = nil
+      check_day_reset!
+      @downscales
+    end
+
+    def set metric, to
+      check_day_reset!
+
+      metric = normalize_metric(metric)
+      ptime, _ = provisioned_for(metric).last
+
+      if ptime and ptime > 2.minutes.ago
+        logger.warn "[actioner] Attempted to scale the same metric more than " +
+          "once in a 2 minute window. Disallowing."
+        return false
+      end
+
+      from = table.last_provisioned_for(metric)
+
+      if from and to > (from * 2)
+        to = from * 2
+
+        logger.warn "[#{metric}] Attempted to scale up " +
+          "more than allowed. Capped scale to #{to}."
+      end
+
+      if to < Actioner.minimum_throughput
+        to = Actioner.minimum_throughput
+
+        logger.warn "[#{metric}] Attempted to scale down to " +
+          "less than minimum throughput. Capped scale to #{to}."
+      end
+
+      if to > Actioner.maximum_throughput
+        to = Actioner.maximum_throughput
+
+        logger.warn "[#{metric}] Attempted to scale up to " +
+          "greater than maximum throughput. Capped scale to #{to}."
+      end
+
+      if from and from == to
+        logger.info "[#{metric}] Attempted to scale to same value. Ignoring..."
+        return false
+      end
+
+      if from and from > to
+        downscale metric, from, to
+      else
+        upscale metric, from, to
+      end
+    end
+
+    def upscale metric, from, to
+      logger.info "[#{metric}][scaling up] " +
+        "#{from ? from.round(2) : "Unknown"} -> #{to.round(2)}"
+
+
+      # Because upscales are not limited, we don't need to queue this operation.
+      if result = scale(metric, to)
+        @provisioned[metric][Time.now.utc] = to
+        @upscales += 1
+      end
+
+      return result
+    end
+
+    def downscale metric, from, to
+      if @downscales >= 4
+        unless @downscale_warn
+          @downscale_warn = true
+          logger.warn "[#{metric.to_s.ljust(6)}][scaling failed]" +
+            " Hit upper limit of downward scales per day."
+        end
+
+        return false
+      end
+
+      if @pending[metric]
+        previous_pending = @pending[metric].last
+        logger.info "[#{metric}][scaling down] " +
+          "#{previous_pending} -> #{to.round(2)} (overwritten pending)"
+      else
+        logger.info "[#{metric}][scaling down] " +
+          "#{from ? from.round(2) : "Unknown"} -> #{to.round(2)}"
+      end
+
+      queue_operation! metric, to
+    end
+
+    def queue_operation! metric, value
+      if @pending[metric]
+        logger.debug "[#{metric}] Overwriting pending op with #{value.round(2)}"
+      end
+
+      @pending[metric] = [Time.now.utc, value]
+
+      try_flush!
+    end
+
+    def try_flush!
+      if should_flush?
+        if flush_operations!
+          @downscales += 1
+          @last_action = Time.now.utc
+          return true
+        else
+          return false
+        end
+      else
+        return false
+      end
+    end
+
+    def flush_operations!
+      result = nil
+
+      if @pending[:writes] and @pending[:reads]
+        _, wvalue = @pending[:writes]
+        _, rvalue = @pending[:reads]
+
+        if result = scale_both(rvalue, wvalue)
+          @provisioned[:writes][Time.now.utc] = wvalue
+          @provisioned[:reads][Time.now.utc] = rvalue
+
+          @pending[:writes] = nil
+          @pending[:reads] = nil
+        end
+      elsif @pending[:writes]
+        time, value = @pending[:writes]
+
+        if result = scale(:writes, value)
+          @provisioned[:writes][Time.now.utc] = value
+
+          @pending[:writes] = nil
+        end
+      elsif @pending[:reads]
+        time, value = @pending[:reads]
+
+        if result = scale(:reads, value)
+          @provisioned[:reads][Time.now.utc] = value
+          @pending[:reads] = nil
+        end
+      end
+
+      logger.info "[flush] All pending downscales have been flushed."
+      return result
+    end
+
+    def should_flush?
+      if @opts[:group_downscales].nil?
+        logger.info "[flush] Downscales are not being grouped. Should flush."
+        return true
+      end
+
+      if @pending[:reads] and @pending[:writes]
+        logger.info "[flush] Both a read and a write are pending. Should flush."
+        return true
+      end
+
+      now = Time.now.utc
+
+      # I know what you're thinking. How would the last action ever be in the
+      # future? Locally, we use Timecop to fake out the time. Unfortunately it
+      # doesn't kick in until after the first data point, so when this object is
+      # created the @last_action is set to Time.now.utc, then the time gets
+      # rolled back, causing the last action to be in the future. This hack
+      # fixes that.
+      @last_action = now if @last_action > now
+
+      if (@opts[:flush_after] and @last_action and
+        (now > @last_action + @opts[:flush_after]))
+
+        logger.info "[flush] Flush timeout of #{@opts[:flush_after]} reached."
+        return true
+      end
+
+      logger.info "[flush] Flushing conditions not met. Pending operations: " +
+        "#{@pending[:reads] ? "1 read" : "no reads"}, " +
+        "#{@pending[:writes] ? "1 write" : "no writes"}"
+
+      return false
+    end
+
+    private
+
+    def normalize_metric metric
+      case metric
+      when :reads, :provisioned_reads, :consumed_reads
+        :reads
+      when :writes, :provisioned_writes, :consumed_writes
+        :writes
+      end
+    end
+  end
+end

data/lib/dynamo-autoscale/cw_poller.rb

@@ -0,0 +1,49 @@
+module DynamoAutoscale
+  class CWPoller < Poller
+    INTERVAL = 1.minute
+
+    def poll tables, &block
+      if tables.nil?
+        tables = AWS::DynamoDB.new.tables.to_a.map(&:name)
+      end
+
+      loop do
+        # Sleep until the next interval occurrs. This calculation ensures that
+        # polling always happens on interval boundaries regardless of how long
+        # polling takes.
+        sleep_duration = INTERVAL - ((Time.now.to_i + INTERVAL) % INTERVAL)
+        logger.debug "Sleeping for #{sleep_duration} seconds..."
+        sleep(sleep_duration)
+
+        do_poll(tables, &block)
+      end
+    end
+
+    def do_poll tables, &block
+      logger.debug "Beginning CloudWatch poll..."
+      now = Time.now
+
+      tables.each do |table|
+        # This code will dispatch a message to the listening table that looks
+        # like this:
+        #
+        #   {
+        #     :consumed_reads=>{
+        #       2013-06-19 12:22:00 UTC=>2.343117697349672
+        #     },
+        #     :consumed_writes=>{
+        #       2013-06-19 12:22:00 UTC=>3.0288461538461537
+        #     }
+        #   }
+        #
+        # There may also be :provisioned_reads and :provisioned_writes
+        # depending on how the CloudWatch API feels.
+        block.call(table, Metrics.all_metrics(table, {
+          period:     5.minutes,
+          start_time: now - 20.minutes,
+          end_time:   now,
+        }))
+      end
+    end
+  end
+end

data/lib/dynamo-autoscale/dispatcher.rb

@@ -0,0 +1,39 @@
+module DynamoAutoscale
+  class Dispatcher
+    def initialize
+      @last_check = {}
+    end
+
+    def dispatch table, time, datum, &block
+      DynamoAutoscale.current_table = table
+      logger.debug "#{time}: Dispatching to #{table.name} with data: #{datum}"
+
+      if datum[:provisioned_reads] and (datum[:consumed_reads] > datum[:provisioned_reads])
+        lost_reads = datum[:consumed_reads] - datum[:provisioned_reads]
+
+        logger.warn "[reads] Lost units: #{lost_reads} " +
+          "(#{datum[:consumed_reads]} - #{datum[:provisioned_reads]})"
+      end
+
+      if datum[:provisioned_writes] and (datum[:consumed_writes] > datum[:provisioned_writes])
+        lost_writes = datum[:consumed_writes] - datum[:provisioned_writes]
+
+        logger.warn "[writes] Lost units: #{lost_writes} " +
+          "(#{datum[:consumed_writes]} - #{datum[:provisioned_writes]})"
+      end
+
+      table.tick(time, datum)
+      block.call(table, time, datum) if block
+
+      if @last_check[table.name].nil? or @last_check[table.name] < time
+        DynamoAutoscale.rules.test(table)
+        @last_check[table.name] = time
+      else
+        logger.debug "#{table.name}: Skipped rule check, already checked for " +
+          "a later data point."
+      end
+
+      DynamoAutoscale.current_table = nil
+    end
+  end
+end

data/lib/dynamo-autoscale/dynamo_actioner.rb

@@ -0,0 +1,59 @@
+module DynamoAutoscale
+  class DynamoActioner < Actioner
+    def scale metric, value
+      aws_throughput_key = case metric
+      when :reads
+        :read_capacity_units
+      when :writes
+        :write_capacity_units
+      end
+
+      dynamo_scale(table, aws_throughput_key => metric)
+    end
+
+    def scale_both reads, writes
+      dynamo_scale(table, {
+        read_capacity_units: reads, write_capacity_units: writes
+      })
+    end
+
+    private
+
+    def self.dynamo_scale opts
+      aws_table = AWS::DynamoDB.new.tables[table.name]
+
+      if aws_table.status == :updating
+        logger.warn "Cannot scale throughputs. Table is updating."
+        return false
+      end
+
+      aws_table.provision_throughput(opts)
+      return true
+    rescue AWS::DynamoDB::Errors::ValidationException => e
+      # When you try to set throughput to a negative value or the same value it
+      # was previously you get this.
+      logger.warn "[#{e.class}] #{e.message}"
+      return false
+    rescue AWS::DynamoDB::Errors::ResourceInUseException => e
+      # When you try to update a table that is being updated you get this.
+      logger.warn "[#{e.class}] #{e.message}"
+      return false
+    rescue AWS::DynamoDB::Errors::LimitExceededException => e
+      # When you try to increase throughput greater than 2x or you try to
+      # decrease more than 4 times per day you get this.
+
+      aws_description = self.describe_table(table)
+      decreases_today = aws_description[:provisioned_throughput][:number_of_decreases_today]
+
+      downscales(table, decreases_today)
+      logger.warn "[#{e.class}] #{e.message}"
+      return false
+    end
+
+    def self.describe_table
+      data = AWS::DynamoDB::ClientV2.new.describe_table(table_name: table.name)
+
+      data[:table]
+    end
+  end
+end

data/lib/dynamo-autoscale/ext/active_support/duration.rb

@@ -0,0 +1,7 @@
+module ActiveSupport
+  class Duration
+    def inspect
+      "#{self.to_i}.seconds"
+    end
+  end
+end

data/lib/dynamo-autoscale/local_actioner.rb

@@ -0,0 +1,39 @@
+module DynamoAutoscale
+  class LocalActioner < Actioner
+    # Dummy scaling method.
+    def scale metric, value
+      return true
+    end
+
+    def scale_both reads, writes
+      return true
+    end
+
+    # These filters use the arrays inside the local actioner to fake the
+    # provisioned reads and writes when the local data enters the system. It
+    # makes it look like we're actually modifying the provisioned numbers.
+    def self.faux_provisioning_filters
+      [Proc.new do |table, time, datum|
+        actioner = DynamoAutoscale.actioners[table]
+
+        actioner.provisioned_reads.reverse_each do |rtime, reads|
+          logger.debug "Checking if #{time} > #{rtime}"
+          if time > rtime
+            logger.debug "[filter] Faked provisioned_reads to be #{reads} at #{time}"
+            datum[:provisioned_reads] = reads
+            break
+          end
+        end
+
+        actioner.provisioned_writes.reverse_each do |wtime, writes|
+          logger.debug "Checking if #{time} > #{wtime}"
+          if time > wtime
+            logger.debug "[filter] Faked provisioned_writes to be #{writes} at #{time}"
+            datum[:provisioned_writes] = writes
+            break
+          end
+        end
+      end]
+    end
+  end
+end