chore: define explicit state containers for providers (#13423)

bentsku · web-flow · commit ff6e5e08d24e · 2025-12-01T18:32:55.000+01:00
diff --git a/localstack-core/localstack/services/acm/provider.py b/localstack-core/localstack/services/acm/provider.py
@@ -10,6 +10,7 @@
     RequestCertificateResponse,
 )
 from localstack.services import moto
+from localstack.state import StateVisitor
 from localstack.utils.patch import patch
 
 # reduce the validation wait time from 60 (default) to 10 seconds
@@ -100,6 +101,9 @@ def describe(describe_orig, self):
 
 
 class AcmProvider(AcmApi):
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(acm_models.acm_backends)
+
     @handler("RequestCertificate", expand=False)
     def request_certificate(
         self,
diff --git a/localstack-core/localstack/services/apigateway/legacy/provider.py b/localstack-core/localstack/services/apigateway/legacy/provider.py
@@ -118,14 +118,19 @@
 from localstack.services.apigateway.legacy.helpers import multi_value_dict_for_list
 from localstack.services.apigateway.legacy.invocations import invoke_rest_api_from_request
 from localstack.services.apigateway.legacy.router_asf import ApigatewayRouter, to_invocation_context
-from localstack.services.apigateway.models import ApiGatewayStore, RestApiContainer
+from localstack.services.apigateway.models import (
+    ApiGatewayStore,
+    RestApiContainer,
+    apigateway_stores,
+)
 from localstack.services.apigateway.next_gen.execute_api.router import (
     ApiGatewayRouter as ApiGatewayRouterNextGen,
 )
 from localstack.services.apigateway.patches import apply_patches
 from localstack.services.edge import ROUTER
 from localstack.services.moto import call_moto, call_moto_with_request
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws.arns import InvalidArnException, get_partition, parse_arn
 from localstack.utils.collections import (
     DelSafeDict,
@@ -192,6 +197,12 @@ def on_after_init(self):
         apply_patches()
         self.router.register_routes()
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.apigateway import apigateway_backends
+
+        visitor.visit(apigateway_backends)
+        visitor.visit(apigateway_stores)
+
     @handler("TestInvokeMethod", expand=False)
     def test_invoke_method(
         self, context: RequestContext, request: TestInvokeMethodRequest
diff --git a/localstack-core/localstack/services/cloudwatch/provider.py b/localstack-core/localstack/services/cloudwatch/provider.py
@@ -35,6 +35,7 @@
 from localstack.services.cloudwatch.alarm_scheduler import AlarmScheduler
 from localstack.services.edge import ROUTER
 from localstack.services.plugins import SERVICE_PLUGINS, ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws import arns
 from localstack.utils.aws.arns import extract_account_id_from_arn, lambda_function_name
 from localstack.utils.aws.request_context import (
@@ -306,8 +307,13 @@ def __init__(self):
         self.tags = TaggingService()
         self.alarm_scheduler = None
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(cloudwatch_backends)
+
     def on_after_init(self):
         ROUTER.add(PATH_GET_RAW_METRICS, self.get_raw_metrics)
+
+    def on_before_start(self):
         self.start_alarm_scheduler()
 
     def on_before_state_reset(self):
@@ -337,9 +343,10 @@ def start_alarm_scheduler(self):
             self.alarm_scheduler = AlarmScheduler()
 
     def shutdown_alarm_scheduler(self):
-        LOG.debug("stopping cloudwatch scheduler")
-        self.alarm_scheduler.shutdown_scheduler()
-        self.alarm_scheduler = None
+        if self.alarm_scheduler:
+            LOG.debug("stopping cloudwatch scheduler")
+            self.alarm_scheduler.shutdown_scheduler()
+            self.alarm_scheduler = None
 
     def delete_alarms(self, context: RequestContext, alarm_names: AlarmNames, **kwargs) -> None:
         moto.call_moto(context)
diff --git a/localstack-core/localstack/services/cloudwatch/provider_v2.py b/localstack-core/localstack/services/cloudwatch/provider_v2.py
@@ -161,6 +161,8 @@ def accept_state_visitor(self, visitor: StateVisitor):
 
     def on_after_init(self):
         ROUTER.add(PATH_GET_RAW_METRICS, self.get_raw_metrics)
+
+    def on_before_start(self):
         self.start_alarm_scheduler()
 
     def on_before_state_reset(self):
@@ -192,9 +194,10 @@ def start_alarm_scheduler(self):
             self.alarm_scheduler = AlarmScheduler()
 
     def shutdown_alarm_scheduler(self):
-        LOG.debug("stopping cloudwatch scheduler")
-        self.alarm_scheduler.shutdown_scheduler()
-        self.alarm_scheduler = None
+        if self.alarm_scheduler:
+            LOG.debug("stopping cloudwatch scheduler")
+            self.alarm_scheduler.shutdown_scheduler()
+            self.alarm_scheduler = None
 
     def delete_alarms(self, context: RequestContext, alarm_names: AlarmNames, **kwargs) -> None:
         """
diff --git a/localstack-core/localstack/services/configservice/provider.py b/localstack-core/localstack/services/configservice/provider.py
@@ -1,5 +1,9 @@
 from localstack.aws.api.config import ConfigApi
+from localstack.state import StateVisitor
 
 
 class ConfigProvider(ConfigApi):
-    pass
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.config.models import config_backends
+
+        visitor.visit(config_backends)
diff --git a/localstack-core/localstack/services/dynamodb/provider.py b/localstack-core/localstack/services/dynamodb/provider.py
@@ -535,6 +535,7 @@ def __init__(self):
         self.server = self._new_dynamodb_server()
         self._expired_items_worker = ExpiredItemsWorker()
         self._router_rules = []
+        # TODO: fix _event_forwarder to have lazy instantiation of the ThreadPoolExecutor
         self._event_forwarder = EventForwarder()
 
     def on_before_start(self):
diff --git a/localstack-core/localstack/services/dynamodb/v2/provider.py b/localstack-core/localstack/services/dynamodb/v2/provider.py
@@ -392,6 +392,7 @@ def on_before_stop(self):
 
     def accept_state_visitor(self, visitor: StateVisitor):
         visitor.visit(dynamodb_stores)
+        # FIXME: DDB v2 does not depend on dynamodbstreams_stores as DynamoDB Local holds all the state
         visitor.visit(dynamodbstreams_stores)
         visitor.visit(AssetDirectory(self.service, os.path.join(config.dirs.data, self.service)))
 
diff --git a/localstack-core/localstack/services/dynamodbstreams/provider.py b/localstack-core/localstack/services/dynamodbstreams/provider.py
@@ -37,6 +37,7 @@
     table_name_from_stream_arn,
 )
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.collections import select_from_typed_dict
 
 LOG = logging.getLogger(__name__)
@@ -57,6 +58,11 @@ class DynamoDBStreamsProvider(DynamodbstreamsApi, ServiceLifecycleHook):
     def __init__(self):
         self.shard_to_region = {}
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from localstack.services.dynamodbstreams.models import dynamodbstreams_stores
+
+        visitor.visit(dynamodbstreams_stores)
+
     def describe_stream(
         self,
         context: RequestContext,
diff --git a/localstack-core/localstack/services/dynamodbstreams/v2/provider.py b/localstack-core/localstack/services/dynamodbstreams/v2/provider.py
@@ -18,6 +18,7 @@
 from localstack.services.dynamodb.v2.provider import DynamoDBProvider, modify_context_region
 from localstack.services.dynamodbstreams.dynamodbstreams_api import get_original_region
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws.arns import parse_arn
 
 LOG = logging.getLogger(__name__)
@@ -32,6 +33,11 @@ def __init__(self):
         self.server = DynamodbServer.get()
         self.shard_to_region = {}
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        # DynamoDB Streams state is entirely dependent on DynamoDB Local state, and does not hold state itself
+        # the DynamoDB provider is responsible for the persistence of DDB Streams
+        pass
+
     def on_after_init(self):
         # add response processor specific to ddblocal
         handlers.modify_service_response.append(self.service, modify_ddblocal_arns)
diff --git a/localstack-core/localstack/services/ec2/provider.py b/localstack-core/localstack/services/ec2/provider.py
@@ -94,6 +94,7 @@
 from localstack.services.ec2.patches import apply_patches
 from localstack.services.moto import call_moto, call_moto_with_request
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.patch import patch
 from localstack.utils.strings import first_char_to_upper, long_uid, short_uid
 
@@ -107,6 +108,11 @@ class Ec2Provider(Ec2Api, ABC, ServiceLifecycleHook):
     def on_after_init(self):
         apply_patches()
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.ec2.models import ec2_backends
+
+        visitor.visit(ec2_backends)
+
     @handler("DescribeAvailabilityZones", expand=False)
     def describe_availability_zones(
         self,
diff --git a/localstack-core/localstack/services/es/provider.py b/localstack-core/localstack/services/es/provider.py
@@ -68,6 +68,7 @@
 )
 from localstack.aws.connect import connect_to
 from localstack.services.opensearch.packages import ELASTICSEARCH_DEFAULT_VERSION
+from localstack.state import StateVisitor
 
 
 def _version_to_opensearch(
@@ -208,6 +209,11 @@ def exception_mapper():
 
 
 class EsProvider(EsApi):
+    def accept_state_visitor(self, visitor: StateVisitor):
+        # ES state entirely depends on `opensearch`, and delegates its entire state to it
+        # we do not need to manage state in ES
+        pass
+
     def create_elasticsearch_domain(
         self,
         context: RequestContext,
diff --git a/localstack-core/localstack/services/events/provider.py b/localstack-core/localstack/services/events/provider.py
@@ -168,6 +168,7 @@
     recursive_remove_none_values_from_dict,
 )
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.common import truncate
 from localstack.utils.event_matcher import matches_event
 from localstack.utils.strings import long_uid
@@ -246,6 +247,9 @@ def __init__(self):
         self._connection_service_store: ConnectionServiceDict = {}
         self._api_destination_service_store: ApiDestinationServiceDict = {}
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(events_stores)
+
     def on_before_start(self):
         JobScheduler.start()
 
diff --git a/localstack-core/localstack/services/events/v1/provider.py b/localstack-core/localstack/services/events/v1/provider.py
@@ -45,6 +45,7 @@
 from localstack.services.events.v1.models import EventsStore, events_stores
 from localstack.services.moto import call_moto
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws.arns import event_bus_arn, parse_arn
 from localstack.utils.aws.client_types import ServicePrincipal
 from localstack.utils.aws.message_forwarding import send_event_to_target
@@ -83,6 +84,14 @@ def on_before_start(self):
     def on_before_stop(self):
         JobScheduler.shutdown()
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.events.models import events_backends
+
+        from localstack.services.events.v1.models import events_stores
+
+        visitor.visit(events_backends)
+        visitor.visit(events_stores)
+
     @route("/_aws/events/rules/<path:rule_arn>/trigger")
     def trigger_scheduled_rule(self, request: Request, rule_arn: str):
         """Developer endpoint to trigger a scheduled rule."""
diff --git a/localstack-core/localstack/services/firehose/provider.py b/localstack-core/localstack/services/firehose/provider.py
@@ -94,6 +94,7 @@
     convert_source_config_to_desc,
 )
 from localstack.services.firehose.models import FirehoseStore, firehose_stores
+from localstack.state import StateVisitor
 from localstack.utils.aws.arns import (
     extract_account_id_from_arn,
     extract_region_from_arn,
@@ -251,8 +252,12 @@ class FirehoseProvider(FirehoseApi):
 
     def __init__(self) -> None:
         super().__init__()
+        # TODO: stop/restart the kinesis listeners when stopping the service / reset the state / restore the state
         self.kinesis_listeners = {}
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(firehose_stores)
+
     @staticmethod
     def get_store(account_id: str, region_name: str) -> FirehoseStore:
         return firehose_stores[account_id][region_name]
diff --git a/localstack-core/localstack/services/iam/provider.py b/localstack-core/localstack/services/iam/provider.py
@@ -75,6 +75,7 @@
 )
 from localstack.services.iam.resources.service_linked_roles import SERVICE_LINKED_ROLES
 from localstack.services.moto import call_moto
+from localstack.state import StateVisitor
 from localstack.utils.aws.request_context import extract_access_key_id_from_auth_header
 
 LOG = logging.getLogger(__name__)
@@ -110,6 +111,9 @@ def __init__(self):
         apply_iam_patches()
         self.policy_simulator = BasicIAMPolicySimulator()
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(iam_backends)
+
     @handler("CreateRole", expand=False)
     def create_role(
         self, context: RequestContext, request: CreateRoleRequest
diff --git a/localstack-core/localstack/services/kms/provider.py b/localstack-core/localstack/services/kms/provider.py
@@ -138,6 +138,7 @@
     validate_alias_name,
 )
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws.arns import get_partition, kms_alias_arn, parse_arn
 from localstack.utils.collections import PaginatedList
 from localstack.utils.common import select_attributes
@@ -201,6 +202,9 @@ class KmsProvider(KmsApi, ServiceLifecycleHook):
     - VerifyMac
     """
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        visitor.visit(kms_stores)
+
     #
     # Helpers
     #
diff --git a/localstack-core/localstack/services/logs/provider.py b/localstack-core/localstack/services/logs/provider.py
@@ -43,6 +43,7 @@
 from localstack.services.logs.models import get_moto_logs_backend, logs_stores
 from localstack.services.moto import call_moto
 from localstack.services.plugins import ServiceLifecycleHook
+from localstack.state import StateVisitor
 from localstack.utils.aws import arns
 from localstack.utils.aws.client_types import ServicePrincipal
 from localstack.utils.bootstrap import is_api_enabled
@@ -57,6 +58,12 @@ def __init__(self):
         super().__init__()
         self.cw_client = connect_to().cloudwatch
 
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.logs.models import logs_backends
+
+        visitor.visit(logs_backends)
+        visitor.visit(logs_stores)
+
     def put_log_events(
         self,
         context: RequestContext,
diff --git a/localstack-core/localstack/services/resource_groups/provider.py b/localstack-core/localstack/services/resource_groups/provider.py
@@ -1,5 +1,9 @@
 from localstack.aws.api.resource_groups import ResourceGroupsApi
+from localstack.state import StateVisitor
 
 
 class ResourceGroupsProvider(ResourceGroupsApi):
-    pass
+    def accept_state_visitor(self, visitor: StateVisitor):
+        from moto.resourcegroups.models import resourcegroups_backends
+
+        visitor.visit(resourcegroups_backends)
diff --git a/localstack-core/localstack/services/resourcegroupstaggingapi/provider.py b/localstack-core/localstack/services/resourcegroupstaggingapi/provider.py
@@ -1,7 +1,12 @@
 from abc import ABC
 
 from localstack.aws.api.resourcegroupstaggingapi import ResourcegroupstaggingapiApi
+from localstack.state import StateVisitor
 
 
 class ResourcegroupstaggingapiProvider(ResourcegroupstaggingapiApi, ABC):
-    pass
+    def accept_state_visitor(self, visitor: StateVisitor):
+        # currently, Moto resourcegroupstaggingapi stores all tags into the other services backend, so their backend
+        # does not hold any state and is not worth saving. It only holds direct references to other services
+        # It only holds pagination tokens that are not worth keeping
+        pass
