elasticsearch-drain 0.0.2 → 0.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: f9322921e73eb7ffbce7f78845f5cc43e1fe63cb
-  data.tar.gz: 7eee908fc33ff6fdf6633902f4b6efb0095b8e52
+  metadata.gz: eb751f23604391d421be5540c9006217eb0d91b3
+  data.tar.gz: deb14051a07a1efd61162ff3f489c80c8513a54f
 SHA512:
-  metadata.gz: 50d8dd4fa38c7015a08ed115173e0af9eb8221f8c238cd792990e5a46cb86da5cf933d5ddc9b53fad86b822bdb17105ec9bfd6d4bbff10495410dde431401288
-  data.tar.gz: fa1af902021af37359f2f14cc4fb88817ed99242902fd7c6044aec2880aec74d998fbf6f690f9f621b51f99461f74b256fa23d43aaf7b6156f84bc4122568c70
+  metadata.gz: 3935ee1cabb06227d9f21375e90deabf10f7c552089c84ee014ef78a1527d8b3c5a65c17126696824f3cc21574c6fee9e178e6dc1da52326328a7ee43ac7b6c4
+  data.tar.gz: 263a73535b0908e021da6f4f4028efd1ac4e7fdb0c6f7be70e2ad2905c4a77ee511d09d9170f84bfb7d9ae0d1c9e5550f5026bd90641cd9986c965772edab166

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    elasticsearch-drain (0.0.2)
+    elasticsearch-drain (0.0.1)
       aws-sdk (~> 2)
       elasticsearch (~> 1.0)
       thor (~> 0.19)
@@ -14,25 +14,25 @@ GEM
     ast (2.1.0)
     astrolabe (1.3.1)
       parser (~> 2.2)
-    aws-sdk (2.1.34)
-      aws-sdk-resources (= 2.1.34)
-    aws-sdk-core (2.1.34)
+    aws-sdk (2.6.14)
+      aws-sdk-resources (= 2.6.14)
+    aws-sdk-core (2.6.14)
       jmespath (~> 1.0)
-    aws-sdk-resources (2.1.34)
-      aws-sdk-core (= 2.1.34)
+    aws-sdk-resources (2.6.14)
+      aws-sdk-core (= 2.6.14)
     coderay (1.1.0)
     crack (0.4.2)
       safe_yaml (~> 1.0.0)
     docile (1.1.5)
-    elasticsearch (1.0.14)
-      elasticsearch-api (= 1.0.14)
-      elasticsearch-transport (= 1.0.14)
-    elasticsearch-api (1.0.14)
+    elasticsearch (1.1.0)
+      elasticsearch-api (= 1.1.0)
+      elasticsearch-transport (= 1.1.0)
+    elasticsearch-api (1.1.0)
       multi_json
     elasticsearch-extensions (0.0.18)
       ansi
       ruby-prof
-    elasticsearch-transport (1.0.14)
+    elasticsearch-transport (1.1.0)
       faraday
       multi_json
     faraday (0.9.2)
@@ -53,7 +53,7 @@ GEM
       guard-compat (~> 1.2)
       minitest (>= 3.0)
     hashdiff (0.2.2)
-    jmespath (1.1.3)
+    jmespath (1.3.1)
     json (1.8.3)
     listen (3.0.3)
       rb-fsevent (>= 0.9.3)
@@ -62,7 +62,7 @@ GEM
     method_source (0.8.2)
     minitest (5.8.1)
     mixlib-shellout (2.2.3)
-    multi_json (1.11.2)
+    multi_json (1.12.1)
     multipart-post (2.0.0)
     nenv (0.2.0)
     notiffany (0.0.8)
@@ -127,3 +127,6 @@ DEPENDENCIES
   vcr (~> 2.9)
   webmock (~> 1.21)
   yard (~> 0.8)
+
+BUNDLED WITH
+   1.13.2

data/README.md

@@ -46,6 +46,8 @@ gem install bundler
 bundle install
 ```
 
+To enable the tests that will hit the AWS APIs pass `ALLOW_DISABLED_VCR=true`
+
 Run test tests (unit and style):
 ```bash
 rake

data/Rakefile

@@ -14,6 +14,8 @@ Rake::TestTask.new do |t|
   t.pattern = 'test/**/test_*.rb'
 end
 
+ENV['ES_VERSION'] ||= '1.7.2'
+
 def elasticsearch_command
   path = "tmp/elasticsearch-#{ENV['ES_VERSION']}/bin/elasticsearch"
   path = ::File.expand_path(path, __dir__)
@@ -22,7 +24,6 @@ end
 
 ENV['TEST_CLUSTER_NODES'] = '1'
 ENV['TEST_CLUSTER_COMMAND'] = elasticsearch_command
