feat(cbrs): Time Window Routing (#7418)

This PR implements a routing strategy for routing to TIER_1 and instead
of downsampling storage tiers, shrinking time windows.

# Design decisions

## Window Sizing Algorithm

This is the easiest thing to change in this PR (and likely it will
change). At the moment, it looks up the amount of outcomes for the
requested time range and shrinks the time window down assuming that the
distribution of datapoints is uniform across time (which is not true).
Most recent datapoints are prioritized first

## Pagination

Only the TraceItemTable endpoint makes use of the recommendations of
this routing strategy, the endpoint and routing strategy interact across
queries. This is to facilitate a simple client side UX where all the
client has to do is pass the `page_token` across their requests and not
worry about anything else.

Here's a diagram explaining the flow:

```
┌─────────────────┐              ┌────────────────┐                               
│                 │              │                │                               
│                 │              │                │                               
│                 │              │   Routing      │                               
│   Client        ├─────────────►│   Strategy     ┼───────────────────┐           
│                 │  page_token  │                │                   │           
│                 │              │                │                   │           
│                 │              │                │            Narrows│time window
└─────────────────┘              └────────────────┘            For Endpoint       
        ▲                                                             │           
        │                                                             │           
        │                         ┌────────────────┐                  │           
        │                         │                │                  │           
        │                         │                │                  │           
        │                         │ TraceItemTable │◄─────────────────┘           
        └─────────────────────────┼ Endpoint       │                              
             Encodes              │                │                              
             Time Window          └────────────────┘                              
             In Page Token                                                        

```

In order to facilitate pagination, the TraceItemTable endpoint now
queries for `limit + 1` rows in order to know if there are more items in
this current window or if we can move on to the next one

# What's missing

1. This functionality can be tested more rigorously, I deliberately did
not spend too much time on it because I know it will change and the
priority is to get something out there to try
2. We could probably have more observability into what the strategy is
doing and understanding our success metrics better. These will be added
as we understand the problem more

---------

Co-authored-by: getsantry[bot] <66042841+getsantry[bot]@users.noreply.github.com>

Author: volokluev

.vscode/settings.json

@@ -69,6 +69,7 @@
   "mypy-type-checker.args": [
     "--strict"
   ],
-  "makefile.configureOnOpen": false
+  "makefile.configureOnOpen": false,
+  "cursorpyright.analysis.autoImportCompletions": true
 
 }

snuba/web/rpc/common/common.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Callable, TypeVar, cast
 
 from google.protobuf.message import Message as ProtobufMessage
@@ -56,9 +56,9 @@ def truncate_request_meta_to_day(meta: RequestMeta) -> None:
     start_timestamp = start_timestamp.replace(
         hour=0, minute=0, second=0, microsecond=0
     ) - timedelta(days=1)
-    end_timestamp = end_timestamp.replace(
-        hour=0, minute=0, second=0, microsecond=0
-    ) + timedelta(days=1)
+    end_timestamp = end_timestamp.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(
+        days=1
+    )
 
     meta.start_timestamp.seconds = int(start_timestamp.timestamp())
     meta.end_timestamp.seconds = int(end_timestamp.timestamp())
