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>

10 files changed, 181 insertions, 63 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
diff --git a/spec/integration/parser/functions/include.rb b/spec/integration/parser/functions/include.rb
deleted file mode 100755
index f84d43276..000000000
--- a/spec/integration/parser/functions/include.rb
+++ /dev/null
@@ -1,25 +0,0 @@
-#!/usr/bin/env ruby
-
-require File.dirname(__FILE__) + '/../../../spec_helper'
-
-describe "The include function" do
-    before :each do
-        @node = Puppet::Node.new("mynode")
-        @compiler = Puppet::Parser::Compiler.new(@node)
-        @compiler.send(:evaluate_main)
-        @scope = @compiler.topscope
-        # preload our functions
-        Puppet::Parser::Functions.function(:include)
-        Puppet::Parser::Functions.function(:require)
-    end
-
-    it "should add a containment relationship between the 'included' class and our class" do
-        @compiler.known_resource_types.add Puppet::Resource::Type.new(:hostclass, "includedclass")
-
-        @scope.function_include("includedclass")
-
-        klass_resource = @compiler.findresource(:class,"includedclass")
-        klass_resource.should be_instance_of(Puppet::Parser::Resource)
-        @compiler.catalog.should be_edge(@scope.resource, klass_resource)
-    end
-end
diff --git a/spec/unit/parser/compiler.rb b/spec/unit/parser/compiler.rb
index d5d46b1f3..e7d336583 100755
--- a/spec/unit/parser/compiler.rb
+++ b/spec/unit/parser/compiler.rb
@@ -10,6 +10,11 @@ class CompilerTestResource
         @title = title
     end
 
+    def [](attr)
+        return nil if attr == :stage
+        :main
+    end
+
     def ref
         "%s[%s]" % [type.to_s.capitalize, title]
     end
@@ -31,12 +36,15 @@ class CompilerTestResource
 end
 
 describe Puppet::Parser::Compiler do
+    def resource(type, title)
+        Puppet::Parser::Resource.new(type, title, :scope => @scope)
+    end
+
     before :each do
         @node = Puppet::Node.new "testnode"
         @known_resource_types = Puppet::Resource::TypeCollection.new "development"
-
         @compiler = Puppet::Parser::Compiler.new(@node)
-        @scope = Puppet::Parser::Scope.new(:compiler => @compiler, :source => "fake_source")
+        @scope = Puppet::Parser::Scope.new(:compiler => @compiler, :source => stub('source'))
         @scope_resource = Puppet::Parser::Resource.new(:file, "/my/file", :scope => @scope)
         @scope.resource = @scope_resource
         @compiler.environment.stubs(:known_resource_types).returns @known_resource_types
@@ -100,6 +108,10 @@ describe Puppet::Parser::Compiler do
             compiler.classlist.should include("foo")
             compiler.classlist.should include("bar")
         end
+
+        it "should add a 'main' stage to the catalog" do
+            @compiler.catalog.resource(:stage, :main).should be_instance_of(Puppet::Parser::Resource)
+        end
     end
 
     describe "when managing scopes" do
@@ -209,7 +221,7 @@ describe Puppet::Parser::Compiler do
         end
 
         it "should ignore builtin resources" do
-            resource = stub 'builtin', :ref => "File[testing]", :builtin? => true, :type => "file"
+            resource = resource(:file, "testing")
 
             @compiler.add_resource(@scope, resource)
             resource.expects(:evaluate).never
@@ -229,7 +241,9 @@ describe Puppet::Parser::Compiler do
         end
 
         it "should not evaluate already-evaluated resources" do
-            resource = stub 'already_evaluated', :ref => "File[testing]", :builtin? => false, :evaluated? => true, :virtual? => false, :type => "file"
+            resource = resource(:file, "testing")
+            resource.stubs(:evaluated?).returns true
+
             @compiler.add_resource(@scope, resource)
             resource.expects(:evaluate).never
 
@@ -350,16 +364,16 @@ describe Puppet::Parser::Compiler do
         end
 
         it "should return added resources in add order" do
-            resource1 = stub "1", :ref => "File[yay]", :type => "file"
+            resource1 = resource(:file, "yay")
             @compiler.add_resource(@scope, resource1)
-            resource2 = stub "2", :ref => "File[youpi]", :type => "file"
+            resource2 = resource(:file, "youpi")
             @compiler.add_resource(@scope, resource2)
 
             @compiler.resources.should == [resource1, resource2]
         end
 
         it "should add resources that do not conflict with existing resources" do
-            resource = CompilerTestResource.new(:file, "yay")
+            resource = resource(:file, "yay")
             @compiler.add_resource(@scope, resource)
 
             @compiler.catalog.should be_vertex(resource)
@@ -374,42 +388,90 @@ describe Puppet::Parser::Compiler do
         end
 
         it "should add an edge from the scope resource to the added resource" do
-            resource = stub "noconflict", :ref => "File[yay]", :type => "file"
+            resource = resource(:file, "yay")
             @compiler.add_resource(@scope, resource)
 
             @compiler.catalog.should be_edge(@scope.resource, resource)
         end
 
