Fix QueueFile to support read-to-end semantics and add some tests

This commit is contained in:
Joseph Schorr 2016-03-29 17:01:15 -04:00
parent 2c7aae10a9
commit 6251e63e0e
3 changed files with 121 additions and 16 deletions

View file

@ -1,7 +1,8 @@
from multiprocessing import Process, Queue
from collections import namedtuple
import logging
import multiprocessing
import os
import time
import sys
import traceback
@ -46,23 +47,27 @@ class QueueProcess(object):
args = (self._get_producer, self._queues, self._chunk_size, self._args)
QueueProcess.run_process(_run, args, finished=self._finished)
QueueResult = namedtuple('QueueResult', ['data', 'exception'])
def _run(get_producer, queues, chunk_size, args):
producer = get_producer(*args)
while True:
try:
data = producer(chunk_size) or None
result = QueueResult(producer(chunk_size) or None, None)
except Exception as ex:
message = '%s\n%s' % (ex.message, "".join(traceback.format_exception(*sys.exc_info())))
data = Exception(message)
result = QueueResult(None, Exception(message))
for queue in queues:
try:
queue.put(data, block=True)
queue.put(result, block=True)
except Exception as ex:
logger.exception('Exception writing to queue.')
return
if data is None or isinstance(data, Exception):
# Terminate the producer loop if the data produced is empty or an exception occurred.
if result.data is None or result.exception is not None:
break
# Important! This allows the thread that writes the queue data to the pipe