ecs_deploy 0.2.0 → 1.0.1

data/lib/ecs_deploy/auto_scaler/auto_scaling_group_config.rb

@@ -0,0 +1,209 @@
+require "aws-sdk-autoscaling"
+require "aws-sdk-ec2"
+
+require "ecs_deploy"
+require "ecs_deploy/auto_scaler/config_base"
+require "ecs_deploy/auto_scaler/cluster_resource_manager"
+
+module EcsDeploy
+  module AutoScaler
+    AutoScalingGroupConfig = Struct.new(:name, :region, :cluster, :buffer, :service_configs) do
+      include ConfigBase
+
+      MAX_DETACHABLE_INSTANCE_COUNT = 20
+
+      def initialize(attributes = {}, logger)
+        attributes = attributes.dup
+        services = attributes.delete("services")
+        super(attributes, logger)
+        self.service_configs = services.map do |s|
+          ServiceConfig.new(s.merge("cluster" => cluster, "region" => region), logger)
+        end
+      end
+
+      def update_desired_capacity(required_capacity)
+        detach_and_terminate_orphan_instances
+
+        desired_capacity = (required_capacity + buffer.to_f).ceil
+
+        current_asg = client.describe_auto_scaling_groups({
+          auto_scaling_group_names: [name],
+        }).auto_scaling_groups[0]
+
+        if current_asg.desired_capacity > desired_capacity
+          decreased_capacity = decrease_desired_capacity(current_asg.desired_capacity - desired_capacity)
+          if decreased_capacity > 0
+            new_desired_capacity = current_asg.desired_capacity - decreased_capacity
+            cluster_resource_manager.trigger_capacity_update(current_asg.desired_capacity, new_desired_capacity)
+            @logger.info "#{log_prefix} Update desired_capacity to #{new_desired_capacity}"
+          else
+            @logger.info "#{log_prefix} Tried to Update desired_capacity but there were no deregisterable instances"
+          end
+        elsif current_asg.desired_capacity < desired_capacity
+          client.update_auto_scaling_group(
+            auto_scaling_group_name: name,
+            min_size: 0,
+            max_size: [current_asg.max_size, desired_capacity].max,
+            desired_capacity: desired_capacity,
+          )
+          cluster_resource_manager.trigger_capacity_update(current_asg.desired_capacity, desired_capacity)
+          @logger.info "#{log_prefix} Update desired_capacity to #{desired_capacity}"
+        end
+      rescue => e
+        AutoScaler.error_logger.error(e)
+      end
+
+      def cluster_resource_manager
+        @cluster_resource_manager ||= EcsDeploy::AutoScaler::ClusterResourceManager.new(
+          region: region,
+          cluster: cluster,
+          service_configs: service_configs,
+          capacity_based_on: "instances",
+          logger: @logger,
+        )
+      end
+
+      def detach_instances(instance_ids:, should_decrement_desired_capacity:)
+        return if instance_ids.empty?
+
+        instance_ids.each_slice(MAX_DETACHABLE_INSTANCE_COUNT) do |ids|
+          client.detach_instances(
+            auto_scaling_group_name: name,
+            instance_ids: ids,
+            should_decrement_desired_capacity: should_decrement_desired_capacity,
+          )
+        end
+
+        @logger.info "#{log_prefix} Detach instances from ASG: #{instance_ids.inspect}"
+      end
+
+      private
+
+      def decrease_desired_capacity(count)
+        container_instance_arns_in_service = cluster_resource_manager.fetch_container_instance_arns_in_service
+        container_instances_in_cluster = cluster_resource_manager.fetch_container_instances_in_cluster
+        auto_scaling_group_instances = instances(reload: true)
+        deregisterable_instances = container_instances_in_cluster.select do |i|
+          i.pending_tasks_count == 0 &&
+            !running_essential_task?(i, container_instance_arns_in_service) &&
+            auto_scaling_group_instances.any? {|instance| instance.instance_id == i.ec2_instance_id }
+        end
+
+        @logger.info "#{log_prefix} Fetch deregisterable instances: #{deregisterable_instances.map(&:ec2_instance_id).inspect}"
+
+        az_to_instance_count = auto_scaling_group_instances.each_with_object(Hash.new(0)) { |i, h| h[i.availability_zone] += 1 }
+        az_to_deregisterable_instances = deregisterable_instances.group_by do |i|
+          i.attributes.find { |a| a.name == "ecs.availability-zone" }.value
+        end
+
+        deregistered_instance_ids = []
+        prev_max_count = nil
+        # Select instances to be deregistered keeping the balance of instance count per availability zone
+        while deregistered_instance_ids.size < count
+          max_count = az_to_instance_count.each_value.max
+          break if max_count == prev_max_count # No more deregistable instances with keeping the balance
+
+          azs = az_to_instance_count.select { |_, c| c == max_count }.keys
+          azs.each do |az|
+            instance = az_to_deregisterable_instances[az]&.pop
+            next if instance.nil?
+            begin
+              cluster_resource_manager.deregister_container_instance(instance.container_instance_arn)
+              deregistered_instance_ids << instance.ec2_instance_id
+              az_to_instance_count[az] -= 1
+            rescue EcsDeploy::AutoScaler::ClusterResourceManager::DeregisterContainerInstanceFailed
+            end
+            break if deregistered_instance_ids.size >= count
+          end
+          prev_max_count = max_count
+        end
+
+        @logger.info "#{log_prefix} Deregistered instances: #{deregistered_instance_ids.inspect}"
+
+        detach_and_terminate_instances(deregistered_instance_ids)
+
+        deregistered_instance_ids.size
+      end
+
+      def detach_and_terminate_instances(instance_ids)
+        return if instance_ids.empty?
+
+        detach_instances(
+          instance_ids: instance_ids,
+          should_decrement_desired_capacity: true
+        )
+
+        sleep 3
+
+        ec2_client.terminate_instances(instance_ids: instance_ids)
+
+        @logger.info "#{log_prefix} Terminated instances: #{instance_ids.inspect}"
+      rescue => e
+        AutoScaler.error_logger.error(e)
+      end
+
+      def detach_and_terminate_orphan_instances
+        container_instance_ids = cluster_resource_manager.fetch_container_instances_in_cluster.map(&:ec2_instance_id)
+        orphans = instances(reload: true).reject do |i|
+          next true if container_instance_ids.include?(i.instance_id)
+
+          # The lifecycle state of terminated instances becomes "Terminating", "Terminating:Wait", or "Terminating:Proceed",
+          # and we can't detach instances in such a state.
+          if i.lifecycle_state.start_with?("Terminating")
+            AutoScaler.error_logger.warn("#{log_prefix} The lifesycle state of #{i.instance_id} is \"#{i.lifecycle_state}\", so ignore it")
+            next true
+          end
+        end.map(&:instance_id)
+
+        return if orphans.empty?
+
+        targets = ec2_client.describe_instances(instance_ids: orphans).reservations.flat_map(&:instances).select do |i|
+          (Time.now - i.launch_time) > 600
+        end
+
+        detach_and_terminate_instances(targets.map(&:instance_id))
+      rescue => e
+        AutoScaler.error_logger.error(e)
+      end
+
+      def client
+        Aws::AutoScaling::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: region,
+          logger: logger
+        )
+      end
+
+      def ec2_client
+        Aws::EC2::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: region,
+          logger: logger
+        )
+      end
+
+      def instances(reload: false)
+        if reload || @instances.nil?
+          resp = client.describe_auto_scaling_groups({
+            auto_scaling_group_names: [name],
+          })
+          @instances = resp.auto_scaling_groups[0].instances
+        else
+          @instances
+        end
+      end
+
+      def running_essential_task?(instance, container_instance_arns_in_service)
+        return false if instance.running_tasks_count == 0
+
+        container_instance_arns_in_service.include?(instance.container_instance_arn)
+      end
+
+      def log_prefix
+        "[#{self.class.to_s.sub(/\AEcsDeploy::AutoScaler::/, "")} #{name} #{region}]"
+      end
+    end
+  end
+end

