Natural span context continuation for boto3 sqs receive_message (#152)

Author: Michele Mancioppi

.dockerignore

@@ -0,0 +1 @@
+.gitignore

.pre-commit-config.yaml

@@ -8,14 +8,23 @@ repos:
       additional_dependencies: ['click==8.0.4']
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.971
+    rev: v0.982
     hooks:
     -   id: mypy
-        entry: venv/bin/mypy
+        entry: mypy
+        additional_dependencies:
+        - types-attrs==19.1.0
+        - types-boto==2.49.17
+        - types-PyMySQL==1.0.19
+        - types-PyYAML==6.0.11
+        - types-redis==4.3.20
+        - types-requests==2.28.11.2
+        - types-setuptools==65.5.0.2
+        - types-urllib3<1.27
         exclude: '(src/test/|noxfile\.py)|(src/lumigo_opentelemetry/external/.*)'
         args: ['--explicit-package-bases',  # Needed for src/ci
                '--namespace-packages',  # Needed for opentelemetry package
-               '--disallow-any-generics', '--disallow-untyped-defs', '--disallow-incomplete-defs', '--check-untyped-defs', '--no-implicit-optional', '--warn-redundant-casts', '--warn-unused-ignores', '--warn-return-any', '--no-implicit-reexport', '--strict-equality', '--ignore-missing-imports']
+               '--disallow-any-generics', '--disallow-untyped-defs', '--disallow-incomplete-defs', '--check-untyped-defs', '--no-implicit-optional', '--warn-redundant-casts', '--warn-return-any', '--no-implicit-reexport', '--strict-equality', '--ignore-missing-imports']
 
 -   repo: https://gitlab.com/pycqa/flake8
     rev: 5.0.4

README.md

