Parallel: Support queue in TaskManager

This change implements queue for TaskManager and
use it in test results queuing.

- Replace thread-based test message processing
  with TaskManager in testsuite.py
- Enhance TaskManager with pending task queue and
  improved worker management
- Add proper task queuing mechanism to handle test
  result messages in order

Author: squirrelsc

lisa/testsuite.py

@@ -5,12 +5,10 @@
 
 import copy
 import logging
-import threading
 import traceback
 from dataclasses import dataclass, field
-from functools import wraps
+from functools import partial, wraps
 from pathlib import Path
-from queue import Queue
 from time import sleep
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
@@ -44,6 +42,7 @@
     get_logger,
     remove_handler,
 )
+from lisa.util.parallel import Task, TaskManager
 from lisa.util.perf_timer import Timer, create_timer
 
 _all_suites: Dict[str, TestSuiteMetadata] = {}
@@ -372,7 +371,7 @@ def _send_result_message(self, stacktrace: Optional[str] = None) -> None:
         result_message.suite_full_name = self.runtime_data.metadata.suite.full_name
         result_message.stacktrace = stacktrace
 
-        _queue_test_message(result_message=result_message)
+        _queue_test_message(result=self, result_message=result_message)
 
     @retry(exceptions=FileExistsError, tries=30, delay=0.1)  # type: ignore
     def __create_case_log_path(self) -> Path:
@@ -895,24 +894,18 @@ def get_cases_metadata() -> Dict[str, TestCaseMetadata]:
 
 
 def start_test_result_message_processing() -> None:
-    global __message_thread
-    if not __message_thread:
-        # clear any left messages for UT
-        global __message_queue
-        __message_queue = Queue()
-
-        __message_thread = threading.Thread(target=_message_worker)
-        __message_thread.start()
+    global __test_message_manager
+    if __test_message_manager is None:
+        __test_message_manager = TaskManager(max_workers=1)
     else:
-        raise LisaException("message thread is already created and started")
+        raise LisaException("test message tasks is already initialized")
 
 
 def wait_for_test_result_messages() -> None:
-    __message_queue.put(None)
-    global __message_thread
-    assert __message_thread is not None
-    __message_thread.join()
-    __message_thread = None
+    global __test_message_manager
+    if __test_message_manager is not None:
+        __test_message_manager.wait_for_all_workers()
+        __test_message_manager = None
 
 
 def _add_suite_metadata(metadata: TestSuiteMetadata) -> None:
@@ -972,27 +965,24 @@ def _add_case_to_suite(
     test_suite.cases.append(test_case)
 
 
-def _message_worker() -> None:
-    while True:
-        message = __message_queue.get()
-        if message is None:
-            break
-        _send_result_message(message)
-
-
-# process test results message in an order and in background
-__message_thread: Optional[threading.Thread] = None
-__message_queue: Queue[Optional[TestResultMessage]] = Queue()
+# process test results message in an order, so the max_workers is 1
+__test_message_manager: Optional[TaskManager[None]] = None
 
 
 def _send_result_message(result_message: TestResultMessage) -> None:
     plugin_manager.hook.update_test_result_message(message=result_message)
     notifier.notify(message=result_message)
 
 
-def _queue_test_message(result_message: TestResultMessage) -> None:
-    assert result_message is not None
-    __message_queue.put(result_message)
+def _queue_test_message(result: TestResult, result_message: TestResultMessage) -> None:
+    assert __test_message_manager, "test message manager is not initialized"
+    __test_message_manager.submit_task(
+        Task(
+            task_id=0,
+            task=partial(_send_result_message, result_message),
+            parent_logger=result.log,
+        )
+    )
 
 
 plugin_manager.add_hookspecs(TestResult)

lisa/util/parallel.py

@@ -1,12 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
-from concurrent.futures import (
-    ALL_COMPLETED,
-    FIRST_COMPLETED,
-    Future,
-    ThreadPoolExecutor,
-    wait,
-)
+import time
+from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
+from queue import Queue
 from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar
 
 from assertpy import assert_that
@@ -85,6 +81,7 @@ def __init__(
         self._cancelled = False
         self._future_task_map: Dict[Future[T_RESULT], Task[T_RESULT]] = {}
         self._is_verbose = is_verbose
+        self._pending_tasks: Queue[Task[T_RESULT]] = Queue()
 
     def __enter__(self) -> Any:
         return self._pool.__enter__()
@@ -97,9 +94,8 @@ def running_count(self) -> int:
         return len(self._futures)
 
     def submit_task(self, task: Task[T_RESULT]) -> None:
-        future: Future[T_RESULT] = self._pool.submit(task)
-        self._future_task_map[future] = task
-        self._futures.append(future)
+        self._pending_tasks.put(task)
+        self._process_pending_tasks()
 
     def cancel(self) -> None:
         self._log.info("Called to cancel all tasks.")
@@ -137,8 +133,22 @@ def _process_done_futures(self) -> None:
                 self._future_task_map.pop(future)
 
     def wait_for_all_workers(self) -> None:
-        remaining_worker_count = self.wait_worker(return_condition=ALL_COMPLETED)
-        assert_that(remaining_worker_count).is_zero()
+        while True:
+            self._process_pending_tasks()
+            has_remaining = not self._pending_tasks.empty() or self.wait_worker()
+            if not has_remaining:
+                break
+            time.sleep(0)
+
+        assert_that(has_remaining).is_false()
+
+    def _process_pending_tasks(self) -> None:
+        while not self._pending_tasks.empty() and self.has_idle_worker():
+            self.check_cancelled()
+            task = self._pending_tasks.get()
+            future: Future[T_RESULT] = self._pool.submit(task)
+            self._future_task_map[future] = task
+            self._futures.append(future)
 
 
 _default_task_manager: Optional[TaskManager[Any]] = None

selftests/runners/test_lisa_runner.py

@@ -62,9 +62,11 @@ class RunnerTestCase(TestCase):
 
     def setUp(self) -> None:
         lisa.environment._global_environment_id = 0
+        start_test_result_message_processing()
 
     def tearDown(self) -> None:
         test_testsuite.cleanup_cases_metadata()  # Necessary side effects!
+        wait_for_test_result_messages()
 
     def test_merge_req_create_on_new(self) -> None:
         # if no predefined envs, can generate from requirement
@@ -744,8 +746,6 @@ def _run_all_tests(self, runner: LisaRunner) -> List[TestResultMessage]:
         results_collector = RunnerResult(schema.Notifier())
         register_notifier(results_collector)
 
-        start_test_result_message_processing()
-
         try:
             runner.initialize()
 

selftests/test_testsuite.py

@@ -23,6 +23,8 @@
     get_suites_metadata,
     node_requirement,
     simple_requirement,
+    start_test_result_message_processing,
+    wait_for_test_result_messages,
 )
 from lisa.util.logger import Logger
 from selftests.test_environment import generate_runbook
@@ -171,6 +173,10 @@ def generate_suite_instance(self) -> MockTestSuite:
 
     def setUp(self) -> None:
         cleanup_cases_metadata()
+        start_test_result_message_processing()
+
+    def tearDown(self) -> None:
+        wait_for_test_result_messages()
 
     def test_expanded_nodespace(self) -> None:
         cases = generate_cases_metadata()