(puppetlabs#3910) Make puppet agent use environment specified in catalog

If there is an environment specified in the catalog to be applied (i.e. an
external node classifier was used and gave an environment for this node) the
puppet agent will switch to that environment.

When switching environment, the agent needs to get facts and plugins again
using the new environment, and then try fetching the catalog again using the
new facts, since this may change the catalog. It loops in this process several
times, until the environment doesn't change. If it fails to stabilize in a
limited number of tries it will raise an error and abort the run.

This change also means that the prerun command runs after fetching catalog but before
applying. It previously ran before fetching the catalog.

lib/puppet/configurer.rb

@@ -97,14 +97,29 @@ def convert_catalog(result, duration)
     catalog
   end
 
-  # Retrieve (optionally) and apply a catalog. If a catalog is passed in
-  # the options, then apply that one, otherwise retrieve it.
-  def retrieve_and_apply_catalog(options, fact_options)
+  def prepare_and_retrieve_catalog(options)
+    prepare(options)
+
+    if Puppet::Resource::Catalog.indirection.terminus_class == :rest
+      # This is a bit complicated.  We need the serialized and escaped facts,
+      # and we need to know which format they're encoded in.  Thus, we
+      # get a hash with both of these pieces of information.
+      fact_options = facts_for_uploading
+    end
+
+    # set report host name now that we have the fact
+    options[:report].host = Puppet[:node_name_value]
+
     unless catalog = (options.delete(:catalog) || retrieve_catalog(fact_options))
       Puppet.err "Could not retrieve catalog; skipping run"
       return
     end
+    catalog
+  end
 
+  # Retrieve (optionally) and apply a catalog. If a catalog is passed in
+  # the options, then apply that one, otherwise retrieve it.
+  def apply_catalog(catalog, options)
     report = options[:report]
     report.configuration_version = catalog.version
     report.environment = Puppet[:environment]
@@ -126,25 +141,29 @@ def run(options = {})
 
     Puppet::Util::Log.newdestination(report)
     begin
-      prepare(options)
-
-      if Puppet::Resource::Catalog.indirection.terminus_class == :rest
-        # This is a bit complicated.  We need the serialized and escaped facts,
-        # and we need to know which format they're encoded in.  Thus, we
-        # get a hash with both of these pieces of information.
-        fact_options = facts_for_uploading
-      end
+      begin
+        unless catalog = prepare_and_retrieve_catalog(options)
+          return nil
+        end
 
-      # set report host name now that we have the fact
-      report.host = Puppet[:node_name_value]
+        # Here we set the local environment based on what we get from the
+        # catalog. Since a change in environment means a change in facts, and
+        # facts may be used to determine which catalog we get, we need to
+        # rerun the process if the environment is changed.
+        tries = 0
+        while catalog.environment and not catalog.environment.empty? and catalog.environment != Puppet[:environment]
+          if tries > 3
+            raise Puppet::Error, "Catalog environment didn't stabilize after #{tries} fetches, aborting run"
+          end
+          Puppet.warning "Local environment: \"#{Puppet[:environment]}\" doesn't match server specified environment \"#{catalog.environment}\", restarting agent run with new environment"
+          Puppet[:environment] = catalog.environment
+          return nil unless catalog = prepare_and_retrieve_catalog(options)
+          tries += 1
+        end
 
-      begin
         execute_prerun_command or return nil
-        if retrieve_and_apply_catalog(options, fact_options)
-          report.exit_status
-        else
-          nil
-        end
+        apply_catalog(catalog, options)
+        report.exit_status
       rescue SystemExit,NoMemoryError
         raise
       rescue => detail

lib/puppet/resource/catalog.rb

@@ -454,7 +454,7 @@ def to_pson_data_hash
         'tags'      => tags,
         'name'      => name,
         'version'   => version,
-        'environment' => environment,
+        'environment' => environment.to_s,
         'resources' => vertices.collect { |v| v.to_pson_data_hash },
         'edges'     => edges.   collect { |e| e.to_pson_data_hash },
         'classes'   => classes

spec/unit/configurer_spec.rb

@@ -249,9 +249,8 @@
       Puppet.settings[:prerun_command] = "/my/command"
       Puppet::Util::Execution.expects(:execute).with(["/my/command"]).raises(Puppet::ExecutionFailure, "Failed")
 
-      report.expects(:<<).with { |log| log.message.include?("Could not run command from prerun_command") }
-
       @agent.run.should be_nil
+      report.logs.find { |x| x.message =~ /Could not run command from prerun_command/ }.should be
     end
 
     it "should send the transaction report even if the post-run command fails" do
@@ -320,6 +319,22 @@
       @agent.run.should be_nil
     end
 
+    it "should refetch the catalog if the server specifies a new environment in the catalog" do
+      @catalog.stubs(:environment).returns("second_env")
+      @agent.expects(:prepare).twice
+      @agent.expects(:retrieve_catalog).returns(@catalog).twice
+
+      @agent.run
+    end
+
+    it "should change the environment setting if the server specifies a new environment in the catalog" do
+      @catalog.stubs(:environment).returns("second_env")
+
+      @agent.run
+
+      Puppet[:environment].should == "second_env"
+    end
+
     describe "when not using a REST terminus for catalogs" do
       it "should not pass any facts when retrieving the catalog" do
         Puppet::Resource::Catalog.indirection.terminus_class = :compiler