Initial commit
[trapthecap.git] / systemu.rb
1 # vim: ts=2:sw=2:sts=2:et:fdm=marker
2 require 'tmpdir'
3 require 'socket'
4 require 'fileutils'
5 require 'rbconfig'
6 require 'thread'
7 require 'yaml'
8
9 class Object
10 def systemu(*a, &b) SystemUniversal.new(*a, &b).systemu end
11 end
12
13 class SystemUniversal
14 #
15 # constants
16 #
17 SystemUniversal::VERSION = '1.2.0' unless defined? SystemUniversal::VERSION
18 def version() SystemUniversal::VERSION end
19 #
20 # class methods
21 #
22
23 @host = Socket.gethostname
24 @ppid = Process.ppid
25 @pid = Process.pid
26 @turd = ENV['SYSTEMU_TURD']
27
28 c = ::Config::CONFIG
29 ruby = File.join(c['bindir'], c['ruby_install_name']) << c['EXEEXT']
30 @ruby = if system('%s -e 42' % ruby)
31 ruby
32 else
33 system('%s -e 42' % 'ruby') ? 'ruby' : warn('no ruby in PATH/CONFIG')
34 end
35
36 class << self
37 %w( host ppid pid ruby turd ).each{|a| attr_accessor a}
38 end
39
40 #
41 # instance methods
42 #
43
44 def initialize argv, opts = {}, &block
45 getopt = getopts opts
46
47 @argv = argv
48 @block = block
49
50 @stdin = getopt[ ['stdin', 'in', '0', 0] ]
51 @stdout = getopt[ ['stdout', 'out', '1', 1] ]
52 @stderr = getopt[ ['stderr', 'err', '2', 2] ]
53 @env = getopt[ 'env' ]
54 @cwd = getopt[ 'cwd' ]
55
56 @host = getopt[ 'host', self.class.host ]
57 @ppid = getopt[ 'ppid', self.class.ppid ]
58 @pid = getopt[ 'pid', self.class.pid ]
59 @ruby = getopt[ 'ruby', self.class.ruby ]
60 end
61
62 def systemu
63 tmpdir do |tmp|
64 c = child_setup tmp
65 status = nil
66
67 begin
68 thread = nil
69
70 quietly{
71 IO.popen "#{ @ruby } #{ c['program'] }", 'r+' do |pipe|
72 line = pipe.gets
73 case line
74 when %r/^pid: \d+$/
75 cid = Integer line[%r/\d+/]
76 else
77 begin
78 buf = pipe.read
79 buf = "#{ line }#{ buf }"
80 e = Marshal.load buf
81 raise unless Exception === e
82 raise e
83 rescue
84 raise "wtf?\n#{ buf }\n"
85 end
86 end
87 thread = new_thread cid, @block if @block
88 pipe.read rescue nil
89 end
90 }
91 status = $?
92 ensure
93 if thread
94 begin
95 class << status
96 attr 'thread'
97 end
98 status.instance_eval{ @thread = thread }
99 rescue
100 42
101 end
102 end
103 end
104
105 if @stdout or @stderr
106 open(c['stdout']){|f| relay f => @stdout} if @stdout
107 open(c['stderr']){|f| relay f => @stderr} if @stderr
108 status
109 else
110 [status, IO.read(c['stdout']), IO.read(c['stderr'])]
111 end
112 end
113 end
114
115 def new_thread cid, block
116 q = Queue.new
117 Thread.new(cid) do |cid|
118 current = Thread.current
119 current.abort_on_exception = true
120 q.push current
121 block.call cid
122 end
123 q.pop
124 end
125
126 def child_setup tmp
127 stdin = File.expand_path(File.join(tmp, 'stdin'))
128 stdout = File.expand_path(File.join(tmp, 'stdout'))
129 stderr = File.expand_path(File.join(tmp, 'stderr'))
130 program = File.expand_path(File.join(tmp, 'program'))
131 config = File.expand_path(File.join(tmp, 'config'))
132
133 if @stdin
134 open(stdin, 'w'){|f| relay @stdin => f}
135 else
136 FileUtils.touch stdin
137 end
138 FileUtils.touch stdout
139 FileUtils.touch stderr
140
141 c = {}
142 c['argv'] = @argv
143 c['env'] = @env
144 c['cwd'] = @cwd
145 c['stdin'] = stdin
146 c['stdout'] = stdout
147 c['stderr'] = stderr
148 c['program'] = program
149 open(config, 'w'){|f| YAML.dump c, f}
150
151 open(program, 'w'){|f| f.write child_program(config)}
152
153 c
154 end
155
156 def quietly
157 v = $VERBOSE
158 $VERBOSE = nil
159 yield
160 ensure
161 $VERBOSE = v
162 end
163
164 def child_program config
165 <<-program
166 PIPE = STDOUT.dup
167 begin
168 require 'yaml'
169
170 config = YAML.load(IO.read('#{ config }'))
171
172 argv = config['argv']
173 env = config['env']
174 cwd = config['cwd']
175 stdin = config['stdin']
176 stdout = config['stdout']
177 stderr = config['stderr']
178
179 Dir.chdir cwd if cwd
180 env.each{|k,v| ENV[k.to_s] = v.to_s} if env
181
182 STDIN.reopen stdin
183 STDOUT.reopen stdout
184 STDERR.reopen stderr
185
186 PIPE.puts "pid: \#{ Process.pid }"
187 PIPE.flush ### the process is ready yo!
188 PIPE.close
189
190 exec *argv
191 rescue Exception => e
192 PIPE.write Marshal.dump(e) rescue nil
193 exit 42
194 end
195 program
196 end
197
198 def relay srcdst
199 src, dst, ignored = srcdst.to_a.first
200 if src.respond_to? 'read'
201 while((buf = src.read(8192))); dst << buf; end
202 else
203 src.each{|buf| dst << buf}
204 end
205 end
206
207 def tmpdir d = Dir.tmpdir, max = 42, &b
208 i = -1 and loop{
209 i += 1
210
211 tmp = File.join d, "systemu_#{ @host }_#{ @ppid }_#{ @pid }_#{ rand }_#{ i += 1 }"
212
213 begin
214 Dir.mkdir tmp
215 rescue Errno::EEXIST
216 raise if i >= max
217 next
218 end
219
220 break(
221 if b
222 begin
223 b.call tmp
224 ensure
225 FileUtils.rm_rf tmp unless SystemU.turd
226 end
227 else
228 tmp
229 end
230 )
231 }
232 end
233
234 def getopts opts = {}
235 lambda do |*args|
236 keys, default, ignored = args
237 catch('opt') do
238 [keys].flatten.each do |key|
239 [key, key.to_s, key.to_s.intern].each do |key|
240 throw 'opt', opts[key] if opts.has_key?(key)
241 end
242 end
243 default
244 end
245 end
246 end
247 end
248
249 SystemU = SystemUniversal unless defined? SystemU
250
251
252
253
254
255
256
257
258
259
260
261
262
263 if $0 == __FILE__
264 #
265 # date
266 #
267 date = %q( ruby -e" t = Time.now; STDOUT.puts t; STDERR.puts t " )
268
269 status, stdout, stderr = systemu date
270 p [status, stdout, stderr]
271
272 status = systemu date, 1=>(stdout = '')
273 p [status, stdout]
274
275 status = systemu date, 2=>(stderr = '')
276 p [status, stderr]
277 #
278 # sleep
279 #
280 sleep = %q( ruby -e" p(sleep(1)) " )
281 status, stdout, stderr = systemu sleep
282 p [status, stdout, stderr]
283
284 sleep = %q( ruby -e" p(sleep(42)) " )
285 status, stdout, stderr = systemu(sleep){|cid| Process.kill 9, cid}
286 p [status, stdout, stderr]
287 #
288 # env
289 #
290 env = %q( ruby -e" p ENV['A'] " )
291 status, stdout, stderr = systemu env, :env => {'A' => 42}
292 p [status, stdout, stderr]
293 #
294 # cwd
295 #
296 env = %q( ruby -e" p Dir.pwd " )
297 status, stdout, stderr = systemu env, :cwd => Dir.tmpdir
298 p [status, stdout, stderr]
299 end