Skip to content

Commit

Permalink
Ensures that QueueService instances are unique within types (#16751)
Browse files Browse the repository at this point in the history
  • Loading branch information
chrisguidry authored Jan 16, 2025
1 parent 0e91ca1 commit 5b5043c
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/prefect/_internal/concurrency/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def __init__(self, *args: Hashable) -> None:
self._task: Optional[asyncio.Task[None]] = None
self._stopped: bool = False
self._started: bool = False
self._key = hash(args)
self._key = hash((self.__class__, *args))
self._lock = threading.Lock()
self._queue_get_thread = WorkerThread(
# TODO: This thread should not need to be a daemon but when it is not, it
Expand Down Expand Up @@ -256,7 +256,7 @@ def instance(cls, *args: Hashable) -> Self:
If an instance already exists with the given arguments, it will be returned.
"""
with cls._instance_lock:
key = hash(args)
key = hash((cls, *args))
if key not in cls._instances:
cls._instances[key] = cls._new_instance(*args)

Expand Down
8 changes: 8 additions & 0 deletions tests/_internal/concurrency/test_services.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,14 @@ def test_instance_returns_new_instance_with_unique_key():
assert isinstance(new_instance, MockService)


def test_different_subclasses_have_unique_instances():
instance = MockService.instance()
assert isinstance(instance, MockService)
new_instance = MockBatchedService.instance()
assert new_instance is not instance
assert isinstance(new_instance, MockBatchedService)


def test_instance_returns_same_instance_after_error():
event = threading.Event()

Expand Down

0 comments on commit 5b5043c

Please sign in to comment.