@@ -142,51 +142,32 @@ def trace_item_filters_to_expression(
         if len(filters) == 0:
             return literal(True)
         elif len(filters) == 1:
-            return trace_item_filters_to_expression(
-                filters[0], attribute_key_to_expression
-            )
+            return trace_item_filters_to_expression(filters[0], attribute_key_to_expression)
         return and_cond(
-            *(
-                trace_item_filters_to_expression(x, attribute_key_to_expression)
-                for x in filters
-            )
+            *(trace_item_filters_to_expression(x, attribute_key_to_expression) for x in filters)
         )
 
     if item_filter.HasField("or_filter"):
         filters = item_filter.or_filter.filters
         if len(filters) == 0:
-            raise BadSnubaRPCRequestException(
-                "Invalid trace item filter, empty 'or' clause"
-            )
+            raise BadSnubaRPCRequestException("Invalid trace item filter, empty 'or' clause")
         elif len(filters) == 1:
-            return trace_item_filters_to_expression(
-                filters[0], attribute_key_to_expression
-            )
+            return trace_item_filters_to_expression(filters[0], attribute_key_to_expression)
         return or_cond(
-            *(
-                trace_item_filters_to_expression(x, attribute_key_to_expression)
-                for x in filters
-            )
+            *(trace_item_filters_to_expression(x, attribute_key_to_expression) for x in filters)
         )
 
     if item_filter.HasField("not_filter"):
         filters = item_filter.not_filter.filters
         if len(filters) == 0:
-            raise BadSnubaRPCRequestException(
-                "Invalid trace item filter, empty 'not' clause"
-            )
+            raise BadSnubaRPCRequestException("Invalid trace item filter, empty 'not' clause")
         elif len(filters) == 1:
             return not_cond(
-                trace_item_filters_to_expression(
-                    filters[0], attribute_key_to_expression
-                )
+                trace_item_filters_to_expression(filters[0], attribute_key_to_expression)
             )
         return not_cond(
             and_cond(
-                *(
-                    trace_item_filters_to_expression(x, attribute_key_to_expression)
-                    for x in filters
-                )
+                *(trace_item_filters_to_expression(x, attribute_key_to_expression) for x in filters)
             )
         )
 
@@ -198,9 +179,7 @@ def trace_item_filters_to_expression(
 
         value_type = v.WhichOneof("value")
         if value_type is None:
-            raise BadSnubaRPCRequestException(
-                "comparison does not have a right hand side"
-            )
+            raise BadSnubaRPCRequestException("comparison does not have a right hand side")
 
         if v.is_null:
             v_expression: Expression = literal(None)
@@ -246,9 +225,7 @@ def trace_item_filters_to_expression(
             )
             # we redefine the way equals works for nulls
             # now null=null is true
-            expr_with_null = or_cond(
-                expr, and_cond(f.isNull(k_expression), f.isNull(v_expression))
-            )
+            expr_with_null = or_cond(expr, and_cond(f.isNull(k_expression), f.isNull(v_expression)))
             return expr_with_null
         if op == ComparisonFilter.OP_NOT_EQUALS:
             _check_non_string_values_cannot_ignore_case(item_filter.comparison_filter)
@@ -259,9 +236,7 @@ def trace_item_filters_to_expression(
             )
             # we redefine the way not equals works for nulls
             # now null!=null is true
-            expr_with_null = or_cond(
-                expr, f.xor(f.isNull(k_expression), f.isNull(v_expression))
-            )
+            expr_with_null = or_cond(expr, f.xor(f.isNull(k_expression), f.isNull(v_expression)))
             return expr_with_null
         if op == ComparisonFilter.OP_LIKE:
             if k.type != AttributeKey.Type.TYPE_STRING:
@@ -354,11 +329,15 @@ def timestamp_in_range_condition(start_ts: int, end_ts: int) -> Expression:
     return and_cond(
         f.less(
             column("timestamp"),
-            f.toDateTime(end_ts),
+            f.toDateTime(
+                datetime.fromtimestamp(end_ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
+            ),
         ),
         f.greaterOrEquals(
             column("timestamp"),
-            f.toDateTime(start_ts),
+            f.toDateTime(
+                datetime.fromtimestamp(start_ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
+            ),
         ),
     )
 
@@ -372,9 +351,7 @@ def base_conditions_and(meta: RequestMeta, *other_exprs: Expression) -> Expressi
     """
     return and_cond(
         project_id_and_org_conditions(meta),
-        timestamp_in_range_condition(
-            meta.start_timestamp.seconds, meta.end_timestamp.seconds
-        ),
+        timestamp_in_range_condition(meta.start_timestamp.seconds, meta.end_timestamp.seconds),
         *other_exprs,
     )
 

snuba/web/rpc/common/pagination.py

@@ -0,0 +1,90 @@
+"""
+This file contains functionality to encode and decode custom page tokens
+"""
+
+from google.protobuf.timestamp_pb2 import Timestamp
+from sentry_protos.snuba.v1.request_common_pb2 import PageToken
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
+    AndFilter,
+    ComparisonFilter,
+    TraceItemFilter,
+)
+
+
+class FlexibleTimeWindowPage:
+    _START_TIMESTAMP_KEY = "sentry.start_timestamp"
+    _END_TIMESTAMP_KEY = "sentry.end_timestamp"
+    _OFFSET_KEY = "sentry.offset"
+
+    def __init__(
+        self,
+        start_timestamp: Timestamp | None,
+        end_timestamp: Timestamp | None,
+        offset: int | None = None,
+    ):
+        self.start_timestamp = start_timestamp
+        self.end_timestamp = end_timestamp
+        self.offset = offset
+
+    def encode(self) -> PageToken:
+        if self.start_timestamp is not None and self.end_timestamp is not None:
+            return PageToken(
+                filter_offset=TraceItemFilter(
+                    and_filter=AndFilter(
+                        filters=[
+                            TraceItemFilter(
+                                comparison_filter=ComparisonFilter(
+                                    key=AttributeKey(name=self._START_TIMESTAMP_KEY),
+                                    op=ComparisonFilter.OP_GREATER_THAN_OR_EQUALS,
+                                    value=AttributeValue(val_int=self.start_timestamp.seconds),
+                                )
+                            ),
+                            TraceItemFilter(
+                                comparison_filter=ComparisonFilter(
+                                    key=AttributeKey(name=self._END_TIMESTAMP_KEY),
+                                    op=ComparisonFilter.OP_LESS_THAN,
+                                    value=AttributeValue(val_int=self.end_timestamp.seconds),
+                                )
+                            ),
+                            TraceItemFilter(
+                                comparison_filter=ComparisonFilter(
+                                    key=AttributeKey(name=self._OFFSET_KEY),
+                                    op=ComparisonFilter.OP_EQUALS,
+                                    value=AttributeValue(
+                                        val_int=self.offset if self.offset is not None else 0
+                                    ),
+                                )
+                            ),
+                        ]
+                    )
+                )
+            )
+        else:
+            return PageToken(offset=self.offset if self.offset is not None else 0)
+
+    @classmethod
+    def decode(cls, page_token: PageToken) -> "FlexibleTimeWindowPage":
+        start_timestamp = None
+        end_timestamp = None
+        offset = None
+        if page_token.filter_offset.HasField("and_filter"):
+            for filter in page_token.filter_offset.and_filter.filters:
+                if (
+                    filter.HasField("comparison_filter")
+                    and filter.comparison_filter.key.name == cls._START_TIMESTAMP_KEY
+                ):
+                    start_timestamp = Timestamp(seconds=filter.comparison_filter.value.val_int)
+                if (
+                    filter.HasField("comparison_filter")
+                    and filter.comparison_filter.key.name == cls._END_TIMESTAMP_KEY
+                ):
+                    end_timestamp = Timestamp(seconds=filter.comparison_filter.value.val_int)
+                if (
+                    filter.HasField("comparison_filter")
+                    and filter.comparison_filter.key.name == cls._OFFSET_KEY
+                ):
+                    offset = filter.comparison_filter.value.val_int
+        elif page_token.HasField("offset"):
+            offset = page_token.offset
+        return cls(start_timestamp, end_timestamp, offset)

snuba/web/rpc/storage_routing/routing_strategies/common.py

@@ -0,0 +1,17 @@
+from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
+
+
+class OutcomeCategory:
+    SPAN_INDEXED = 16
+    LOG_ITEM = 23
+
+
+class Outcome:
+    ACCEPTED = 0
+
+
+ITEM_TYPE_TO_OUTCOME_CATEGORY = {
+    TraceItemType.TRACE_ITEM_TYPE_UNSPECIFIED: OutcomeCategory.SPAN_INDEXED,
+    TraceItemType.TRACE_ITEM_TYPE_SPAN: OutcomeCategory.SPAN_INDEXED,
+    TraceItemType.TRACE_ITEM_TYPE_LOG: OutcomeCategory.LOG_ITEM,
+}

snuba/web/rpc/storage_routing/routing_strategies/outcomes_based.py

@@ -4,7 +4,7 @@
 
 import sentry_sdk
 from google.protobuf.json_format import MessageToDict
-from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta, TraceItemType
+from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta
 
 from snuba import state
 from snuba.attribution.appid import AppID
@@ -28,29 +28,17 @@
     treeify_or_and_conditions,
 )
 from snuba.web.rpc.storage_routing.common import extract_message_meta
+from snuba.web.rpc.storage_routing.routing_strategies.common import (
+    ITEM_TYPE_TO_OUTCOME_CATEGORY,
+    Outcome,
+)
 from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
     BaseRoutingStrategy,
     RoutingContext,
     RoutingDecision,
 )
 
 
-# TODO import these from sentry-relay
-class OutcomeCategory:
-    SPAN_INDEXED = 16
-    LOG_ITEM = 23
-
-
-class Outcome:
-    ACCEPTED = 0
-
-
-_ITEM_TYPE_TO_OUTCOME = {
-    TraceItemType.TRACE_ITEM_TYPE_SPAN: OutcomeCategory.SPAN_INDEXED,
-    TraceItemType.TRACE_ITEM_TYPE_LOG: OutcomeCategory.LOG_ITEM,
-}
-
-
 def project_id_and_org_conditions(meta: RequestMeta) -> Expression:
     return and_cond(
         in_cond(
@@ -111,11 +99,7 @@ def get_ingested_items_for_timerange(self, routing_context: RoutingContext) -> i
                 ),
                 f.equals(column("outcome"), Outcome.ACCEPTED),
                 f.equals(
-                    column("category"),
-                    _ITEM_TYPE_TO_OUTCOME.get(
-                        in_msg_meta.trace_item_type,
-                        OutcomeCategory.SPAN_INDEXED,
-                    ),
+                    column("category"), ITEM_TYPE_TO_OUTCOME_CATEGORY[in_msg_meta.trace_item_type]
                 ),
             ),
         )
@@ -186,8 +170,8 @@ def _get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecis
             # for GetTraces, there is no type specified so we assume spans because
             # that is necessary for traces anyways
             # if the type is specified and we don't know its outcome, route to Tier_1
-            in_msg_meta.trace_item_type != TraceItemType.TRACE_ITEM_TYPE_UNSPECIFIED
-            and in_msg_meta.trace_item_type not in _ITEM_TYPE_TO_OUTCOME
+            in_msg_meta.trace_item_type
+            not in ITEM_TYPE_TO_OUTCOME_CATEGORY
         ):
             return routing_decision
         # if we're querying a short enough timeframe, don't bother estimating, route to tier 1 and call it a day