Improve SuspendObjectStream concurrency and expand test coverage#11
Merged
Conversation
…n version support to 3.15
…or, multi-waiter resume, and wait_to_suspend cycles
Member
Author
|
@sourcery-ai title |
Contributor
审阅者指南重构 SuspendObjectStream 的并发和队列处理逻辑，使其基于 state-lock 驱动并支持多等待者安全；围绕挂起、队列、回调和生成器生命周期添加全面测试，并更新项目元数据以支持 Python 3.15 和新的补丁版本发布。 SuspendObjectStream 中多等待者挂起与恢复的时序图sequenceDiagram
actor SuspendCaller
participant SuspendObjectStream
actor Worker1
actor Worker2
SuspendCaller->>SuspendObjectStream: wait_to_suspend(tags, timeout)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
SuspendObjectStream->>SuspendObjectStream: set _suspend_tags, __suspend_signal
SuspendObjectStream-->>SuspendCaller: await __suspend_signal
Worker1->>SuspendObjectStream: _wait_for_continue(tag)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
alt first waiter
SuspendObjectStream->>SuspendObjectStream: set_result(__suspend_signal)
SuspendObjectStream->>SuspendObjectStream: create __resume_signal
SuspendObjectStream-->>Worker1: await shared_fut
end
Worker2->>SuspendObjectStream: _wait_for_continue(tag)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
alt additional waiter
SuspendObjectStream->>SuspendObjectStream: add_done_callback(shared_fut)
SuspendObjectStream-->>Worker2: await fut
end
SuspendCaller->>SuspendObjectStream: resume()
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
SuspendObjectStream->>SuspendObjectStream: set_result(__resume_signal)
Worker1-->>SuspendObjectStream: fut completed
Worker2-->>SuspendObjectStream: fut completed
SuspendObjectStream->>SuspendObjectStream: clear __resume_signal under _state_lock
文件级变更
提示与命令与 Sourcery 交互
自定义你的使用体验访问你的控制面板 以：
获取帮助Original review guide in EnglishReviewer's GuideRefactors SuspendObjectStream’s concurrency and queue-handling logic to be state-lock–driven and multi-waiter safe, adds comprehensive tests around suspension, queue, callbacks, and generator lifecycle, and updates project metadata for Python 3.15 support and a new patch release. Sequence diagram for multi-waiter suspend and resume in SuspendObjectStreamsequenceDiagram
actor SuspendCaller
participant SuspendObjectStream
actor Worker1
actor Worker2
SuspendCaller->>SuspendObjectStream: wait_to_suspend(tags, timeout)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
SuspendObjectStream->>SuspendObjectStream: set _suspend_tags, __suspend_signal
SuspendObjectStream-->>SuspendCaller: await __suspend_signal
Worker1->>SuspendObjectStream: _wait_for_continue(tag)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
alt first waiter
SuspendObjectStream->>SuspendObjectStream: set_result(__suspend_signal)
SuspendObjectStream->>SuspendObjectStream: create __resume_signal
SuspendObjectStream-->>Worker1: await shared_fut
end
Worker2->>SuspendObjectStream: _wait_for_continue(tag)
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
alt additional waiter
SuspendObjectStream->>SuspendObjectStream: add_done_callback(shared_fut)
SuspendObjectStream-->>Worker2: await fut
end
SuspendCaller->>SuspendObjectStream: resume()
SuspendObjectStream->>SuspendObjectStream: acquire _state_lock
SuspendObjectStream->>SuspendObjectStream: set_result(__resume_signal)
Worker1-->>SuspendObjectStream: fut completed
Worker2-->>SuspendObjectStream: fut completed
SuspendObjectStream->>SuspendObjectStream: clear __resume_signal under _state_lock
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
sourcery-ai
Bot
changed the title
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - 我发现了 3 个问题，并给出了一些高层次的反馈：
- 新的
_state_lock是一个异步的aiologic.Lock，但在resume、set_callback_func、set_callback_fun_sending和get_response_generator这些同步方法里用的是普通的with，这会在运行时出错；要么把这些方法改成 async 并配合异步锁使用，要么在这些方法里避免使用异步锁。 - 在
set_queue_done中，_queue_done会在_put_to_queue运行之前就被设置为True；如果_put_to_queue因为非BrokenResourceError的原因失败，队列会被标记为完成，但结束标记并没有被放入队列——建议只在发送成功后再将_queue_done置为True，或者显式处理非BrokenResourceError的异常。 - 新增的
_wait_for_continue/wait_to_suspend同步逻辑更复杂，使用了共享的 future 加_state_lock；建议再仔细检查一下潜在的竞态条件（例如：wait_to_suspend在某个等待者获取锁的同时清空__suspend_signal），并在这些边界情况附近加上简短注释，说明预期的执行顺序。
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new `_state_lock` is an async `aiologic.Lock`, but it's being used with plain `with` in sync methods like `resume`, `set_callback_func`, `set_callback_fun_sending`, and `get_response_generator`, which will fail at runtime; either make these methods async or avoid using the async lock in them.
- In `set_queue_done`, `_queue_done` is set to `True` before `_put_to_queue` runs; if `_put_to_queue` fails for reasons other than `BrokenResourceError`, the queue will be marked done without the done marker being enqueued—consider only flipping `_queue_done` after a successful send or handling non-`BrokenResourceError` exceptions explicitly.
- The new `_wait_for_continue` / `wait_to_suspend` synchronization logic is more complex and uses shared futures plus `_state_lock`; it would be helpful to double-check for potential races (e.g., `wait_to_suspend` clearing `__suspend_signal` while a waiter is acquiring the lock) and possibly add brief comments around those edge cases to explain the intended ordering.
## Individual Comments
### Comment 1
<location path="src/amrita_sense/streaming.py" line_range="121-130" />
<code_context>
+ shared_fut = self.__resume_signal
+ is_first = True
+
+ await asyncio.sleep(0)
+
+ if is_first:
+ fut = shared_fut
+ else:
</code_context>
<issue_to_address>
**issue (bug_risk):** Cancellation before `fut` initialization leads to `UnboundLocalError` in `finally`.
If the task is cancelled during `await asyncio.sleep(0)`, control goes to `finally` before `fut` is assigned in either branch, so referencing `fut` there raises `UnboundLocalError` instead of doing the intended cleanup. You can prevent this by initializing `fut = None` before the `async with` and guarding the cleanup (e.g. only resetting `__resume_signal` if `is_first` and `fut is not None`), or by moving `await asyncio.sleep(0)` to a point after `fut` is always assigned.
</issue_to_address>
### Comment 2
<location path="src/amrita_sense/streaming.py" line_range="183-159" />
<code_context>
"""Resume to run when suspend."""
- if self.__resume_signal and not self.__resume_signal.done():
- self.__resume_signal.set_result(True)
+ with self._state_lock:
+ if self.__resume_signal and not self.__resume_signal.done():
+ self.__resume_signal.set_result(True)
</code_context>
<issue_to_address>
**issue (bug_risk):** `resume` uses an async lock with a synchronous `with`, which will fail at runtime.
`_state_lock` is an `aiologic.Lock`, which must be used with `async with`, not a synchronous `with` block. In its current form, `resume()` will raise a `TypeError` on first use. Update this to either make `resume` async and use `async with self._state_lock`, or change `_state_lock` to a synchronous lock type and use it consistently across the class.
</issue_to_address>
### Comment 3
<location path="tests/test_object_stream.py" line_range="26" />
<code_context>
+ yield item
@pytest.mark.asyncio
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a test for `wait_to_suspend` timeout behavior to ensure state is correctly reset after a timeout.
The new `wait_to_suspend` uses `_state_lock` and resets `__suspend_signal` / `_suspend_tags` in a `finally` block, and we already cover the "already waiting" and full suspend/resume cases. What’s missing is a test where `wait_to_suspend(..., timeout=...)` actually times out. Please add a test that:
1. Calls `await obj.wait_to_suspend("tag", timeout=small_value)` and asserts `asyncio.TimeoutError`.
2. Verifies a subsequent `wait_to_suspend("tag", timeout=...)` still works, confirming the timeout didn’t leave `__suspend_signal` / `_suspend_tags` in a bad state.
This would help catch any regressions where a timeout leaves the object stuck in an "already waiting" or inconsistent state.
Suggested implementation:
```python
# A. suspend decorator (4 tests)
# ============================================================================
=======
async def _async_gen(items: list):
"""Create an async generator from a list."""
for item in items:
yield item
@pytest.mark.asyncio
assert suspend, "Suspend not called"
finally:
hd.cancel()
# ============================================================================
# A. suspend decorator (4 tests)
# ============================================================================
@pytest.mark.asyncio
async def test_wait_to_suspend_timeout_resets_state(suspendable_obj):
"""
wait_to_suspend(...) should reset internal state when it times out so that a
subsequent wait_to_suspend on the same tag still works.
"""
# 1) First call: let it time out and ensure we get asyncio.TimeoutError
with pytest.raises(asyncio.TimeoutError):
await suspendable_obj.wait_to_suspend("timeout-tag", timeout=0.01)
# 2) Second call: verify state was reset and we can wait again successfully
async def trigger_suspend():
# Give wait_to_suspend a chance to start before triggering suspend
await asyncio.sleep(0.01)
await suspendable_obj.suspend("timeout-tag")
task = asyncio.create_task(trigger_suspend())
try:
suspend = await suspendable_obj.wait_to_suspend("timeout-tag", timeout=0.5)
assert suspend, "Suspend not called after timeout reset"
finally:
task.cancel()
```
To wire this test into your existing codebase you will likely need to:
1. Replace the `suspendable_obj` fixture parameter with whatever fixture/object you already use to exercise `wait_to_suspend` in your existing tests (e.g. `obj`, `handler`, `stream`, etc.).
2. Adjust the `suspendable_obj.suspend("timeout-tag")` call if your API uses a different method name or signature to trigger suspension.
3. Ensure `asyncio` and `pytest` are imported at the top of `tests/test_object_stream.py` (they almost certainly already are given other async tests in this file).
4. Optionally tune the timeout values (`0.01` / `0.5`) if your CI environments are slow; e.g. use `0.05` / `1.0` while still ensuring the first call reliably times out and the second reliably completes.
</issue_to_address>Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Original comment in English
Hey - I've found 3 issues, and left some high level feedback:
- The new
_state_lockis an asyncaiologic.Lock, but it's being used with plainwithin sync methods likeresume,set_callback_func,set_callback_fun_sending, andget_response_generator, which will fail at runtime; either make these methods async or avoid using the async lock in them. - In
set_queue_done,_queue_doneis set toTruebefore_put_to_queueruns; if_put_to_queuefails for reasons other thanBrokenResourceError, the queue will be marked done without the done marker being enqueued—consider only flipping_queue_doneafter a successful send or handling non-BrokenResourceErrorexceptions explicitly. - The new
_wait_for_continue/wait_to_suspendsynchronization logic is more complex and uses shared futures plus_state_lock; it would be helpful to double-check for potential races (e.g.,wait_to_suspendclearing__suspend_signalwhile a waiter is acquiring the lock) and possibly add brief comments around those edge cases to explain the intended ordering.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new `_state_lock` is an async `aiologic.Lock`, but it's being used with plain `with` in sync methods like `resume`, `set_callback_func`, `set_callback_fun_sending`, and `get_response_generator`, which will fail at runtime; either make these methods async or avoid using the async lock in them.
- In `set_queue_done`, `_queue_done` is set to `True` before `_put_to_queue` runs; if `_put_to_queue` fails for reasons other than `BrokenResourceError`, the queue will be marked done without the done marker being enqueued—consider only flipping `_queue_done` after a successful send or handling non-`BrokenResourceError` exceptions explicitly.
- The new `_wait_for_continue` / `wait_to_suspend` synchronization logic is more complex and uses shared futures plus `_state_lock`; it would be helpful to double-check for potential races (e.g., `wait_to_suspend` clearing `__suspend_signal` while a waiter is acquiring the lock) and possibly add brief comments around those edge cases to explain the intended ordering.
## Individual Comments
### Comment 1
<location path="src/amrita_sense/streaming.py" line_range="121-130" />
<code_context>
+ shared_fut = self.__resume_signal
+ is_first = True
+
+ await asyncio.sleep(0)
+
+ if is_first:
+ fut = shared_fut
+ else:
</code_context>
<issue_to_address>
**issue (bug_risk):** Cancellation before `fut` initialization leads to `UnboundLocalError` in `finally`.
If the task is cancelled during `await asyncio.sleep(0)`, control goes to `finally` before `fut` is assigned in either branch, so referencing `fut` there raises `UnboundLocalError` instead of doing the intended cleanup. You can prevent this by initializing `fut = None` before the `async with` and guarding the cleanup (e.g. only resetting `__resume_signal` if `is_first` and `fut is not None`), or by moving `await asyncio.sleep(0)` to a point after `fut` is always assigned.
</issue_to_address>
### Comment 2
<location path="src/amrita_sense/streaming.py" line_range="183-159" />
<code_context>
"""Resume to run when suspend."""
- if self.__resume_signal and not self.__resume_signal.done():
- self.__resume_signal.set_result(True)
+ with self._state_lock:
+ if self.__resume_signal and not self.__resume_signal.done():
+ self.__resume_signal.set_result(True)
</code_context>
<issue_to_address>
**issue (bug_risk):** `resume` uses an async lock with a synchronous `with`, which will fail at runtime.
`_state_lock` is an `aiologic.Lock`, which must be used with `async with`, not a synchronous `with` block. In its current form, `resume()` will raise a `TypeError` on first use. Update this to either make `resume` async and use `async with self._state_lock`, or change `_state_lock` to a synchronous lock type and use it consistently across the class.
</issue_to_address>
### Comment 3
<location path="tests/test_object_stream.py" line_range="26" />
<code_context>
+ yield item
@pytest.mark.asyncio
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a test for `wait_to_suspend` timeout behavior to ensure state is correctly reset after a timeout.
The new `wait_to_suspend` uses `_state_lock` and resets `__suspend_signal` / `_suspend_tags` in a `finally` block, and we already cover the "already waiting" and full suspend/resume cases. What’s missing is a test where `wait_to_suspend(..., timeout=...)` actually times out. Please add a test that:
1. Calls `await obj.wait_to_suspend("tag", timeout=small_value)` and asserts `asyncio.TimeoutError`.
2. Verifies a subsequent `wait_to_suspend("tag", timeout=...)` still works, confirming the timeout didn’t leave `__suspend_signal` / `_suspend_tags` in a bad state.
This would help catch any regressions where a timeout leaves the object stuck in an "already waiting" or inconsistent state.
Suggested implementation:
```python
# A. suspend decorator (4 tests)
# ============================================================================
=======
async def _async_gen(items: list):
"""Create an async generator from a list."""
for item in items:
yield item
@pytest.mark.asyncio
assert suspend, "Suspend not called"
finally:
hd.cancel()
# ============================================================================
# A. suspend decorator (4 tests)
# ============================================================================
@pytest.mark.asyncio
async def test_wait_to_suspend_timeout_resets_state(suspendable_obj):
"""
wait_to_suspend(...) should reset internal state when it times out so that a
subsequent wait_to_suspend on the same tag still works.
"""
# 1) First call: let it time out and ensure we get asyncio.TimeoutError
with pytest.raises(asyncio.TimeoutError):
await suspendable_obj.wait_to_suspend("timeout-tag", timeout=0.01)
# 2) Second call: verify state was reset and we can wait again successfully
async def trigger_suspend():
# Give wait_to_suspend a chance to start before triggering suspend
await asyncio.sleep(0.01)
await suspendable_obj.suspend("timeout-tag")
task = asyncio.create_task(trigger_suspend())
try:
suspend = await suspendable_obj.wait_to_suspend("timeout-tag", timeout=0.5)
assert suspend, "Suspend not called after timeout reset"
finally:
task.cancel()
```
To wire this test into your existing codebase you will likely need to:
1. Replace the `suspendable_obj` fixture parameter with whatever fixture/object you already use to exercise `wait_to_suspend` in your existing tests (e.g. `obj`, `handler`, `stream`, etc.).
2. Adjust the `suspendable_obj.suspend("timeout-tag")` call if your API uses a different method name or signature to trigger suspension.
3. Ensure `asyncio` and `pytest` are imported at the top of `tests/test_object_stream.py` (they almost certainly already are given other async tests in this file).
4. Optionally tune the timeout values (`0.01` / `0.5`) if your CI environments are slow; e.g. use `0.05` / `1.0` while still ensuring the first call reliably times out and the second reliably completes.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Member
Author
|
@sourcery-ai title |
|
Using #!/usr/bin/env python3
import asyncio
import aiologic
lock = aiologic.Lock()
async def work():
async with lock:
await asyncio.sleep(0)
async def main():
async with lock:
task = asyncio.create_task(work())
await asyncio.sleep(0) # switch to the task
with lock: # deadlock!
print("never")
if __name__ == "__main__":
asyncio.run(main())For example, |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by Sourcery
改进 SuspendObjectStream 的并发语义并扩展测试覆盖率。
Bug 修复：
功能增强：
构建：
CI：
测试：
Original summary in English
Summary by Sourcery
Improve SuspendObjectStream concurrency semantics and expand coverage.
Bug Fixes:
Enhancements:
Build:
CI:
Tests: