feat(api): set up google cloud profiler (#2936)

#2793 This PR attempts to set up Google Cloud Profiler following this [guide](https://cloud.google.com/profiler/docs/profiling-python).
google · Nov 29, 2024 · 05fa231 · 05fa231
1 parent 8a8febc
commit 05fa231
Show file tree

Hide file tree

Showing 3 changed files with 98 additions and 13 deletions.
diff --git a/gcp/api/poetry.lock b/gcp/api/poetry.lock
diff --git a/gcp/api/pyproject.toml b/gcp/api/pyproject.toml
@@ -12,6 +12,7 @@ requests = "==2.32.3"
 grpcio = "==1.64.1"
 grpcio-reflection = "==1.62.2"
 grpcio-health-checking = "==1.62.2"
+google-cloud-profiler = "==4.1.0"
 
 osv = { path = "../../", develop = true }
 

diff --git a/gcp/api/server.py b/gcp/api/server.py
@@ -50,6 +50,8 @@
 
 from gcp.api.cursor import QueryCursor
 
+import googlecloudprofiler
+
 _SHUTDOWN_GRACE_DURATION = 5
 
 _MAX_SINGLE_QUERY_TIME = timedelta(seconds=20)
@@ -1395,6 +1397,13 @@ def main():
     setup_gcp_logging('api-backend')
     logging.getLogger().addFilter(trace_filter)
 
+    # Profiler initialization. It starts a daemon thread which continuously
+    # collects and uploads profiles. Best done as early as possible.
+    try:
+      googlecloudprofiler.start(service="osv-api-profiler")
+    except (ValueError, NotImplementedError) as e:
+      logging.error(e)
+
   logging.getLogger().setLevel(logging.INFO)
 
   parser = argparse.ArgumentParser(