fix: Align trace_context with middleware and add tests

Refactors the existing `trace_context` manager to align its behavior with the
`TraceContextMiddleware`, ensuring consistent and safe observability for
non-web tasks.

This commit addresses these issues by:
- Modifying `trace_context` to only accept `origin` and `trace_id`.
- Adding validation to ensure any provided `trace_id` is a valid UUID, matching the middleware's logic.
- Removing the ability to create arbitrary context variables via `**kwargs`.

Additionally, this change introduces:
- Tests for the `trace_context` manager, covering its functionality, validation, thread-safety, and use as a decorator.
- Documentation in the `README.md` with usage instructions and an example for background tasks.

Author: kdelee

ansible_base/lib/logging/context.py

@@ -13,23 +13,29 @@ class trace_context:
     A context manager and decorator to set the trace context for non-web operations.
     """
 
-    def __init__(self, origin=None, **kwargs):
+    def __init__(self, origin=None, trace_id=None):
         self.origin = origin
-        self.kwargs = kwargs
         self.tokens = []
 
+        if trace_id:
+            try:
+                # Validate that the provided header is a valid UUID
+                uuid.UUID(trace_id)
+                self.trace_id = trace_id
+            except (ValueError, TypeError):
+                # If it's not a valid UUID, discard it and we'll generate a new one
+                self.trace_id = str(uuid.uuid4())
+        else:
+            self.trace_id = str(uuid.uuid4())
+
     def __enter__(self):
-        # Set a new trace ID for this context
-        self.tokens.append(trace_id_var.set(str(uuid.uuid4())))
+        # Set the trace ID for this context
+        self.tokens.append(trace_id_var.set(self.trace_id))
 
         # Set the origin (e.g., 'dispatcher')
         if self.origin:
             self.tokens.append(origin_var.set(self.origin))
 
-        for key, value in self.kwargs.items():
-            var = contextvars.ContextVar(key)
-            self.tokens.append(var.set(value))
-
     def __exit__(self, exc_type, exc_value, traceback):
         # Reset the context variables to their previous state
         for token in self.tokens:

ansible_base/lib/middleware/profiling/README.md

@@ -31,6 +31,8 @@ To enable cProfile support, set the following in your Django settings:
 ANSIBLE_BASE_CPROFILE_REQUESTS = True
 ```
 
+> **Note:** Enabling cProfile has significant performance implications and is intended for temporary, live debugging sessions, not for permanent use in production environments.
+
 ### SQL Profiling Support
 
 When the `ANSIBLE_BASE_SQL_PROFILING` setting is enabled, the middleware provides insights into the database queries executed during a request. It adds the following headers to the response:
@@ -48,6 +50,8 @@ To enable SQL profiling, set the following in your Django settings:
 ANSIBLE_BASE_SQL_PROFILING = True
 ```
 
+> **Note:** This feature is most effective when used in combination with your database's slow query logging capabilities. For high-traffic environments, consider configuring your database to log only a percentage of queries to manage logging overhead.
+
 ## `DABProfiler`
 
 For profiling non-HTTP contexts, such as background tasks or gRPC services, the `DABProfiler` class can be used directly.
@@ -76,6 +80,34 @@ def my_background_task():
     print(f"Task took {elapsed:.3f}s to complete.")
 ```
 
+## `trace_context` for Background Tasks
+
+For adding observability to non-HTTP contexts without the overhead of the `DABProfiler`, the `trace_context` context manager is the ideal tool. It ensures that background tasks can be traced with a unique request ID, just like the `ObservabilityMiddleware` does for web requests.
+
+This is particularly useful for background tasks, such as those initiated by the controller's dispatcher, where you want to correlate all log messages for a specific operation.
+
+### Example Usage
+
+Here's how you might use the `trace_context` manager in the controller's dispatcher to ensure that all work related to a specific job has a consistent trace ID.
+
+```python
+# In a hypothetical controller dispatcher task
+from ansible_base.lib.logging.context import trace_context
+
+def run_job(job_id, parent_trace_id=None):
+    """
+    A background task that runs a job.
+    """
+    # Use the parent_trace_id if it exists; otherwise, a new one will be generated.
+    # The origin is a string that identifies the source of the trace.
+    with trace_context(origin='controller_dispatcher', trace_id=parent_trace_id):
+        # All logging within this block will now have the same trace_id.
+        # logger.info(f"Starting job {job_id}")
+        # ... do work ...
+        # logger.info(f"Finished job {job_id}")
+        pass
+```
+
 ## Visualizing Profile Data
 
 The `.prof` files generated by the cProfile support can be analyzed with a variety of tools.

test_app/tests/lib/logging/test_context.py

@@ -2,27 +2,31 @@
 import threading
 import time
 import unittest
+import uuid
 
-from ansible_base.lib.logging.context import trace_id_var
+import pytest
 
+from ansible_base.lib.logging.context import origin_var, trace_context, trace_id_var
+
+
+class TestTraceContextThreadSafety(unittest.TestCase):
+    """
+    Tests the thread safety of context variables.
+    """
 
-class TestContextSafety(unittest.TestCase):
     def test_trace_id_is_thread_safe(self):
         """
