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.

(cherry picked from commit 83ce6ce)

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()