Merge pull request #8914 from joshcooper/splay_then_disable_9998

(PUP-9998) Check agent disabled status after sleeping

Author: cthorn42

lib/puppet/agent.rb

@@ -38,15 +38,25 @@ def needing_restart?
   # Perform a run with our client.
   def run(client_options = {})
     if disabled?
-      Puppet.notice _("Skipping run of %{client_class}; administratively disabled (Reason: '%{disable_message}');\nUse 'puppet agent --enable' to re-enable.") % { client_class: client_class, disable_message: disable_message }
+      log_disabled_message
       return
     end
 
     result = nil
     wait_for_lock_deadline = nil
     block_run = Puppet::Application.controlled_run do
-      # splay may sleep for awhile!
-      splay(client_options.fetch(:splay, Puppet[:splay]))
+      # splay may sleep for awhile when running onetime! If not onetime, then
+      # the job scheduler splays (only once) so that agents assign themselves a
+      # slot within the splay interval.
+      do_splay = client_options.fetch(:splay, Puppet[:splay])
+      if do_splay
+        splay(do_splay)
+
+        if disabled?
+          log_disabled_message
+          break
+        end
+      end
 
       # waiting for certs may sleep for awhile depending on onetime, waitforcert and maxwaitforcert!
       # this needs to happen before forking so that if we fail to obtain certs and try to exit, then
@@ -59,14 +69,19 @@ def run(client_options = {})
           begin
             # lock may sleep for awhile depending on waitforlock and maxwaitforlock!
             lock do
-              # NOTE: Timeout is pretty heinous as the location in which it
-              # throws an error is entirely unpredictable, which means that
-              # it can interrupt code blocks that perform cleanup or enforce
-              # sanity. The only thing a Puppet agent should do after this
-              # error is thrown is die with as much dignity as possible.
-              Timeout.timeout(Puppet[:runtimeout], RunTimeoutError) do
-                Puppet.override(ssl_context: ssl_context) do
-                  client.run(client_args)
+              if disabled?
+                log_disabled_message
+                nil
+              else
+                # NOTE: Timeout is pretty heinous as the location in which it
+                # throws an error is entirely unpredictable, which means that
+                # it can interrupt code blocks that perform cleanup or enforce
+                # sanity. The only thing a Puppet agent should do after this
+                # error is thrown is die with as much dignity as possible.
+                Timeout.timeout(Puppet[:runtimeout], RunTimeoutError) do
+                  Puppet.override(ssl_context: ssl_context) do
+                    client.run(client_args)
+                  end
                 end
               end
             end
@@ -88,8 +103,7 @@ def run(client_options = {})
             end
           rescue RunTimeoutError => detail
             Puppet.log_exception(detail, _("Execution of %{client_class} did not complete within %{runtimeout} seconds and was terminated.") %
-              {client_class: client_class,
-              runtimeout: Puppet[:runtimeout]})
+              {client_class: client_class, runtimeout: Puppet[:runtimeout]})
             nil
           rescue StandardError => detail
             Puppet.log_exception(detail, _("Could not run %{client_class}: %{detail}") % { client_class: client_class, detail: detail })
@@ -155,4 +169,8 @@ def wait_for_certificates(options)
     sm = Puppet::SSL::StateMachine.new(waitforcert: waitforcert, onetime: Puppet[:onetime])
     sm.ensure_client_certificate
   end
+
+  def log_disabled_message
+    Puppet.notice _("Skipping run of %{client_class}; administratively disabled (Reason: '%{disable_message}');\nUse 'puppet agent --enable' to re-enable.") % { client_class: client_class, disable_message: disable_message }
+  end
 end

spec/unit/agent_spec.rb

@@ -103,6 +103,8 @@ def controlled_run(&block)
     end
 
     it "should splay" do
+      Puppet[:splay] = true
+
       expect(@agent).to receive(:splay)
 
       @agent.run
@@ -185,6 +187,26 @@ def controlled_run(&block)
       expect(@agent.run).to eq(:result)
     end
 
+    it "should check if it's disabled after splaying and log a message" do
+      Puppet[:splay] = true
+      Puppet[:splaylimit] = '5s'
+      Puppet[:onetime] = true
+
+      expect(@agent).to receive(:disabled?).and_return(false, true)
+
+      allow(Puppet).to receive(:notice).and_call_original
+      expect(Puppet).to receive(:notice).with(/Skipping run of .*; administratively disabled.*/)
+      @agent.run
+    end
+
+    it "should check if it's disabled after acquiring the lock and log a message" do
+      expect(@agent).to receive(:disabled?).and_return(false, true)
+
+      allow(Puppet).to receive(:notice).and_call_original
+      expect(Puppet).to receive(:notice).with(/Skipping run of .*; administratively disabled.*/)
+      @agent.run
+    end
+
     describe "and a puppet agent is already running" do
       before(:each) do
         allow_any_instance_of(Object).to receive(:sleep)