tractor/tests/test_resource_cache.py

'''
Suites for our `.trionics.maybe_open_context()` multi-task
shared-cached `@acm` API.

'''
from contextlib import asynccontextmanager as acm
import platform
from typing import Awaitable

import pytest
import trio
from trio.testing import wait_all_tasks_blocked
import tractor
from tractor.trionics import (
    maybe_open_context,
    collapse_eg,
)
from tractor.log import (
    get_console_log,
    get_logger,
)

log = get_logger()

_resource: int = 0


@acm
async def maybe_increment_counter(task_name: str):
    global _resource

    _resource += 1
    await trio.lowlevel.checkpoint()
    yield _resource
    await trio.lowlevel.checkpoint()
    _resource -= 1


@pytest.mark.parametrize(
    'key_on',
    ['key_value', 'kwargs'],
    ids="key_on={}".format,
)
def test_resource_only_entered_once(key_on):
    global _resource
    _resource = 0

    key = None
    if key_on == 'key_value':
        key = 'some_common_key'

    async def main():
        cache_active: bool = False

        async def enter_cached_mngr(name: str):
            nonlocal cache_active

            if key_on == 'kwargs':
                # make a common kwargs input to key on it
                kwargs = {'task_name': 'same_task_name'}
                assert key is None
            else:
                # different task names per task will be used
                kwargs = {'task_name': name}

            async with maybe_open_context(
                maybe_increment_counter,
                kwargs=kwargs,
                key=key,

            ) as (cache_hit, resource):
                if cache_hit:
                    try:
                        cache_active = True
                        assert resource == 1
                        await trio.sleep_forever()
                    finally:
                        cache_active = False
                else:
                    assert resource == 1
                    await trio.sleep_forever()

        with trio.move_on_after(0.5):
            async with (
                tractor.open_root_actor(),
                trio.open_nursery() as tn,
            ):
                for i in range(10):
                    tn.start_soon(
                        enter_cached_mngr,
                        f'task_{i}',
                    )
                    await trio.sleep(0.001)

    trio.run(main)


def test_last_moc_user_waits_for_resource_exit():
    '''
    Verify the final user cannot return before resource teardown.

    Previously the final `maybe_open_context()` user only signalled
    `_Cache.run_ctx()` through its `no_more_users` event. The user
    then returned while the service task was still running the
    resource's `__aexit__()`, so callers could observe stale external
    state immediately after their `async with` block.

    The resource sets `exit_started` before blocking on
    `allow_exit`. The user task must remain inside MOC until the test
    releases that deterministic checkpoint and `__aexit__()` sets
    `exit_finished`.

    '''
    async def main():
        exit_started = trio.Event()
        allow_exit = trio.Event()
        exit_finished = trio.Event()
        user_returned = trio.Event()

        @acm
        async def open_resource():
            try:
                yield
            finally:
                exit_started.set()
                await allow_exit.wait()
                exit_finished.set()

        async def use_resource():
            async with maybe_open_context(open_resource):
                pass

            assert exit_finished.is_set()
            user_returned.set()

        async with (
            tractor.open_root_actor(),
            trio.open_nursery() as tn,
        ):
            tn.start_soon(use_resource)
            await exit_started.wait()
            assert not user_returned.is_set()
            allow_exit.set()
            await user_returned.wait()

    trio.run(main)


def test_moc_delivers_resource_exit_error():
    '''
    Verify a resource exit error reaches the final MOC user.

    Previously `_Cache.run_ctx()` executed the cached resource's
    `__aexit__()` after the final user had returned. An exit failure
    therefore surfaced later through the actor service nursery rather
    than at the user's `async with maybe_open_context()` boundary.

    This resource raises a unique `ResourceExitError` during exit.
    Catching that exact instance around MOC proves the service task
    delivered the failure to the final user without replacing it.

    '''
    class ResourceExitError(Exception):
        pass

    exit_error = ResourceExitError('resource exit failed')

    async def main():
        @acm
        async def open_resource():
            yield
            raise exit_error

        async with tractor.open_root_actor():
            with pytest.raises(ResourceExitError) as exc_info:
                async with maybe_open_context(open_resource):
                    pass

            assert exc_info.value is exit_error

    trio.run(main)


