Add Throttler API (#330)

* Add Throttler API

Signed-off-by: Yuri Shkuro <[email protected]>

* fix

Signed-off-by: Yuri Shkuro <[email protected]>

Author: yurishkuro

jaeger_client/throttler.py

@@ -28,7 +28,21 @@
 default_logger = logging.getLogger('jaeger_tracing')
 
 
-class RemoteThrottler(object):
+class Throttler(object):
+    def set_client_id(self, client_id: int) -> None:
+        """
+        Called by tracer to set client ID of throttler.
+        """
+        pass
+
+    def is_allowed(self, operation: str) -> bool:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        pass
+
+
+class RemoteThrottler(Throttler):
     """
     RemoteThrottler controls the flow of spans emitted from client to prevent
     flooding. RemoteThrottler requests credits from the throttling service
@@ -78,10 +92,7 @@ def is_allowed(self, operation):
             self.credits[operation] = value - MINIMUM_CREDITS
             return True
 
-    def _set_client_id(self, client_id):
-        """
-        Method for tracer to set client ID of throttler.
-        """
+    def set_client_id(self, client_id):
         with self.lock:
             if self.client_id is None:
                 self.client_id = client_id

jaeger_client/tracer.py

@@ -39,6 +39,7 @@
 from .utils import local_ip
 from .sampler import Sampler
 from .reporter import BaseReporter
+from .throttler import Throttler
 
 logger = logging.getLogger('jaeger_tracing')
 
@@ -63,7 +64,7 @@ def __init__(
         tags: Optional[Dict[str, Any]] = None,
         max_tag_value_length: int = constants.MAX_TAG_VALUE_LENGTH,
         max_traceback_length: int = constants.MAX_TRACEBACK_LENGTH,
-        throttler: Optional[Any] = None,
+        throttler: Optional[Throttler] = None,
         scope_manager: Optional[ScopeManager] = None,
     ) -> None:
         self.service_name = service_name
@@ -115,7 +116,7 @@ def __init__(
         self.throttler = throttler
         if self.throttler:
             client_id = random.randint(0, sys.maxsize)
-            self.throttler._set_client_id(client_id)
+            self.throttler.set_client_id(client_id)
             self.tags[constants.CLIENT_UUID_TAG_KEY] = client_id
 
         self.reporter.set_process(

tests/test_config.py

@@ -24,7 +24,7 @@
 from jaeger_client.config import DEFAULT_THROTTLER_PORT
 from jaeger_client.metrics import MetricsFactory
 from jaeger_client.reporter import NullReporter
-from tests.test_utils import TestSampler
+from tests.test_utils import MockSampler
 
 
 class ConfigTests(unittest.TestCase):
@@ -85,7 +85,7 @@ def test_bad_sampler(self):
             c.sampler.is_sampled(0)
 
     def test_object_sampler_sampler(self):
-        sampler = TestSampler()
+        sampler = MockSampler()
         c = Config({'sampler': sampler}, service_name='x')
         assert c.sampler is sampler
 

tests/test_utils.py

@@ -84,5 +84,6 @@ def test_get_local_ip_by_socket_does_not_blow_up():
     jaeger_client.utils.get_local_ip_by_socket()
 
 
-class TestSampler(Sampler):
-    pass
+class MockSampler(Sampler):
+    def __init__(self):
+        pass