Spike: add profiling support with functiontrace

localstack/config.py

@@ -421,6 +421,12 @@ def in_docker():
 # whether to enable debugpy
 DEVELOP = is_env_true("DEVELOP")
 
+# whether to enable profiling
+ENABLE_PROFILING = is_env_true("ENABLE_PROFILING")
+
+# Whether to automatically upload the captured profile to the firefox profiler
+ENABLE_PROFILING_REPORT_UPLOAD = is_env_true("ENABLE_PROFILING_REPORT_UPLOAD")
+
 # PORT FOR DEBUGGER
 DEVELOP_PORT = int(os.environ.get("DEVELOP_PORT", "").strip() or DEFAULT_DEVELOP_PORT)
 
@@ -1126,6 +1132,8 @@ def use_custom_dns():
  "DYNAMODB_WRITE_ERROR_PROBABILITY",
  "EAGER_SERVICE_LOADING",
  "ENABLE_CONFIG_UPDATES",
+ "ENABLE_PROFILING",
+ "ENABLE_PROFILING_REPORT_UPLOAD",
  "EXTRA_CORS_ALLOWED_HEADERS",
  "EXTRA_CORS_ALLOWED_ORIGINS",
  "EXTRA_CORS_EXPOSE_HEADERS",

localstack/packages/functiontrace.py

@@ -0,0 +1,79 @@
+import logging
+from typing import List
+
+from localstack.constants import ARTIFACTS_REPO
+from localstack.packages import InstallTarget, Package, PackageInstaller
+from localstack.packages.core import PermissionDownloadInstaller
+from localstack.utils.files import chmod_r, mkdir
+from localstack.utils.http import download_github_artifact
+from localstack.utils.run import run
+
+LOG = logging.getLogger(__name__)
+
+
+class FunctionTracePackage(Package):
+ def __init__(self):
+ super().__init__("FunctionTrace", "latest")
+
+ def get_versions(self) -> List[str]:
+ return ["latest"]
+
+ def _get_installer(self, version: str) -> PackageInstaller:
+ return FunctionTracePackageInstaller("functiontrace", version)
+
+
+class FunctionTracePackageInstaller(PackageInstaller):
+ # TODO: migrate this to the upcoming pip installer
+
+ def is_installed(self) -> bool:
+ try:
+ import functiontrace
+
+ assert functiontrace
+ return True
+ except ModuleNotFoundError:
+ return False
+
+ def _get_install_marker_path(self, install_dir: str) -> str:
+ # TODO: This method currently does not provide the actual install_marker.
+ # Since we overwrote is_installed(), this installer does not install anything under
+ # var/static libs, and we also don't need an executable, we don't need it to operate the installer.
+ # fix with migration to pip installer
+ return install_dir
+
+ def _install(self, target: InstallTarget) -> None:
+ cmd = "pip install functiontrace"
+ run(cmd)
+
+
+class FunctionTraceServerPackage(Package):
+ def __init__(self):
+ super().__init__("FunctionTraceServer", "latest")
+
+ def get_versions(self) -> List[str]:
+ return ["latest"]
+
+ def _get_installer(self, version: str) -> PackageInstaller:
+ return FunctionTraceServerPackageInstaller("functiontrace-server", version)
+
+
+class FunctionTraceServerPackageInstaller(PermissionDownloadInstaller):
+ # TODO: migrate this to the upcoming pip installer
+
+ def _get_download_url(self) -> str:
+ return (
+ ARTIFACTS_REPO
+ + "/raw/79c588fe97515a5124f58c4ae28201938084cc64/functiontrace/amd64/functiontrace-server"
+ )
+
+ def _install(self, target: InstallTarget) -> None:
+ target_directory = self._get_install_dir(target)
+ mkdir(target_directory)
+ download_url = self._get_download_url()
+ target_path = self._get_install_marker_path(target_directory)
+ download_github_artifact(download_url, target_path)
+ chmod_r(self.get_executable_path(), 0o777)
+
+
+functiontrace_package = FunctionTracePackage()
+functiontrace_server_package = FunctionTraceServerPackage()

localstack/profiler/plugin.py

