Skip to content

fix(decisioning): supervise timed and canceled work #1000

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bokelley wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(decisioning): supervise timed and canceled work #1000

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 189 additions & 23 deletions src/adcp/decisioning/dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
TaskHandoffContext,
TaskRegistry,
)
from adcp.decisioning.time_budget import SyncExecutorAdmission, _bind_routed_sync_execution
from adcp.decisioning.types import (
AdcpError,
TaskHandoff,
Expand All @@ -85,6 +86,11 @@

logger = logging.getLogger(__name__)

# Strong references for synchronous adopter lifecycles that outlive a
# cancelled request. A Python thread cannot be cancelled; its completion hooks
# must still settle durable proposal/idempotency state.
_SUPERVISED_SYNC_LIFECYCLES: set[asyncio.Task[Any]] = set()

# ---------------------------------------------------------------------------
# Specialism enum — spec slugs known to the framework
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -1313,6 +1319,7 @@ async def _invoke_platform_method(
webhook_target: WebhookDeliveryTarget | None = None,
webhook_auto_emit: bool = True,
pre_handoff_reject: Callable[[], None] | None = None,
sync_admission: SyncExecutorAdmission | None = None,
) -> Any:
"""Invoke a platform method, projecting hybrid returns.

Expand Down Expand Up @@ -1383,12 +1390,16 @@ async def _invoke_platform_method(
off on a ``wholesale`` request is rejected cleanly instead of
leaking a task the buyer was told was rejected. Runs only on the
``TaskHandoff`` arm; sync / workflow-handoff returns ignore it.
:param sync_admission: Optional bounded admission controller for a sync
method. Its permit remains held until the underlying thread future
actually completes, including after caller cancellation.
"""
# pydantic is a required dep; import here (not at module level) to mirror
# the lazy-import discipline used throughout this module.
from pydantic import ValidationError as _ValidationError # noqa: PLC0415

method = getattr(platform, method_name)
sync_lifecycle_continues = False
# Re-validate through the platform method's own annotation when it's a
# stricter subclass of the shim's already-deserialized type. Skipped
# when arg_projector is set — that path replaces positional args entirely.
Expand All @@ -1407,31 +1418,74 @@ async def _invoke_platform_method(

try:
if asyncio.iscoroutinefunction(method):
if arg_projector is not None:
result = await method(**arg_projector, ctx=ctx)
elif extra_kwargs:
result = await method(params, ctx, **extra_kwargs)
else:
result = await method(params, ctx)
# Async router delegates may resolve to synchronous tenant
# children only after account routing. Propagate the same bounded
# admission controller and configured executor through ContextVars
# so that path cannot bypass the timed-sync limit.
with _bind_routed_sync_execution(sync_admission, executor):
if arg_projector is not None:
result = await method(**arg_projector, ctx=ctx)
elif extra_kwargs:
result = await method(params, ctx, **extra_kwargs)
else:
result = await method(params, ctx)
else:
ctx_snapshot = contextvars.copy_context()
loop = asyncio.get_running_loop()
if arg_projector is not None:
projected_kwargs = {**arg_projector, "ctx": ctx}
result = await loop.run_in_executor(
executor,
functools.partial(ctx_snapshot.run, method, **projected_kwargs),
)
worker_call = functools.partial(ctx_snapshot.run, method, **projected_kwargs)
elif extra_kwargs:
result = await loop.run_in_executor(
executor,
functools.partial(ctx_snapshot.run, method, params, ctx, **extra_kwargs),
worker_call = functools.partial(
ctx_snapshot.run, method, params, ctx, **extra_kwargs
)
else:
result = await loop.run_in_executor(
executor,
functools.partial(ctx_snapshot.run, method, params, ctx),
)
worker_call = functools.partial(ctx_snapshot.run, method, params, ctx)

if sync_admission is not None:
await sync_admission.acquire()
try:
worker_future = executor.submit(worker_call)
except BaseException:
if sync_admission is not None:
sync_admission.release()
raise

if sync_admission is not None:

def _release_admission(_future: object) -> None:
try:
loop.call_soon_threadsafe(sync_admission.release)
except RuntimeError:
# Event loop already closed during process teardown.
pass

worker_future.add_done_callback(_release_admission)

worker_async_future = asyncio.wrap_future(worker_future, loop=loop)
try:
result = await asyncio.shield(worker_async_future)
except asyncio.CancelledError:
if on_failure is not None:
sync_lifecycle_continues = True
lifecycle = asyncio.create_task(
_settle_cancelled_sync_lifecycle(
worker_async_future,
ctx=ctx,
method_name=method_name,
registry=registry,
executor=executor,
on_complete=on_complete,
on_failure=on_failure,
pre_handoff_reject=pre_handoff_reject,
request_params=params,
webhook_target=webhook_target,
webhook_auto_emit=webhook_auto_emit,
)
)
_SUPERVISED_SYNC_LIFECYCLES.add(lifecycle)
lifecycle.add_done_callback(_SUPERVISED_SYNC_LIFECYCLES.discard)
raise
except AdcpError as exc:
# Adopter raised structured error — propagate verbatim. The
# outer middleware projects to the wire envelope. Fire
Expand Down Expand Up @@ -1543,15 +1597,77 @@ async def _invoke_platform_method(
if on_failure is not None:
await _safe_on_failure_call(on_failure, wrapped, method_name)
raise wrapped from exc
except BaseException as exc:
# ``asyncio.CancelledError`` (and shutdown BaseExceptions) bypass the
# wire-error wrapping above, but must still release framework state
# reserved before adapter dispatch. Preserve the exact exception.
nested_sync_future = getattr(exc, "_adcp_sync_worker_future", None)
if isinstance(nested_sync_future, asyncio.Future) and on_failure is not None:
sync_lifecycle_continues = True
lifecycle = asyncio.create_task(
_settle_cancelled_sync_lifecycle(
nested_sync_future,
ctx=ctx,
method_name=method_name,
registry=registry,
executor=executor,
on_complete=on_complete,
on_failure=on_failure,
pre_handoff_reject=pre_handoff_reject,
request_params=params,
webhook_target=webhook_target,
webhook_auto_emit=webhook_auto_emit,
)
)
_SUPERVISED_SYNC_LIFECYCLES.add(lifecycle)
lifecycle.add_done_callback(_SUPERVISED_SYNC_LIFECYCLES.discard)
if on_failure is not None and not sync_lifecycle_continues:
await _safe_on_failure_call(on_failure, exc, method_name)
raise

return await _project_invocation_result(
result,
ctx=ctx,
method_name=method_name,
registry=registry,
executor=executor,
on_complete=on_complete,
on_failure=on_failure,
pre_handoff_reject=pre_handoff_reject,
request_params=params,
webhook_target=webhook_target,
webhook_auto_emit=webhook_auto_emit,
)


async def _project_invocation_result(
result: Any,
*,
ctx: RequestContext[Any],
method_name: str,
registry: TaskRegistry,
executor: ThreadPoolExecutor,
on_complete: Callable[[Any], Awaitable[None]] | None,
on_failure: Callable[[BaseException], Awaitable[None]] | None,
pre_handoff_reject: Callable[[], None] | None,
request_params: BaseModel,
webhook_target: WebhookDeliveryTarget | None,
webhook_auto_emit: bool,
) -> Any:
"""Project a raw adopter result and settle its framework lifecycle hooks."""
if is_task_handoff(result):
# Reject before any side effect (registry row, background task,
# completion webhook) is created. The wholesale discovery guard
# uses this so an adopter handing off on a synchronous-only
# wholesale request never leaks a task the buyer is told is
# rejected.
if pre_handoff_reject is not None:
pre_handoff_reject()
try:
pre_handoff_reject()
except BaseException as exc:
if on_failure is not None:
await _safe_on_failure_call(on_failure, exc, method_name)
raise
return await _project_handoff(
result,
ctx,
Expand All @@ -1560,7 +1676,7 @@ async def _invoke_platform_method(
executor=executor,
on_complete=on_complete,
on_failure=on_failure,
request_params=params,
request_params=request_params,
webhook_target=webhook_target,
webhook_auto_emit=webhook_auto_emit,
)
Expand All @@ -1571,7 +1687,7 @@ async def _invoke_platform_method(
method_name=method_name,
registry=registry,
executor=executor,
request_params=params,
request_params=request_params,
)

# Sync return path. Fire on_complete with the typed result before
Expand All @@ -1596,6 +1712,51 @@ async def _invoke_platform_method(
return strip_credentials_from_wire_result(method_name, result)


async def _settle_cancelled_sync_lifecycle(
worker_future: asyncio.Future[Any],
*,
ctx: RequestContext[Any],
method_name: str,
registry: TaskRegistry,
executor: ThreadPoolExecutor,
on_complete: Callable[[Any], Awaitable[None]] | None,
on_failure: Callable[[BaseException], Awaitable[None]] | None,
pre_handoff_reject: Callable[[], None] | None,
request_params: BaseModel,
webhook_target: WebhookDeliveryTarget | None,
webhook_auto_emit: bool,
) -> None:
"""Settle a sync worker after its request task has been cancelled."""
try:
result = await worker_future
except BaseException as exc:
if on_failure is not None:
await _safe_on_failure_call(on_failure, exc, method_name)
return
try:
await _project_invocation_result(
result,
ctx=ctx,
method_name=method_name,
registry=registry,
executor=executor,
on_complete=on_complete,
on_failure=on_failure,
pre_handoff_reject=pre_handoff_reject,
request_params=request_params,
webhook_target=webhook_target,
webhook_auto_emit=webhook_auto_emit,
)
except BaseException:
# Lifecycle hooks already apply their own rollback semantics. There is
# no request waiter left to receive this exception, so retain it in
# server logs rather than producing an unhandled-task warning.
logger.exception(
"Cancelled request's synchronous %s lifecycle failed while settling",
method_name,
)


async def _safe_on_failure_call(
on_failure: Callable[[BaseException], Awaitable[None]],
exc: BaseException,
Expand All @@ -1610,10 +1771,9 @@ async def _safe_on_failure_call(
"""
try:
await on_failure(exc)
except Exception:
except BaseException:
logger.exception(
"on_failure hook raised while handling %s for %s — original "
"exception still propagates",
"on_failure hook raised while handling %s for %s — original exception still propagates",
type(exc).__name__,
method_name,
)
Expand Down Expand Up @@ -1809,6 +1969,12 @@ async def _run() -> None:
)
await _fail(wrapped)
return
except BaseException as exc:
# Background task cancellation must release any proposal
# reservation, then remain cancellation to the task scheduler.
if on_failure is not None:
await _safe_on_failure_call(on_failure, exc, method_name)
raise

# Framework completion hook (e.g., proposal_store.commit for
# finalize, mark_proposal_consumed for create_media_buy). Runs
Expand Down
Loading
Loading