Collect metrics once per lambda invocation (#1643)

* Collect metrics once before ending lambda function

* Allow a MetricsSet instance to be added to the registry

* Add a test

* Add first pass at docs

* Apply suggestions from code review

Co-authored-by: Brandon Morelli <[email protected]>
Co-authored-by: Benjamin Wohlwend <[email protected]>

* More review suggestions

* Simplify API to accept string or class object

* Remove some unnecessary checks (type hints are enough)

* Rename MetricsSet to MetricSet for more consistency

* Add an alias for backwards compat

* Document metrics_sets and fix up custom metrics docs

* Rename Client._metrics to Client.metrics

* Add __module__ to metricset identifier

Co-authored-by: Benjamin Wohlwend <[email protected]>

* Add a deprecation warning for client._metrics

* CHANGELOG

Co-authored-by: Brandon Morelli <[email protected]>
Co-authored-by: Benjamin Wohlwend <[email protected]>

Author: 3 people

CHANGELOG.asciidoc

@@ -37,6 +37,7 @@ endif::[]
 ===== Features
  * Add backend granularity data to SQL backends as well as Cassandra and pymongo {pull}1585[#1585], {pull}1639[#1639]
  * Add support for instrumenting the Elasticsearch 8 Python client {pull}1642[#1642]
+ * Add docs and better support for custom metrics, including in AWS Lambda {pull}1643[#1643]
 
 [float]
 ===== Bug fixes

docs/configuration.asciidoc

@@ -1113,6 +1113,20 @@ See <<prometheus-metricset>> for more information.
 
 NOTE: This feature is currently in beta status.
 
+[float]
+[[config-metrics_sets]]
+==== `metrics_sets`
+
+[options="header"]
+|============
+| Environment                  | Django/Flask   | Default
+| `ELASTIC_APM_METRICS_SETS`   | `METRICS_SETS` | ["elasticapm.metrics.sets.cpu.CPUMetricSet"]
+|============
+
+List of import paths for the MetricSets that should be used to collect metrics.
+
+See <<custom-metrics>> for more information.
+
 [float]
 [[config-central_config]]
 ==== `central_config`

docs/metrics.asciidoc

@@ -11,6 +11,7 @@ These metrics will be sent regularly to the APM Server and from there to Elastic
 * <<cpu-memory-metricset>>
 * <<breakdown-metricset>>
 * <<prometheus-metricset>>
+* <<custom-metrics>>
 
 [float]
 [[cpu-memory-metricset]]
@@ -160,3 +161,55 @@ All metrics collected from `prometheus_client` are prefixed with `"prometheus.me
 [[prometheus-metricset-beta]]
 ===== Beta limitations
  * The metrics format may change without backwards compatibility in future releases.
+
+[float]
+[[custom-metrics]]
+=== Custom Metrics
+
+Custom metrics allow you to send your own metrics to Elasticsearch.
+
+The most common way to send custom metrics is with the
+<<prometheus-metricset,Prometheus metric set>>.  However, you can also use your
+own metric set. If you collect the metrics manually in your code, you can use
+the base `MetricSet` class:
+
+[source,python]
+----
+from elasticapm.metrics.base_metrics import MetricSet
+
+client = elasticapm.Client()
+metricset = client.metrics.register(MetricSet)
+
+for x in range(10):
+    metricset.counter("my_counter").inc()
+----
+
+Alternatively, you can create your own MetricSet class which inherits from the
+base class. In this case, you'll usually want to override the `before_collect`
+method, where you can gather and set metrics before they are collected and sent
+to Elasticsearch.
+
+You can add your `MetricSet` class as shown in the example above, or you can
+add an import string for your class to the <<config-metrics_sets,`metrics_sets`>>
+configuration option:
+
+[source,bash]
+----
+ELASTIC_APM_METRICS_SETS="elasticapm.metrics.sets.cpu.CPUMetricSet,myapp.metrics.MyMetricSet"
+----
+
+Your MetricSet might look something like this:
+
+[source,python]
+----
+from elasticapm.metrics.base_metrics import MetricSet
+
+class MyAwesomeMetricSet(MetricSet):
+    def before_collect(self):
+        self.gauge("my_gauge").set(myapp.some_value)
+----
+
+In the example above, the MetricSet would look up `myapp.some_value` and set
+the metric `my_gauge` to that value. This would happen whenever metrics are
+collected/sent, which is controlled by the
+<<config-metrics_interval,`metrics_interval`>> setting.

elasticapm/base.py

@@ -198,15 +198,15 @@ def __init__(self, config=None, **inline):
         )
         self.include_paths_re = stacks.get_path_regex(self.config.include_paths) if self.config.include_paths else None
         self.exclude_paths_re = stacks.get_path_regex(self.config.exclude_paths) if self.config.exclude_paths else None
-        self._metrics = MetricsRegistry(self)
+        self.metrics = MetricsRegistry(self)
         for path in self.config.metrics_sets:
-            self._metrics.register(path)
+            self.metrics.register(path)
         if self.config.breakdown_metrics:
-            self._metrics.register("elasticapm.metrics.sets.breakdown.BreakdownMetricSet")
+            self.metrics.register("elasticapm.metrics.sets.breakdown.BreakdownMetricSet")
         if self.config.prometheus_metrics:
-            self._metrics.register("elasticapm.metrics.sets.prometheus.PrometheusMetrics")
+            self.metrics.register("elasticapm.metrics.sets.prometheus.PrometheusMetrics")
         if self.config.metrics_interval:
-            self._thread_managers["metrics"] = self._metrics
+            self._thread_managers["metrics"] = self.metrics
         compat.atexit_register(self.close)
         if self.config.central_config:
             self._thread_managers["config"] = self.config
@@ -682,6 +682,11 @@ def check_server_version(
         lte = lte or (2**32,)  # let's assume APM Server version will never be greater than 2^32
         return bool(gte <= self.server_version <= lte)
 
+    @property
+    def _metrics(self):
+        warnings.warn(DeprecationWarning("Use `client.metrics` instead"))
+        return self.metrics
+
 
 class DummyClient(Client):
     """Sends messages into an empty void"""

elasticapm/contrib/serverless/aws.py

@@ -81,6 +81,10 @@ def __init__(self, name: Optional[str] = None, elasticapm_client: Optional[Clien
 
         # Disable all background threads except for transport
         kwargs["metrics_interval"] = "0ms"
+        kwargs["breakdown_metrics"] = False
+        if "metrics_sets" not in kwargs and "ELASTIC_APM_METRICS_SETS" not in os.environ:
+            # Allow users to override metrics sets
+            kwargs["metrics_sets"] = []
         kwargs["central_config"] = False
         kwargs["cloud_provider"] = "none"
         kwargs["framework_name"] = "AWS Lambda"
@@ -220,6 +224,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 elasticapm.set_transaction_outcome(outcome="failure", override=False)
 
         self.client.end_transaction()
+        # Collect any custom+prometheus metrics if enabled
+        self.client.metrics.collect()
 
         try:
             logger.debug("flushing elasticapm")

elasticapm/metrics/base_metrics.py

@@ -31,6 +31,7 @@
 import threading
 import time
 from collections import defaultdict
+from typing import Union
 
 from elasticapm.conf import constants
 from elasticapm.utils.logging import get_logger
@@ -56,25 +57,36 @@ def __init__(self, client, tags=None):
         self._collect_timer = None
         super(MetricsRegistry, self).__init__()
 
-    def register(self, class_path):
+    def register(self, metricset: Union[str, type]) -> "MetricSet":
         """
         Register a new metric set
-        :param class_path: a string with the import path of the metricset class
+
+        :param metricset: a string with the import path of the metricset class,
+            or a class object that can be used to instantiate the metricset.
+            If a class object is used, you can use the class object or
+            `metricset.__name__` to retrieve the metricset using `get_metricset`.
+        :return: the metricset instance
         """
-        if class_path in self._metricsets:
-            return
+        class_id = metricset if isinstance(metricset, str) else f"{metricset.__module__}.{metricset.__name__}"
+        if class_id in self._metricsets:
+            return self._metricsets[class_id]
         else:
-            try:
-                class_obj = import_string(class_path)
-                self._metricsets[class_path] = class_obj(self)
-            except ImportError as e:
-                logger.warning("Could not register %s metricset: %s", class_path, str(e))
-
-    def get_metricset(self, class_path):
+            if isinstance(metricset, str):
+                try:
+                    class_obj = import_string(metricset)
+                    self._metricsets[metricset] = class_obj(self)
+                except ImportError as e:
+                    logger.warning("Could not register %s metricset: %s", metricset, str(e))
+            else:
+                self._metricsets[class_id] = metricset(self)
+        return self._metricsets.get(class_id)
+
+    def get_metricset(self, metricset: Union[str, type]) -> "MetricSet":
+        metricset = metricset if isinstance(metricset, str) else f"{metricset.__module__}.{metricset.__name__}"
         try:
-            return self._metricsets[class_path]
+            return self._metricsets[metricset]
         except KeyError:
-            raise MetricSetNotFound(class_path)
+            raise MetricSetNotFound(metricset)
 
     def collect(self):
         """
@@ -114,7 +126,7 @@ def ignore_patterns(self):
         return self.client.config.disable_metrics or []
 
 
-class MetricsSet(object):
+class MetricSet(object):
     def __init__(self, registry):
         self._lock = threading.Lock()
         self._counters = {}
@@ -282,13 +294,17 @@ def before_collect(self):
         pass
 
     def before_yield(self, data):
+        """
+        A method that is called right before the data is yielded to be sent
+        to Elasticsearch. Can be used to modify the data.
+        """
         return data
 
     def _labels_to_key(self, labels):
         return tuple((k, str(v)) for k, v in sorted(labels.items()))
 
 
-class SpanBoundMetricSet(MetricsSet):
+class SpanBoundMetricSet(MetricSet):
     def before_yield(self, data):
         tags = data.get("tags", None)
         if tags:
@@ -495,3 +511,9 @@ def reset(self):
 class MetricSetNotFound(LookupError):
     def __init__(self, class_path):
         super(MetricSetNotFound, self).__init__("%s metric set not found" % class_path)
+
+
+# This is for backwards compatibility for the brave souls who were using
+# the undocumented system for custom metrics before we fixed it up and
+# documented it.
+MetricsSet = MetricSet

elasticapm/metrics/sets/cpu_linux.py

@@ -34,7 +34,7 @@
 import resource
 import threading
 
-from elasticapm.metrics.base_metrics import MetricsSet
+from elasticapm.metrics.base_metrics import MetricSet
 
 SYS_STATS = "/proc/stat"
 MEM_STATS = "/proc/meminfo"
@@ -72,7 +72,7 @@ def __init__(self, limit, usage, stat):
         self.stat = stat if os.access(stat, os.R_OK) else None
 
 
-class CPUMetricSet(MetricsSet):
+class CPUMetricSet(MetricSet):
     def __init__(
         self,
         registry,

elasticapm/metrics/sets/cpu_psutil.py

@@ -28,15 +28,15 @@
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from elasticapm.metrics.base_metrics import MetricsSet
+from elasticapm.metrics.base_metrics import MetricSet
 
 try:
     import psutil
 except ImportError:
     raise ImportError("psutil not found. Install it to get system and process metrics")
 
 
-class CPUMetricSet(MetricsSet):
+class CPUMetricSet(MetricSet):
     def __init__(self, registry):
         psutil.cpu_percent(interval=None)
         self._process = psutil.Process()

elasticapm/metrics/sets/prometheus.py

@@ -33,10 +33,10 @@
 
 import prometheus_client
 
-from elasticapm.metrics.base_metrics import MetricsSet
+from elasticapm.metrics.base_metrics import MetricSet
 
 
-class PrometheusMetrics(MetricsSet):
+class PrometheusMetrics(MetricSet):
     def __init__(self, registry):
         super(PrometheusMetrics, self).__init__(registry)
         self._prometheus_registry = prometheus_client.REGISTRY

elasticapm/traces.py

@@ -233,7 +233,7 @@ def __init__(
         self._span_timers_lock = threading.Lock()
         self._dropped_span_statistics = defaultdict(lambda: {"count": 0, "duration.sum.us": 0})
         try:
-            self._breakdown = self.tracer._agent._metrics.get_metricset(
+            self._breakdown = self.tracer._agent.metrics.get_metricset(
                 "elasticapm.metrics.sets.breakdown.BreakdownMetricSet"
             )
         except (LookupError, AttributeError):