def setUp(self):
self.config = TestConfig()
self.config.make_default()
self.redis = self.config.get_redis()
self.redis.flushdb()
self.manager = ShireManager(self.config)
self.pool_manager = PoolStatusManager(self.redis)
def __init__(self, config, name, sleep_time=None, max_workhorses=None, verbose=False):
super(Pool, self).__init__()
self.config = config
self.verbose = verbose
self.redis = config.get_redis()
self.name = name
self.pool_logger = create_logger('shire.pool.{}'.format(name))
if verbose:
create_console_handler(self.pool_logger)
self.sleep_time = sleep_time if sleep_time else self.SLEEP_TIME
self.section = self.config.section_getter(self.config.POOL_SECTION)
self.check_time = int(self.section.get(self.config.POOL_CHECK_TIME, self.config.POOL_CHECK_TIME_DEFAULT))
self.max_workhorses = max_workhorses if max_workhorses else self.MAX_WORKHORSES
self._queue_manager = QueueManager(connection=self.redis)
self._children = []
self.log = RedisLogger(config=self.config, pool=self.name)
self._ps = PoolStatusManager(connection=self.redis)
# когда форкуется workhorse, то она наследует signal.signal(signal.SIGTERM, self.terminate)
# и надо проверять в self.terminate кто сейчас завершается
self.i_am_pool = True
class TestPoolBase(TestBase):
def setUp(self):
self.redis.flushall()
self.config[self.config.POOL_SECTION][self.config.POOL_CHECK_TIME] = '1'
self.config.save(self.config_path)
self.pool_status_manager = PoolStatusManager(connection=self.redis)
self.subprocesses = []
def tearDown(self):
for subproc in self.subprocesses:
if hasattr(subproc, 'pid') and self.check_pid(subproc.pid):
subproc.kill()
def _start_subproc(self, *args):
proc = self.popen_cli(*args)
self.subprocesses.append(proc)
return proc
def _job_to_queue(self, pool_name, job_sleep_time):
job = TestSleepJob.delay(config=self.config, pool=pool_name, queue='abc',
kwargs=dict(sleep=job_sleep_time))
redis_queue = QueueManager(connection=self.redis)
redis_queue.push(job.pool, job.id)
return job
def _get_pool_uuid(self, pool_name, timeout=60):
_uuid, iteration_count = None, 0
while not _uuid:
for row in self.pool_status_manager.get_all(pool=pool_name):
_uuid = row.split(':')[-1]
break
time.sleep(0.1)
iteration_count += 1
if iteration_count > timeout / 0.1:
raise Exception("Can't get pool_uuid after {} secs".format(timeout))
return _uuid
class TestShireManager(TestCase):
FIRST_TEST_POOL = 'test1'
SECOND_TEST_POOL = 'test2'
def setUp(self):
self.config = TestConfig()
self.config.make_default()
self.redis = self.config.get_redis()
self.redis.flushdb()
self.manager = ShireManager(self.config)
self.pool_manager = PoolStatusManager(self.redis)
def emulate_pools(self, status=PoolStatusManager.STATUS_ACTIVE):
# Эмулируем пулы в нужном статусе
uuid1, uuid2 = uuid.uuid4(), uuid.uuid4()
self.pool_manager.set_status(self.FIRST_TEST_POOL, uuid1, status)
self.pool_manager.set_status(self.SECOND_TEST_POOL, uuid2, status)
return str(uuid1), str(uuid2)
def pool_status(self, pool, _uuid):
return self.pool_manager.get_status(pool, _uuid)
def test_terminate_once(self):
uuid1, uuid2 = self.emulate_pools()
self.manager.terminate_pools(pools=(self.FIRST_TEST_POOL, ))
self.assertEqual(self.pool_status(self.FIRST_TEST_POOL, uuid1),
PoolStatusManager.STATUS_DEAD,
u'Первый пул в статусе DEAD')
self.assertEqual(self.pool_status(self.SECOND_TEST_POOL, uuid2),
PoolStatusManager.STATUS_ACTIVE,
u'Второй пул работает как обычно')
def test_terminate(self):
uuid1, uuid2 = self.emulate_pools()
self.manager.terminate_pools()
self.assertEqual((
self.pool_status(self.FIRST_TEST_POOL, uuid1),
self.pool_status(self.SECOND_TEST_POOL, uuid2),
), (
PoolStatusManager.STATUS_DEAD,
PoolStatusManager.STATUS_DEAD,
), u'Оба пула в статусе DEAD')
def test_clear_by_status(self):
self.emulate_pools()
self.manager.clean_pools(
from_statuses=(PoolStatusManager.STATUS_ACTIVE, ))
self.assertEqual(self.pool_manager.get_all(), [],
u'Пулы вычищены из redis')
def test_clear(self):
self.emulate_pools(PoolStatusManager.STATUS_TERMITATED)
self.manager.clean_pools()
pools = self.pool_manager.get_all()
self.assertEqual(pools, [], u'Пулы вычищены из redis')
def test_kill_not_dead(self):
# По умолчанию не должны убиваться пулы, которые не имеют статуса DEAD
uuid1, uuid2 = self.emulate_pools()
self.manager.kill_pools(pools=(self.FIRST_TEST_POOL, ))
self.assertEqual(self.pool_status(self.FIRST_TEST_POOL, uuid1),
PoolStatusManager.STATUS_ACTIVE,
u'Первый пул все ещё активен')
def test_kill(self):
uuid1, uuid2 = self.emulate_pools(PoolStatusManager.STATUS_DEAD)
self.manager.kill_pools()
self.assertEqual((
self.pool_status(self.FIRST_TEST_POOL, uuid1),
self.pool_status(self.SECOND_TEST_POOL, uuid2),
), (
PoolStatusManager.STATUS_KILL,
PoolStatusManager.STATUS_KILL,
), u'Оба пула в статусе KILL')
def test_status(self):
uuid1, uuid2 = self.emulate_pools()
self.pool_manager.set_status(self.SECOND_TEST_POOL, uuid2,
PoolStatusManager.STATUS_DEAD)
by_pool = {x['name']: x for x in self.manager.get_status()}
self.assertEqual(
by_pool[self.FIRST_TEST_POOL], {
'status': PoolStatusManager.STATUS_ACTIVE,
'name': self.FIRST_TEST_POOL,
'uuid': uuid1
}, u'Информация по первому пулу верна')
self.assertEqual(
by_pool[self.SECOND_TEST_POOL], {
'status': PoolStatusManager.STATUS_DEAD,
'name': self.SECOND_TEST_POOL,
'uuid': uuid2
}, u'Информация по второму пулу верна')
class Pool(Daemon):
SLEEP_TIME = 1
MAX_WORKHORSES = None # без ограничения
def __init__(self, config, name, sleep_time=None, max_workhorses=None, verbose=False):
super(Pool, self).__init__()
self.config = config
self.verbose = verbose
self.redis = config.get_redis()
self.name = name
self.pool_logger = create_logger('shire.pool.{}'.format(name))
if verbose:
create_console_handler(self.pool_logger)
self.sleep_time = sleep_time if sleep_time else self.SLEEP_TIME
self.section = self.config.section_getter(self.config.POOL_SECTION)
self.check_time = int(self.section.get(self.config.POOL_CHECK_TIME, self.config.POOL_CHECK_TIME_DEFAULT))
self.max_workhorses = max_workhorses if max_workhorses else self.MAX_WORKHORSES
self._queue_manager = QueueManager(connection=self.redis)
self._children = []
self.log = RedisLogger(config=self.config, pool=self.name)
self._ps = PoolStatusManager(connection=self.redis)
# когда форкуется workhorse, то она наследует signal.signal(signal.SIGTERM, self.terminate)
# и надо проверять в self.terminate кто сейчас завершается
self.i_am_pool = True
def pool_log(self, msg, level='info'):
if not self.verbose:
return
getattr(self.pool_logger, level)(msg)
def run(self):
self.pool_log('Pool "{}" started'.format(self.name))
self.pool_log('Max workhorses: {}'.format('unlimited' if self.max_workhorses is None else self.max_workhorses))
self.pool_log('Redis check timeout: {}s'.format(self.check_time))
self._ps.set_status(pool=self.name, _uuid=self.uuid, status=PoolStatusManager.STATUS_ACTIVE)
signal.signal(signal.SIGTERM, self.terminate)
while True:
if not self.can_start_new_workhorse():
time.sleep(self.sleep_time)
continue
while True:
job_id = self._queue_manager.pop(pool=self.name, timeout=self.check_time)
if self.status in (PoolStatusManager.STATUS_DEAD, PoolStatusManager.STATUS_KILL):
if job_id:
self._queue_manager.push(pool=self.name, job_id=job_id, to_tail=True)
if self.status == PoolStatusManager.STATUS_KILL:
self.kill_children()
self.terminate()
if job_id:
# бывает None при выходе из brpop по таймауту
self._start_workhorse(job_id) # тут пролетает sys.exit() от workhorse
if not self.can_start_new_workhorse():
# проверяем после запуска задачи, что бы понять - можно ли выбирать другие задачи
break
def _start_workhorse(self, job_id):
self.pool_log('Workhorse for job #{} started'.format(job_id))
workhorse = Workhorse(pool=self, job_id=job_id)
pid = workhorse.fork()
self._children.append(pid)
@property
def status(self):
return self._ps.get_status(pool=self.name, _uuid=self.uuid)
def wait_children(self):
while self._get_children_count() > 0:
time.sleep(self.sleep_time)
if self.status == PoolStatusManager.STATUS_KILL:
self.kill_children()
return
def kill_children(self):
for child_pid in self._children:
if check_pid_is_shire(child_pid):
os.kill(child_pid, signal.SIGKILL)
self._children = []
def can_start_new_workhorse(self):
children_count = self._get_children_count()
if self.max_workhorses:
return children_count < self.max_workhorses
return True
def _get_children_count(self):
ended_pids = []
try:
while True:
pid, ret_code = os.waitpid(-1, os.WNOHANG)
if not pid:
break
ended_pids.append(pid)
except OSError:
# нет child-процессов
pass
for child_pid in self._children[:]:
if child_pid in ended_pids:
self._children.remove(child_pid)
return len(self._children)
def terminate(self, sign=None, frame=None):
if self.i_am_pool:
self.wait_children()
self._ps.set_status(pool=self.name, _uuid=self.uuid, status=PoolStatusManager.STATUS_TERMITATED)
sys.exit(0)
class ShireManager(object):
CONST_ALL = '__all__'
EXPORT_METHODS = ('get_status', 'terminate_pools', 'kill_pools',
'clean_pools')
def __init__(self, config):
self.config = config
self.status_manager = PoolStatusManager(self.config.get_redis())
def _prepare_statuses(self, statuses):
if statuses == self.CONST_ALL:
statuses = (PoolStatusManager.STATUS_DEAD,
PoolStatusManager.STATUS_KILL,
PoolStatusManager.STATUS_ACTIVE,
PoolStatusManager.STATUS_TERMITATED)
return set(statuses)
def get_pools(self):
# TODO: Добавить опциональный параметр get_extra=False, при получении которого загружать так же из
# базы данных кол-во активных, ожидающих и выполненных задач, а так же время последнего обновления задачи
for key in self.status_manager.get_all():
try:
prefix, pool, uuid = key.rsplit(':', 2)
except ValueError:
# Плохое значение в get_all. Кейс скорее невозможный, проверка на всякий
continue
yield {
'name': pool,
'uuid': uuid,
'status': self.status_manager.get_status(pool, uuid)
}
def get_status(self, pools=CONST_ALL, from_statuses=CONST_ALL):
statuses = self._prepare_statuses(from_statuses)
check_name = (lambda x: True) if pools == self.CONST_ALL else (
lambda x: x in pools)
for pool in self.get_pools():
if pool['status'] in statuses and check_name(pool['name']):
yield pool
def update_pool_statuses(self,
status,
pools=CONST_ALL,
from_statuses=CONST_ALL):
statuses = self._prepare_statuses(from_statuses)
check_name = (lambda x: True) if pools == self.CONST_ALL else (
lambda x: x in pools)
for pool in self.get_pools():
if pool['status'] in statuses and check_name(pool['name']):
self.status_manager.set_status(pool['name'], pool['uuid'],
status)
def terminate_pools(self,
pools=CONST_ALL,
from_statuses=(PoolStatusManager.STATUS_ACTIVE, )):
# Останавливаем пулы, посылая им статус DEAD
self.update_pool_statuses(PoolStatusManager.STATUS_DEAD,
pools=pools,
from_statuses=from_statuses)
def kill_pools(self,
pools=CONST_ALL,
from_statuses=(PoolStatusManager.STATUS_DEAD, )):
# Приказывает пулам принудительно уничтожить всех своих потомков, посылая им статус KILL
self.update_pool_statuses(PoolStatusManager.STATUS_KILL,
pools=pools,
from_statuses=from_statuses)
def clean_pools(self,
from_statuses=(PoolStatusManager.STATUS_TERMITATED, ),
**kwargs):
# Очищает записи в redis от устаревших uuid
statuses = self._prepare_statuses(from_statuses)
for pool in self.get_pools():
if pool['status'] in statuses:
self.status_manager.del_status(pool['name'], pool['uuid'])
def run_command(self, command, *args, **kwargs):
if command in self.EXPORT_METHODS:
return getattr(self, command)(*args, **kwargs)
def __init__(self, config):
self.config = config
self.status_manager = PoolStatusManager(self.config.get_redis())
def setUp(self):
self.redis.flushall()
self.config[self.config.POOL_SECTION][self.config.POOL_CHECK_TIME] = '1'
self.config.save(self.config_path)
self.pool_status_manager = PoolStatusManager(connection=self.redis)
self.subprocesses = []