def test_moc_final_user_cancellation_waits_for_exit():
    '''
    Verify final-user cancellation still waits for successful exit.

    Previously cancellation escaped the final MOC user immediately
    after it signalled `_Cache.run_ctx()`, leaving resource exit to
    finish later in the actor service task. This violated the context
    manager boundary even when cleanup itself succeeded.

    The consumer cancels its own scope while holding the sole cached
    resource. The resource sets `exit_finished` from its `finally`
    block, and the consumer checks that event immediately after its
    cancel scope catches `trio.Cancelled`. This proves MOC's
    completion wait is shielded without suppressing the original
    cancellation.

    '''
    async def main():
        exit_finished = trio.Event()

        @acm
        async def open_resource():
            try:
                yield
            finally:
                exit_finished.set()

        async with tractor.open_root_actor():
            with trio.CancelScope() as cs:
                async with maybe_open_context(open_resource):
                    cs.cancel()
                    await trio.sleep_forever()

            assert cs.cancelled_caught
            assert exit_finished.is_set()

    trio.run(main)


def test_moc_exit_error_masks_final_user_cancellation():
    '''
    Verify cleanup errors survive final-user cancellation.

    A cancelled final user previously signalled `no_more_users` and
    propagated `trio.Cancelled` before `_Cache.run_ctx()` completed
    resource exit. If `__aexit__()` then failed, its error was
    detached from the API call which caused teardown.

    The consumer cancels its own scope at a deterministic checkpoint
    inside MOC. Resource exit raises `ResourceExitError`; observing
    that exact error outside the cancel scope proves MOC shields the
    completion wait and applies normal context-manager masking, where
    a cleanup failure replaces the active cancellation.

    '''
    class ResourceExitError(Exception):
        pass

    exit_error = ResourceExitError('resource exit failed')

    async def main():
        @acm
        async def open_resource():
            yield
            raise exit_error

        async with tractor.open_root_actor():
            with pytest.raises(ResourceExitError) as exc_info:
                with trio.CancelScope() as cs:
                    async with maybe_open_context(open_resource):
                        cs.cancel()
                        await trio.sleep_forever()

            assert exc_info.value is exit_error

    trio.run(main)


def test_moc_service_nursery_cancellation_completes_exit():
    '''
    Verify service-nursery cancellation cannot strand a final user.

    `_Cache.run_ctx()` and an MOC consumer may share a
    caller-provided service nursery. Cancelling that nursery
    interrupts the service task's `no_more_users` wait and the
    consumer body together. A shielded final-user wait would deadlock
    if `run_ctx()` failed to publish completion while propagating its
    own `trio.Cancelled`.

    The outer task waits for resource entry, then cancels the exact
    nursery containing both tasks. The resource shields one cleanup
    checkpoint and sets `exit_finished`; observing both it and
    `service_finished` proves cancellation propagated normally while
    MOC's completion handshake terminated deterministically.

    '''
    async def main():
        resource_entered = trio.Event()
        exit_finished = trio.Event()
        service_finished = trio.Event()
        service_tn: trio.Nursery|None = None

        @acm
        async def open_resource():
            try:
                resource_entered.set()
                yield
            finally:
                with trio.CancelScope(shield=True):
                    await trio.lowlevel.checkpoint()
                    exit_finished.set()

        async def use_resource(tn: trio.Nursery):
            async with maybe_open_context(
                open_resource,
                tn=tn,
            ):
                await trio.sleep_forever()

        async def run_service():
            nonlocal service_tn

            async with trio.open_nursery() as tn:
                service_tn = tn
                tn.start_soon(use_resource, tn)

            service_finished.set()

        async with trio.open_nursery() as outer_tn:
            outer_tn.start_soon(run_service)
            await resource_entered.wait()
            assert service_tn is not None
            service_tn.cancel_scope.cancel()
            await service_finished.wait()

        assert exit_finished.is_set()

    trio.run(main)


@tractor.context
async def streamer(
    ctx: tractor.Context,
    seq: list[int] = list(range(1000)),
) -> None:

    await ctx.started()
    async with ctx.open_stream() as stream:
        for val in seq:
            await stream.send(val)
            await trio.sleep(0.001)

    print('producer finished')


@acm
async def open_stream() -> Awaitable[
    tuple[
        tractor.ActorNursery,
        tractor.MsgStream,
    ]
]:
    try:
        async with tractor.open_nursery() as an:
            portal = await an.start_actor(
                'streamer',
                enable_modules=[__name__],
            )
            try:
                async with (
                    portal.open_context(streamer) as (ctx, first),
                    ctx.open_stream() as stream,
                ):
                    print('Entered open_stream() caller')
                    yield an, stream
                    print('Exited open_stream() caller')

            finally:
                print(
                    'Cancelling streamer with,\n'
                    '=> `Portal.cancel_actor()`'
                )
                await portal.cancel_actor()
                print('Cancelled streamer')

    except Exception as err:
        print(
            f'`open_stream()` errored?\n'
            f'{err!r}\n'
        )
        await tractor.pause(shield=True)
        raise err


