gh-3108: Avoid materializing f_locals for KI protection (#3110)

* gh-3108: avoid materializing f_locals by using weakrefs to code objects instead

* enable ki protection on async_generator objects

* avoid adding an extra coroutine wrapper to Task coros

* fix returning the wrong object in (enable|disable)_ki_protection

* remove KIProtectionSignature from _check_type_completeness.json

* fix refcycles

* add newsfragment

* fix mypy

* now that the type annotation for enable_ki_protection is fixed, we need to fix the use of Any

* pre-commit

* add test for ki protection leaking accross local functions

* add fix for ki protection leaking accross local functions

* do slots properly

* python 3.8 support

* test reading currently_ki_protected doesn't freeze locals

* cover some tricky bits of ki.py

* cover a potentially impossible scenario

* eek out some last coverage of the eeking out coverage tests

* even more partial coverage

* Update src/trio/_core/_ki.py

* cleaner _IdRef.__eq__

* if the current_task().coro.cr_frame is in the stack ki_protection_enabled is current_task()._ki_protected

* Update newsfragments/3108.bugfix.rst

* avoid copying code objects for ki protected function

* Update src/trio/_core/_ki.py

* Update src/trio/_core/_ki.py

Co-authored-by: EXPLOSION <[email protected]>

* remove workaround for 3.8

* Add docs and update news

Co-Authored-By: oremanj <[email protected]>

* wrap ki protection locals demos in async def so they work

* add newsfragment for 2670

* Apply suggestions from code review

* use peo614

* add tests for passing on inspect flags

* 'return; yield' isn't considered covered

* Update newsfragments/3108.bugfix.rst

* [pre-commit.ci] auto fixes from pre-commit.com hooks

---------

Co-authored-by: EXPLOSION <[email protected]>
Co-authored-by: oremanj <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/source/reference-lowlevel.rst

@@ -377,6 +377,46 @@ These transitions are accomplished using two function decorators:
    poorly-timed :exc:`KeyboardInterrupt` could leave the lock in an
    inconsistent state and cause a deadlock.
 
+   Since KeyboardInterrupt protection is tracked per code object, any attempt to
+   conditionally protect the same block of code in different ways is unlikely to behave
+   how you expect. If you try to conditionally protect a closure, it will be
+   unconditionally protected instead::
+
+       def example(protect: bool) -> bool:
+           def inner() -> bool:
+               return trio.lowlevel.currently_ki_protected()
+           if protect:
+               inner = trio.lowlevel.enable_ki_protection(inner)
+           return inner()
+
+       async def amain():
+           assert example(False) == False
+           assert example(True) == True  # once protected ...
+           assert example(False) == True  # ... always protected
+
+       trio.run(amain)
+
+   If you really need conditional protection, you can achieve it by giving each
+   KI-protected instance of the closure its own code object::
+
+       def example(protect: bool) -> bool:
+           def inner() -> bool:
+               return trio.lowlevel.currently_ki_protected()
+           if protect:
+               inner.__code__ = inner.__code__.replace()
+               inner = trio.lowlevel.enable_ki_protection(inner)
+           return inner()
+
+       async def amain():
+           assert example(False) == False
+           assert example(True) == True
+           assert example(False) == False
+
+       trio.run(amain)
+
+   (This isn't done by default because it carries some memory overhead and reduces
+   the potential for specializing optimizations in recent versions of CPython.)
+
 .. autofunction:: currently_ki_protected
 
 

newsfragments/2670.bugfix.rst

@@ -0,0 +1,2 @@
+:func:`inspect.iscoroutinefunction` and the like now give correct answers when
+called on KI-protected functions.

newsfragments/3108.bugfix.rst

@@ -0,0 +1,26 @@
+Rework KeyboardInterrupt protection to track code objects, rather than frames,
+as protected or not. The new implementation no longer needs to access
+``frame.f_locals`` dictionaries, so it won't artificially extend the lifetime of
+local variables. Since KeyboardInterrupt protection is now imposed statically
+(when a protected function is defined) rather than each time the function runs,
+its previously-noticeable performance overhead should now be near zero.
+The lack of a call-time wrapper has some other benefits as well:
+
+* :func:`inspect.iscoroutinefunction` and the like now give correct answers when
+  called on KI-protected functions.
+
+* Calling a synchronous KI-protected function no longer pushes an additional stack
+  frame, so tracebacks are clearer.
+
+* A synchronous KI-protected function invoked from C code (such as a weakref
+  finalizer) is now guaranteed to start executing; previously there would be a brief
+  window in which KeyboardInterrupt could be raised before the protection was
+  established.
+
+One minor drawback of the new approach is that multiple instances of the same
+closure share a single KeyboardInterrupt protection state (because they share a
+single code object). That means that if you apply
+`@enable_ki_protection <trio.lowlevel.enable_ki_protection>` to some of them
+and not others, you won't get the protection semantics you asked for. See the
+documentation of `@enable_ki_protection <trio.lowlevel.enable_ki_protection>`
+for more details and a workaround.

src/trio/_core/_generated_instrumentation.py

@@ -3,10 +3,9 @@
 # *************************************************************
 from __future__ import annotations
 
-import sys
 from typing import TYPE_CHECKING
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -15,6 +14,7 @@
 __all__ = ["add_instrument", "remove_instrument"]
 
 
+@enable_ki_protection
 def add_instrument(instrument: Instrument) -> None:
     """Start instrumenting the current run loop with the given instrument.
 
@@ -24,13 +24,13 @@ def add_instrument(instrument: Instrument) -> None:
     If ``instrument`` is already active, does nothing.
 
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.instruments.add_instrument(instrument)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def remove_instrument(instrument: Instrument) -> None:
     """Stop instrumenting the current run loop with the given instrument.
 
@@ -44,7 +44,6 @@ def remove_instrument(instrument: Instrument) -> None:
           deactivated.
 
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.instruments.remove_instrument(instrument)
     except AttributeError:

src/trio/_core/_generated_io_epoll.py

@@ -6,7 +6,7 @@
 import sys
 from typing import TYPE_CHECKING
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -18,6 +18,7 @@
 __all__ = ["notify_closing", "wait_readable", "wait_writable"]
 
 
+@enable_ki_protection
 async def wait_readable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is readable.
 
@@ -40,13 +41,13 @@ async def wait_readable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_readable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_writable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is writable.
 
@@ -59,13 +60,13 @@ async def wait_writable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_writable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def notify_closing(fd: int | _HasFileNo) -> None:
     """Notify waiters of the given object that it will be closed.
 
@@ -91,7 +92,6 @@ def notify_closing(fd: int | _HasFileNo) -> None:
     step, so other tasks won't be able to tell what order they happened
     in anyway.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.notify_closing(fd)
     except AttributeError:

src/trio/_core/_generated_io_kqueue.py

@@ -6,7 +6,7 @@
 import sys
 from typing import TYPE_CHECKING
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -31,18 +31,19 @@
 ]
 
 
+@enable_ki_protection
 def current_kqueue() -> select.kqueue:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.current_kqueue()
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def monitor_kevent(
     ident: int,
     filter: int,
@@ -51,13 +52,13 @@ def monitor_kevent(
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.monitor_kevent(ident, filter)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_kevent(
     ident: int,
     filter: int,
@@ -67,7 +68,6 @@ async def wait_kevent(
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_kevent(
             ident,
@@ -78,6 +78,7 @@ async def wait_kevent(
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_readable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is readable.
 
@@ -100,13 +101,13 @@ async def wait_readable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_readable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_writable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is writable.
 
@@ -119,13 +120,13 @@ async def wait_writable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_writable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def notify_closing(fd: int | _HasFileNo) -> None:
     """Notify waiters of the given object that it will be closed.
 
@@ -151,7 +152,6 @@ def notify_closing(fd: int | _HasFileNo) -> None:
     step, so other tasks won't be able to tell what order they happened
     in anyway.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.notify_closing(fd)
     except AttributeError: