Support for Jaeger Remote Sampling addition to contrib repo

opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

@@ -1246,6 +1246,8 @@ def add_span_processor(self, span_processor: SpanProcessor) -> None:
     def shutdown(self):
         """Shut down the span processors added to the tracer provider."""
         self._active_span_processor.shutdown()
+        if self.sampler is not None:
+            self.sampler.close()
         if self._atexit_handler is not None:
             atexit.unregister(self._atexit_handler)
             self._atexit_handler = None

opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py

@@ -136,7 +136,7 @@ class CustomSamplerFactory:
 import os
 from logging import getLogger
 from types import MappingProxyType
-from typing import Optional, Sequence
+from typing import Any, Optional, Sequence
 
 # pylint: disable=unused-import
 from opentelemetry.context import Context
@@ -150,6 +150,42 @@ class CustomSamplerFactory:
 
 _logger = getLogger(__name__)
 
+# TODO(sconover): not sure what to do w/ these constants + associated attrs, 
+# the don't seem like a "fit" in this codebase. However these are central to the workings
+# of the ported tests.
+SAMPLER_TYPE_TAG_KEY = 'sampler.type'
+SAMPLER_PARAM_TAG_KEY = 'sampler.param'
+
+SAMPLER_TYPE_RATE_LIMITING = 'ratelimiting'
+SAMPLER_TYPE_LOWER_BOUND = 'lowerbound'
+SAMPLER_TYPE_TRACE_ID_RATIO = 'traceidratio'
+SAMPLER_TYPE_PARENT_BASED_TRACE_ID_RATIO = 'parentbased_traceidratio'
+SAMPLER_TYPE_GUARANTEED_THROUGHPUT = 'guaranteedthroughput'
+SAMPLER_TYPE_ADAPTIVE = 'adaptive'
+SAMPLER_TYPE_REMOTE_CONTROLLED = 'remotecontrolled'
+
+# How often remotely controlled sampler polls for sampling strategy
+DEFAULT_SAMPLING_INTERVAL = 60
+
+
+DEFAULT_SAMPLING_PROBABILITY = 0.001
+DEFAULT_LOWER_BOUND = 1.0 / (10.0 * 60.0)  # sample once every 10 minutes
+DEFAULT_MAX_OPERATIONS = 2000
+
+STRATEGIES_STR = 'perOperationStrategies'
+OPERATION_STR = 'operation'
+DEFAULT_LOWER_BOUND_STR = 'defaultLowerBoundTracesPerSecond'
+PROBABILISTIC_SAMPLING_STR = 'probabilisticSampling'
+SAMPLING_RATE_STR = 'samplingRate'
+DEFAULT_SAMPLING_PROBABILITY_STR = 'defaultSamplingProbability'
+OPERATION_SAMPLING_STR = 'operationSampling'
+MAX_TRACES_PER_SECOND_STR = 'maxTracesPerSecond'
+RATE_LIMITING_SAMPLING_STR = 'rateLimitingSampling'
+STRATEGY_TYPE_STR = 'strategyType'
+PROBABILISTIC_SAMPLING_STRATEGY = 'PROBABILISTIC'
+RATE_LIMITING_SAMPLING_STRATEGY = 'RATE_LIMITING'
+
+
 
 class Decision(enum.Enum):
     # IsRecording() == false, span will not be recorded and all events and attributes will be dropped.
