IWF-466 add state_execution_locals and record_events (#87)

* IWF-466 add state_execution_locals and record_events

* IWF-466 formatting

* IWF-466 remove debug printing

* IWF-466 improve readability

---------

Co-authored-by: Nathan Givens <[email protected]>

Author: N-give

iwf/persistence.py

@@ -1,20 +1,24 @@
-from typing import Any, Union
+from typing import Any, Tuple, Union
 
 from iwf.data_attributes import DataAttributes
 from iwf.search_attributes import SearchAttributes
+from iwf.state_execution_locals import StateExecutionLocals
 
 
 class Persistence:
     _data_attributes: DataAttributes
     _search_attributes: SearchAttributes
+    _state_execution_locals: StateExecutionLocals
 
     def __init__(
         self,
         data_attributes: DataAttributes,
         search_attributes: SearchAttributes,
+        state_execution_locals: StateExecutionLocals,
     ):
         self._data_attributes = data_attributes
         self._search_attributes = search_attributes
+        self._state_execution_locals = state_execution_locals
 
     def get_data_attribute(self, key: str) -> Any:
         return self._data_attributes.get_data_attribute(key)
@@ -65,3 +69,12 @@ def set_search_attribute_keyword_array(
         self, key: str, value: Union[None, list[str]]
     ):
         self._search_attributes.set_search_attribute_keyword_array(key, value)
+
+    def get_state_execution_local(self, key: str) -> Any:
+        return self._state_execution_locals.get_state_execution_local(key)
+
+    def set_state_execution_local(self, key: str, value: Any):
+        self._state_execution_locals.set_state_execution_local(key, value)
+
+    def record_event(self, key: str, *event_data: Tuple[Any, ...]):
+        self._state_execution_locals.record_event(key, *event_data)

iwf/state_execution_locals.py

@@ -0,0 +1,56 @@
+from typing import Any, List, Tuple
+
+from iwf.errors import WorkflowDefinitionError
+from iwf.iwf_api.models import EncodedObject, KeyValue
+from iwf.object_encoder import ObjectEncoder
+
+
+class StateExecutionLocals:
+    _record_events: dict[str, EncodedObject]
+    _attribute_name_to_encoded_object_map: dict[str, EncodedObject]
+    _upsert_attributes_to_return_to_server: dict[str, EncodedObject]
+    _object_encoder: ObjectEncoder
+
+    def __init__(
+        self,
+        attribute_name_to_encoded_object_map: dict[str, EncodedObject],
+        object_encoder: ObjectEncoder,
+    ):
+        self._object_encoder = object_encoder
+        self._attribute_name_to_encoded_object_map = (
+            attribute_name_to_encoded_object_map
+        )
+        self._upsert_attributes_to_return_to_server = {}
+        self._record_events = {}
+
+    def set_state_execution_local(self, key: str, value: Any):
+        encoded_data = self._object_encoder.encode(value)
+        self._attribute_name_to_encoded_object_map[key] = encoded_data
+        self._upsert_attributes_to_return_to_server[key] = encoded_data
+
+    def get_state_execution_local(self, key: str) -> Any:
+        encoded_object = self._attribute_name_to_encoded_object_map.get(key)
+        if encoded_object is None:
+            return None
+        return self._object_encoder.decode(encoded_object)
+
+    def record_event(self, key: str, *event_data: Tuple[Any, ...]):
+        if key in self._record_events:
+            raise WorkflowDefinitionError("Cannot record the same event more than once")
+
+        if event_data is not None and len(event_data) == 1:
+            self._record_events[key] = self._object_encoder.encode(event_data[0])
+
+        self._record_events[key] = self._object_encoder.encode(event_data)
+
+    def get_upsert_state_execution_local_attributes(self) -> List[KeyValue]:
+        return [
+            KeyValue(item_key, item_value)
+            for item_key, item_value in self._upsert_attributes_to_return_to_server.items()
+        ]
+
+    def get_record_events(self) -> List[KeyValue]:
+        return [
+            KeyValue(item_key, item_value)
+            for item_key, item_value in self._record_events.items()
+        ]

iwf/tests/test_persistence_state_execution_locals.py

@@ -0,0 +1,82 @@
+import inspect
+import time
+import unittest
+
+from iwf.rpc import rpc
+from iwf.client import Client
+from iwf.command_request import CommandRequest
+from iwf.command_results import CommandResults
+from iwf.communication import Communication
+from iwf.persistence import Persistence
+from iwf.persistence_schema import PersistenceField, PersistenceSchema
+from iwf.state_decision import StateDecision
+from iwf.state_schema import StateSchema
+from iwf.tests.worker_server import registry
+from iwf.workflow import ObjectWorkflow
+from iwf.workflow_context import WorkflowContext
+from iwf.workflow_options import WorkflowOptions
+from iwf.workflow_state import T, WorkflowState
+
+PERSISTENCE_LOCAL_KEY = "persistence-test-key"
+PERSISTENCE_LOCAL_VALUE = "persistence-test-value"
+PERSISTENCE_DATA_ATTRIBUTE_KEY = "persistence-data-attribute-key"
+
+
+class PersistenceStateExecutionLocalRWState(WorkflowState[None]):
+    def wait_until(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        persistence: Persistence,
+        communication: Communication,
+    ):
+        persistence.set_state_execution_local(
+            PERSISTENCE_LOCAL_KEY, PERSISTENCE_LOCAL_VALUE
+        )
+        return CommandRequest.empty()
+
+    def execute(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        command_results: CommandResults,
+        persistence: Persistence,
+        communication: Communication,
+    ):
+        value = persistence.get_state_execution_local(PERSISTENCE_LOCAL_KEY)
+        persistence.set_data_attribute(PERSISTENCE_DATA_ATTRIBUTE_KEY, value)
+        return StateDecision.graceful_complete_workflow()
+
+
+class PersistenceStateExecutionLocalWorkflow(ObjectWorkflow):
+    def get_workflow_states(self) -> StateSchema:
+        return StateSchema.with_starting_state(PersistenceStateExecutionLocalRWState())
+
+    def get_persistence_schema(self) -> PersistenceSchema:
+        return PersistenceSchema.create(
+            PersistenceField.data_attribute_def(PERSISTENCE_DATA_ATTRIBUTE_KEY, str)
+        )
+
+    @rpc()
+    def test_persistence_read(self, persistence: Persistence):
+        return persistence.get_data_attribute(PERSISTENCE_DATA_ATTRIBUTE_KEY)
+
+
+class TestPersistenceExecutionLocalRead(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        wf = PersistenceStateExecutionLocalWorkflow()
+        registry.add_workflow(wf)
+        cls.client = Client(registry)
+
+    def test_persistence_execution_local_workflow(self):
+        wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
+        start_options = WorkflowOptions()
+        self.client.start_workflow(
+            PersistenceStateExecutionLocalWorkflow, wf_id, 200, None, start_options
+        )
+        self.client.wait_for_workflow_completion(wf_id, None)
+        res = self.client.invoke_rpc(
+            wf_id, PersistenceStateExecutionLocalWorkflow.test_persistence_read
+        )
+        assert res == PERSISTENCE_LOCAL_VALUE

iwf/worker_service.py

@@ -1,7 +1,7 @@
 import traceback
 import typing
 from dataclasses import dataclass
-from typing import Union
+from typing import List, Union
 
 from iwf.command_request import _to_idl_command_request
 from iwf.command_results import from_idl_command_results
@@ -25,6 +25,7 @@
 from iwf.registry import Registry
 from iwf.search_attributes import SearchAttributes
 from iwf.state_decision import StateDecision, _to_idl_state_decision
+from iwf.state_execution_locals import StateExecutionLocals
 from iwf.utils.iwf_typing import assert_not_unset, unset_to_none
 from iwf.workflow_context import WorkflowContext, _from_idl_context
 from iwf.workflow_state import get_input_type
@@ -99,8 +100,13 @@ def handle_workflow_worker_rpc(
         search_attributes = SearchAttributes(
             search_attributes_types, unset_to_none(request.search_attributes)
         )
+        state_execution_locals = StateExecutionLocals(
+            to_map(None), self._options.object_encoder
+        )
 
-        persistence = Persistence(data_attributes, search_attributes)
+        persistence = Persistence(
+            data_attributes, search_attributes, state_execution_locals
+        )
 
         communication = Communication(
             internal_channel_types,
@@ -145,6 +151,9 @@ def handle_workflow_worker_rpc(
         )
         if upsert_sas:
             response.upsert_search_attributes = upsert_sas
+        record_events = state_execution_locals.get_record_events()
+        if len(record_events) > 0:
+            response.record_events = record_events
         if len(communication.get_to_trigger_state_movements()) > 0:
             movements = communication.get_to_trigger_state_movements()
             decision = StateDecision.multi_next_states(*movements)
@@ -188,8 +197,13 @@ def handle_workflow_state_wait_until(
         search_attributes = SearchAttributes(
             search_attributes_types, unset_to_none(request.search_attributes)
         )
+        state_execution_locals = StateExecutionLocals(
+            to_map(None), self._options.object_encoder
+        )
 
-        persistence = Persistence(data_attributes, search_attributes)
+        persistence = Persistence(
+            data_attributes, search_attributes, state_execution_locals
+        )
 
         communication = Communication(
             internal_channel_types,
@@ -211,13 +225,20 @@ def handle_workflow_state_wait_until(
             search_attributes.get_upsert_to_server_string_array_attribute_map(),
         )
 
+        upsert_state_locals = (
+            state_execution_locals.get_upsert_state_execution_local_attributes()
+        )
+        record_events = state_execution_locals.get_record_events()
+
         response = WorkflowStateWaitUntilResponse(
             command_request=_to_idl_command_request(command_request),
             publish_to_inter_state_channel=pubs,
             upsert_data_objects=[
                 KeyValue(k, v)
                 for (k, v) in data_attributes.get_updated_values_to_return().items()
             ],
+            upsert_state_locals=upsert_state_locals,
+            record_events=record_events,
         )
 
         if upsert_sas:
@@ -257,8 +278,13 @@ def handle_workflow_state_execute(
         search_attributes = SearchAttributes(
             search_attributes_types, unset_to_none(request.search_attributes)
         )
+        state_execution_locals = StateExecutionLocals(
+            to_map(request.state_locals), self._options.object_encoder
+        )
 
-        persistence = Persistence(data_attributes, search_attributes)
+        persistence = Persistence(
+            data_attributes, search_attributes, state_execution_locals
+        )
 
         communication = Communication(
             internal_channel_types,
@@ -288,6 +314,10 @@ def handle_workflow_state_execute(
             search_attributes.get_upsert_to_server_double_attribute_map(),
             search_attributes.get_upsert_to_server_string_array_attribute_map(),
         )
+        upsert_state_locals = (
+            state_execution_locals.get_upsert_state_execution_local_attributes()
+        )
+        record_events = state_execution_locals.get_record_events()
 
         response = WorkflowStateExecuteResponse(
             state_decision=_to_idl_state_decision(
@@ -301,6 +331,8 @@ def handle_workflow_state_execute(
                 KeyValue(k, v)
                 for (k, v) in data_attributes.get_updated_values_to_return().items()
             ],
+            upsert_state_locals=upsert_state_locals,
+            record_events=record_events,
         )
 
         if upsert_sas:
@@ -367,3 +399,15 @@ def _create_upsert_search_attributes(
         sas.append(sa)
 
     return sas
+
+
+def to_map(key_values: Union[None, Unset, List[KeyValue]]) -> dict[str, EncodedObject]:
+    key_values = unset_to_none(key_values) or []
+    kvs = {}
+    for kv in key_values:
+        k = unset_to_none(kv.key)
+        v = unset_to_none(kv.value)
+        if k and v:
+            kvs[k] = v
+
+    return kvs