Clarify

Author: kcons

src/sentry/utils/function_cache.py

@@ -202,26 +202,20 @@ def build_batched_cache_func[
     """
 
     def inner(args: list[Key]) -> list[R]:
-        result = []
-        missing_indexes = []
-        for i, arg in enumerate(args):
-            cache_key = cache_key_for_batched_func(func_to_cache, arg)
-            cached_val = cache.get(cache_key, None)
-            if cached_val is None:
-                missing_indexes.append(i)
-            result.append(cached_val)
-        if missing_indexes:
-            timeout = cache_ttl.total_seconds()
-            missing_keys = [args[i] for i in missing_indexes]
-            missing_vals = func_to_cache(missing_keys)
-            if len(missing_vals) != len(missing_keys):
+        cache_keys = [cache_key_for_batched_func(func_to_cache, arg) for arg in args]
+        values = cache.get_many(cache_keys)
+        missing_keys = [key for key in cache_keys if key not in values]
+        if missing_keys:
+            key_to_arg = {key: arg for arg, key in zip(args, cache_keys)}
+            missing_args = [key_to_arg[key] for key in missing_keys]
+            missing_vals = func_to_cache(missing_args)
+            if len(missing_vals) != len(missing_args):
                 raise ValueError(
                     f"Function {func_to_cache.__qualname__} returned {len(missing_vals)} results for {len(missing_keys)} keys"
                 )
-            for out_idx, key, val in zip(missing_indexes, missing_keys, missing_vals):
-                result[out_idx] = val
-                cache_key = cache_key_for_batched_func(func_to_cache, key)
-                cache.set(cache_key, val, timeout=timeout)
-        return result
+            fetched = dict(zip(missing_keys, missing_vals))
+            cache.set_many(fetched, timeout=cache_ttl.total_seconds())
+            values.update(fetched)
+        return [values[key] for key in cache_keys]
 
     return inner

src/sentry/workflow_engine/processors/data_condition_group.py

@@ -49,9 +49,7 @@ def get_data_conditions_for_groups(
     by_id = DefaultDict[int, list[DataCondition]](list)
     for cond in DataCondition.objects.filter(condition_group_id__in=data_condition_group_ids):
         by_id[cond.condition_group_id].append(cond)
-    result = [by_id[dcg] for dcg in data_condition_group_ids]
-    assert len(result) == len(data_condition_group_ids)
-    return result
+    return [by_id[dcg] for dcg in data_condition_group_ids]
 
 
 def batch_get_slow_conditions(

tests/sentry/utils/test_function_cache.py

@@ -131,6 +131,9 @@ def batch_arg_extractor(instance: CacheModel) -> str:
             mock_test_func
         )
 
+        # Empty list should work fine.
+        assert decorated_test_func([]) == []
+
         # First call should hit the function
         self.assert_called_with_count(mock_test_func, ["test"], 0)
         assert decorated_test_func(["test"]) == [0]
@@ -140,6 +143,10 @@ def batch_arg_extractor(instance: CacheModel) -> str:
         assert decorated_test_func(["test"]) == [0]
         self.assert_called_with_count(mock_test_func, ["test"], 1)
 
+        # Duplicate key should work fine.
+        assert decorated_test_func(["test", "test"]) == [0, 0]
+        self.assert_called_with_count(mock_test_func, ["test"], 1)
+
         # Create a model, should invalidate cache
         CacheModel.objects.create(some_field="test")
         assert decorated_test_func(["test"]) == [1]