data/lib/ecs_deploy/auto_scaler/cluster_resource_manager.rb

@@ -0,0 +1,149 @@
+require "timeout"
+
+require "aws-sdk-ecs"
+
+module EcsDeploy
+  module AutoScaler
+    class ClusterResourceManager
+      class DeregisterContainerInstanceFailed < StandardError; end
+
+      MAX_DESCRIBABLE_SERVICE_COUNT = 10
+
+      def initialize(region:, cluster:, service_configs:, logger: nil, capacity_based_on:)
+        @region = region
+        @cluster = cluster
+        @logger = logger
+        @service_configs = service_configs
+        @capacity_based_on = capacity_based_on
+        if @capacity_based_on != "instances" && @capacity_based_on != "vCPUs"
+          raise ArgumentError, 'capacity_based_on should be either "instances" or "vCPUs"'
+        end
+
+        @mutex = Mutex.new
+        @resource = ConditionVariable.new
+        @used_capacity = @service_configs.sum { |s| s.desired_count * s.required_capacity }
+        @capacity = calculate_active_instance_capacity
+      end
+
+      def acquire(capacity, timeout: nil)
+        @mutex.synchronize do
+          @logger&.debug("#{log_prefix} Try to acquire #{capacity} capacity (capacity: #{@capacity}, used_capacity: #{@used_capacity})")
+          Timeout.timeout(timeout) do
+            while @capacity - @used_capacity < capacity
+              @resource.wait(@mutex)
+            end
+          end
+          @used_capacity += capacity
+          @logger&.debug("#{log_prefix} Acquired #{capacity} capacity (capacity: #{@capacity}, used_capacity: #{@used_capacity})")
+        end
+        true
+      rescue Timeout::Error
+        false
+      end
+
+      def release(capacity)
+        @mutex.synchronize do
+          @used_capacity -= capacity
+          @resource.broadcast
+        end
+        @logger&.debug("#{log_prefix} Released #{capacity} capacity (capacity: #{@capacity}, used_capacity: #{@used_capacity})")
+        true
+      end
+
+      def fetch_container_instances_in_cluster
+        cl = ecs_client
+        resp = cl.list_container_instances(cluster: @cluster)
+        if resp.container_instance_arns.empty?
+          []
+        else
+          resp.flat_map do |resp|
+            cl.describe_container_instances(cluster: @cluster, container_instances: resp.container_instance_arns).container_instances
+          end
+        end
+      end
+
+      def fetch_container_instance_arns_in_service
+        task_groups = @service_configs.map { |s| "service:#{s.name}" }
+        ecs_client.list_container_instances(cluster: @cluster, filter: "task:group in [#{task_groups.join(",")}]").flat_map(&:container_instance_arns)
+      end
+
+      def deregister_container_instance(container_instance_arn)
+        ecs_client.deregister_container_instance(cluster: @cluster, container_instance: container_instance_arn, force: true)
+      rescue Aws::ECS::Errors::InvalidParameterException
+        raise DeregisterContainerInstanceFailed
+      end
+
+      def trigger_capacity_update(old_desired_capacity, new_desired_capacity, interval: 5, wait_until_capacity_updated: false)
+        th = Thread.new do
+          @logger&.info "#{log_prefix} Start updating capacity: #{old_desired_capacity} -> #{new_desired_capacity}"
+          Timeout.timeout(180) do
+            until @capacity == new_desired_capacity || (new_desired_capacity >= old_desired_capacity && @capacity > new_desired_capacity)
+              @mutex.synchronize do
+                begin
+                  @capacity = calculate_active_instance_capacity
+                  @resource.broadcast
+                rescue => e
+                  AutoScaler.error_logger.warn("#{log_prefix} `#{__method__}': #{e} (#{e.class})")
+                end
+              end
+
+              sleep interval
+            end
+            @logger&.info "#{log_prefix} capacity is updated to #{@capacity}"
+          end
+        end
+
+        if wait_until_capacity_updated
+          @logger&.info "#{log_prefix} Wait for the capacity of active instances to become #{new_desired_capacity} from #{old_desired_capacity}"
+          begin
+            th.join
+          rescue Timeout::Error => e
+            msg = "#{log_prefix} `#{__method__}': #{e} (#{e.class})"
+            if @capacity_based_on == "vCPUs"
+              # Timeout::Error sometimes occur.
+              # For example, @capacity won't be new_desired_capacity if new_desired_capacity is odd and all instances have 2 vCPUs
+              AutoScaler.error_logger.warn(msg)
+            else
+              AutoScaler.error_logger.error(msg)
+            end
+          end
+        end
+      end
+
+      def calculate_active_instance_capacity
+        cl = ecs_client
+
+        if @capacity_based_on == "instances"
+          return cl.list_container_instances(cluster: @cluster, status: "ACTIVE").sum do |resp|
+            resp.container_instance_arns.size
+          end
+        end
+
+        total_cpu = cl.list_container_instances(cluster: @cluster, status: "ACTIVE").sum do |resp|
+          next 0 if resp.container_instance_arns.empty?
+          ecs_client.describe_container_instances(
+            cluster: @cluster,
+            container_instances: resp.container_instance_arns,
+          ).container_instances.sum { |ci| ci.registered_resources.find { |r| r.name == "CPU" }.integer_value }
+        end
+
+        total_cpu / 1024
+      end
+
+      private
+
+      def ecs_client
+        Aws::ECS::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: @region,
+          logger: @logger,
+        )
+      end
+
+      def log_prefix
+        "[#{self.class.to_s.gsub(/\AEcsDeploy::AutoScaler::/, "")} #{@region} #{@cluster}]"
+      end
+    end
+  end
+end

