Isolated dependency context / local exit stacks

[thomas-touhey]

Hi,

In the context of an event consumer, I am defining an asynchronous, injectable function that is called every time an event is consumed; let's call it process_event. Taking full advantage of asyncio's capabilities, in order not to finish processing the message (which can take seconds) before starting to process another one, I am making a task of this function using TaskGroup.create_task. The structure of this program resembles the following:

async def process_event(...):  # Deps, message contents, etc
    ...

async def main():
    async with TaskGroup() as tg:
        async for event in event_queue:
            tg.create_task(process_event(...))  # Include event/message contents.

However, note that this does not include cleanup. Just like FastAPI routers, I would like to have cleanup for every event processing task, in an isolated manner so that cleanup for event A processing does not impact event B processing that is happening in parallel.

Currently, the options I have at my disposal with fastapi-injectable, from what I get, are the following:

• Calling cleanup_all_exit_stacks() after the task group is done. This is a memory leak, connection/socket leak, etc for long-running event consumers.
• Calling cleanup_exit_stack_of_func in an intermediate function that calls cleanup_exit_stack_of_func, e.g.:

async def process_event(...):  # Deps, message contents, etc
    ...

async def call_process_event(*args, **kwargs):
    await process_event(*args, **kwargs)
    await cleanup_exit_stack_of_func(process_event, raise_exception=True)

async def main():
    async with TaskGroup() as tg:
        async for event in event_queue:
            tg.create_task(call_process_event(...))  # Include event/message contents.

For some reason when trying this the cleanup part of my dependencies doesn't seem called, but even if they were, this does not pair well with parallel executions of the event processing function.

For my use case, what would be preferrable, rather than using global exit stacks, would be to instantiate and use a local exit stack. Something like:

async with AsyncExitStack() as exit_stack:
    await process_event(..., exit_stack=exit_stack)

# or something specific to fastapi-injectable:
async with fastapi_injectable.create_exit_stack() as exit_stack:
    await process_event(..., exit_stack=exit_stack)

Does something like this exist in fastapi-injectable, or is planned? Thanks :)

[sueun-dev]

I don't see a first-class per-invocation/local exit stack API in 1.4.7 rn.

The current stack is keyed by callable, not by each call. resolve_dependencies() gets the stack for root_dep.call, and cleanup_exit_stack_of_func(process_event) closes that function-level stack.

I tested this with two overlapping tasks and use_cache=False. The two dependency instances were separate, but when the fast task called cleanup, the slow task's dependency was already closed before the slow task finished. So yeah, I wouldn't call cleanup_exit_stack_of_func(process_event) from inside each parallel task.

Cleaning after the whole TaskGroup is safe for correctness, but it has the long-running worker leak problem you mentioned.

The workaround I could make behave correctly was using a fresh injectable wrapper per event, ideally around the raw body:

async def call_process_event(event):
    @injectable(use_cache=False)
    async def run_once(dep: Annotated[Dep, Depends(get_dep)]):
        await process_event_body(event, dep)

    try:
        await run_once()
    finally:
        await cleanup_exit_stack_of_func(run_once, raise_exception=True)

That gave me per-task cleanup in the test, but it's still a workaround. A real async with create_exit_stack(): ... style API would fit this use case much better.

[thomas-touhey]

Thanks for the answer! I will use that as a hotfix first version, but waiting for the maintainer(s?) for a long-term answer if they have one :)

[JasperSui]

Thanks @sueun-dev for the precise analysis and workaround, the default exit stack is keyed by the function, so concurrent calls of the same injectable share one stack, which is why per-event cleanup_exit_stack_of_func() races with parallel work (one event's cleanup tears down another's in-flight resources), and why cleaning only after the TaskGroup leaks for long-running consumers.

So I made a injectable_scope in v1.5.0, injectable_scope() gives each unit of work its own exit stack and cache, with contextvars keeping parallel tasks isolated — the same request-scoped model FastAPI uses internally. Cleanup runs when each scope's async with block exits, so there's no global accumulation. Your TaskGroup example maps directly:

from fastapi_injectable import async_get_injected_obj, injectable_scope


async def process_event(event):
    async with injectable_scope():
        dep = await async_get_injected_obj(get_dep)
        await handle(event, dep)
    # only THIS event's resources are cleaned up here — siblings are untouched


async def main():
    async with asyncio.TaskGroup() as tg:
        async for event in event_queue:
            tg.create_task(process_event(event))

If you'd rather own the stack yourself (closer to your create_exit_stack() sketch), InjectableScope is usable directly, and you can route individual resolutions into it with scope=:

from fastapi_injectable import InjectableScope, async_get_injected_obj

scope = InjectableScope()
async with scope:
    dep = await async_get_injected_obj(get_dep, scope=scope)
    scope.exit_stack.push_async_callback(my_cleanup)  # rides the same lifecycle
# dep and my_cleanup are torn down together when the block exits

One caveat: it's async-first, under the background_thread loop strategy contextvars don't propagate across threads, so stick to the async API (async_get_injected_obj / injectable_scope) for scoped resolution. For your pure-async TaskGroup setup that's exactly the right path.

Available in v1.5.0:

pip install "fastapi-injectable>=1.5.0"

See the new Isolated dependency scopes for parallel work section in the README for details. Give it a try and let me know if it fits your consumer, thanks @thomas-touhey!

[thomas-touhey]

Thank you so much @JasperSui! I'll try it out :)