21
21
22
22
# TODO: pick what public namespace to re-export this from.
23
23
24
+ # This is one of the few things that should unreservedly be reimplemented natively.
25
+ # Atomic ops allow the design here to be lockfree.
26
+ # Actually using a futex here might improve performance further, though I doubt it.
27
+
28
+ # We can also do a little bit better prior to making it native by only locking
29
+ # if we observe multiple threads. The gain for this is minor, but not insignificant.
30
+
31
+ # This particular lock implementation optimizes for the uncontested lock case.
32
+ # This is the most important optimization for well-designed parallel code
33
+ # with fine-grained lock use. It's also barging, which is known to be better
34
+ # than FIFO in the case of highly contested locks.
35
+
24
36
25
37
class AsyncLock :
26
38
"""An async lock that doesn't bind to an event loop."""
@@ -36,36 +48,26 @@ def __init__(self) -> None:
36
48
self ._lockv : bool = False
37
49
self ._internal_lock : threading .RLock = threading .RLock ()
38
50
39
- def __locked (self ) -> bool :
40
- with self ._internal_lock :
41
- return self ._lockv or (any (not w .cancelled () for w in (self ._waiters )))
42
-
43
51
async def __aenter__ (self ) -> None :
44
- # placing the body of acquire here
45
- # causes problems with eager tasks factories.
46
- await self .__acquire ()
47
-
48
- async def __acquire (self ) -> None :
49
52
with self ._internal_lock :
50
- if not self .__locked () :
53
+ if not self ._lockv :
51
54
self ._lockv = True
52
55
return
53
56
54
57
fut : cf .Future [None ] = cf .Future ()
55
58
56
- with self ._internal_lock :
57
- self ._waiters .append (fut )
59
+ self ._waiters .append (fut )
58
60
59
61
try :
60
62
await asyncio .wrap_future (fut )
61
63
except asyncio .CancelledError :
62
- if fut .done () and not fut .cancelled ():
63
- self ._lockv = False
64
+ with self ._internal_lock :
65
+ if not self ._lockv :
66
+ self ._maybe_wake ()
64
67
raise
65
68
66
- finally :
67
- self ._maybe_wake ()
68
- return
69
+ with self ._internal_lock :
70
+ self ._lockv = True
69
71
70
72
def _maybe_wake (self ) -> None :
71
73
with self ._internal_lock :
0 commit comments