data/lib/ecs_deploy/auto_scaler/config_base.rb

@@ -0,0 +1,16 @@
+module EcsDeploy
+  module AutoScaler
+    module ConfigBase
+      def initialize(attributes = {}, logger)
+        attributes.each do |key, val|
+          send("#{key}=", val)
+        end
+        @logger = logger
+      end
+
+      def logger
+        @logger
+      end
+    end
+  end
+end

data/lib/ecs_deploy/auto_scaler/instance_drainer.rb

@@ -0,0 +1,134 @@
+require "aws-sdk-ec2"
+require "aws-sdk-ecs"
+require "aws-sdk-sqs"
+
+require "ecs_deploy"
+
+module EcsDeploy
+  module AutoScaler
+    class InstanceDrainer
+      def initialize(auto_scaling_group_configs:, spot_fleet_request_configs:, logger:)
+        @auto_scaling_group_configs = auto_scaling_group_configs || []
+        @spot_fleet_request_configs = spot_fleet_request_configs || []
+        @logger = logger
+        @stop = false
+      end
+
+      def poll_spot_instance_interruption_warnings(queue_url)
+        @logger.debug "Start polling spot instance interruption warnings of #{queue_url}"
+
+        # cf. https://docs.aws.amazon.com/general/latest/gr/rande.html#sqs_region
+        region = URI.parse(queue_url).host.split(".")[1]
+
+        poller = Aws::SQS::QueuePoller.new(queue_url, client: sqs_client(region))
+        poller.before_request do |stats|
+          throw :stop_polling if @stop
+        end
+
+        until @stop
+          begin
+            poller.poll(max_number_of_messages: 10, visibility_timeout: 15) do |messages, _|
+              instance_ids = messages.map do |msg|
+                JSON.parse(msg.body).dig("detail", "instance-id")
+              end
+
+              config_to_instance_ids = build_config_to_instance_ids(instance_ids, region)
+              set_instance_state_to_draining(config_to_instance_ids, region)
+              # Detach the instances to launch other instances
+              detach_instances_from_auto_scaling_groups(config_to_instance_ids, region)
+            end
+          rescue => e
+            AutoScaler.error_logger.error(e)
+          end
+        end
+
+        @logger.debug "Stop polling spot instance interruption warnings of #{queue_url}"
+      end
+
+      def stop
+        @stop = true
+      end
+
+      private
+
+      def build_config_to_instance_ids(instance_ids, region)
+        config_to_instance_ids = Hash.new{ |h, k| h[k] = [] }
+        ec2_client(region).describe_instances(instance_ids: instance_ids).each do |resp|
+          resp.reservations.each do |reservation|
+            reservation.instances.each do |i|
+              sfr_id = i.tags.find { |t| t.key == "aws:ec2spot:fleet-request-id" }&.value
+              if sfr_id
+                config = @spot_fleet_request_configs.find { |c| c.id == sfr_id && c.region == region }
+                config_to_instance_ids[config] << i.instance_id if config
+                next
+              end
+
+              asg_name = i.tags.find { |t| t.key == "aws:autoscaling:groupName" }&.value
+              if asg_name
+                config = @auto_scaling_group_configs.find { |c| c.name == asg_name && c.region == region }
+                config_to_instance_ids[config] << i.instance_id if config
+              end
+            end
+          end
+        end
+
+        config_to_instance_ids
+      end
+
+      def set_instance_state_to_draining(config_to_instance_ids, region)
+        cl = ecs_client(region)
+        config_to_instance_ids.each do |config, instance_ids|
+          arns = cl.list_container_instances(
+            cluster: config.cluster,
+            filter: "ec2InstanceId in [#{instance_ids.join(",")}]",
+          ).container_instance_arns
+
+          if instance_ids.size != arns.size
+            AutoScaler.error_logger.warn("The number of ARNs differs from the number of instance IDs: instance_ids: #{instance_ids.inspect}, container_instance_arns: #{arns.inspect}")
+          end
+          next if arns.empty?
+
+          cl.update_container_instances_state(
+            cluster: config.cluster,
+            container_instances: arns,
+            status: "DRAINING",
+          )
+          @logger.info "Draining instances: region: #{region}, cluster: #{config.cluster}, instance_ids: #{instance_ids.inspect}, container_instance_arns: #{arns.inspect}"
+        end
+      end
+
+      def detach_instances_from_auto_scaling_groups(config_to_instance_ids, region)
+        @auto_scaling_group_configs.each do |config|
+          config.detach_instances(instance_ids: config_to_instance_ids[config], should_decrement_desired_capacity: false)
+        end
+      end
+
+      def ec2_client(region)
+        Aws::EC2::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: region,
+          logger: @logger,
+        )
+      end
+
+      def ecs_client(region)
+        Aws::ECS::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: region,
+          logger: @logger,
+        )
+      end
+
+      def sqs_client(region)
+        Aws::SQS::Client.new(
+          access_key_id: EcsDeploy.config.access_key_id,
+          secret_access_key: EcsDeploy.config.secret_access_key,
+          region: region,
+          logger: @logger,
+        )
+      end
+    end
+  end
+end