workflow_runtime.py

# -*- coding: utf-8 -*-

"""
Copyright 2023 The Dapr Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import inspect
import time
from functools import wraps
from typing import Any, Awaitable, Callable, Optional, Sequence, TypeVar, Union

import grpc

from dapr.clients import DaprInternalError
from dapr.clients.grpc.interceptors import DaprClientTimeoutInterceptor
from dapr.clients.http.client import DAPR_API_TOKEN_HEADER
from dapr.conf import settings
from dapr.conf.helpers import GrpcEndpoint
from dapr.ext.workflow._durabletask import task, worker
from dapr.ext.workflow._durabletask.internal.shared import is_async_callable as _is_async_callable
from dapr.ext.workflow.dapr_workflow_context import DaprWorkflowContext
from dapr.ext.workflow.logger import Logger, LoggerOptions
from dapr.ext.workflow.util import getAddress
from dapr.ext.workflow.workflow_activity_context import Activity, WorkflowActivityContext
from dapr.ext.workflow.workflow_context import Workflow

from . import _model_protocol

T = TypeVar('T')
TInput = TypeVar('TInput')
TOutput = TypeVar('TOutput')

ClientInterceptor = Union[
    grpc.UnaryUnaryClientInterceptor,
    grpc.UnaryStreamClientInterceptor,
    grpc.StreamUnaryClientInterceptor,
    grpc.StreamStreamClientInterceptor,
]

# Durabletask returns decoded JSON, so we type the input as ``object | None`` and let the
# wrapper narrow it via the activity's declared model.
SyncActivityWrapper = Callable[[task.ActivityContext, object | None], object]
AsyncActivityWrapper = Callable[[task.ActivityContext, object | None], Awaitable[object]]
ActivityWrapper = SyncActivityWrapper | AsyncActivityWrapper


def _coerce_activity_input(inp: object | None, input_model: type | None) -> object | None:
    """Coerce the raw input to the activity's declared model, if it has one."""
    if inp is None or input_model is None or isinstance(inp, input_model):
        return inp
    return _model_protocol.coerce_to_model(inp, input_model)


def _make_activity_wrapper(fn: Activity, logger: Logger) -> ActivityWrapper:
    """Wrap a user activity for the durabletask worker.

    Returns:
        An ``async def`` wrapper for async activities, a plain ``def`` for sync.
    """
    accepts_input, input_model = _model_protocol.resolve_input(fn)

    def _call_args(ctx: task.ActivityContext, inp: object | None) -> tuple:
        wf_ctx = WorkflowActivityContext(ctx)
        if not accepts_input:
            return (wf_ctx,)
        return (wf_ctx, _coerce_activity_input(inp, input_model))

    def _log_failure(ctx: task.ActivityContext, exc: Exception) -> None:
        activity_id = getattr(ctx, 'task_id', 'unknown')
        logger.warning(f'Activity execution failed - task_id: {activity_id}, error: {exc}')

    is_async = _is_async_callable(fn)
    activity_name = getattr(fn, '__name__', repr(fn))
    kind = 'async' if is_async else 'sync'
    logger.debug(f"Registering activity '{activity_name}' on the {kind} dispatch path.")
    if is_async:

        async def async_activity_wrapper(
            ctx: task.ActivityContext, inp: object | None = None
        ) -> object:
            try:
                return await fn(*_call_args(ctx, inp))
            except Exception as exc:
                _log_failure(ctx, exc)
                raise

        return async_activity_wrapper

    def sync_activity_wrapper(ctx: task.ActivityContext, inp: object | None = None) -> object:
        try:
            return fn(*_call_args(ctx, inp))
        except Exception as exc:
            _log_failure(ctx, exc)
            raise

    return sync_activity_wrapper


