Enrich example handler with correlation and duration metrics

src/lambda_obs/handler.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import time
 from typing import Any
 
 from .logging import Logger
@@ -11,21 +12,38 @@
 metrics = Metrics(namespace="LambdaObsStarter", service="lambda-obs-starter")
 
 
+def _request_id(event: dict[str, Any], context: Any) -> str | None:
+    if context is not None and getattr(context, "aws_request_id", None):
+        return str(context.aws_request_id)
+    headers = (event or {}).get("headers") or {}
+    for key in ("x-correlation-id", "X-Correlation-Id", "x-request-id"):
+        if key in headers and headers[key]:
+            return str(headers[key])
+    return (event or {}).get("correlation_id")
+
+
 def handler(event: dict[str, Any], context: Any = None) -> dict[str, Any]:
     """Handle a simple greeting event.
 
     Expected event shape: ``{"name": "world"}``.
+    Optional correlation via ``correlation_id`` or ``headers.x-correlation-id``.
     """
+    started = time.perf_counter()
     name = (event or {}).get("name") or "world"
-    bound = logger.bind(function="handler")
+    cid = _request_id(event or {}, context)
+    bound = logger.with_correlation_id(cid).bind(function="handler")
+    if cid:
+        metrics.add_dimension("correlation_present", "true")
     bound.info("invocation started", name=name)
 
     metrics.add_metric("Invocations", 1, unit="Count")
     try:
         greeting = f"hello {name}"
         metrics.add_metric("Success", 1, unit="Count")
-        bound.info("invocation succeeded", greeting=greeting)
-        return {"ok": True, "message": greeting}
+        elapsed_ms = (time.perf_counter() - started) * 1000.0
+        metrics.timing("HandlerDuration", elapsed_ms)
+        bound.info("invocation succeeded", greeting=greeting, duration_ms=round(elapsed_ms, 3))
+        return {"ok": True, "message": greeting, "correlation_id": bound._base.get("correlation_id")}
     except Exception as exc:  # pragma: no cover - defensive
         metrics.add_metric("Errors", 1, unit="Count")
         bound.error("invocation failed", error=str(exc))

tests/test_handler.py

@@ -1,12 +1,34 @@
+from types import SimpleNamespace
+
 from lambda_obs.handler import handler
 
 
 def test_handler_greets_name():
     result = handler({"name": "lei"}, None)
     assert result["ok"] is True
     assert result["message"] == "hello lei"
+    assert "correlation_id" in result
 
 
 def test_handler_default_name():
     result = handler({}, None)
     assert result["message"] == "hello world"
+
+
+def test_handler_uses_event_correlation_id():
+    result = handler({"name": "lei", "correlation_id": "abc-999"}, None)
+    assert result["correlation_id"] == "abc-999"
+
+
+def test_handler_uses_header_correlation_id():
+    result = handler(
+        {"name": "lei", "headers": {"x-correlation-id": "hdr-1"}},
+        None,
+    )
+    assert result["correlation_id"] == "hdr-1"
+
+
+def test_handler_uses_context_request_id():
+    ctx = SimpleNamespace(aws_request_id="req-42")
+    result = handler({"name": "lei"}, ctx)
+    assert result["correlation_id"] == "req-42"