feat(profiling): add serverless scheduler and conditionally use it (#3870)

Based on Julien's feedback in our RFC, I've added a small scheduler which flushes after the periodic function is called 60 times. This ensures that profiles are sampled over the course of 60s of runtime instead of wallclock time.

Author: astuyve

ddtrace/profiling/profiler.py

@@ -131,7 +131,14 @@ class _ProfilerInstance(service.Service):
 
     _recorder = attr.ib(init=False, default=None)
     _collectors = attr.ib(init=False, default=None)
-    _scheduler = attr.ib(init=False, default=None)
+    _scheduler = attr.ib(
+        init=False,
+        default=None,
+        type=scheduler.Scheduler,
+    )
+    _lambda_function_name = attr.ib(
+        init=False, factory=lambda: os.environ.get("AWS_LAMBDA_FUNCTION_NAME"), type=Optional[str]
+    )
 
     ENDPOINT_TEMPLATE = "https://intake.profile.{}"
 
@@ -165,6 +172,9 @@ def _build_default_exporters(self):
             # to the agent base path.
             endpoint_path = "profiling/v1/input"
 
+        if self._lambda_function_name is not None:
+            self.tags.update({"functionname": self._lambda_function_name.encode("utf-8")})
+
         return [
             http.PprofHTTPExporter(
                 service=self.service,
@@ -209,9 +219,11 @@ def __attrs_post_init__(self):
         exporters = self._build_default_exporters()
 
         if exporters:
-            self._scheduler = scheduler.Scheduler(
-                recorder=r, exporters=exporters, before_flush=self._collectors_snapshot
-            )
+            if self._lambda_function_name is None:
+                scheduler_class = scheduler.Scheduler
+            else:
+                scheduler_class = scheduler.ServerlessScheduler
+            self._scheduler = scheduler_class(recorder=r, exporters=exporters, before_flush=self._collectors_snapshot)
 
         self.set_asyncio_event_loop_policy()
 

ddtrace/profiling/scheduler.py

@@ -65,3 +65,35 @@ def periodic(self):
             self.flush()
         finally:
             self.interval = max(0, self._configured_interval - (compat.monotonic() - start_time))
+
+
+@attr.s
+class ServerlessScheduler(Scheduler):
+    """Serverless scheduler that works on, e.g., AWS Lambda.
+
+    The idea with this scheduler is to not sleep 60s, but to sleep 1s and flush out profiles after 60 sleeping period.
+    As the service can be frozen a few seconds after flushing out a profile, we want to make sure the next flush is not
+    > 60s later, but after at least 60 periods of 1s.
+
+    """
+
+    # We force this interval everywhere
+    FORCED_INTERVAL = 1.0
+    FLUSH_AFTER_INTERVALS = 60.0
+
+    _interval = attr.ib(default=FORCED_INTERVAL, type=float)
+    _profiled_intervals = attr.ib(init=False, default=0)
+
+    def periodic(self):
+        # Check both the number of intervals and time frame to be sure we don't flush, e.g., empty profiles
+        if self._profiled_intervals >= self.FLUSH_AFTER_INTERVALS and (compat.time_ns() - self._last_export) >= (
+            self.FORCED_INTERVAL * self.FLUSH_AFTER_INTERVALS
+        ):
+            try:
+                super(ServerlessScheduler, self).periodic()
+            finally:
+                # Override interval so it's always back to the value we n
+                self.interval = self.FORCED_INTERVAL
+                self._profiled_intervals = 0
+        else:
+            self._profiled_intervals += 1

docs/spelling_wordlist.txt

@@ -186,3 +186,5 @@ yaaredis
 Kinesis
 AppSec
 libddwaf
+Serverless
+serverless

releasenotes/notes/feat-serverless-scheduler-31a819bc9eb4f332.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Adds support for Lambda profiling, which can be enabled by starting the profiler outside of the handler (on cold start).

tests/profiling/test_profiler.py

@@ -9,6 +9,7 @@
 from ddtrace.profiling import event
 from ddtrace.profiling import exporter
 from ddtrace.profiling import profiler
+from ddtrace.profiling import scheduler
 from ddtrace.profiling.collector import asyncio
 from ddtrace.profiling.collector import memalloc
 from ddtrace.profiling.collector import stack
@@ -378,3 +379,11 @@ def test_default_collectors():
     else:
         assert any(isinstance(c, asyncio.AsyncioLockCollector) for c in p._profiler._collectors)
     p.stop(flush=False)
+
+
+def test_profiler_serverless(monkeypatch):
+    # type: (...) -> None
+    monkeypatch.setenv("AWS_LAMBDA_FUNCTION_NAME", "foobar")
+    p = profiler.Profiler()
+    assert isinstance(p._scheduler, scheduler.ServerlessScheduler)
+    assert p.tags["functionname"] == b"foobar"

tests/profiling/test_scheduler.py

@@ -1,6 +1,9 @@
 # -*- encoding: utf-8 -*-
 import logging
 
+import mock
+
+from ddtrace.internal import compat
 from ddtrace.profiling import event
 from ddtrace.profiling import exporter
 from ddtrace.profiling import recorder
@@ -54,3 +57,20 @@ def call_me():
     assert caplog.record_tuples == [
         (("ddtrace.profiling.scheduler", logging.ERROR, "Scheduler before_flush hook failed"))
     ]
+
+
+@mock.patch("ddtrace.profiling.scheduler.Scheduler.periodic")
+def test_serverless_periodic(mock_periodic):
+    r = recorder.Recorder()
+    s = scheduler.ServerlessScheduler(r, [exporter.NullExporter()])
+    # Fake start()
+    s._last_export = compat.time_ns()
+    s.periodic()
+    assert s._profiled_intervals == 1
+    mock_periodic.assert_not_called()
+    s._last_export = compat.time_ns() - 65
+    s._profiled_intervals = 65
+    s.periodic()
+    assert s._profiled_intervals == 0
+    assert s.interval == 1
+    mock_periodic.assert_called()