8000 Implement account ID namespacing for legacy stepfunctions provider · localstack/localstack@c6c71f9 · GitHub
[go: up one dir, main page]
Skip to content

Commit c6c71f9

Browse files
viren-nadkarniviren-nadkarni
committed
Implement account ID namespacing for legacy stepfunctions provider
1 parent 2682b40 commit c6c71f9

File tree

4 files changed

+138
-133
lines changed

4 files changed

+138
-133
lines changed

‎localstack/config.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -941,8 +941,6 @@ def legacy_fallback(envar_name: str, default: T) -> T:
941941
# DEV: sbx_user1051 (default when not provided) Alternative system user or empty string to skip dropping privileges.
942942
LAMBDA_INIT_USER = os.environ.get("LAMBDA_INIT_USER")
943943

944-
# Adding Stepfunctions default port
945-
LOCAL_PORT_STEPFUNCTIONS = int(os.environ.get("LOCAL_PORT_STEPFUNCTIONS") or 8083)
946944
# Stepfunctions lambda endpoint override
947945
STEPFUNCTIONS_LAMBDA_ENDPOINT = os.environ.get("STEPFUNCTIONS_LAMBDA_ENDPOINT", "").strip()
948946

‎localstack/services/stepfunctions/provider.py

Lines changed: 12 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import threading
44

