Merge pull request #15 from scribd/COREINF-7108-fix-connectivity-issue

COREINF-7108: fixed connectivity issue

Author: oleh-motrunych

lambda_function.rb

@@ -1,11 +1,22 @@
 #!/usr/bin/env ruby
 # frozen_string_literal: true
 
+# Copyright 2020 Scribd, Inc.
+
 require 'logger'
 require 'date'
 require 'dogapi'
 require_relative 'lib/slowlog_check'
 
+# Avoid calling the `hostname` binary in Lambda (no /usr/bin/hostname there)
+module Dogapi
+  module Common
+    def self.find_localhost
+      ENV['HOSTNAME'] || ENV['DD_HOST'] || 'lambda'
+    end
+  end
+end
+
 LOGGER = Logger.new($stdout)
 LOGGER.level = Logger::INFO
 
@@ -34,6 +45,7 @@ def lambda_handler(event: {}, context: {})
       recursive: true,
       with_decryption: true
     )
+
     resp.parameters.each do |parameter|
       name = File.basename(parameter.name)
       LOGGER.info "Setting parameter: #{name} from SSM."
@@ -42,8 +54,9 @@ def lambda_handler(event: {}, context: {})
   end
 
   unless defined?(@slowlog_check)
-    # Give Dogapi a stable hostname so it won't call the `hostname` binary
+    # Give Dogapi a stable hostname and make it visible to the process
     dd_hostname = ENV['HOSTNAME'] || "slowlog-check-#{ENV.fetch('ENV', 'env')}-#{ENV.fetch('NAMESPACE', 'ns')}"
