def testCreateTask(self):
"""Tests the CreateTask function."""
manager = task_manager.TaskManager()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(manager._total_number_of_tasks, 0)
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(manager._total_number_of_tasks, 1)
self.assertIsInstance(task, tasks.Task)
self.assertEqual(task.session_identifier,
self._TEST_SESSION_IDENTIFIER)
self.assertIsNotNone(task.identifier)
self.assertIsNone(task.last_processing_time)
def testAbandonInactiveProcessingTasks(self):
"""Tests the _AbandonInactiveProcessingTasks function."""
manager = task_manager.TaskManager()
test_tasks = []
for _ in range(2):
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertIsNotNone(task.last_processing_time)
test_tasks.append(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 2)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 2)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
for task in test_tasks:
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 2)
def testCreateRetryTask(self):
"""Tests the CreateRetryTask function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(manager._total_number_of_tasks, 1)
retry_task = manager.CreateRetryTask()
self.assertIsNone(retry_task)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
retry_task = manager.CreateRetryTask()
self.assertIsNotNone(retry_task)
self.assertNotEqual(retry_task, task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_abandoned), 1)
self.assertEqual(manager._total_number_of_tasks, 2)
def testUpdateLatestProcessingTime(self):
"""Tests the _UpdateLatestProcessingTime function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.UpdateProcessingTime()
manager._UpdateLatestProcessingTime(task)
def testGetStatusInformation(self):
"""Tests the GetStatusInformation function."""
manager = task_manager.TaskManager()
manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
result_status = manager.GetStatusInformation()
self.assertIsNotNone(result_status)
self.assertEqual(result_status.number_of_abandoned_tasks, 0)
self.assertEqual(result_status.number_of_queued_tasks, 1)
self.assertEqual(result_status.number_of_tasks_pending_merge, 0)
self.assertEqual(result_status.number_of_tasks_processing, 0)
self.assertEqual(result_status.total_number_of_tasks, 1)
def testGetProcessedTaskByIdentifier(self):
"""Tests the GetProcessedTaskByIdentifier function."""
manager = task_manager.TaskManager()
# Test with queued task.
test_tasks = []
for _ in range(3):
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
test_tasks.append(task)
self.assertEqual(len(manager._tasks_queued), 3)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
task = test_tasks[0]
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test with processing task.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test with abandoned task.
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test status of task is unknown.
task = tasks.Task()
with self.assertRaises(KeyError):
manager.GetProcessedTaskByIdentifier(task.identifier)
def testHasTasksPendingRetry(self):
"""Tests the _HasTasksPendingRetry function."""
manager = task_manager.TaskManager()
for _ in range(2):
manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
result = manager._HasTasksPendingRetry()
self.assertFalse(result)
manager._AbandonQueuedTasks()
result = manager._HasTasksPendingRetry()
self.assertTrue(result)
def testGetTaskPendingMerge(self):
"""Tests the GetTaskPendingMerge function."""
current_task = tasks.Task()
current_task.storage_file_size = 10
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge when _tasks_pending_merge is
# empty.
result_task = manager.GetTaskPendingMerge(current_task)
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge when there is a current task
# with a higher merge priority.
current_task.merge_priority = 1
task.merge_priority = 1000
result_task = manager.GetTaskPendingMerge(current_task)
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge.
result_task = manager.GetTaskPendingMerge(None)
self.assertIsNotNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 1)
def testHasTasksPendingMerge(self):
"""Tests the _HasTasksPendingMerge function."""
manager = task_manager.TaskManager()
test_tasks = []
for _ in range(2):
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
test_tasks.append(task)
result = manager._HasTasksPendingMerge()
self.assertFalse(result)
for task in test_tasks:
manager.UpdateTaskAsPendingMerge(task)
result = manager._HasTasksPendingMerge()
self.assertTrue(result)
def testGetFailedTasks(self):
"""Tests the GetFailedTasks function."""
manager = task_manager.TaskManager()
test_tasks = [
manager.CreateTask(self._TEST_SESSION_IDENTIFIER) for _ in range(2)
]
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_abandoned), 0)
result_tasks = manager.GetFailedTasks()
self.assertEqual(result_tasks, [])
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 2)
result_tasks = manager.GetFailedTasks()
self.assertEqual(set(result_tasks), set(test_tasks))
def testAbandonQueuedTasks(self):
"""Tests the _AbandonQueuedTasks function."""
manager = task_manager.TaskManager()
test_tasks = [
manager.CreateTask(self._TEST_SESSION_IDENTIFIER) for _ in range(2)
]
for task in test_tasks:
self.assertIsNotNone(task.start_time)
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 2)
def testRemoveTask(self):
"""Tests the RemoveTask function."""
manager = task_manager.TaskManager()
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
with self.assertRaises(KeyError):
manager.RemoveTask(task)
# Test with abandoned task that was not retried (has no retry task).
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
with self.assertRaises(KeyError):
manager.RemoveTask(task)
# Test with abandoned task that was retried (has retry task).
task.has_retry = True
manager.RemoveTask(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
def testGetTaskPendingRetry(self):
"""Tests the _GetTaskPendingRetry function."""
manager = task_manager.TaskManager()
test_tasks = [
manager.CreateTask(self._TEST_SESSION_IDENTIFIER) for _ in range(2)
]
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_abandoned), 0)
result_task = manager._GetTaskPendingRetry()
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 2)
result_task = manager._GetTaskPendingRetry()
self.assertIn(result_task, test_tasks)
def testUpdateTaskAsPendingMerge(self):
"""Tests the UpdateTaskAsPendingMerge function."""
manager = task_manager.TaskManager()
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test with processing task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 2)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test with abandoned task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 2)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 2)
self.assertEqual(len(manager._tasks_abandoned), 1)
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 3)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test status of task is unknown.
task = tasks.Task()
task.storage_file_size = 10
with self.assertRaises(KeyError):
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
# Test with abandoned task that is retried.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 3)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 3)
self.assertEqual(len(manager._tasks_abandoned), 1)
retry_task = manager.CreateRetryTask()
self.assertIsNotNone(retry_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 3)
self.assertEqual(len(manager._tasks_abandoned), 1)
with self.assertRaises(KeyError):
manager.UpdateTaskAsPendingMerge(task)
def testHasPendingTasks(self):
"""Tests the HasPendingTasks function."""
manager = task_manager.TaskManager()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.HasPendingTasks()
self.assertFalse(result)
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.HasPendingTasks()
self.assertTrue(result)
# Test with processing task.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.HasPendingTasks()
self.assertTrue(result)
# Test with abandoned task.
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result = manager.HasPendingTasks()
self.assertTrue(result)
# Test with task pending merge.
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.HasPendingTasks()
self.assertTrue(result)
# Test with merging task.
result_task = manager.GetTaskPendingMerge(None)
self.assertIsNotNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.HasPendingTasks()
self.assertTrue(result)
def testUpdateTaskAsProcessingByIdentifier(self):
"""Tests the UpdateTaskAsProcessingByIdentifier function."""
manager = task_manager.TaskManager()
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertIsNone(task.last_processing_time)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Indicate to the task manager that the task is processing.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Check if the last processing time was set.
self.assertIsNotNone(task.last_processing_time)
last_processing_time = task.last_processing_time
time.sleep(0.01)
# Indicate to the task manager that the task is still processing.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Check if the last processing time was updated.
self.assertGreater(task.last_processing_time, last_processing_time)
# Test with abandoned task.
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
# Indicate to the task manager that the task is still processing.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Indicate to the task manager that the task is pending merge.
task.storage_file_size = 10
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test status of task is unknown.
task = tasks.Task()
with self.assertRaises(KeyError):
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
def testCheckTaskToMerge(self):
"""Tests the CheckTaskToMerge function."""
manager = task_manager.TaskManager()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
result = manager.CheckTaskToMerge(task)
self.assertTrue(result)
# Test with processing task.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
result = manager.CheckTaskToMerge(task)
self.assertTrue(result)
# Test with abandoned task.
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result = manager.CheckTaskToMerge(task)
self.assertTrue(result)
# Test with abandoned task that is retried.
retry_task = manager.CreateRetryTask()
self.assertIsNotNone(retry_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result = manager.CheckTaskToMerge(task)
self.assertFalse(result)
# Test status of task is unknown.
del manager._tasks_abandoned[task.identifier]
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
with self.assertRaises(KeyError):
manager.CheckTaskToMerge(task)
def __init__(self,
maximum_number_of_tasks=None,
number_of_worker_processes=0,
worker_memory_limit=None,
worker_timeout=None):
"""Initializes an engine.
Args:
maximum_number_of_tasks (Optional[int]): maximum number of concurrent
tasks, where 0 represents no limit.
number_of_worker_processes (Optional[int]): number of worker processes.
worker_memory_limit (Optional[int]): maximum amount of memory a worker is
allowed to consume, where None represents the default memory limit
and 0 represents no limit.
worker_timeout (Optional[float]): number of minutes before a worker
process that is not providing status updates is considered inactive,
where None or 0.0 represents the default timeout.
"""
if maximum_number_of_tasks is None:
maximum_number_of_tasks = self._MAXIMUM_NUMBER_OF_TASKS
if number_of_worker_processes < 1:
# One worker for each "available" CPU (minus other processes).
# The number here is derived from the fact that the engine starts up:
# * A main process.
#
# If we want to utilize all CPUs on the system we therefore need to start
# up workers that amounts to the total number of CPUs - the other
# processes.
try:
cpu_count = multiprocessing.cpu_count() - 1
if cpu_count <= self._WORKER_PROCESSES_MINIMUM:
cpu_count = self._WORKER_PROCESSES_MINIMUM
elif cpu_count >= self._WORKER_PROCESSES_MAXIMUM:
cpu_count = self._WORKER_PROCESSES_MAXIMUM
except NotImplementedError:
logger.error((
'Unable to determine number of CPUs defaulting to {0:d} worker '
'processes.').format(self._WORKER_PROCESSES_MINIMUM))
cpu_count = self._WORKER_PROCESSES_MINIMUM
number_of_worker_processes = cpu_count
if worker_memory_limit is None:
worker_memory_limit = definitions.DEFAULT_WORKER_MEMORY_LIMIT
if not worker_timeout:
worker_timeout = definitions.DEFAULT_WORKER_TIMEOUT
super(ExtractionMultiProcessEngine, self).__init__()
self._enable_sigsegv_handler = False
self._last_worker_number = 0
self._maximum_number_of_containers = 50
self._maximum_number_of_tasks = maximum_number_of_tasks
self._merge_task = None
self._merge_task_on_hold = None
self._number_of_consumed_events = 0
self._number_of_consumed_event_tags = 0
self._number_of_consumed_extraction_warnings = 0
self._number_of_consumed_reports = 0
self._number_of_consumed_sources = 0
self._number_of_produced_events = 0
self._number_of_produced_event_tags = 0
self._number_of_produced_extraction_warnings = 0
self._number_of_produced_reports = 0
self._number_of_produced_sources = 0
self._number_of_worker_processes = number_of_worker_processes
self._parsers_counter = None
self._path_spec_extractor = extractors.PathSpecExtractor()
self._resolver_context = context.Context()
self._status = definitions.STATUS_INDICATOR_IDLE
self._storage_merge_reader = None
self._storage_merge_reader_on_hold = None
self._task_manager = task_manager.TaskManager()
self._task_queue = None
self._task_queue_port = None
self._task_storage_format = None
self._worker_memory_limit = worker_memory_limit
self._worker_timeout = worker_timeout
def testCompleteTask(self):
"""Tests the CompleteTask function."""
manager = task_manager.TaskManager()
# Test with queued task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
with self.assertRaises(KeyError):
manager.CompleteTask(task)
# Test with processing task.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
with self.assertRaises(KeyError):
manager.CompleteTask(task)
# Test with task pending merge.
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
with self.assertRaises(KeyError):
manager.CompleteTask(task)
# Test with merging task.
result_task = manager.GetTaskPendingMerge(None)
self.assertIsNotNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager.CompleteTask(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
# Test with abandoned task.
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
with self.assertRaises(KeyError):
manager.CompleteTask(task)
# Test with retry task.
retry_task = manager.CreateRetryTask()
self.assertIsNotNone(retry_task)
manager.UpdateTaskAsPendingMerge(retry_task)
result_task = manager.GetTaskPendingMerge(None)
self.assertIsNotNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 1)
self.assertEqual(len(manager._tasks_abandoned), 1)
manager.CompleteTask(retry_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
