Adding Correlation Context API and propagator

opentelemetry-api/src/opentelemetry/correlationcontext/__init__.py

@@ -0,0 +1,100 @@
+# Copyright 2020, OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import typing
+
+from opentelemetry.context import get_value, set_value
+from opentelemetry.context.context import Context
+
+_CORRELATION_CONTEXT_KEY = "correlation-context"
+
+
+def get_correlations(
+    context: typing.Optional[Context] = None,
+) -> typing.Dict[str, object]:
+    """ Returns the name/value pairs in the CorrelationContext
+
+    Args:
+        context: The Context to use. If not set, uses current Context
+
+    Returns:
+        Name/value pairs in the CorrelationContext
+    """
+    correlations = get_value(_CORRELATION_CONTEXT_KEY, context=context)
+    if isinstance(correlations, dict):
+        return correlations.copy()
+    return {}
+
+
+def get_correlation(
+    name: str, context: typing.Optional[Context] = None
+) -> typing.Optional[object]:
+    """ Provides access to the value for a name/value pair in the CorrelationContext
+
+    Args:
+        name: The name of the value to retrieve
+        context: The Context to use. If not set, uses current Context
+
+    Returns:
+        The value associated with the given name, or null if the given name is
+        not present.
+    """
+    return get_correlations(context=context).get(name)
+
+
+def set_correlation(
+    name: str, value: object, context: typing.Optional[Context] = None
+) -> Context:
+    """Sets a value in the CorrelationContext
+
+    Args:
+        name: The name of the value to set
+        value: The value to set
+        context: The Context to use. If not set, uses current Context
+
+    Returns:
+        A Context with the value updated
+    """
+    correlations = get_correlations(context=context)
+    correlations[name] = value
+    return set_value(_CORRELATION_CONTEXT_KEY, correlations, context=context)
+
+
+def remove_correlation(
+    name: str, context: typing.Optional[Context] = None
+) -> Context:
+    """Removes a value from the CorrelationContext
+    Args:
+        name: The name of the value to remove
+        context: The Context to use. If not set, uses current Context
+
+    Returns:
+        A Context with the name/value removed
+    """
+    correlations = get_correlations(context=context)
+    correlations.pop(name, None)
+
+    return set_value(_CORRELATION_CONTEXT_KEY, correlations, context=context)
+
+
+def clear_correlations(context: typing.Optional[Context] = None) -> Context:
+    """Removes all values from the CorrelationContext
+    Args:
+        context: The Context to use. If not set, uses current Context
+
+    Returns:
+        A Context with all correlations removed
+    """
+    return set_value(_CORRELATION_CONTEXT_KEY, {}, context=context)

opentelemetry-api/src/opentelemetry/correlationcontext/propagation/__init__.py

@@ -0,0 +1,108 @@
+# Copyright 2020, OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import re
+import typing
+import urllib.parse
+
+from opentelemetry import correlationcontext
+from opentelemetry.context import get_current
+from opentelemetry.context.context import Context
+from opentelemetry.trace.propagation import httptextformat
+
+
+class CorrelationContextPropagator(httptextformat.HTTPTextFormat):
+    MAX_HEADER_LENGTH = 8192
+    MAX_PAIR_LENGTH = 4096
+    MAX_PAIRS = 180
+    _CORRELATION_CONTEXT_HEADER_NAME = "otcorrelationcontext"
+
+    def extract(
+        self,
+        get_from_carrier: httptextformat.Getter[
+            httptextformat.HTTPTextFormatT
+        ],
+        carrier: httptextformat.HTTPTextFormatT,
+        context: typing.Optional[Context] = None,
+    ) -> Context:
+        """ Extract CorrelationContext from the carrier.
+
+        See `opentelemetry.trace.propagation.httptextformat.HTTPTextFormat.extract`
+        """
+
+        if context is None:
+            context = get_current()
+
+        header = _extract_first_element(
+            get_from_carrier(carrier, self._CORRELATION_CONTEXT_HEADER_NAME)
+        )
+
+        if not header or len(header) > self.MAX_HEADER_LENGTH:
+            return context
+
+        correlations = header.split(",")
+        total_correlations = self.MAX_PAIRS
+        for correlation in correlations:
+            if total_correlations <= 0:
+                return context
+            total_correlations -= 1
+            if len(correlation) > self.MAX_PAIR_LENGTH:
+                continue
+            try:
+                name, value = correlation.split("=", 1)
+            except Exception:  # pylint: disable=broad-except
+                continue
+            context = correlationcontext.set_correlation(
+                urllib.parse.unquote(name).strip(),
+                urllib.parse.unquote(value).strip(),
+                context=context,
+            )
+
+        return context
+
+    def inject(
+        self,
+        set_in_carrier: httptextformat.Setter[httptextformat.HTTPTextFormatT],
+        carrier: httptextformat.HTTPTextFormatT,
+        context: typing.Optional[Context] = None,
+    ) -> None:
+        """Injects CorrelationContext into the carrier.
+
+        See `opentelemetry.trace.propagation.httptextformat.HTTPTextFormat.inject`
+        """
+        correlations = correlationcontext.get_correlations(context=context)
+        if not correlations:
+            return
+
+        correlation_context_string = _format_correlations(correlations)
+        set_in_carrier(
+            carrier,
+            self._CORRELATION_CONTEXT_HEADER_NAME,
+            correlation_context_string,
+        )
+
+
+def _format_correlations(correlations: typing.Dict[str, object]) -> str:
+    return ",".join(
+        key + "=" + urllib.parse.quote_plus(str(value))
+        for key, value in correlations.items()
+    )
+
+
+def _extract_first_element(
+    items: typing.Iterable[httptextformat.HTTPTextFormatT],
+) -> typing.Optional[httptextformat.HTTPTextFormatT]:
+    if items is None:
+        return None
+    return next(iter(items), None)

opentelemetry-api/src/opentelemetry/propagators/__init__.py

@@ -58,6 +58,10 @@ def example_route():
 import opentelemetry.trace as trace
 from opentelemetry.context import get_current
 from opentelemetry.context.context import Context
+from opentelemetry.correlationcontext.propagation import (
+    CorrelationContextPropagator,
+)
+from opentelemetry.propagators import composite
 from opentelemetry.trace.propagation import httptextformat
 from opentelemetry.trace.propagation.tracecontexthttptextformat import (
     TraceContextHTTPTextFormat,
@@ -106,8 +110,8 @@ def inject(
     get_global_httptextformat().inject(set_in_carrier, carrier, context)
 
 
-_HTTP_TEXT_FORMAT = (
-    TraceContextHTTPTextFormat()
+_HTTP_TEXT_FORMAT = composite.CompositeHTTPPropagator(
+    [TraceContextHTTPTextFormat(), CorrelationContextPropagator()],
 )  # type: httptextformat.HTTPTextFormat