-ENV['ES_VERSION'] ||= '1.7.2'
 
 namespace :elasticsearch do
   task :clean do

data/lib/elasticsearch/drain.rb

@@ -38,12 +38,12 @@ module Elasticsearch
       @asg_client ||= AutoScaling.new(@asg_name, @region)
     end
 
-    # Convience method to access {Elasticsearch::Drain::Nodes}
+    # Convenience method to access {Elasticsearch::Drain::Nodes}
     def nodes
-      Nodes.new(client, asg)
+      @nodes ||= Nodes.new(client, asg)
     end
 
-    # Convience method to access {Elasticsearch::Drain::Cluster#cluster}
+    # Convenience method to access {Elasticsearch::Drain::Cluster#cluster}
     #
     # @return [Elasticsearch::API::Cluster] Elasticsearch cluster client
     def cluster
@@ -52,11 +52,21 @@ module Elasticsearch
 
     def active_nodes_in_asg
       instances = asg.instances
-      nodes.nodes_in_asg(reload: true, instances: instances)
+      nodes.filter_nodes(instances, true)
+    end
+
+    module Errors
+      class WaiterExpired < RuntimeError
+        def new(_msg = nil)
+          'Waiter Expired' + $ERROR_INFO
+        end
+      end
+      class NodeNotFound < RuntimeError; end
     end
   end
 end
 
+require_relative 'drain/util'
 require_relative 'drain/autoscaling'
 require_relative 'drain/version'
 require_relative 'drain/base'

data/lib/elasticsearch/drain/autoscaling.rb

@@ -1,26 +1,34 @@
 module Elasticsearch
   class Drain
     class AutoScaling
+      include Drain::Util
 
       # @attribute [r]
       # EC2 AutoScaling Group name
       attr_reader :asg
 
       # @attribute [r]
-      # EC2 Client
-      attr_reader :ec2_client
+      # AWS region
+      attr_reader :region
 
       def initialize(asg, region)
         @asg = asg
-        @asg_client = Aws::AutoScaling::Client.new(region: region)
-        @ec2_client = Aws::EC2::Client.new(region: region)
+        @region = region
         @instances = nil
         @instance_ids = nil
       end
 
+      def asg_client
+        Aws::AutoScaling::Client.new(region: region)
+      end
+
+      def ec2_client
+        Aws::EC2::Client.new(region: region)
+      end
+
       def find_instances_in_asg
         instances = []
-        @asg_client.describe_auto_scaling_instances.each do |page|
+        asg_client.describe_auto_scaling_instances.each do |page|
           instances << page.auto_scaling_instances.map do |i|
             i.instance_id if i.auto_scaling_group_name == asg
           end
@@ -30,17 +38,30 @@ module Elasticsearch
         @instance_ids = instances
       end
 
+      # Get instances in an AutoScaling Group
+      #
+      # @return [Array<Aws::EC2::Types::Instance>] EC2 Instance objects
       def describe_instances
         instances = []
         find_instances_in_asg if @instance_ids.nil?
         return [] if @instance_ids.empty?
-        @ec2_client.describe_instances(instance_ids: @instance_ids).each do |page|
+        ec2_client.describe_instances(instance_ids: @instance_ids).each do |page|
           instances << page.reservations.map(&:instances)
         end
         instances.flatten!
         @instances = instances
       end
 
+      # Describe an AutoScaling Group
+      #
+      # @return [Struct] AutoScaling Group
+      def describe_autoscaling_group
+        groups = asg_client.describe_auto_scaling_groups(
+          auto_scaling_group_names: [asg]
+        )
+        groups.auto_scaling_groups.first
+      end
+
       def find_private_ips
         instances = describe_instances.clone
         return [] if instances.nil?
@@ -57,27 +78,51 @@ module Elasticsearch
       def instance(ipaddress)
         describe_instances if @instances.nil?
         instances = @instances.clone
-        instances.find { |i| i.private_ip_address == ipaddress }
+        instance = instances.find { |i| i.private_ip_address == ipaddress }
+        fail Errors::NodeNotFound if instance.nil?
+        instance
       end
 
       # Sets the MinSize of an AutoScalingGroup
       #
       # @option [FixNum] count (0) The new MinSize of the AutoScalingGroup
       # @return [Struct] Empty response from the sdk
