145 lines
3.4 KiB
Python
145 lines
3.4 KiB
Python
|
#!/usr/bin/python
|
||
|
|
#
|
||
|
|
# beware the jobberwack
|
||
|
|
#
|
||
|
|
import sys, os, errno, select, subprocess, fcntl
|
||
|
|
import options
|
||
|
|
from helpers import *
|
||
|
|
|
||
|
|
optspec = """
|
||
|
|
jwack [-j maxjobs] -- <command...>
|
||
|
|
--
|
||
|
|
j,jobs= maximum jobs to run at once
|
||
|
|
"""
|
||
|
|
|
||
|
|
_fds = None
|
||
|
|
_tokens = {}
|
||
|
|
_waitfds = {}
|
||
|
|
|
||
|
|
|
||
|
|
def setup(maxjobs):
|
||
|
|
global _fds
|
||
|
|
if _fds:
|
||
|
|
return # already set up
|
||
|
|
flags = ' ' + os.getenv('MAKEFLAGS', '') + ' '
|
||
|
|
FIND = ' --jobserver-fds='
|
||
|
|
ofs = flags.find(FIND)
|
||
|
|
if ofs >= 0:
|
||
|
|
s = flags[ofs+len(FIND):]
|
||
|
|
(arg,junk) = s.split(' ', 1)
|
||
|
|
(a,b) = arg.split(',', 1)
|
||
|
|
a = atoi(a)
|
||
|
|
b = atoi(b)
|
||
|
|
if a <= 0 or b <= 0:
|
||
|
|
raise ValueError('invalid --jobserver-fds: %r' % arg)
|
||
|
|
_fds = (a,b)
|
||
|
|
if maxjobs and not _fds:
|
||
|
|
# need to start a new server
|
||
|
|
_fds = os.pipe()
|
||
|
|
os.write(_fds[1], 't' * maxjobs)
|
||
|
|
os.putenv('MAKEFLAGS',
|
||
|
|
'%s --jobserver-fds=%d,%d -j ' % (os.getenv('MAKEFLAGS'),
|
||
|
|
_fds[0], _fds[1]))
|
||
|
|
|
||
|
|
|
||
|
|
def wait(want_token):
|
||
|
|
rfds = _waitfds.keys()
|
||
|
|
if _fds and want_token:
|
||
|
|
rfds.append(_fds[0])
|
||
|
|
r,w,x = select.select(rfds, [], [])
|
||
|
|
#print 'readable: %r' % r
|
||
|
|
for fd in r:
|
||
|
|
if _fds and fd == _fds[0]:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
p = _waitfds[fd]
|
||
|
|
os.write(_fds[1], 't' * _tokens[fd])
|
||
|
|
b = os.read(fd, 1)
|
||
|
|
#print 'read: %r' % b
|
||
|
|
if b:
|
||
|
|
#print 'giving up %d tokens for child' % _tokens[fd]
|
||
|
|
_tokens[fd] = 0
|
||
|
|
else:
|
||
|
|
os.close(fd)
|
||
|
|
del _waitfds[fd]
|
||
|
|
del _tokens[fd]
|
||
|
|
p.wait()
|
||
|
|
|
||
|
|
|
||
|
|
def wait_for_token():
|
||
|
|
pfd = atoi(os.getenv('JWACK_PARENT_FD', ''))
|
||
|
|
#print 'pfd is %d' % pfd
|
||
|
|
if pfd:
|
||
|
|
#print 'wrote to pfd'
|
||
|
|
os.write(pfd, 'j') # tell parent to give back his token
|
||
|
|
while 1:
|
||
|
|
print 'waiting for tokens...'
|
||
|
|
wait(want_token=1)
|
||
|
|
if _fds:
|
||
|
|
fcntl.fcntl(_fds[0], fcntl.F_SETFL, os.O_NONBLOCK)
|
||
|
|
try:
|
||
|
|
b = os.read(_fds[0], 1) # FIXME try: block
|
||
|
|
except OSError, e:
|
||
|
|
if e.errno == errno.EAGAIN:
|
||
|
|
b = ''
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
raise
|
||
|
|
if b:
|
||
|
|
break
|
||
|
|
print 'got a token.'
|
||
|
|
|
||
|
|
|
||
|
|
def wait_all():
|
||
|
|
while _waitfds:
|
||
|
|
wait(want_token=0)
|
||
|
|
|
||
|
|
|
||
|
|
def force_return_tokens():
|
||
|
|
n = sum(_tokens.values())
|
||
|
|
print 'returning %d tokens' % n
|
||
|
|
if _fds:
|
||
|
|
os.write(_fds[1], 't' * n)
|
||
|
|
for k in _tokens.keys():
|
||
|
|
_tokens[k] = 0
|
||
|
|
|
||
|
|
|
||
|
|
def _pre_job(r,w):
|
||
|
|
os.putenv('JWACK_PARENT_FD', str(w))
|
||
|
|
os.close(r)
|
||
|
|
|
||
|
|
|
||
|
|
def start_job(argv, stdout=None):
|
||
|
|
global _mytokens
|
||
|
|
setup(1)
|
||
|
|
if stdout:
|
||
|
|
argx = dict(stdout=stdout)
|
||
|
|
else:
|
||
|
|
argx = dict()
|
||
|
|
wait_for_token()
|
||
|
|
r,w = os.pipe()
|
||
|
|
p = subprocess.Popen(argv, preexec_fn=lambda: _pre_job(r,w), **argx)
|
||
|
|
os.close(w)
|
||
|
|
_waitfds[r] = p
|
||
|
|
_tokens[r] = 1
|
||
|
|
return p
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
o = options.Options('jwack', optspec)
|
||
|
|
(opt, flags, extra) = o.parse(sys.argv[1:])
|
||
|
|
|
||
|
|
if not extra:
|
||
|
|
o.fatal("no command line given")
|
||
|
|
|
||
|
|
setup(opt.jobs)
|
||
|
|
try:
|
||
|
|
p = start_job(extra)
|
||
|
|
wait_all()
|
||
|
|
return p.wait()
|
||
|
|
finally:
|
||
|
|
force_return_tokens()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|