class _Client(object):
class RunTask(object):
def __init__(self):
self.pid = Promise()
self.term_info = Promise()
def list_promises(self):
return [self.pid, self.term_info]
class SignalTask(object):
def __init__(self):
self.ack = Promise()
def list_promises(self):
return [self.ack]
def __init__(self, server_pid, channel): # executor_pid, executor_stderr
self._server_pid = server_pid
self._server_exit_status = None
self._channel = channel
self._channel_in = dupfdopen(channel, 'r')
self._channel_out = dupfdopen(channel, 'w')
self._errored = Bool()
self._input_queue = deque()
self._lock = threading.Lock()
self._queue_not_empty = threading.Condition(self._lock)
self._should_stop = Bool()
self._next_task_id = 1
self._tasks = {}
self._fail = self._create_fail()
self._server_stop = Promise()
self._read_thread_inited = threading.Event()
self._write_thread_inited = threading.Event()
write_thread = ProfiledThread(
target=_weak_method(self._write_loop),
name_prefix='RunnerClnWr')
self._write_thread = weakref.ref(write_thread)
read_thread = ProfiledThread(
target=_weak_method(self._read_loop),
name_prefix='RunnerClnRd')
self._read_thread = weakref.ref(read_thread)
write_thread.daemon = True
read_thread.daemon = True
write_thread.start()
read_thread.start()
self._read_thread_inited.wait()
self._write_thread_inited.wait()
def _create_fail(self):
errored = self._errored
queue_not_empty = self._queue_not_empty
channel = self._channel
tasks = self._tasks
# Don't closure on self to prevent cycle
def fail():
with queue_not_empty:
if errored:
return
errored.set()
try:
channel.shutdown(socket.SHUT_RDWR)
except:
pass
exc = ServiceUnavailable("Runner abnormal termination")
tasks_values = tasks.values()
tasks.clear()
for task in tasks_values:
for p in task.list_promises():
if not p.is_set():
try:
p.set(None, exc)
except:
pass
queue_not_empty.notify_all()
return fail
def _wait_stop(self):
for w in [self._read_thread, self._write_thread]:
t = w()
if t:
t.join()
_, status = os.waitpid(self._server_pid, 0)
logging.info('_Server exited with %s' % status)
if status:
raise RuntimeError("Runner.Server process exited abnormally %d" % status)
def stop(self):
do_wait = False
with self._lock:
if not self._should_stop:
self._should_stop.set()
self._queue_not_empty.notify()
do_wait = True
if do_wait:
self._server_stop.run_and_set(self._wait_stop)
self._server_stop.to_future().get()
def __del__(self):
self.stop()
def _get_next_task_id(self):
ret = self._next_task_id
self._next_task_id += 1
return ret
def _enqueue(self, task, msg):
with self._lock:
if self._errored:
raise ServiceUnavailable("Runner in malformed state")
task_id = self._get_next_task_id()
msg.task_id = task_id
self._tasks[task_id] = task
self._input_queue.append(msg)
self._queue_not_empty.notify()
return task_id
def start(self, *pargs, **pkwargs):
task = self.RunTask()
task_id = self._enqueue(task, NewTaskParamsMessage(*pargs, **pkwargs))
def join_pid():
pid = task.pid.to_future().get()
return _Popen(
pid=pid,
task_id=task_id,
exit_status=task.term_info.to_future(),
client=self # FIXME weak
)
return join_pid
def Popen(self, *pargs, **pkwargs):
def start():
return self.start(*pargs, **pkwargs)()
stdin_content = pkwargs.pop('stdin_content', None)
if stdin_content is None:
return start()
# FIXME Why 'stdin_content' supported only in Popen?
# TODO Reimplement: use named pipes?
with NamedTemporaryFile() as tmp:
tmp.write(stdin_content)
tmp.flush()
pkwargs['stdin'] = tmp.name
return start()
def call(self, *args, **kwargs):
return self.Popen(*args, **kwargs).wait()
def check_call(self, *args, **kwargs):
return check_process_call(self.call, args, kwargs) # TODO
def _send_signal(self, target_task, sig, group):
task = self.SignalTask()
self._enqueue(task, SendSignalRequestMessage(target_task._task_id, sig, group))
return task.ack.to_future()
@fail_on_error
def _write_loop(self):
queue = self._input_queue
channel_out = self._channel_out
lock = self._lock
queue_not_empty = self._queue_not_empty
channel = self._channel
should_stop = self._should_stop
errored = self._errored
self._write_thread_inited.set()
del self # don't use self
logging.debug('_Client._write_loop started')
while True:
with lock:
while not (should_stop or queue or errored):
queue_not_empty.wait()
if errored:
raise RuntimeError("Errored in another thread")
messages = []
while queue:
messages.append(queue.popleft())
if should_stop:
messages.append(StopServiceRequestMessage())
logging.debug('_Client._write_loop append(StopServiceRequestMessage)')
for msg in messages:
serialize(channel_out, msg)
channel_out.flush()
if should_stop:
break
channel.shutdown(socket.SHUT_WR)
logging.debug('_Client._write_loop finished')
@fail_on_error
def _read_loop(self):
channel_in = self._channel_in
lock = self._lock
tasks = self._tasks
should_stop = self._should_stop
channel = self._channel
errored = self._errored
self._read_thread_inited.set()
del self # don't use self
stop_response_received = False
logging.debug('_Client._read_loop started')
while True:
try:
msg = deserialize(channel_in)
except EOFError:
logging.debug('_Client._read_loop EOFError')
break
#logging.debug('_Client._read_loop %s' % msg)
if isinstance(msg, ProcessStartMessage):
#logging.debug('_Client._read_loop ProcessStartMessage for task_id=%d, pid=%d, received' % (msg.task_id, msg.pid))
pid, error = msg.pid, msg.error
if error:
#error = RuntimeError(error) # FIXME pickle Exception?
pass # TODO Check! :)
with lock:
try:
task = tasks[msg.task_id] if pid else tasks.pop(msg.task_id)
except KeyError:
if errored:
continue
else:
raise
task.pid.set(pid, error)
elif isinstance(msg, ProcessTerminationMessage):
#logging.info('_Client._read_loop ProcessTerminationMessage for task_id=%d, received' % msg.task_id)
with lock:
try:
task = tasks.pop(msg.task_id)
except KeyError:
if errored: # FIXME Don't remember
continue
else:
raise
# FIXME Better?
p = task.term_info
if isinstance(msg.exit_status, Exception):
p.set(None, msg.exit_status)
else:
p.set(msg.exit_status)
elif isinstance(msg, SendSignalResponseMessage):
with lock:
try:
task = tasks.pop(msg.task_id)
except KeyError:
if errored:
continue
else:
raise
task.ack.set(msg.was_sent, msg.error)
elif isinstance(msg, StopServiceResponseMessage):
logging.debug('_Client._read_loop StopServiceResponseMessage received')
if not should_stop: # FIXME more fine check
raise RuntimeError('StopServiceResponseMessage received but StopServiceRequestMessage was not send')
if tasks:
raise RuntimeError('StopServiceResponseMessage received but not for all process replices was received')
stop_response_received = True
else:
raise RuntimeError('Unknown message type')
if not stop_response_received:
raise RuntimeError('Socket closed without StopServiceResponseMessage')
channel.shutdown(socket.SHUT_RD)
logging.debug('_Client._read_loop finished')
