lib/mcollective/shell.rb in mcollective-client-2.2.4 vs lib/mcollective/shell.rb in mcollective-client-2.4.0

@@ -15,22 +15,26 @@
   #   stdin       - a string that will be sent to stdin of the program
   #   stdout      - a variable that will receive stdout, must support <<
   #   stderr      - a variable that will receive stdin, must support <<
   #   environment - the shell environment, defaults to include LC_ALL=C
   #                 set to nil to clear the environment even of LC_ALL
+  #   timeout     - a timeout in seconds after which the subprocess is killed,
+  #                 the special value :on_thread_exit kills the subprocess
+  #                 when the invoking thread (typically the agent) has ended
   #
   class Shell
-    attr_reader :environment, :command, :status, :stdout, :stderr, :stdin, :cwd
+    attr_reader :environment, :command, :status, :stdout, :stderr, :stdin, :cwd, :timeout
 
     def initialize(command, options={})
       @environment = {"LC_ALL" => "C"}
       @command = command
       @status = nil
       @stdout = ""
       @stderr = ""
       @stdin = nil
       @cwd = Dir.tmpdir
+      @timeout = nil
 
       options.each do |opt, val|
         case opt.to_s
           when "stdout"
             raise "stdout should support <<" unless val.respond_to?("<<")
@@ -52,10 +56,13 @@
             if val.nil?
               @environment = {}
             else
               @environment.merge!(val.dup)
             end
+          when "timeout"
+            raise "timeout should be a positive integer or the symbol :on_thread_exit symbol" unless val.eql?(:on_thread_exit) || ( val.is_a?(Fixnum) && val>0 )
+            @timeout = val
         end
       end
     end
 
     # Actually does the systemu call passing in the correct environment, stdout and stderr
@@ -65,23 +72,53 @@
               "stderr" => @stderr,
               "cwd"    => @cwd}
 
       opts["stdin"] = @stdin if @stdin
 
-      # Running waitpid on the cid here will start a thread
-      # with the waitpid in it, this way even if the thread
-      # that started this process gets killed due to agent
-      # timeout or such there will still be a waitpid waiting
-      # for the child to exit and not leave zombies.
+
+      thread = Thread.current
+      # Start a double fork and exec with systemu which implies a guard thread.
+      # If a valid timeout is configured the guard thread will terminate the
+      # executing process and reap the pid.
+      # If no timeout is specified the process will run to completion with the
+      # guard thread reaping the pid on completion.
       @status = systemu(@command, opts) do |cid|
         begin
-          sleep 1
-          Process::waitpid(cid)
+          if timeout.is_a?(Fixnum)
+            # wait for the specified timeout
+            sleep timeout
+          else
+            # sleep while the agent thread is still alive
+            while(thread.alive?)
+              sleep 0.1
+            end
+          end
+
+          # if the process is still running
+          if (Process.kill(0, cid))
+            # and a timeout was specified
+            if timeout
+              if Util.windows?
+                Process.kill('KILL', cid)
+              else
+                # Kill the process
+                Process.kill('TERM', cid)
+                sleep 2
+                Process.kill('KILL', cid) if (Process.kill(0, cid))
+              end
+            end
+            # only wait if the parent thread is dead
+            Process.waitpid(cid) unless thread.alive?
+          end
         rescue SystemExit
+        rescue Errno::ESRCH
         rescue Errno::ECHILD
+          Log.warn("Could not reap process '#{cid}'.")
         rescue Exception => e
           Log.info("Unexpected exception received while waiting for child process: #{e.class}: #{e}")
         end
       end
+      @status.thread.kill
+      @status
     end
   end
 end