forked from goodboy/tractor
Rework interface: pass func and kwargs
After more extensive testing I realized that keying on the context manager *instance id* isn't going to work since each entering task is going to create a unique key XD Instead pass the manager function as `acm_func` and optionally allow keying the resource on the passed `kwargs` (if hashable) or the `key:str`. Further, pass the key to the enterer task and avoid a separate keying scheme for the manager versus the value it delivers. Don't bother with checking and releasing the lock in `finally:` block, it should be an error if it's still locked.context_caching
parent
3826bc9972
commit
52627a6326
|
@ -24,6 +24,7 @@ from typing import (
|
|||
AsyncContextManager,
|
||||
AsyncGenerator,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Hashable,
|
||||
Optional,
|
||||
Sequence,
|
||||
|
@ -131,69 +132,76 @@ class _Cache:
|
|||
async def run_ctx(
|
||||
cls,
|
||||
mng,
|
||||
key,
|
||||
ctx_key: tuple,
|
||||
task_status: TaskStatus[T] = trio.TASK_STATUS_IGNORED,
|
||||
|
||||
) -> None:
|
||||
async with mng as value:
|
||||
|
||||
_, no_more_users = cls.resources[id(mng)]
|
||||
cls.values[key] = value
|
||||
_, no_more_users = cls.resources[ctx_key]
|
||||
cls.values[ctx_key] = value
|
||||
task_status.started(value)
|
||||
try:
|
||||
await no_more_users.wait()
|
||||
finally:
|
||||
value = cls.values.pop(key)
|
||||
# discard nursery ref so it won't be re-used (an error)
|
||||
cls.resources.pop(id(mng))
|
||||
# discard nursery ref so it won't be re-used (an error)?
|
||||
value = cls.values.pop(ctx_key)
|
||||
cls.resources.pop(ctx_key)
|
||||
|
||||
|
||||
@acm
|
||||
async def maybe_open_context(
|
||||
|
||||
key: Hashable,
|
||||
mngr: AsyncContextManager[T],
|
||||
acm_func: Callable[..., AsyncContextManager[T]],
|
||||
|
||||
# XXX: used as cache key after conversion to tuple
|
||||
# and all embedded values must also be hashable
|
||||
kwargs: Optional[dict] = {},
|
||||
key: Hashable = None,
|
||||
|
||||
) -> AsyncIterator[tuple[bool, T]]:
|
||||
'''
|
||||
Maybe open a context manager if there is not already a _Cached
|
||||
version for the provided ``key``. Return the _Cached instance on
|
||||
a _Cache hit.
|
||||
version for the provided ``key`` for *this* actor. Return the
|
||||
_Cached instance on a _Cache hit.
|
||||
|
||||
'''
|
||||
# lock resource acquisition around task racing / ``trio``'s
|
||||
# scheduler protocol
|
||||
await _Cache.lock.acquire()
|
||||
|
||||
ctx_key = id(mngr)
|
||||
ctx_key = (id(acm_func), key or tuple(kwargs.items()))
|
||||
value = None
|
||||
|
||||
try:
|
||||
value = _Cache.values[key]
|
||||
|
||||
except KeyError:
|
||||
log.info(f'Allocating new resource for {key}')
|
||||
|
||||
# **critical section** that should prevent other tasks from
|
||||
# checking the _Cache until complete otherwise the scheduler
|
||||
# may switch and by accident we create more then one feed.
|
||||
# may switch and by accident we create more then one resource.
|
||||
value = _Cache.values[ctx_key]
|
||||
|
||||
except KeyError:
|
||||
log.info(f'Allocating new resource for {ctx_key}')
|
||||
|
||||
mngr = acm_func(**kwargs)
|
||||
# TODO: avoid pulling from ``tractor`` internals and
|
||||
# instead offer a "root nursery" in piker actors?
|
||||
service_n = current_actor()._service_n
|
||||
|
||||
# TODO: does this need to be a tractor "root nursery"?
|
||||
assert not _Cache.resources.get(ctx_key), f'Resource exists? {ctx_key}'
|
||||
ln, _ = _Cache.resources[ctx_key] = (service_n, trio.Event())
|
||||
resources = _Cache.resources
|
||||
assert not resources.get(ctx_key), f'Resource exists? {ctx_key}'
|
||||
ln, _ = resources[ctx_key] = (service_n, trio.Event())
|
||||
|
||||
value = await ln.start(_Cache.run_ctx, mngr, key)
|
||||
value = await ln.start(
|
||||
_Cache.run_ctx,
|
||||
mngr,
|
||||
ctx_key,
|
||||
)
|
||||
_Cache.users += 1
|
||||
_Cache.lock.release()
|
||||
|
||||
yield False, value
|
||||
|
||||
else:
|
||||
log.info(f'Reusing _Cached resource for {key}')
|
||||
log.info(f'Reusing _Cached resource for {ctx_key}')
|
||||
_Cache.users += 1
|
||||
_Cache.lock.release()
|
||||
yield True, value
|
||||
|
@ -204,13 +212,8 @@ async def maybe_open_context(
|
|||
if value is not None:
|
||||
# if no more consumers, teardown the client
|
||||
if _Cache.users <= 0:
|
||||
log.info(f'De-allocating resource for {key}')
|
||||
|
||||
if _Cache.lock.locked():
|
||||
_Cache.lock.release()
|
||||
log.info(f'De-allocating resource for {ctx_key}')
|
||||
|
||||
# terminate mngr nursery
|
||||
entry = _Cache.resources.get(ctx_key)
|
||||
if entry:
|
||||
_, no_more_users = entry
|
||||
_, no_more_users = _Cache.resources[ctx_key]
|
||||
no_more_users.set()
|
||||
|
|
Loading…
Reference in New Issue