Skip to content

Commit

Permalink
Make trio.lowlevel.checkpoint much faster
Browse files Browse the repository at this point in the history
This is a bit of a kluge, and will hopefully be cleaned up in the
future when we overhaul KeyboardInterrupt (python-triogh-733) and/or the
cancellation checking APIs (python-triogh-961). But 'checkpoint()' is a common
operation, and this speeds it up by ~7x, so... not bad for a ~4 line
change.

Before this change:

- `await checkpoint()`: ~24k/second
- `await cancel_shielded_checkpoint()`: ~180k/second

After this change:

- `await checkpoint()`: ~170k/second
- `await cancel_shielded_checkpoint()`: ~180k/second

Benchmark script:

```python
import time
import trio

LOOP_SIZE = 1_000_000

async def main():
    start = time.monotonic()
    for _ in range(LOOP_SIZE):
        await trio.lowlevel.checkpoint()
        #await trio.lowlevel.cancel_shielded_checkpoint()
    end = time.monotonic()
    print(f"{LOOP_SIZE / (end - start):.2f} schedules/second")

trio.run(main)
```
  • Loading branch information
njsmith committed Jun 14, 2020
1 parent 2d76954 commit 880e3af
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions trio/_core/_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
from ._traps import (
Abort,
wait_task_rescheduled,
cancel_shielded_checkpoint,
CancelShieldedCheckpoint,
PermanentlyDetachCoroutineObject,
WaitTaskRescheduled,
Expand Down Expand Up @@ -2284,8 +2285,17 @@ async def checkpoint():
:func:`checkpoint`.)
"""
with CancelScope(deadline=-inf):
await _core.wait_task_rescheduled(lambda _: _core.Abort.SUCCEEDED)
# The scheduler is what checks timeouts and converts them into
# cancellations. So by doing the schedule point first, we ensure that the
# cancel point has the most up-to-date info.
await cancel_shielded_checkpoint()
task = current_task()
task._cancel_points += 1
if task._cancel_status.effectively_cancelled or (
task is task._runner.main_task and task._runner.ki_pending
):
with CancelScope(deadline=-inf):
await _core.wait_task_rescheduled(lambda _: _core.Abort.SUCCEEDED)


async def checkpoint_if_cancelled():
Expand Down

0 comments on commit 880e3af

Please sign in to comment.