class WorkflowRuntime:
    """WorkflowRuntime is the entry point for registering workflows and activities."""

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[str] = None,
        logger_options: Optional[LoggerOptions] = None,
        interceptors: Optional[Sequence[ClientInterceptor]] = None,
        maximum_concurrent_activity_work_items: Optional[int] = None,
        maximum_concurrent_orchestration_work_items: Optional[int] = None,
        maximum_thread_pool_workers: Optional[int] = None,
        worker_ready_timeout: Optional[float] = None,
    ):
        self._logger = Logger('WorkflowRuntime', logger_options)
        self._worker_ready_timeout = 30.0 if worker_ready_timeout is None else worker_ready_timeout

        metadata = tuple()
        if settings.DAPR_API_TOKEN:
            metadata = ((DAPR_API_TOKEN_HEADER, settings.DAPR_API_TOKEN),)
        address = getAddress(host, port)

        try:
            uri = GrpcEndpoint(address)
        except ValueError as error:
            raise DaprInternalError(f'{error}') from error

        options = self._logger.get_options()
        all_interceptors = []
        if interceptors:
            all_interceptors.extend(interceptors)
        all_interceptors.append(DaprClientTimeoutInterceptor())
        self.__worker = worker.TaskHubGrpcWorker(
            host_address=uri.endpoint,
            metadata=metadata,
            secure_channel=uri.tls,
            log_handler=options.log_handler,
            log_formatter=options.log_formatter,
            interceptors=all_interceptors,
            concurrency_options=worker.ConcurrencyOptions(
                maximum_concurrent_activity_work_items=maximum_concurrent_activity_work_items,
                maximum_concurrent_orchestration_work_items=maximum_concurrent_orchestration_work_items,
                maximum_thread_pool_workers=maximum_thread_pool_workers,
            ),
        )

    def register_workflow(self, fn: Workflow, *, name: Optional[str] = None):
        effective_name = name or fn.__name__
        self._logger.info(f"Registering workflow '{effective_name}' with runtime")

        accepts_input, input_model = _model_protocol.resolve_input(fn)

        def orchestrationWrapper(ctx: task.OrchestrationContext, inp: Optional[TInput] = None):
            """Responsible to call Workflow function in orchestrationWrapper"""
            instance_id = getattr(ctx, 'instance_id', 'unknown')

            try:
                daprWfContext = DaprWorkflowContext(ctx, self._logger.get_options())
                if not accepts_input:
                    result = fn(daprWfContext)
                else:
                    if (
                        (inp is not None)
                        and (input_model is not None)
                        and not isinstance(inp, input_model)
                    ):
                        inp = _model_protocol.coerce_to_model(inp, input_model)
                    result = fn(daprWfContext, inp)
                return result
            except Exception as e:
                self._logger.exception(
                    f'Workflow execution failed - instance_id: {instance_id}, error: {e}'
                )
                raise

        if hasattr(fn, '_workflow_registered'):
            # whenever a workflow is registered, it has a _dapr_alternate_name attribute
            alt_name = fn.__dict__['_dapr_alternate_name']
            raise ValueError(f'Workflow {fn.__name__} already registered as {alt_name}')
        if hasattr(fn, '_dapr_alternate_name'):
            alt_name = fn._dapr_alternate_name
            if name is not None:
                m = f'Workflow {fn.__name__} already has an alternate name {alt_name}'
                raise ValueError(m)
        else:
            fn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__

        self.__worker._registry.add_named_orchestrator(
            fn.__dict__['_dapr_alternate_name'], orchestrationWrapper
        )
        fn.__dict__['_workflow_registered'] = True

    def register_versioned_workflow(
        self, fn: Workflow, *, name: str, version_name: Optional[str] = None, is_latest: bool
    ):
        effective_name = name or fn.__name__
        self._logger.info(
            f"Registering version {version_name} of workflow '{effective_name}' with runtime"
        )

        accepts_input, input_model = _model_protocol.resolve_input(fn)

        def orchestrationWrapper(ctx: task.OrchestrationContext, inp: Optional[TInput] = None):
            """Responsible to call Workflow function in orchestrationWrapper"""
            daprWfContext = DaprWorkflowContext(ctx, self._logger.get_options())
            if not accepts_input:
                return fn(daprWfContext)
            if (inp is not None) and (input_model is not None) and not isinstance(inp, input_model):
                inp = _model_protocol.coerce_to_model(inp, input_model)
            return fn(daprWfContext, inp)

        if hasattr(fn, '_workflow_registered'):
            # whenever a workflow is registered, it has a _dapr_alternate_name attribute
            alt_name = fn.__dict__['_dapr_alternate_name']
            raise ValueError(f'Workflow {fn.__name__} already registered as {alt_name}')
        if hasattr(fn, '_dapr_alternate_name'):
            alt_name = fn._dapr_alternate_name
            if name is not None:
                m = f'Workflow {fn.__name__} already has an alternate name {alt_name}'
                raise ValueError(m)
        else:
            fn.__dict__['_dapr_alternate_name'] = name

        actual_version_name = version_name if version_name is not None else fn.__name__

        self.__worker._registry.add_named_orchestrator(
            name,
            orchestrationWrapper,
            version_name=actual_version_name,
            is_latest=is_latest,
        )
        fn.__dict__['_workflow_registered'] = True

    def register_activity(self, fn: Activity, *, name: Optional[str] = None):
        """Register a workflow activity. ``def`` and ``async def`` are both supported.
        Async activities run on the worker's event loop. Sync activities run in the
        thread pool sized by ``maximum_thread_pool_workers``.
        """
        effective_name = name or fn.__name__
        self._logger.info(f"Registering activity '{effective_name}' with runtime")

        activity_wrapper = _make_activity_wrapper(fn, self._logger)

        if hasattr(fn, '_activity_registered'):
            # whenever an activity is registered, it has a _dapr_alternate_name attribute
            alt_name = fn.__dict__['_dapr_alternate_name']
            raise ValueError(f'Activity {fn.__name__} already registered as {alt_name}')
        if hasattr(fn, '_dapr_alternate_name'):
            alt_name = fn._dapr_alternate_name
            if name is not None:
                m = f'Activity {fn.__name__} already has an alternate name {alt_name}'
                raise ValueError(m)
        else:
            fn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__

        self.__worker._registry.add_named_activity(
            fn.__dict__['_dapr_alternate_name'], activity_wrapper
        )
        fn.__dict__['_activity_registered'] = True

    def wait_for_worker_ready(self, timeout: float = 30.0) -> bool:
        """
        Wait for the worker's gRPC stream to become ready to receive work items.
        This method polls the worker's is_worker_ready() method until it returns True
        or the timeout is reached.

        Args:
            timeout: Maximum time in seconds to wait for the worker to be ready.
                    Defaults to 30 seconds.

        Returns:
            True if the worker's gRPC stream is ready to receive work items, False if timeout.
        """
        if not hasattr(self.__worker, 'is_worker_ready'):
            return False

        elapsed = 0.0
        poll_interval = 0.1  # 100ms

        while elapsed < timeout:
            if self.__worker.is_worker_ready():
                return True
            time.sleep(poll_interval)
            elapsed += poll_interval

        self._logger.warning(
            f'WorkflowRuntime worker readiness check timed out after {timeout} seconds'
        )
        return False

    def start(self):
        """Starts the listening for work items on a background thread.
        This method waits for the worker's gRPC stream to be fully initialized
        before returning, ensuring that workflows can be scheduled immediately
        after start() completes.
        """
        try:
            try:
                self.__worker.start()
            except Exception as start_error:
                self._logger.exception(f'WorkflowRuntime worker did not start: {start_error}')
                raise

            # Verify the worker and its stream reader are ready
            if hasattr(self.__worker, 'is_worker_ready'):
                try:
                    is_ready = self.wait_for_worker_ready(timeout=self._worker_ready_timeout)
                    if not is_ready:
                        raise RuntimeError('WorkflowRuntime worker and its stream are not ready')
                    else:
                        self._logger.debug(
                            'WorkflowRuntime worker is ready and its stream can receive work items'
                        )
                except Exception as ready_error:
                    self._logger.exception(
                        f'WorkflowRuntime wait_for_worker_ready() raised exception: {ready_error}'
                    )
                    raise ready_error
            else:
                self._logger.warning(
                    'Unable to verify stream readiness. Workflows scheduled immediately may not be received.'
                )
        except Exception:
            raise

    def shutdown(self):
        """Stops the listening for work items on a background thread."""
        try:
            self.__worker.stop()
        except Exception:
            raise

    def versioned_workflow(
        self,
        __fn: Workflow = None,
        *,
        name: str,
        version_name: Optional[str] = None,
        is_latest: bool,
    ):
        def wrapper(fn: Workflow):
            self.register_versioned_workflow(
                fn, name=name, version_name=version_name, is_latest=is_latest
            )

            @wraps(fn)
            def innerfn():
                return fn

            if hasattr(fn, '_dapr_alternate_name'):
                innerfn.__dict__['_dapr_alternate_name'] = fn.__dict__['_dapr_alternate_name']
            else:
                innerfn.__dict__['_dapr_alternate_name'] = name

            innerfn.__signature__ = inspect.signature(fn)
            return innerfn

        if __fn:
            # This case is true when the decorator is used without arguments
            # and the function to be decorated is passed as the first argument.
            return wrapper(__fn)

        return wrapper

    def workflow(self, __fn: Workflow = None, *, name: Optional[str] = None):
        """Decorator to register a workflow function.

        This example shows how to register a workflow function with a name:

                from dapr.ext.workflow import WorkflowRuntime
                wfr = WorkflowRuntime()

                @wfr.workflow(name="add")
                def add(ctx, x: int, y: int) -> int:
                    return x + y

        This example shows how to register a workflow function without
        an alternate name:

                    from dapr.ext.workflow import WorkflowRuntime
                    wfr = WorkflowRuntime()

                    @wfr.workflow
                    def add(ctx, x: int, y: int) -> int:
                        return x + y

        Args:
            name (Optional[str], optional): Name to identify the workflow function as in
            the workflow runtime. Defaults to None.
        """

        def wrapper(fn: Workflow):
            self.register_workflow(fn, name=name)

            @wraps(fn)
            def innerfn():
                return fn

            if hasattr(fn, '_dapr_alternate_name'):
                innerfn.__dict__['_dapr_alternate_name'] = fn.__dict__['_dapr_alternate_name']
            else:
                innerfn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__
            innerfn.__signature__ = inspect.signature(fn)
            return innerfn

        if __fn:
            # This case is true when the decorator is used without arguments
            # and the function to be decorated is passed as the first argument.
            return wrapper(__fn)

        return wrapper

    def activity(self, __fn: Activity = None, *, name: Optional[str] = None):
        """Decorator to register an activity function.

        This example shows how to register an activity function with an alternate name:

            from dapr.ext.workflow import WorkflowRuntime
            wfr = WorkflowRuntime()

            @wfr.activity(name="add")
            def add(ctx, x: int, y: int) -> int:
                return x + y

        This example shows how to register an activity function without an alternate name:

                from dapr.ext.workflow import WorkflowRuntime
                wfr = WorkflowRuntime()

                @wfr.activity
                def add(ctx, x: int, y: int) -> int:
                    return x + y

        Args:
            name (Optional[str], optional): Name to identify the activity function as in
            the workflow runtime. Defaults to None.
        """

        def wrapper(fn: Activity):
            self.register_activity(fn, name=name)

            @wraps(fn)
            def innerfn():
                return fn

            if hasattr(fn, '_dapr_alternate_name'):
                innerfn.__dict__['_dapr_alternate_name'] = fn.__dict__['_dapr_alternate_name']
            else:
                innerfn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__
            innerfn.__signature__ = inspect.signature(fn)
            return innerfn

        if __fn:
            # This case is true when the decorator is used without arguments
            # and the function to be decorated is passed as the first argument.
            return wrapper(__fn)

        return wrapper


def alternate_name(name: Optional[str] = None):
    """Decorator to register a workflow or activity function with an alternate name.

    This example shows how to register a workflow function with an alternate name:

            from dapr.ext.workflow import WorkflowRuntime
            wfr = WorkflowRuntime()

            @wfr.workflow
            @alternate_name(add")
            def add(ctx, x: int, y: int) -> int:
                return x + y

    Args:
        name (Optional[str], optional): Name to identify the workflow or activity function as in
        the workflow runtime. Defaults to None.
    """

    def wrapper(fn: Any):
        if hasattr(fn, '_dapr_alternate_name'):
            raise ValueError(
                f'Function {fn.__name__} already has an alternate name {fn._dapr_alternate_name}'
            )
        fn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__

        if _is_async_callable(fn):

            @wraps(fn)
            async def innerfn(*args, **kwargs):
                return await fn(*args, **kwargs)
        else:

            @wraps(fn)
            def innerfn(*args, **kwargs):
                return fn(*args, **kwargs)

        innerfn.__dict__['_dapr_alternate_name'] = name if name else fn.__name__
        innerfn.__signature__ = inspect.signature(fn)
        return innerfn

    return wrapper