Merge pull request #238 from rtyler/elb-rebase

Elb rebase

README.md

@@ -238,6 +238,20 @@ Vagrant.configure("2") do |config|
 end
 ```
 
+### Elastic Load Balancers
+
+You can automatically attach an instance to an ELB during boot and detach on destroy.
+
+```ruby
+Vagrant.configure("2") do |config|
+  # ... other stuff
+
+  config.vm.provider "aws" do |aws|
+    aws.elb = "production-web"
+  end
+end
+```
+
 ## Development
 
 To work on the `vagrant-aws` plugin, clone this repository out, and use

lib/vagrant-aws/action.rb

@@ -37,6 +37,7 @@ def self.action_destroy
                 end
               end
               b2.use ConnectAWS
+              b2.use ElbDeregisterInstance
               b2.use TerminateInstance
               b2.use ProvisionerCleanup if defined?(ProvisionerCleanup)
             else
@@ -113,13 +114,9 @@ def self.action_ssh_run
         end
       end
 
-      def self.action_prepare_boot
-        Vagrant::Action::Builder.new.tap do |b|
-          b.use Provision
-          b.use SyncFolders
-          b.use WarnNetworks
-        end
-      end
+      def self.action_prepare_boot Vagrant::Action::Builder.new.tap do |b|
+        b.use Provision b.use SyncFolders b.use WarnNetworks b.use
+        ElbRegisterInstance end end
 
       # This action is called to bring the box up from nothing.
       def self.action_up
@@ -185,6 +182,8 @@ def self.action_reload
       autoload :TimedProvision, action_root.join("timed_provision") # some plugins now expect this action to exist
       autoload :WaitForState, action_root.join("wait_for_state")
       autoload :WarnNetworks, action_root.join("warn_networks")
+      autoload :ElbRegisterInstance, action_root.join("elb_register_instance")
+      autoload :ElbDeregisterInstance, action_root.join("elb_deregister_instance")
     end
   end
 end

lib/vagrant-aws/action/connect_aws.rb

@@ -37,6 +37,7 @@ def call(env)
 
           @logger.info("Connecting to AWS...")
           env[:aws_compute] = Fog::Compute.new(fog_config)
+          env[:aws_elb]     = Fog::AWS::ELB.new(fog_config.except(:provider))
 
           @app.call(env)
         end

lib/vagrant-aws/action/elb_deregister_instance.rb

@@ -0,0 +1,24 @@
+require 'vagrant-aws/util/elb'
+
+module VagrantPlugins
+  module AWS
+    module Action
+      # This registers instance in ELB
+      class ElbDeregisterInstance
+        include ElasticLoadBalancer
+
+        def initialize(app, env)
+          @app    = app
+          @logger = Log4r::Logger.new("vagrant_aws::action::elb_deregister_instance")
+        end
+
+        def call(env)
+          if elb_name = env[:machine].provider_config.elb
+            deregister_instance env, elb_name, env[:machine].id
+          end
+          @app.call(env)
+        end
+      end
+    end
+  end
+end

lib/vagrant-aws/action/elb_register_instance.rb

@@ -0,0 +1,24 @@
+require 'vagrant-aws/util/elb'
+
+module VagrantPlugins
+  module AWS
+    module Action
+      # This registers instance in ELB
+      class ElbRegisterInstance
+        include ElasticLoadBalancer
+
+        def initialize(app, env)
+          @app    = app
+          @logger = Log4r::Logger.new("vagrant_aws::action::elb_register_instance")
+        end
+
+        def call(env)
+          @app.call(env)
+          if elb_name = env[:machine].provider_config.elb
+            register_instance env, elb_name, env[:machine].id
+          end
+        end
+      end
+    end
+  end
+end

lib/vagrant-aws/config.rb

@@ -138,6 +138,12 @@ class Config < Vagrant.plugin("2", :config)
       # @return [Boolean]
       attr_accessor :associate_public_ip
 
+      # The name of ELB, which an instance should be
+      # attached to
+      #
+      # @return [String]
+      attr_accessor :elb
+
       def initialize(region_specific=false)
         @access_key_id          = UNSET_VALUE
         @ami                    = UNSET_VALUE
@@ -164,6 +170,7 @@ def initialize(region_specific=false)
         @monitoring             = UNSET_VALUE
         @ebs_optimized          = UNSET_VALUE
         @associate_public_ip    = UNSET_VALUE
+        @elb                    = UNSET_VALUE
 
         # Internal state (prefix with __ so they aren't automatically
         # merged)
@@ -305,6 +312,9 @@ def finalize!
         # default false
         @associate_public_ip = false if @associate_public_ip == UNSET_VALUE
 
+        # Don't attach instance to any ELB by default
+        @elb = nil if @elb == UNSET_VALUE
+
         # Compile our region specific configurations only within
         # NON-REGION-SPECIFIC configurations.
         if !@__region_specific

lib/vagrant-aws/errors.rb

@@ -25,6 +25,9 @@ class RsyncError < VagrantAWSError
 
       class MkdirError < VagrantAWSError
         error_key(:mkdir_error)
+
+      class ElbDoesNotExistError < VagrantAWSError
+        error_key("elb_does_not_exist")
       end
     end
   end

lib/vagrant-aws/util/elb.rb

@@ -0,0 +1,56 @@
+module VagrantPlugins
+  module AWS
+    module ElasticLoadBalancer
+
+      def register_instance(env, elb_name, instance_id)
+        env[:ui].info I18n.t("vagrant_aws.elb.registering", instance_id: instance_id, elb_name: elb_name), :new_line => false
+        elb = get_load_balancer(env[:aws_elb], elb_name)
+        unless elb.instances.include? instance_id
+          elb.register_instances([instance_id])
+          env[:ui].info I18n.t("vagrant_aws.elb.ok"), :prefix => false
+          adjust_availability_zones env, elb
+        else
+          env[:ui].info I18n.t("vagrant_aws.elb.skipped"), :prefix => false
+        end
+      end
+
+      def deregister_instance(env, elb_name, instance_id)
+        env[:ui].info I18n.t("vagrant_aws.elb.deregistering", instance_id: instance_id, elb_name: elb_name), :new_line => false
+        elb = get_load_balancer(env[:aws_elb], elb_name)
+        if elb.instances.include? instance_id
+          elb.deregister_instances([instance_id])
+          env[:ui].info I18n.t("vagrant_aws.elb.ok"), :prefix => false
+          adjust_availability_zones env, elb
+        else
+          env[:ui].info I18n.t("vagrant_aws.elb.skipped"), :prefix => false
+        end
+      end
+
+      def adjust_availability_zones(env, elb)
+        env[:ui].info I18n.t("vagrant_aws.elb.adjusting", elb_name: elb.id), :new_line => false
+
+        instances = env[:aws_compute].servers.all("instance-id" => elb.instances)
+
+        azs = if instances.empty?
+                ["#{env[:machine].provider_config.region}a"]
+              else
+                instances.map(&:availability_zone).uniq
+              end
+
+        az_to_disable = elb.availability_zones - azs
+        az_to_enable = azs - elb.availability_zones
+
+        elb.enable_availability_zones az_to_enable unless az_to_enable.empty?
+        elb.disable_availability_zones az_to_disable unless az_to_disable.empty?
+
+        env[:ui].info I18n.t("vagrant_aws.elb.ok"), :prefix => false
+      end
+
+      private
+
+      def get_load_balancer(aws, name)
+        aws.load_balancers.find { |lb| lb.id == name } or raise Errors::ElbDoesNotExistError
+      end
+    end
+  end
+end

locales/en.yml

@@ -2,6 +2,18 @@ en:
   vagrant_aws:
     already_status: |-
       The machine is already %{status}.
+    elb:
+      adjusting: |-
+        Adjusting availability zones of ELB %{elb_name}...
+      registering: |-
+        Registering %{instance_id} at ELB %{elb_name}...
+      deregistering: |-
+        Deregistering %{instance_id} from ELB %{elb_name}...
+      ok: |-
+        ok
+      skipped: |-
+        skipped
+
     launching_instance: |-
       Launching an instance with the following settings...
     launch_no_keypair: |-
@@ -86,6 +98,8 @@ en:
 
         Host path: %{hostpath}
         Error: %{err}
+      elb_does_not_exist: |-
+        ELB configured for the instance does not exist
 
     states:
       short_not_created: |-