Adding #2658 - Adding support for run stages

This allows you to specify a run stage for either
a class or a resource.

By default, all classes get directly added to the
'main' stage.  You can create new stages as resources:

    stage { [pre, post]: }

To order stages, use standard relationships:

    stage { pre: before => Stage[main] }

Or use the new relationship syntax:

    stage { pre: } -> Stage[main] -> stage { post: }

Then use the new class parameters to specify a stage:

    class { foo: stage => pre }

If you set a stage on an individual resource, it will
fail; stages can only be set on class resources.

Signed-off-by: Luke Kanies <luke@puppetlabs.com>

5 files changed, 80 insertions, 17 deletions

diff --git a/lib/puppet/parser/compiler.rb b/lib/puppet/parser/compiler.rb
index ae4af7622..6cc71a62e 100644
--- a/lib/puppet/parser/compiler.rb
+++ b/lib/puppet/parser/compiler.rb
@@ -51,16 +51,31 @@ class Puppet::Parser::Compiler
         # Note that this will fail if the resource is not unique.
         @catalog.add_resource(resource)
 
-        # And in the resource graph.  At some point, this might supercede
-        # the global resource table, but the table is a lot faster
-        # so it makes sense to maintain for now.
-        if resource.type.to_s.downcase == "class" and main = @catalog.resource(:class, :main)
-            @catalog.add_edge(main, resource)
-        else
-            @catalog.add_edge(scope.resource, resource)
+        set_container_resource(scope, resource)
+    end
+
+    def set_container_resource(scope, resource)
+        # Add our container edge.  If we're a class, then we get treated specially - we can
+        # control the stage that the class is applied in.  Otherwise, we just
+        # get added to our parent container.
+        return if resource.type.to_s.downcase == "stage"
+
+        if resource.type.to_s.downcase != "class"
+            if resource[:stage]
+                raise ArgumentError, "Only classes can set 'stage'; normal resources like #{resource} cannot change run stage"
+            end
+            return @catalog.add_edge(scope.resource, resource)
         end
+
+        unless stage = @catalog.resource(:stage, resource[:stage] || :main)
+            raise ArgumentError, "Could not find stage #{resource[:stage] || :main} specified by #{resource}"
+        end
+
+        @catalog.add_edge(stage, resource)
     end
 
+    private :set_container_resource
+
     # Do we use nodes found in the code, vs. the external node sources?
     def ast_nodes?
         known_resource_types.nodes?
@@ -169,7 +184,6 @@ class Puppet::Parser::Compiler
         end
 
         initvars()
-        init_main()
     end
 
     # Create a new scope, with either a specified parent scope or
@@ -407,12 +421,6 @@ class Puppet::Parser::Compiler
         data
     end
 
-    # Initialize the top-level scope, class, and resource.
-    def init_main
-        # Create our initial scope and a resource that will evaluate main.
-        @topscope = Puppet::Parser::Scope.new(:compiler => self)
-    end
-
     # Set up all of our internal variables.
     def initvars
         # The list of objects that will available for export.
@@ -435,6 +443,12 @@ class Puppet::Parser::Compiler
         @catalog = Puppet::Resource::Catalog.new(@node.name)
         @catalog.version = known_resource_types.version
 
+        # Create our initial scope and a resource that will evaluate main.
+        @topscope = Puppet::Parser::Scope.new(:compiler => self)
+
+        @main_stage_resource = Puppet::Parser::Resource.new("stage", :main, :scope => @topscope)
+        @catalog.add_resource(@main_stage_resource)
+
         # local resource array to maintain resource ordering
         @resources = []
 
diff --git a/lib/puppet/parser/resource.rb b/lib/puppet/parser/resource.rb
index a4587723e..78d90d4f1 100644
--- a/lib/puppet/parser/resource.rb
+++ b/lib/puppet/parser/resource.rb
@@ -22,9 +22,9 @@ class Puppet::Parser::Resource < Puppet::Resource
     include Puppet::Parser::YamlTrimmer
 
     attr_accessor :source, :scope, :rails_id
-    attr_accessor :virtual, :override, :translated, :catalog
+    attr_accessor :virtual, :override, :translated, :catalog, :evaluated
 
-    attr_reader :exported, :evaluated, :parameters
+    attr_reader :exported, :parameters
 
     # Determine whether the provided parameter name is a relationship parameter.
     def self.relationship_parameter?(name)
diff --git a/lib/puppet/simple_graph.rb b/lib/puppet/simple_graph.rb
index 91603945c..cf0eff38e 100644
--- a/lib/puppet/simple_graph.rb
+++ b/lib/puppet/simple_graph.rb
@@ -323,7 +323,8 @@ class Puppet::SimpleGraph
         # graph.  We could get a similar affect by only setting relationships
         # to container leaves, but that would result in many more
         # relationships.
-        containers = other.topsort.find_all { |v| v.is_a?(type) and vertex?(v) }
+        stage_class = Puppet::Type.type(:stage)
+        containers = other.topsort.find_all { |v| (v.is_a?(type) or v.is_a?(stage_class)) and vertex?(v) }
         containers.each do |container|
             # Get the list of children from the other graph.
             children = other.adjacent(container, :direction => :out)
diff --git a/lib/puppet/type.rb b/lib/puppet/type.rb
index 33d2d0390..7f8fb09f3 100644
--- a/lib/puppet/type.rb
+++ b/lib/puppet/type.rb
@@ -1402,6 +1402,36 @@ class Type
             This will restart the sshd service if the sshd config file changes.}
     end
 
+    newmetaparam(:stage) do
+        desc %{Which run stage a given resource should reside in.  This just creates
+            a dependency on or from the named milestone.  For instance, saying that
+            this is in the 'bootstrap' stage creates a dependency on the 'bootstrap'
+            milestone.
+        
+            By default, all classes get directly added to the
+            'main' stage.  You can create new stages as resources:
+            
+                stage { [pre, post]: }
+            
+            To order stages, use standard relationships:
+            
+                stage { pre: before => Stage[main] }
+            
+            Or use the new relationship syntax:
+            
+                Stage[pre] -> Stage[main] -> Stage[post]
+            
+            Then use the new class parameters to specify a stage:
+            
+                class { foo: stage => pre }
+            
+            Stages can only be set on classes, not individual resources.  This will
+            fail::
+
+                file { '/foo': stage => pre, ensure => file }
+        }
+    end
+
     ###############################
     # All of the provider plumbing for the resource types.
     require 'puppet/provider'
diff --git a/lib/puppet/type/stage.rb b/lib/puppet/type/stage.rb
new file mode 100644
index 000000000..fba78764d
--- /dev/null
+++ b/lib/puppet/type/stage.rb
@@ -0,0 +1,18 @@
+Puppet::Type.newtype(:stage) do
+    desc "A resource type for specifying run stages.  The actual stage should
+    be specified on resources::
+        class { foo: stage => pre }
+        
+    And you must manually control stage order::
+    
+        stage { pre: before => Stage[main] }
+        
+    You automatically get a 'main' stage created, and by default all resources
+    get inserted into that stage.
+
+    You can only set stages on class resources, not normal builtin resources."
+
+    newparam :name do
+        desc "The name of the stage. This will be used as the 'stage' for each resource."
+    end
+end