|
| 1 | +# Copyright 2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Caching for values that belong to the event loop that created them.""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import asyncio |
| 20 | +from collections.abc import Callable |
| 21 | +from functools import cached_property |
| 22 | +import threading |
| 23 | +from typing import Any |
| 24 | +from typing import cast |
| 25 | +from typing import Generic |
| 26 | +from typing import TYPE_CHECKING |
| 27 | +from typing import TypeVar |
| 28 | +import weakref |
| 29 | + |
| 30 | +_T = TypeVar('_T') |
| 31 | + |
| 32 | + |
| 33 | +class _NoRunningLoop: |
| 34 | + """Stands in for "read from outside any event loop" as a cache key. |
| 35 | +
|
| 36 | + Weak keys cannot be ``None``, and every read made with no loop running has to |
| 37 | + share one entry. This module holds the single instance for the life of the |
| 38 | + process, so that entry is never evicted. |
| 39 | + """ |
| 40 | + |
| 41 | + |
| 42 | +_NO_RUNNING_LOOP = _NoRunningLoop() |
| 43 | + |
| 44 | + |
| 45 | +class _PerLoopCache(Generic[_T]): |
| 46 | + """Holds one value per event loop for a single attribute of a single owner. |
| 47 | +
|
| 48 | + An async client holds sockets, locks and tasks that belong to the event loop |
| 49 | + that was running when they were opened. Caching such a client once for the |
| 50 | + lifetime of its owner therefore breaks as soon as a second loop uses it: the |
| 51 | + second loop reaches into a loop that has already been closed. Owners outlive |
| 52 | + loops routinely -- the synchronous runner entry points run each call in a new |
| 53 | + loop, and servers that dispatch requests to a thread pool do the same. |
| 54 | +
|
| 55 | + Values are keyed weakly by the running loop, so every loop builds its own and |
| 56 | + an entry vanishes once its loop is collected, whether or not that loop was |
| 57 | + ever closed. Reading a value for a new loop additionally discards the entries |
| 58 | + whose loop reports itself closed, which releases the value promptly for a |
| 59 | + closed loop that something else still holds. Reads with no loop running share |
| 60 | + a single entry, which is never discarded. Sharing that one entry is safe |
| 61 | + because a reader with no loop running goes on to use the synchronous surface, |
| 62 | + which has no loop affinity. A caller that instead read the value off-loop and |
| 63 | + drove it from a loop would keep using it after that loop closed, since this |
| 64 | + is the one key the sweep never reaches. |
| 65 | +
|
| 66 | + Keying weakly on the loop object is safe against address reuse. A lookup hits |
| 67 | + only when the stored key both hashes equal to and compares equal to the key |
| 68 | + being looked up, and a ``weakref.ref`` whose referent has died compares equal |
| 69 | + to nothing but itself. A reference to a new loop that happens to land on a |
| 70 | + collected loop's address is therefore a miss even though the two addresses, |
| 71 | + and so the two hashes, agree. The reference's own callback has usually |
| 72 | + dropped the entry by then in any case. |
| 73 | + """ |
| 74 | + |
| 75 | + def __init__(self) -> None: |
| 76 | + self._lock = threading.Lock() |
| 77 | + self._values: weakref.WeakKeyDictionary[Any, _T] = ( |
| 78 | + weakref.WeakKeyDictionary() |
| 79 | + ) |
| 80 | + self._build_locks: weakref.WeakKeyDictionary[Any, threading.Lock] = ( |
| 81 | + weakref.WeakKeyDictionary() |
| 82 | + ) |
| 83 | + |
| 84 | + def get(self, build: Callable[[], _T]) -> _T: |
| 85 | + """Returns the value for the running loop, calling ``build`` on a miss.""" |
| 86 | + try: |
| 87 | + loop: Any = asyncio.get_running_loop() |
| 88 | + except RuntimeError: |
| 89 | + loop = _NO_RUNNING_LOOP |
| 90 | + |
| 91 | + # Membership rather than a default, so a cached value of None is a hit. |
| 92 | + # The key is held by `loop` for as long as this call runs, so the entry |
| 93 | + # cannot be collected between the two lookups. |
| 94 | + if loop in self._values: |
| 95 | + return self._values[loop] |
| 96 | + |
| 97 | + # Only the miss path is guarded. A hit is an unlocked dictionary lookup, |
| 98 | + # and the entry it finds belongs to the running loop, which cannot be |
| 99 | + # closed. The shared lock covers the eviction sweep and the insert, which |
| 100 | + # would otherwise raise "dictionary changed size during iteration", and |
| 101 | + # hands out one build lock per key. Building is left to that per-key lock: |
| 102 | + # a build opens sockets and takes tens of milliseconds, and every fresh |
| 103 | + # loop misses cold, so building under the shared lock would make loops |
| 104 | + # that share nothing queue behind each other. |
| 105 | + with self._lock: |
| 106 | + if loop in self._values: |
| 107 | + return self._values[loop] |
| 108 | + # Dropping an entry runs no disposal hook. An entry is only dropped once |
| 109 | + # its loop is closed or already collected, and neither state can run the |
| 110 | + # asynchronous teardown a client would need; the value releases its |
| 111 | + # sockets when it is itself collected. |
| 112 | + # |
| 113 | + # Iterate the mapping itself rather than its underlying table: that |
| 114 | + # defers the collection callbacks, which fire on whichever thread |
| 115 | + # happened to drop the last reference to a loop. |
| 116 | + for stale in [ |
| 117 | + key |
| 118 | + for key in self._values |
| 119 | + if key is not _NO_RUNNING_LOOP and key.is_closed() |
| 120 | + ]: |
| 121 | + del self._values[stale] |
| 122 | + build_lock = self._build_locks.setdefault(loop, threading.Lock()) |
| 123 | + |
| 124 | + # Two threads missing the same key at once would otherwise each build a |
| 125 | + # value and only one would be kept: the other is live, holds open sockets, |
| 126 | + # and nothing goes on to close it. Only reads made with no loop running |
| 127 | + # can collide that way, since they are the only ones that share a key. |
| 128 | + with build_lock: |
| 129 | + if loop in self._values: |
| 130 | + return self._values[loop] |
| 131 | + built = build() |
| 132 | + with self._lock: |
| 133 | + self._values[loop] = built |
| 134 | + return built |
| 135 | + |
| 136 | + |
| 137 | +def _cache_for(owner: Any, key: str) -> _PerLoopCache[Any]: |
| 138 | + """Returns ``owner``'s cache for ``key``, creating it on first use. |
| 139 | +
|
| 140 | + The cache lives in ``owner.__dict__``, the way |
| 141 | + ``functools.cached_property`` stores its own value, so an owner that is |
| 142 | + copied or discarded takes its values with it. |
| 143 | + """ |
| 144 | + cache = owner.__dict__.get(key) |
| 145 | + if cache is None: |
| 146 | + # setdefault so that two threads racing the first read share one cache. |
| 147 | + cache = owner.__dict__.setdefault(key, _PerLoopCache()) |
| 148 | + return cast('_PerLoopCache[Any]', cache) |
| 149 | + |
| 150 | + |
| 151 | +def per_loop_value(owner: Any, key: str, build: Callable[[], _T]) -> _T: |
| 152 | + """Returns a value for ``owner``, built once per running event loop. |
| 153 | +
|
| 154 | + Args: |
| 155 | + owner: The object the value belongs to. |
| 156 | + key: The ``owner.__dict__`` entry holding the cache. Must not collide with a |
| 157 | + field name. |
| 158 | + build: Called to produce the value when the running loop has none. |
| 159 | +
|
| 160 | + Returns: |
| 161 | + The value belonging to the running event loop. |
| 162 | + """ |
| 163 | + cache: _PerLoopCache[_T] = _cache_for(owner, key) |
| 164 | + return cache.get(build) |
| 165 | + |
| 166 | + |
| 167 | +if TYPE_CHECKING: |
| 168 | + # A type checker reads this as the decorator it stands in for, rather than as |
| 169 | + # the subclass below. Subclasses out in the wild override a decorated |
| 170 | + # attribute with a plain ``functools.cached_property``, and a checker rejects |
| 171 | + # an override whose descriptor type differs from the one it overrides -- it |
| 172 | + # recognizes only ``functools.cached_property`` itself, not a subclass of it. |
| 173 | + # Declaring the alias keeps every such override valid. |
| 174 | + PerLoopCachedProperty = cached_property |
| 175 | +else: |
| 176 | + |
| 177 | + class PerLoopCachedProperty(cached_property): |
| 178 | + """A ``cached_property`` that keeps one value per running event loop. |
| 179 | +
|
| 180 | + Reads behave like ``cached_property`` except that the value is cached |
| 181 | + against the running loop rather than once for the owner. See |
| 182 | + ``_PerLoopCache`` for why a client has to be cached that way. |
| 183 | +
|
| 184 | + Like ``cached_property`` this defines no ``__set__``, so assigning over the |
| 185 | + attribute still writes to the owner's ``__dict__`` and shadows the |
| 186 | + descriptor from then on. The cache lives under a separate key, leaving the |
| 187 | + attribute's own slot free for those assignments. |
| 188 | +
|
| 189 | + Deletion is where the two part company. ``cached_property`` invalidates |
| 190 | + itself through ``del owner.attr``; here a read writes only the cache, so |
| 191 | + that raises ``AttributeError`` unless a value was assigned over the |
| 192 | + attribute first, and deleting an assigned value uncovers the descriptor |
| 193 | + rather than clearing what it has cached. |
| 194 | +
|
| 195 | + A Pydantic owner must additionally list this class in |
| 196 | + ``model_config['ignored_types']``, because Pydantic exempts |
| 197 | + ``cached_property`` by module name rather than by type. |
| 198 | + """ |
| 199 | + |
| 200 | + def __get__(self, instance, owner=None): |
| 201 | + if instance is None: |
| 202 | + return self |
| 203 | + return per_loop_value( |
| 204 | + instance, f'_per_loop_{self.attrname}', lambda: self.func(instance) |
| 205 | + ) |
0 commit comments