Merge pull request #8 from scribd/jimp/issue/5

feat: Add support for the 4 permutations of Elasticache redis endpoints

Author: jim80net

README.md

@@ -24,7 +24,7 @@ The following parameters need to be supplied to the lambda function.
 
 Environment Variable | Type   | Description                                                   | Required | Default
 ---------------------|--------|---------------------------------------------------------------|----------|---------
-REDIS_HOST           | string | FQDN of the elasticache redis endpoint                        | yes      | -
+REDIS_HOST           | string | FQDN or URI of the elasticache redis endpoint                 | yes      | -
 DATADOG_API_KEY      | string | Datadog API Key                                               | no       | -
 DATADOG_APP_KEY      | string | Datadog API Key                                               | no       | -
 NAMESPACE            | string | "namespace" tag to apply to Datadog metric                    | yes      | -
@@ -50,7 +50,7 @@ DO NOT use this script in production environments, as it will CPU thrash the tar
 
 Environment Variable | Type   | Description                                                   | Required | Default
 ---------------------|--------|---------------------------------------------------------------|----------|---------
-REDIS_HOST           | string | FQDN of the elasticache redis endpoint                        | yes      | -
+REDIS_HOST           | string | FQDN or URI of the elasticache redis endpoint                 | yes      | -
 
 
 # Requirements

inject_slow_query.rb

@@ -5,14 +5,14 @@
 
 require 'logger'
 require 'redis'
+require_relative 'lib/slowlog_check'
 
 LOGGER = Logger.new($stdout)
 LOGGER.level = Logger::WARN
 
-REDIS = Redis.new(
-  host: ENV.fetch('REDIS_HOST'),
-  ssl: :true
-  )
+REDIS = SlowlogCheck::Redis.new(
+  host: ENV.fetch('REDIS_HOST')
+).redis
 
 if ARGV[0].nil?
   raise "Specify milliseconds to inject as the first positional argument to `#{__FILE__}`"

lambda_function.rb

@@ -3,13 +3,11 @@
 
 require 'logger'
 require 'date'
-require 'redis'
 require 'dogapi'
 require_relative 'lib/slowlog_check'
 
 LOGGER = Logger.new($stdout)
 LOGGER.level = Logger::INFO
-LOGGER.freeze
 
 def event_time
   # DateTime because Time does not natively parse AWS CloudWatch Event time
@@ -24,6 +22,7 @@ def log_context
   LOGGER.info "Event time: #{event_time}."
 end
 
+
 def lambda_handler(event: {}, context: {})
   @event = event
   log_context
@@ -49,10 +48,9 @@ def lambda_handler(event: {}, context: {})
         ENV.fetch('DATADOG_API_KEY'),
         ENV.fetch('DATADOG_APP_KEY')
       ),
-      redis: Redis.new(
-        host: ENV.fetch('REDIS_HOST'),
-        ssl: :true
-      ),
+      redis: {
+        host: ENV.fetch('REDIS_HOST')
+      },
       namespace: ENV.fetch('NAMESPACE'),
       env: ENV.fetch('ENV'),
       metricname: ENV.fetch('METRICNAME', 'elasticache.slowlog')

lib/slowlog_check.rb

@@ -2,27 +2,18 @@
 require 'logger'
 
 class SlowlogCheck
+  require_relative 'slowlog_check/redis'
+
   ::LOGGER ||= ::Logger.new($stdout)
-  MAXLENGTH = 1048576 #255 levels of recursion for #
 
   def initialize(params = {})
     @ddog = params.fetch(:ddog)
-    @redis = params.fetch(:redis)
+    @redis= SlowlogCheck::Redis.new(params.fetch(:redis))
     @metricname = params.fetch(:metricname)
     @namespace = params.fetch(:namespace)
     @env = params.fetch(:env)
   end
 
-  def replication_group
-    host = @redis.connection.fetch(:host)
-    matches = /\w\.(?<replication_group>[\w-]+)\.\w+\.\w+\.cache\.amazonaws\.com/.match(host)
-    if matches
-      matches[:replication_group]
-    else
-      raise "Unable to parse REDIS_HOST. Is #{host} a valid elasticache endpoint?"
-    end
-  end
-
   def status_or_error(resp)
     return resp[1].fetch("status") if resp[1].key?("status")
     return resp[1].fetch("errors") if resp[1].key?("errors")
@@ -31,7 +22,7 @@ def status_or_error(resp)
 
   def last_datadog_metrics_submitted_by_me_in_the_last_2_hours
     resp = @ddog.get_points(
-      "#{@metricname}.95percentile{replication_group:#{replication_group}}",
+      "#{@metricname}.95percentile{replication_group:#{@redis.replication_group}}",
       Time.now - 7200,
       Time.now
     )
