From 2e442953e6898aaa6d6f350375b84422c9bb2118 Mon Sep 17 00:00:00 2001 From: Giovanni Grano Date: Thu, 13 Nov 2025 18:24:59 +0100 Subject: [PATCH 1/3] use utility function to access --- .../stepfunctions/backend/execution.py | 7 ++++ .../services/stepfunctions/provider.py | 32 +++++++++++-------- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/localstack-core/localstack/services/stepfunctions/backend/execution.py b/localstack-core/localstack/services/stepfunctions/backend/execution.py index 2798bf7d479ee..64960a179a06c 100644 --- a/localstack-core/localstack/services/stepfunctions/backend/execution.py +++ b/localstack-core/localstack/services/stepfunctions/backend/execution.py @@ -94,6 +94,13 @@ def terminated(self) -> None: self.execution.publish_execution_status_change_event() +EXEC_ARN_TO_WORKER: dict[str, ExecutionWorker] = {} + + +def get_exec_worker(arn: str) -> ExecutionWorker | None: + return EXEC_ARN_TO_WORKER.get(arn) + + class Execution: name: Final[str] sm_type: Final[StateMachineType] diff --git a/localstack-core/localstack/services/stepfunctions/provider.py b/localstack-core/localstack/services/stepfunctions/provider.py index 19e3e68e07603..16831779355be 100644 --- a/localstack-core/localstack/services/stepfunctions/provider.py +++ b/localstack-core/localstack/services/stepfunctions/provider.py @@ -139,7 +139,11 @@ ) from localstack.services.stepfunctions.backend.activity import Activity, ActivityTask from localstack.services.stepfunctions.backend.alias import Alias -from localstack.services.stepfunctions.backend.execution import Execution, SyncExecution +from localstack.services.stepfunctions.backend.execution import ( + Execution, + SyncExecution, + get_exec_worker, +) from localstack.services.stepfunctions.backend.state_machine import ( StateMachineInstance, StateMachineRevision, @@ -710,7 +714,7 @@ def send_task_heartbeat( running_executions: list[Execution] = self._get_executions(context, ExecutionStatus.RUNNING) for execution in running_executions: try: - if execution.exec_worker.env.callback_pool_manager.heartbeat( + if get_exec_worker(execution.exec_arn).env.callback_pool_manager.heartbeat( callback_id=task_token ): return SendTaskHeartbeatOutput() @@ -732,7 +736,7 @@ def send_task_success( running_executions: list[Execution] = self._get_executions(context, ExecutionStatus.RUNNING) for execution in running_executions: try: - if execution.exec_worker.env.callback_pool_manager.notify( + if get_exec_worker(execution.exec_arn).env.callback_pool_manager.notify( callback_id=task_token, outcome=outcome ): return SendTaskSuccessOutput() @@ -755,7 +759,7 @@ def send_task_failure( store = self.get_store(context) for execution in store.executions.values(): try: - if execution.exec_worker.env.callback_pool_manager.notify( + if get_exec_worker(execution.exec_arn).env.callback_pool_manager.notify( callback_id=task_token, outcome=outcome ): return SendTaskFailureOutput() @@ -1427,9 +1431,9 @@ def describe_map_run( ) -> DescribeMapRunOutput: store = self.get_store(context) for execution in store.executions.values(): - map_run_record: MapRunRecord | None = ( - execution.exec_worker.env.map_run_record_pool_manager.get(map_run_arn) - ) + map_run_record: MapRunRecord | None = get_exec_worker( + execution.exec_arn + ).env.map_run_record_pool_manager.get(map_run_arn) if map_run_record is not None: return map_run_record.describe() raise ResourceNotFound() @@ -1444,9 +1448,9 @@ def list_map_runs( ) -> ListMapRunsOutput: # TODO: add support for paging. execution = self._get_execution(context=context, execution_arn=execution_arn) - map_run_records: list[MapRunRecord] = ( - execution.exec_worker.env.map_run_record_pool_manager.get_all() - ) + map_run_records: list[MapRunRecord] = get_exec_worker( + execution.exec_arn + ).env.map_run_record_pool_manager.get_all() return ListMapRunsOutput( mapRuns=[map_run_record.list_item() for map_run_record in map_run_records] ) @@ -1467,9 +1471,9 @@ def update_map_run( # TODO: investigate behaviour of empty requests. store = self.get_store(context) for execution in store.executions.values(): - map_run_record: MapRunRecord | None = ( - execution.exec_worker.env.map_run_record_pool_manager.get(map_run_arn) - ) + map_run_record: MapRunRecord | None = get_exec_worker( + execution.exec_arn + ).env.map_run_record_pool_manager.get(map_run_arn) if map_run_record is not None: map_run_record.update( max_concurrency=max_concurrency, @@ -1585,7 +1589,7 @@ def _send_activity_task_started( ) -> None: executions: list[Execution] = self._get_executions(context) for execution in executions: - callback_endpoint = execution.exec_worker.env.callback_pool_manager.get( + callback_endpoint = get_exec_worker(execution.exec_arn).env.callback_pool_manager.get( callback_id=task_token ) if isinstance(callback_endpoint, ActivityCallbackEndpoint): From c98424bce11258f847f4f66fa6a29d9030b15cf7 Mon Sep 17 00:00:00 2001 From: Giovanni Grano Date: Thu, 13 Nov 2025 18:35:05 +0100 Subject: [PATCH 2/3] remove usage in execution --- .../stepfunctions/backend/execution.py | 25 +++---------- .../backend/test_state/execution.py | 4 ++- .../services/stepfunctions/provider.py | 35 ++++++++++++++++--- 3 files changed, 39 insertions(+), 25 deletions(-) diff --git a/localstack-core/localstack/services/stepfunctions/backend/execution.py b/localstack-core/localstack/services/stepfunctions/backend/execution.py index 64960a179a06c..779d56131f7b5 100644 --- a/localstack-core/localstack/services/stepfunctions/backend/execution.py +++ b/localstack-core/localstack/services/stepfunctions/backend/execution.py @@ -15,7 +15,6 @@ ExecutionStatus, GetExecutionHistoryOutput, HistoryEventList, - InvalidName, SensitiveCause, SensitiveError, StartExecutionOutput, @@ -71,11 +70,12 @@ def __init__(self, execution: Execution): self.execution = execution def _reflect_execution_status(self): - exit_program_state: ProgramState = self.execution.exec_worker.env.program_state() + exec_worker = get_exec_worker(self.execution.exec_arn) + exit_program_state: ProgramState = exec_worker.env.program_state() self.execution.stop_date = datetime.datetime.now(tz=datetime.UTC) if isinstance(exit_program_state, ProgramEnded): self.execution.exec_status = ExecutionStatus.SUCCEEDED - self.execution.output = self.execution.exec_worker.env.states.get_input() + self.execution.output = exec_worker.env.states.get_input() elif isinstance(exit_program_state, ProgramStopped): self.execution.exec_status = ExecutionStatus.ABORTED elif isinstance(exit_program_state, ProgramError): @@ -132,8 +132,6 @@ class Execution: error: SensitiveError | None cause: SensitiveCause | None - exec_worker: ExecutionWorker | None - _activity_store: dict[Arn, Activity] def __init__( @@ -264,7 +262,8 @@ def to_execution_list_item(self) -> ExecutionListItem: return item def to_history_output(self) -> GetExecutionHistoryOutput: - env = self.exec_worker.env + exec_worker = get_exec_worker(self.exec_arn) + env = exec_worker.env event_history: HistoryEventList = [] if env is not None: # The execution has not started yet. @@ -314,20 +313,6 @@ def _get_start_execution_worker(self) -> ExecutionWorker: mock_test_case=self.mock_test_case, ) - def start(self) -> None: - # TODO: checks exec_worker does not exists already? - if self.exec_worker: - raise InvalidName() # TODO. - self.exec_worker = self._get_start_execution_worker() - self.exec_status = ExecutionStatus.RUNNING - self.publish_execution_status_change_event() - self.exec_worker.start() - - def stop(self, stop_date: datetime.datetime, error: str | None, cause: str | None): - exec_worker: ExecutionWorker | None = self.exec_worker - if exec_worker: - exec_worker.stop(stop_date=stop_date, cause=cause, error=error) - def publish_execution_status_change_event(self): input_value = ( {} if not self.input_data else to_json_str(self.input_data, separators=(",", ":")) diff --git a/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py b/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py index 92fcb8546c0d6..58755c64af6a7 100644 --- a/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py +++ b/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py @@ -26,6 +26,7 @@ from localstack.services.stepfunctions.backend.execution import ( BaseExecutionWorkerCommunication, Execution, + get_exec_worker, ) from localstack.services.stepfunctions.backend.state_machine import StateMachineInstance from localstack.services.stepfunctions.backend.test_state.execution_worker import ( @@ -43,7 +44,8 @@ class TestCaseExecutionWorkerCommunication(BaseExecutionWorkerCommunication): _execution: TestStateExecution def terminated(self) -> None: - exit_program_state: ProgramState = self.execution.exec_worker.env.program_state() + exec_worker = get_exec_worker(self.execution.exec_arn) + exit_program_state: ProgramState = exec_worker.env.program_state() if isinstance(exit_program_state, ProgramChoiceSelected): self.execution.exec_status = ExecutionStatus.SUCCEEDED self.execution.output = self.execution.exec_worker.env.states.get_input() diff --git a/localstack-core/localstack/services/stepfunctions/provider.py b/localstack-core/localstack/services/stepfunctions/provider.py index 16831779355be..595c39a029dc8 100644 --- a/localstack-core/localstack/services/stepfunctions/provider.py +++ b/localstack-core/localstack/services/stepfunctions/provider.py @@ -140,6 +140,7 @@ from localstack.services.stepfunctions.backend.activity import Activity, ActivityTask from localstack.services.stepfunctions.backend.alias import Alias from localstack.services.stepfunctions.backend.execution import ( + EXEC_ARN_TO_WORKER, Execution, SyncExecution, get_exec_worker, @@ -881,7 +882,15 @@ def start_execution( store.executions[exec_arn] = execution - execution.start() + if get_exec_worker(exec_arn): + raise InvalidName() # TODO + + exec_worker = execution._get_start_execution_worker() + execution.exec_status = ExecutionStatus.RUNNING + execution.publish_execution_status_change_event() + exec_worker.start() + EXEC_ARN_TO_WORKER[exec_arn] = exec_worker + return execution.to_start_output() def start_sync_execution( @@ -955,7 +964,15 @@ def start_sync_execution( ) self.get_store(context).executions[exec_arn] = execution - execution.start() + if get_exec_worker(exec_arn): + raise InvalidName() # TODO + + exec_worker = execution._get_start_execution_worker() + execution.exec_status = ExecutionStatus.RUNNING + execution.publish_execution_status_change_event() + exec_worker.start() + EXEC_ARN_TO_WORKER[exec_arn] = exec_worker + return execution.to_start_sync_execution_output() def describe_execution( @@ -1272,7 +1289,9 @@ def stop_execution( self._raise_resource_type_not_in_context(resource_type=execution.sm_type) stop_date = datetime.datetime.now(tz=datetime.UTC) - execution.stop(stop_date=stop_date, cause=cause, error=error) + + if exec_worker := get_exec_worker(execution.exec_arn): + exec_worker.stop(stop_date=stop_date, cause=cause, error=error) return StopExecutionOutput(stopDate=stop_date) def update_state_machine( @@ -1525,7 +1544,15 @@ def test_state( input_data=input_json, activity_store=self.get_store(context).activities, ) - execution.start() + + if get_exec_worker(exec_arn): + raise InvalidName() # TODO + + exec_worker = execution._get_start_execution_worker() + execution.exec_status = ExecutionStatus.RUNNING + execution.publish_execution_status_change_event() + exec_worker.start() + EXEC_ARN_TO_WORKER[exec_arn] = exec_worker test_state_output = execution.to_test_state_output( inspection_level=inspection_level or InspectionLevel.INFO From 2fb78d095a9b1a7fdcf1b63785a7a13d39cb8adc Mon Sep 17 00:00:00 2001 From: Giovanni Grano Date: Thu, 13 Nov 2025 19:11:27 +0100 Subject: [PATCH 3/3] drop exec_worker from attributes --- .../localstack/services/stepfunctions/backend/execution.py | 1 - .../services/stepfunctions/backend/test_state/execution.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/localstack-core/localstack/services/stepfunctions/backend/execution.py b/localstack-core/localstack/services/stepfunctions/backend/execution.py index 779d56131f7b5..24f9a3746dca9 100644 --- a/localstack-core/localstack/services/stepfunctions/backend/execution.py +++ b/localstack-core/localstack/services/stepfunctions/backend/execution.py @@ -174,7 +174,6 @@ def __init__( self.stop_date = None self.output = None self.output_details = CloudWatchEventsExecutionDataDetails(included=True) - self.exec_worker = None self.error = None self.cause = None self._activity_store = activity_store diff --git a/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py b/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py index 58755c64af6a7..d7b28691b86c8 100644 --- a/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py +++ b/localstack-core/localstack/services/stepfunctions/backend/test_state/execution.py @@ -37,7 +37,6 @@ class TestStateExecution(Execution): - exec_worker: TestStateExecutionWorker | None next_state: str | None class TestCaseExecutionWorkerCommunication(BaseExecutionWorkerCommunication): @@ -48,7 +47,7 @@ def terminated(self) -> None: exit_program_state: ProgramState = exec_worker.env.program_state() if isinstance(exit_program_state, ProgramChoiceSelected): self.execution.exec_status = ExecutionStatus.SUCCEEDED - self.execution.output = self.execution.exec_worker.env.states.get_input() + self.execution.output = exec_worker.env.states.get_input() self.execution.next_state = exit_program_state.next_state_name else: self._reflect_execution_status()