diff --git a/core/queue.py b/core/queue.py index cc06468..c6925f4 100644 --- a/core/queue.py +++ b/core/queue.py @@ -1,4 +1,5 @@ import logging +import threading from pydantic import BaseModel from typing import Any, Optional from gitea.models import IssueModel, PullRequestModel @@ -19,43 +20,53 @@ class WorkQueue: """Thread-safe work queue grouped by repo.""" def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() self._queue: list[WorkItem] = [] self._enqueued_repos: set[str] = set() def enqueue(self, item: WorkItem) -> None: - self._queue.append(item) - self._enqueued_repos.add(item.repo_full_name) - logger.info(f"Enqueued work item: {item.task_type} #{item.task_number} for {item.repo_full_name}") + with self._lock: + self._queue.append(item) + self._enqueued_repos.add(item.repo_full_name) + logger.info(f"Enqueued work item: {item.task_type} #{item.task_number} for {item.repo_full_name}") def enqueue_batch(self, items: list[WorkItem]) -> None: - for item in items: - self.enqueue(item) + with self._lock: + for item in items: + self._queue.append(item) + self._enqueued_repos.add(item.repo_full_name) + logger.info(f"Enqueued work item: {item.task_type} #{item.task_number} for {item.repo_full_name}") def get_repo_work(self, repo: str) -> list[WorkItem]: """Get all work items for a specific repo.""" - items: list[WorkItem] = [ - item for item in self._queue if item.repo_full_name == repo - ] - logger.info(f"Retrieved {len(items)} work items for repository: {repo}") - return items + with self._lock: + items: list[WorkItem] = [ + item for item in self._queue if item.repo_full_name == repo + ] + logger.info(f"Retrieved {len(items)} work items for repository: {repo}") + return items def remove_repo_work(self, repo: str) -> None: """Remove all work items for a specific repo.""" - self._queue = [ - item for item in self._queue if item.repo_full_name != repo - ] - self._enqueued_repos.discard(repo) - logger.info(f"Removed all work items for repository: {repo}") + with self._lock: + self._queue = [ + item for item in self._queue if item.repo_full_name != repo + ] + self._enqueued_repos.discard(repo) + logger.info(f"Removed all work items for repository: {repo}") def get_next_repo(self) -> str | None: """Get the next repo with work, or None if empty.""" - if not self._enqueued_repos: - return None - return next(iter(self._enqueued_repos)) + with self._lock: + if not self._enqueued_repos: + return None + return next(iter(self._enqueued_repos)) @property def is_empty(self) -> bool: - return len(self._queue) == 0 + with self._lock: + return len(self._queue) == 0 def __len__(self) -> int: - return len(self._queue) + with self._lock: + return len(self._queue) diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 0000000..97e8b48 --- /dev/null +++ b/tests/test_queue.py @@ -0,0 +1,102 @@ +import threading +from core.queue import WorkQueue, WorkItem +from gitea.models import IssueModel + + +def test_work_queue_basic_operations() -> None: + queue: WorkQueue = WorkQueue() + assert queue.is_empty + assert len(queue) == 0 + assert queue.get_next_repo() is None + + item1 = WorkItem( + repo_full_name="meeks/repo1", + task_type="issue", + task_number=1, + task_info=IssueModel(number=1), + ) + item2 = WorkItem( + repo_full_name="meeks/repo1", + task_type="issue", + task_number=2, + task_info=IssueModel(number=2), + ) + item3 = WorkItem( + repo_full_name="meeks/repo2", + task_type="issue", + task_number=3, + task_info=IssueModel(number=3), + ) + + queue.enqueue(item1) + assert not queue.is_empty + assert len(queue) == 1 + assert queue.get_next_repo() == "meeks/repo1" + + queue.enqueue_batch([item2, item3]) + assert len(queue) == 3 + + # get_repo_work + repo1_work = queue.get_repo_work("meeks/repo1") + assert len(repo1_work) == 2 + assert repo1_work[0].task_number == 1 + assert repo1_work[1].task_number == 2 + + # remove_repo_work + queue.remove_repo_work("meeks/repo1") + assert len(queue) == 1 + assert queue.get_next_repo() == "meeks/repo2" + + queue.remove_repo_work("meeks/repo2") + assert queue.is_empty + assert len(queue) == 0 + assert queue.get_next_repo() is None + + +def test_work_queue_thread_safety() -> None: + queue: WorkQueue = WorkQueue() + num_threads: int = 10 + items_per_thread: int = 100 + barrier = threading.Barrier(num_threads) + + def worker(thread_idx: int) -> None: + barrier.wait() # synchronize start + for i in range(items_per_thread): + item = WorkItem( + repo_full_name=f"meeks/repo_{thread_idx}", + task_type="issue", + task_number=i, + task_info=IssueModel(number=i), + ) + queue.enqueue(item) + + threads: list[threading.Thread] = [] + for idx in range(num_threads): + t = threading.Thread(target=worker, args=(idx,)) + threads.append(t) + t.start() + + for t in threads: + t.join() + + # Verify that all items are enqueued + assert len(queue) == num_threads * items_per_thread + + # Concurrently remove repo work + barrier_remove = threading.Barrier(num_threads) + + def remover(thread_idx: int) -> None: + barrier_remove.wait() + queue.remove_repo_work(f"meeks/repo_{thread_idx}") + + remove_threads: list[threading.Thread] = [] + for idx in range(num_threads): + t = threading.Thread(target=remover, args=(idx,)) + remove_threads.append(t) + t.start() + + for t in remove_threads: + t.join() + + assert queue.is_empty + assert len(queue) == 0