def test_basic_retry_policy_with_supression():
class ExpectedError(Exception):
pass
class UnexpectedError(Exception):
pass
def function():
function.call_count += 1
if function.call_count % 2 == 0:
raise UnexpectedError()
else:
raise ExpectedError()
function.call_count = 0
def suppression_test(exception: Exception) -> bool:
return isinstance(exception, ExpectedError)
clock = TestingClock()
policy = BasicRetryPolicy(3,
constant_delay(1),
suppression_test=suppression_test,
clock=clock)
try:
policy.call(function)
except UnexpectedError:
assert function.call_count == 2
assert clock.time() == 1
def test_basic_retry_policy_with_delay(delay):
value = object()
def good_function():
good_function.call_count += 1
return value
good_function.call_count = 0
class CustomException(Exception):
pass
def flaky_function():
flaky_function.call_count += 1
if flaky_function.call_count % 2 == 0:
return value
else:
raise CustomException()
flaky_function.call_count = 0
def bad_function():
bad_function.call_count += 1
raise CustomException()
bad_function.call_count = 0
clock = TestingClock()
policy = BasicRetryPolicy(3, delay, clock=clock)
assert policy.call(good_function) is value
assert good_function.call_count == 1
assert clock.time() == 0
clock = TestingClock()
policy = BasicRetryPolicy(3, delay, clock=clock)
assert policy.call(flaky_function) is value
assert flaky_function.call_count == 2
assert clock.time() == 1 # one retry
clock = TestingClock()
policy = BasicRetryPolicy(3, delay, clock=clock)
try:
policy.call(bad_function)
except Exception as e:
assert isinstance(e, RetryException)
assert isinstance(e.__cause__, CustomException)
assert bad_function.call_count == 3
assert clock.time() == 2 # two retries
def __init__(
self,
storage_key: StorageKey,
raw_topic: Optional[str],
replacements_topic: Optional[str],
max_batch_size: int,
max_batch_time_ms: int,
bootstrap_servers: Sequence[str],
group_id: str,
commit_log_topic: Optional[str],
auto_offset_reset: str,
queued_max_messages_kbytes: int,
queued_min_messages: int,
processes: Optional[int],
input_block_size: Optional[int],
output_block_size: Optional[int],
commit_retry_policy: Optional[RetryPolicy] = None,
profile_path: Optional[str] = None,
) -> None:
self.storage = get_writable_storage(storage_key)
self.bootstrap_servers = bootstrap_servers
self.broker_config = get_default_kafka_configuration(
storage_key, bootstrap_servers=bootstrap_servers
)
self.producer_broker_config = build_kafka_producer_configuration(
storage_key,
bootstrap_servers=bootstrap_servers,
override_params={
"partitioner": "consistent",
"message.max.bytes": 50000000, # 50MB, default is 1MB
},
)
stream_loader = self.storage.get_table_writer().get_stream_loader()
self.raw_topic: Topic
if raw_topic is not None:
self.raw_topic = Topic(raw_topic)
else:
self.raw_topic = Topic(stream_loader.get_default_topic_spec().topic_name)
self.replacements_topic: Optional[Topic]
if replacements_topic is not None:
self.replacements_topic = Topic(replacements_topic)
else:
replacement_topic_spec = stream_loader.get_replacement_topic_spec()
if replacement_topic_spec is not None:
self.replacements_topic = Topic(replacement_topic_spec.topic_name)
else:
self.replacements_topic = None
self.commit_log_topic: Optional[Topic]
if commit_log_topic is not None:
self.commit_log_topic = Topic(commit_log_topic)
else:
commit_log_topic_spec = stream_loader.get_commit_log_topic_spec()
if commit_log_topic_spec is not None:
self.commit_log_topic = Topic(commit_log_topic_spec.topic_name)
else:
self.commit_log_topic = None
# XXX: This can result in a producer being built in cases where it's
# not actually required.
self.producer = Producer(self.producer_broker_config)
self.metrics = MetricsWrapper(
environment.metrics,
"consumer",
tags={"group": group_id, "storage": storage_key.value},
)
self.max_batch_size = max_batch_size
self.max_batch_time_ms = max_batch_time_ms
self.group_id = group_id
self.auto_offset_reset = auto_offset_reset
self.queued_max_messages_kbytes = queued_max_messages_kbytes
self.queued_min_messages = queued_min_messages
self.processes = processes
self.input_block_size = input_block_size
self.output_block_size = output_block_size
self.__profile_path = profile_path
if commit_retry_policy is None:
commit_retry_policy = BasicRetryPolicy(
3,
constant_delay(1),
lambda e: isinstance(e, KafkaException)
and e.args[0].code()
in (
KafkaError.REQUEST_TIMED_OUT,
KafkaError.NOT_COORDINATOR,
KafkaError._WAIT_COORD,
),
)
self.__commit_retry_policy = commit_retry_policy
def __init__(
self,
dataset_name: str,
raw_topic: Optional[str],
replacements_topic: Optional[str],
max_batch_size: int,
max_batch_time_ms: int,
bootstrap_servers: Sequence[str],
group_id: str,
commit_log_topic: Optional[str],
auto_offset_reset: str,
queued_max_messages_kbytes: int,
queued_min_messages: int,
rapidjson_deserialize: bool,
rapidjson_serialize: bool,
commit_retry_policy: Optional[RetryPolicy] = None,
) -> None:
self.dataset = get_dataset(dataset_name)
self.dataset_name = dataset_name
if not bootstrap_servers:
self.bootstrap_servers = settings.DEFAULT_DATASET_BROKERS.get(
dataset_name, settings.DEFAULT_BROKERS,
)
else:
self.bootstrap_servers = bootstrap_servers
stream_loader = enforce_table_writer(self.dataset).get_stream_loader()
self.raw_topic: Topic
if raw_topic is not None:
self.raw_topic = Topic(raw_topic)
else:
self.raw_topic = Topic(stream_loader.get_default_topic_spec().topic_name)
self.replacements_topic: Optional[Topic]
if replacements_topic is not None:
self.replacements_topic = Topic(replacements_topic)
else:
replacement_topic_spec = stream_loader.get_replacement_topic_spec()
if replacement_topic_spec is not None:
self.replacements_topic = Topic(replacement_topic_spec.topic_name)
else:
self.replacements_topic = None
self.commit_log_topic: Optional[Topic]
if commit_log_topic is not None:
self.commit_log_topic = Topic(commit_log_topic)
else:
commit_log_topic_spec = stream_loader.get_commit_log_topic_spec()
if commit_log_topic_spec is not None:
self.commit_log_topic = Topic(commit_log_topic_spec.topic_name)
else:
self.commit_log_topic = None
# XXX: This can result in a producer being built in cases where it's
# not actually required.
self.producer = Producer(
{
"bootstrap.servers": ",".join(self.bootstrap_servers),
"partitioner": "consistent",
"message.max.bytes": 50000000, # 50MB, default is 1MB
}
)
self.metrics = util.create_metrics(
"snuba.consumer", tags={"group": group_id, "dataset": self.dataset_name},
)
self.max_batch_size = max_batch_size
self.max_batch_time_ms = max_batch_time_ms
self.group_id = group_id
self.auto_offset_reset = auto_offset_reset
self.queued_max_messages_kbytes = queued_max_messages_kbytes
self.queued_min_messages = queued_min_messages
if commit_retry_policy is None:
commit_retry_policy = BasicRetryPolicy(
3,
constant_delay(1),
lambda e: isinstance(e, KafkaException)
and e.args[0].code()
in (
KafkaError.REQUEST_TIMED_OUT,
KafkaError.NOT_COORDINATOR_FOR_GROUP,
KafkaError._WAIT_COORD,
),
)
self.__commit_retry_policy = commit_retry_policy
self.__rapidjson_deserialize = rapidjson_deserialize
self.__rapidjson_serialize = rapidjson_serialize