local.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. import asyncio
  2. import contextlib
  3. import contextvars
  4. import threading
  5. from typing import Any, Union
  6. class _Storage:
  7. """Thread-tagged storage for a non-thread-critical ``Local``.
  8. The data is tagged with the identity of the thread that owns it. This lets
  9. ``_CVar`` ignore data that leaked into an unrelated thread.
  10. Python 3.14 added ``sys.flags.thread_inherit_context``, which is enabled by
  11. default on free-threaded builds. When set, a new thread starts with a copy
  12. of the spawning thread's context instead of an empty one, so the contextvar
  13. backing a ``Local`` would otherwise be visible in any thread spawned from
  14. one that had set it -- breaking the documented "thread-local in sync
  15. threads" behaviour. asgiref re-homes the storage to the current thread at
  16. the points where it *intentionally* moves work between threads (see
  17. ``asgiref.sync._restore_context``); data merely inherited by an unrelated
  18. thread is never re-homed and so stays isolated.
  19. """
  20. __slots__ = ("thread_id", "data")
  21. def __init__(self, thread_id: int, data: dict[str, Any]) -> None:
  22. self.thread_id = thread_id
  23. self.data = data
  24. def _rehome(storage: "_Storage") -> "_Storage":
  25. """Return a copy of *storage* owned by the current thread."""
  26. return _Storage(threading.get_ident(), storage.data)
  27. class _CVar:
  28. """Storage utility for Local."""
  29. def __init__(self) -> None:
  30. self._data: "contextvars.ContextVar[_Storage]" = contextvars.ContextVar(
  31. "asgiref.local"
  32. )
  33. def _storage(self) -> "_Storage":
  34. # Only return storage that belongs to the current thread. Storage with
  35. # a different thread id was inherited by this thread (rather than
  36. # intentionally moved here by asgiref) and must not be visible.
  37. storage = self._data.get(None)
  38. if storage is None or storage.thread_id != threading.get_ident():
  39. return _Storage(threading.get_ident(), {})
  40. return storage
  41. def __getattr__(self, key):
  42. try:
  43. return self._storage().data[key]
  44. except KeyError:
  45. raise AttributeError(f"{self!r} object has no attribute {key!r}")
  46. def __setattr__(self, key: str, value: Any) -> None:
  47. if key == "_data":
  48. return super().__setattr__(key, value)
  49. data = self._storage().data.copy()
  50. data[key] = value
  51. self._data.set(_Storage(threading.get_ident(), data))
  52. def __delattr__(self, key: str) -> None:
  53. data = self._storage().data.copy()
  54. if key in data:
  55. del data[key]
  56. self._data.set(_Storage(threading.get_ident(), data))
  57. else:
  58. raise AttributeError(f"{self!r} object has no attribute {key!r}")
  59. class Local:
  60. """Local storage for async tasks.
  61. This is a namespace object (similar to `threading.local`) where data is
  62. also local to the current async task (if there is one).
  63. In async threads, local means in the same sense as the `contextvars`
  64. module - i.e. a value set in an async frame will be visible:
  65. - to other async code `await`-ed from this frame.
  66. - to tasks spawned using `asyncio` utilities (`create_task`, `wait_for`,
  67. `gather` and probably others).
  68. - to code scheduled in a sync thread using `sync_to_async`
  69. In "sync" threads (a thread with no async event loop running), the
  70. data is thread-local, but additionally shared with async code executed
  71. via the `async_to_sync` utility, which schedules async code in a new thread
  72. and copies context across to that thread.
  73. If `thread_critical` is True, then the local will only be visible per-thread,
  74. behaving exactly like `threading.local` if the thread is sync, and as
  75. `contextvars` if the thread is async. This allows genuinely thread-sensitive
  76. code (such as DB handles) to be kept strictly to their initial thread and
  77. disable the sharing across `sync_to_async` and `async_to_sync` wrapped calls.
  78. Unlike plain `contextvars` objects, this utility is threadsafe.
  79. """
  80. def __init__(self, thread_critical: bool = False) -> None:
  81. self._thread_critical = thread_critical
  82. self._thread_lock = threading.RLock()
  83. self._storage: "Union[threading.local, _CVar]"
  84. if thread_critical:
  85. # Thread-local storage
  86. self._storage = threading.local()
  87. else:
  88. # Contextvar storage
  89. self._storage = _CVar()
  90. @contextlib.contextmanager
  91. def _lock_storage(self):
  92. # Thread safe access to storage
  93. if self._thread_critical:
  94. is_async = True
  95. try:
  96. # this is a test for are we in a async or sync
  97. # thread - will raise RuntimeError if there is
  98. # no current loop
  99. asyncio.get_running_loop()
  100. except RuntimeError:
  101. is_async = False
  102. if not is_async:
  103. # We are in a sync thread, the storage is
  104. # just the plain thread local (i.e, "global within
  105. # this thread" - it doesn't matter where you are
  106. # in a call stack you see the same storage)
  107. yield self._storage
  108. else:
  109. # We are in an async thread - storage is still
  110. # local to this thread, but additionally should
  111. # behave like a context var (is only visible with
  112. # the same async call stack)
  113. # Ensure context exists in the current thread
  114. if not hasattr(self._storage, "cvar"):
  115. self._storage.cvar = _CVar()
  116. # self._storage is a thread local, so the members
  117. # can't be accessed in another thread (we don't
  118. # need any locks)
  119. yield self._storage.cvar
  120. else:
  121. # Lock for thread_critical=False as other threads
  122. # can access the exact same storage object
  123. with self._thread_lock:
  124. yield self._storage
  125. def __getattr__(self, key):
  126. with self._lock_storage() as storage:
  127. return getattr(storage, key)
  128. def __setattr__(self, key, value):
  129. if key in ("_local", "_storage", "_thread_critical", "_thread_lock"):
  130. return super().__setattr__(key, value)
  131. with self._lock_storage() as storage:
  132. setattr(storage, key, value)
  133. def __delattr__(self, key):
  134. with self._lock_storage() as storage:
  135. delattr(storage, key)