def submit(self, executor, task, tag=None):
"""Submits a task to a provided executor
:type executor: ibm_s3transfer.futures.BoundedExecutor
:param executor: The executor to submit the callable to
:type task: ibm_s3transfer.tasks.Task
:param task: The task to submit to the executor
:type tag: ibm_s3transfer.futures.TaskTag
:param tag: A tag to associate to the submitted task
:rtype: concurrent.futures.Future
:returns: A future representing the submitted task
"""
logger.debug(
"Submitting task %s to executor %s for transfer request: %s." %
(task, executor, self.transfer_id))
future = executor.submit(task, tag=tag)
# Add this created future to the list of associated future just
# in case it is needed during cleanups.
self.add_associated_future(future)
future.add_done_callback(
FunctionContainer(self.remove_associated_future, future))
return future
def test_done_callbacks_only_ran_once_on_exception(self):
# We want to be able to handle the case where the final task completes
# and anounces done but there is an error in the submission task
# which will cause it to need to anounce done as well. In this case,
# we do not want the failure cleanups to be invoked more than once.
final_task = self.get_task(FailureTask, is_final=True)
self.main_kwargs['executor'] = self.executor
self.main_kwargs['tasks_to_submit'] = [final_task]
submission_task = self.get_task(ExceptionSubmissionTask,
main_kwargs=self.main_kwargs)
# Add the callback to the callbacks to be invoked when the
# transfer fails.
invocations_of_cleanup = []
cleanup_callback = FunctionContainer(invocations_of_cleanup.append,
'cleanup happened')
self.transfer_coordinator.add_failure_cleanup(cleanup_callback)
submission_task()
# Make sure the task failed to start
self.assertEqual(self.transfer_coordinator.status, 'failed')
# Make sure the cleanup was called only onece.
self.assertEqual(invocations_of_cleanup, ['cleanup happened'])
def test_handles_cleanups_submitted_in_other_tasks(self):
invocations_of_cleanup = []
event = Event()
cleanup_callback = FunctionContainer(invocations_of_cleanup.append,
'cleanup happened')
# We want the cleanup to be added in the execution of the task and
# still be executed by the submission task when it fails.
task = self.get_task(SuccessTask,
main_kwargs={
'callbacks': [event.set],
'failure_cleanups': [cleanup_callback]
})
self.main_kwargs['executor'] = self.executor
self.main_kwargs['tasks_to_submit'] = [task]
self.main_kwargs['additional_callbacks'] = [event.wait]
submission_task = self.get_task(ExceptionSubmissionTask,
main_kwargs=self.main_kwargs)
submission_task()
self.assertEqual(self.transfer_coordinator.status, 'failed')
# Make sure the cleanup was called even though the callback got
# added in a completely different task.
self.assertEqual(invocations_of_cleanup, ['cleanup happened'])
def test_add_done_callback(self):
done_callbacks = []
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(return_call_args, 'foo', biz='baz')
wrapped_future = ExecutorFuture(future)
wrapped_future.add_done_callback(
FunctionContainer(done_callbacks.append, 'called'))
self.assertEqual(done_callbacks, ['called'])
def test_repr(self):
func_container = FunctionContainer(self.get_args_kwargs,
'foo',
bar='baz')
self.assertEqual(
str(func_container),
'Function: %s with args %s and kwargs %s' % (self.get_args_kwargs,
('foo', ), {
'bar': 'baz'
}))
def test_waits_for_tasks_submitted_by_other_tasks_on_exception(self):
# In this test, we want to make sure that any tasks that may be
# submitted in another task complete before we start performing
# cleanups.
#
# This is tested by doing the following:
#
# ExecutionSubmissionTask
# |
# +--submits-->SubmitMoreTasksTask
# |
# +--submits-->SuccessTask
# |
# +-->sleeps-->adds failure cleanup
#
# In the end, the failure cleanup of the SuccessTask should be ran
# when the ExecutionSubmissionTask fails. If the
# ExeceptionSubmissionTask did not run the failure cleanup it is most
# likely that it did not wait for the SuccessTask to complete, which
# it needs to because the ExeceptionSubmissionTask does not know
# what failure cleanups it needs to run until all spawned tasks have
# completed.
invocations_of_cleanup = []
event = Event()
cleanup_callback = FunctionContainer(invocations_of_cleanup.append,
'cleanup happened')
cleanup_task = self.get_task(SuccessTask,
main_kwargs={
'callbacks': [event.set],
'failure_cleanups':
[cleanup_callback]
})
task_for_submitting_cleanup_task = self.get_task(SubmitMoreTasksTask,
main_kwargs={
'executor':
self.executor,
'tasks_to_submit':
[cleanup_task]
})
self.main_kwargs['executor'] = self.executor
self.main_kwargs['tasks_to_submit'] = [
task_for_submitting_cleanup_task
]
self.main_kwargs['additional_callbacks'] = [event.wait]
submission_task = self.get_task(ExceptionSubmissionTask,
main_kwargs=self.main_kwargs)
submission_task()
self.assertEqual(self.transfer_coordinator.status, 'failed')
self.assertEqual(invocations_of_cleanup, ['cleanup happened'])
def test_submission_task_announces_done_if_cancelled_before_main(self):
invocations_of_done = []
done_callback = FunctionContainer(invocations_of_done.append,
'done announced')
self.transfer_coordinator.add_done_callback(done_callback)
self.transfer_coordinator.cancel()
submission_task = self.get_task(NOOPSubmissionTask,
main_kwargs=self.main_kwargs)
submission_task()
# Because the submission task was cancelled before being run
# it did not submit any extra tasks so a result it is responsible
# for making sure it announces done as nothing else will.
self.assertEqual(invocations_of_done, ['done announced'])
def test_calls_failure_cleanups_on_exception(self):
submission_task = self.get_task(ExceptionSubmissionTask,
main_kwargs=self.main_kwargs)
# Add the callback to the callbacks to be invoked when the
# transfer fails.
invocations_of_cleanup = []
cleanup_callback = FunctionContainer(invocations_of_cleanup.append,
'cleanup happened')
self.transfer_coordinator.add_failure_cleanup(cleanup_callback)
submission_task()
# Make sure the task failed to start
self.assertEqual(self.transfer_coordinator.status, 'failed')
# Make sure the cleanup was called.
self.assertEqual(invocations_of_cleanup, ['cleanup happened'])
def test_failure_cleanups_on_done(self):
cleanup_invocations = []
callback = FunctionContainer(cleanup_invocations.append,
'cleanup called')
# Add the failure cleanup to the transfer.
self.transfer_coordinator.add_failure_cleanup(callback)
# Announce that the transfer is done. This should invoke the failure
# cleanup.
self.transfer_coordinator.announce_done()
self.assertEqual(cleanup_invocations, ['cleanup called'])
# If done is announced again, we should not invoke the cleanup again
# because done has already been announced and thus the cleanup has
# been ran as well.
self.transfer_coordinator.announce_done()
self.assertEqual(cleanup_invocations, ['cleanup called'])
def test_done_callbacks_on_done(self):
done_callback_invocations = []
callback = FunctionContainer(done_callback_invocations.append,
'done callback called')
# Add the done callback to the transfer.
self.transfer_coordinator.add_done_callback(callback)
# Announce that the transfer is done. This should invoke the done
# callback.
self.transfer_coordinator.announce_done()
self.assertEqual(done_callback_invocations, ['done callback called'])
# If done is announced again, we should not invoke the callback again
# because done has already been announced and thus the callback has
# been ran as well.
self.transfer_coordinator.announce_done()
self.assertEqual(done_callback_invocations, ['done callback called'])
def submit(self, task, tag=None, block=True):
"""Submit a task to complete
:type task: ibm_s3transfer.tasks.Task
:param task: The task to run __call__ on
:type tag: ibm_s3transfer.futures.TaskTag
:param tag: An optional tag to associate to the task. This
is used to override which semaphore to use.
:type block: boolean
:param block: True if to wait till it is possible to submit a task.
False, if not to wait and raise an error if not able to submit
a task.
:returns: The future assocaited to the submitted task
"""
semaphore = self._semaphore
# If a tag was provided, use the semaphore associated to that
# tag.
if tag:
semaphore = self._tag_semaphores[tag]
# Call acquire on the semaphore.
acquire_token = semaphore.acquire(task.transfer_id, block)
# Create a callback to invoke when task is done in order to call
# release on the semaphore.
release_callback = FunctionContainer(semaphore.release,
task.transfer_id, acquire_token)
# Submit the task to the underlying executor.
future = ExecutorFuture(self._executor.submit(task))
# Add the Semaphore.release() callback to the future such that
# it is invoked once the future completes.
future.add_done_callback(release_callback)
return future
def test_call(self):
func_container = FunctionContainer(self.get_args_kwargs,
'foo',
bar='baz')
self.assertEqual(func_container(), (('foo', ), {'bar': 'baz'}))
def _get_final_io_task_submission_callback(self, download_manager,
io_executor):
final_task = download_manager.get_final_io_task()
return FunctionContainer(self._transfer_coordinator.submit,
io_executor, final_task)
def add_failure_cleanup(self, function, *args, **kwargs):
"""Adds a callback to call upon failure"""
with self._failure_cleanups_lock:
self._failure_cleanups.append(
FunctionContainer(function, *args, **kwargs))
def add_done_callback(self, function, *args, **kwargs):
"""Add a done callback to be invoked when transfer is done"""
with self._done_callbacks_lock:
self._done_callbacks.append(
FunctionContainer(function, *args, **kwargs))
def add_done_callback_to_future(self, future, fn, *args, **kwargs):
callback_for_future = FunctionContainer(fn, *args, **kwargs)
future.add_done_callback(callback_for_future)
