feat(slices): Make consumers "slice-aware"

snuba/cli/consumer.py

@@ -50,6 +50,12 @@
     help="The storage to target",
     required=True,
 )
+@click.option(
+    "--slice-id",
+    "slice_id",
+    type=int,
+    help="The slice id for the storage",
+)
 @click.option(
     "--max-batch-size",
     default=settings.DEFAULT_MAX_BATCH_SIZE,
@@ -121,6 +127,7 @@ def consumer(
     consumer_group: str,
     bootstrap_server: Sequence[str],
     storage_name: str,
+    slice_id: Optional[int],
     max_batch_size: int,
     max_batch_time_ms: int,
     auto_offset_reset: str,
@@ -176,10 +183,11 @@ def stats_callback(stats_json: str) -> None:
         profile_path=profile_path,
         stats_callback=stats_callback,
         parallel_collect=parallel_collect,
+        slice_id=slice_id,
         cooperative_rebalancing=cooperative_rebalancing,
     )
 
-    consumer = consumer_builder.build_base_consumer()
+    consumer = consumer_builder.build_base_consumer(slice_id)
 
     def handler(signum: int, frame: Any) -> None:
         consumer.signal_shutdown()

snuba/consumers/consumer.py

@@ -334,6 +334,7 @@ def build_batch_writer(
     metrics: MetricsBackend,
     replacements_producer: Optional[ConfluentKafkaProducer] = None,
     replacements_topic: Optional[Topic] = None,
+    slice_id: Optional[int] = None,
 ) -> Callable[[], ProcessedMessageBatchWriter]:
 
     assert not (replacements_producer is None) ^ (replacements_topic is None)
@@ -342,6 +343,7 @@ def build_batch_writer(
     writer = table_writer.get_batch_writer(
         metrics,
         {"load_balancing": "in_order", "insert_distributed_sync": 1},
+        slice_id=slice_id,
     )
 
     def build_writer() -> ProcessedMessageBatchWriter:

snuba/consumers/consumer_builder.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from typing import Callable, Optional, Sequence
 
-from arroyo import Topic
+from arroyo import Topic as ArroyoTopic
 from arroyo.backends.kafka import KafkaConsumer, KafkaPayload
 from arroyo.commit import IMMEDIATE
 from arroyo.processing import StreamProcessor
@@ -18,6 +18,7 @@
     build_mock_batch_writer,
     process_message,
 )
+from snuba.datasets.partitioning import validate_passed_slice
 from snuba.datasets.storages.factory import get_writable_storage
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.environment import setup_sentry
@@ -77,6 +78,7 @@ def __init__(
         max_batch_time_ms: int,
         metrics: MetricsBackend,
         parallel_collect: bool,
+        slice_id: Optional[int],
         stats_callback: Optional[Callable[[str], None]] = None,
         commit_retry_policy: Optional[RetryPolicy] = None,
         profile_path: Optional[str] = None,
@@ -93,12 +95,15 @@ def __init__(
             .topic
         )
 
+        validate_passed_slice(storage_key, slice_id)
+
         self.broker_config = get_default_kafka_configuration(
-            topic, bootstrap_servers=kafka_params.bootstrap_servers
+            topic, slice_id, bootstrap_servers=kafka_params.bootstrap_servers
         )
         logger.info(f"librdkafka log level: {self.broker_config.get('log_level', 6)}")
         self.producer_broker_config = build_kafka_producer_configuration(
             topic,
+            slice_id,
             bootstrap_servers=kafka_params.bootstrap_servers,
             override_params={
                 "partitioner": "consistent",
@@ -108,29 +113,37 @@ def __init__(
 
         stream_loader = self.storage.get_table_writer().get_stream_loader()
 
-        self.raw_topic: Topic
+        self.raw_topic: ArroyoTopic
         if kafka_params.raw_topic is not None:
-            self.raw_topic = Topic(kafka_params.raw_topic)
+            self.raw_topic = ArroyoTopic(kafka_params.raw_topic)
         else:
-            self.raw_topic = Topic(stream_loader.get_default_topic_spec().topic_name)
+            default_topic_spec = stream_loader.get_default_topic_spec()
+            self.raw_topic = ArroyoTopic(
+                default_topic_spec.get_physical_topic_name(slice_id)
+            )
 
-        self.replacements_topic: Optional[Topic]
+        self.replacements_topic: Optional[ArroyoTopic]
         if kafka_params.replacements_topic is not None:
-            self.replacements_topic = Topic(kafka_params.replacements_topic)
+            self.replacements_topic = ArroyoTopic(kafka_params.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)
+                self.replacements_topic = ArroyoTopic(
+                    replacement_topic_spec.get_physical_topic_name(slice_id)
+                )
             else:
                 self.replacements_topic = None
 
-        self.commit_log_topic: Optional[Topic]
+        self.commit_log_topic: Optional[ArroyoTopic]
         if kafka_params.commit_log_topic is not None:
-            self.commit_log_topic = Topic(kafka_params.commit_log_topic)
+            self.commit_log_topic = ArroyoTopic(kafka_params.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)
+                self.commit_log_topic = ArroyoTopic(
+                    commit_log_topic_spec.get_physical_topic_name(slice_id)
+                )
             else:
                 self.commit_log_topic = None
 
@@ -172,15 +185,24 @@ def __init__(
         self.__commit_retry_policy = commit_retry_policy
 
     def __build_consumer(
-        self, strategy_factory: ProcessingStrategyFactory[KafkaPayload]
+        self,
+        strategy_factory: ProcessingStrategyFactory[KafkaPayload],
+        slice_id: Optional[int] = None,
     ) -> StreamProcessor[KafkaPayload]:
-        configuration = build_kafka_consumer_configuration(
+
+        # retrieves the default logical topic
+        topic = (
             self.storage.get_table_writer()
             .get_stream_loader()
             .get_default_topic_spec()
-            .topic,
+            .topic
+        )
+
+        configuration = build_kafka_consumer_configuration(
+            topic,
             bootstrap_servers=self.bootstrap_servers,
             group_id=self.group_id,
+            slice_id=slice_id,
             auto_offset_reset=self.auto_offset_reset,
             strict_offset_reset=self.strict_offset_reset,
             queued_max_messages_kbytes=self.queued_max_messages_kbytes,
@@ -220,6 +242,7 @@ def __build_consumer(
 
     def __build_streaming_strategy_factory(
         self,
+        slice_id: Optional[int] = None,
     ) -> ProcessingStrategyFactory[KafkaPayload]:
         table_writer = self.storage.get_table_writer()
         stream_loader = table_writer.get_stream_loader()
@@ -240,6 +263,7 @@ def __build_streaming_strategy_factory(
                     self.producer if self.replacements_topic is not None else None
                 ),
                 replacements_topic=self.replacements_topic,
+                slice_id=slice_id,
             )
             if self.__mock_parameters is None
             else build_mock_batch_writer(
@@ -267,8 +291,12 @@ def __build_streaming_strategy_factory(
 
         return strategy_factory
 
-    def build_base_consumer(self) -> StreamProcessor[KafkaPayload]:
+    def build_base_consumer(
+        self, slice_id: Optional[int] = None
+    ) -> StreamProcessor[KafkaPayload]:
         """
         Builds the consumer.
         """
-        return self.__build_consumer(self.__build_streaming_strategy_factory())
+        return self.__build_consumer(
+            self.__build_streaming_strategy_factory(slice_id), slice_id
+        )

snuba/datasets/partitioning.py

@@ -3,6 +3,8 @@
 should be stored. These do not require individual physical partitions but allow
 for repartitioning with less code changes per physical change.
 """
+from typing import Optional
+
 from snuba.datasets.storages.storage_key import StorageKey
 
 SENTRY_LOGICAL_PARTITIONS = 256
@@ -39,3 +41,16 @@ def is_storage_partitioned(storage: StorageKey) -> bool:
     from snuba.settings import SLICED_STORAGES
 
     return True if storage.value in SLICED_STORAGES.keys() else False
+
+
+def validate_passed_slice(storage_key: StorageKey, slice_id: Optional[int]) -> None:
+    """
+    Verifies that the given storage can be sliced
+    and that the slice_id passed in is within the range
+    of the total number of slices for the given storage
+    """
+    from snuba.settings import SLICED_STORAGES
+
+    if slice_id is not None:
+        assert storage_key.value in SLICED_STORAGES
+        assert slice_id < SLICED_STORAGES[storage_key.value]

snuba/datasets/table_storage.py

@@ -17,6 +17,7 @@
 from snuba.datasets.schemas.tables import WritableTableSchema, WriteFormat
 from snuba.processor import MessageProcessor
 from snuba.replacers.replacer_processor import ReplacerProcessor
+from snuba.settings import SLICED_KAFKA_TOPIC_MAP
 from snuba.snapshots import BulkLoadSource
 from snuba.snapshots.loaders import BulkLoader
 from snuba.snapshots.loaders.single_table import RowProcessor, SingleTableBulkLoader
@@ -57,6 +58,14 @@ def topic_creation_config(self) -> Mapping[str, str]:
     def __eq__(self, other: object) -> bool:
         return isinstance(other, KafkaTopicSpec) and self.topic == other.topic
 
+    def get_physical_topic_name(self, slice_id: Optional[int] = None) -> str:
+        if slice_id is not None:
+            physical_topic = SLICED_KAFKA_TOPIC_MAP[(self.topic_name, slice_id)]
+        else:
+            physical_topic = self.topic_name
+
+        return physical_topic
+
 
 def get_topic_name(topic: Topic) -> str:
     return settings.KAFKA_TOPIC_MAP.get(topic.value, topic.value)
@@ -225,6 +234,7 @@ def get_batch_writer(
         options: ClickhouseWriterOptions = None,
         table_name: Optional[str] = None,
         chunk_size: int = settings.CLICKHOUSE_HTTP_CHUNK_SIZE,
+        slice_id: Optional[int] = None,
     ) -> BatchWriter[JSONRow]:
         table_name = table_name or self.__table_schema.get_table_name()
         if self.__write_format == WriteFormat.JSON:
@@ -240,7 +250,7 @@ def get_batch_writer(
             raise TypeError("unknown table format", self.__write_format)
         options = self.__update_writer_options(options)
 
-        return get_cluster(self.__storage_set).get_batch_writer(
+        return get_cluster(self.__storage_set, slice_id).get_batch_writer(
             metrics,
             insert_statement,
             encoding=None,

snuba/utils/streams/configuration_builder.py

@@ -10,19 +10,28 @@
 from snuba.utils.streams.types import KafkaBrokerConfig
 
 
-def _get_default_topic_configuration(topic: Optional[Topic]) -> Mapping[str, Any]:
+def _get_default_topic_configuration(
+    topic: Optional[Topic], slice_id: Optional[int] = None
+) -> Mapping[str, Any]:
+
     if topic is not None:
-        return settings.KAFKA_BROKER_CONFIG.get(topic.value, settings.BROKER_CONFIG)
+        if slice_id is not None:
+            return settings.SLICED_KAFKA_BROKER_CONFIG.get(
+                (topic.value, slice_id), settings.BROKER_CONFIG
+            )
+        else:
+            return settings.KAFKA_BROKER_CONFIG.get(topic.value, settings.BROKER_CONFIG)
     else:
         return settings.BROKER_CONFIG
 
 
 def get_default_kafka_configuration(
     topic: Optional[Topic] = None,
+    slice_id: Optional[int] = None,
     bootstrap_servers: Optional[Sequence[str]] = None,
     override_params: Optional[Mapping[str, Any]] = None,
 ) -> KafkaBrokerConfig:
-    default_topic_config = _get_default_topic_configuration(topic)
+    default_topic_config = _get_default_topic_configuration(topic, slice_id)
 
     return build_kafka_configuration(
         default_topic_config, bootstrap_servers, override_params
@@ -32,14 +41,15 @@ def get_default_kafka_configuration(
 def build_kafka_consumer_configuration(
     topic: Optional[Topic],
     group_id: str,
+    slice_id: Optional[int] = None,
     auto_offset_reset: Optional[str] = None,
     queued_max_messages_kbytes: Optional[int] = None,
     queued_min_messages: Optional[int] = None,
     bootstrap_servers: Optional[Sequence[str]] = None,
     override_params: Optional[Mapping[str, Any]] = None,
     strict_offset_reset: Optional[bool] = None,
 ) -> KafkaBrokerConfig:
-    default_topic_config = _get_default_topic_configuration(topic)
+    default_topic_config = _get_default_topic_configuration(topic, slice_id)
 
     return _build_kafka_consumer_configuration(
         default_topic_config,
@@ -55,11 +65,13 @@ def build_kafka_consumer_configuration(
 
 def build_kafka_producer_configuration(
     topic: Optional[Topic],
+    slice_id: Optional[int] = None,
     bootstrap_servers: Optional[Sequence[str]] = None,
     override_params: Optional[Mapping[str, Any]] = None,
 ) -> KafkaBrokerConfig:
     broker_config = get_default_kafka_configuration(
         topic=topic,
+        slice_id=slice_id,
         bootstrap_servers=bootstrap_servers,
         override_params=override_params,
     )

tests/datasets/test_table_storage.py

@@ -0,0 +1,21 @@
+from snuba.datasets.storages.factory import get_writable_storage
+from snuba.datasets.storages.storage_key import StorageKey
+from snuba.settings import SLICED_KAFKA_TOPIC_MAP
+
+
+def test_get_physical_topic_name(monkeypatch) -> None:  # type: ignore
+
+    monkeypatch.setitem(
+        SLICED_KAFKA_TOPIC_MAP, ("ingest-replay-events", 2), "ingest-replay-events-2"
+    )
+
+    storage_key = StorageKey.REPLAYS
+    storage = get_writable_storage(storage_key)
+
+    stream_loader = storage.get_table_writer().get_stream_loader()
+
+    default_topic_spec = stream_loader.get_default_topic_spec()
+
+    physical_topic_name = default_topic_spec.get_physical_topic_name(slice_id=2)
+
+    assert physical_topic_name == "ingest-replay-events-2"