55
from localstack import config
6+
from localstack.aws.accounts import get_aws_account_id
67
from localstack.aws.api import RequestContext, handler
78
from localstack.aws.api.stepfunctions import (
89
CreateStateMachineInput,
@@ -16,11 +17,7 @@
1617
from localstack.aws.forwarder import get_request_forwarder_http
1718
from localstack.constants import LOCALHOST
1819
from localstack.services.plugins import ServiceLifecycleHook
19-
from localstack.services.stepfunctions.stepfunctions_starter import (
20-
start_stepfunctions,
21-
stop_stepfunctions,
22-
wait_for_stepfunctions,
23-
)
20+
from localstack.services.stepfunctions.stepfunctions_starter import StepFunctionsServerManager
2421
from localstack.state import AssetDirectory, StateVisitor
2522

2623
# lock to avoid concurrency issues when creating state machines in parallel (required for StepFunctions-Local)
@@ -30,33 +27,28 @@
3027

3128

3229
class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
30+
server_manager = StepFunctionsServerManager()
31+
3332
def __init__(self):
3433
self.forward_request = get_request_forwarder_http(self.get_forward_url)
3534

3635
def get_forward_url(self) -> str:
3736
"""Return the URL of the backend StepFunctions server to forward requests to"""
38-
return f"http://{LOCALHOST}:{config.LOCAL_PORT_STEPFUNCTIONS}"
37+
account_id = get_aws_account_id()
< 57AE code>38+
server = self.server_manager.get_server_for_account(account_id)
39+
return f"http://{LOCALHOST}:{server.port}"
3940

4041
def accept_state_visitor(self, visitor: StateVisitor):
4142
visitor.visit(AssetDirectory(os.path.join(config.dirs.data, self.service)))
4243

43-
def on_before_start(self):
44-
start_stepfunctions()
45-
wait_for_stepfunctions()
46-
47-
def on_before_state_reset(self):
48-
stop_stepfunctions()
49-
5044
def on_before_state_load(self):
51-
stop_stepfunctions()
45+
self.server_manager. B41A shutdown_all()
5246

53-
def on_after_state_reset(self):
54-
start_stepfunctions()
55-
wait_for_stepfunctions()
47+
def on_before_state_reset(self):
48+
self.server_manager.shutdown_all()
5649

57-
def on_after_state_load(self):
58-
start_stepfunctions()
59-
wait_for_stepfunctions()
50+
def on_before_stop(self):
51+
self.server_manager.shutdown_all()
6052

6153
def create_state_machine(
6254
self, context: RequestContext, request: CreateStateMachineInput

‎localstack/services/stepfunctions/stepfunctions_starter.py

Lines changed: 121 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -1,125 +1,135 @@
11
import logging
2-
import subprocess
2+
import threading
3+
from typing import Any, Dict
34

45
from localstack import config
5-
from localstack.aws.accounts import get_aws_account_id
6-
from localstack.aws.connect import connect_to
7-
from localstack.services.infra import do_run, log_startup_message
86
from localstack.services.stepfunctions.packages import stepfunctions_local_package
97
from localstack.utils.aws import aws_stack
10-
from localstack.utils.common import wait_for_port_open
11-
from localstack.utils.net import wait_for_port_closed
12-
from localstack.utils.run import ShellCommandThread, wait_for_process_to_be_killed
13-
from localstack.utils.sync import retry
8+
from localstack.utils.net import get_free_tcp_port
9+
from localstack.utils.run import ShellCommandThread
10+
from localstack.utils.serving import Server
11+
from localstack.utils.threads import TMP_THREADS, FuncThread
1412

1513
LOG = logging.getLogger(__name__)
1614

1715
# max heap size allocated for the Java process
1816
MAX_HEAP_SIZE = "256m"
1917

20-
# todo: will be replaced with plugin mechanism
21-
PROCESS_THREAD: ShellCommandThread | subprocess.Popen | None = None
22-
23-
24-
# TODO: pass env more explicitly
25-
def get_command(backend_port):
26-
install_dir_stepfunctions = stepfunctions_local_package.get_installed_dir()
27-
cmd = (
28-
"cd %s; PORT=%s java "
29-
"-javaagent:aspectjweaver-1.9.7.jar "
30-
"-Dorg.aspectj.weaver.loadtime.configuration=META-INF/aop.xml "
31-
"-Dcom.amazonaws.sdk.disableCertChecking -Xmx%s "
32-
"-jar StepFunctionsLocal.jar --aws-account %s"
33-
) % (
34-
install_dir_stepfunctions,
35-
backend_port,
36-
MAX_HEAP_SIZE,
37-
get_aws_account_id(),
38-
)
39-
if config.STEPFUNCTIONS_LAMBDA_ENDPOINT.lower() != "default":
40-
lambda_endpoint = config.STEPFUNCTIONS_LAMBDA_ENDPOINT or aws_stack.get_local_service_url(
41-
"lambda"
18+
19+
class StepFunctionsServer(Server):
20+
def __init__(self, port: int, account_id: str, host: str = "localstack") -> None:
21+
self.account_id = account_id
22+
super().__init__(port, host)
23+
24+
def do_start_thread(self) -> FuncThread:
25+
cmd = self.generate_shell_command()
26+
env_vars = self.generate_env_vars()
27+
LOG.debug("Starting StepFunctions process %s with env vars %s", cmd, env_vars)
28+
t = ShellCommandThread(
29+
cmd,
30+
strip_color=True,
31+
env_vars=env_vars,
32+
log_listener=self._log_listener,
33+
name="stepfunctions",
4234
)
43-
cmd += f" --lambda-endpoint {lambda_endpoint}"
44-
# add service endpoint flags
45-
services = [
46-
"athena",
47-
"batch",
48-
"dynamodb",
49-
"ecs",
50-
"eks",
51-
"events",
52-
"glue",
53-
"sagemaker",
54-
"sns",
55-
"sqs",
56-
"stepfunctions",
57-
]
58-
for service in services:
59-
flag = f"--{service}-endpoint"
60-
if service == "stepfunctions":
61-
flag = "--step-functions-endpoint"
62-
elif service == "events":
63-
flag = "--eventbridge-endpoint"
64-
elif service in ["athena", "eks"]:
65-
flag = f"--step-functions-{service}"
66-
endpoint = aws_stack.get_local_service_url(service)
67-
cmd += f" {flag} {endpoint}"
68-
69-
return cmd
70-
71-
72-
def start_stepfunctions(asynchronous: bool = True, persistence_path: str | None = None):
73-
# TODO: introduce Server abstraction for StepFunctions process
74-
global PROCESS_THREAD
75-
backend_port = config.LOCAL_PORT_STEPFUNCTIONS
76-
stepfunctions_local_package.install()
77-
cmd = get_command(backend_port)
78-< F438 div class="diff-text-inner"> log_startup_message("StepFunctions")
79-
# TODO: change ports in stepfunctions.jar, then update here
80-
PROCESS_THREAD = do_run(
81-
cmd,
82-
asynchronous,
83-
strip_color=True,
84-
env_vars={
35+
TMP_THREADS.append(t)
36+
t.start()
37+
return t
38+
39+
def generate_env_vars(self) -> Dict[str, Any]:
40+
return {
8541
"EDGE_PORT": config.EDGE_PORT_HTTP or config.EDGE_PORT,
8642
"EDGE_PORT_HTTP": config.EDGE_PORT_HTTP or config.EDGE_PORT,
87-
"DATA_DIR": persistence_path or config.dirs.data,
88-
},
89-
)
90-
return PROCESS_THREAD
91-
92-
93-
def wait_for_stepfunctions():
94-
retry(check_stepfunctions, sleep=0.5, retries=15)
95-
96-
97-
def stop_stepfunctions():
98-
if not PROCESS_THREAD or not PROCESS_THREAD.process:
99-
return
100-
LOG.debug("Restarting StepFunctions process ...")
101-
102-
pid = PROCESS_THREAD.process.pid
103-
PROCESS_THREAD.stop()
104-
wait_for_port_closed(config.LOCAL_PORT_STEPFUNCTIONS, sleep_time=0.5, retries=15)
105-
try:
106-
# TODO: currently failing in CI (potentially due to a defunct process) - need to investigate!
107-
wait_for_process_to_be_killed(pid, sleep=0.3, retries=10)
108-
except Exception as e:
109-
LOG.warning("StepFunctions process not properly terminated: %s", e)
110-
111-
112-
def check_stepfunctions(expect_shutdown: bool = False, print_error: bool = False) -> None:
113-
out = None
114-
try:
115-
wait_for_port_open(config.LOCAL_PORT_STEPFUNCTIONS, sleep_time=2)
116-
endpoint_url = f"http://127.0.0.1:{config.LOCAL_PORT_STEPFUNCTIONS}"
117-
out = connect_to(endpoint_url=endpoint_url).stepfunctions.list_state_machines()
118-
except Exception:
119-
if print_error:
120-
LOG.exception("StepFunctions health check failed")
121-
122-
if expect_shutdown:
123-
assert out is None
124-
else:
125-
assert out and isinstance(out.get("stateMachines"), list)
43+
"DATA_DIR": config.dirs.data,
44+
}
45+
46+
def generate_shell_command(self, port: int, account_id: str) -> str:
47+
install_dir_stepfunctions = stepfunctions_local_package.get_installed_dir()
48+
49+
cmd = (
50+
"cd %s; PORT=%s java "
51+
"-javaagent:aspectjweaver-1.9.7.jar "
52+
"-Dorg.aspectj.weaver.loadtime.configuration=META-INF/aop.xml "
53+
"-Dcom.amazonaws.sdk.disableCertChecking -Xmx%s "
54+
"-jar StepFunctionsLocal.jar --aws-account %s"
55+
) % (
56+
install_dir_stepfunctions,
57+
port,
58+
MAX_HEAP_SIZE,
59+
account_id,
60+
)
61+
62+
if config.STEPFUNCTIONS_LAMBDA_ENDPOINT.lower() != "default":
63+
lambda_endpoint = (
64+
config.STEPFUNCTIONS_LAMBDA_ENDPOINT or aws_stack.get_local_service_url("lambda")
65+
)
66+
cmd += f" --lambda-endpoint {lambda_endpoint}"
67+
68+
# add service endpoint flags
69+
services = [
70+
"athena",
71+
"batch",
72+
"dynamodb",
73+
"ecs",
74+
"eks",
75+
"events",
76+
"glue",
77+
"sagemaker",
78+
"sns",
79+
"sqs",
80+
"stepfunctions",
81+
]
82+
83+
for service in services:
84+
flag = f"--{service}-endpoint"
85+
if service == "stepfunctions":
86+
flag = "--step-functions-endpoint"
87+
elif service == "events":
88+
flag = "--eventbridge-endpoint"
89+
elif service in ["athena", "eks"]:
90+
flag = f"--step-functions-{service}"
91+
endpoint = aws_stack.get_local_service_url(service)
92+
cmd += f" {flag} {endpoint}"
93+
94+
return cmd
95+
96+
def _log_listener(self, line, **kwargs):
97+
LOG.debug(line.rstrip())
98+
99+
100+
class StepFunctionsServerManager:
101+
def __init__(self):
102+
self._lock = threading.RLock()
103+
self._servers = dict[str, StepFunctionsServer] = {}
104+
105+
def get_server_for_account(self, account_id: str) -> StepFunctionsServer:
106+
if account_id in self._servers:
107+
return self._servers[account_id]
108+
109+
with self._lock:
110+
if account_id in self._servers:
111+
return self._servers[account_id]
112+
113+
LOG.info("Creating StepFunctions server for account %s", account_id)
114+
self._servers[account_id] = self._create_stepfunctions_server(account_id)
115+
self._servers[account_id].start()
116+
if not self._servers[account_id].wait_is_up(timeout=self.default_startup_timeout):
117+
raise TimeoutError("gave up waiting for StepFunctions server to start up")
118+
return self._servers[account_id]
119+
120+
def shutdown_all(self):
121+
with self._lock:
122+
while self._servers:
123+
account_id, server = self._servers.popitem()
124+
LOG.info("Shutting down StepFunctions for account %s", account_id)
125+
server.shutdown()
126+
127+
def _create_stepfunctions_server(self, account_id: str) -> StepFunctionsServer:
128+
port = get_free_tcp_port()
129+
stepfunctions_local_package.install()
130+
131+
server = StepFunctionsServer(
132+
port=port,
133+
account_id=account_id,
134+
)
135+
return server

‎tests/aws/services/stepfunctions/v2/test_stepfunctions_v2.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,29 +195,34 @@ def setup_and_tear_down(aws_client):
195195
zip_file=zip_file,
196196
envvars={"Hello": TEST_RESULT_VALUE},
197197
client=aws_client.lambda_,
198+
s3_client=aws_client.s3,
198199
)
199200
testutil.create_lambda_function(
200201
func_name=TEST_LAMBDA_NAME_2,
201202
zip_file=zip_file,
202203
envvars={"Hello": TEST_RESULT_VALUE_2},
203204
client=aws_client.lambda_,
205+
s3_client=aws_client.s3,
204206
)
205207
testutil.create_lambda_function(
206208
func_name=TEST_LAMBDA_NAME_3,
207209
zip_file=zip_file,
208210
envvars={"Hello": "Replace Value"},
209211
client=aws_client.lambda_,
212+
s3_client=aws_client.s3,
210213
)
211214
testutil.create_lambda_function(
212215
func_name=TEST_LAMBDA_NAME_4,
213216
zip_file=zip_file,
214217
envvars={"Hello": TEST_RESULT_VALUE_4},
215218
client=aws_client.lambda_,
219+
s3_client=aws_client.s3,
216220
)
217221
testutil.create_lambda_function(
218222
func_name=TEST_LAMBDA_NAME_5,
219223
zip_file=zip_file2,
220224
client=aws_client.lambda_,
225+
s3_client=aws_client.s3,
221226
)
222227

223228
active_waiter = lambda_client.get_waiter("function_active_v2")

0 commit comments

Comments
 (0)
0