def testGetTaskPendingRetry(self):
"""Tests the _GetTaskPendingRetry 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)
result_task = manager._GetTaskPendingRetry()
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._latest_task_processing_time -= (
2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result_task = manager._GetTaskPendingRetry()
self.assertEqual(result_task, task)
def testGetTasksCheckMerge(self):
"""Tests the GetTasksCheckMerge function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
result_tasks = manager.GetTasksCheckMerge()
self.assertEqual(result_tasks, [task])
def testUpdateTasksAsPendingMerge(self):
"""Tests the UpdateTasksAsPendingMerge function."""
manager = task_manager.TaskManager()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
task1 = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task1.storage_file_size = 10
task2 = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task2.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
manager.UpdateTasksAsPendingMerge([task1, task2])
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(len(manager._tasks_pending_merge), 2)
def testAbandonInactiveProcessingTasks(self):
"""Tests the _AbandonInactiveProcessingTasks function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
# Indicate to the task manager that the task is processing.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertIsNotNone(task.last_processing_time)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
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_abandoned), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
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 testTaskAbandonment(self):
"""Tests the abandoning and adoption of tasks"""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertEqual(manager.GetAbandonedTasks(), [])
self.assertTrue(manager.HasPendingTasks())
self.assertIsNone(manager.GetRetryTask())
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
timestamp = int(time.time() * 1000000)
inactive_time = timestamp - task_manager.TaskManager._TASK_INACTIVE_TIME
task.last_processing_time = inactive_time - 1
# HasPendingTasks is responsible for marking tasks as abandoned, and it
# should still return True if the task is abandoned, as it needs to be
# retried.
self.assertTrue(manager.HasPendingTasks())
abandoned_tasks = manager.GetAbandonedTasks()
self.assertIn(task, abandoned_tasks)
retry_task = manager.GetRetryTask()
self.assertIsNotNone(retry_task)
self.assertEqual(task.identifier, retry_task.original_task_identifier)
self.assertTrue(task.retried)
manager.CompleteTask(retry_task)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(manager.GetAbandonedTasks(), [])
self.assertTrue(manager.HasPendingTasks())
def testUpdateLatestProcessingTime(self):
"""Tests the _UpdateLatestProcessingTime function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.UpdateProcessingTime()
manager._UpdateLatestProcessingTime(task)
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 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 __init__(
self, maximum_number_of_tasks=_MAXIMUM_NUMBER_OF_TASKS):
"""Initializes an engine.
Args:
maximum_number_of_tasks (Optional[int]): maximum number of concurrent
tasks, where 0 represents no limit.
"""
super(TaskMultiProcessEngine, self).__init__()
self._enable_sigsegv_handler = False
self._last_worker_number = 0
self._maximum_number_of_tasks = maximum_number_of_tasks
self._merge_task = None
self._merge_task_on_hold = None
self._number_of_consumed_event_tags = 0
self._number_of_consumed_events = 0
self._number_of_consumed_reports = 0
self._number_of_consumed_sources = 0
self._number_of_consumed_warnings = 0
self._number_of_produced_event_tags = 0
self._number_of_produced_events = 0
self._number_of_produced_reports = 0
self._number_of_produced_sources = 0
self._number_of_produced_warnings = 0
self._number_of_worker_processes = 0
self._path_spec_extractor = extractors.PathSpecExtractor()
self._processing_configuration = None
self._resolver_context = context.Context()
self._session_identifier = None
self._status = definitions.STATUS_INDICATOR_IDLE
self._storage_merge_reader = None
self._storage_merge_reader_on_hold = None
self._task_queue = None
self._task_queue_port = None
self._task_manager = task_manager.TaskManager()
def testCreateTask(self):
"""Tests the CreateTask methods."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertIsInstance(task, tasks.Task)
self.assertEqual(task.session_identifier,
self._TEST_SESSION_IDENTIFIER)
self.assertIsNotNone(task.identifier)
def testCompleteTask(self):
"""Tests the CompleteTask function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertEqual(len(manager._tasks_queued), 1)
manager.CompleteTask(task)
self.assertEqual(len(manager._tasks_queued), 0)
def testHasPendingTasks(self):
"""Tests the HasPendingTasks function."""
manager = task_manager.TaskManager()
result = manager.HasPendingTasks()
self.assertFalse(result)
manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
result = manager.HasPendingTasks()
self.assertTrue(result)
def testTaskIsRetriable(self):
"""Tests the _TaskIsRetriable function."""
task = tasks.Task()
manager = task_manager.TaskManager()
result = manager._TaskIsRetriable(task)
self.assertTrue(result)
task.retried = True
result = manager._TaskIsRetriable(task)
self.assertFalse(result)
def testHasTasksPendingRetry(self):
"""Tests the _HasTasksPendingRetry function."""
manager = task_manager.TaskManager()
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()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
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 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 testHasTasksPendingMerge(self):
"""Tests the _HasTasksPendingMerge function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
result = manager._HasTasksPendingMerge()
self.assertFalse(result)
manager.UpdateTaskAsPendingMerge(task)
result = manager._HasTasksPendingMerge()
self.assertTrue(result)
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 testTaskProcessing(self):
"""Tests the UpdateTaskAsProcessing and GetProcessingTasks methods."""
manager = task_manager.TaskManager()
tasks_processing = manager.GetProcessingTasks()
self.assertEqual(0, len(tasks_processing))
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertIsNone(task.last_processing_time)
manager.UpdateTaskAsProcessing(task)
tasks_processing = manager.GetProcessingTasks()
self.assertEqual(1, len(tasks_processing))
task = tasks_processing[0]
self.assertIsNotNone(task.last_processing_time)
def testPendingMerge(self):
"""Tests the UpdateTaskPendingMerge and GetTaskPending merge methods."""
manager = task_manager.TaskManager()
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertIsNone(task_pending_merge)
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
with self.assertRaises(ValueError):
manager.UpdateTaskAsPendingMerge(task)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertIsNone(task_pending_merge)
task.storage_file_size = 10
manager.UpdateTaskAsPendingMerge(task)
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertEqual(task, task_pending_merge)
manager.CompleteTask(task)
small_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
small_task.storage_file_size = 100
manager.UpdateTaskAsProcessingByIdentifier(small_task.identifier)
large_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
large_task.storage_file_size = 1000
manager.UpdateTaskAsProcessingByIdentifier(large_task.identifier)
directory_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
directory_task.file_entry_type = dfvfs_definitions.FILE_ENTRY_TYPE_DIRECTORY
directory_task.storage_file_size = 1000
manager.UpdateTaskAsProcessingByIdentifier(directory_task.identifier)
manager.UpdateTaskAsPendingMerge(small_task)
# Simulate a delayed update from a worker that the task is still processing
manager.UpdateTaskAsProcessingByIdentifier(small_task.identifier)
manager.UpdateTaskAsPendingMerge(directory_task)
# Test that a directory task preempts a small task.
merging_task = manager.GetTaskPendingMerge(small_task)
self.assertEqual(merging_task, directory_task)
self.assertEqual(manager._tasks_merging.keys(),
[directory_task.identifier])
manager.CompleteTask(directory_task)
# Test that a small task is not preempted by a large task.
manager.UpdateTaskAsPendingMerge(large_task)
merging_task = manager.GetTaskPendingMerge(small_task)
self.assertEqual(merging_task, small_task)
self.assertEqual(manager._tasks_merging.keys(),
[small_task.identifier])
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 testAbandonQueuedTasks(self):
"""Tests the _AbandonQueuedTasks function."""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertIsNotNone(task.start_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)
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)
def testTaskProcessing(self):
"""Tests the UpdateTaskAsProcessing and GetTasksCheckMerge methods."""
manager = task_manager.TaskManager()
tasks_processing = manager.GetTasksCheckMerge()
self.assertEqual(0, len(tasks_processing))
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertIsNone(task.last_processing_time)
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
tasks_processing = manager.GetTasksCheckMerge()
self.assertEqual(1, len(tasks_processing))
task = tasks_processing[0]
self.assertIsNotNone(task.last_processing_time)
# Simulate the foreman marking a task as complete while the manager
# believed it to still be processing.
manager.CompleteTask(task)
def testGetFailedTasks(self):
"""Tests the GetFailedTasks 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)
result_tasks = manager.GetFailedTasks()
self.assertEqual(result_tasks, [])
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result_tasks = manager.GetFailedTasks()
self.assertEqual(result_tasks, [task])
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 __init__(self,
maximum_number_of_tasks=_MAXIMUM_NUMBER_OF_TASKS,
use_zeromq=True):
"""Initializes an engine object.
Args:
maximum_number_of_tasks (Optional[int]): maximum number of concurrent
tasks, where 0 represents no limit.
use_zeromq (Optional[bool]): True if ZeroMQ should be used for queuing
instead of Python's multiprocessing queue.
"""
super(TaskMultiProcessEngine, self).__init__()
self._enable_sigsegv_handler = False
self._filter_find_specs = None
self._last_worker_number = 0
self._maximum_number_of_tasks = maximum_number_of_tasks
self._memory_profiler = None
self._merge_task = None
self._merge_task_on_hold = None
self._number_of_consumed_errors = 0
self._number_of_consumed_event_tags = 0
self._number_of_consumed_events = 0
self._number_of_consumed_reports = 0
self._number_of_consumed_sources = 0
self._number_of_produced_errors = 0
self._number_of_produced_event_tags = 0
self._number_of_produced_events = 0
self._number_of_produced_reports = 0
self._number_of_produced_sources = 0
self._number_of_worker_processes = 0
self._path_spec_extractor = extractors.PathSpecExtractor()
self._processing_configuration = None
self._processing_profiler = None
self._resolver_context = context.Context()
self._serializers_profiler = None
self._session_identifier = None
self._status = definitions.PROCESSING_STATUS_IDLE
self._storage_merge_reader = None
self._storage_merge_reader_on_hold = None
self._task_queue = None
self._task_queue_port = None
self._task_manager = task_manager.TaskManager()
self._use_zeromq = use_zeromq
def testTaskAbandonment(self):
"""Tests the abandoning and adoption of tasks"""
manager = task_manager.TaskManager()
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
self.assertEqual(manager.GetAbandonedTasks(), [])
self.assertTrue(manager.HasActiveTasks())
manager.UpdateTaskAsProcessing(task)
timestamp = int(time.time() * 1000000)
# pylint: disable=protected-access
inactive_time = timestamp - task_manager.TaskManager._TASK_INACTIVE_TIME
task.last_processing_time = inactive_time - 1
# HasActiveTasks is responsible for marking tasks as abandoned.
self.assertFalse(manager.HasActiveTasks())
abandoned_tasks = manager.GetAbandonedTasks()
self.assertIn(task, abandoned_tasks)
manager.AdoptTask(task)
self.assertEqual(manager.GetAbandonedTasks(), [])
self.assertTrue(manager.HasActiveTasks())
def testPendingMerge(self):
"""Tests the UpdateTaskPendingMerge and GetTaskPending merge methods."""
manager = task_manager.TaskManager()
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertIsNone(task_pending_merge)
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
with self.assertRaises(KeyError):
manager.UpdateTaskAsPendingMerge(task)
manager.UpdateTaskAsProcessing(task)
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertIsNone(task_pending_merge)
task.storage_file_size = 10
manager.UpdateTaskAsPendingMerge(task)
task_pending_merge = manager.GetTaskPendingMerge(None)
self.assertEqual(task, task_pending_merge)
manager.CompleteTask(task)
small_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
small_task.storage_file_size = 100
manager.UpdateTaskAsProcessing(small_task)
large_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
large_task.storage_file_size = 1000
manager.UpdateTaskAsProcessing(large_task)
directory_task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
directory_task.file_entry_type = dfvfs_definitions.FILE_ENTRY_TYPE_DIRECTORY
directory_task.storage_file_size = 1000
manager.UpdateTaskAsProcessing(directory_task)
manager.UpdateTaskAsPendingMerge(small_task)
manager.UpdateTaskAsPendingMerge(directory_task)
# Test that a directory task preempts a small task.
merging_task = manager.GetTaskPendingMerge(small_task)
self.assertEqual(merging_task, directory_task)
# Test that a small task is not preempted by a large task.
manager.UpdateTaskAsPendingMerge(large_task)
merging_task = manager.GetTaskPendingMerge(small_task)
self.assertEqual(merging_task, small_task)
def testGetTaskPendingRetry(self):
"""Tests the _GetTaskPendingRetry 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)
result_task = manager._GetTaskPendingRetry()
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_abandoned), 0)
manager._AbandonQueuedTasks()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
result_task = manager._GetTaskPendingRetry()
self.assertEqual(result_task, task)