@acm
async def maybe_open_stream(taskname: str):
    async with maybe_open_context(
        # NOTE: all secondary tasks should cache hit on the same key
        acm_func=open_stream,
    ) as (
        cache_hit,
        (an, stream)
    ):
        # when the actor + portal + ctx + stream has already been
        # allocated we want to just bcast to this task.
        if cache_hit:
            print(f'{taskname} loaded from cache')

            # add a new broadcast subscription for the quote stream
            # if this feed is already allocated by the first
            # task that entereed
            async with stream.subscribe() as bstream:
                yield an, bstream
                print(
                    f'cached task exited\n'
                    f')>\n'
                    f' |_{taskname}\n'
                )

            # we should always unreg the "cloned" bcrc for this
            # consumer-task
            assert id(bstream) not in bstream._state.subs

        else:
            # yield the actual stream
            try:
                yield an, stream
            finally:
                print(
                    f'NON-cached task exited\n'
                    f')>\n'
                    f' |_{taskname}\n'
                )

        first_bstream = stream._broadcaster
        bcrx_state = first_bstream._state
        subs: dict[int, int] = bcrx_state.subs
        if len(subs) == 1:
            assert id(first_bstream) in subs
            # ^^TODO! the bcrx should always de-allocate all subs,
            # including the implicit first one allocated on entry
            # by the first subscribing peer task, no?
            #
            # -[ ] adjust `MsgStream.subscribe()` to do this mgmt!
            #  |_ allows reverting `MsgStream.receive()` to the
            #    non-bcaster method.
            #  |_ we can decide whether to reset `._broadcaster`?
            #
            # await tractor.pause(shield=True)


def test_open_local_sub_to_stream(
    debug_mode: bool,
):
    '''
    Verify a single inter-actor stream can can be fanned-out shared to
    N local tasks using `trionics.maybe_open_context()`.

    '''
    from .conftest import cpu_perf_headroom
    timeout: float = (
        4
        if not platform.system() == "Windows"
        else 10
    ) * cpu_perf_headroom()

    if debug_mode:
        timeout = 999
        print(f'IN debug_mode, setting large timeout={timeout!r}..')

    async def main():

        full = list(range(1000))
        an: tractor.ActorNursery|None = None
        num_tasks: int = 10

        async def get_sub_and_pull(taskname: str):

            nonlocal an

            stream: tractor.MsgStream
            async with (
                maybe_open_stream(taskname) as (
                    an,
                    stream,
                ),
            ):
                if '0' in taskname:
                    assert isinstance(stream, tractor.MsgStream)
                else:
                    assert isinstance(
                        stream,
                        tractor.trionics.BroadcastReceiver
                    )

                first = await stream.receive()
                print(f'{taskname} started with value {first}')
                seq: list[int] = []
                async for msg in stream:
                    seq.append(msg)

                assert set(seq).issubset(set(full))

            # end of @acm block
            print(f'{taskname} finished')

        root: tractor.Actor
        with trio.fail_after(timeout) as cs:
            # TODO: turns out this isn't multi-task entrant XD
            # We probably need an indepotent entry semantic?
            async with tractor.open_root_actor(
                debug_mode=debug_mode,
                # maybe_enable_greenback=True,
                #
                # ^TODO? doesn't seem to mk breakpoint() usage work
                # bc each bg task needs to open a portal??
                # - [ ] we should consider making this part of
                #      our taskman defaults?
                #   |_see https://github.com/goodboy/tractor/pull/363
                #
            ) as root:
                assert root.is_registrar

                async with (
                    trio.open_nursery() as tn,
                ):
                    for i in range(num_tasks):
                        tn.start_soon(
                            get_sub_and_pull,
                            f'task_{i}',
                        )
                        await trio.sleep(0.001)

                print('all consumer tasks finished!')

                # ?XXX, ensure actor-nursery is shutdown or we might
                # hang here due to a minor task deadlock/race-condition?
                #
                # - seems that all we need is a checkpoint to ensure
                #   the last suspended task, which is inside
                #   `.maybe_open_context()`, can do the
                #   `Portal.cancel_actor()` call?
                #
                # - if that bg task isn't resumed, then this blocks
                #   timeout might hit before that?
                #
                if root.ipc_server.has_peers():
                    await trio.lowlevel.checkpoint()

                    # alt approach, cancel the entire `an`
                    # await tractor.pause()
                    # await an.cancel()

            # end of runtime scope
            print('root actor terminated.')

        if cs.cancelled_caught:
            pytest.fail(
                'Should NOT time out in `open_root_actor()` ?'
            )

        print('exiting main.')

    trio.run(main)



