LockStore: fix acquiring a lock via LockStore.try_acquire_lock (#12832)

Signed-off-by: Sumner Evans <sumner@beeper.com>
This commit is contained in:
Sumner Evans 2022-05-30 02:41:13 -06:00 committed by GitHub
parent 28989cb301
commit bda4600399
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 73 additions and 1 deletions

View file

@ -13,7 +13,7 @@
# limitations under the License.
import logging
from types import TracebackType
from typing import TYPE_CHECKING, Optional, Tuple, Type
from typing import TYPE_CHECKING, Optional, Set, Tuple, Type
from weakref import WeakValueDictionary
from twisted.internet.interfaces import IReactorCore
@ -84,6 +84,8 @@ class LockStore(SQLBaseStore):
self._on_shutdown,
)
self._acquiring_locks: Set[Tuple[str, str]] = set()
@wrap_as_background_process("LockStore._on_shutdown")
async def _on_shutdown(self) -> None:
"""Called when the server is shutting down"""
@ -103,6 +105,21 @@ class LockStore(SQLBaseStore):
context manager if the lock is successfully acquired, which *must* be
used (otherwise the lock will leak).
"""
if (lock_name, lock_key) in self._acquiring_locks:
return None
try:
self._acquiring_locks.add((lock_name, lock_key))
return await self._try_acquire_lock(lock_name, lock_key)
finally:
self._acquiring_locks.discard((lock_name, lock_key))
async def _try_acquire_lock(
self, lock_name: str, lock_key: str
) -> Optional["Lock"]:
"""Try to acquire a lock for the given name/key. Will return an async
context manager if the lock is successfully acquired, which *must* be
used (otherwise the lock will leak).
"""
# Check if this process has taken out a lock and if it's still valid.
lock = self._live_tokens.get((lock_name, lock_key))