feat: profile memory usage with memray

Makefile

@@ -58,5 +58,10 @@ run-api-server:
 	cd gcp/api && docker build -f Dockerfile.esp -t osv/esp:latest .
 	cd gcp/api && $(install-cmd) && GOOGLE_CLOUD_PROJECT=oss-vdb $(run-cmd) python test_server.py $(HOME)/.config/gcloud/application_default_credentials.json $(ARGS)
 
+profile-api-server:
+	test -f $(HOME)/.config/gcloud/application_default_credentials.json || (echo "GCP Application Default Credentials not set, try 'gcloud auth login --update-adc'"; exit 1)
+	cd gcp/api && docker build -f Dockerfile.esp -t osv/esp:latest .
+	cd gcp/api && $(install-cmd) && GOOGLE_CLOUD_PROJECT=oss-vdb-test $(run-cmd) python test_server.py $(HOME)/.config/gcloud/application_default_credentials.json $(ARGS)
+
 # TODO: API integration tests.
 all-tests: lib-tests worker-tests importer-tests alias-tests website-tests vulnfeed-tests

docker/staging_api_test/perform_api_calls.py

@@ -30,7 +30,7 @@
 import osv
 import osv.logs
 
-BASE_URL = 'https://api.test.osv.dev/v1'
+BASE_URL = 'http://localhost:8080/v1'
 GCP_PROJECT = 'oss-vdb-test'
 BUG_DIR = './all_bugs'
 

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"
+memray = "==1.15.0"
 
 osv = { path = "../../", develop = true }
 

gcp/api/server.py

@@ -1411,13 +1411,28 @@ def main():
       action='store_true',
       default=False,
       help='If set reflection is enabled to allow debugging with grpcurl.')
+  parser.add_argument(
+      '--memray',
+      action='store_true',
+      default=False,
+      help='If set, memray is enabled to profile memory usage.')
 
   args = parser.parse_args()
   port = args.port
   if not port:
     port = int(os.environ.get('PORT', '8000'))
 
-  serve(port, args.local)
+  if args.memray:
+    import memray
+
+    file_path = 'memray-' + datetime.now().strftime('%Y%m%d%H%M%S') + '.bin'
+
+    with memray.Tracker(file_path):
+      logging.info('Starting with memory profiler; output in -> %s', file_path)
+      serve(port, args.local)
+
+  else:
+    serve(port, args.local)
 
 
 if __name__ == '__main__':

gcp/api/test_server.py

@@ -20,7 +20,8 @@
 
 _ESP_PORT = 8080
 _BACKEND_PORT = 8000
-_GCP_PROJECT = 'oss-vdb'
+_GCP_PROJECT = 'oss-vdb-test'
+_USE_MEMRAY = False
 
 
 class ServerInstance:
@@ -45,17 +46,18 @@ def stop(self):
       self.backend.kill()
 
 
-def start_backend(port, log_path):
+def start_backend(port, log_path, memray):
   """Start backend server."""
   log_handle = open(log_path, 'w')
   env = os.environ.copy()
   env['GOOGLE_CLOUD_PROJECT'] = _GCP_PROJECT
 
+  commands = [sys.executable, 'server.py', f'--port={port}']
+  if memray:
+    commands.append('--memray')
+
   backend_proc = subprocess.Popen(
-      [sys.executable, 'server.py', f'--port={port}'],
-      env=env,
-      stdout=log_handle,
-      stderr=subprocess.STDOUT)
+      commands, env=env, stdout=log_handle, stderr=subprocess.STDOUT)
 
   return backend_proc
 
@@ -164,7 +166,7 @@ def start(credential_path,
   esp = None
   try:
     if not no_backend:
-      backend = start_backend(_BACKEND_PORT, 'backend.log')
+      backend = start_backend(_BACKEND_PORT, 'backend.log', _USE_MEMRAY)
     esp = start_esp(port, backend_port, credential_path, 'esp.log')
   except Exception:
     if esp: