def test_datadog_sampler_tracer_rate_limited(dummy_tracer):
rule = SamplingRule(sample_rate=1.0, name='test.span')
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
limiter_spy.is_allowed.return_value = False # Have the limiter deny the span
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
dummy_tracer.configure(sampler=sampler_spy)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace('test.span') as span:
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_called_once_with()
# We know it was not sampled because of our limiter
assert span.sampled is False
assert span._context.sampling_priority is AUTO_REJECT
assert_sampling_decision_tags(span, rule=1.0, limit=None)
def test_datadog_sampler_tracer(dummy_tracer):
rule = SamplingRule(sample_rate=1.0, name='test.span')
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
dummy_tracer.configure(sampler=sampler_spy)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace('test.span') as span:
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_called_once_with()
# We know it was sampled because we have a sample rate of 1.0
assert span.sampled is True
assert span._context.sampling_priority is AUTO_KEEP
assert_sampling_decision_tags(span, rule=1.0)
def test_datadog_sampler_tracer_rate_0(dummy_tracer):
rule = SamplingRule(
sample_rate=0, name='test.span') # Sample rate of 0 means never sample
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace('test.span') as span:
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_not_called()
# We know it was not sampled because we have a sample rate of 0.0
assert span.sampled is False
assert span._context.sampling_priority is AUTO_REJECT
assert_sampling_decision_tags(span, rule=0)
def test_datadog_sampler_sample_no_rules(mock_sample, dummy_tracer):
sampler = DatadogSampler()
span = create_span(tracer=dummy_tracer)
# Default RateByServiceSampler() is applied
# No rules configured
# No global rate limit
# No rate limit configured
# RateByServiceSampler.sample(span) returns True
mock_sample.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
span = create_span(tracer=dummy_tracer)
# Default RateByServiceSampler() is applied
# No rules configured
# No global rate limit
# No rate limit configured
# RateByServiceSampler.sample(span) returns False
mock_sample.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
def test_datadog_sampler_tracer_child(dummy_tracer):
rule = SamplingRule(
sample_rate=1.0) # No rules means it gets applied to every span
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
dummy_tracer.configure(sampler=sampler_spy)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace("parent.span") as parent:
with dummy_tracer.trace("child.span"):
# Assert all of our expected functions were called
# DEV: `assert_called_once_with` ensures we didn't also call with the child span
sampler_spy.sample.assert_called_once_with(parent)
rule_spy.matches.assert_called_once_with(parent)
rule_spy.sample.assert_called_once_with(parent)
limiter_spy.is_allowed.assert_called_once_with()
spans = dummy_tracer.pop()
assert len(spans) == 2, "Trace should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is AUTO_KEEP
assert spans[0].get_metric(SAMPLING_RULE_DECISION) == 1.0
def test_datadog_sampler_tracer_rate_0(dummy_tracer):
rule = SamplingRule(
sample_rate=0, name="test.span") # Sample rate of 0 means never sample
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
dummy_tracer.configure(sampler=sampler_spy)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace("test.span") as span:
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_not_called()
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is AUTO_REJECT
assert spans[0].get_metric(SAMPLING_RULE_DECISION) == 0
def test_datadog_sampler_tracer_child(dummy_tracer):
rule = SamplingRule(
sample_rate=1.0) # No rules means it gets applied to every span
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
assert dummy_tracer.sampler is sampler_spy
with dummy_tracer.trace('parent.span') as parent:
with dummy_tracer.trace('child.span') as child:
# Assert all of our expected functions were called
# DEV: `assert_called_once_with` ensures we didn't also call with the child span
sampler_spy.sample.assert_called_once_with(parent)
rule_spy.matches.assert_called_once_with(parent)
rule_spy.sample.assert_called_once_with(parent)
limiter_spy.is_allowed.assert_called_once_with()
# We know it was sampled because we have a sample rate of 1.0
assert parent.sampled is True
assert parent._context.sampling_priority is AUTO_KEEP
assert_sampling_decision_tags(parent, rule=1.0)
assert child.sampled is True
assert child._parent is parent
assert child._context.sampling_priority is AUTO_KEEP
def test_datadog_sampler_sample_no_rules(mock_is_allowed, dummy_tracer):
sampler = DatadogSampler()
span = create_span(tracer=dummy_tracer)
# Default SamplingRule(sample_rate=1.0) is applied
# No priority sampler configured
# No rules configured
# RateLimiter is allowed, it is sampled
mock_is_allowed.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
assert_sampling_decision_tags(span, rule=1.0, limit=1.0)
mock_is_allowed.assert_called_once_with()
mock_is_allowed.reset_mock()
span = create_span(tracer=dummy_tracer)
# Default SamplingRule(sample_rate=1.0) is applied
# No priority sampler configured
# No rules configured
# RateLimit not allowed, it is not sampled
mock_is_allowed.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
# DEV: Is `None` since we only add tag to non-rate limited traces
assert_sampling_decision_tags(span, rule=1.0, limit=None)
mock_is_allowed.assert_called_once_with()
def test_datadog_sampler_tracer_start_span(dummy_tracer):
rule = SamplingRule(
sample_rate=1.0) # No rules means it gets applied to every span
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
dummy_tracer.configure(sampler=sampler_spy)
assert dummy_tracer.sampler is sampler_spy
span = dummy_tracer.start_span("test.span")
span.finish()
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_called_once_with()
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is AUTO_KEEP
assert spans[0].get_metric(SAMPLING_RULE_DECISION) == 1.0
def test_datadog_sampler_init():
# No args
sampler = DatadogSampler()
assert sampler.rules == []
assert isinstance(sampler.limiter, RateLimiter)
assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
# With rules
rule = SamplingRule(sample_rate=1)
sampler = DatadogSampler(rules=[rule])
assert sampler.rules == [rule]
assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
# With rate limit
sampler = DatadogSampler(rate_limit=10)
assert sampler.limiter.rate_limit == 10
# Invalid rules
for val in (None, True, False, object(), 1, Exception()):
with pytest.raises(TypeError):
DatadogSampler(rules=[val])
# Ensure rule order
rule_1 = SamplingRule(sample_rate=1)
rule_2 = SamplingRule(sample_rate=0.5, service='test')
rule_3 = SamplingRule(sample_rate=0.25, name='flask.request')
sampler = DatadogSampler(rules=[rule_1, rule_2, rule_3])
assert sampler.rules == [rule_1, rule_2, rule_3]
def test_datadog_sampler_tracer_start_span(dummy_tracer):
rule = SamplingRule(
sample_rate=1.0) # No rules means it gets applied to every span
rule_spy = mock.Mock(spec=rule, wraps=rule)
rule_spy.sample_rate = rule.sample_rate
sampler = DatadogSampler(rules=[rule_spy])
limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
sampler.limiter = limiter_spy
sampler_spy = mock.Mock(spec=sampler, wraps=sampler)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
assert dummy_tracer.sampler is sampler_spy
span = dummy_tracer.start_span('test.span')
# Assert all of our expected functions were called
sampler_spy.sample.assert_called_once_with(span)
rule_spy.matches.assert_called_once_with(span)
rule_spy.sample.assert_called_once_with(span)
limiter_spy.is_allowed.assert_called_once_with()
# We know it was sampled because we have a sample rate of 1.0
assert span.sampled is True
assert span._context.sampling_priority is AUTO_KEEP
assert_sampling_decision_tags(span, rule=1.0)
def test_datadog_sampler_init():
# No args
sampler = DatadogSampler()
assert sampler.rules == []
assert isinstance(sampler.limiter, RateLimiter)
assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
assert isinstance(sampler.default_sampler, RateByServiceSampler)
# With rules
rule = SamplingRule(sample_rate=1)
sampler = DatadogSampler(rules=[rule])
assert sampler.rules == [rule]
assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
assert isinstance(sampler.default_sampler, RateByServiceSampler)
# With rate limit
sampler = DatadogSampler(rate_limit=10)
assert sampler.limiter.rate_limit == 10
assert isinstance(sampler.default_sampler, RateByServiceSampler)
# With default_sample_rate
sampler = DatadogSampler(default_sample_rate=0.5)
assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
assert isinstance(sampler.default_sampler, SamplingRule)
assert sampler.default_sampler.sample_rate == 0.5
# From env variables
with override_env(
dict(DD_TRACE_SAMPLE_RATE="0.5", DD_TRACE_RATE_LIMIT="10")):
sampler = DatadogSampler()
assert sampler.limiter.rate_limit == 10
assert isinstance(sampler.default_sampler, SamplingRule)
assert sampler.default_sampler.sample_rate == 0.5
# Invalid env vars
with override_env(dict(DD_TRACE_SAMPLE_RATE="asdf")):
with pytest.raises(ValueError):
DatadogSampler()
# Invalid env vars
with override_env(dict(DD_TRACE_RATE_LIMIT="asdf")):
with pytest.raises(ValueError):
DatadogSampler()
# Invalid rules
for val in (None, True, False, object(), 1, Exception()):
with pytest.raises(TypeError):
DatadogSampler(rules=[val])
# Ensure rule order
rule_1 = SamplingRule(sample_rate=1)
rule_2 = SamplingRule(sample_rate=0.5, service="test")
rule_3 = SamplingRule(sample_rate=0.25, name="flask.request")
sampler = DatadogSampler(rules=[rule_1, rule_2, rule_3])
assert sampler.rules == [rule_1, rule_2, rule_3]
def test_datadog_sampler_sample_no_rules(mock_sample, dummy_tracer):
sampler = DatadogSampler()
dummy_tracer.configure(sampler=sampler)
# Default RateByServiceSampler() is applied
# No rules configured
# No global rate limit
# No rate limit configured
# RateByServiceSampler.sample(span) returns True
mock_sample.return_value = True
with dummy_tracer.trace("test"):
pass
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is AUTO_KEEP
# Default RateByServiceSampler() is applied
# No rules configured
# No global rate limit
# No rate limit configured
# RateByServiceSampler.sample(span) returns False
mock_sample.return_value = False
with dummy_tracer.trace("test"):
pass
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is AUTO_REJECT
def test_datadog_sampler_tracer_rate_limited(dummy_tracer):
rule = SamplingRule(sample_rate=1.0, name="test.span")
sampler = DatadogSampler(rules=[rule], rate_limit=0)
dummy_tracer.configure(sampler=sampler)
with dummy_tracer.trace("test.span"):
pass
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is USER_REJECT
assert_sampling_decision_tags(spans[0], rule=1.0, limit=0.0)
def test_datadog_sampler_tracer_start_span(dummy_tracer):
# No rules means it gets applied to every span
rule = SamplingRule(sample_rate=1.0)
sampler = DatadogSampler(rules=[rule])
dummy_tracer.configure(sampler=sampler)
span = dummy_tracer.start_span("test.span")
span.finish()
spans = dummy_tracer.pop()
assert len(spans) == 1, "Span should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is USER_KEEP
assert_sampling_decision_tags(spans[0], rule=1.0, limit=None)
def test_datadog_sampler_tracer_child(dummy_tracer):
# No rules means it gets applied to every span
rule = SamplingRule(sample_rate=1.0)
sampler = DatadogSampler(rules=[rule])
dummy_tracer.configure(sampler=sampler)
with dummy_tracer.trace("parent.span"):
with dummy_tracer.trace("child.span"):
pass
spans = dummy_tracer.pop()
assert len(spans) == 2, "Trace should have been sampled and written"
assert spans[0].get_metric(SAMPLING_PRIORITY_KEY) is USER_KEEP
assert_sampling_decision_tags(spans[0], rule=1.0, limit=None)
assert_sampling_decision_tags(spans[1], agent=None, rule=None, limit=None)
def test_sampling(self, tracer):
with tracer.trace("trace1"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1.0)
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace2"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=0.000001)
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace3"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1,
rules=[SamplingRule(1.0)])
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace4"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1,
rules=[SamplingRule(0)])
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace5"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1)
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace6"):
with tracer.trace("child") as span:
span.set_tag(MANUAL_DROP_KEY)
sampler = DatadogSampler(default_sample_rate=1)
tracer.configure(sampler=sampler, writer=tracer.writer)
with tracer.trace("trace7"):
with tracer.trace("child") as span:
span.set_tag(MANUAL_KEEP_KEY)
sampler = RateSampler(0.0000000001)
tracer.configure(sampler=sampler, writer=tracer.writer)
# This trace should not appear in the snapshot
with tracer.trace("trace8"):
with tracer.trace("child"):
pass
tracer.shutdown()
def test_datadog_sampler_update_rate_by_service_sample_rates(dummy_tracer):
cases = [
{
"service:,env:": 1,
},
{
"service:,env:": 1,
"service:mcnulty,env:dev": 0.33,
"service:postgres,env:dev": 0.7,
},
{
"service:,env:": 1,
"service:mcnulty,env:dev": 0.25,
"service:postgres,env:dev": 0.5,
"service:redis,env:prod": 0.75,
},
]
# By default sampler sets it's default sampler to RateByServiceSampler
sampler = DatadogSampler()
for case in cases:
sampler.update_rate_by_service_sample_rates(case)
rates = {}
for k, v in iteritems(sampler.default_sampler._by_service_samplers):
rates[k] = v.sample_rate
assert case == rates, "%s != %s" % (case, rates)
# It's important to also test in reverse mode for we want to make sure key deletion
# works as well as key insertion (and doing this both ways ensures we trigger both cases)
cases.reverse()
for case in cases:
sampler.update_rate_by_service_sample_rates(case)
rates = {}
for k, v in iteritems(sampler.default_sampler._by_service_samplers):
rates[k] = v.sample_rate
assert case == rates, "%s != %s" % (case, rates)
def test_sampling_rule_init_via_env():
# Testing single sampling rule
with override_env(
dict(DD_TRACE_SAMPLING_RULES=
'[{"sample_rate":1.0,"service":"xyz","name":"abc"}]')):
sampling_rule = DatadogSampler().rules
assert sampling_rule[0].sample_rate == 1.0
assert sampling_rule[0].service == "xyz"
assert sampling_rule[0].name == "abc"
assert len(sampling_rule) == 1
# Testing multiple sampling rules
with override_env(
dict(DD_TRACE_SAMPLING_RULES=
'[{"sample_rate":1.0,"service":"xyz","name":"abc"}, \
{"sample_rate":0.5,"service":"my-service","name":"my-name"}]')):
sampling_rule = DatadogSampler().rules
assert sampling_rule[0].sample_rate == 1.0
assert sampling_rule[0].service == "xyz"
assert sampling_rule[0].name == "abc"
assert sampling_rule[1].sample_rate == 0.5
assert sampling_rule[1].service == "my-service"
assert sampling_rule[1].name == "my-name"
assert len(sampling_rule) == 2
# Testing for only Sample rate being set
with override_env(dict(DD_TRACE_SAMPLING_RULES='[{"sample_rate":1.0}]')):
sampling_rule = DatadogSampler().rules
assert sampling_rule[0].sample_rate == 1.0
assert sampling_rule[0].service == SamplingRule.NO_RULE
assert sampling_rule[0].name == SamplingRule.NO_RULE
assert len(sampling_rule) == 1
# Testing for no name being set
with override_env(
dict(
DD_TRACE_SAMPLING_RULES='[{"sample_rate":1.0,"service":"xyz"}]'
)):
sampling_rule = DatadogSampler().rules
assert sampling_rule[0].sample_rate == 1.0
assert sampling_rule[0].service == "xyz"
assert sampling_rule[0].name == SamplingRule.NO_RULE
assert len(sampling_rule) == 1
# Testing for no service being set
with override_env(
dict(
DD_TRACE_SAMPLING_RULES='[{"sample_rate":1.0,"name":"abc"}]')):
sampling_rule = DatadogSampler().rules
assert sampling_rule[0].sample_rate == 1.0
assert sampling_rule[0].service == SamplingRule.NO_RULE
assert sampling_rule[0].name == "abc"
assert len(sampling_rule) == 1
# The Following error handling test use assertions on the json items instead of asserting on
# the returned stringdue to older version of python not keeping load order in dictionaires
# Testing for Sample rate greater than 1.0
with pytest.raises(ValueError) as excinfo:
with override_env(
dict(DD_TRACE_SAMPLING_RULES=
'[{"sample_rate":2.0,"service":"xyz","name":"abc"}]')):
sampling_rule = DatadogSampler().rules
assert str(excinfo.value).endswith(
"SamplingRule(sample_rate=2.0) must be greater than or equal to 0.0 and less than or equal to 1.0"
)
assert '"sample_rate": 2.0' in str(excinfo.value)
assert '"service": "xyz"' in str(excinfo.value)
assert '"name": "abc"' in str(excinfo.value)
# Testing for no Sample rate
with pytest.raises(KeyError) as excinfo:
with override_env(
dict(DD_TRACE_SAMPLING_RULES='[{"service":"xyz","name":"abc"}]'
)):
sampling_rule = DatadogSampler().rules
assert str(excinfo.value).startswith(
"'No sample_rate provided for sampling rule: ")
assert '"service": "xyz"' in str(excinfo.value)
assert '"name": "abc"' in str(excinfo.value)
# Testing for Invalid JSON
with pytest.raises(ValueError) as excinfo:
with override_env(
dict(DD_TRACE_SAMPLING_RULES=
'["sample_rate":1.0,"service":"xyz","name":"abc"]')):
sampling_rule = DatadogSampler().rules
assert 'Unable to parse DD_TRACE_SAMPLING_RULES=["sample_rate":1.0,"service":"xyz","name":"abc"]' == str(
excinfo.value)
# # Testing invalid rule with multiple rules defined
with pytest.raises(KeyError) as excinfo:
with override_env(
dict(DD_TRACE_SAMPLING_RULES=
'[{"sample_rate":1.0,"service":"xyz","name":"abc"},' +
'{"service":"my-service","name":"my-name"}]')):
sampling_rule = DatadogSampler().rules
assert str(excinfo.value).startswith(
"'No sample_rate provided for sampling rule: ")
assert '"service": "my-service"' in str(excinfo.value)
assert '"name": "my-name"' in str(excinfo.value)
def test_datadog_sampler_sample_rules(mock_is_allowed, dummy_tracer):
# Do not let the limiter get in the way of our test
mock_is_allowed.return_value = True
rules = [
mock.Mock(spec=SamplingRule),
mock.Mock(spec=SamplingRule),
mock.Mock(spec=SamplingRule),
]
sampler = DatadogSampler(rules=rules)
sampler.default_sampler = mock.Mock(spec=SamplingRule)
sampler.default_sampler.return_value = True
# Reset all of our mocks
@contextlib.contextmanager
def reset_mocks():
def reset():
mock_is_allowed.reset_mock()
for rule in rules:
rule.reset_mock()
rule.sample_rate = 0.5
sampler.default_sampler.reset_mock()
sampler.default_sampler.sample_rate = 1.0
reset() # Reset before, just in case
try:
yield
finally:
reset() # Must reset after
# No rules want to sample
# It is allowed because of default rate sampler
# All rules SamplingRule.matches are called
# No calls to SamplingRule.sample happen
with reset_mocks():
span = create_span(tracer=dummy_tracer)
for rule in rules:
rule.matches.return_value = False
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
for rule in rules:
rule.matches.assert_called_once_with(span)
rule.sample.assert_not_called()
sampler.default_sampler.matches.assert_not_called()
sampler.default_sampler.sample.assert_called_once_with(span)
assert_sampling_decision_tags(span, rule=1.0, limit=1.0)
# One rule thinks it should be sampled
# All following rule's SamplingRule.matches are not called
# It goes through limiter
# It is allowed
with reset_mocks():
span = create_span(tracer=dummy_tracer)
rules[1].matches.return_value = True
rules[1].sample.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5, limit=1.0)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_not_called()
rules[1].matches.assert_called_once_with(span)
rules[1].sample.assert_called_once_with(span)
rules[2].matches.assert_not_called()
rules[2].sample.assert_not_called()
# All rules think it should be sampled
# The first rule's SamplingRule.matches is called
# It goes through limiter
# It is allowed
with reset_mocks():
span = create_span(tracer=dummy_tracer)
for rule in rules:
rule.matches.return_value = True
rules[0].sample.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5, limit=1.0)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_called_once_with(span)
for rule in rules[1:]:
rule.matches.assert_not_called()
rule.sample.assert_not_called()
# Rule matches but does not think it should be sampled
# The rule's SamplingRule.matches is called
# The rule's SamplingRule.sample is called
# Rate limiter is not called
# The span is rejected
with reset_mocks():
span = create_span(tracer=dummy_tracer)
rules[0].matches.return_value = False
rules[2].matches.return_value = False
rules[1].matches.return_value = True
rules[1].sample.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
mock_is_allowed.assert_not_called()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_not_called()
rules[1].matches.assert_called_once_with(span)
rules[1].sample.assert_called_once_with(span)
rules[2].matches.assert_not_called()
rules[2].sample.assert_not_called()
# No rules match and priority sampler is defined
# All rules SamplingRule.matches are called
# Priority sampler's `sample` method is called
# Result of priority sampler is returned
# Rate limiter is not called
# TODO: Remove this case when we remove fallback to priority sampling
with reset_mocks():
span = create_span(tracer=dummy_tracer)
# Configure mock priority sampler
priority_sampler = RateByServiceSampler()
for rate_sampler in priority_sampler._by_service_samplers.values():
rate_sampler.set_sample_rate(1)
spy_sampler = mock.Mock(spec=RateByServiceSampler,
wraps=priority_sampler)
sampler._priority_sampler = spy_sampler
for rule in rules:
rule.matches.return_value = False
rule.sample.return_value = False
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_not_called()
sampler.default_sampler.sample.assert_not_called()
spy_sampler.sample.assert_called_once_with(span)
assert_sampling_decision_tags(span, agent=1)
[r.matches.assert_called_once_with(span) for r in rules]
[r.sample.assert_not_called() for r in rules]
# Reset priority sampler property
sampler._priority_sampler = None
# No rules match and priority sampler is defined
# All rules SamplingRule.matches are called
# Priority sampler's `sample` method is called
# Result of priority sampler is returned
# Rate limiter is not called
# TODO: Remove this case when we remove fallback to priority sampling
with reset_mocks():
span = create_span(tracer=dummy_tracer)
# Configure mock priority sampler
priority_sampler = RateByServiceSampler()
for rate_sampler in priority_sampler._by_service_samplers.values():
rate_sampler.set_sample_rate(0)
spy_sampler = mock.Mock(spec=RateByServiceSampler,
wraps=priority_sampler)
sampler._priority_sampler = spy_sampler
for rule in rules:
rule.matches.return_value = False
rule.sample.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
mock_is_allowed.assert_not_called()
sampler.default_sampler.sample.assert_not_called()
spy_sampler.sample.assert_called_once_with(span)
assert_sampling_decision_tags(span, agent=0)
[r.matches.assert_called_once_with(span) for r in rules]
[r.sample.assert_not_called() for r in rules]
# Reset priority sampler property
sampler._priority_sampler = None
def test_datadog_sampler_sample_rules(mock_is_allowed, dummy_tracer):
# Do not let the limiter get in the way of our test
mock_is_allowed.return_value = True
rules = [
mock.Mock(spec=SamplingRule),
mock.Mock(spec=SamplingRule),
mock.Mock(spec=SamplingRule),
]
sampler = DatadogSampler(rules=rules)
sampler.default_sampler = mock.Mock(spec=SamplingRule)
sampler.default_sampler.return_value = True
# Reset all of our mocks
@contextlib.contextmanager
def reset_mocks():
def reset():
mock_is_allowed.reset_mock()
for rule in rules:
rule.reset_mock()
rule.sample_rate = 0.5
sampler.default_sampler.reset_mock()
sampler.default_sampler.sample_rate = 1.0
reset() # Reset before, just in case
try:
yield
finally:
reset() # Must reset after
# No rules want to sample
# It is allowed because of default rate sampler
# All rules SamplingRule.matches are called
# No calls to SamplingRule.sample happen
with reset_mocks():
span = create_span(tracer=dummy_tracer)
for rule in rules:
rule.matches.return_value = False
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
for rule in rules:
rule.matches.assert_called_once_with(span)
rule.sample.assert_not_called()
sampler.default_sampler.matches.assert_not_called()
sampler.default_sampler.sample.assert_called_once_with(span)
assert_sampling_decision_tags(span, rule=1.0, limit=1.0)
# One rule thinks it should be sampled
# All following rule's SamplingRule.matches are not called
# It goes through limiter
# It is allowed
with reset_mocks():
span = create_span(tracer=dummy_tracer)
rules[1].matches.return_value = True
rules[1].sample.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5, limit=1.0)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_not_called()
rules[1].matches.assert_called_once_with(span)
rules[1].sample.assert_called_once_with(span)
rules[2].matches.assert_not_called()
rules[2].sample.assert_not_called()
# All rules think it should be sampled
# The first rule's SamplingRule.matches is called
# It goes through limiter
# It is allowed
with reset_mocks():
span = create_span(tracer=dummy_tracer)
for rule in rules:
rule.matches.return_value = True
rules[0].sample.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
mock_is_allowed.assert_called_once_with()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5, limit=1.0)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_called_once_with(span)
for rule in rules[1:]:
rule.matches.assert_not_called()
rule.sample.assert_not_called()
# Rule matches but does not think it should be sampled
# The rule's SamplingRule.matches is called
# The rule's SamplingRule.sample is called
# Rate limiter is not called
# The span is rejected
with reset_mocks():
span = create_span(tracer=dummy_tracer)
rules[0].matches.return_value = False
rules[2].matches.return_value = False
rules[1].matches.return_value = True
rules[1].sample.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
mock_is_allowed.assert_not_called()
sampler.default_sampler.sample.assert_not_called()
assert_sampling_decision_tags(span, rule=0.5)
rules[0].matches.assert_called_once_with(span)
rules[0].sample.assert_not_called()
rules[1].matches.assert_called_once_with(span)
rules[1].sample.assert_called_once_with(span)
rules[2].matches.assert_not_called()
rules[2].sample.assert_not_called()
def test_sampling(writer, tracer):
if writer == "sync":
writer = AgentWriter(
tracer.writer.agent_url,
priority_sampler=tracer.priority_sampler,
sync_mode=True,
)
# Need to copy the headers which contain the test token to associate
# traces with this test case.
writer._headers = tracer.writer._headers
else:
writer = tracer.writer
tracer.configure(writer=writer)
with tracer.trace("trace1"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1.0)
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace2"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=0.000001)
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace3"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1, rules=[SamplingRule(1.0)])
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace4"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1, rules=[SamplingRule(0)])
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace5"):
with tracer.trace("child"):
pass
sampler = DatadogSampler(default_sample_rate=1)
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace6"):
with tracer.trace("child") as span:
span.set_tag(MANUAL_DROP_KEY)
sampler = DatadogSampler(default_sample_rate=1)
tracer.configure(sampler=sampler, writer=writer)
with tracer.trace("trace7"):
with tracer.trace("child") as span:
span.set_tag(MANUAL_KEEP_KEY)
sampler = RateSampler(0.0000000001)
tracer.configure(sampler=sampler, writer=writer)
# This trace should not appear in the snapshot
with tracer.trace("trace8"):
with tracer.trace("child"):
pass
tracer.shutdown()
class MatchNoSample(SamplingRule):
def matches(self, span):
return True
def sample(self, span):
return False
@pytest.mark.parametrize(
"sampler, sampling_priority, rule, limit",
[
(
DatadogSampler(
default_sample_rate=1.0,
rules=[
NoMatch(0.5),
NoMatch(0.5),
NoMatch(0.5),
],
),
AUTO_KEEP,
1.0,
1.0,
),
(
DatadogSampler(
default_sample_rate=1.0,
rules=[
NoMatch(0.5),
NoMatch(0.5),
MatchSample(0.5),
],