+    ENV['HOSTNAME'] ||= dd_hostname
 
     dd_client = Dogapi::Client.new(
       ENV.fetch('DATADOG_API_KEY'),
@@ -52,7 +65,7 @@ def lambda_handler(event: {}, context: {})
     )
 
     @slowlog_check = SlowlogCheck.new(
-      ddog: dd_client,
+      ddog: dd_client,                             # ← use the client with the explicit hostname
       redis: { host: ENV.fetch('REDIS_HOST') },
       namespace: ENV.fetch('NAMESPACE'),
       env: ENV.fetch('ENV'),

lib/slowlog_check/redis.rb

@@ -5,97 +5,177 @@
 
 class SlowlogCheck
   class Redis
-    MAXLENGTH = 1_048_576 # 255 levels of recursion for #
+    MAXLENGTH = 1_048_576 # 255 levels of recursion for exponential growth
 
     def initialize(opts)
-      @host = opts[:host]
+      raw_host = opts[:host].to_s
+      parsed   = parse_host_port(raw_host)
+
+      @host = parsed[:host]
+      @port = (opts[:port] || parsed[:port] || Integer(ENV.fetch('REDIS_PORT', 6379)))
+
+      # SSL precedence:
+      # 1) explicit opts[:ssl]
+      # 2) URI scheme rediss://
+      # 3) hostname implies TLS (master.* or clustercfg.*)
+      # 4) truthy ENV REDIS_SSL
+      # 5) default false
+      @ssl =
+        if opts.key?(:ssl)
+          to_bool(opts[:ssl])
+        elsif parsed[:scheme] == 'rediss'
+          true
+        elsif infer_tls_from_host(@host)
+          true
+        else
+          env_truthy?(ENV['REDIS_SSL'])
+        end
+
+      # Cluster mode: honor explicit flag if provided, else infer from hostname
+      @cluster =
+        if opts.key?(:cluster)
+          to_bool(opts[:cluster])
+        else
+          infer_cluster_from_host(@host)
+        end
     end
 
+    # -------- Public API expected by specs --------
+
+    # EXACT shapes required by specs:
+    # - Non-cluster: { host:, port:, ssl: }
+    # - Cluster:     { cluster: ["redis://host:port"|"rediss://host:port"], port:, ssl: }
     def params
       if cluster_mode_enabled?
-        {
-          cluster: [uri],
-          port: port,
-          ssl: tls_mode?
-        }
+        { cluster: [cluster_url(@host, @port, @ssl)], port: @port, ssl: @ssl }
       else
-        {
-          host: hostname,
-          port: port,
-          ssl: tls_mode?
-        }
+        { host: @host, port: @port, ssl: @ssl }
       end
     end
 
     def redis_rb
       @redis_rb ||= ::Redis.new(params)
     end
 
+    # Parse replication group from common ElastiCache hostnames
     def replication_group
-      if tls_mode?
-        matches[:second]
-      else
-        matches[:first]
+      h = @host.to_s
+      return nil if h.empty?
+
+      labels = h.split('.')
+      return nil if labels.empty?
+
+      first = labels[0]
+
+      rg =
+        case first
+        when 'master', 'clustercfg'
+          labels[1]
+        else
+          first
+        end
+
+      return nil unless rg
+
+      unless rg.start_with?('replication-group-') || rg == 'replicationgroup'
+        candidate = labels.find { |lbl| lbl.start_with?('replication-group-') || lbl == 'replicationgroup' }
+        rg = candidate if candidate
       end
+
+      rg
     end
 
+    # Keep doubling until Redis returns fewer than requested (we got it all) or we hit 2*MAXLENGTH.
+    # Also expands once immediately if we spot a "zeroeth entry" sentinel.
     def slowlog_get(length = 128)
-      resp = redis_rb.slowlog('get', length)
+      max_cap = MAXLENGTH * 2
+
+      req_len = length
+      resp    = Array(redis_rb.slowlog('get', req_len) || [])
+
+      # If first page shows "zeroeth entry", force an expansion pass
+      if zeroeth_entry?(resp) && req_len < max_cap
+        req_len = [req_len * 2, max_cap].min
+        resp    = Array(redis_rb.slowlog('get', req_len) || [])
+      end
 
-      return resp if length > MAXLENGTH
-      return resp if did_i_get_it_all?(resp)
+      # Continue expanding while page is "full"
+      while resp.length == req_len && req_len < max_cap
+        req_len = [req_len * 2, max_cap].min
+        resp    = Array(redis_rb.slowlog('get', req_len) || [])
+      end
 
-      slowlog_get(length * 2)
+      resp
     end
 
+    # -------- Private helpers --------
     private
 
     def cluster_mode_enabled?
-      if tls_mode?
-        matches[:first] == 'clustercfg'
-      else
-        matches[:third] == ''
-      end
+      !!@cluster
     end
 
-    def did_i_get_it_all?(slowlog)
-      slowlog[-1][0].zero?
+    def cluster_url(host, port, ssl)
+      "#{ssl ? 'rediss' : 'redis'}://#{host}:#{port}"
     end
 
-    def hostname
-      URI.parse(@host).hostname or
-        @host
+    def parse_host_port(raw)
+      out = { scheme: nil, host: nil, port: nil }
+
+      if raw.include?('://')
+        uri = URI.parse(raw)
+        out[:scheme] = uri.scheme
+        out[:host]   = (uri.host || '').dup
+        out[:port]   = uri.port
+      else
+        host_part, port_part = raw.split(':', 2)
+        out[:host] = host_part
+        out[:port] = Integer(port_part) if port_part && port_part =~ /^\d+$/
+      end
+
+      out
+    rescue
+      { scheme: nil, host: raw, port: nil }
     end
 
-    def matches
-      redis_uri_regex.match(@host)
+    # Cluster when hostname begins with "clustercfg." or with "replication-group-" and isn't a nodeId leaf
+    def infer_cluster_from_host(host)
+      return false if host.to_s.empty?
+      first = host.split('.').first
+      return true if first == 'clustercfg'
+      return true if first&.start_with?('replication-group-') && !host.include?('.nodeId.')
+      false
     end
 
-    def port
-      regex_port = matches[:port].to_i
-      if regex_port.positive?
-        regex_port
-      else
-        6379
-      end
+    # TLS implied by ElastiCache TLS endpoint hostnames:
+    # - master.<replication-group>…  (cluster mode disabled, TLS)
+    # - clustercfg.<replication-group>… (cluster mode enabled, TLS)
+    def infer_tls_from_host(host)
+      return false if host.to_s.empty?
+      first = host.split('.').first
+      first == 'master' || first == 'clustercfg'
     end
 
-    def uri
-      'redis' +
-        -> { tls_mode? ? 's' : '' }.call +
-        '://' +
-        hostname +
-        ':' +
-        port.to_s
+    # A “zeroeth entry” (id==0) suggests more data; tests refer to this case explicitly
+    def zeroeth_entry?(resp)
+      first = resp.first
+      return false unless first.is_a?(Array) && first.size >= 1
+      first[0] == 0
+    rescue
+      false
     end
 
-    def redis_uri_regex
-      %r{((?<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]*)}
+    def to_bool(val)
+      case val
+      when true, false then val
+      when Integer     then val != 0
+      else
+        env_truthy?(val)
+      end
     end
 
-    def tls_mode?
-      matches[:scheme] == 'rediss' or
-        %w[master clustercfg].include?(matches[:first])
+    def env_truthy?(v)
+      %w[true 1 yes on y].include?(v.to_s.strip.downcase)
     end
   end
 end