-        it "should add edges from the class resources to the main class" do
-            main = CompilerTestResource.new(:class, :main)
-            @compiler.add_resource(@scope, main)
-            resource = CompilerTestResource.new(:class, "foo")
+        it "should add an edge to any specified stage for class resources" do
+            other_stage = resource(:stage, "other")
+            @compiler.add_resource(@scope, other_stage)
+            resource = resource(:class, "foo")
+            resource[:stage] = 'other'
+
+            @compiler.add_resource(@scope, resource)
+
+            @compiler.catalog.edge?(other_stage, resource).should be_true
+        end
+
+        it "should fail if a non-class resource attempts to set a stage" do
+            other_stage = resource(:stage, "other")
+            @compiler.add_resource(@scope, other_stage)
+            resource = resource(:file, "foo")
+            resource[:stage] = 'other'
+
+            lambda { @compiler.add_resource(@scope, resource) }.should raise_error(ArgumentError)
+        end
+
+        it "should fail if an unknown stage is specified" do
+            resource = resource(:class, "foo")
+            resource[:stage] = 'other'
+
+            lambda { @compiler.add_resource(@scope, resource) }.should raise_error(ArgumentError)
+        end
+
+        it "should add edges from the class resources to the main stage if no stage is specified" do
+            main = @compiler.catalog.resource(:stage, :main)
+            resource = resource(:class, "foo")
             @compiler.add_resource(@scope, resource)
 
             @compiler.catalog.should be_edge(main, resource)
         end
 
-        it "should just add edges to the scope resource for the class resources when no main class can be found" do
-            resource = CompilerTestResource.new(:class, "foo")
+        it "should not add non-class resources that don't specify a stage to the 'main' stage" do
+            main = @compiler.catalog.resource(:stage, :main)
+            resource = resource(:file, "foo")
             @compiler.add_resource(@scope, resource)
 
-            @compiler.catalog.should be_edge(@scope.resource, resource)
+            @compiler.catalog.should_not be_edge(main, resource)
+        end
+
+        it "should not add any parent-edges to stages" do
+            stage = resource(:stage, "other")
+            @compiler.add_resource(@scope, stage)
+
+            @scope.resource = resource(:class, "foo")
+
+            @compiler.catalog.edge?(@scope.resource, stage).should be_false
+        end
+
+        it "should not attempt to add stages to other stages" do
+            other_stage = resource(:stage, "other")
+            second_stage = resource(:stage, "second")
+            @compiler.add_resource(@scope, other_stage)
+            @compiler.add_resource(@scope, second_stage)
+
+            second_stage[:stage] = "other"
+
+            @compiler.catalog.edge?(other_stage, second_stage).should be_false
         end
 
         it "should have a method for looking up resources" do
-            resource = stub 'resource', :ref => "Yay[foo]", :type => "file"
+            resource = resource(:yay, "foo")
             @compiler.add_resource(@scope, resource)
             @compiler.findresource("Yay[foo]").should equal(resource)
         end
 
         it "should be able to look resources up by type and title" do
-            resource = stub 'resource', :ref => "Yay[foo]", :type => "file"
+            resource = resource(:yay, "foo")
             @compiler.add_resource(@scope, resource)
             @compiler.findresource("Yay", "foo").should equal(resource)
         end
 
         it "should not evaluate virtual defined resources" do
-            resource = stub 'notevaluated', :ref => "File[testing]", :builtin? => false, :evaluated? => false, :virtual? => true, :type => "file"
+            resource = resource(:file, "testing")
+            resource.virtual = true
             @compiler.add_resource(@scope, resource)
 
             resource.expects(:evaluate).never
@@ -643,8 +705,8 @@ describe Puppet::Parser::Compiler do
     describe "when managing resource overrides" do
 
         before do
-            @override = stub 'override', :ref => "My[ref]", :type => "my"
-            @resource = stub 'resource', :ref => "My[ref]", :builtin? => true, :type => "my"
+            @override = stub 'override', :ref => "File[/foo]", :type => "my"
+            @resource = resource(:file, "/foo")
         end
 
         it "should be able to store overrides" do
diff --git a/spec/unit/simple_graph.rb b/spec/unit/simple_graph.rb
index f8596c7d6..234a2fc1f 100755
--- a/spec/unit/simple_graph.rb
+++ b/spec/unit/simple_graph.rb
@@ -440,6 +440,8 @@ describe Puppet::SimpleGraph do
             @top = Container.new("top", ["g", "h", @middle, @one, @three])
             @empty = Container.new("empty", [])
 
+            @stage = Puppet::Type.type(:stage).new(:name => "foo")
+
             @contgraph = @top.to_graph
 
             # We have to add the container to the main graph, else it won't
@@ -477,6 +479,12 @@ describe Puppet::SimpleGraph do
             @depgraph.vertices.find_all { |v| v.is_a?(Container) }.should be_empty
         end
 
+        # This is a bit hideous, but required to make stages work with relationships - they're
+        # the top of the graph.
+        it "should remove all Stage resources from the dependency graph" do
+            @depgraph.vertices.find_all { |v| v.is_a?(Puppet::Type.type(:stage)) }.should be_empty
+        end
+
         it "should add container relationships to contained objects" do
             @contgraph.leaves(@middle).each do |leaf|
                 @depgraph.should be_edge("h", leaf)
diff --git a/spec/unit/type.rb b/spec/unit/type.rb
index 9381c0a83..a17c9a8fa 100755
--- a/spec/unit/type.rb
+++ b/spec/unit/type.rb
@@ -445,7 +445,8 @@ describe Puppet::Type do
         end
     end
 
-    describe "when managing relationships" do
+    it "should have a 'stage' metaparam" do
+        Puppet::Type.metaparamclass(:stage).should be_instance_of(Class)
     end
 end
 
diff --git a/spec/unit/type/stage.rb b/spec/unit/type/stage.rb
new file mode 100644
index 000000000..68846520d
--- /dev/null
+++ b/spec/unit/type/stage.rb
@@ -0,0 +1,9 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../../spec_helper'
+
+describe Puppet::Type.type(:stage) do
+    it "should have a 'name' parameter'" do
+        Puppet::Type.type(:stage).new(:name => :foo)[:name].should == :foo
+    end
+end