@@ -212,6 +248,27 @@ def should_sample(
     def get_description(self) -> str:
         pass
 
+    # TODO(sconover) added close to all samplers, because of cleanup needed
+    # for RemoteControlledSampler, in the contrib repo
+    # Q: Where should sampler.close() be called? I believe it might be 
+    # TracerProvider#shutdown (but not entirely sure)
+    def close(self) -> None:
+        pass
+
+    # TODO(sconover): the adaptive sampler in the contrib repo makes use of sampler equality
+    def __eq__(self, other: Any) -> bool:
+        return (
+            isinstance(other, self.__class__) and self.__dict__ == other.__dict__
+        )
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
+
+    # TODO(sconover): the jaeger tests (in the contrib repo) assert against the string value of a sampler,
+    # in order to detect what sampler is "in effect", e.g. in AdaptiveSampler
+    def __str__(self) -> str:
+        return self.get_description()
+
 
 class StaticSampler(Sampler):
     """Sampler that always returns the same decision."""
@@ -231,6 +288,7 @@ def should_sample(
     ) -> "SamplingResult":
         if self._decision is Decision.DROP:
             attributes = None
+
         return SamplingResult(
             self._decision,
             attributes,
@@ -263,6 +321,10 @@ def __init__(self, rate: float):
             raise ValueError("Probability must be in range [0.0, 1.0].")
         self._rate = rate
         self._bound = self.get_bound_for_rate(self._rate)
+        self._attributes = {
+            SAMPLER_TYPE_TAG_KEY: SAMPLER_TYPE_TRACE_ID_RATIO,
+            SAMPLER_PARAM_TAG_KEY: rate
+        }
 
     # For compatibility with 64 bit trace IDs, the sampler checks the 64
     # low-order bits of the trace ID to decide whether to sample a given trace.
@@ -295,6 +357,16 @@ def should_sample(
             decision = Decision.RECORD_AND_SAMPLE
         if decision is Decision.DROP:
             attributes = None
+
+        # TODO(sconover): the jaeger tests (in the contrib repo) really really want this probabilistic sampler 
+        # to indicate key elements of internal state via attributes
+        # I expect this might be a controversial issue in code review 
+        # (perhaps unacceptable aspect of the "port", especially since its existence service to support
+        # tests in an entirely different repo (opentelementry-python-contrib)
+        if attributes == None:
+            attributes = {}
+        attributes = {**self._attributes, **attributes}
+
         return SamplingResult(
             decision,
             attributes,
@@ -413,7 +485,7 @@ class _ParentBasedAlwaysOn(ParentBased):
     def __init__(self, _):
         super().__init__(ALWAYS_ON)
 
-
+# TODO(sconover): Something must be done to allow samplers from the contrib repo to be specified
 _KNOWN_SAMPLERS = {
     "always_on": ALWAYS_ON,
     "always_off": ALWAYS_OFF,

opentelemetry-sdk/tests/trace/test_sampling.py

@@ -220,7 +220,7 @@ def test_probability_sampler(self):
             attributes={"sampled.expect": "true"},
         )
         self.assertTrue(sampled_result.decision.is_sampled())
-        self.assertEqual(sampled_result.attributes, {"sampled.expect": "true"})
+        self.assertEqual(sampled_result.attributes, {"sampler.type": "traceidratio", "sampler.param": 0.5, "sampled.expect": "true"})
         self.assertIsNone(sampled_result.trace_state)
 
         not_sampled_result = sampler.should_sample(
@@ -231,7 +231,7 @@ def test_probability_sampler(self):
             attributes={"sampled.expect": "false"},
         )
         self.assertFalse(not_sampled_result.decision.is_sampled())
-        self.assertEqual(not_sampled_result.attributes, {})
+        self.assertEqual(not_sampled_result.attributes, {"sampler.type": "traceidratio", "sampler.param": 0.5})
         self.assertIsNone(sampled_result.trace_state)
 
     def test_probability_sampler_zero(self):
@@ -537,3 +537,17 @@ def implicit_parent_context(span: trace.Span):
             context_api.detach(token)
 
         self.exec_parent_based(implicit_parent_context)
+
+    def test_sampler_equality(self):
+        const1 = sampling.StaticSampler(True)
+        const2 = sampling.StaticSampler(True)
+        const3 = sampling.StaticSampler(False)
+        self.assertEqual(const1, const2)
+        self.assertNotEqual(const1, const3)
+
+        prob1 = sampling.TraceIdRatioBased(rate=0.01)
+        prob2 = sampling.TraceIdRatioBased(rate=0.01)
+        prob3 = sampling.TraceIdRatioBased(rate=0.02)
+        self.assertEqual(prob1, prob2)
+        self.assertNotEqual(prob1, prob3)
+        self.assertNotEqual(const1, prob1)

opentelemetry-sdk/tests/trace/test_trace.py

@@ -50,6 +50,7 @@
     Decision,
     ParentBased,
     StaticSampler,
+    Sampler,
 )
 from opentelemetry.sdk.util import BoundedDict, ns_to_iso_str
 from opentelemetry.sdk.util.instrumentation import InstrumentationInfo
@@ -83,7 +84,8 @@ def test_extends_api(self):
         self.assertIsInstance(tracer, trace_api.Tracer)
 
     def test_shutdown(self):
-        tracer_provider = trace.TracerProvider()
+        mock_sampler = Mock(spec=Sampler)
+        tracer_provider = trace.TracerProvider(sampler=mock_sampler)
 
         mock_processor1 = mock.Mock(spec=trace.SpanProcessor)
         tracer_provider.add_span_processor(mock_processor1)
@@ -96,6 +98,8 @@ def test_shutdown(self):
         self.assertEqual(mock_processor1.shutdown.call_count, 1)
         self.assertEqual(mock_processor2.shutdown.call_count, 1)
 
+        self.assertEqual(mock_sampler.close.call_count, 1)
+
         shutdown_python_code = """
 import atexit
 from unittest import mock