-      def min_size(count = 0)
-        @asg_client.update_auto_scaling_group(
+      def min_size=(count = 0)
+        asg_client.update_auto_scaling_group(
           auto_scaling_group_name: asg,
           min_size: count
         )
+        wait_until(count) do
+          min_size
+        end
+      end
+
+      # Gets the MinSize of an AutoScalingGroup
+      #
+      # @return [Integer] Value of MinSize of an AutoScalingGroup
+      def min_size
+        group = describe_autoscaling_group
+        group.min_size
+      end
+
+      # Gets the DesiredCapacity of an AutoScalingGroup
+      #
+      # @return [Integer] Value of DesiredCapacity of an AutoScalingGroup
+      def desired_capacity
+        group = describe_autoscaling_group
+        group.desired_capacity
       end
 
       def detach_instance(instance_id)
-        resp = @asg_client.detach_instances(
+        current_desired_capacity = desired_capacity
+        asg_client.detach_instances(
           instance_ids: [instance_id],
           auto_scaling_group_name: asg,
           should_decrement_desired_capacity: true
         )
-        resp.activities.first.status_code == 'Successful'
+        wait_until(current_desired_capacity - 1) do
+          desired_capacity
+        end
       end
     end
   end

data/lib/elasticsearch/drain/cli.rb

@@ -2,27 +2,93 @@ require 'thor'
 
 module Elasticsearch
   class Drain
-    class CLI < ::Thor
+    class CLI < ::Thor # rubocop:disable Metrics/ClassLength
       package_name :elasticsearch
 
       attr_reader :drainer
       attr_accessor :active_nodes
 
+      # rubocop:disable Metrics/LineLength
       desc 'asg', 'Drain all documents from all nodes in an EC2 AutoScaling Group'
       option :host, default: 'localhost:9200'
       option :asg, required: true
       option :region, required: true
+      option :nodes, type: :array, desc: 'A comma separated list of node IDs to drain. If specified, the --number option has no effect'
+      option :number, type: :numeric, desc: 'The number of nodes to drain'
+      option :continue, type: :boolean, default: true, desc: 'Whether to continue draining nodes once the first iteration of --number is complete'
+      # rubocop:enable Metrics/LineLength
       def asg # rubocop:disable Metrics/MethodLength
         @drainer = Elasticsearch::Drain.new(options[:host],
                                             options[:asg],
                                             options[:region])
+
         ensure_cluster_healthy
         @active_nodes = drainer.active_nodes_in_asg
+
+        # If :nodes are specified, :number has no effect
+        if options[:nodes]
+          say "Nodes #{options[:nodes].join(', ')} have been specified, the --number option has no effect"
+          number_to_drain = nil
+          currently_draining_nodes = nil
+        else
+          number_to_drain = options[:number]
+          currently_draining_nodes = drainer.cluster.currently_draining('_id')
+        end
+
+        # If a node or nodes are specified, only drain the requested node(s)
+        @active_nodes = active_nodes.find_all do |n|
+          instance_id = drainer.asg.instance(n.ipaddress).instance_id
+          options[:nodes].include?(instance_id)
+        end if options[:nodes]
+
         do_exit { say_status 'Complete', 'Nothing to do', :green } if active_nodes.empty?
-        say_status 'Found Nodes', "AutoScalingGroup: #{instances}", :magenta
-        ensure_cluster_healthy
-        drain_nodes
-        remove_nodes
+        say_status 'Found Nodes', "AutoScalingGroup: #{instances(active_nodes)}", :magenta
+
+        until active_nodes.empty?
+          ensure_cluster_healthy
+
+          nodes = active_nodes
+
+          # If there are nodes in cluster settings "transient.cluster.routing.allocation.exclude"
+          # test if those nodes are still in the ASG. If so, work on them first unless nodes are
+          # specified.
+          if currently_draining_nodes
+            nodes_to_drain = active_nodes.find_all { |n| currently_draining_nodes.split(',').include?(n.id) }
+
+            # If the list of nodes_to_drain isn't empty, we want to set nodes to the list of nodes
+            # we've already been working on.
+            unless nodes_to_drain.empty?
+              nodes = nodes_to_drain
+
+              say_status 'Active Nodes', "Resuming drain process on #{instances(nodes)}", :magenta
+            end
+
+            # We should only process currently_draining_nodes once
+            currently_draining_nodes = nil
+          end
+
+          # If we specify a number but DON'T specify nodes, sample the active_nodes.
+          if number_to_drain
+            nodes = nodes.sample(number_to_drain.to_i)
+            say_status 'Active Nodes', "Sampled #{number_to_drain} nodes and got #{instances(nodes)}", :magenta
+          end
+
+          @active_nodes = nodes unless options[:continue]
+
+          drain_nodes(nodes)
+          remove_nodes(nodes)
+
+          # Remove the drained nodes from the list of active_nodes
+          @active_nodes -= nodes
+
+          unless active_nodes.empty?
+            say_status 'Drain Nodes', "#{active_nodes.length} nodes remaining", :green
+
+            sleep_time = wait_sleep_time
+            say_status 'Waiting', "Sleeping for #{sleep_time} seconds before the next iteration", :green
+            sleep sleep_time
+          end
+        end
         say_status 'Complete', 'Draining nodes complete!', :green
       end
 
@@ -40,32 +106,54 @@ module Elasticsearch
           exit code
         end
 
-        def instances
-          instances = active_nodes.map(&:ipaddress)
+        def instances(nodes)
+          instances = nodes.map(&:ipaddress)
           instances.join(' ')
         end
 
-        def drain_nodes
-          drainer.asg.min_size(0)
-          nodes_to_drain = active_nodes.map(&:id).join(',')
+        def adjusted_min_size(nodes)
+          min_size = drainer.asg.min_size
+          desired_capacity = drainer.asg.desired_capacity
+          desired_min_size = if (desired_capacity - nodes.length) >= min_size # Removing the nodes won't violate the min_size
+                               # Reduce the asg min_size proportionally
+                               (min_size - nodes.length) <= 0 ? 0 : (min_size - nodes.length)
+                             else
+                               # Removing the nodes will result in the min_size being violated
+                               desired_capacity - nodes.length
+                             end
+          say_status 'Debug', 'min_size = #{min_size}, desired_capacity = #{desired_capacity}, desired_min_size = #{desired_min_size}', :magenta
+          desired_min_size
+        end
+
+        def drain_nodes(nodes)
+          drainer.asg.min_size = adjusted_min_size(nodes)
+          nodes_to_drain = nodes.map(&:id).join(',')
           say_status 'Drain Nodes', "Draining nodes: #{nodes_to_drain}", :magenta
           drainer.cluster.drain_nodes(nodes_to_drain, '_id')
         end
 
-        def remove_nodes # rubocop:disable Metrics/MethodLength
-          while active_nodes.length > 0
-            active_nodes.each do |instance|
-              instance_id = drainer.asg.instance(instance.ipaddress).instance_id
-              instance.instance_id = instance_id
+        def wait_sleep_time
+          ips = active_nodes.map(&:ipaddress)
+          bytes = drainer.nodes.filter_nodes(ips).map(&:bytes_stored)
+          sleep_time = 10
+          sleep_time = 30 if bytes.any? { |b| b >= 100_000 }
+          sleep_time = 60 if bytes.any? { |b| b >= 1_000_000 }
+          sleep_time = 120 if bytes.any? { |b| b >= 10_000_000_000 }
+          sleep_time
+        end
 
-              self.active_nodes = drainer.active_nodes_in_asg
+        def remove_nodes(nodes) # rubocop:disable Metrics/MethodLength
+          while nodes.length > 0
+            sleep_time = wait_sleep_time
+            nodes.each do |instance|
+              instance = drainer.nodes.filter_nodes([instance.ipaddress], true).first
               if instance.bytes_stored > 0
                 say_status 'Drain Status', "Node #{instance.ipaddress} has #{instance.bytes_stored} bytes to move", :blue
-                sleep 2
+                sleep sleep_time
               else
-                remove_node(instance)
-                self.active_nodes = drainer.active_nodes_in_asg
-                break if active_nodes.length < 1
+                next unless remove_node(instance)
+                nodes.delete_if { |n| n.ipaddress == instance.ipaddress }
+                break if nodes.length < 1
                 say_status 'Waiting', 'Sleeping for 1 minute before removing the next node', :green
                 sleep 60
               end
@@ -74,18 +162,24 @@ module Elasticsearch
         end
 
         def remove_node(instance) # rubocop:disable Metrics/MethodLength
+          instance_id = drainer.asg.instance(instance.ipaddress).instance_id
+          instance.instance_id = instance_id
           say_status(
             'Removing Node',
             "Removing #{instance.ipaddress} from Elasticsearch cluster and #{drainer.asg.asg} AutoScalingGroup",
-            :magenta)
+            :magenta
+          )
           sleep 5 unless instance.in_recovery?
           node = "#{instance.instance_id}(#{instance.ipaddress})"
           ensure_cluster_healthy
           say_status 'ASG Remove Node', "Removing node: #{node} from AutoScalingGroup: #{drainer.asg.asg}", :magenta
           drainer.asg.detach_instance(instance.instance_id)
+          sleep 2
           ensure_cluster_healthy
           say_status 'Terminate Instance', "Terminating instance: #{node}", :magenta
           instance.terminate
+        rescue Errors::NodeNotFound
+          false
         end
       end
     end