def poll_task_completion(self,
timeout,
timestamp_identifier,
task_log=None):
# type: (float, str, str) -> None
"""
Poll for the task to complete
:param timeout: Stop polling after a set timeout
:type timeout: float
:param task_log: Task logging string. Default to 'Task <task_name>'
:type task_log: str
:param timestamp_identifier: Identifier with a timestamp. Format <timestamp>_<identifier>
Used to check if the task to poll still exists
:type timestamp_identifier: str
:return:
"""
# @todo better to use a thread for this?
task_log = task_log or 'Task {0}'.format(
self.ensure_single_container.task_name)
# Let's wait for 2nd job in queue to have finished if no callback provided
slept = 0
while slept < timeout:
self.logger.info(
'{0} is waiting for similar tasks to finish - ({1})'.format(
task_log, slept + self.poll_sleep_time))
if not self.is_task_still_registered(timestamp_identifier):
self.discard_task_similar_jobs()
slept += self.poll_sleep_time
time.sleep(self.poll_sleep_time)
if slept >= timeout:
self.logger.error(
'{0} waited {1}s for similar tasks to finish, but timeout was reached'
.format(task_log, slept))
exception_message = 'Could not start within timeout of {0}s while waiting for other tasks'.format(
timeout)
self.unittest_set_state_exception(exception_message)
timeout_message = '{0} - {1} could not be started within timeout of {2}s'.format(
self.message, task_log, timeout)
raise EnsureSingleTimeoutReached(timeout_message)
self.unittest_set_state_waiting()
def new_function(*args, **kwargs):
"""
Wrapped function
:param args: Arguments without default values
:param kwargs: Arguments with default values
"""
def log_message(message, level='info'):
"""
Log a message with some additional information
:param message: Message to log
:param level: Log level
:return: None
"""
if level not in ('info', 'warning', 'debug', 'error'):
raise ValueError(
'Unsupported log level "{0}" specified'.format(level))
complete_message = 'Ensure single {0} mode - ID {1} - {2}'.format(
mode, now, message)
getattr(logger, level)(complete_message)
def update_value(key, append, value_to_store=None):
"""
Store the specified value in the PersistentFactory
:param key: Key to store the value for
:param append: If True, the specified value will be appended else element at index 0 will be popped
:param value_to_store: Value to append to the list
:return: Updated value
"""
with VolatileMutex(name=key, wait=5):
if persistent_client.exists(key):
val = persistent_client.get(key)
if append is True and value_to_store is not None:
val['values'].append(value_to_store)
elif append is False and len(val['values']) > 0:
val['values'].pop(0)
else:
log_message('Setting initial value for key {0}'.format(
persistent_key))
val = {'mode': mode, 'values': []}
persistent_client.set(key, val)
return val
now = '{0}_{1}'.format(
int(time.time()), ''.join(
random.choice(string.ascii_letters + string.digits)
for _ in range(10)))
task_names = [
task_name
] if extra_task_names is None else [task_name] + extra_task_names
persistent_key = '{0}_{1}'.format(ENSURE_SINGLE_KEY, task_name)
persistent_client = PersistentFactory.get_client()
if mode == 'DEFAULT':
with VolatileMutex(persistent_key, wait=5):
for task in task_names:
key_to_check = '{0}_{1}'.format(
ENSURE_SINGLE_KEY, task)
if persistent_client.exists(key_to_check):
log_message(
'Execution of task {0} discarded'.format(
task_name))
return None
log_message('Setting key {0}'.format(persistent_key))
persistent_client.set(persistent_key, {'mode': mode})
try:
output = function(*args, **kwargs)
log_message(
'Task {0} finished successfully'.format(task_name))
return output
finally:
with VolatileMutex(persistent_key, wait=5):
if persistent_client.exists(persistent_key):
log_message(
'Deleting key {0}'.format(persistent_key))
persistent_client.delete(persistent_key)
elif mode == 'CHAINED':
if extra_task_names is not None:
log_message('Extra tasks are not allowed in this mode',
level='error')
raise ValueError(
'Ensure single {0} mode - ID {1} - Extra tasks are not allowed in this mode'
.format(mode, now))
# 1. Create key to be stored in arakoon and update kwargs with args
timeout = kwargs.pop(
'chain_timeout'
) if 'chain_timeout' in kwargs else global_timeout
function_info = inspect.getargspec(function)
kwargs_dict = {}
for index, arg in enumerate(args):
kwargs_dict[function_info.args[index]] = arg
kwargs_dict.update(kwargs)
params_info = 'with params {0}'.format(
kwargs_dict) if kwargs_dict else 'with default params'
# 2. Set the key in arakoon if non-existent
value = update_value(key=persistent_key, append=True)
# 3. Validate whether another job with same params is being executed, skip if so
for item in value['values'][
1:]: # 1st element is processing job, we check all other queued jobs for identical params
if item['kwargs'] == kwargs_dict:
log_message(
'Execution of task {0} {1} discarded because of identical parameters'
.format(task_name, params_info))
return None
log_message('New task {0} {1} scheduled for execution'.format(
task_name, params_info))
update_value(key=persistent_key,
append=True,
value_to_store={
'kwargs': kwargs_dict,
'timestamp': now
})
# 4. Poll the arakoon to see whether this call is the first in list, if so --> execute, else wait
first_element = None
counter = 0
while first_element != now and counter < timeout:
if persistent_client.exists(persistent_key):
value = persistent_client.get(persistent_key)
first_element = value['values'][0]['timestamp']
if first_element == now:
try:
if counter != 0:
current_time = int(time.time())
starting_time = int(now.split('_')[0])
log_message(
'Task {0} {1} had to wait {2} seconds before being able to start'
.format(task_name, params_info,
current_time - starting_time))
output = function(*args, **kwargs)
log_message(
'Task {0} finished successfully'.format(
task_name))
finally:
update_value(key=persistent_key, append=False)
return output
counter += 1
time.sleep(1)
if counter == timeout:
update_value(key=persistent_key, append=False)
log_message(
'Could not start task {0} {1}, within expected time ({2}s). Removed it from queue'
.format(task_name, params_info, timeout),
level='error')
raise EnsureSingleTimeoutReached(
'Ensure single {0} mode - ID {1} - Task {2} could not be started within timeout of {3}s'
.format(mode, now, task_name, timeout))
else:
raise ValueError(
'Unsupported mode "{0}" provided'.format(mode))
def new_function(self, *args, **kwargs):
"""
Wrapped function
:param self: With bind=True, the celery task result itself is passed in
:param args: Arguments without default values
:param kwargs: Arguments with default values
"""
def log_message(message, level='info'):
"""
Log a message with some additional information
:param message: Message to log
:param level: Log level
:return: None
"""
if level not in ('info', 'warning', 'debug', 'error', 'exception'):
raise ValueError('Unsupported log level "{0}" specified'.format(level))
if unittest_mode is False:
complete_message = 'Ensure single {0} mode - ID {1} - {2}'.format(mode, now, message)
else:
complete_message = 'Ensure single {0} mode - ID {1} - {2} - {3}'.format(mode, now, threading.current_thread().getName(), message)
getattr(logger, level)(complete_message)
def update_value(key, append, value_to_update=None):
"""
Store the specified value in the PersistentFactory
:param key: Key to store the value for
:param append: If True, the specified value will be appended else element at index 0 will be popped
:param value_to_update: Value to append to the list or remove from the list
:return: Updated value
"""
with volatile_mutex(name=key, wait=5):
vals = list(persistent_client.get_multi([key], must_exist=False))
if vals[0] is not None:
val = vals[0]
if append is True and value_to_update is not None:
val['values'].append(value_to_update)
elif append is False and value_to_update is not None:
for value_item in val['values']:
if value_item == value_to_update:
val['values'].remove(value_item)
break
elif append is False and len(val['values']) > 0:
val['values'].pop(0)
else:
log_message('Setting initial value for key {0}'.format(key))
val = {'mode': mode,
'values': []}
persistent_client.set(key, val)
return val
if not hasattr(self, 'request'):
raise RuntimeError('The decorator ensure_single can only be applied to bound tasks (with bind=True argument)')
now = '{0}_{1}'.format(int(time.time()), ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(10)))
task_id = self.request.id
async_task = task_id is not None # Async tasks have an ID, inline executed tasks have None as ID
task_names = [task_name] if extra_task_names is None else [task_name] + extra_task_names
thread_name = threading.current_thread().getName()
unittest_mode = os.environ.get('RUNNING_UNITTESTS') == 'True'
persistent_key = '{0}_{1}'.format(ENSURE_SINGLE_KEY, task_name)
persistent_client = PersistentFactory.get_client()
if mode == 'DEFAULT':
with volatile_mutex(persistent_key, wait=5):
for task in task_names:
key_to_check = '{0}_{1}'.format(ENSURE_SINGLE_KEY, task)
if persistent_client.exists(key_to_check):
if async_task is True or callback is None:
log_message('Execution of task {0} discarded'.format(task_name))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('DISCARDED', None)
return None
else:
log_message('Execution of task {0} in progress, executing callback function'.format(task_name))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('CALLBACK', None)
return callback(*args, **kwargs)
log_message('Setting key {0}'.format(persistent_key))
persistent_client.set(persistent_key, {'mode': mode,
'values': [{'task_id': task_id}]})
try:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXECUTING', None)
output = f(*args, **kwargs)
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('FINISHED', None)
Decorators.unittest_thread_info_by_state['FINISHED'].append(thread_name)
log_message('Task {0} finished successfully'.format(task_name))
return output
finally:
with volatile_mutex(persistent_key, wait=5):
log_message('Deleting key {0}'.format(persistent_key))
persistent_client.delete(persistent_key, must_exist=False)
elif mode == 'DEDUPED':
if extra_task_names is not None:
log_message('Extra tasks are not allowed in this mode',
level='error')
raise ValueError('Ensure single {0} mode - ID {1} - Extra tasks are not allowed in this mode'.format(mode, now))
# Update kwargs with args
sleep = 1 if unittest_mode is False else 0.1
timeout = kwargs.pop('ensure_single_timeout', 10 if unittest_mode is True else global_timeout)
function_info = inspect.getargspec(f)
kwargs_dict = {}
for index, arg in enumerate(args):
kwargs_dict[function_info.args[index]] = arg
kwargs_dict.update(kwargs)
params_info = 'with params {0}'.format(kwargs_dict) if kwargs_dict else 'with default params'
# Set the key in arakoon if non-existent
value = update_value(key=persistent_key,
append=True)
# Validate whether another job with same params is being executed
job_counter = 0
for item in value['values']:
if item['kwargs'] == kwargs_dict:
job_counter += 1
if job_counter == 2: # 1st job with same params is being executed, 2nd is scheduled for execution ==> Discard current
if async_task is True: # Not waiting for other jobs to finish since asynchronously
log_message('Execution of task {0} {1} discarded because of identical parameters'.format(task_name, params_info))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('DISCARDED', None)
return None
# If executed inline (sync), execute callback if any provided
if callback is not None:
log_message('Execution of task {0} {1} in progress, executing callback function'.format(task_name, params_info))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('CALLBACK', None)
return callback(*args, **kwargs)
# Let's wait for 2nd job in queue to have finished if no callback provided
slept = 0
while slept < timeout:
log_message('Task {0} {1} is waiting for similar tasks to finish - ({2})'.format(task_name, params_info, slept + sleep))
values = list(persistent_client.get_multi([persistent_key], must_exist=False))
if values[0] is None:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITED', None)
return None # All pending jobs have been deleted in the meantime, no need to wait
if item['timestamp'] not in [value['timestamp'] for value in values[0]['values']]:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITED', None)
return None # Similar tasks have been executed, so sync task currently waiting can return without having been executed
slept += sleep
time.sleep(sleep)
if slept >= timeout:
log_message('Task {0} {1} waited {2}s for similar tasks to finish, but timeout was reached'.format(task_name, params_info, slept),
level='error')
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXCEPTION', 'Could not start within timeout of {0}s while waiting for other tasks'.format(timeout))
raise EnsureSingleTimeoutReached('Ensure single {0} mode - ID {1} - Task {2} could not be started within timeout of {3}s'.format(mode,
now,
task_name,
timeout))
if unittest_mode is True:
if thread_name not in Decorators.unittest_thread_info_by_state['WAITING']:
Decorators.unittest_thread_info_by_state['WAITING'].append(thread_name)
log_message('New task {0} {1} scheduled for execution'.format(task_name, params_info))
update_value(key=persistent_key,
append=True,
value_to_update={'kwargs': kwargs_dict,
'task_id': task_id,
'timestamp': now})
# Poll the arakoon to see whether this call is the only in list, if so --> execute, else wait
slept = 0
while slept < timeout:
values = list(persistent_client.get_multi([persistent_key], must_exist=False))
if values[0] is not None:
queued_jobs = [v for v in values[0]['values'] if v['kwargs'] == kwargs_dict]
if len(queued_jobs) != 1:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITING', None)
if thread_name not in Decorators.unittest_thread_info_by_state['WAITING']:
Decorators.unittest_thread_info_by_state['WAITING'].append(thread_name)
else:
try:
if slept != 0:
log_message('Task {0} {1} had to wait {2} seconds before being able to start'.format(task_name,
params_info,
slept))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXECUTING', None)
output = f(*args, **kwargs)
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('FINISHED', None)
Decorators.unittest_thread_info_by_state['FINISHED'].append(thread_name)
log_message('Task {0} finished successfully'.format(task_name))
return output
finally:
update_value(key=persistent_key,
append=False,
value_to_update=queued_jobs[0])
slept += sleep
time.sleep(sleep)
if slept >= timeout:
update_value(key=persistent_key,
append=False,
value_to_update={'kwargs': kwargs_dict,
'task_id': task_id,
'timestamp': now})
log_message('Could not start task {0} {1}, within expected time ({2}s). Removed it from queue'.format(task_name, params_info, timeout),
level='error')
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXCEPTION', 'Could not start within timeout of {0}s while queued'.format(timeout))
raise EnsureSingleTimeoutReached('Ensure single {0} mode - ID {1} - Task {2} could not be started within timeout of {3}s'.format(mode,
now,
task_name,
timeout))
elif mode == 'CHAINED':
if extra_task_names is not None:
log_message('Extra tasks are not allowed in this mode',
level='error')
raise ValueError('Ensure single {0} mode - ID {1} - Extra tasks are not allowed in this mode'.format(mode, now))
# Update kwargs with args
sleep = 1 if unittest_mode is False else 0.1
timeout = kwargs.pop('ensure_single_timeout', 10 if unittest_mode is True else global_timeout)
function_info = inspect.getargspec(f)
kwargs_dict = {}
for index, arg in enumerate(args):
kwargs_dict[function_info.args[index]] = arg
kwargs_dict.update(kwargs)
params_info = 'with params {0}'.format(kwargs_dict) if kwargs_dict else 'with default params'
# Set the key in arakoon if non-existent
value = update_value(key=persistent_key,
append=True)
# Validate whether another job with same params is being executed, skip if so
for item in value['values'][1:]: # 1st element is processing job, we check all other queued jobs for identical params
if item['kwargs'] == kwargs_dict:
if async_task is True: # Not waiting for other jobs to finish since asynchronously
log_message('Execution of task {0} {1} discarded because of identical parameters'.format(task_name, params_info))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('DISCARDED', None)
return None
# If executed inline (sync), execute callback if any provided
if callback is not None:
log_message('Execution of task {0} {1} in progress, executing callback function'.format(task_name, params_info))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('CALLBACK', None)
return callback(*args, **kwargs)
# Let's wait for 2nd job in queue to have finished if no callback provided
slept = 0
while slept < timeout:
log_message('Task {0} {1} is waiting for similar tasks to finish - ({2})'.format(task_name, params_info, slept + sleep))
values = list(persistent_client.get_multi([persistent_key], must_exist=False))
if values[0] is None:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITED', None)
return None # All pending jobs have been deleted in the meantime, no need to wait
if item['timestamp'] not in [value['timestamp'] for value in values[0]['values']]:
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITED', None)
return None # Similar tasks have been executed, so sync task currently waiting can return without having been executed
slept += sleep
time.sleep(sleep)
if slept >= timeout:
log_message('Task {0} {1} waited {2}s for similar tasks to finish, but timeout was reached'.format(task_name, params_info, slept),
level='error')
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXCEPTION', 'Could not start within timeout of {0}s while waiting for other tasks'.format(timeout))
raise EnsureSingleTimeoutReached('Ensure single {0} mode - ID {1} - Task {2} could not be started within timeout of {3}s'.format(mode,
now,
task_name,
timeout))
if unittest_mode is True:
if thread_name not in Decorators.unittest_thread_info_by_state['WAITING']:
Decorators.unittest_thread_info_by_state['WAITING'].append(thread_name)
log_message('New task {0} {1} scheduled for execution'.format(task_name, params_info))
update_value(key=persistent_key,
append=True,
value_to_update={'kwargs': kwargs_dict,
'task_id': task_id,
'timestamp': now})
# Poll the arakoon to see whether this call is the first in list, if so --> execute, else wait
first_element = None
slept = 0
while slept < timeout:
values = list(persistent_client.get_multi([persistent_key], must_exist=False))
if values[0] is not None:
value = values[0]
first_element = value['values'][0]['timestamp'] if len(value['values']) > 0 else None
if first_element == now:
try:
if slept > 0:
log_message('Task {0} {1} had to wait {2} seconds before being able to start'.format(task_name,
params_info,
slept))
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXECUTING', None)
output = f(*args, **kwargs)
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('FINISHED', None)
Decorators.unittest_thread_info_by_state['FINISHED'].append(thread_name)
log_message('Task {0} finished successfully'.format(task_name))
return output
finally:
update_value(key=persistent_key,
append=False)
else:
if unittest_mode is True:
if thread_name not in Decorators.unittest_thread_info_by_state['WAITING']:
Decorators.unittest_thread_info_by_name[thread_name] = ('WAITING', None)
Decorators.unittest_thread_info_by_state['WAITING'].append(thread_name)
slept += sleep
time.sleep(sleep)
if slept >= timeout:
update_value(key=persistent_key,
append=False,
value_to_update={'kwargs': kwargs_dict,
'task_id': task_id,
'timestamp': now})
log_message('Could not start task {0} {1}, within expected time ({2}s). Removed it from queue'.format(task_name, params_info, timeout),
level='error')
if unittest_mode is True:
Decorators.unittest_thread_info_by_name[thread_name] = ('EXCEPTION', 'Could not start within timeout of {0}s while queued'.format(timeout))
raise EnsureSingleTimeoutReached('Ensure single {0} mode - ID {1} - Task {2} could not be started within timeout of {3}s'.format(mode,
now,
task_name,
timeout))
else:
raise ValueError('Unsupported mode "{0}" provided'.format(mode))
def lock_deduped_mode(self, kwargs, timeout):
# type: (dict, float) -> None
"""
Lock the function in deduped mode
:param kwargs: Dict containing all arguments as key word arguments
:type kwargs: dict
:param timeout: Polling timeout in seconds
:type timeout: float
:raises: EnsureSingleTaskDiscarded: If the task was discarded
:raises: EnsureSingleDoCallBack: If the task is required to call the callback function instead of the task function
:raises: EnsureSingleSimilarJobsCompleted: If the task was waiting on a similar task to end but those tasks have finished
"""
self.validate_no_extra_names()
kwargs_dict = self._filter_ignorable_arguments(kwargs)
params_info = 'with params {0} (all passed arguments: {1}'.format(
kwargs_dict, kwargs) if kwargs_dict else 'with default params'
task_log_format = 'task {0} {1}'
task_log_name = task_log_format.format(
self.ensure_single_container.task_name, params_info)
task_log_id = task_log_format.format(self.task_id, params_info)
# Acquire
try:
self.run_hook('before_validation')
# Validate whether another job with same params is being executed
# 1st job with same params is being executed, 2nd is scheduled for execution ==> Discard current
self._ensure_job_limit(kwargs_dict,
timeout, (task_log_name, task_log_id),
job_limit=2)
self.logger.info(
'New {0} scheduled for execution'.format(task_log_name))
new_task_data = {
'kwargs': kwargs_dict,
'task_id': self.task_id,
'timestamp': self.now
}
self._register_task(new_task_data)
# Poll the arakoon to see whether this call is the only in list, if so --> execute, else wait
slept = 0
while slept < timeout:
current_registrations, initial_registrations = self.get_task_registrations(
)
queued_jobs = [
t_d for t_d in current_registrations
if t_d['kwargs'] == kwargs_dict
]
if len(queued_jobs) == 1:
# The only queued job. No more need to poll
break
self.unittest_set_state_waiting()
slept += self.poll_sleep_time
time.sleep(self.poll_sleep_time)
finally:
self.run_hook('after_validation')
successful = True
try:
if slept:
if slept >= timeout:
self.logger.error(
'Could not start {0}, within expected time ({1}s). Removed it from queue'
.format(task_log_name, timeout))
self.unittest_set_state_exception(
'Could not start within timeout of {0}s while queued'.
format(timeout))
timeout_message = '{0} - Task {1} could not be started within timeout of {2}s'.format(
self.message, self.ensure_single_container.task_name,
timeout)
raise EnsureSingleTimeoutReached(timeout_message)
else:
self.logger.info(
'{0} had to wait {1} seconds before being able to start'
.format(task_log_name, slept))
self.run_hook('before_execution')
self.unittest_set_state_executing()
yield
# Release
except:
successful = False
raise
finally:
if successful:
self.unittest_set_state_finished()
self.logger.info('Task {0} finished successfully'.format(
self.ensure_single_container.task_name))
self._unregister_task(new_task_data)
self.run_hook('after_execution')