diff --git a/ddtrace/profiling/profiler.py b/ddtrace/profiling/profiler.py index 9888ea31f1f..918d0f74889 100644 --- a/ddtrace/profiling/profiler.py +++ b/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() diff --git a/ddtrace/profiling/scheduler.py b/ddtrace/profiling/scheduler.py index 323f4fba24e..9bf2d3d139a 100644 --- a/ddtrace/profiling/scheduler.py +++ b/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 diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 90ad67bc9a9..8d64c69c9d6 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -186,3 +186,5 @@ yaaredis Kinesis AppSec libddwaf +Serverless +serverless diff --git a/releasenotes/notes/feat-serverless-scheduler-31a819bc9eb4f332.yaml b/releasenotes/notes/feat-serverless-scheduler-31a819bc9eb4f332.yaml new file mode 100644 index 00000000000..0059f160364 --- /dev/null +++ b/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). diff --git a/tests/profiling/test_profiler.py b/tests/profiling/test_profiler.py index 7da3d630c2d..39dd413728d 100644 --- a/tests/profiling/test_profiler.py +++ b/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" diff --git a/tests/profiling/test_scheduler.py b/tests/profiling/test_scheduler.py index 49effaa4beb..483ec643527 100644 --- a/tests/profiling/test_scheduler.py +++ b/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()