Add LocalStack Prometheus Metrics extension (#92)

Author: gregfurman

prometheus/Makefile

@@ -0,0 +1,38 @@
+VENV_BIN = python3 -m venv
+VENV_DIR ?= .venv
+VENV_ACTIVATE = $(VENV_DIR)/bin/activate
+VENV_RUN = . $(VENV_ACTIVATE)
+
+venv: $(VENV_ACTIVATE)
+
+$(VENV_ACTIVATE): pyproject.toml
+	test -d .venv || $(VENV_BIN) .venv
+	$(VENV_RUN); pip install --upgrade pip setuptools plux wheel
+	$(VENV_RUN); pip install --upgrade black isort pyproject-flake8 flake8-black flake8-isort
+	touch $(VENV_DIR)/bin/activate
+
+clean:
+	rm -rf .venv/
+	rm -rf build/
+	rm -rf .eggs/
+	rm -rf *.egg-info/
+
+lint: venv
+	$(VENV_RUN); python -m pflake8 --show-source
+
+format: venv
+	$(VENV_RUN); python -m isort .; python -m black .
+
+install: venv
+	$(VENV_RUN); python -m pip install -e .[dev]
+
+dist: venv
+	$(VENV_RUN); python setup.py sdist bdist_wheel
+
+publish: clean-dist venv dist
+	$(VENV_RUN); pip install --upgrade twine; twine upload dist/*
+
+clean-dist: clean
+	rm -rf dist/
+
+.PHONY: clean clean-dist dist install publish

prometheus/README.md

@@ -0,0 +1,130 @@
+# LocalStack Prometheus Metrics
+[![Install LocalStack Extension](https://localstack.cloud/gh/extension-badge.svg)](https://app.localstack.cloud/extensions/remote?url=git+https://github.com/localstack/localstack-extensions/#egg=localstack-extension-prometheus-metrics&subdirectory=prometheus)
+
+Instruments, collects, and exposes LocalStack metrics via a [Prometheus](https://prometheus.io/) endpoint.
+
+## Installing
+
+```bash
+localstack extensions install localstack-extension-prometheus-metrics
+```
+
+**Note**: This plugin only supports LocalStack `>=v4.2`
+
+## Usage
+
+Scrape metrics via the endpoint:
+```bash
+curl localhost.localstack.cloud:4566/_extension/metrics
+```
+
+## Quickstart (Docker-Compose)
+
+See the documentation on [Automating extension installation](https://docs.localstack.cloud/user-guide/extensions/managing-extensions/#automating-extensions-installation) for more details.
+
+First, enable the extension by adding it to your LocalStack environment:
+
+```yaml
+services:
+  localstack:
+    environment:
+      - EXTENSION_AUTO_INSTALL=localstack-extension-prometheus-metrics
+```
+
+Next, you'll need to spin up a Prometheus instance to run alongside your LocalStack container. A [configuration file](https://prometheus.io/docs/prometheus/latest/configuration/configuration/#configuration-file) is required.
+
+### Option 1: Using a Volume Mount (Recommended)
+
+Create `prometheus_config.yml`:
+```yaml
+global:
+  scrape_interval: 15s # Set the scrape interval to every 15 seconds
+  scrape_timeout: 5s   # Set the scrape request timeout to 5 seconds
+# Scrape configuration for LocalStack metrics
+scrape_configs:
+  - job_name: 'localstack'
+    static_configs:
+        # Note: The target needs to match the LocalStack container name for the Prometheus container to resolve the endpoint.
+      - targets: ['localstack:4566']    # Target the LocalStack Gateway. 
+    metrics_path: '/_extension/metrics' # Metrics are exposed via `/_extension/metrics` endpoint
+```
+
+And mount it on startup in your `docker-compose.yml`:
+```yaml
+services:
+  # ... LocalStack container should be defined
+  prometheus:
+    image: prom/prometheus
+    ports:
+      - "9090:9090"
+    volumes:
+      - "./prometheus_config.yml:/etc/prometheus/prometheus.yml"
+```
+
+### Option 2: Inline Configuration
+
+Using the Docker Compose top-level [configs](https://docs.docker.com/reference/compose-file/configs/):
+```yaml
+services:
+  prometheus:
+    image: prom/prometheus
+    ports:
+      - "9090:9090"
+    configs:
+      - source: prometheus_config
+        target: /etc/prometheus/prometheus.yml
+
+configs:
+  prometheus_config:
+    content: |
+      global:
+        scrape_interval: 15s
+        scrape_timeout: 5s
+      scrape_configs:
+        - job_name: 'localstack'
+          static_configs:
+            - targets: ['localstack:4566']
+          metrics_path: '/_extension/metrics'
+```
+
+### Full Example
+
+```yaml
+services:
+  localstack:
+    container_name: "${LOCALSTACK_DOCKER_NAME:-localstack-main}"
+    image: localstack/localstack-pro  # required for Pro
+    ports:
+      - "4566:4566"            # LocalStack Gateway
+      - "4510-4559:4510-4559"  # external services port range
+      - "443:443"              # LocalStack HTTPS Gateway (Pro)
+    environment:
+      - LOCALSTACK_AUTH_TOKEN=${LOCALSTACK_AUTH_TOKEN:?}  # required for Pro
+      - DEBUG=${DEBUG:-0}
+      - PERSISTENCE=${PERSISTENCE:-0}
+      - EXTENSION_AUTO_INSTALL=localstack-extension-prometheus-metrics
+    volumes:
+      - "${LOCALSTACK_VOLUME_DIR:-./volume}:/var/lib/localstack"
+      - "/var/run/docker.sock:/var/run/docker.sock"
+
+  prometheus:
+    image: prom/prometheus
+    ports:
+      - "9090:9090"
+    volumes:
+      - "./prometheus_config.yml:/etc/prometheus/prometheus.yml" # Assumes prometheus_config.yml exists in your CWD
+```
+
+## Available Metrics
+
+The Prometheus extension exposes various LocalStack metrics through the `/_extension/metrics` endpoint, including:
+- Request counts by service
+- Request latencies
+- Resource utilization
+- Error rates
+
+For a complete list of available metrics, visit the endpoint directly at `localhost.localstack.cloud:4566/_extension/metrics` when LocalStack is running.
+
+## Licensing
+
+* [client_python](https://github.com/prometheus/client_python) is licensed under the Apache License version 2.

prometheus/localstack_prometheus/__init__.py

@@ -0,0 +1 @@
+name = "localstack_prometheus"

prometheus/localstack_prometheus/expose.py

@@ -0,0 +1,11 @@
+from localstack.extensions.api import http
+from prometheus_client.exposition import choose_encoder
+
+
+def retrieve_metrics(request: http.Request):
+    """Expose the Prometheus metrics"""
+    _generate_latest_metrics, content_type = choose_encoder(
+        request.headers.get("Content-Type", "")
+    )
+    data = _generate_latest_metrics()
+    return http.Response(response=data, status=200, mimetype=content_type)

prometheus/localstack_prometheus/extension.py

@@ -0,0 +1,44 @@
+import logging
+
+from localstack.aws.chain import (
+    CompositeExceptionHandler,
+    CompositeHandler,
+    CompositeResponseHandler,
+)
+from localstack.extensions.api import Extension, http
+
+from localstack_prometheus.expose import retrieve_metrics
+from localstack_prometheus.handler import RequestMetricsHandler, ResponseMetricsHandler
+from localstack_prometheus.instruments.patch import apply_poller_tracking_patches
+
+LOG = logging.getLogger(__name__)
+
+
+class PrometheusMetricsExtension(Extension):
+    name = "prometheus"
+
+    def on_extension_load(self):
+        apply_poller_tracking_patches()
+        LOG.debug("PrometheusMetricsExtension: extension is loaded")
+
+    def on_platform_start(self):
+        LOG.debug("PrometheusMetricsExtension: localstack is starting")
+
+    def on_platform_ready(self):
+        LOG.debug("PrometheusMetricsExtension: localstack is running")
+
+    def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
+        router.add("/_extension/metrics", retrieve_metrics)
+        LOG.debug("Added /metrics endpoint for Prometheus metrics")
+
+    def update_request_handlers(self, handlers: CompositeHandler):
+        # Prepend the RequestMetricsHandler to handlers ensuring it runs first
+        handlers.handlers.insert(0, RequestMetricsHandler())
+
+    def update_response_handlers(self, handlers: CompositeResponseHandler):
+        # Insert the ResponseMetricsHandler as the final handler in the chain.
+        handlers.handlers.append(ResponseMetricsHandler())
+
+    def update_exception_handlers(self, handlers: CompositeExceptionHandler):
+        # TODO
+        pass

prometheus/localstack_prometheus/handler.py

@@ -0,0 +1,80 @@
+import logging
+import time
+
+from localstack.aws.api import RequestContext
+from localstack.aws.chain import Handler, HandlerChain
+from localstack.http import Response
+
+from localstack_prometheus.metrics.core import (
+    LOCALSTACK_IN_FLIGHT_REQUESTS_GAUGE,
+    LOCALSTACK_REQUEST_PROCESSING_DURATION_SECONDS,
+)
+
+LOG = logging.getLogger(__name__)
+
+
+class TimedRequestContext(RequestContext):
+    start_time: float | None
+
+
+class RequestMetricsHandler(Handler):
+    """
+    Handler that records the start time of incoming requests
+    """
+
+    def __call__(
+        self, chain: HandlerChain, context: TimedRequestContext, response: Response
+    ):
+        # Record the start time
+        context.start_time = time.perf_counter()
+
+        # Do not record metrics if no service operation information is found
+        if not context.service_operation:
+            return
+
+        service, operation = context.service_operation
+        LOCALSTACK_IN_FLIGHT_REQUESTS_GAUGE.labels(
+            service=service, operation=operation
+        ).inc()
+
+
+class ResponseMetricsHandler(Handler):
+    """
+    Handler that records metrics when a response is ready
+    """
+
+    def __call__(
+        self, chain: HandlerChain, context: TimedRequestContext, response: Response
+    ):
+        # Do not record metrics if no service operation information is found
+        if not context.service_operation:
+            return
+
+        service, operation = context.service_operation
+        LOCALSTACK_IN_FLIGHT_REQUESTS_GAUGE.labels(
+            service=service, operation=operation
+        ).dec()
+
+        # Do not record if response is None
+        if response is None:
+            return
+
+        # Do not record if no start_time attribute is found
+        if not hasattr(context, "start_time") or context.start_time is None:
+            return
+
+        duration = time.perf_counter() - context.start_time
+
+        if (ex := context.service_exception) is not None:
+            status = ex.code
+        else:
+            status = "success"
+
+        status_code = str(response.status_code)
+
+        LOCALSTACK_REQUEST_PROCESSING_DURATION_SECONDS.labels(
+            service=service,
+            operation=operation,
+            status=status,
+            status_code=status_code,
+        ).observe(duration)

prometheus/localstack_prometheus/instruments/__init__.py

@@ -0,0 +1,52 @@
+import logging
+
+from localstack.services.lambda_.event_source_mapping.pollers.dynamodb_poller import (
+    DynamoDBPoller,
+)
+from localstack.services.lambda_.event_source_mapping.pollers.kinesis_poller import (
+    KinesisPoller,
+)
+from localstack.services.lambda_.event_source_mapping.pollers.sqs_poller import (
+    SqsPoller,
+)
+from localstack.services.lambda_.event_source_mapping.senders.lambda_sender import (
+    LambdaSender,
+)
+from localstack.utils.patch import Patch, Patches
+
+from localstack_prometheus.instruments.poller import tracked_poll_events
+from localstack_prometheus.instruments.sender import tracked_send_events
+from localstack_prometheus.instruments.sqs_poller import tracked_sqs_handle_messages
+from localstack_prometheus.instruments.stream_poller import tracked_get_records
+
+LOG = logging.getLogger(__name__)
+
+
+def apply_poller_tracking_patches():
+    """Apply all poller metrics tracking patches in one call"""
+    patches = Patches(
+        [
+            # Track entire poll_events function
+            Patch.function(target=SqsPoller.poll_events, fn=tracked_poll_events),
+            Patch.function(target=KinesisPoller.poll_events, fn=tracked_poll_events),
+            Patch.function(target=DynamoDBPoller.poll_events, fn=tracked_poll_events),
+            # Track when events get sent to the target lambda
+            Patch.function(target=LambdaSender.send_events, fn=tracked_send_events),
+            # TODO: Standardise a single abstract method that all Poller subclasses can use to fetch records
+            # SQS-specific patches
+            Patch.function(
+                target=SqsPoller.handle_messages, fn=tracked_sqs_handle_messages
+            ),
+            # Stream-specific patches
+            Patch.function(target=KinesisPoller.get_records, fn=tracked_get_records),
+            Patch.function(target=DynamoDBPoller.get_records, fn=tracked_get_records),
+            # TODO: How should KafkaPollers be handled?
+        ]
+    )
+
+    # TODO: Investigate patching subclasses of Poller and Sender to ensure all children have changes
+    # since currently, Pipes Senders and Kafka Pollers are unsupported.
+
+    patches.apply()
+    LOG.debug("Applied all poller event and latency tracking patches")
+    return patches