Skip to content

Nexus cancellation types #981

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.

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
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions temporalio/worker/_interceptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,7 @@ class StartNexusOperationInput(Generic[InputT, OutputT]):
operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]]
input: InputT
schedule_to_close_timeout: Optional[timedelta]
cancellation_type: temporalio.workflow.NexusOperationCancellationType
headers: Optional[Mapping[str, str]]
output_type: Optional[Type[OutputT]] = None

Expand Down
26 changes: 14 additions & 12 deletions temporalio/worker/_workflow_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,13 @@
import temporalio.bridge.proto.activity_result
import temporalio.bridge.proto.child_workflow
import temporalio.bridge.proto.common
import temporalio.bridge.proto.nexus
import temporalio.bridge.proto.workflow_activation
import temporalio.bridge.proto.workflow_commands
import temporalio.bridge.proto.workflow_completion
import temporalio.common
import temporalio.converter
import temporalio.exceptions
import temporalio.nexus
import temporalio.workflow
from temporalio.service import __version__

Expand Down Expand Up @@ -1502,9 +1502,10 @@ async def workflow_start_nexus_operation(
service: str,
operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
input: Any,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
headers: Optional[Mapping[str, str]] = None,
output_type: Optional[Type[OutputT]],
schedule_to_close_timeout: Optional[timedelta],
cancellation_type: temporalio.workflow.NexusOperationCancellationType,
headers: Optional[Mapping[str, str]],
) -> temporalio.workflow.NexusOperationHandle[OutputT]:
# start_nexus_operation
return await self._outbound.start_nexus_operation(
Expand All @@ -1515,6 +1516,7 @@ async def workflow_start_nexus_operation(
input=input,
output_type=output_type,
schedule_to_close_timeout=schedule_to_close_timeout,
cancellation_type=cancellation_type,
headers=headers,
)
)
Expand Down Expand Up @@ -2757,7 +2759,7 @@ def _apply_schedule_command(
if self._input.retry_policy:
self._input.retry_policy.apply_to_proto(v.retry_policy)
v.cancellation_type = cast(
"temporalio.bridge.proto.workflow_commands.ActivityCancellationType.ValueType",
temporalio.bridge.proto.workflow_commands.ActivityCancellationType.ValueType,
int(self._input.cancellation_type),
)

Expand Down Expand Up @@ -2893,7 +2895,7 @@ def _apply_start_command(self) -> None:
if self._input.task_timeout:
v.workflow_task_timeout.FromTimedelta(self._input.task_timeout)
v.parent_close_policy = cast(
"temporalio.bridge.proto.child_workflow.ParentClosePolicy.ValueType",
temporalio.bridge.proto.child_workflow.ParentClosePolicy.ValueType,
int(self._input.parent_close_policy),
)
v.workflow_id_reuse_policy = cast(
Expand All @@ -2915,7 +2917,7 @@ def _apply_start_command(self) -> None:
self._input.search_attributes, v.search_attributes
)
v.cancellation_type = cast(
"temporalio.bridge.proto.child_workflow.ChildWorkflowCancellationType.ValueType",
temporalio.bridge.proto.child_workflow.ChildWorkflowCancellationType.ValueType,
int(self._input.cancellation_type),
)
if self._input.versioning_intent:
Expand Down Expand Up @@ -3011,11 +3013,6 @@ def __init__(

@property
def operation_token(self) -> Optional[str]:
# TODO(nexus-preview): How should this behave?
# Java has a separate class that only exists if the operation token exists:
# https://github.com/temporalio/sdk-java/blob/master/temporal-sdk/src/main/java/io/temporal/internal/sync/NexusOperationExecutionImpl.java#L26
# And Go similar:
# https://github.com/temporalio/sdk-go/blob/master/internal/workflow.go#L2770-L2771
try:
return self._start_fut.result()
except BaseException:
Expand Down Expand Up @@ -3064,6 +3061,11 @@ def _apply_schedule_command(self) -> None:
v.schedule_to_close_timeout.FromTimedelta(
self._input.schedule_to_close_timeout
)
v.cancellation_type = cast(
temporalio.bridge.proto.nexus.NexusOperationCancellationType.ValueType,
int(self._input.cancellation_type),
)

if self._input.headers:
for key, val in self._input.headers.items():
v.nexus_header[key] = val
Expand Down
63 changes: 60 additions & 3 deletions temporalio/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -858,9 +858,10 @@ async def workflow_start_nexus_operation(
service: str,
operation: Union[nexusrpc.Operation[InputT, OutputT], str, Callable[..., Any]],
input: Any,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
headers: Optional[Mapping[str, str]] = None,
output_type: Optional[Type[OutputT]],
schedule_to_close_timeout: Optional[timedelta],
cancellation_type: temporalio.workflow.NexusOperationCancellationType,
headers: Optional[Mapping[str, str]],
) -> NexusOperationHandle[OutputT]: ...

@abstractmethod
Expand Down Expand Up @@ -5137,6 +5138,46 @@ def _to_proto(self) -> temporalio.bridge.proto.common.VersioningIntent.ValueType
ServiceT = TypeVar("ServiceT")


class NexusOperationCancellationType(IntEnum):
"""Defines behavior of a Nexus operation when the caller workflow initiates cancellation.

Pass one of these values to :py:meth:`NexusClient.start_operation` to define cancellation
behavior.

To initiate cancellation, use :py:meth:`NexusOperationHandle.cancel` and then `await` the
operation handle. This will result in a :py:class:`exceptions.NexusOperationError`. The values
of this enum define what is guaranteed to have happened by that point.
"""

ABANDON = int(temporalio.bridge.proto.nexus.NexusOperationCancellationType.ABANDON)
"""Do not send any cancellation request to the operation handler; just report cancellation to the caller"""

TRY_CANCEL = int(
temporalio.bridge.proto.nexus.NexusOperationCancellationType.TRY_CANCEL
)
"""Send a cancellation request but immediately report cancellation to the caller. Note that this
does not guarantee that cancellation is delivered to the operation handler if the caller exits
before the delivery is done.
"""

# TODO(nexus-preview): core needs to be updated to handle
# NexusOperationCancelRequestCompleted and NexusOperationCancelRequestFailed
# see https://github.com/temporalio/sdk-core/issues/911
# WAIT_REQUESTED = int(
# temporalio.bridge.proto.nexus.NexusOperationCancellationType.WAIT_CANCELLATION_REQUESTED
# )
# """Send a cancellation request and wait for confirmation that the request was received.
# Does not wait for the operation to complete.
# """

WAIT_COMPLETED = int(
temporalio.bridge.proto.nexus.NexusOperationCancellationType.WAIT_CANCELLATION_COMPLETED
)
"""Send a cancellation request and wait for the operation to complete.
Note that the operation may not complete as cancelled (for example, if it catches the
:py:exc:`asyncio.CancelledError` resulting from the cancellation request)."""


class NexusClient(ABC, Generic[ServiceT]):
"""A client for invoking Nexus operations.

Expand Down Expand Up @@ -5167,6 +5208,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> NexusOperationHandle[OutputT]: ...

Expand All @@ -5180,6 +5222,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> NexusOperationHandle[OutputT]: ...

Expand All @@ -5196,6 +5239,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> NexusOperationHandle[OutputT]: ...

Expand All @@ -5212,6 +5256,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> NexusOperationHandle[OutputT]: ...

Expand All @@ -5228,6 +5273,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> NexusOperationHandle[OutputT]: ...

Expand All @@ -5239,6 +5285,7 @@ async def start_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> Any:
"""Start a Nexus operation and return its handle.
Expand Down Expand Up @@ -5268,6 +5315,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> OutputT: ...

Expand All @@ -5281,6 +5329,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> OutputT: ...

Expand All @@ -5297,6 +5346,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> OutputT: ...

Expand All @@ -5316,6 +5366,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> OutputT: ...

Expand All @@ -5332,6 +5383,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> OutputT: ...

Expand All @@ -5343,6 +5395,7 @@ async def execute_operation(
*,
output_type: Optional[Type[OutputT]] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> Any:
"""Execute a Nexus operation and return its result.
Expand Down Expand Up @@ -5394,6 +5447,7 @@ async def start_operation(
*,
output_type: Optional[Type] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> Any:
return (
Expand All @@ -5404,6 +5458,7 @@ async def start_operation(
input=input,
output_type=output_type,
schedule_to_close_timeout=schedule_to_close_timeout,
cancellation_type=cancellation_type,
headers=headers,
)
)
Expand All @@ -5415,13 +5470,15 @@ async def execute_operation(
*,
output_type: Optional[Type] = None,
schedule_to_close_timeout: Optional[timedelta] = None,
cancellation_type: NexusOperationCancellationType = NexusOperationCancellationType.WAIT_COMPLETED,
headers: Optional[Mapping[str, str]] = None,
) -> Any:
handle = await self.start_operation(
operation,
input,
output_type=output_type,
schedule_to_close_timeout=schedule_to_close_timeout,
cancellation_type=cancellation_type,
headers=headers,
)
return await handle
Expand Down
39 changes: 39 additions & 0 deletions tests/helpers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@
from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar

from temporalio.api.common.v1 import WorkflowExecution
from temporalio.api.enums.v1 import EventType as EventType
from temporalio.api.enums.v1 import IndexedValueType
from temporalio.api.history.v1 import HistoryEvent
from temporalio.api.operatorservice.v1 import (
AddSearchAttributesRequest,
ListSearchAttributesRequest,
Expand Down Expand Up @@ -287,3 +289,40 @@ async def check_unpaused() -> bool:
return not info.paused

await assert_eventually(check_unpaused)


async def print_history(handle: WorkflowHandle):
i = 1
async for evt in handle.fetch_history_events():
event = EventType.Name(evt.event_type).removeprefix("EVENT_TYPE_")
print(f"{i:2}: {event}")
i += 1


async def print_interleaved_histories(*handles: WorkflowHandle) -> None:
"""
Print the interleaved history events from multiple workflow handles in columns.
"""
all_events: list[tuple[WorkflowHandle, HistoryEvent, int]] = []
for handle in handles:
event_num = 1
async for event in handle.fetch_history_events():
all_events.append((handle, event, event_num))
event_num += 1
all_events.sort(key=lambda item: item[1].event_time.ToDatetime())
col_width = 40

def _format_row(items: list[str], truncate: bool = False) -> str:
if truncate:
items = [item[: col_width - 3] for item in items]
return " | ".join(f"{item:<{col_width - 3}}" for item in items)

headers = [handle.id for handle in handles]
print("\n" + _format_row(headers, truncate=True))
print("-" * (col_width * len(handles) + len(handles) - 1))
for handle, event, event_num in all_events:
event_type = EventType.Name(event.event_type).removeprefix("EVENT_TYPE_")
row = [""] * len(handles)
col_idx = handles.index(handle)
row[col_idx] = f"{event_num:2}: {event_type[: col_width - 5]}"
print(_format_row(row))
Loading
Loading