@acm
async def maybe_cancel_outer_cs(
    cs: trio.CancelScope|None = None,
    delay: float = 0,
):
    # on first task delay this enough to block
    # the 2nd task but then cancel it mid sleep
    # so that the tn.start() inside the key-err handler block
    # is cancelled and would previously corrupt the
    # mutext state.
    log.info(f'task entering sleep({delay})')
    await trio.sleep(delay)
    if cs:
        log.info('task calling cs.cancel()')
        cs.cancel()

    yield

    if cs:
        await trio.sleep_forever()

    # XXX, if not cancelled we'll leak this inf-blocking
    # subtask to the actor's service tn..
    else:
        await trio.lowlevel.checkpoint()


@pytest.mark.parametrize(
    'delay',
    [0.05, 0.5, 1],
    ids="pre_sleep_delay={}".format,
)
@pytest.mark.parametrize(
    'cancel_by_cs',
    [True, False],
    ids="cancel_by_cs={}".format,
)
def test_lock_not_corrupted_on_fast_cancel(
    delay: float,
    cancel_by_cs: bool,
    debug_mode: bool,
    loglevel: str,
):
    '''
    Verify that if the caching-task (the first to enter
    `maybe_open_context()`) is cancelled mid-cache-miss, the embedded
    mutex can never be left in a corrupted state.

    That is, the lock is always eventually released ensuring a peer
    (cache-hitting) task will never,

    - be left to inf-block/hang on the `lock.acquire()`.
    - try to release the lock when still owned by the caching-task
      due to it having erronously exited without calling
      `lock.release()`.

    '''
    async def use_moc(
        delay: float,
        cs: trio.CancelScope|None = None,
    ):
        log.info('task entering moc')
        async with maybe_open_context(
            maybe_cancel_outer_cs,
            kwargs={
                'cs': cs,
                'delay': delay,
            },
        ) as (cache_hit, _null):
            if cache_hit:
                log.info('2nd task entered')
            else:
                log.info('1st task entered')

            if cs:
                await trio.sleep_forever()

            else:
                await trio.sleep(delay)

        # ^END, exit shared ctx.

    async def main():
        with trio.fail_after(delay + 2):
            async with (
                tractor.open_root_actor(
                    debug_mode=debug_mode,
                    loglevel=loglevel,
                ),
                # ?TODO, pass this as the parent tn?
                trio.open_nursery() as tn,
            ):
                get_console_log('info')
                log.info('yo starting')
                cs = tn.cancel_scope
                tn.start_soon(
                    use_moc,
                    delay,
                    cs if cancel_by_cs else None,
                    name='child',
                )
                with trio.CancelScope() as rent_cs:
                    await use_moc(
                        delay=delay,
                        cs=rent_cs if cancel_by_cs else None,
                    )

    trio.run(main)


@acm
async def acm_with_resource(resource_id: str):
    '''
    Yield `resource_id` as the cached value.

    Used to verify per-`ctx_key` isolation when the same
    `acm_func` is called with different kwargs.

    '''
    yield resource_id


def test_per_ctx_key_resource_lifecycle(
    debug_mode: bool,
    loglevel: str,
):
    '''
    Verify that `maybe_open_context()` correctly isolates resource
    lifecycle **per `ctx_key`** when the same `acm_func` is called
    with different kwargs.

    Previously `_Cache.users` was a single global `int` and
    `_Cache.locks` was keyed on `fid` (function ID), so calling
    the same `acm_func` with different kwargs (producing different
    `ctx_key`s) meant:

    - teardown for one key was skipped bc the *other* key's users
      kept the global count > 0,
    - and re-entry could hit the old
      `assert not resources.get(ctx_key)` crash during the
      teardown window.

    This was the root cause of a long-standing bug in piker's
    `brokerd.kraken` backend.

    '''
    timeout: float = 6
    if debug_mode:
        timeout = 999

    async def main():
        a_ready = trio.Event()
        a_exit = trio.Event()

        async def hold_resource_a():
            '''
            Open resource 'a' and keep it alive until signalled.

            '''
            async with maybe_open_context(
                acm_with_resource,
                kwargs={'resource_id': 'a'},
            ) as (cache_hit, value):
                assert not cache_hit
                assert value == 'a'
                log.info("resource 'a' entered (holding)")
                a_ready.set()
                await a_exit.wait()
                log.info("resource 'a' exiting")

        with trio.fail_after(timeout):
            async with (
                tractor.open_root_actor(
                    debug_mode=debug_mode,
                    loglevel=loglevel,
                ),
                trio.open_nursery() as tn,
            ):
                # Phase 1: bg task holds resource 'a' open.
                tn.start_soon(hold_resource_a)
                await a_ready.wait()

                # Phase 2: open resource 'b' (different kwargs,
                # same acm_func) then exit it while 'a' is still
                # alive.
                async with maybe_open_context(
                    acm_with_resource,
                    kwargs={'resource_id': 'b'},
                ) as (cache_hit, value):
                    assert not cache_hit
                    assert value == 'b'
                    log.info("resource 'b' entered")

                log.info("resource 'b' exited, waiting for teardown")
                await trio.lowlevel.checkpoint()

                # Phase 3: re-open 'b'; must be a fresh cache MISS
                # proving 'b' was torn down independently of 'a'.
                #
                # With the old global `_Cache.users` counter this
                # would be a stale cache HIT (leaked resource) or
                # trigger `assert not resources.get(ctx_key)`.
                async with maybe_open_context(
                    acm_with_resource,
                    kwargs={'resource_id': 'b'},
                ) as (cache_hit, value):
                    assert not cache_hit, (
                        "resource 'b' was NOT torn down despite "
                        "having zero users! (global user count bug)"
                    )
                    assert value == 'b'
                    log.info(
                        "resource 'b' re-entered "
                        "(cache miss, correct)"
                    )

                # Phase 4: let 'a' exit, clean shutdown.
                a_exit.set()

    trio.run(main)