@@ -0,0 +1,22 @@
+from localstack import config
+from localstack.runtime import hooks
+
+
+@hooks.on_infra_start()
+def start_profiling():
+ if not config.ENABLE_PROFILING:
+ return
+
+ from localstack.profiler.start import start_profiling
+
+ start_profiling()
+
+
+@hooks.on_infra_shutdown()
+def stop_profiling():
+ if not config.ENABLE_PROFILING:
+ return
+
+ from localstack.profiler.start import stop_profiling
+
+ stop_profiling()

localstack/profiler/start.py

@@ -0,0 +1,52 @@
+import logging
+import os
+from pathlib import Path
+
+from localstack import config
+from localstack.packages.functiontrace import functiontrace_package, functiontrace_server_package
+from localstack.profiler.upload import upload_profile
+
+LOG = logging.getLogger(__name__)
+
+OUTPUT_DIR = Path(config.dirs.cache) / "profiles"
+
+
+def start_profiling():
+ functiontrace_server_package.install()
+
+ # functiontrace-server must be on the PATH
+ os.environ["PATH"] += os.pathsep + functiontrace_server_package.get_installed_dir()
+
+ functiontrace_package.install()
+
+ import _functiontrace
+ import functiontrace
+
+ LOG.debug("storing profiles to %s", OUTPUT_DIR)
+ OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+ functiontrace.setup_dependencies()
+
+ LOG.info("Starting profiling")
+ _functiontrace.begin_tracing(str(OUTPUT_DIR))
+
+
+def stop_profiling():
+ import _functiontrace
+
+ LOG.info("Stopping profiling")
+ _functiontrace.terminate()
+
+ # only enable profile uploading if user opts in
+ if not config.ENABLE_PROFILING_REPORT_UPLOAD:
+ return
+
+ # determine last profile
+ file_candidates = (
+ path for path in Path(OUTPUT_DIR).glob("functiontrace*.json*") if "latest" not in str(path)
+ )
+ profiles = sorted(file_candidates, key=lambda file: file.stat().st_mtime)
+ latest_profile = profiles[0]
+ url = upload_profile(latest_profile)
+ # TODO: add flag to turn off automatic profile uploading
+ LOG.info("LocalStack profiling report uploaded and is viewable at: '%s'", url)

localstack/profiler/upload.py

@@ -0,0 +1,68 @@
+import base64
+import json
+from pathlib import Path
+
+import requests
+
+
+def jwt_base64_decode(payload):
+ """Decode a Base64 encoded string from a JWT token.
+
+ JWT encodes using the URLSafe base64 algorithm and then removes the
+ padding. This function does the opposite: adds the padding back and then
+ uses the URLSafe base64 algorithm to decode the string.
+ """
+ # Thanks Simon Sapin
+ # (https://stackoverflow.com/questions/2941995/python-ignore-incorrect-padding-error-when-base64-decoding)
+ missing_padding = len(payload) % 4
+ if missing_padding:
+ payload += "=" * (4 - missing_padding)
+
+ decoded_bytes = base64.urlsafe_b64decode(payload)
+ decoded_str = decoded_bytes.decode("utf-8")
+ return decoded_str
+
+
+def handle_input_line(line):
+ """This handles one line of input, that should be a JWT token.
+
+ This will first split the token in its 3 components, base64 decode the 2nd
+ one that's the payload, json-parse it, and finally print the key
+ "profileToken" from that JSON payload.
+ """
+ _, payload, _ = line.strip().split(".")
+
+ decoded_str = jwt_base64_decode(payload)
+ json_payload = json.loads(decoded_str)
+ token = json_payload["profileToken"]
+ return token
+
+
+def upload_profile(path: Path) -> str:
+ """
+ Upload the profile to a public place, and return the URL the profile is viewable from.
+ """
+ with path.open("rb") as infile:
+ r = requests.post(
+ "https://api.profiler.firefox.com/compressed-store",
+ headers={
+ "Accept": "application/vnd.firefox-profiler+json;version=1.0",
+ },
+ data=infile,
+ )
+ r.raise_for_status()
+
+ token = r.text
+
+ profile_token = handle_input_line(token)
+ return f"https://profiler.firefox.com/public/{profile_token}"
+
+
+# Execute it only when run directly
+if __name__ == "__main__":
+ import sys
+
+ for file_path in sys.argv[1:]:
+ file_path = Path(file_path)
+ url = upload_profile(file_path)
+ print(f"Hosted at: {url}")