@@ -193,3 +193,20 @@ tracer_provider.force_flush()
 
 # Now the Python process can terminate, with all the spans closed so far sent to Lumigo
 ```
+
+### Consuming SQS messages with Boto3 receive_message
+
+Messaging instrumentations that retrieve messages from queues tend to be counter-intuitive for end-users: when retrivieng one of more messages from the queue, one would natutally expect that all calls done _using data from those messages_, e.g., sending their content to a database or another queue, would result in spans that are children of the describing the retrivieng of those messages.
+
+Consider the following scenario, which is supported by the `boto3` SQS `receive_message` instrumentation of the Lumigo OpenTelemetry Distro for Python:
+
+```python
+response = client.receive_message(...)  # Instrumentation creates a `span_0` span
+
+for message in response.get("Messages", []):
+  # The SQS.ReceiveMessage span is active in this scope
+  with trace.start_as_current_span("span_1"):  # span_0 is the parent of span_1
+    do_something()
+```
+
+Without the scope provided by the iterator over `response["Messages"]`, `span_1` would be without a parent span, and that would result in a separate invocation and a separate transaction in Lumigo.

noxfile.py

@@ -1,15 +1,21 @@
 from __future__ import annotations
 import os
+import sys
 import tempfile
-from xml.etree import ElementTree
 import time
 from typing import List, Union, Optional
+from xml.etree import ElementTree
 
 import nox
 import requests
 import yaml
 
-from src.ci.tested_versions_utils import (
+# Ensure nox can load local packages
+repo_dir = os.path.dirname(__file__)
+if repo_dir not in sys.path:
+    sys.path.append(repo_dir)
+
+from src.ci.tested_versions_utils import (  # noqa: E402
     NonSemanticVersion,
     SemanticVersion,
     TestedVersions,
@@ -117,6 +123,61 @@ def dependency_versions_to_be_tested(
     ]
 
 
+@nox.session(python=python_versions())
+@nox.parametrize(
+    "boto3_version",
+    dependency_versions_to_be_tested(
+        directory="boto3",
+        dependency_name="boto3",
+        test_untested_versions=should_test_only_untested_versions(),
+    ),
+)
+def integration_tests_boto3_sqs(
+    session,
+    boto3_version,
+):
+    with TestedVersions.save_tests_result("botocore", "boto3", boto3_version):
+        install_package("boto3", boto3_version, session)
+
+        session.install(".")
+
+        abs_path = os.path.abspath("src/test/integration/boto3-sqs/")
+        with tempfile.NamedTemporaryFile(suffix=".txt", prefix=abs_path) as temp_file:
+            full_path = f"{temp_file}.txt"
+
+            with session.chdir("src/test/integration/boto3-sqs"):
+                session.install("-r", "requirements_others.txt")
+
+                try:
+                    session.run(
+                        "sh",
+                        "./scripts/run_app",
+                        env={
+                            "LUMIGO_DEBUG_SPANDUMP": full_path,
+                        },
+                        external=True,
+                    )  # One happy day we will have https://github.com/wntrblm/nox/issues/198
+
+                    # TODO Make this deterministic
+                    # Give time for app to start
+                    time.sleep(8)
+
+                    session.run(
+                        "pytest",
+                        "--tb",
+                        "native",
+                        "--log-cli-level=INFO",
+                        "--color=yes",
+                        "-v",
+                        "./tests/test_boto3_sqs.py",
+                        env={
+                            "LUMIGO_DEBUG_SPANDUMP": full_path,
+                        },
+                    )
+                finally:
+                    kill_process_and_clean_outputs(full_path, "run_app", session)
+
+
 @nox.session(python=python_versions())
 @nox.parametrize(
     "boto3_version",

requirements.txt

@@ -1,14 +1,12 @@
 pre-commit==2.20.0 
 nox==2022.1.7
-black==22.3.0
-click==8.0.4
-flake8==4.0.1
 pytest==7.1.1
 mock==4.0.3
 pytest-cov==3.0.0
 capturer==3.0
 attrs==21.4.0
 requests==2.27.1
+httpretty==1.1.4
 moto==3.1.5
 itsdangerous==2.0.1
 Jinja2==3.0.3
@@ -19,9 +17,11 @@ zipp==3.8.0
 psutil==5.9.1
 pyyaml==6.0
 attrs==21.4.0
-mypy==0.971
-types-requests==2.28.10
-types-PyMySQL==1.0.19
+types-attrs==19.1.0
 types-boto==2.49.17
-types-redis==4.3.20
+types-PyMySQL==1.0.19
 types-PyYAML==6.0.11
+types-redis==4.3.20
+types-requests==2.28.11.2
+types-setuptools==65.5.0.2
+types-urllib3<1.27

src/lumigo_opentelemetry/external/botocore/extensions/sqs.py

@@ -17,13 +17,18 @@
     _AwsSdkExtension,
     _BotoResultT,
 )
+from opentelemetry import context as context_api
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace.span import Span
+from opentelemetry.trace import get_current_span
+from opentelemetry.trace.propagation import _SPAN_KEY
+from opentelemetry.trace.span import NonRecordingSpan, Span
+from typing import Generator
 
 _SUPPORTED_OPERATIONS = ["SendMessage", "SendMessageBatch", "ReceiveMessage"]
 
 
 class _SqsExtension(_AwsSdkExtension):
+
     def extract_attributes(self, attributes: _AttributeMapT):
         queue_url = self._call_context.params.get("QueueUrl")
         if queue_url:
@@ -52,4 +57,68 @@ def on_success(self, span: Span, result: _BotoResultT):
                 span.set_attribute(
                     SpanAttributes.MESSAGING_MESSAGE_ID,
                     result["Messages"][0]["MessageId"],
-                )
+                )
+
+                # Replace the 'Messages' list with a list that restores `span` as trace
+                # context over iterations.
+                ctx = _SqsExtension.ScopeContext(span)
+                result["Messages"] = _SqsExtension.ContextableList(result["Messages"], ctx)
+
+    class ScopeContext:
+        def __init__(self, span: Span):
+            self._span = span
+            self._scope_depth = 0
+            self._token = None
+
+        def _enter(self):
+            # NonRecordingSpan is what we get when there is no span in the context
+            if self._scope_depth == 0 and isinstance(get_current_span(), NonRecordingSpan):
+                # If there is not current active span (e.g. in case of an overarching entry span),
+                # we restore as active the "ReceiveMessage" span. We neither record exceptions on it
+                # nor set statuses, as the "ReceiveMessage" span is already ended when it is used as
+                # parent by spans created while iterating on the messages.
+                self._token = context_api.attach(context_api.set_value(_SPAN_KEY, self._span))
+                self._scope_depth = 1
+            else:
+                self._scope_depth = self._scope_depth + 1
+
+        def _exit(self):
+            if not self._scope_depth:
+                # Something very fishy going on!
+                return
+
+            self._scope_depth = self._scope_depth - 1
+
+            if self._scope_depth == 0 and self._token:
+                # Clear the context outside of the iteration, so that a second iteration would
+                # revert to the previous context
+                context_api.detach(self._token)
+                self._token = None
+
+    class ContextableList(list):
+        """
+        Since the classic way to process SQS messages is using a `for` loop, without a well defined scope like a
+        callback - we are doing something similar to the instrumentation of Kafka-python and instrumenting the
+        `__iter__` functions and the `__getitem__` functions to set the span context of the "ReceiveMessage" span.
+
+        Since the return value from an `SQS.ReceiveMessage` returns a builtin list, we cannot wrap it and change
+        all of the calls for `list.__iter__` and `list.__getitem__` - therefore we use ContextableList.
+
+        This object also implements the `__enter__` and `__exit__` objects for context managers.
+        """
+        def __init__(self, l, scopeContext):
+            super().__init__(l)
+            self._scopeContext = scopeContext
+
+        # Iterator support
+        def __iter__(self) -> Generator:
+            # Try / finally ensures that we detach the context over a `break` issues from inside an iteration
+            try:
+                self._scopeContext._enter()
+
+                index = 0
+                while index < len(self):
+                    yield self[index]
+                    index = index + 1
+            finally:
+                self._scopeContext._exit()

src/lumigo_opentelemetry/libs/general_utils.py

@@ -1,5 +1,5 @@
 try:
-    from collections import Iterable
+    from collections import Iterable  # type: ignore
 except ImportError:
     # collections.Iterables was removed in Python 3.10
     from collections.abc import Iterable

src/test/integration/boto3-sqs/.gitignore

@@ -0,0 +1,6 @@
+bin
+include
+lib
+lib64
+share
+pyvenv.cfg

src/test/integration/boto3-sqs/app/main.py

@@ -0,0 +1,66 @@
+import boto3
+
+from lumigo_opentelemetry import tracer_provider
+from opentelemetry.trace import SpanKind
+
+from moto import mock_sqs
+
+
+@mock_sqs
+def run():
+    tracer = tracer_provider.get_tracer(__name__)
+
+    client = boto3.client("sqs", region_name="eu-central-1")
+    queue = client.create_queue(QueueName="test")
+
+    client.send_message(
+        QueueUrl=queue["QueueUrl"],
+        MessageBody="Message_1",
+    )
+    client.send_message(
+        QueueUrl=queue["QueueUrl"],
+        MessageBody="Message_2",
+    )
+
+    messages_copy = None
+
+    response = client.receive_message(QueueUrl=queue["QueueUrl"], MaxNumberOfMessages=1)
+
+    if response:
+        if "Messages" in response and response["Messages"]:
+            for message in response["Messages"]:
+                # In the iterator, we restore the span for client.receive_message
+                with tracer.start_as_current_span(
+                    "consuming_message_1", kind=SpanKind.INTERNAL
+                ) as processing_span:
+                    processing_span.set_attribute("message.body", message["Body"])
+
+                    with tracer.start_as_current_span(
+                        "consuming_message_2", kind=SpanKind.INTERNAL
+                    ) as processing_span:
+                        processing_span.set_attribute("foo", "bar")
+
+            messages_copy = response["Messages"].copy()
+
+    for message in messages_copy:
+        with tracer.start_as_current_span(
+            "iterator_on_copy", kind=SpanKind.INTERNAL
+        ) as span:
+            span.set_attribute("foo", "bar")
+
+    response = client.receive_message(QueueUrl=queue["QueueUrl"], MaxNumberOfMessages=1)
+    if response:
+        for message in response["Messages"]:
+            # In the iterator, we must restore the span for client.receive_message
+            break  # Here we are supposed to lose the trace context
+
+    with tracer.start_as_current_span(
+        "after_iterator_break", kind=SpanKind.INTERNAL
+    ) as span:
+        span.set_attribute("foo", "bar")
+
+    tracer_provider.force_flush()
+
+
+if __name__ == "__main__":
+    run()

src/test/integration/boto3-sqs/conftest.py

@@ -0,0 +1,8 @@
+import pytest
+
+from test.test_utils.spans_parser import SpansContainer
+
+
+@pytest.fixture(autouse=True)
+def increment_spans_counter():
+    SpansContainer.increment_spans()

src/test/integration/boto3-sqs/requirements_others.txt

@@ -0,0 +1,3 @@
+pytest
+psutil
+moto

src/test/integration/boto3-sqs/scripts/run_app

@@ -0,0 +1,6 @@
+#!/bin/sh -e
+
+script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
+
+
+python3 $(dirname $script_dir)/app/main.py