def test_moc_reentry_during_teardown(
    debug_mode: bool,
    loglevel: str,
):
    '''
    Reproduce re-entry while an identical cached context exits.

    - multiple tasks use the same `acm_func` with no kwargs,
      producing an identical `ctx_key`;
    - all users leave and the final user starts resource teardown;
    - `_Cache.run_ctx()` removes the cached value and resource entry
      before entering the resource's blocking `__aexit__()` body;
    - a new task attempts to enter that same `ctx_key` during exit;
    - the per-key lock keeps that entrant queued until exit
      completes;
    - the entrant then receives a fresh cache miss and resource.

    Without teardown sharing the registration lock, re-entry could
    race resource replacement while the prior generation was still
    exiting. The final user could also return before that exit
    completed.

    The first resource generation signals `in_aexit` and waits on
    `allow_aexit`. The re-entry task signals `reentry_started` and
    blocks inside MOC; only after `wait_all_tasks_blocked()` confirms
    that ordering does the coordinator release cleanup. The entrant
    must then receive a fresh cache miss. `first_done` additionally
    proves the first MOC user observed completed teardown before
    returning.

    '''
    async def main():
        in_aexit = trio.Event()
        allow_aexit = trio.Event()
        reentry_started = trio.Event()
        generation: int = 0

        @acm
        async def cached_client():
            '''
            Simulate a no-argument `kraken.api.get_client()`.

            '''
            nonlocal generation

            generation += 1
            resource_generation: int = generation
            yield 'the-client'
            if resource_generation == 1:
                in_aexit.set()
                await allow_aexit.wait()

        first_done = trio.Event()

        async def use_and_exit():
            async with maybe_open_context(
                cached_client,
            ) as (cache_hit, value):
                assert value == 'the-client'
            first_done.set()

        async def reenter_during_teardown():
            '''
            Wait for the acm's `__aexit__` to start (meaning
            the cached value is no longer available), then re-enter.

            '''
            await in_aexit.wait()

            # Tell the coordinator this task is about to enter MOC.
            # `Event.set()` is not a checkpoint. Though `async with`
            # awaits MOC's `__aenter__()`, its async generator runs
            # synchronously until the held per-key `lock.acquire()`
            # actually suspends this task.
            reentry_started.set()
            async with maybe_open_context(
                cached_client,
            ) as (cache_hit, value):
                assert not cache_hit
                assert value == 'the-client'

            await first_done.wait()

        with trio.fail_after(5):
            async with (
                tractor.open_root_actor(
                    debug_mode=debug_mode,
                    loglevel=loglevel,
                ),
                collapse_eg(),
                trio.open_nursery() as tn,
            ):
                tn.start_soon(use_and_exit)
                tn.start_soon(reenter_during_teardown)
                await reentry_started.wait()

                # Wait until the re-entry task is queued on MOC's
                # per-key lock while `_Cache.run_ctx()` remains
                # blocked in the first generation's `__aexit__()`.
                # Only then release cleanup, making the intended
                # enter-during-sibling-exit ordering deterministic.
                await wait_all_tasks_blocked()
                assert not first_done.is_set()
                allow_aexit.set()

    trio.run(main)