@@ -114,17 +105,6 @@ def add_metric_to_bucket(prior, new)
     }
   end
 
-  def did_i_get_it_all?(slowlog)
-    slowlog[-1][0] == 0
-  end
-
-  def redis_slowlog(length = 128)
-    resp = @redis.slowlog('get', length)
-
-    return resp if length > MAXLENGTH
-    return resp if did_i_get_it_all?(resp)
-    return redis_slowlog(length * 2)
-  end
 
   def empty_values
     {
@@ -169,7 +149,7 @@ def pad_results_with_zero(report)
 
   def slowlogs_by_flush_interval
     result = reporting_interval
-    redis_slowlog.each do |slowlog|
+    @redis.slowlog.each do |slowlog|
       time = slowlog_time(slowlog)
       break if minute_precision(time) <= minute_precision(last_time_submitted)
 
@@ -211,8 +191,8 @@ def slowlogs_by_flush_interval
 
   def default_tags
     {
-      replication_group: replication_group,
-      service: replication_group,
+      replication_group: @redis.replication_group,
+      service: @redis.replication_group,
       namespace: @namespace,
       aws: 'true',
       env: @env
@@ -224,7 +204,7 @@ def emit_point(params)
     type = params.fetch(:type, 'gauge')
     interval = params.fetch(:interval, 60)
     points = params.fetch(:points)
-    host = params.fetch(:host, replication_group)
+    host = params.fetch(:host, @redis.replication_group)
     tags = params.fetch(:tags, default_tags)
 
     LOGGER.info "Sending slowlog entry: #{metric}: #{points.first[1]}µs executing #{tags[:command]} at #{points.first[0]}."
@@ -238,7 +218,7 @@ def emit_point(params)
         tags: tags
       }
     )
-    raise "Error submitting metric for #{replication_group}" unless status_or_error(resp) == "ok"
+    raise "Error submitting metric for #{@redis.replication_group}" unless status_or_error(resp) == "ok"
 
     # Sigh. After doing all the work to pass around Time objects, dogapi-rb changes this to an integer.
     @last_time_submitted = Time.at(points.first[0])

lib/slowlog_check/redis.rb

@@ -0,0 +1,98 @@
+require 'redis'
+require 'uri'
+
+class SlowlogCheck::Redis
+  MAXLENGTH = 1048576 #255 levels of recursion for #
+
+  def initialize(opts)
+    @host = opts[:host]
+  end
+
+  def params
+    if cluster_mode_enabled?
+      {
+        cluster: [uri],
+        port: port,
+        ssl: tls_mode?
+      }
+    else
+      {
+        host: hostname,
+        port: port,
+        ssl: tls_mode?
+      }
+    end
+  end
+
+  def redis
+    @redis ||= Redis.new(params)
+  end
+
+  def replication_group
+    if tls_mode?
+      matches[:second]
+    else
+      matches[:first]
+    end
+  end
+
+
+  def slowlog(length = 128)
+    resp = redis.slowlog('get', length)
+
+    return resp if length > MAXLENGTH
+    return resp if did_i_get_it_all?(resp)
+    return slowlog(length * 2)
+  end
+
+  private
+
+  def cluster_mode_enabled?
+    if tls_mode?
+      matches[:first] == 'clustercfg'
+    else
+      matches[:third] == ''
+    end
+  end
+
+  def did_i_get_it_all?(slowlog)
+    slowlog[-1][0] == 0
+  end
+
+  def hostname
+    URI.parse(@host).hostname or
+      @host
+  end
+
+  def matches
+    redis_uri_regex.match(@host)
+  end
+
+  def port
+    regex_port = matches[:port].to_i
+    if regex_port > 0
+      regex_port
+    else
+      6379
+    end
+  end
+
+  def uri
+    'redis' +
+      lambda { tls_mode? ? 's' : '' }.call +
+      '://' +
+      hostname +
+      ':' +
+      port.to_s
+  end
+
+  def redis_uri_regex
+    /((?<scheme>redi[s]+)\:\/\/){0,1}(?<first>[0-9A-Za-z_-]+)\.(?<second>[0-9A-Za-z_-]+)\.{0,1}(?<third>[0-9A-Za-z_]*)\.(?<region>[0-9A-Za-z_-]+)\.cache\.amazonaws\.com:{0,1}(?<port>[0-9]*)/
+  end
+
+  def tls_mode?
+    matches[:scheme] == 'rediss' or
+      %w[master clustercfg].include?(matches[:first])
+  end
+
+end