From e2e87bcf76a8b9c5d1aabc070bb53c3277d58f7e Mon Sep 17 00:00:00 2001 From: Guo Xiang Tan Date: Mon, 17 Apr 2017 15:30:39 +0800 Subject: [PATCH] Allow workers to timeout. --- lib/image_optim.rb | 4 +++ lib/image_optim/cmd.rb | 64 +++++++++++++++++++++++++++++++++ lib/image_optim/config.rb | 5 +++ lib/image_optim/worker.rb | 13 ++++++- spec/image_optim/cmd_spec.rb | 26 ++++++++++++++ spec/image_optim/config_spec.rb | 16 +++++++++ 6 files changed, 127 insertions(+), 1 deletion(-) diff --git a/lib/image_optim.rb b/lib/image_optim.rb index 3a3ff80c..8a2cd682 100644 --- a/lib/image_optim.rb +++ b/lib/image_optim.rb @@ -44,6 +44,9 @@ class ImageOptim # Cache worker digests attr_reader :cache_worker_digests + # Timeout + attr_reader :timeout + # Initialize workers, specify options using worker underscored name: # # pass false to disable worker @@ -76,6 +79,7 @@ def initialize(options = {}) allow_lossy cache_dir cache_worker_digests + timeout ].each do |name| instance_variable_set(:"@#{name}", config.send(name)) $stderr << "#{name}: #{send(name)}\n" if verbose diff --git a/lib/image_optim/cmd.rb b/lib/image_optim/cmd.rb index decfabf4..e9bc1915 100644 --- a/lib/image_optim/cmd.rb +++ b/lib/image_optim/cmd.rb @@ -1,4 +1,5 @@ require 'English' +require 'timeout' class ImageOptim # Helper for running commands @@ -15,6 +16,32 @@ def run(*args) success end + # Run commands using `Process.spawn` + # Return success status + # Will raise Timeout::Error when command timeouts + def run_with_timeout(timeout, *args) + success = false + + if timeout > 0 + pid = spawn_process(*args) + + begin + Timeout.timeout(timeout) do + Process.wait(pid, 0) + check_status! + success = $CHILD_STATUS.exitstatus.zero? + end + rescue Timeout::Error => e + cleanup_process(pid) + raise e + end + else + success = run(*args) + end + + success + end + # Run using backtick # Return captured output # Will raise SignalException if process was interrupted @@ -44,6 +71,43 @@ def check_status! fail SignalException, status.termsig end + + def cleanup_process(pid) + Thread.new do + Process.kill('-TERM', pid) + Process.detach(pid) + + begin + Timeout.timeout(10) do + begin + Process.getpgid(pid) + rescue Errno::ESRCH + sleep 0.001 + retry + end + end + rescue Timeout::Error + Process.kill('-KILL', pid) + end + end + end + + def spawn_process(*args) + pgroup_opt = Gem.win_platform? ? :new_pgroup : :pgroup + + if args.last.is_a?(Hash) + args.last[pgroup_opt] = true + else + args.push(pgroup_opt => true) + end + + if Process.respond_to?(:spawn) + Process.spawn(*args) + else + args.pop if RUBY_VERSION < '1.9' + Process.fork{ exec(*args) } + end + end end end end diff --git a/lib/image_optim/config.rb b/lib/image_optim/config.rb index 80d73922..a9d36855 100644 --- a/lib/image_optim/config.rb +++ b/lib/image_optim/config.rb @@ -157,6 +157,11 @@ def cache_worker_digests !!get!(:cache_worker_digests) end + def timeout + timeout = get!(:timeout) + timeout ? timeout.to_i : 0 + end + # Options for worker class by its `bin_sym`: # * `Hash` passed as is # * `{}` for `true` or `nil` diff --git a/lib/image_optim/worker.rb b/lib/image_optim/worker.rb index 75f53723..1ca83f71 100644 --- a/lib/image_optim/worker.rb +++ b/lib/image_optim/worker.rb @@ -12,6 +12,8 @@ class ImageOptim class Worker extend ClassMethods + class Timeout < Timeout::Error; end + class << self # Default init for worker is new # Check example of override in gifsicle worker @@ -151,7 +153,16 @@ def run_command(cmd_args) {:out => Path::NULL, :err => Path::NULL}, ].flatten end - Cmd.run(*args) + + if @image_optim.timeout > 0 + begin + Cmd.run_with_timeout(@image_optim.timeout, *args) + rescue ::Timeout::Error + raise ImageOptim::Worker::Timeout + end + else + Cmd.run(*args) + end end end end diff --git a/spec/image_optim/cmd_spec.rb b/spec/image_optim/cmd_spec.rb index 7640fc50..0e39c98e 100644 --- a/spec/image_optim/cmd_spec.rb +++ b/spec/image_optim/cmd_spec.rb @@ -41,6 +41,32 @@ def expect_int_exception(&block) end end + describe '.run_with_timeout' do + it 'calls spawn and returns status' do + expect(Cmd.run_with_timeout(20, 'sh -c "exit 0"')).to eq(true) + expect($CHILD_STATUS.exitstatus).to eq(0) + + [1, 66].each do |status| + expect(Cmd.run_with_timeout(20, "sh -c \"exit #{status}\"")). + to eq(false) + + expect($CHILD_STATUS.exitstatus).to eq(status) + end + end + + it 'raises Timeout::Error if process timeouts' do + expect{ Cmd.run_with_timeout(0.001, 'sleep 1') }. + to raise_error(Timeout::Error) + end + + it 'calls system if timeout is <= zero' do + expect(Cmd.run_with_timeout(0, 'sh -c "sleep 0.001; exit 0"')).to eq(true) + + expect(Cmd.run_with_timeout(-1, 'sh -c "sleep 0.001; exit 1"')). + to eq(false) + end + end + describe '.capture' do it 'calls ` and returns result' do output = double diff --git a/spec/image_optim/config_spec.rb b/spec/image_optim/config_spec.rb index f6bd20f8..0a7b288a 100644 --- a/spec/image_optim/config_spec.rb +++ b/spec/image_optim/config_spec.rb @@ -94,6 +94,22 @@ end end + describe '#timeout' do + before do + allow(IOConfig).to receive(:read_options).and_return({}) + end + + it 'is 0 by default' do + config = IOConfig.new({}) + expect(config.timeout).to eq(0) + end + + it 'converts value to number' do + config = IOConfig.new(:timeout => '15') + expect(config.timeout).to eq(15) + end + end + describe '#for_worker' do before do allow(IOConfig).to receive(:read_options).and_return({})