-        Verify that the trace_id context variable is thread-safe.
+        Verify that the trace_id context variable is thread-safe and does not leak between threads.
         """
         results = []
 
         def target_function(thread_id):
             # Set a unique trace ID for this thread
             trace_id_var.set(f"trace-id-{thread_id}")
-
             # Sleep for a random, short duration to encourage thread interleaving
             time.sleep(random.uniform(0.01, 0.05))
-
             # Get the trace ID and verify it has not been changed by another thread
             retrieved_id = trace_id_var.get()
-
             # Store the result of the check for the main thread to verify
             results.append(retrieved_id == f"trace-id-{thread_id}")
 
@@ -38,3 +42,102 @@ def target_function(thread_id):
         # Verify that all threads successfully retrieved their own context
         self.assertEqual(len(results), 10, "Not all threads completed successfully.")
         self.assertTrue(all(results), "Context leaked between threads.")
+
+
+class TestTraceContext:
+    """
+    Tests the functionality of the trace_context context manager and decorator.
+    """
+
+    def test_generates_trace_id(self):
+        """
+        Test that the context manager generates a new trace_id when none is provided.
+        """
+        assert trace_id_var.get() is None
+        with trace_context(origin='test_origin'):
+            generated_id = trace_id_var.get()
+            assert generated_id is not None
+            assert isinstance(uuid.UUID(generated_id), uuid.UUID)
+        assert trace_id_var.get() is None
+
+    def test_uses_provided_trace_id(self):
+        """
+        Test that the context manager uses the trace_id that is passed to it.
+        """
+        provided_id = str(uuid.uuid4())
+        assert trace_id_var.get() is None
+        with trace_context(origin='test_origin', trace_id=provided_id):
+            assert trace_id_var.get() == provided_id
+        assert trace_id_var.get() is None
+
+    def test_handles_invalid_trace_id(self):
+        """
+        Test that the context manager generates a new trace_id if the provided one is invalid.
+        """
+        invalid_id = 'not-a-uuid'
+        assert trace_id_var.get() is None
+        with trace_context(origin='test_origin', trace_id=invalid_id):
+            generated_id = trace_id_var.get()
+            assert generated_id is not None
+            assert generated_id != invalid_id
+            assert isinstance(uuid.UUID(generated_id), uuid.UUID)
+        assert trace_id_var.get() is None
+
+    def test_resets_context_on_exception(self):
+        """
+        Test that context variables are reset even if an exception is raised.
+        """
+        assert trace_id_var.get() is None
+        with pytest.raises(ValueError):
+            with trace_context(origin='test_exception'):
+                raise ValueError("Test exception")
+        assert trace_id_var.get() is None
+
+    def test_as_decorator(self):
+        """
+        Test that the trace_context decorator sets and clears context correctly.
+        """
+
+        @trace_context(origin='test_decorator')
+        def my_function():
+            assert trace_id_var.get() is not None
+            assert origin_var.get() == 'test_decorator'
+
+        assert trace_id_var.get() is None
+        my_function()
+        assert trace_id_var.get() is None
+
+    def test_decorator_with_provided_id(self):
+        """
+        Test that the trace_context decorator uses a provided trace_id.
+        """
+        provided_id = str(uuid.uuid4())
+
+        @trace_context(origin='test_decorator_id', trace_id=provided_id)
+        def my_function():
+            assert trace_id_var.get() == provided_id
+            assert origin_var.get() == 'test_decorator_id'
+
+        assert trace_id_var.get() is None
+        my_function()
+        assert trace_id_var.get() is None
+
+    def test_nested_trace_context(self):
+        """
+        Test that nested trace_context managers work correctly, restoring the previous context.
+        """
+        outer_id = str(uuid.uuid4())
+        with trace_context(origin='outer', trace_id=outer_id):
+            assert trace_id_var.get() == outer_id
+            assert origin_var.get() == 'outer'
+
+            with trace_context(origin='inner'):
+                inner_id = trace_id_var.get()
+                assert inner_id is not None
+                assert inner_id != outer_id
+                assert origin_var.get() == 'inner'
+
+            assert trace_id_var.get() == outer_id
+            assert origin_var.get() == 'outer'
+
+        assert trace_id_var.get() is None