From 1077ec2d886b9201bca504782a25636bcf904036 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 17:03:49 +0000 Subject: [PATCH 01/41] Deduplicate merger WIP --- smartfeed/examples/example_client.py | 13 +- smartfeed/manager.py | 5 +- smartfeed/schemas.py | 297 +++++++++++++++++++++-- tests/fixtures/configs.py | 30 +++ tests/fixtures/redis.py | 28 ++- tests/test_merger_append.py | 5 +- tests/test_merger_append_distribute.py | 5 +- tests/test_merger_deduplication.py | 271 +++++++++++++++++++++ tests/test_merger_percentage.py | 3 +- tests/test_merger_percentage_gradient.py | 5 +- tests/test_merger_positional.py | 7 +- tests/test_merger_view_session.py | 11 +- tests/test_parsing_config.py | 14 +- tests/test_redis_live.py | 5 +- tests/test_sub_feed.py | 8 +- tests/utils.py | 17 ++ 16 files changed, 670 insertions(+), 54 deletions(-) create mode 100644 tests/test_merger_deduplication.py create mode 100644 tests/utils.py diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index 9a421ff..8b9d6a7 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -2,7 +2,7 @@ import json from typing import Optional, Union -from pydantic import BaseModel, Field, validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside @@ -18,13 +18,16 @@ class TestClientRequest(BaseModel): base64.urlsafe_b64encode(json.dumps({"data": {}}).encode()).decode() ) - class Config: - validate_all = True + model_config = ConfigDict(validate_default=True) - @validator("next_page") + @field_validator("next_page") + @classmethod def validate_next_page(cls, value: Union[str, FeedResultNextPage]) -> Union[str, FeedResultNextPage]: if isinstance(value, str): - return FeedResultNextPage.parse_obj(json.loads(base64.urlsafe_b64decode(value))) + payload = json.loads(base64.urlsafe_b64decode(value)) + if hasattr(FeedResultNextPage, "model_validate"): + return FeedResultNextPage.model_validate(payload) # type: ignore[attr-defined] + return FeedResultNextPage.parse_obj(payload) return value diff --git a/smartfeed/manager.py b/smartfeed/manager.py index e91bbe9..7ac06f9 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -20,7 +20,10 @@ def __init__(self, config: Dict, methods_dict: Dict, redis_client: Optional[Unio :param redis_client: объект клиента Redis (для конфигурации с view_session = True). """ - self.feed_config = FeedConfig.parse_obj(config) + if hasattr(FeedConfig, "model_validate"): + self.feed_config = FeedConfig.model_validate(config) # type: ignore[attr-defined] + else: + self.feed_config = FeedConfig.parse_obj(config) self.methods_dict = methods_dict self.redis_client = redis_client diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 45df221..a9ef784 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,18 +1,21 @@ +import base64 import inspect import json import logging +import zlib from abc import ABC, abstractmethod from collections import defaultdict, deque from random import shuffle from typing import Annotated, Any, Callable, Dict, List, Literal, Optional, Union, no_type_check import redis -from pydantic import BaseModel, Field, root_validator +from pydantic import BaseModel, Field, model_validator from redis.asyncio import Redis as AsyncRedis from redis.asyncio import RedisCluster as AsyncRedisCluster FeedTypes = Annotated[ Union[ + "DeduplicationMerger", "MergerAppend", "MergerAppendDistribute", "MergerPositional", @@ -501,17 +504,17 @@ class MergerPositional(BaseFeedConfigModel): positional: FeedTypes default: FeedTypes - @root_validator(skip_on_failure=True) - def validate_merger_positional(cls, values: Dict[str, Any]) -> Dict[str, Any]: - if not values["positions"] and not all((values["start"], values["end"], values["step"])): + @model_validator(mode="after") + def validate_merger_positional(self) -> "MergerPositional": + if not self.positions and not all((self.start, self.end, self.step)): raise ValueError('Either "positions" or "start", "end", and "step" must be provided') - if values["start"] and values["positions"]: - if isinstance(values["start"], int) and values["start"] <= max(values["positions"]): + if self.start and self.positions: + if isinstance(self.start, int) and self.start <= max(self.positions): raise ValueError('"start" must be bigger than maximum value of "positions"') - if isinstance(values["start"], int) and isinstance(values["end"], int): - if values["end"] <= values["start"]: + if isinstance(self.start, int) and isinstance(self.end, int): + if self.end <= self.start: raise ValueError('"end" must be bigger than "start"') - return values + return self async def get_data( self, @@ -757,13 +760,13 @@ class MergerPercentageGradient(BaseFeedConfigModel): size_to_step: int shuffle: bool = False - @root_validator(skip_on_failure=True) - def validate_merger_percentage_gradient(cls, values: Dict[str, Any]) -> Dict[str, Any]: - if values["step"] < 1 or values["step"] > 100: + @model_validator(mode="after") + def validate_merger_percentage_gradient(self) -> "MergerPercentageGradient": + if self.step < 1 or self.step > 100: raise ValueError('"step" must be in range from 1 to 100') - if values["size_to_step"] < 1: + if self.size_to_step < 1: raise ValueError('"size_to_step" must be bigger than 1') - return values + return self async def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: """ @@ -1018,6 +1021,247 @@ async def get_data( return result +class DeduplicationMergerItem(BaseModel): + """Configuration item for DeduplicationMerger.""" + + priority: int = 0 + data: FeedTypes + + +class DeduplicationMerger(BaseFeedConfigModel): + """Merger that deduplicates items and refills to the requested limit. + + Key properties: + - Always tries to return exactly `limit` unique items if they exist upstream. + - Supports cross-page deduplication using either cursor state or Redis. + - Supports explicit per-source priority; higher priority wins on same dedup key. + """ + + merger_id: str + type: Literal["merger_deduplication"] + items: List[DeduplicationMergerItem] + + dedup_key: Optional[str] = None + missing_key_policy: Literal["error", "keep", "drop"] = "error" + + state_backend: Literal["cursor", "redis"] = "cursor" + state_ttl_seconds: int = 3600 + cursor_compress: bool = True + cursor_max_keys: Optional[int] = None + + overfetch_factor: int = 2 + max_refill_loops: int = 20 + + def _normalize_key(self, value: Any) -> str: + if isinstance(value, (str, int)): + return str(value) + if isinstance(value, (dict, list)): + return json.dumps(value, sort_keys=True, default=str) + return str(value) + + def _extract_dedup_value(self, item: Any) -> Any: + if not self.dedup_key: + return item + + try: + value = item.get(self.dedup_key) + except AttributeError: + value = getattr(item, self.dedup_key, None) + + if value is None and self.missing_key_policy == "error": + raise AssertionError( + f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}" + ) + return value + + def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> List[str]: + entry = next_page.data.get(self.merger_id) + if not entry or entry.after is None: + return [] + + after = entry.after + if isinstance(after, dict) and "z" in after: + payload = base64.urlsafe_b64decode(after["z"].encode()) + raw = zlib.decompress(payload).decode() + return list(json.loads(raw)) + if isinstance(after, dict) and "seen" in after: + return list(after["seen"]) + if isinstance(after, list): + return list(after) + return [] + + def _encode_seen_for_cursor(self, seen_keys_in_order: List[str]) -> Any: + if self.cursor_max_keys is not None: + seen_keys_in_order = seen_keys_in_order[-self.cursor_max_keys :] + + if not self.cursor_compress: + return {"v": 1, "seen": seen_keys_in_order} + + raw = json.dumps(seen_keys_in_order).encode() + compressed = zlib.compress(raw) + return { + "v": 1, + "c": "zlib+base64", + "n": len(seen_keys_in_order), + "z": base64.urlsafe_b64encode(compressed).decode(), + } + + async def _redis_sismember(self, redis_client: Union[redis.Redis, AsyncRedis], key: str, member: str) -> bool: + res = redis_client.sismember(key, member) + if inspect.iscoroutine(res): + res = await res + return bool(res) + + async def _redis_sadd_and_expire( + self, + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + members: List[str], + ) -> None: + if not members: + return + res = redis_client.sadd(key, *members) + if inspect.iscoroutine(res): + await res + await redis_client.expire(key, self.state_ttl_seconds) + else: + redis_client.expire(key, self.state_ttl_seconds) + + def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: + suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") + if suffix: + return f"dedup:{self.merger_id}:{user_id}:{suffix}" + return f"dedup:{self.merger_id}:{user_id}" + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + if limit <= 0: + return FeedResult(data=[], next_page=next_page, has_next_page=False) + + # Treat an explicit "page 0" (or missing cursor for this merger) as a fresh session. + # This allows clients to restart the feed (e.g., full reload) without carrying over seen state. + requested_page = next_page.data.get(self.merger_id).page if self.merger_id in next_page.data else None + is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) + + if self.state_backend == "redis" and not redis_client: + raise ValueError("Redis client must be provided if using DeduplicationMerger with state_backend=redis") + + if hasattr(next_page, "model_copy"): + working_next_page = next_page.model_copy(deep=True) # type: ignore[attr-defined] + else: + working_next_page = next_page.copy(deep=True) + sorted_items = sorted(self.items, key=lambda x: x.priority, reverse=True) + + seen_keys_in_order: List[str] = [] + seen_cursor_set: set[str] = set() + if self.state_backend == "cursor" and not is_fresh_session: + seen_keys_in_order = self._decode_seen_from_cursor(next_page) + seen_cursor_set = set(seen_keys_in_order) + + redis_state_key = "" + if self.state_backend == "redis" and redis_client: + redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) + if is_fresh_session: + # Drop state for a full restart. + deleted = redis_client.delete(redis_state_key) + if inspect.iscoroutine(deleted): + await deleted + + result_items: List[Any] = [] + accepted: Dict[str, Dict[str, Any]] = {} + redis_new_members: List[str] = [] + any_has_next_page = False + + loops = 0 + while len(result_items) < limit and loops < self.max_refill_loops: + loops += 1 + before_len = len(result_items) + + for item in sorted_items: + remaining = limit - len(result_items) + if remaining <= 0: + break + + request_limit = max(1, remaining * max(1, self.overfetch_factor)) + item_result = await item.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=request_limit, + next_page=working_next_page, + redis_client=redis_client, + **params, + ) + + any_has_next_page = any_has_next_page or item_result.has_next_page + working_next_page.data.update(item_result.next_page.data) + + for entity in item_result.data: + raw_value = self._extract_dedup_value(entity) + if raw_value is None: + if self.missing_key_policy == "drop": + continue + if self.missing_key_policy == "keep": + # Make a unique key per object instance representation. + raw_value = ("__missing__", id(entity)) + + key = self._normalize_key(raw_value) + + if key in accepted: + if item.priority > accepted[key]["priority"]: + result_items[accepted[key]["index"]] = entity + accepted[key]["priority"] = item.priority + continue + + if self.state_backend == "cursor": + if key in seen_cursor_set: + continue + else: + assert redis_client is not None + if await self._redis_sismember(redis_client, redis_state_key, key): + continue + + accepted[key] = {"priority": item.priority, "index": len(result_items)} + result_items.append(entity) + + if self.state_backend == "cursor": + seen_cursor_set.add(key) + seen_keys_in_order.append(key) + else: + redis_new_members.append(key) + + if len(result_items) >= limit: + break + + if len(result_items) >= limit: + break + + if len(result_items) == before_len: + break + + if self.state_backend == "redis" and redis_client: + await self._redis_sadd_and_expire(redis_client, redis_state_key, redis_new_members) + + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + merger_after: Any = None + if self.state_backend == "cursor": + merger_after = self._encode_seen_for_cursor(seen_keys_in_order) + + if hasattr(working_next_page, "model_copy"): + result_next_page = working_next_page.model_copy(deep=True) # type: ignore[attr-defined] + else: + result_next_page = working_next_page.copy(deep=True) + result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) + + return FeedResult(data=result_items, next_page=result_next_page, has_next_page=any_has_next_page) + + class SubFeed(BaseFeedConfigModel): """ Модель субфида. @@ -1122,11 +1366,20 @@ class FeedConfig(BaseModel): # Update Forward Refs -MergerPositional.update_forward_refs() -MergerPercentage.update_forward_refs() -SubFeed.update_forward_refs() -MergerPercentageItem.update_forward_refs() -MergerAppend.update_forward_refs() -MergerAppendDistribute.update_forward_refs() -MergerPercentageGradient.update_forward_refs() -MergerViewSession.update_forward_refs() +def _rebuild_model(model: Any) -> None: + if hasattr(model, "model_rebuild"): + model.model_rebuild() # type: ignore[attr-defined] + else: + model.update_forward_refs() # type: ignore[attr-defined] + + +_rebuild_model(MergerPositional) +_rebuild_model(MergerPercentage) +_rebuild_model(SubFeed) +_rebuild_model(MergerPercentageItem) +_rebuild_model(MergerAppend) +_rebuild_model(MergerAppendDistribute) +_rebuild_model(MergerPercentageGradient) +_rebuild_model(MergerViewSession) +_rebuild_model(DeduplicationMergerItem) +_rebuild_model(DeduplicationMerger) diff --git a/tests/fixtures/configs.py b/tests/fixtures/configs.py index 8c96e4e..a982e3d 100644 --- a/tests/fixtures/configs.py +++ b/tests/fixtures/configs.py @@ -86,3 +86,33 @@ }, }, } + + +PARSING_DEDUP_CONFIG_FIXTURE = { + "version": "1", + "feed": { + "merger_id": "merger_deduplication_parsing_example", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "items": [ + { + "priority": 100, + "data": { + "subfeed_id": "subfeed_dedup_priority_high", + "type": "subfeed", + "method_name": "posted", + }, + }, + { + "priority": 0, + "data": { + "subfeed_id": "subfeed_dedup_priority_low", + "type": "subfeed", + "method_name": "posted", + }, + }, + ], + }, +} diff --git a/tests/fixtures/redis.py b/tests/fixtures/redis.py index b98695e..2c07678 100644 --- a/tests/fixtures/redis.py +++ b/tests/fixtures/redis.py @@ -1,10 +1,30 @@ import pytest +import pytest_asyncio import redis from redis.asyncio import Redis as AsyncRedis -@pytest.fixture(scope="function") -def redis_client(request): +@pytest_asyncio.fixture(scope="function") +async def redis_client(request): + """Provide a Redis client for tests. + + If Redis is not available on localhost:6379, skip tests that depend on it. + """ + if request.param == "async": - return AsyncRedis(host="localhost", port=6379) - return redis.Redis(host="localhost", port=6379, db=0) + client = AsyncRedis(host="localhost", port=6379) + try: + await client.ping() + except Exception: # pragma: no cover + pytest.skip("Redis is not available on localhost:6379") + yield client + await client.aclose() + return + + client = redis.Redis(host="localhost", port=6379, db=0) + try: + client.ping() + except Exception: # pragma: no cover + pytest.skip("Redis is not available on localhost:6379") + yield client + client.close() diff --git a/tests/test_merger_append.py b/tests/test_merger_append.py index e9db5c7..309ea82 100644 --- a/tests/test_merger_append.py +++ b/tests/test_merger_append.py @@ -3,6 +3,7 @@ from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerAppend from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_APPEND_CONFIG +from tests.utils import parse_model @pytest.mark.asyncio @@ -11,7 +12,7 @@ async def test_merger_append() -> None: Тест для проверки получения данных из append мерджера. """ - merger_append = MergerAppend.parse_obj(MERGER_APPEND_CONFIG) + merger_append = parse_model(MergerAppend, MERGER_APPEND_CONFIG) merger_append_res = await merger_append.get_data( methods_dict=METHODS_DICT, limit=11, @@ -28,7 +29,7 @@ async def test_merger_append_with_item_1_page_2() -> None: Тест для проверки получения данных из append мерджера с курсором пагинации первого субфида. """ - merger_append = MergerAppend.parse_obj(MERGER_APPEND_CONFIG) + merger_append = parse_model(MergerAppend, MERGER_APPEND_CONFIG) merger_append_res = await merger_append.get_data( methods_dict=METHODS_DICT, limit=11, diff --git a/tests/test_merger_append_distribute.py b/tests/test_merger_append_distribute.py index 6bb1782..bc4878b 100644 --- a/tests/test_merger_append_distribute.py +++ b/tests/test_merger_append_distribute.py @@ -3,6 +3,7 @@ from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerAppendDistribute from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_APPEND_DISTRIBUTE_CONFIG +from tests.utils import parse_model @pytest.mark.asyncio @@ -11,7 +12,7 @@ async def test_merger_disturbed_append() -> None: Тест для проверки получения данных из append мерджера. """ - merger_distributed = MergerAppendDistribute.parse_obj(MERGER_APPEND_DISTRIBUTE_CONFIG) + merger_distributed = parse_model(MergerAppendDistribute, MERGER_APPEND_DISTRIBUTE_CONFIG) merger_distributed_res = await merger_distributed.get_data( methods_dict=METHODS_DICT, limit=20, @@ -31,7 +32,7 @@ async def test_merger_append_with_item_1_page_2() -> None: """ Тест для проверки получения данных из append мерджера с курсором пагинации первого субфида. """ - merger_distributed = MergerAppendDistribute.parse_obj(MERGER_APPEND_DISTRIBUTE_CONFIG) + merger_distributed = parse_model(MergerAppendDistribute, MERGER_APPEND_DISTRIBUTE_CONFIG) merger_distributed_res = await merger_distributed.get_data( methods_dict=METHODS_DICT, limit=11, diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py new file mode 100644 index 0000000..5e16c45 --- /dev/null +++ b/tests/test_merger_deduplication.py @@ -0,0 +1,271 @@ +import inspect + +import pytest + +from smartfeed.schemas import ( + DeduplicationMerger, + FeedResultClient, + FeedResultNextPage, + FeedResultNextPageInside, +) + +from tests.fixtures.redis import redis_client # noqa: F401 +from tests.utils import parse_model + + +def make_offset_paged_method(items): + async def _method(user_id, limit, next_page): # pylint: disable=unused-argument + offset = int(next_page.after or 0) + result_data = items[offset : offset + limit] + next_page.after = offset + len(result_data) + next_page.page += 1 + has_next_page = (offset + len(result_data)) < len(items) + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) + + return _method + + +@pytest.mark.asyncio +async def test_deduplication_merger_cursor_priority_and_cross_page() -> None: + low_items = [ + {"id": 1, "src": "low"}, + {"id": 2, "src": "low"}, + {"id": 3, "src": "low"}, + {"id": 4, "src": "low"}, + {"id": 5, "src": "low"}, + # repeats later (cross-page duplicates) + {"id": 3, "src": "low"}, + {"id": 4, "src": "low"}, + {"id": 6, "src": "low"}, + {"id": 7, "src": "low"}, + {"id": 8, "src": "low"}, + {"id": 9, "src": "low"}, + {"id": 10, "src": "low"}, + ] + high_items = [ + {"id": 3, "src": "high"}, + {"id": 4, "src": "high"}, + ] + + methods_dict = { + "low": make_offset_paged_method(low_items), + "high": make_offset_paged_method(high_items), + } + + config = { + "merger_id": "dedup_example", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 3, + "items": [ + { + "priority": 100, + "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high"}, + }, + { + "priority": 0, + "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low"}, + }, + ], + } + + merger = parse_model(DeduplicationMerger, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res_1.data) == 5 + ids_1 = [x["id"] for x in res_1.data] + assert len(ids_1) == len(set(ids_1)) + # Priority: id 3 and 4 must come from high + for x in res_1.data: + if x["id"] in {3, 4}: + assert x["src"] == "high" + + # Next page should not repeat 3/4 even though low repeats them later. + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=res_1.next_page, + ) + + ids_2 = [x["id"] for x in res_2.data] + assert not (set(ids_1) & set(ids_2)) + + # Ensure merger stores cursor state (compressed) in its own after. + assert "dedup_example" in res_2.next_page.data + assert isinstance(res_2.next_page.data["dedup_example"].after, dict) + assert "z" in res_2.next_page.data["dedup_example"].after + + +@pytest.mark.asyncio +async def test_deduplication_merger_refill_to_limit() -> None: + dup_items = [ + {"id": 1}, + {"id": 1}, + {"id": 1}, + {"id": 1}, + {"id": 1}, + {"id": 2}, + {"id": 3}, + {"id": 4}, + {"id": 5}, + {"id": 6}, + ] + + methods_dict = { + "dups": make_offset_paged_method(dup_items), + } + + config = { + "merger_id": "dedup_refill", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "overfetch_factor": 4, + "max_refill_loops": 10, + "items": [ + { + "priority": 0, + "data": {"subfeed_id": "sf_dups", "type": "subfeed", "method_name": "dups"}, + } + ], + } + + merger = parse_model(DeduplicationMerger, config) + + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage(data={}), + ) + + assert [x["id"] for x in res.data] == [1, 2, 3, 4, 5] + + +@pytest.mark.asyncio +async def test_deduplication_merger_page_zero_resets_cursor_state() -> None: + items = [{"id": i} for i in range(1, 50)] + methods_dict = {"stream": make_offset_paged_method(items)} + + config = { + "merger_id": "dedup_reset", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 2, + "items": [ + { + "priority": 0, + "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "stream"}, + } + ], + } + + merger = parse_model(DeduplicationMerger, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage(data={}), + ) + assert [x["id"] for x in res_1.data] == [1, 2, 3, 4, 5] + + # Simulate a full reload: page 0 requested again. Even if the client mistakenly + # keeps the previous "after" payload, we start a new session. + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage( + data={ + "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after) + } + ), + ) + + assert [x["id"] for x in res_2.data] == [1, 2, 3, 4, 5] + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_redis_backend(redis_client) -> None: + # This dataset repeats ids across pages (sliding window style) + items = [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + {"id": 2}, + {"id": 3}, + {"id": 4}, + {"id": 5}, + {"id": 6}, + {"id": 4}, + {"id": 7}, + {"id": 8}, + ] + + methods_dict = {"stream": make_offset_paged_method(items)} + + config = { + "merger_id": "dedup_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 4, + "items": [ + { + "priority": 0, + "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "stream"}, + } + ], + } + + merger = parse_model(DeduplicationMerger, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=4, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client, + custom_deduplication_key="t1", + ) + + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=4, + next_page=res_1.next_page, + redis_client=redis_client, + custom_deduplication_key="t1", + ) + + ids_1 = [x["id"] for x in res_1.data] + ids_2 = [x["id"] for x in res_2.data] + + assert len(ids_1) == len(set(ids_1)) + assert len(ids_2) == len(set(ids_2)) + assert not (set(ids_1) & set(ids_2)) + + # Redis backend should not store seen ids in cursor after. + assert "dedup_redis" in res_2.next_page.data + assert res_2.next_page.data["dedup_redis"].after is None + + # Ensure fixture works for both sync/async redis. + key = "dedup:dedup_redis:u:t1" + members = redis_client.smembers(key) + if inspect.iscoroutine(members): + members = await members + assert len(members) >= len(set(ids_1 + ids_2)) diff --git a/tests/test_merger_percentage.py b/tests/test_merger_percentage.py index e5ab76e..89e225e 100644 --- a/tests/test_merger_percentage.py +++ b/tests/test_merger_percentage.py @@ -3,6 +3,7 @@ from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPercentage from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_PERCENTAGE_CONFIG +from tests.utils import parse_model @pytest.mark.asyncio @@ -11,7 +12,7 @@ async def test_merger_percentage() -> None: Тест для проверки получения данных из процентного мерджера. """ - merger_percentage = MergerPercentage.parse_obj(MERGER_PERCENTAGE_CONFIG) + merger_percentage = parse_model(MergerPercentage, MERGER_PERCENTAGE_CONFIG) merger_percentage_res = await merger_percentage.get_data( methods_dict=METHODS_DICT, limit=10, diff --git a/tests/test_merger_percentage_gradient.py b/tests/test_merger_percentage_gradient.py index eaaba9c..73bc769 100644 --- a/tests/test_merger_percentage_gradient.py +++ b/tests/test_merger_percentage_gradient.py @@ -3,6 +3,7 @@ from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPercentageGradient from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_PERCENTAGE_GRADIENT_CONFIG +from tests.utils import parse_model @pytest.mark.asyncio @@ -11,7 +12,7 @@ async def test_merger_percentage_gradient() -> None: Тест для проверки получения данных из процентного мерджера с градиентом. """ - merger_percentage_gradient = MergerPercentageGradient.parse_obj(MERGER_PERCENTAGE_GRADIENT_CONFIG) + merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) merger_percentage_gradient_res = await merger_percentage_gradient.get_data( methods_dict=METHODS_DICT, limit=10, @@ -44,7 +45,7 @@ async def test_merger_percentage_gradient_next_page() -> None: Тест для проверки получения данных из процентного мерджера с градиентом после изменения процента на другой странице. """ - merger_percentage_gradient = MergerPercentageGradient.parse_obj(MERGER_PERCENTAGE_GRADIENT_CONFIG) + merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) merger_percentage_gradient_res = await merger_percentage_gradient.get_data( methods_dict=METHODS_DICT, limit=10, diff --git a/tests/test_merger_positional.py b/tests/test_merger_positional.py index c0f3815..370f770 100644 --- a/tests/test_merger_positional.py +++ b/tests/test_merger_positional.py @@ -3,6 +3,7 @@ from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPositional from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_POSITIONAL_CONFIG +from tests.utils import parse_model @pytest.mark.asyncio @@ -11,7 +12,7 @@ async def test_merger_positional_with_positions() -> None: Тест для проверки получения данных из позиционного мерджера на основе позиций в конфигурации. """ - merger_positional = MergerPositional.parse_obj(MERGER_POSITIONAL_CONFIG) + merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) merger_positional_res = await merger_positional.get_data( methods_dict=METHODS_DICT, limit=9, @@ -33,7 +34,7 @@ async def test_merger_positional_with_step() -> None: Тест для проверки получения данных из позиционного мерджера на основе шагов в конфигурации. """ - merger_positional = MergerPositional.parse_obj(MERGER_POSITIONAL_CONFIG) + merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) merger_positional_res = await merger_positional.get_data( methods_dict=METHODS_DICT, limit=10, @@ -56,7 +57,7 @@ async def test_merger_positional_with_empty_default() -> None: Тест для проверки получения данных из позиционного мерджера на основе шагов в конфигурации. """ - merger_positional = MergerPositional.parse_obj(MERGER_POSITIONAL_CONFIG) + merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) merger_positional.default.method_name = "empty" merger_positional_res = await merger_positional.get_data( methods_dict=METHODS_DICT, diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py index 78a0566..bd02bfe 100644 --- a/tests/test_merger_view_session.py +++ b/tests/test_merger_view_session.py @@ -7,6 +7,7 @@ from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG, MERGER_VIEW_SESSION_DUPS_CONFIG from tests.fixtures.redis import redis_client +from tests.utils import parse_model @pytest.mark.asyncio @@ -15,7 +16,7 @@ async def test_merger_view_session_no_redis() -> None: Тест для проверки получения данных из мерджера с кэшированием без клиента Redis. """ - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) with pytest.raises(ValueError): await merger_vs.get_data( methods_dict=METHODS_DICT, @@ -32,7 +33,7 @@ async def test_merger_view_session(redis_client) -> None: Тест для проверки получения данных из мерджера с кэшированием. """ - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) merger_vs_res = await merger_vs.get_data( methods_dict=METHODS_DICT, limit=10, @@ -59,7 +60,7 @@ async def test_merger_view_session_custom_key(redis_client) -> None: Тест для проверки получения данных из мерджера с кэшированием по ключу с кастомным постфиксом. """ - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) # Даем дополнительный параметр, который мерджер добавит в ключ кэша. merger_vs_res = await merger_vs.get_data( methods_dict=METHODS_DICT, @@ -88,7 +89,7 @@ async def test_merger_view_session_next_page(redis_client) -> None: Тест для проверки получения данных следующей страницы из мерджера с кэшированием. """ - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) merger_vs_res = await merger_vs.get_data( methods_dict=METHODS_DICT, limit=10, @@ -113,7 +114,7 @@ async def test_merger_view_session_next_page(redis_client) -> None: @pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio async def test_merger_view_session_deduplication(redis_client) -> None: - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_DUPS_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_DUPS_CONFIG) merger_vs_res = await merger_vs.get_data( methods_dict=METHODS_DICT, limit=10, diff --git a/tests/test_parsing_config.py b/tests/test_parsing_config.py index 4d789cb..10cc8f1 100644 --- a/tests/test_parsing_config.py +++ b/tests/test_parsing_config.py @@ -2,6 +2,7 @@ from smartfeed.manager import FeedManager from smartfeed.schemas import ( + DeduplicationMerger, FeedConfig, MergerAppend, MergerPercentage, @@ -11,7 +12,7 @@ MergerViewSession, SubFeed, ) -from tests.fixtures.configs import METHODS_DICT, PARSING_CONFIG_FIXTURE +from tests.fixtures.configs import METHODS_DICT, PARSING_CONFIG_FIXTURE, PARSING_DEDUP_CONFIG_FIXTURE @pytest.mark.asyncio @@ -45,3 +46,14 @@ async def test_parsing_config() -> None: # SubFeed with Raise Exception False. assert isinstance(feed_manager.feed_config.feed.default.items[0].data, SubFeed) assert feed_manager.feed_config.feed.default.items[0].data.raise_error is False + + +@pytest.mark.asyncio +async def test_parsing_config_deduplication_merger() -> None: + feed_manager = FeedManager(config=PARSING_DEDUP_CONFIG_FIXTURE, methods_dict=METHODS_DICT) + + assert isinstance(feed_manager.feed_config, FeedConfig) + assert isinstance(feed_manager.feed_config.feed, DeduplicationMerger) + assert len(feed_manager.feed_config.feed.items) == 2 + assert feed_manager.feed_config.feed.items[0].priority == 100 + assert isinstance(feed_manager.feed_config.feed.items[0].data, SubFeed) diff --git a/tests/test_redis_live.py b/tests/test_redis_live.py index 1a23839..85effa2 100644 --- a/tests/test_redis_live.py +++ b/tests/test_redis_live.py @@ -8,6 +8,7 @@ from smartfeed.schemas import FeedResultNextPage, MergerViewSession from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG +from tests.utils import parse_model class RedisReplicationSimulator: @@ -63,7 +64,7 @@ async def test_redis_replication_delay_problem(): # Используем симулятор задержки репликации redis_client = RedisReplicationSimulator(real_client) - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) print("\n=== Демонстрация проблемы с задержкой репликации ===") @@ -117,7 +118,7 @@ async def test_redis_multiple_requests(): real_client.delete(test_key) redis_client = RedisReplicationSimulator(real_client) - merger_vs = MergerViewSession.parse_obj(MERGER_VIEW_SESSION_CONFIG) + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) print("\n=== Тест множественных запросов ===") diff --git a/tests/test_sub_feed.py b/tests/test_sub_feed.py index 7da1924..11645d4 100644 --- a/tests/test_sub_feed.py +++ b/tests/test_sub_feed.py @@ -16,7 +16,7 @@ async def test_sub_feed() -> None: Тест для проверки получения данных из субфида (без параметров). """ - sub_feed = SubFeed.parse_obj(SUBFEED_CONFIG) + sub_feed = SubFeed.model_validate(SUBFEED_CONFIG) sub_feed_data = await sub_feed.get_data( methods_dict=METHODS_DICT, limit=15, @@ -33,7 +33,7 @@ async def test_sub_feed_with_params() -> None: Тест для проверки получения данных из субфида (с параметрами). """ - sub_feed = SubFeed.parse_obj(SUBFEED_WITH_PARAMS_CONFIG) + sub_feed = SubFeed.model_validate(SUBFEED_WITH_PARAMS_CONFIG) sub_feed_data = await sub_feed.get_data( methods_dict=METHODS_DICT, limit=15, @@ -50,7 +50,7 @@ async def test_sub_feed_raise_error() -> None: Тест для проверки получения данных из субфида (без параметров). """ - sub_feed = SubFeed.parse_obj(SUBFEED_CONFIG_RAISE_ERROR) + sub_feed = SubFeed.model_validate(SUBFEED_CONFIG_RAISE_ERROR) with pytest.raises(Exception): await sub_feed.get_data( @@ -67,7 +67,7 @@ async def test_sub_feed_no_raise_error() -> None: Тест для проверки получения данных из субфида (без параметров). """ - sub_feed = SubFeed.parse_obj(SUBFEED_CONFIG_NO_RAISE_ERROR) + sub_feed = SubFeed.model_validate(SUBFEED_CONFIG_NO_RAISE_ERROR) sub_feed_data = await sub_feed.get_data( methods_dict=METHODS_DICT, limit=15, diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..f38689f --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from typing import Any, Dict, Type, TypeVar + + +T = TypeVar("T") + + +def parse_model(model_cls: Type[T], obj: Dict[str, Any]) -> T: + """Parse a dict into a Pydantic model. + + Uses Pydantic v2 `model_validate` when available, otherwise falls back to v1 `parse_obj`. + """ + + if hasattr(model_cls, "model_validate"): + return model_cls.model_validate(obj) # type: ignore[attr-defined] + return model_cls.parse_obj(obj) # type: ignore[attr-defined] From 1967175f8dd9bafe6682b89444ac56d3aba840e6 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 17:07:25 +0000 Subject: [PATCH 02/41] Name fix. --- smartfeed/schemas.py | 16 ++++++++-------- tests/test_merger_deduplication.py | 10 +++++----- tests/test_parsing_config.py | 4 ++-- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index a9ef784..5a86a3a 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -15,7 +15,7 @@ FeedTypes = Annotated[ Union[ - "DeduplicationMerger", + "MergerDeduplication", "MergerAppend", "MergerAppendDistribute", "MergerPositional", @@ -1021,14 +1021,14 @@ async def get_data( return result -class DeduplicationMergerItem(BaseModel): - """Configuration item for DeduplicationMerger.""" +class MergerDeduplicationItem(BaseModel): + """Configuration item for MergerDeduplication.""" priority: int = 0 data: FeedTypes -class DeduplicationMerger(BaseFeedConfigModel): +class MergerDeduplication(BaseFeedConfigModel): """Merger that deduplicates items and refills to the requested limit. Key properties: @@ -1039,7 +1039,7 @@ class DeduplicationMerger(BaseFeedConfigModel): merger_id: str type: Literal["merger_deduplication"] - items: List[DeduplicationMergerItem] + items: List[MergerDeduplicationItem] dedup_key: Optional[str] = None missing_key_policy: Literal["error", "keep", "drop"] = "error" @@ -1151,7 +1151,7 @@ async def get_data( is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) if self.state_backend == "redis" and not redis_client: - raise ValueError("Redis client must be provided if using DeduplicationMerger with state_backend=redis") + raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") if hasattr(next_page, "model_copy"): working_next_page = next_page.model_copy(deep=True) # type: ignore[attr-defined] @@ -1381,5 +1381,5 @@ def _rebuild_model(model: Any) -> None: _rebuild_model(MergerAppendDistribute) _rebuild_model(MergerPercentageGradient) _rebuild_model(MergerViewSession) -_rebuild_model(DeduplicationMergerItem) -_rebuild_model(DeduplicationMerger) +_rebuild_model(MergerDeduplicationItem) +_rebuild_model(MergerDeduplication) \ No newline at end of file diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 5e16c45..26b4c78 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -3,7 +3,7 @@ import pytest from smartfeed.schemas import ( - DeduplicationMerger, + MergerDeduplication, FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, @@ -71,7 +71,7 @@ async def test_deduplication_merger_cursor_priority_and_cross_page() -> None: ], } - merger = parse_model(DeduplicationMerger, config) + merger = parse_model(MergerDeduplication, config) res_1 = await merger.get_data( methods_dict=methods_dict, @@ -139,7 +139,7 @@ async def test_deduplication_merger_refill_to_limit() -> None: ], } - merger = parse_model(DeduplicationMerger, config) + merger = parse_model(MergerDeduplication, config) res = await merger.get_data( methods_dict=methods_dict, @@ -171,7 +171,7 @@ async def test_deduplication_merger_page_zero_resets_cursor_state() -> None: ], } - merger = parse_model(DeduplicationMerger, config) + merger = parse_model(MergerDeduplication, config) res_1 = await merger.get_data( methods_dict=methods_dict, @@ -232,7 +232,7 @@ async def test_deduplication_merger_redis_backend(redis_client) -> None: ], } - merger = parse_model(DeduplicationMerger, config) + merger = parse_model(MergerDeduplication, config) res_1 = await merger.get_data( methods_dict=methods_dict, diff --git a/tests/test_parsing_config.py b/tests/test_parsing_config.py index 10cc8f1..291971b 100644 --- a/tests/test_parsing_config.py +++ b/tests/test_parsing_config.py @@ -2,7 +2,7 @@ from smartfeed.manager import FeedManager from smartfeed.schemas import ( - DeduplicationMerger, + MergerDeduplication, FeedConfig, MergerAppend, MergerPercentage, @@ -53,7 +53,7 @@ async def test_parsing_config_deduplication_merger() -> None: feed_manager = FeedManager(config=PARSING_DEDUP_CONFIG_FIXTURE, methods_dict=METHODS_DICT) assert isinstance(feed_manager.feed_config, FeedConfig) - assert isinstance(feed_manager.feed_config.feed, DeduplicationMerger) + assert isinstance(feed_manager.feed_config.feed, MergerDeduplication) assert len(feed_manager.feed_config.feed.items) == 2 assert feed_manager.feed_config.feed.items[0].priority == 100 assert isinstance(feed_manager.feed_config.feed.items[0].data, SubFeed) From 60f1a77f963c4a3b08aacca9b6987bedf5d78ca9 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 17:47:46 +0000 Subject: [PATCH 03/41] More tests. --- smartfeed/schemas.py | 45 ++ tests/test_merger_deduplication.py | 754 ++++++++++++++++++++++++++++- 2 files changed, 795 insertions(+), 4 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 5a86a3a..c5e58bb 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1052,6 +1052,43 @@ class MergerDeduplication(BaseFeedConfigModel): overfetch_factor: int = 2 max_refill_loops: int = 20 + def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: + keys: set[str] = set() + + subfeed_id = getattr(feed, "subfeed_id", None) + if isinstance(subfeed_id, str) and subfeed_id: + keys.add(subfeed_id) + + merger_id = getattr(feed, "merger_id", None) + if isinstance(merger_id, str) and merger_id: + keys.add(merger_id) + + # Recurse into known child containers across existing feed types. + child: Any + for attr_name in ("data", "positional", "default"): + child = getattr(feed, attr_name, None) + if isinstance(child, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(child)) + + for attr_name in ("item_from", "item_to"): + child = getattr(feed, attr_name, None) + inner = getattr(child, "data", None) + if isinstance(inner, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(inner)) + + items = getattr(feed, "items", None) + if isinstance(items, list): + for item in items: + if isinstance(item, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(item)) + continue + + inner = getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(inner)) + + return keys + def _normalize_key(self, value: Any) -> str: if isinstance(value, (str, int)): return str(value) @@ -1157,6 +1194,14 @@ async def get_data( working_next_page = next_page.model_copy(deep=True) # type: ignore[attr-defined] else: working_next_page = next_page.copy(deep=True) + + if is_fresh_session: + # Reset cursors for all descendants under this merger so upstream nodes also restart. + descendant_keys: set[str] = set() + for item in self.items: + descendant_keys.update(self._collect_descendant_cursor_keys(item.data)) + for key in descendant_keys: + working_next_page.data.pop(key, None) sorted_items = sorted(self.items, key=lambda x: x.priority, reverse=True) seen_keys_in_order: List[str] = [] diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 26b4c78..cede8f3 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -13,10 +13,13 @@ from tests.utils import parse_model -def make_offset_paged_method(items): +def make_offset_paged_method(items, *, max_per_call=None): async def _method(user_id, limit, next_page): # pylint: disable=unused-argument offset = int(next_page.after or 0) - result_data = items[offset : offset + limit] + effective_limit = limit + if isinstance(max_per_call, int) and max_per_call > 0: + effective_limit = min(effective_limit, max_per_call) + result_data = items[offset : offset + effective_limit] next_page.after = offset + len(result_data) next_page.page += 1 has_next_page = (offset + len(result_data)) < len(items) @@ -25,6 +28,90 @@ async def _method(user_id, limit, next_page): # pylint: disable=unused-argument return _method +async def _run_two_pages( + *, + config, + methods_dict, + user_id, + limit, + redis_client_instance=None, + **params, +): + merger = parse_model(MergerDeduplication, config) + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client_instance, + **params, + ) + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=res_1.next_page, + redis_client=redis_client_instance, + **params, + ) + return res_1, res_2 + + +def _assert_dedup_backend_state(*, res, merger_id: str, state_backend: str) -> None: + assert merger_id in res.next_page.data + if state_backend == "cursor": + assert isinstance(res.next_page.data[merger_id].after, dict) + else: + assert res.next_page.data[merger_id].after is None + + +def _ids(data): + return [x["id"] for x in data] + + +def _assert_two_pages_no_overlap(res_1, res_2): + ids_1 = set(_ids(res_1.data)) + ids_2 = set(_ids(res_2.data)) + assert len(ids_1) == len(res_1.data) + assert len(ids_2) == len(res_2.data) + assert not (ids_1 & ids_2) + + +def _assert_cursor_monotonic_if_present(res_1, res_2, keys): + """Assert that cursor values monotonically advance for keys that are present. + + MergerDeduplication may stop early once it has enough unique items, so a + descendant might not be called on a given page. This helper only asserts + monotonicity when the cursor key exists in `res_1`. + """ + + for key in keys: + if key not in res_1.next_page.data: + continue + + assert key in res_2.next_page.data + + after_1 = res_1.next_page.data[key].after + after_2 = res_2.next_page.data[key].after + + if after_1 is None or after_2 is None: + continue + + if isinstance(after_1, int) and isinstance(after_2, int): + assert after_2 >= after_1 + continue + + # Merger cursors can be structured (dict), just require presence. + if isinstance(after_1, dict) and isinstance(after_2, dict): + continue + + # If values are comparable, enforce monotonicity; otherwise don't fail. + try: + assert after_2 >= after_1 + except TypeError: + pass + + @pytest.mark.asyncio async def test_deduplication_merger_cursor_priority_and_cross_page() -> None: low_items = [ @@ -182,14 +269,15 @@ async def test_deduplication_merger_page_zero_resets_cursor_state() -> None: assert [x["id"] for x in res_1.data] == [1, 2, 3, 4, 5] # Simulate a full reload: page 0 requested again. Even if the client mistakenly - # keeps the previous "after" payload, we start a new session. + # keeps the previous cursor payloads (including subfeed cursors), we start a new session. res_2 = await merger.get_data( methods_dict=methods_dict, user_id="u", limit=5, next_page=FeedResultNextPage( data={ - "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after) + "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after), + "sf_stream": res_1.next_page.data["sf_stream"], } ), ) @@ -269,3 +357,661 @@ async def test_deduplication_merger_redis_backend(redis_client) -> None: if inspect.iscoroutine(members): members = await members assert len(members) >= len(set(ids_1 + ids_2)) + + +@pytest.mark.asyncio +async def test_deduplication_merger_priority_replacement_across_loops_cursor_backend() -> None: + # This test forces the higher-priority source to surface a duplicate only on a later call, + # so we exercise the in-page replacement logic. + # Important: dedup calls sources in descending priority. To ensure we exercise + # replacement, we need a lower-priority source to introduce id=5 *before* + # the high-priority source sees id=5 on a later refill loop. + low_items = [ + {"id": 5, "src": "low"}, + {"id": 6, "src": "low"}, + {"id": 7, "src": "low"}, + {"id": 99, "src": "low"}, + ] + mid_items = [ + {"id": 5, "src": "mid"}, + {"id": 98, "src": "mid"}, + {"id": 8, "src": "mid"}, + {"id": 9, "src": "mid"}, + ] + high_items = [ + {"id": 1, "src": "high"}, + {"id": 5, "src": "high"}, + {"id": 2, "src": "high"}, + {"id": 3, "src": "high"}, + ] + + methods_dict = { + "low": make_offset_paged_method(low_items, max_per_call=1), + "mid": make_offset_paged_method(mid_items, max_per_call=1), + "high": make_offset_paged_method(high_items, max_per_call=1), + } + + config = { + "merger_id": "dedup_priority_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 1, + "max_refill_loops": 10, + "items": [ + {"priority": 100, "data": {"subfeed_id": "sf_high_p", "type": "subfeed", "method_name": "high"}}, + {"priority": 50, "data": {"subfeed_id": "sf_mid_p", "type": "subfeed", "method_name": "mid"}}, + {"priority": 0, "data": {"subfeed_id": "sf_low_p", "type": "subfeed", "method_name": "low"}}, + ], + } + + res_1 = await parse_model(MergerDeduplication, config).get_data( + methods_dict=methods_dict, + user_id="u", + limit=4, + next_page=FeedResultNextPage(data={}), + ) + + # Ensure id=5 is present and comes from highest priority, even though low/mid can surface it earlier. + winners = {x["id"]: x["src"] for x in res_1.data} + assert winners[5] == "high" + _assert_dedup_backend_state(res=res_1, merger_id="dedup_priority_cursor", state_backend="cursor") + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_priority_replacement_across_loops_redis_backend(redis_client) -> None: + low_items = [ + {"id": 5, "src": "low"}, + {"id": 6, "src": "low"}, + {"id": 7, "src": "low"}, + {"id": 99, "src": "low"}, + ] + mid_items = [ + {"id": 5, "src": "mid"}, + {"id": 98, "src": "mid"}, + {"id": 8, "src": "mid"}, + {"id": 9, "src": "mid"}, + ] + high_items = [ + {"id": 1, "src": "high"}, + {"id": 5, "src": "high"}, + {"id": 2, "src": "high"}, + {"id": 3, "src": "high"}, + ] + + methods_dict = { + "low": make_offset_paged_method(low_items, max_per_call=1), + "mid": make_offset_paged_method(mid_items, max_per_call=1), + "high": make_offset_paged_method(high_items, max_per_call=1), + } + + config = { + "merger_id": "dedup_priority_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 1, + "max_refill_loops": 10, + "items": [ + {"priority": 100, "data": {"subfeed_id": "sf_high_pr", "type": "subfeed", "method_name": "high"}}, + {"priority": 50, "data": {"subfeed_id": "sf_mid_pr", "type": "subfeed", "method_name": "mid"}}, + {"priority": 0, "data": {"subfeed_id": "sf_low_pr", "type": "subfeed", "method_name": "low"}}, + ], + } + + res_1 = await parse_model(MergerDeduplication, config).get_data( + methods_dict=methods_dict, + user_id="u", + limit=4, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client, + custom_deduplication_key="priority", + ) + + winners = {x["id"]: x["src"] for x in res_1.data} + assert winners[5] == "high" + _assert_dedup_backend_state(res=res_1, merger_id="dedup_priority_redis", state_backend="redis") + + +@pytest.mark.asyncio +async def test_deduplication_merger_with_append_and_three_sources_cursor_backend() -> None: + # Inner MergerAppend (two subfeeds) + two extra subfeeds as separate dedup items. + a_items = [{"id": i, "src": "a"} for i in range(1, 30)] + b_items = [{"id": i, "src": "b"} for i in range(10, 40)] + c_items = [{"id": i, "src": "c"} for i in range(20, 60)] + d_items = [{"id": i, "src": "d"} for i in range(25, 70)] + + # Cap each subfeed to 1 item per call so dedup must invoke all children + # (and therefore exercise nested cursor propagation). + methods_dict = { + "a": make_offset_paged_method(a_items, max_per_call=1), + "b": make_offset_paged_method(b_items, max_per_call=1), + "c": make_offset_paged_method(c_items, max_per_call=1), + "d": make_offset_paged_method(d_items, max_per_call=1), + } + + append_config = { + "merger_id": "inner_append_unused", + "type": "merger_append", + "items": [ + {"subfeed_id": "sf_a_append", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b_append", "type": "subfeed", "method_name": "b"}, + ], + } + + config = { + "merger_id": "dedup_with_append_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": append_config}, + {"priority": 5, "data": {"subfeed_id": "sf_c", "type": "subfeed", "method_name": "c"}}, + {"priority": 0, "data": {"subfeed_id": "sf_d", "type": "subfeed", "method_name": "d"}}, + ], + } + + res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=15) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_with_append_cursor", state_backend="cursor") + + # Cursor correctness: descendant subfeed cursors exist and advance. + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_a_append", "sf_b_append", "sf_c", "sf_d"]) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_append_and_three_sources_redis_backend(redis_client) -> None: + a_items = [{"id": i, "src": "a"} for i in range(1, 30)] + b_items = [{"id": i, "src": "b"} for i in range(10, 40)] + c_items = [{"id": i, "src": "c"} for i in range(20, 60)] + d_items = [{"id": i, "src": "d"} for i in range(25, 70)] + + methods_dict = { + "a": make_offset_paged_method(a_items, max_per_call=1), + "b": make_offset_paged_method(b_items, max_per_call=1), + "c": make_offset_paged_method(c_items, max_per_call=1), + "d": make_offset_paged_method(d_items, max_per_call=1), + } + + append_config = { + "merger_id": "inner_append_unused_r", + "type": "merger_append", + "items": [ + {"subfeed_id": "sf_a_append_r", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b_append_r", "type": "subfeed", "method_name": "b"}, + ], + } + + config = { + "merger_id": "dedup_with_append_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": append_config}, + {"priority": 5, "data": {"subfeed_id": "sf_c_r", "type": "subfeed", "method_name": "c"}}, + {"priority": 0, "data": {"subfeed_id": "sf_d_r", "type": "subfeed", "method_name": "d"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=15, + redis_client_instance=redis_client, + custom_deduplication_key="append", + ) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_with_append_redis", state_backend="redis") + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_a_append_r", "sf_b_append_r", "sf_c_r", "sf_d_r"]) + + +@pytest.mark.asyncio +async def test_deduplication_merger_with_percentage_cursor_backend() -> None: + a_items = [{"id": i, "src": "pa"} for i in range(1, 60)] + b_items = [{"id": i, "src": "pb"} for i in range(30, 90)] + c_items = [{"id": i, "src": "pc"} for i in range(40, 120)] + + methods_dict = { + "pa": make_offset_paged_method(a_items, max_per_call=1), + "pb": make_offset_paged_method(b_items, max_per_call=1), + "pc": make_offset_paged_method(c_items, max_per_call=1), + } + + percentage_config = { + "merger_id": "inner_percentage_unused", + "type": "merger_percentage", + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_pa", "type": "subfeed", "method_name": "pa"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_pb", "type": "subfeed", "method_name": "pb"}}, + ], + } + + config = { + "merger_id": "dedup_percentage_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "overfetch_factor": 2, + "items": [ + {"priority": 0, "data": percentage_config}, + {"priority": 10, "data": {"subfeed_id": "sf_pc", "type": "subfeed", "method_name": "pc"}}, + {"priority": 5, "data": {"subfeed_id": "sf_pd", "type": "subfeed", "method_name": "pa"}}, + ], + } + + res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=20) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_percentage_cursor", state_backend="cursor") + + for key in ("sf_pa", "sf_pb", "sf_pc"): + assert key in res_1.next_page.data + assert isinstance(res_1.next_page.data[key].after, int) + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_pa", "sf_pb", "sf_pc", "sf_pd"]) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_percentage_redis_backend(redis_client) -> None: + a_items = [{"id": i, "src": "pa"} for i in range(1, 60)] + b_items = [{"id": i, "src": "pb"} for i in range(30, 90)] + c_items = [{"id": i, "src": "pc"} for i in range(40, 120)] + + methods_dict = { + "pa": make_offset_paged_method(a_items, max_per_call=1), + "pb": make_offset_paged_method(b_items, max_per_call=1), + "pc": make_offset_paged_method(c_items, max_per_call=1), + } + + percentage_config = { + "merger_id": "inner_percentage_unused_r", + "type": "merger_percentage", + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_pa_r", "type": "subfeed", "method_name": "pa"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_pb_r", "type": "subfeed", "method_name": "pb"}}, + ], + } + + config = { + "merger_id": "dedup_percentage_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 0, "data": percentage_config}, + {"priority": 10, "data": {"subfeed_id": "sf_pc_r", "type": "subfeed", "method_name": "pc"}}, + {"priority": 5, "data": {"subfeed_id": "sf_pd_r", "type": "subfeed", "method_name": "pa"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=20, + redis_client_instance=redis_client, + custom_deduplication_key="percentage", + ) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_percentage_redis", state_backend="redis") + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_pa_r", "sf_pb_r", "sf_pc_r", "sf_pd_r"]) + + +@pytest.mark.asyncio +async def test_deduplication_merger_with_positional_cursor_backend() -> None: + # MergerPositional carries its own merger cursor; verify it survives nesting in dedup. + pos_items = [{"id": i, "src": "pos"} for i in range(1, 100)] + def_items = [{"id": i, "src": "def"} for i in range(50, 140)] + extra_items = [{"id": i, "src": "extra"} for i in range(80, 180)] + + methods_dict = { + "pos": make_offset_paged_method(pos_items, max_per_call=1), + "def": make_offset_paged_method(def_items, max_per_call=1), + "extra": make_offset_paged_method(extra_items, max_per_call=1), + } + + positional_config = { + "merger_id": "inner_positional", + "type": "merger_positional", + "positions": [0, 2, 4, 6, 8], + "positional": {"subfeed_id": "sf_positional", "type": "subfeed", "method_name": "pos"}, + "default": {"subfeed_id": "sf_default", "type": "subfeed", "method_name": "def"}, + } + + config = { + "merger_id": "dedup_positional_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "overfetch_factor": 2, + "items": [ + # Positional must run; it owns its own merger cursor entry. + {"priority": 10, "data": positional_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra2", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=20) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_positional_cursor", state_backend="cursor") + assert "inner_positional" in res_1.next_page.data + assert "inner_positional" in res_2.next_page.data + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_positional", "sf_default", "sf_extra", "sf_extra2"]) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_positional_redis_backend(redis_client) -> None: + pos_items = [{"id": i, "src": "pos"} for i in range(1, 100)] + def_items = [{"id": i, "src": "def"} for i in range(50, 140)] + extra_items = [{"id": i, "src": "extra"} for i in range(80, 180)] + + methods_dict = { + "pos": make_offset_paged_method(pos_items, max_per_call=1), + "def": make_offset_paged_method(def_items, max_per_call=1), + "extra": make_offset_paged_method(extra_items, max_per_call=1), + } + + positional_config = { + "merger_id": "inner_positional_r", + "type": "merger_positional", + "positions": [0, 2, 4, 6, 8], + "positional": {"subfeed_id": "sf_positional_r", "type": "subfeed", "method_name": "pos"}, + "default": {"subfeed_id": "sf_default_r", "type": "subfeed", "method_name": "def"}, + } + + config = { + "merger_id": "dedup_positional_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": positional_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_r", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra2_r", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=20, + redis_client_instance=redis_client, + custom_deduplication_key="positional", + ) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_positional_redis", state_backend="redis") + assert "inner_positional_r" in res_1.next_page.data + + _assert_cursor_monotonic_if_present( + res_1, + res_2, + ["sf_positional_r", "sf_default_r", "sf_extra_r", "sf_extra2_r"], + ) + + +@pytest.mark.asyncio +async def test_deduplication_merger_with_percentage_gradient_cursor_backend() -> None: + from_items = [{"id": i, "src": "from"} for i in range(1, 140)] + to_items = [{"id": i, "src": "to"} for i in range(60, 200)] + extra_items = [{"id": i, "src": "extra"} for i in range(120, 300)] + + methods_dict = { + "from": make_offset_paged_method(from_items, max_per_call=1), + "to": make_offset_paged_method(to_items, max_per_call=1), + "extra": make_offset_paged_method(extra_items, max_per_call=1), + } + + gradient_config = { + "merger_id": "inner_gradient", + "type": "merger_percentage_gradient", + "item_from": {"percentage": 80, "data": {"subfeed_id": "sf_from", "type": "subfeed", "method_name": "from"}}, + "item_to": {"percentage": 20, "data": {"subfeed_id": "sf_to", "type": "subfeed", "method_name": "to"}}, + "step": 10, + "size_to_step": 10, + } + + config = { + "merger_id": "dedup_gradient_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": gradient_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_g", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra_g2", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=25) + _assert_two_pages_no_overlap(res_1, res_2) + assert "inner_gradient" in res_1.next_page.data + assert "inner_gradient" in res_2.next_page.data + _assert_dedup_backend_state(res=res_2, merger_id="dedup_gradient_cursor", state_backend="cursor") + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_from", "sf_to", "sf_extra_g", "sf_extra_g2"]) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_percentage_gradient_redis_backend(redis_client) -> None: + from_items = [{"id": i, "src": "from"} for i in range(1, 140)] + to_items = [{"id": i, "src": "to"} for i in range(60, 200)] + extra_items = [{"id": i, "src": "extra"} for i in range(120, 300)] + + methods_dict = { + "from": make_offset_paged_method(from_items, max_per_call=1), + "to": make_offset_paged_method(to_items, max_per_call=1), + "extra": make_offset_paged_method(extra_items, max_per_call=1), + } + + gradient_config = { + "merger_id": "inner_gradient_r", + "type": "merger_percentage_gradient", + "item_from": {"percentage": 80, "data": {"subfeed_id": "sf_from_r", "type": "subfeed", "method_name": "from"}}, + "item_to": {"percentage": 20, "data": {"subfeed_id": "sf_to_r", "type": "subfeed", "method_name": "to"}}, + "step": 10, + "size_to_step": 10, + } + + config = { + "merger_id": "dedup_gradient_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": gradient_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_gr", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra_gr2", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=25, + redis_client_instance=redis_client, + custom_deduplication_key="gradient", + ) + _assert_two_pages_no_overlap(res_1, res_2) + assert "inner_gradient_r" in res_1.next_page.data + _assert_dedup_backend_state(res=res_2, merger_id="dedup_gradient_redis", state_backend="redis") + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_from_r", "sf_to_r", "sf_extra_gr", "sf_extra_gr2"]) + + +@pytest.mark.parametrize("state_backend", ["cursor", "redis"]) +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_view_session_child(state_backend, redis_client) -> None: + # MergerViewSession always requires Redis, so this test always uses redis_client. + # We still validate both dedup state backends. + base_items = [{"id": i, "src": "vs"} for i in range(1, 200)] + extra_items = [{"id": i, "src": "extra"} for i in range(50, 260)] + + methods_dict = { + "vs": make_offset_paged_method(base_items), + "extra": make_offset_paged_method(extra_items), + } + + view_session_config = { + "merger_id": "inner_view_session", + "type": "merger_view_session", + "session_size": 60, + "session_live_time": 60, + "data": {"subfeed_id": "sf_vs", "type": "subfeed", "method_name": "vs"}, + "deduplicate": True, + "dedup_key": "id", + } + + config = { + "merger_id": f"dedup_vs_{state_backend}", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": state_backend, + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": view_session_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_vs", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra_vs2", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=20, + redis_client_instance=redis_client, + custom_deduplication_key=f"vs_{state_backend}", + custom_view_session_key=f"vs_{state_backend}", + ) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id=f"dedup_vs_{state_backend}", state_backend=state_backend) + assert "inner_view_session" in res_1.next_page.data + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_vs", "sf_extra_vs", "sf_extra_vs2"]) + + +@pytest.mark.asyncio +async def test_deduplication_merger_with_append_distribute_cursor_backend() -> None: + # MergerAppendDistribute (type merger_distribute) + two extra subfeeds. + s1 = [{"id": i, "src": "s1", "group": "g1" if i % 2 == 0 else "g2"} for i in range(1, 120)] + s2 = [{"id": i, "src": "s2", "group": "g2" if i % 3 == 0 else "g3"} for i in range(60, 200)] + extra = [{"id": i, "src": "extra", "group": "g9"} for i in range(100, 240)] + + methods_dict = { + "s1": make_offset_paged_method(s1, max_per_call=1), + "s2": make_offset_paged_method(s2, max_per_call=1), + "extra": make_offset_paged_method(extra, max_per_call=1), + } + + distribute_config = { + "merger_id": "inner_distribute_unused", + "type": "merger_distribute", + "distribution_key": "group", + "items": [ + {"subfeed_id": "sf_s1", "type": "subfeed", "method_name": "s1"}, + {"subfeed_id": "sf_s2", "type": "subfeed", "method_name": "s2"}, + ], + } + + config = { + "merger_id": "dedup_distribute_cursor", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": distribute_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_dist", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra_dist2", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=25) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_distribute_cursor", state_backend="cursor") + for key in ("sf_s1", "sf_s2"): + assert key in res_1.next_page.data + + _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_s1", "sf_s2", "sf_extra_dist", "sf_extra_dist2"]) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_deduplication_merger_with_append_distribute_redis_backend(redis_client) -> None: + s1 = [{"id": i, "src": "s1", "group": "g1" if i % 2 == 0 else "g2"} for i in range(1, 120)] + s2 = [{"id": i, "src": "s2", "group": "g2" if i % 3 == 0 else "g3"} for i in range(60, 200)] + extra = [{"id": i, "src": "extra", "group": "g9"} for i in range(100, 240)] + + methods_dict = { + "s1": make_offset_paged_method(s1, max_per_call=1), + "s2": make_offset_paged_method(s2, max_per_call=1), + "extra": make_offset_paged_method(extra, max_per_call=1), + } + + distribute_config = { + "merger_id": "inner_distribute_unused_r", + "type": "merger_distribute", + "distribution_key": "group", + "items": [ + {"subfeed_id": "sf_s1_r", "type": "subfeed", "method_name": "s1"}, + {"subfeed_id": "sf_s2_r", "type": "subfeed", "method_name": "s2"}, + ], + } + + config = { + "merger_id": "dedup_distribute_redis", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "overfetch_factor": 2, + "items": [ + {"priority": 10, "data": distribute_config}, + {"priority": 5, "data": {"subfeed_id": "sf_extra_dist_r", "type": "subfeed", "method_name": "extra"}}, + {"priority": 0, "data": {"subfeed_id": "sf_extra_dist2_r", "type": "subfeed", "method_name": "extra"}}, + ], + } + + res_1, res_2 = await _run_two_pages( + config=config, + methods_dict=methods_dict, + user_id="u", + limit=25, + redis_client_instance=redis_client, + custom_deduplication_key="distribute", + ) + _assert_two_pages_no_overlap(res_1, res_2) + _assert_dedup_backend_state(res=res_2, merger_id="dedup_distribute_redis", state_backend="redis") + + _assert_cursor_monotonic_if_present( + res_1, + res_2, + ["sf_s1_r", "sf_s2_r", "sf_extra_dist_r", "sf_extra_dist2_r"], + ) From 29bd63607de450b3073c594b5e6a51ed91ad88e7 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 21:12:06 +0000 Subject: [PATCH 04/41] dedup WIP. --- smartfeed/schemas.py | 635 ++++++++++----- tests/fixtures/configs.py | 37 +- tests/test_merger_deduplication.py | 1172 ++++++++++------------------ tests/test_parsing_config.py | 5 +- 4 files changed, 867 insertions(+), 982 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index c5e58bb..3320fb7 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -87,6 +87,10 @@ class BaseFeedConfigModel(ABC, BaseModel): Абстрактный класс для мерджера / субфида конфигурации. """ + # Higher value means the item should "win" deduplication when duplicates exist. + # This is primarily used by MergerDeduplication and by mergers when a dedup wrapper is active. + dedup_priority: int = 0 + @abstractmethod async def get_data( self, @@ -441,37 +445,60 @@ async def get_data( :return: список данных методом append. """ - # Формируем результат append мерджера. + # When a MergerDeduplication wrapper is active, we may need to respect dedup_priority + # across children without changing the append output order. In that mode we fetch in + # priority order, then concatenate in the configured order and trim to `limit`. + dedup_active = bool(params.pop("_sf_dedup_active", False)) + result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - result_limit = limit - for item in self.items: - # Получаем данные из позиции мерджера. - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) + if dedup_active: + indexed_items = list(enumerate(self.items)) + fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) + fetched: Dict[int, FeedResult] = {} - # Добавляем данные позиции к общему результату процентного мерджера. - result.data.extend(item_result.data) + for idx, item in fetch_order: + fetched[idx] = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) - # Обновляем result_limit - result_limit -= len(item_result.data) + for idx, _item in indexed_items: + item_result = fetched[idx] + result.data.extend(item_result.data) + result.next_page.data.update(item_result.next_page.data) + if item_result.has_next_page: + result.has_next_page = True - # Если has_next_page = False, то проверяем has_next_page у позиции и, если необходимо, обновляем. - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True + if len(result.data) > limit: + result.data = result.data[:limit] + else: + result_limit = limit + for item in self.items: + item_result = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=result_limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) - # Обновляем next_page. - result.next_page.data.update(item_result.next_page.data) + result.data.extend(item_result.data) + result_limit -= len(item_result.data) + + if not result.has_next_page and item_result.has_next_page: + result.has_next_page = True - # Если полученных данных хватает, то прерываем итерацию и возвращаем результат. - if result_limit <= 0: - break + result.next_page.data.update(item_result.next_page.data) + + if result_limit <= 0: + break # Если в конфигурации указано "смешать" данные. if self.shuffle: @@ -537,37 +564,14 @@ async def get_data( :return: список данных в процентном соотношении. """ - # Получаем данные "default". - default_res = await self.default.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) + dedup_active = bool(params.pop("_sf_dedup_active", False)) - # Формируем результат позиционного мерджера. - result = FeedResult( - data=default_res.data, - next_page=FeedResultNextPage( - data={ - self.merger_id: FeedResultNextPageInside( - page=next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1, - after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, - ) - }, - ), - has_next_page=default_res.has_next_page, - ) + # Determine the merger page first (independent of children). + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - # Получаем список позиций с учетом текущей страницы. positional_has_next_page = True - page_positions = [] - available_positions = range( - (result.next_page.data[self.merger_id].page - 1) * limit, - (result.next_page.data[self.merger_id].page * limit) + 1, - ) + page_positions: List[int] = [] + available_positions = range((page - 1) * limit, (page * limit) + 1) for position in self.positions: if position in available_positions: page_positions.append(available_positions.index(position)) @@ -584,14 +588,59 @@ async def get_data( if position in available_positions: page_positions.append(available_positions.index(position)) - # Получаем данные "positional". - pos_res = await self.positional.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=len(page_positions), - next_page=next_page, - redis_client=redis_client, - **params, + default_res: FeedResult + pos_res: FeedResult + + if dedup_active and getattr(self.positional, "dedup_priority", 0) > getattr(self.default, "dedup_priority", 0): + pos_res = await self.positional.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=len(page_positions), + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + default_res = await self.default.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + else: + default_res = await self.default.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + pos_res = await self.positional.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=len(page_positions), + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + + result = FeedResult( + data=default_res.data, + next_page=FeedResultNextPage( + data={ + self.merger_id: FeedResultNextPageInside( + page=page, + after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, + ) + }, + ), + has_next_page=default_res.has_next_page, ) # Если has_next_page = False, то проверяем has_next_page у позиции и, если необходимо, обновляем. @@ -706,26 +755,39 @@ async def get_data( # Формируем результат процентного мерджера. result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - items_data: List = [] - for item in self.items: - # Получаем данные из позиций процентного мерджера. + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + items_data: List = [None] * len(self.items) + results: List[Optional[FeedResult]] = [None] * len(self.items) + + indexed_items = list(enumerate(self.items)) + fetch_order = indexed_items + if dedup_active: + fetch_order = sorted( + indexed_items, + key=lambda p: (getattr(p[1].data, "dedup_priority", 0), -p[0]), + reverse=True, + ) + + for idx, item in fetch_order: item_result = await item.data.get_data( methods_dict=methods_dict, user_id=user_id, limit=limit * item.percentage // 100, next_page=next_page, redis_client=redis_client, + _sf_dedup_active=dedup_active, **params, ) - # Добавляем данные позиции в список данных позиций. - items_data.append(item_result.data) + results[idx] = item_result + + for idx, item_result in enumerate(results): + assert item_result is not None + items_data[idx] = item_result.data - # Если has_next_page = False, то проверяем has_next_page у позиции и, если необходимо, обновляем. if not result.has_next_page and item_result.has_next_page: result.has_next_page = True - - # Обновляем next_page. result.next_page.data.update(item_result.next_page.data) # Добавляем данные позиции к общему результату процентного мерджера. @@ -866,23 +928,49 @@ async def get_data( limit=limit, ) - # Получаем данные из позиций в процентном соотношений. - item_from = await self.item_from.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_from"], - next_page=next_page, - redis_client=redis_client, - **params, - ) - item_to = await self.item_to.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_to"], - next_page=next_page, - redis_client=redis_client, - **params, - ) + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + from_priority = getattr(self.item_from.data, "dedup_priority", 0) + to_priority = getattr(self.item_to.data, "dedup_priority", 0) + + if dedup_active and to_priority > from_priority: + item_to = await self.item_to.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_to"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + item_from = await self.item_from.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_from"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + else: + item_from = await self.item_from.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_from"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + item_to = await self.item_to.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_to"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) from_start_index = 0 to_start_index = 0 @@ -984,62 +1072,75 @@ async def get_data( :return: список данных методом append. """ - # Формируем результат append мерджера. + dedup_active = bool(params.pop("_sf_dedup_active", False)) + result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - result_limit = limit - for item in self.items: - # Получаем данные из позиции мерджера. - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) + if dedup_active: + indexed_items = list(enumerate(self.items)) + fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) + fetched: Dict[int, FeedResult] = {} - # Добавляем данные позиции к общему результату процентного мерджера. - result.data.extend(item_result.data) + for idx, item in fetch_order: + fetched[idx] = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) - # Обновляем result_limit - result_limit -= len(item_result.data) + for idx, _item in indexed_items: + item_result = fetched[idx] + result.data.extend(item_result.data) + result.next_page.data.update(item_result.next_page.data) + if item_result.has_next_page: + result.has_next_page = True - # Если has_next_page = False, то проверяем has_next_page у позиции и, если необходимо, обновляем. - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True + if len(result.data) > limit: + result.data = result.data[:limit] + else: + result_limit = limit + for item in self.items: + item_result = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=result_limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) - # Обновляем next_page. - result.next_page.data.update(item_result.next_page.data) + result.data.extend(item_result.data) + result_limit -= len(item_result.data) + + if not result.has_next_page and item_result.has_next_page: + result.has_next_page = True + + result.next_page.data.update(item_result.next_page.data) - # Если полученных данных хватает, то прерываем итерацию и возвращаем результат. - if result_limit <= 0: - break + if result_limit <= 0: + break # Распределяем данные равномерно по ключу. result.data = await self._uniform_distribute(result.data) return result -class MergerDeduplicationItem(BaseModel): - """Configuration item for MergerDeduplication.""" - - priority: int = 0 - data: FeedTypes - - class MergerDeduplication(BaseFeedConfigModel): - """Merger that deduplicates items and refills to the requested limit. + """Merger that deduplicates while preserving child mixing/position semantics. - Key properties: - - Always tries to return exactly `limit` unique items if they exist upstream. - - Supports cross-page deduplication using either cursor state or Redis. - - Supports explicit per-source priority; higher priority wins on same dedup key. + This merger acts as a wrapper around exactly one child feed node. + Deduplication is applied at the leaf SubFeed method level with a shared seen-set. + This lets nested mergers (positional/percentage/gradient/etc.) keep their slot rules: + duplicates are skipped by fetching additional items from the *same* leaf source. """ merger_id: str type: Literal["merger_deduplication"] - items: List[MergerDeduplicationItem] + data: FeedTypes dedup_key: Optional[str] = None missing_key_policy: Literal["error", "keep", "drop"] = "error" @@ -1049,9 +1150,18 @@ class MergerDeduplication(BaseFeedConfigModel): cursor_compress: bool = True cursor_max_keys: Optional[int] = None - overfetch_factor: int = 2 + overfetch_factor: int = 1 + max_refill_loops: int = 20 + @model_validator(mode="after") + def validate_merger_deduplication(self) -> "MergerDeduplication": + if self.overfetch_factor < 1: + raise ValueError('"overfetch_factor" must be >= 1') + if self.max_refill_loops < 1: + raise ValueError('"max_refill_loops" must be >= 1') + return self + def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: keys: set[str] = set() @@ -1111,53 +1221,68 @@ def _extract_dedup_value(self, item: Any) -> Any: ) return value - def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> List[str]: + def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> Dict[str, int]: entry = next_page.data.get(self.merger_id) if not entry or entry.after is None: - return [] + return {} after = entry.after if isinstance(after, dict) and "z" in after: payload = base64.urlsafe_b64decode(after["z"].encode()) raw = zlib.decompress(payload).decode() - return list(json.loads(raw)) + decoded = json.loads(raw) + if isinstance(decoded, dict): + return {str(k): int(v) for k, v in decoded.items()} + if isinstance(decoded, list): + # v2: list of [key, priority] entries + seen_map: Dict[str, int] = {} + for entry_item in decoded: + if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: + seen_map[str(entry_item[0])] = int(entry_item[1]) + else: + seen_map[str(entry_item)] = 0 + return seen_map + return {} if isinstance(after, dict) and "seen" in after: - return list(after["seen"]) + return {str(k): 0 for k in list(after["seen"])} if isinstance(after, list): - return list(after) - return [] + return {str(k): 0 for k in list(after)} + if isinstance(after, dict): + # v2 uncompressed map + return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} + return {} - def _encode_seen_for_cursor(self, seen_keys_in_order: List[str]) -> Any: + def _encode_seen_for_cursor(self, seen_updates_in_order: List[tuple[str, int]]) -> Any: if self.cursor_max_keys is not None: - seen_keys_in_order = seen_keys_in_order[-self.cursor_max_keys :] + seen_updates_in_order = seen_updates_in_order[-self.cursor_max_keys :] if not self.cursor_compress: - return {"v": 1, "seen": seen_keys_in_order} + return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} - raw = json.dumps(seen_keys_in_order).encode() + raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() compressed = zlib.compress(raw) return { - "v": 1, + "v": 2, "c": "zlib+base64", - "n": len(seen_keys_in_order), + "n": len(seen_updates_in_order), "z": base64.urlsafe_b64encode(compressed).decode(), } - async def _redis_sismember(self, redis_client: Union[redis.Redis, AsyncRedis], key: str, member: str) -> bool: - res = redis_client.sismember(key, member) + async def _redis_zscore(self, redis_client: Union[redis.Redis, AsyncRedis], key: str, member: str) -> Optional[float]: + res = redis_client.zscore(key, member) if inspect.iscoroutine(res): res = await res - return bool(res) + return None if res is None else float(res) - async def _redis_sadd_and_expire( + async def _redis_zadd_and_expire( self, redis_client: Union[redis.Redis, AsyncRedis], key: str, - members: List[str], + member_scores: Dict[str, int], ) -> None: - if not members: + if not member_scores: return - res = redis_client.sadd(key, *members) + res = redis_client.zadd(key, mapping={m: float(s) for m, s in member_scores.items()}) if inspect.iscoroutine(res): await res await redis_client.expire(key, self.state_ttl_seconds) @@ -1197,18 +1322,18 @@ async def get_data( if is_fresh_session: # Reset cursors for all descendants under this merger so upstream nodes also restart. - descendant_keys: set[str] = set() - for item in self.items: - descendant_keys.update(self._collect_descendant_cursor_keys(item.data)) + descendant_keys = self._collect_descendant_cursor_keys(self.data) for key in descendant_keys: working_next_page.data.pop(key, None) - sorted_items = sorted(self.items, key=lambda x: x.priority, reverse=True) - seen_keys_in_order: List[str] = [] - seen_cursor_set: set[str] = set() + # Shared dedup state (cross-page) + seen_priority_map: Dict[str, int] = {} + seen_updates_in_order: List[tuple[str, int]] = [] if self.state_backend == "cursor" and not is_fresh_session: - seen_keys_in_order = self._decode_seen_from_cursor(next_page) - seen_cursor_set = set(seen_keys_in_order) + seen_priority_map = self._decode_seen_from_cursor(next_page) + + # Always maintain a per-request seen set to prevent duplicates within a single get_data() call. + seen_request_set: set[str] = set(seen_priority_map.keys()) redis_state_key = "" if self.state_backend == "redis" and redis_client: @@ -1219,92 +1344,175 @@ async def get_data( if inspect.iscoroutine(deleted): await deleted - result_items: List[Any] = [] - accepted: Dict[str, Dict[str, Any]] = {} - redis_new_members: List[str] = [] - any_has_next_page = False + redis_new_scores: Dict[str, int] = {} - loops = 0 - while len(result_items) < limit and loops < self.max_refill_loops: - loops += 1 - before_len = len(result_items) + # Preserve inner merger ordering/mixing semantics by deduplicating at the leaf method level + # with a shared seen-set. + original_methods_dict = methods_dict - for item in sorted_items: - remaining = limit - len(result_items) - if remaining <= 0: - break + # Create a deep copy of the child tree and rewrite each SubFeed to call a unique wrapper + # so we can associate a dedup_priority with each leaf. + child = self.data + if hasattr(child, "model_copy"): + child = child.model_copy(deep=True) # type: ignore[attr-defined] + else: + child = child.copy(deep=True) - request_limit = max(1, remaining * max(1, self.overfetch_factor)) - item_result = await item.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=request_limit, - next_page=working_next_page, - redis_client=redis_client, - **params, - ) + def iter_subfeeds(feed: BaseFeedConfigModel) -> List["SubFeed"]: + found: List[SubFeed] = [] - any_has_next_page = any_has_next_page or item_result.has_next_page - working_next_page.data.update(item_result.next_page.data) + if isinstance(feed, SubFeed): + found.append(feed) + return found - for entity in item_result.data: - raw_value = self._extract_dedup_value(entity) - if raw_value is None: - if self.missing_key_policy == "drop": - continue - if self.missing_key_policy == "keep": - # Make a unique key per object instance representation. - raw_value = ("__missing__", id(entity)) + for attr_name in ("data", "positional", "default"): + inner = getattr(feed, attr_name, None) + if isinstance(inner, BaseFeedConfigModel): + found.extend(iter_subfeeds(inner)) - key = self._normalize_key(raw_value) + for attr_name in ("item_from", "item_to"): + wrapper = getattr(feed, attr_name, None) + inner = getattr(wrapper, "data", None) + if isinstance(inner, BaseFeedConfigModel): + found.extend(iter_subfeeds(inner)) - if key in accepted: - if item.priority > accepted[key]["priority"]: - result_items[accepted[key]["index"]] = entity - accepted[key]["priority"] = item.priority + items = getattr(feed, "items", None) + if isinstance(items, list): + for item in items: + if isinstance(item, BaseFeedConfigModel): + found.extend(iter_subfeeds(item)) continue - - if self.state_backend == "cursor": - if key in seen_cursor_set: + inner = getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + found.extend(iter_subfeeds(inner)) + + return found + + rewritten_methods_dict = dict(original_methods_dict) + + def wrap_leaf_method(*, subfeed: "SubFeed") -> None: + original_name = subfeed.method_name + original_method = original_methods_dict[original_name] + unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" + # Idempotency: if the same subfeed id appears multiple times, don't re-wrap. + if unique_name in rewritten_methods_dict: + subfeed.method_name = unique_name + return + subfeed.method_name = unique_name + leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) + + async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPageInside, **kw: Any): + collected: List[Any] = [] + local_seen: set[str] = set() + any_has_next_page = False + + loops = 0 + while len(collected) < limit and loops < self.max_refill_loops: + loops += 1 + before_len = len(collected) + + remaining = limit - len(collected) + # Safe oversampling: only when we can rewind integer-offset cursors. + can_overfetch = isinstance(next_page.after, (int, type(None))) + request_limit = max(1, remaining) + if can_overfetch and self.overfetch_factor > 1: + request_limit = max(1, remaining * self.overfetch_factor) + + start_after = 0 if next_page.after is None else int(next_page.after) + + method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) + if not isinstance(method_result, FeedResultClient): + raise TypeError('SubFeed function must return "FeedResultClient" instance.') + + any_has_next_page = any_has_next_page or method_result.has_next_page + + consumed_in_batch = 0 + + for entity in method_result.data: + consumed_in_batch += 1 + raw_value = self._extract_dedup_value(entity) + if raw_value is None: + if self.missing_key_policy == "drop": + continue + if self.missing_key_policy == "keep": + raw_value = ("__missing__", id(entity)) + + key = self._normalize_key(raw_value) + if key in local_seen: continue - else: - assert redis_client is not None - if await self._redis_sismember(redis_client, redis_state_key, key): + + if key in seen_request_set: continue - accepted[key] = {"priority": item.priority, "index": len(result_items)} - result_items.append(entity) + if self.state_backend == "cursor": + existing_priority = seen_priority_map.get(key) + if existing_priority is not None and leaf_priority <= existing_priority: + continue + else: + assert redis_client is not None + existing_score = await self._redis_zscore(redis_client, redis_state_key, key) + if existing_score is not None and leaf_priority <= int(existing_score): + continue - if self.state_backend == "cursor": - seen_cursor_set.add(key) - seen_keys_in_order.append(key) - else: - redis_new_members.append(key) + local_seen.add(key) + collected.append(entity) - if len(result_items) >= limit: - break + seen_request_set.add(key) - if len(result_items) >= limit: - break + if self.state_backend == "cursor": + seen_priority_map[key] = leaf_priority + seen_updates_in_order.append((key, leaf_priority)) + else: + redis_new_scores[key] = max(redis_new_scores.get(key, 0), leaf_priority) + + if len(collected) >= limit: + break + + if len(collected) == before_len: + # No progress this loop. Stop if upstream is exhausted. + if not method_result.has_next_page: + break + + # If we oversampled with a simple integer cursor, rewind to the point we actually consumed. + # This prevents skipping un-inspected items that were fetched but not needed. + if can_overfetch and request_limit > remaining: + end_after = next_page.after + if isinstance(end_after, int) and end_after == start_after + len(method_result.data): + next_page.after = start_after + consumed_in_batch + + return FeedResultClient(data=collected, next_page=next_page, has_next_page=any_has_next_page) + + setattr(_wrapped_method, "_smartfeed_original", original_method) + rewritten_methods_dict[unique_name] = _wrapped_method - if len(result_items) == before_len: - break + for sf in iter_subfeeds(child): + wrap_leaf_method(subfeed=sf) + + child_result = await child.get_data( + methods_dict=rewritten_methods_dict, + user_id=user_id, + limit=limit, + next_page=working_next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) if self.state_backend == "redis" and redis_client: - await self._redis_sadd_and_expire(redis_client, redis_state_key, redis_new_members) + await self._redis_zadd_and_expire(redis_client, redis_state_key, redis_new_scores) page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 merger_after: Any = None if self.state_backend == "cursor": - merger_after = self._encode_seen_for_cursor(seen_keys_in_order) + merger_after = self._encode_seen_for_cursor(seen_updates_in_order) - if hasattr(working_next_page, "model_copy"): - result_next_page = working_next_page.model_copy(deep=True) # type: ignore[attr-defined] + if hasattr(child_result.next_page, "model_copy"): + result_next_page = child_result.next_page.model_copy(deep=True) # type: ignore[attr-defined] else: - result_next_page = working_next_page.copy(deep=True) + result_next_page = child_result.next_page.copy(deep=True) result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) - return FeedResult(data=result_items, next_page=result_next_page, has_next_page=any_has_next_page) + return FeedResult(data=child_result.data, next_page=result_next_page, has_next_page=child_result.has_next_page) class SubFeed(BaseFeedConfigModel): @@ -1354,7 +1562,9 @@ async def get_data( ) # Формируем params для функции субфида. - method_args = inspect.getfullargspec(methods_dict[self.method_name]).args + method = methods_dict[self.method_name] + method_spec = getattr(method, "_smartfeed_original", method) + method_args = inspect.getfullargspec(method_spec).args method_params: Dict[str, Any] = {} for arg in method_args: if arg in params: @@ -1426,5 +1636,4 @@ def _rebuild_model(model: Any) -> None: _rebuild_model(MergerAppendDistribute) _rebuild_model(MergerPercentageGradient) _rebuild_model(MergerViewSession) -_rebuild_model(MergerDeduplicationItem) _rebuild_model(MergerDeduplication) \ No newline at end of file diff --git a/tests/fixtures/configs.py b/tests/fixtures/configs.py index a982e3d..6aff5cb 100644 --- a/tests/fixtures/configs.py +++ b/tests/fixtures/configs.py @@ -96,23 +96,28 @@ "dedup_key": "id", "state_backend": "cursor", "cursor_compress": True, - "items": [ - { - "priority": 100, - "data": { - "subfeed_id": "subfeed_dedup_priority_high", - "type": "subfeed", - "method_name": "posted", + "data": { + "merger_id": "merger_percentage_inside_dedup_parsing_example", + "type": "merger_percentage", + "shuffle": False, + "items": [ + { + "percentage": 50, + "data": { + "subfeed_id": "subfeed_dedup_a", + "type": "subfeed", + "method_name": "posted", + }, }, - }, - { - "priority": 0, - "data": { - "subfeed_id": "subfeed_dedup_priority_low", - "type": "subfeed", - "method_name": "posted", + { + "percentage": 50, + "data": { + "subfeed_id": "subfeed_dedup_b", + "type": "subfeed", + "method_name": "posted", + }, }, - }, - ], + ], + }, }, } diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index cede8f3..5b046a8 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -3,10 +3,10 @@ import pytest from smartfeed.schemas import ( - MergerDeduplication, FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, + MergerDeduplication, ) from tests.fixtures.redis import redis_client # noqa: F401 @@ -14,7 +14,7 @@ def make_offset_paged_method(items, *, max_per_call=None): - async def _method(user_id, limit, next_page): # pylint: disable=unused-argument + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument offset = int(next_page.after or 0) effective_limit = limit if isinstance(max_per_call, int) and max_per_call > 0: @@ -28,63 +28,7 @@ async def _method(user_id, limit, next_page): # pylint: disable=unused-argument return _method -async def _run_two_pages( - *, - config, - methods_dict, - user_id, - limit, - redis_client_instance=None, - **params, -): - merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=FeedResultNextPage(data={}), - redis_client=redis_client_instance, - **params, - ) - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=res_1.next_page, - redis_client=redis_client_instance, - **params, - ) - return res_1, res_2 - - -def _assert_dedup_backend_state(*, res, merger_id: str, state_backend: str) -> None: - assert merger_id in res.next_page.data - if state_backend == "cursor": - assert isinstance(res.next_page.data[merger_id].after, dict) - else: - assert res.next_page.data[merger_id].after is None - - -def _ids(data): - return [x["id"] for x in data] - - -def _assert_two_pages_no_overlap(res_1, res_2): - ids_1 = set(_ids(res_1.data)) - ids_2 = set(_ids(res_2.data)) - assert len(ids_1) == len(res_1.data) - assert len(ids_2) == len(res_2.data) - assert not (ids_1 & ids_2) - - def _assert_cursor_monotonic_if_present(res_1, res_2, keys): - """Assert that cursor values monotonically advance for keys that are present. - - MergerDeduplication may stop early once it has enough unique items, so a - descendant might not be called on a given page. This helper only asserts - monotonicity when the cursor key exists in `res_1`. - """ - for key in keys: if key not in res_1.next_page.data: continue @@ -101,61 +45,67 @@ def _assert_cursor_monotonic_if_present(res_1, res_2, keys): assert after_2 >= after_1 continue - # Merger cursors can be structured (dict), just require presence. if isinstance(after_1, dict) and isinstance(after_2, dict): continue - # If values are comparable, enforce monotonicity; otherwise don't fail. try: assert after_2 >= after_1 except TypeError: pass +def _sources(data): + return [x.get("src") for x in data] + + +def _ids(data): + return [x.get("id") for x in data] + + +def _assert_no_dupes_in_page(data): + ids = _ids(data) + assert len(ids) == len(set(ids)) + + +def _assert_pages_no_overlap(res_1, res_2): + assert not (set(_ids(res_1.data)) & set(_ids(res_2.data))) + + @pytest.mark.asyncio -async def test_deduplication_merger_cursor_priority_and_cross_page() -> None: - low_items = [ - {"id": 1, "src": "low"}, - {"id": 2, "src": "low"}, - {"id": 3, "src": "low"}, - {"id": 4, "src": "low"}, - {"id": 5, "src": "low"}, - # repeats later (cross-page duplicates) - {"id": 3, "src": "low"}, - {"id": 4, "src": "low"}, - {"id": 6, "src": "low"}, - {"id": 7, "src": "low"}, - {"id": 8, "src": "low"}, - {"id": 9, "src": "low"}, - {"id": 10, "src": "low"}, - ] - high_items = [ - {"id": 3, "src": "high"}, - {"id": 4, "src": "high"}, - ] +async def test_dedup_positional_slot_ownership_cursor_backend() -> None: + """Positional slots must remain owned by the positional branch. + + Deduplication must not drop items *after* the positional merge (which would shift indices). + Instead, duplicates must be skipped inside the leaf source that owns the slot. + """ + + # Default branch has early ids 1..3, which will be seen first. + default_items = [{"id": i, "src": "default"} for i in range(1, 300)] + + # Positional branch starts with duplicates 1..3; it must skip them and fetch 4.. instead. + positional_items = [{"id": i, "src": "pos"} for i in range(1, 300)] methods_dict = { - "low": make_offset_paged_method(low_items), - "high": make_offset_paged_method(high_items), + "default": make_offset_paged_method(default_items), + "pos": make_offset_paged_method(positional_items), } config = { - "merger_id": "dedup_example", + "merger_id": "dedup_wrapper", "type": "merger_deduplication", "dedup_key": "id", "state_backend": "cursor", "cursor_compress": True, - "overfetch_factor": 3, - "items": [ - { - "priority": 100, - "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high"}, - }, - { - "priority": 0, - "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low"}, - }, - ], + "max_refill_loops": 20, + "data": { + "merger_id": "positional_mix", + "type": "merger_positional", + # Ensure positional inserts exist on both pages for limit=6: + # page1 uses (1,3,5), page2 uses (7,9,11) which map to the same in-page slots. + "positions": [1, 3, 5, 7, 9, 11], + "positional": {"subfeed_id": "sf_pos", "type": "subfeed", "method_name": "pos"}, + "default": {"subfeed_id": "sf_default", "type": "subfeed", "method_name": "default"}, + }, } merger = parse_model(MergerDeduplication, config) @@ -163,99 +113,68 @@ async def test_deduplication_merger_cursor_priority_and_cross_page() -> None: res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=5, + limit=6, next_page=FeedResultNextPage(data={}), ) - assert len(res_1.data) == 5 - ids_1 = [x["id"] for x in res_1.data] - assert len(ids_1) == len(set(ids_1)) - # Priority: id 3 and 4 must come from high - for x in res_1.data: - if x["id"] in {3, 4}: - assert x["src"] == "high" + assert len(res_1.data) == 6 + _assert_no_dupes_in_page(res_1.data) - # Next page should not repeat 3/4 even though low repeats them later. + # Slot ownership: configured positions [1,3,5] are the positional branch. + assert _sources(res_1.data)[0] == "pos" + assert _sources(res_1.data)[2] == "pos" + assert _sources(res_1.data)[4] == "pos" + + # Next page: still no overlap across pages, and positional slots remain owned. res_2 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=5, + limit=6, next_page=res_1.next_page, ) - ids_2 = [x["id"] for x in res_2.data] - assert not (set(ids_1) & set(ids_2)) - - # Ensure merger stores cursor state (compressed) in its own after. - assert "dedup_example" in res_2.next_page.data - assert isinstance(res_2.next_page.data["dedup_example"].after, dict) - assert "z" in res_2.next_page.data["dedup_example"].after - - -@pytest.mark.asyncio -async def test_deduplication_merger_refill_to_limit() -> None: - dup_items = [ - {"id": 1}, - {"id": 1}, - {"id": 1}, - {"id": 1}, - {"id": 1}, - {"id": 2}, - {"id": 3}, - {"id": 4}, - {"id": 5}, - {"id": 6}, - ] + assert len(res_2.data) == 6 + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) - methods_dict = { - "dups": make_offset_paged_method(dup_items), - } + assert _sources(res_2.data)[0] == "pos" + assert _sources(res_2.data)[2] == "pos" + assert _sources(res_2.data)[4] == "pos" - config = { - "merger_id": "dedup_refill", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "overfetch_factor": 4, - "max_refill_loops": 10, - "items": [ - { - "priority": 0, - "data": {"subfeed_id": "sf_dups", "type": "subfeed", "method_name": "dups"}, - } - ], - } + _assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"]) - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage(data={}), - ) +@pytest.mark.asyncio +async def test_dedup_percentage_slot_ownership_cursor_backend() -> None: + """Percentage mixing order must be preserved even with duplicates across sources.""" - assert [x["id"] for x in res.data] == [1, 2, 3, 4, 5] + # A is called first by the percentage merger; its ids will be seen before B. + a_items = [{"id": i, "src": "A"} for i in range(1, 300)] + # B starts with duplicates 1..3; it must skip them and fetch unique tail items. + # Same IDs as A to force cross-source duplicates. + b_items = [{"id": i, "src": "B"} for i in range(1, 300)] -@pytest.mark.asyncio -async def test_deduplication_merger_page_zero_resets_cursor_state() -> None: - items = [{"id": i} for i in range(1, 50)] - methods_dict = {"stream": make_offset_paged_method(items)} + methods_dict = { + "a": make_offset_paged_method(a_items), + "b": make_offset_paged_method(b_items), + } config = { - "merger_id": "dedup_reset", + "merger_id": "dedup_wrapper_pct", "type": "merger_deduplication", "dedup_key": "id", "state_backend": "cursor", "cursor_compress": True, - "overfetch_factor": 2, - "items": [ - { - "priority": 0, - "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "stream"}, - } - ], + "data": { + "merger_id": "pct_mix", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, + ], + }, } merger = parse_model(MergerDeduplication, config) @@ -263,61 +182,76 @@ async def test_deduplication_merger_page_zero_resets_cursor_state() -> None: res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=5, + limit=10, next_page=FeedResultNextPage(data={}), ) - assert [x["id"] for x in res_1.data] == [1, 2, 3, 4, 5] - # Simulate a full reload: page 0 requested again. Even if the client mistakenly - # keeps the previous cursor payloads (including subfeed cursors), we start a new session. + assert len(res_1.data) == 10 + _assert_no_dupes_in_page(res_1.data) + + # Slot ownership: percentage merge alternates when list sizes are equal. + sources_1 = _sources(res_1.data) + assert sources_1[0] == "A" + assert sources_1[1] == "B" + assert sources_1[2] == "A" + assert sources_1[3] == "B" + res_2 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=5, - next_page=FeedResultNextPage( - data={ - "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after), - "sf_stream": res_1.next_page.data["sf_stream"], - } - ), + limit=10, + next_page=res_1.next_page, ) - assert [x["id"] for x in res_2.data] == [1, 2, 3, 4, 5] + assert len(res_2.data) == 10 + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + sources_2 = _sources(res_2.data) + assert sources_2[0] == "A" + assert sources_2[1] == "B" + + _assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_a", "sf_b", "pct_mix", "dedup_wrapper_pct"]) -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio -async def test_deduplication_merger_redis_backend(redis_client) -> None: - # This dataset repeats ids across pages (sliding window style) - items = [ - {"id": 1}, - {"id": 2}, - {"id": 3}, - {"id": 2}, - {"id": 3}, - {"id": 4}, - {"id": 5}, - {"id": 6}, - {"id": 4}, - {"id": 7}, - {"id": 8}, - ] +async def test_dedup_deep_tree_cursor_backend() -> None: + """Dedup must work through deep merger trees (wrapping leaf methods).""" - methods_dict = {"stream": make_offset_paged_method(items)} + # Leaf sources: intentionally overlapping ids across different leaves. + p_items = [{"id": i, "src": "P"} for i in range(1, 30)] + d1_items = [{"id": i, "src": "D1"} for i in range(1, 30)] # overlaps P + d2_items = [{"id": 100 + i, "src": "D2"} for i in range(1, 30)] + methods_dict = { + "p": make_offset_paged_method(p_items), + "d1": make_offset_paged_method(d1_items), + "d2": make_offset_paged_method(d2_items), + } + + # Deep tree: Dedup -> Positional(default=Percentage(D1,D2), positional=SubFeed(P)) config = { - "merger_id": "dedup_redis", + "merger_id": "dedup_deep", "type": "merger_deduplication", "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 4, - "items": [ - { - "priority": 0, - "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "stream"}, - } - ], + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "pos_deep", + "type": "merger_positional", + # Ensure positional positions exist on both page 1 (1,4) and page 2 (9,12) for limit=8. + "positions": [1, 4, 9, 12], + "positional": {"subfeed_id": "sf_p", "type": "subfeed", "method_name": "p"}, + "default": { + "merger_id": "pct_deep", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_d1", "type": "subfeed", "method_name": "d1"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_d2", "type": "subfeed", "method_name": "d2"}}, + ], + }, + }, } merger = parse_model(MergerDeduplication, config) @@ -325,693 +259,431 @@ async def test_deduplication_merger_redis_backend(redis_client) -> None: res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=4, + limit=8, next_page=FeedResultNextPage(data={}), - redis_client=redis_client, - custom_deduplication_key="t1", ) + assert len(res_1.data) == 8 + _assert_no_dupes_in_page(res_1.data) + + # Positional ownership must hold even with deep defaults. + assert _sources(res_1.data)[0] == "P" # position 1 + assert _sources(res_1.data)[3] == "P" # position 4 + res_2 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=4, + limit=8, next_page=res_1.next_page, - redis_client=redis_client, - custom_deduplication_key="t1", ) - ids_1 = [x["id"] for x in res_1.data] - ids_2 = [x["id"] for x in res_2.data] + assert len(res_2.data) == 8 + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) - assert len(ids_1) == len(set(ids_1)) - assert len(ids_2) == len(set(ids_2)) - assert not (set(ids_1) & set(ids_2)) + assert _sources(res_2.data)[0] == "P" + assert _sources(res_2.data)[3] == "P" - # Redis backend should not store seen ids in cursor after. - assert "dedup_redis" in res_2.next_page.data - assert res_2.next_page.data["dedup_redis"].after is None - # Ensure fixture works for both sync/async redis. - key = "dedup:dedup_redis:u:t1" - members = redis_client.smembers(key) - if inspect.iscoroutine(members): - members = await members - assert len(members) >= len(set(ids_1 + ids_2)) +@pytest.mark.asyncio +async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cursors() -> None: + """When overfetch_factor>1, leaf cursors must be rewound to inspected count. + This is a regression test for the "safe overfetch" logic: we may request more + than we need from a leaf source, but we must not advance that leaf cursor past + un-inspected items. In a deep tree, this must hold for all descendant SubFeeds. + """ -@pytest.mark.asyncio -async def test_deduplication_merger_priority_replacement_across_loops_cursor_backend() -> None: - # This test forces the higher-priority source to surface a duplicate only on a later call, - # so we exercise the in-page replacement logic. - # Important: dedup calls sources in descending priority. To ensure we exercise - # replacement, we need a lower-priority source to introduce id=5 *before* - # the high-priority source sees id=5 on a later refill loop. - low_items = [ - {"id": 5, "src": "low"}, - {"id": 6, "src": "low"}, - {"id": 7, "src": "low"}, - {"id": 99, "src": "low"}, - ] - mid_items = [ - {"id": 5, "src": "mid"}, - {"id": 98, "src": "mid"}, - {"id": 8, "src": "mid"}, - {"id": 9, "src": "mid"}, - ] - high_items = [ - {"id": 1, "src": "high"}, - {"id": 5, "src": "high"}, - {"id": 2, "src": "high"}, - {"id": 3, "src": "high"}, - ] + p_items = [{"id": 1000 + i, "src": "P"} for i in range(1, 200)] + d1_items = [{"id": i, "src": "D1"} for i in range(1, 200)] + d2_items = [{"id": 500 + i, "src": "D2"} for i in range(1, 200)] methods_dict = { - "low": make_offset_paged_method(low_items, max_per_call=1), - "mid": make_offset_paged_method(mid_items, max_per_call=1), - "high": make_offset_paged_method(high_items, max_per_call=1), + "p": make_offset_paged_method(p_items), + "d1": make_offset_paged_method(d1_items), + "d2": make_offset_paged_method(d2_items), } config = { - "merger_id": "dedup_priority_cursor", + "merger_id": "dedup_overfetch", "type": "merger_deduplication", "dedup_key": "id", "state_backend": "cursor", "cursor_compress": True, - "overfetch_factor": 1, - "max_refill_loops": 10, - "items": [ - {"priority": 100, "data": {"subfeed_id": "sf_high_p", "type": "subfeed", "method_name": "high"}}, - {"priority": 50, "data": {"subfeed_id": "sf_mid_p", "type": "subfeed", "method_name": "mid"}}, - {"priority": 0, "data": {"subfeed_id": "sf_low_p", "type": "subfeed", "method_name": "low"}}, - ], + "overfetch_factor": 3, + "data": { + "merger_id": "pos_overfetch", + "type": "merger_positional", + "positions": [1, 4, 9, 12], + "positional": {"subfeed_id": "sf_p", "type": "subfeed", "method_name": "p"}, + "default": { + "merger_id": "pct_overfetch", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_d1", "type": "subfeed", "method_name": "d1"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_d2", "type": "subfeed", "method_name": "d2"}}, + ], + }, + }, } - res_1 = await parse_model(MergerDeduplication, config).get_data( + merger = parse_model(MergerDeduplication, config) + + # Page 1 + res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=4, + limit=8, next_page=FeedResultNextPage(data={}), ) - # Ensure id=5 is present and comes from highest priority, even though low/mid can surface it earlier. - winners = {x["id"]: x["src"] for x in res_1.data} - assert winners[5] == "high" - _assert_dedup_backend_state(res=res_1, merger_id="dedup_priority_cursor", state_backend="cursor") + assert len(res_1.data) == 8 + _assert_no_dupes_in_page(res_1.data) + # Deep descendant cursors: positional leaf requests 2 items; percentage leaves request 4 each. + # With overfetch_factor=3, internal calls may request 6/12, but cursor must not advance that far. + assert res_1.next_page.data["sf_p"].after == 2 + assert res_1.next_page.data["sf_d1"].after == 4 + assert res_1.next_page.data["sf_d2"].after == 4 -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_deduplication_merger_priority_replacement_across_loops_redis_backend(redis_client) -> None: - low_items = [ - {"id": 5, "src": "low"}, - {"id": 6, "src": "low"}, - {"id": 7, "src": "low"}, - {"id": 99, "src": "low"}, - ] - mid_items = [ - {"id": 5, "src": "mid"}, - {"id": 98, "src": "mid"}, - {"id": 8, "src": "mid"}, - {"id": 9, "src": "mid"}, - ] - high_items = [ - {"id": 1, "src": "high"}, - {"id": 5, "src": "high"}, - {"id": 2, "src": "high"}, - {"id": 3, "src": "high"}, - ] - - methods_dict = { - "low": make_offset_paged_method(low_items, max_per_call=1), - "mid": make_offset_paged_method(mid_items, max_per_call=1), - "high": make_offset_paged_method(high_items, max_per_call=1), - } - - config = { - "merger_id": "dedup_priority_redis", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 1, - "max_refill_loops": 10, - "items": [ - {"priority": 100, "data": {"subfeed_id": "sf_high_pr", "type": "subfeed", "method_name": "high"}}, - {"priority": 50, "data": {"subfeed_id": "sf_mid_pr", "type": "subfeed", "method_name": "mid"}}, - {"priority": 0, "data": {"subfeed_id": "sf_low_pr", "type": "subfeed", "method_name": "low"}}, - ], - } - - res_1 = await parse_model(MergerDeduplication, config).get_data( + # Page 2 (monotonic advancement, still no over-advancement) + res_2 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=4, - next_page=FeedResultNextPage(data={}), - redis_client=redis_client, - custom_deduplication_key="priority", + limit=8, + next_page=res_1.next_page, ) - winners = {x["id"]: x["src"] for x in res_1.data} - assert winners[5] == "high" - _assert_dedup_backend_state(res=res_1, merger_id="dedup_priority_redis", state_backend="redis") + assert len(res_2.data) == 8 + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + assert res_2.next_page.data["sf_p"].after == 4 + assert res_2.next_page.data["sf_d1"].after == 8 + assert res_2.next_page.data["sf_d2"].after == 8 -@pytest.mark.asyncio -async def test_deduplication_merger_with_append_and_three_sources_cursor_backend() -> None: - # Inner MergerAppend (two subfeeds) + two extra subfeeds as separate dedup items. - a_items = [{"id": i, "src": "a"} for i in range(1, 30)] - b_items = [{"id": i, "src": "b"} for i in range(10, 40)] - c_items = [{"id": i, "src": "c"} for i in range(20, 60)] - d_items = [{"id": i, "src": "d"} for i in range(25, 70)] - - # Cap each subfeed to 1 item per call so dedup must invoke all children - # (and therefore exercise nested cursor propagation). - methods_dict = { - "a": make_offset_paged_method(a_items, max_per_call=1), - "b": make_offset_paged_method(b_items, max_per_call=1), - "c": make_offset_paged_method(c_items, max_per_call=1), - "d": make_offset_paged_method(d_items, max_per_call=1), - } + _assert_cursor_monotonic_if_present( + res_1, + res_2, + keys=["sf_p", "sf_d1", "sf_d2", "pos_overfetch", "dedup_overfetch"], + ) - append_config = { - "merger_id": "inner_append_unused", - "type": "merger_append", - "items": [ - {"subfeed_id": "sf_a_append", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b_append", "type": "subfeed", "method_name": "b"}, - ], - } + +@pytest.mark.asyncio +async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: + items = [{"id": i, "src": "S"} for i in range(1, 50)] + methods_dict = {"s": make_offset_paged_method(items)} config = { - "merger_id": "dedup_with_append_cursor", + "merger_id": "dedup_reset", "type": "merger_deduplication", "dedup_key": "id", "state_backend": "cursor", "cursor_compress": True, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": append_config}, - {"priority": 5, "data": {"subfeed_id": "sf_c", "type": "subfeed", "method_name": "c"}}, - {"priority": 0, "data": {"subfeed_id": "sf_d", "type": "subfeed", "method_name": "d"}}, - ], - } - - res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=15) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_with_append_cursor", state_backend="cursor") - - # Cursor correctness: descendant subfeed cursors exist and advance. - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_a_append", "sf_b_append", "sf_c", "sf_d"]) - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_deduplication_merger_with_append_and_three_sources_redis_backend(redis_client) -> None: - a_items = [{"id": i, "src": "a"} for i in range(1, 30)] - b_items = [{"id": i, "src": "b"} for i in range(10, 40)] - c_items = [{"id": i, "src": "c"} for i in range(20, 60)] - d_items = [{"id": i, "src": "d"} for i in range(25, 70)] - - methods_dict = { - "a": make_offset_paged_method(a_items, max_per_call=1), - "b": make_offset_paged_method(b_items, max_per_call=1), - "c": make_offset_paged_method(c_items, max_per_call=1), - "d": make_offset_paged_method(d_items, max_per_call=1), - } - - append_config = { - "merger_id": "inner_append_unused_r", - "type": "merger_append", - "items": [ - {"subfeed_id": "sf_a_append_r", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b_append_r", "type": "subfeed", "method_name": "b"}, - ], + "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "s"}, } - config = { - "merger_id": "dedup_with_append_redis", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": append_config}, - {"priority": 5, "data": {"subfeed_id": "sf_c_r", "type": "subfeed", "method_name": "c"}}, - {"priority": 0, "data": {"subfeed_id": "sf_d_r", "type": "subfeed", "method_name": "d"}}, - ], - } + merger = parse_model(MergerDeduplication, config) - res_1, res_2 = await _run_two_pages( - config=config, + res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=15, - redis_client_instance=redis_client, - custom_deduplication_key="append", + limit=5, + next_page=FeedResultNextPage(data={}), ) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_with_append_redis", state_backend="redis") - - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_a_append_r", "sf_b_append_r", "sf_c_r", "sf_d_r"]) + assert _ids(res_1.data) == [1, 2, 3, 4, 5] + # Simulate client "full reload": page=0 for the dedup merger. + # Also include the stale descendant cursor; dedup should clear it. + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage( + data={ + "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after), + "sf_stream": res_1.next_page.data["sf_stream"], + } + ), + ) -@pytest.mark.asyncio -async def test_deduplication_merger_with_percentage_cursor_backend() -> None: - a_items = [{"id": i, "src": "pa"} for i in range(1, 60)] - b_items = [{"id": i, "src": "pb"} for i in range(30, 90)] - c_items = [{"id": i, "src": "pc"} for i in range(40, 120)] - - methods_dict = { - "pa": make_offset_paged_method(a_items, max_per_call=1), - "pb": make_offset_paged_method(b_items, max_per_call=1), - "pc": make_offset_paged_method(c_items, max_per_call=1), - } - - percentage_config = { - "merger_id": "inner_percentage_unused", - "type": "merger_percentage", - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_pa", "type": "subfeed", "method_name": "pa"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_pb", "type": "subfeed", "method_name": "pb"}}, - ], - } - - config = { - "merger_id": "dedup_percentage_cursor", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "overfetch_factor": 2, - "items": [ - {"priority": 0, "data": percentage_config}, - {"priority": 10, "data": {"subfeed_id": "sf_pc", "type": "subfeed", "method_name": "pc"}}, - {"priority": 5, "data": {"subfeed_id": "sf_pd", "type": "subfeed", "method_name": "pa"}}, - ], - } - - res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=20) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_percentage_cursor", state_backend="cursor") - - for key in ("sf_pa", "sf_pb", "sf_pc"): - assert key in res_1.next_page.data - assert isinstance(res_1.next_page.data[key].after, int) - - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_pa", "sf_pb", "sf_pc", "sf_pd"]) + # Must restart from the beginning. + assert _ids(res_2.data) == [1, 2, 3, 4, 5] @pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio -async def test_deduplication_merger_with_percentage_redis_backend(redis_client) -> None: - a_items = [{"id": i, "src": "pa"} for i in range(1, 60)] - b_items = [{"id": i, "src": "pb"} for i in range(30, 90)] - c_items = [{"id": i, "src": "pc"} for i in range(40, 120)] +async def test_dedup_redis_backend_cross_page(redis_client) -> None: + items_a = [{"id": i, "src": "A"} for i in range(1, 300)] + # Same IDs as A to force cross-source duplicates. + items_b = [{"id": i, "src": "B"} for i in range(1, 300)] methods_dict = { - "pa": make_offset_paged_method(a_items, max_per_call=1), - "pb": make_offset_paged_method(b_items, max_per_call=1), - "pc": make_offset_paged_method(c_items, max_per_call=1), - } - - percentage_config = { - "merger_id": "inner_percentage_unused_r", - "type": "merger_percentage", - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_pa_r", "type": "subfeed", "method_name": "pa"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_pb_r", "type": "subfeed", "method_name": "pb"}}, - ], + "a": make_offset_paged_method(items_a), + "b": make_offset_paged_method(items_b), } config = { - "merger_id": "dedup_percentage_redis", + "merger_id": "dedup_redis", "type": "merger_deduplication", "dedup_key": "id", "state_backend": "redis", "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 0, "data": percentage_config}, - {"priority": 10, "data": {"subfeed_id": "sf_pc_r", "type": "subfeed", "method_name": "pc"}}, - {"priority": 5, "data": {"subfeed_id": "sf_pd_r", "type": "subfeed", "method_name": "pa"}}, - ], + "data": { + "merger_id": "pct_mix", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, + ], + }, } - res_1, res_2 = await _run_two_pages( - config=config, + merger = parse_model(MergerDeduplication, config) + + res_1 = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=20, - redis_client_instance=redis_client, - custom_deduplication_key="percentage", + limit=10, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client, + custom_deduplication_key="t1", ) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_percentage_redis", state_backend="redis") - - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_pa_r", "sf_pb_r", "sf_pc_r", "sf_pd_r"]) - - -@pytest.mark.asyncio -async def test_deduplication_merger_with_positional_cursor_backend() -> None: - # MergerPositional carries its own merger cursor; verify it survives nesting in dedup. - pos_items = [{"id": i, "src": "pos"} for i in range(1, 100)] - def_items = [{"id": i, "src": "def"} for i in range(50, 140)] - extra_items = [{"id": i, "src": "extra"} for i in range(80, 180)] - - methods_dict = { - "pos": make_offset_paged_method(pos_items, max_per_call=1), - "def": make_offset_paged_method(def_items, max_per_call=1), - "extra": make_offset_paged_method(extra_items, max_per_call=1), - } - positional_config = { - "merger_id": "inner_positional", - "type": "merger_positional", - "positions": [0, 2, 4, 6, 8], - "positional": {"subfeed_id": "sf_positional", "type": "subfeed", "method_name": "pos"}, - "default": {"subfeed_id": "sf_default", "type": "subfeed", "method_name": "def"}, - } + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=res_1.next_page, + redis_client=redis_client, + custom_deduplication_key="t1", + ) - config = { - "merger_id": "dedup_positional_cursor", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "overfetch_factor": 2, - "items": [ - # Positional must run; it owns its own merger cursor entry. - {"priority": 10, "data": positional_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra2", "type": "subfeed", "method_name": "extra"}}, - ], - } + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) - res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=20) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_positional_cursor", state_backend="cursor") - assert "inner_positional" in res_1.next_page.data - assert "inner_positional" in res_2.next_page.data + # Redis backend should not store seen ids in cursor after. + assert "dedup_redis" in res_2.next_page.data + assert res_2.next_page.data["dedup_redis"].after is None - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_positional", "sf_default", "sf_extra", "sf_extra2"]) + # Ensure state is persisted in Redis. + key = "dedup:dedup_redis:u:t1" + members = redis_client.zrange(key, 0, -1) + if inspect.iscoroutine(members): + members = await members + assert len(members) >= len(set(_ids(res_1.data) + _ids(res_2.data))) -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio -async def test_deduplication_merger_with_positional_redis_backend(redis_client) -> None: - pos_items = [{"id": i, "src": "pos"} for i in range(1, 100)] - def_items = [{"id": i, "src": "def"} for i in range(50, 140)] - extra_items = [{"id": i, "src": "extra"} for i in range(80, 180)] +async def test_dedup_append_distribute_cursor_backend_no_dupes() -> None: + items_a = [{"id": i, "user_id": f"u{i%3}", "src": "A"} for i in range(1, 200)] + items_b = [{"id": i, "user_id": f"u{i%3}", "src": "B"} for i in range(1, 200)] methods_dict = { - "pos": make_offset_paged_method(pos_items, max_per_call=1), - "def": make_offset_paged_method(def_items, max_per_call=1), - "extra": make_offset_paged_method(extra_items, max_per_call=1), - } - - positional_config = { - "merger_id": "inner_positional_r", - "type": "merger_positional", - "positions": [0, 2, 4, 6, 8], - "positional": {"subfeed_id": "sf_positional_r", "type": "subfeed", "method_name": "pos"}, - "default": {"subfeed_id": "sf_default_r", "type": "subfeed", "method_name": "def"}, + "a": make_offset_paged_method(items_a), + "b": make_offset_paged_method(items_b), } config = { - "merger_id": "dedup_positional_redis", + "merger_id": "dedup_dist", "type": "merger_deduplication", "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": positional_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_r", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra2_r", "type": "subfeed", "method_name": "extra"}}, - ], + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "dist", + "type": "merger_distribute", + "distribution_key": "user_id", + "items": [ + {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + ], + }, } - res_1, res_2 = await _run_two_pages( - config=config, + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=20, - redis_client_instance=redis_client, - custom_deduplication_key="positional", + limit=30, + next_page=FeedResultNextPage(data={}), ) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_positional_redis", state_backend="redis") - assert "inner_positional_r" in res_1.next_page.data - _assert_cursor_monotonic_if_present( - res_1, - res_2, - ["sf_positional_r", "sf_default_r", "sf_extra_r", "sf_extra2_r"], - ) + assert len(res.data) == 30 + _assert_no_dupes_in_page(res.data) @pytest.mark.asyncio -async def test_deduplication_merger_with_percentage_gradient_cursor_backend() -> None: - from_items = [{"id": i, "src": "from"} for i in range(1, 140)] - to_items = [{"id": i, "src": "to"} for i in range(60, 200)] - extra_items = [{"id": i, "src": "extra"} for i in range(120, 300)] - - methods_dict = { - "from": make_offset_paged_method(from_items, max_per_call=1), - "to": make_offset_paged_method(to_items, max_per_call=1), - "extra": make_offset_paged_method(extra_items, max_per_call=1), - } - - gradient_config = { - "merger_id": "inner_gradient", - "type": "merger_percentage_gradient", - "item_from": {"percentage": 80, "data": {"subfeed_id": "sf_from", "type": "subfeed", "method_name": "from"}}, - "item_to": {"percentage": 20, "data": {"subfeed_id": "sf_to", "type": "subfeed", "method_name": "to"}}, - "step": 10, - "size_to_step": 10, - } - - config = { - "merger_id": "dedup_gradient_cursor", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": gradient_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_g", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra_g2", "type": "subfeed", "method_name": "extra"}}, - ], - } - - res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=25) - _assert_two_pages_no_overlap(res_1, res_2) - assert "inner_gradient" in res_1.next_page.data - assert "inner_gradient" in res_2.next_page.data - _assert_dedup_backend_state(res=res_2, merger_id="dedup_gradient_cursor", state_backend="cursor") - - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_from", "sf_to", "sf_extra_g", "sf_extra_g2"]) +async def test_dedup_in_page_deletion_priority_keeps_high_priority_even_if_config_order_is_low_first() -> None: + """High dedup_priority source must not be deleted even if called later in config order. + We use a percentage merger where both branches have overlapping ids. + The "high" branch is second in config, but has higher dedup_priority. + """ -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_deduplication_merger_with_percentage_gradient_redis_backend(redis_client) -> None: - from_items = [{"id": i, "src": "from"} for i in range(1, 140)] - to_items = [{"id": i, "src": "to"} for i in range(60, 200)] - extra_items = [{"id": i, "src": "extra"} for i in range(120, 300)] + low_items = [{"id": i, "src": "low"} for i in range(1, 200)] + high_items = [{"id": i, "src": "high"} for i in range(1, 200)] methods_dict = { - "from": make_offset_paged_method(from_items, max_per_call=1), - "to": make_offset_paged_method(to_items, max_per_call=1), - "extra": make_offset_paged_method(extra_items, max_per_call=1), - } - - gradient_config = { - "merger_id": "inner_gradient_r", - "type": "merger_percentage_gradient", - "item_from": {"percentage": 80, "data": {"subfeed_id": "sf_from_r", "type": "subfeed", "method_name": "from"}}, - "item_to": {"percentage": 20, "data": {"subfeed_id": "sf_to_r", "type": "subfeed", "method_name": "to"}}, - "step": 10, - "size_to_step": 10, + "low": make_offset_paged_method(low_items), + "high": make_offset_paged_method(high_items), } config = { - "merger_id": "dedup_gradient_redis", + "merger_id": "dedup_priority", "type": "merger_deduplication", "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": gradient_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_gr", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra_gr2", "type": "subfeed", "method_name": "extra"}}, - ], + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "pct", + "type": "merger_percentage", + "shuffle": False, + "items": [ + { + "percentage": 50, + "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}, + }, + { + "percentage": 50, + "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high", "dedup_priority": 100}, + }, + ], + }, } - res_1, res_2 = await _run_two_pages( - config=config, + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=25, - redis_client_instance=redis_client, - custom_deduplication_key="gradient", + limit=10, + next_page=FeedResultNextPage(data={}), ) - _assert_two_pages_no_overlap(res_1, res_2) - assert "inner_gradient_r" in res_1.next_page.data - _assert_dedup_backend_state(res=res_2, merger_id="dedup_gradient_redis", state_backend="redis") - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_from_r", "sf_to_r", "sf_extra_gr", "sf_extra_gr2"]) + _assert_no_dupes_in_page(res.data) + # Priority is about which source "wins" for a given dedup_key, not about output order. + # With 50/50 limits, the high-priority branch should supply ids 1..5, while the low-priority + # branch will be advanced to avoid duplicates. + winning = {item["id"]: item["src"] for item in res.data} + assert all(winning[i] == "high" for i in range(1, 6)) -@pytest.mark.parametrize("state_backend", ["cursor", "redis"]) -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio -async def test_deduplication_merger_with_view_session_child(state_backend, redis_client) -> None: - # MergerViewSession always requires Redis, so this test always uses redis_client. - # We still validate both dedup state backends. - base_items = [{"id": i, "src": "vs"} for i in range(1, 200)] - extra_items = [{"id": i, "src": "extra"} for i in range(50, 260)] +async def test_dedup_percentage_gradient_slot_ownership_cursor_backend() -> None: + """Dedup must preserve gradient chunking semantics. - methods_dict = { - "vs": make_offset_paged_method(base_items), - "extra": make_offset_paged_method(extra_items), - } + For limit=10, size_to_step=5, from/to percentages should yield chunks: + - first 5: 3 from A, 2 from B + - next 5: 2 from A, 3 from B + Dedup must refill within each leaf so these chunk sizes remain true. + """ - view_session_config = { - "merger_id": "inner_view_session", - "type": "merger_view_session", - "session_size": 60, - "session_live_time": 60, - "data": {"subfeed_id": "sf_vs", "type": "subfeed", "method_name": "vs"}, - "deduplicate": True, - "dedup_key": "id", + a_items = [{"id": i, "src": "A"} for i in range(1, 300)] + # Start with duplicates, then provide unique tail. + b_items = ([{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)]) + + methods_dict = { + "a": make_offset_paged_method(a_items), + "b": make_offset_paged_method(b_items), } config = { - "merger_id": f"dedup_vs_{state_backend}", + "merger_id": "dedup_gradient", "type": "merger_deduplication", "dedup_key": "id", - "state_backend": state_backend, - "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": view_session_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_vs", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra_vs2", "type": "subfeed", "method_name": "extra"}}, - ], + "state_backend": "cursor", + "cursor_compress": True, + "max_refill_loops": 50, + "data": { + "merger_id": "grad_mix", + "type": "merger_percentage_gradient", + "item_from": { + "percentage": 60, + "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + }, + "item_to": { + "percentage": 40, + "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + }, + "step": 20, + "size_to_step": 5, + "shuffle": False, + }, } - res_1, res_2 = await _run_two_pages( - config=config, + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=20, - redis_client_instance=redis_client, - custom_deduplication_key=f"vs_{state_backend}", - custom_view_session_key=f"vs_{state_backend}", + limit=10, + next_page=FeedResultNextPage(data={}), ) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id=f"dedup_vs_{state_backend}", state_backend=state_backend) - assert "inner_view_session" in res_1.next_page.data - - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_vs", "sf_extra_vs", "sf_extra_vs2"]) - - -@pytest.mark.asyncio -async def test_deduplication_merger_with_append_distribute_cursor_backend() -> None: - # MergerAppendDistribute (type merger_distribute) + two extra subfeeds. - s1 = [{"id": i, "src": "s1", "group": "g1" if i % 2 == 0 else "g2"} for i in range(1, 120)] - s2 = [{"id": i, "src": "s2", "group": "g2" if i % 3 == 0 else "g3"} for i in range(60, 200)] - extra = [{"id": i, "src": "extra", "group": "g9"} for i in range(100, 240)] - methods_dict = { - "s1": make_offset_paged_method(s1, max_per_call=1), - "s2": make_offset_paged_method(s2, max_per_call=1), - "extra": make_offset_paged_method(extra, max_per_call=1), - } - - distribute_config = { - "merger_id": "inner_distribute_unused", - "type": "merger_distribute", - "distribution_key": "group", - "items": [ - {"subfeed_id": "sf_s1", "type": "subfeed", "method_name": "s1"}, - {"subfeed_id": "sf_s2", "type": "subfeed", "method_name": "s2"}, - ], - } + assert len(res.data) == 10 + _assert_no_dupes_in_page(res.data) - config = { - "merger_id": "dedup_distribute_cursor", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": distribute_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_dist", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra_dist2", "type": "subfeed", "method_name": "extra"}}, - ], - } + sources = _sources(res.data) + assert sources[:3] == ["A", "A", "A"] + assert sources[3:5] == ["B", "B"] + assert sources[5:7] == ["A", "A"] + assert sources[7:10] == ["B", "B", "B"] - res_1, res_2 = await _run_two_pages(config=config, methods_dict=methods_dict, user_id="u", limit=25) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_distribute_cursor", state_backend="cursor") - for key in ("sf_s1", "sf_s2"): - assert key in res_1.next_page.data - _assert_cursor_monotonic_if_present(res_1, res_2, ["sf_s1", "sf_s2", "sf_extra_dist", "sf_extra_dist2"]) +@pytest.mark.asyncio +async def test_dedup_preserves_append_priority_and_advances_cursors_cursor_backend() -> None: + """Append order is the priority signal; dedup must not let later sources win duplicates. + Also asserts that a leaf cursor advances even when items are skipped as duplicates. + """ -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_deduplication_merger_with_append_distribute_redis_backend(redis_client) -> None: - s1 = [{"id": i, "src": "s1", "group": "g1" if i % 2 == 0 else "g2"} for i in range(1, 120)] - s2 = [{"id": i, "src": "s2", "group": "g2" if i % 3 == 0 else "g3"} for i in range(60, 200)] - extra = [{"id": i, "src": "extra", "group": "g9"} for i in range(100, 240)] + a_items = [ + {"id": 1, "src": "A"}, + {"id": 2, "src": "A"}, + ] + # B repeats A's ids first, then continues with unique ids. + b_items = [{"id": i, "src": "B"} for i in range(1, 50)] methods_dict = { - "s1": make_offset_paged_method(s1, max_per_call=1), - "s2": make_offset_paged_method(s2, max_per_call=1), - "extra": make_offset_paged_method(extra, max_per_call=1), - } - - distribute_config = { - "merger_id": "inner_distribute_unused_r", - "type": "merger_distribute", - "distribution_key": "group", - "items": [ - {"subfeed_id": "sf_s1_r", "type": "subfeed", "method_name": "s1"}, - {"subfeed_id": "sf_s2_r", "type": "subfeed", "method_name": "s2"}, - ], + "a": make_offset_paged_method(a_items), + "b": make_offset_paged_method(b_items), } config = { - "merger_id": "dedup_distribute_redis", + "merger_id": "dedup_append", "type": "merger_deduplication", "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "overfetch_factor": 2, - "items": [ - {"priority": 10, "data": distribute_config}, - {"priority": 5, "data": {"subfeed_id": "sf_extra_dist_r", "type": "subfeed", "method_name": "extra"}}, - {"priority": 0, "data": {"subfeed_id": "sf_extra_dist2_r", "type": "subfeed", "method_name": "extra"}}, - ], + "state_backend": "cursor", + "cursor_compress": True, + "max_refill_loops": 20, + "data": { + "merger_id": "append_mix", + "type": "merger_append", + "shuffle": False, + "items": [ + {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + ], + }, } - res_1, res_2 = await _run_two_pages( - config=config, + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( methods_dict=methods_dict, user_id="u", - limit=25, - redis_client_instance=redis_client, - custom_deduplication_key="distribute", + limit=5, + next_page=FeedResultNextPage(data={}), ) - _assert_two_pages_no_overlap(res_1, res_2) - _assert_dedup_backend_state(res=res_2, merger_id="dedup_distribute_redis", state_backend="redis") - _assert_cursor_monotonic_if_present( - res_1, - res_2, - ["sf_s1_r", "sf_s2_r", "sf_extra_dist_r", "sf_extra_dist2_r"], - ) + assert _ids(res.data) == [1, 2, 3, 4, 5] + assert _sources(res.data)[:2] == ["A", "A"] + assert _sources(res.data)[2:] == ["B", "B", "B"] + + # B had to scan past duplicated ids 1 and 2, so its cursor should advance + # farther than the number of items it contributed to the final page. + assert "sf_b" in res.next_page.data + assert isinstance(res.next_page.data["sf_b"].after, int) + b_contributed = sum(1 for x in res.data if x.get("src") == "B") + assert res.next_page.data["sf_b"].after > b_contributed diff --git a/tests/test_parsing_config.py b/tests/test_parsing_config.py index 291971b..66c2387 100644 --- a/tests/test_parsing_config.py +++ b/tests/test_parsing_config.py @@ -54,6 +54,5 @@ async def test_parsing_config_deduplication_merger() -> None: assert isinstance(feed_manager.feed_config, FeedConfig) assert isinstance(feed_manager.feed_config.feed, MergerDeduplication) - assert len(feed_manager.feed_config.feed.items) == 2 - assert feed_manager.feed_config.feed.items[0].priority == 100 - assert isinstance(feed_manager.feed_config.feed.items[0].data, SubFeed) + # Deduplication merger is a wrapper around a single child feed. + assert isinstance(feed_manager.feed_config.feed.data, (MergerPercentage, SubFeed)) From 3ed4a094b70fb659243ece10b4bbc3af3841654e Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 22:36:17 +0000 Subject: [PATCH 05/41] More tests. --- tests/test_merger_deduplication.py | 449 +++++++++++++++++++++++++++++ 1 file changed, 449 insertions(+) diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 5b046a8..8347e3d 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -71,6 +71,100 @@ def _assert_pages_no_overlap(res_1, res_2): assert not (set(_ids(res_1.data)) & set(_ids(res_2.data))) +def _inner_append_config(*, merger_id: str, subfeed_id: str, method_name: str, dedup_priority: int): + return { + "merger_id": merger_id, + "type": "merger_append", + # Important: dedup deletion priority must be visible at this node so parent mergers + # can fetch higher-priority subtrees first when a dedup wrapper is active. + "dedup_priority": dedup_priority, + "shuffle": False, + "items": [ + { + "subfeed_id": subfeed_id, + "type": "subfeed", + "method_name": method_name, + "dedup_priority": dedup_priority, + } + ], + } + + +def _build_deep_priority_tree_for_merger_type(*, merger_type: str): + """Return a deep tree config where low/high leaves overlap by id. + + The inner leaves are wrapped into an append merger to ensure a "deep" tree even + when the outer merger is flat. + """ + + low = _inner_append_config(merger_id="inner_low", subfeed_id="sf_low", method_name="low", dedup_priority=0) + high = _inner_append_config(merger_id="inner_high", subfeed_id="sf_high", method_name="high", dedup_priority=100) + + if merger_type == "merger_append": + return { + "merger_id": "outer_append", + "type": "merger_append", + "shuffle": False, + # Put low first intentionally; priority must still make high win for overlapping ids. + "items": [low, high], + } + + if merger_type == "merger_distribute": + return { + "merger_id": "outer_dist", + "type": "merger_distribute", + "distribution_key": "user_id", + # Put low first intentionally. + "items": [low, high], + } + + if merger_type == "merger_percentage": + return { + "merger_id": "outer_pct", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": low}, + {"percentage": 50, "data": high}, + ], + } + + if merger_type == "merger_percentage_gradient": + return { + "merger_id": "outer_grad", + "type": "merger_percentage_gradient", + "item_from": {"percentage": 60, "data": low}, + "item_to": {"percentage": 40, "data": high}, + "step": 20, + "size_to_step": 5, + "shuffle": False, + } + + if merger_type == "merger_positional": + # High priority on positional branch so it must win duplicates. + high_pos = _inner_append_config( + merger_id="inner_pos_high", + subfeed_id="sf_high", + method_name="high", + dedup_priority=100, + ) + low_def = _inner_append_config( + merger_id="inner_def_low", + subfeed_id="sf_low", + method_name="low", + dedup_priority=0, + ) + return { + "merger_id": "outer_pos", + "type": "merger_positional", + "positions": [1, 3, 5, 7, 9, 11], + "positional": high_pos, + "default": low_def, + } + + raise AssertionError(f"Unknown merger_type: {merger_type}") + + @pytest.mark.asyncio async def test_dedup_positional_slot_ownership_cursor_backend() -> None: """Positional slots must remain owned by the positional branch. @@ -285,6 +379,79 @@ async def test_dedup_deep_tree_cursor_backend() -> None: assert _sources(res_2.data)[3] == "P" +@pytest.mark.parametrize( + "merger_type", + [ + "merger_append", + "merger_distribute", + "merger_positional", + "merger_percentage", + "merger_percentage_gradient", + ], +) +@pytest.mark.asyncio +async def test_dedup_deletion_priority_works_for_deep_trees_all_merger_types(merger_type: str) -> None: + """Deletion priority must work even in deep trees, across merger types. + + For overlapping ids, higher dedup_priority leaf must supply the winning entity. + """ + + # For mixing mergers (percentage/gradient/positional), identical id ranges are enough: the + # high-priority leaf should claim the first chunk of ids and the other leaf must skip them. + # + # For append/distribute, we must ensure BOTH branches contribute to the output (otherwise + # "priority" is unobservable because earlier branches can fill the page). We do that by + # making the low branch short and duplicate-heavy. + if merger_type in {"merger_append", "merger_distribute"}: + low_items = [ + {"id": 1, "user_id": "u0", "src": "low"}, + {"id": 2, "user_id": "u1", "src": "low"}, + {"id": 3, "user_id": "u2", "src": "low"}, + {"id": 1000, "user_id": "u0", "src": "low"}, + {"id": 1001, "user_id": "u1", "src": "low"}, + ] + high_items = [{"id": i, "user_id": f"u{i%3}", "src": "high"} for i in range(1, 200)] + else: + low_items = [{"id": i, "user_id": f"u{i%3}", "src": "low"} for i in range(1, 200)] + high_items = [{"id": i, "user_id": f"u{i%3}", "src": "high"} for i in range(1, 200)] + + methods_dict = { + "low": make_offset_paged_method(low_items), + "high": make_offset_paged_method(high_items), + } + + deep_tree = _build_deep_priority_tree_for_merger_type(merger_type=merger_type) + config = { + "merger_id": f"dedup_priority_deep_{merger_type}", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "data": deep_tree, + } + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=FeedResultNextPage(data={}), + ) + + _assert_no_dupes_in_page(res.data) + + # Priority is about which source wins overlapping ids (not about output order). + winning = {item["id"]: item["src"] for item in res.data} + assert all(winning[i] == "high" for i in range(1, 6) if i in winning) + + # Placement invariant for positional: positional slots must still be owned by positional branch. + if merger_type == "merger_positional": + sources = _sources(res.data) + assert sources[0] == "high" + assert sources[2] == "high" + assert sources[4] == "high" + + @pytest.mark.asyncio async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cursors() -> None: """When overfetch_factor>1, leaf cursors must be rewound to inspected count. @@ -341,6 +508,15 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu assert len(res_1.data) == 8 _assert_no_dupes_in_page(res_1.data) + # Dedup merger cursor must exist and advance page. + assert "dedup_overfetch" in res_1.next_page.data + assert res_1.next_page.data["dedup_overfetch"].page == 2 + assert res_1.next_page.data["dedup_overfetch"].after is not None + + # Positional merger cursor must exist and advance page. + assert "pos_overfetch" in res_1.next_page.data + assert res_1.next_page.data["pos_overfetch"].page == 2 + # Deep descendant cursors: positional leaf requests 2 items; percentage leaves request 4 each. # With overfetch_factor=3, internal calls may request 6/12, but cursor must not advance that far. assert res_1.next_page.data["sf_p"].after == 2 @@ -359,6 +535,9 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu _assert_no_dupes_in_page(res_2.data) _assert_pages_no_overlap(res_1, res_2) + assert res_2.next_page.data["dedup_overfetch"].page == 3 + assert res_2.next_page.data["pos_overfetch"].page == 3 + assert res_2.next_page.data["sf_p"].after == 4 assert res_2.next_page.data["sf_d1"].after == 8 assert res_2.next_page.data["sf_d2"].after == 8 @@ -477,6 +656,276 @@ async def test_dedup_redis_backend_cross_page(redis_client) -> None: assert len(members) >= len(set(_ids(res_1.data) + _ids(res_2.data))) +@pytest.mark.asyncio +async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf_cursor_exactly() -> None: + """Append: across pages there is no overlap; refill advances cursors correctly. + + This uses a max_per_call=1 method for the duplicate-heavy leaf so the wrapper + must do multiple client calls (refill loops). + """ + + a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] + b_items = [{"id": i, "src": "B"} for i in range(1, 50)] + + methods_dict = { + "a": make_offset_paged_method(a_items), + # Force multiple internal calls. + "b": make_offset_paged_method(b_items, max_per_call=1), + } + + config = { + "merger_id": "dedup_append_pages", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "max_refill_loops": 50, + "data": { + "merger_id": "append_mix", + "type": "merger_append", + "shuffle": False, + "items": [ + {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage(data={}), + ) + + assert _ids(res_1.data) == [1, 2, 3, 4, 5] + assert res_1.next_page.data["dedup_append_pages"].page == 2 + + # In dedup-active append mode, each child is requested with the full page limit (5). + # B must therefore collect 5 unique items while skipping 2 duplicates -> scan ids 1..7. + assert res_1.next_page.data["sf_b"].after == 7 + + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=res_1.next_page, + ) + + assert len(res_2.data) == 5 + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + +@pytest.mark.asyncio +async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill() -> None: + """Distribute: duplicates skipped per-leaf and page slices don't overlap.""" + + # A is short so B must contribute. + items_a = [{"id": i, "user_id": f"u{i%2}", "src": "A"} for i in range(1, 4)] + # B overlaps A by id and continues. + items_b = [{"id": i, "user_id": f"u{i%2}", "src": "B"} for i in range(1, 200)] + + methods_dict = { + "a": make_offset_paged_method(items_a), + "b": make_offset_paged_method(items_b), + } + + config = { + "merger_id": "dedup_dist_pages", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "dist", + "type": "merger_distribute", + "distribution_key": "user_id", + "items": [ + {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) + res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) + + assert len(res_1.data) == 10 + assert len(res_2.data) == 10 + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + # Placement/refill: B must skip duplicate ids 1..3 and still fill the page. + b_ids_1 = [x["id"] for x in res_1.data if x.get("src") == "B"] + assert b_ids_1 and min(b_ids_1) >= 4 + + +@pytest.mark.asyncio +async def test_dedup_percentage_gradient_cursor_backend_across_pages() -> None: + a_items = [{"id": i, "src": "A"} for i in range(1, 300)] + b_items = ([{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)]) + + methods_dict = { + "a": make_offset_paged_method(a_items), + "b": make_offset_paged_method(b_items), + } + + config = { + "merger_id": "dedup_grad_pages", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "max_refill_loops": 50, + "data": { + "merger_id": "grad_mix", + "type": "merger_percentage_gradient", + "item_from": {"percentage": 60, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, + "item_to": {"percentage": 40, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, + "step": 20, + "size_to_step": 5, + "shuffle": False, + }, + } + + merger = parse_model(MergerDeduplication, config) + res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) + res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) + + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + # Gradient merger cursor should exist and advance. + assert res_1.next_page.data["grad_mix"].page == 2 + assert res_2.next_page.data["grad_mix"].page == 3 + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_dedup_redis_backend_cross_page_append(redis_client) -> None: + items_a = [{"id": i, "src": "A"} for i in range(1, 20)] + items_b = [{"id": i, "src": "B"} for i in range(1, 300)] + + methods_dict = { + "a": make_offset_paged_method(items_a), + "b": make_offset_paged_method(items_b), + } + + config = { + "merger_id": "dedup_redis_append", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "redis", + "state_ttl_seconds": 60, + "data": { + "merger_id": "append_mix", + "type": "merger_append", + "shuffle": False, + "items": [ + {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, + {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client, + custom_deduplication_key="t2", + ) + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=res_1.next_page, + redis_client=redis_client, + custom_deduplication_key="t2", + ) + + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + assert res_2.next_page.data["dedup_redis_append"].after is None + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_dedup_wrapper_with_view_session_merger(redis_client) -> None: + """Dedup wrapper must work when the child is a view_session merger.""" + + # Two leaves with overlapping ids; view_session computes a session once. + items_low = [{"id": i, "src": "low"} for i in range(1, 100)] + items_high = [{"id": i, "src": "high"} for i in range(1, 100)] + + methods_dict = { + "low": make_offset_paged_method(items_low), + "high": make_offset_paged_method(items_high), + } + + config = { + "merger_id": "dedup_vs", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "vs", + "type": "merger_view_session", + "session_size": 30, + "session_live_time": 60, + "deduplicate": False, + "shuffle": False, + "data": { + "merger_id": "pct", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}}, + {"percentage": 50, "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high", "dedup_priority": 100}}, + ], + }, + }, + } + + merger = parse_model(MergerDeduplication, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=FeedResultNextPage(data={}), + redis_client=redis_client, + custom_view_session_key="vs1", + ) + + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=res_1.next_page, + redis_client=redis_client, + custom_view_session_key="vs1", + ) + + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + # Deletion priority: for the overlapping early ids, the winning entity must be from high. + winning = {item["id"]: item["src"] for item in (res_1.data + res_2.data)} + assert all(winning[i] == "high" for i in range(1, 11) if i in winning) + + @pytest.mark.asyncio async def test_dedup_append_distribute_cursor_backend_no_dupes() -> None: items_a = [{"id": i, "user_id": f"u{i%3}", "src": "A"} for i in range(1, 200)] From aa2b608c60c848d9cd309eea3e4d426dd1be20fb Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 22:59:18 +0000 Subject: [PATCH 06/41] More tests and minor after fixes. --- smartfeed/schemas.py | 8 +- tests/test_merger_deduplication.py | 258 +++++++++++++++++++++++++++++ 2 files changed, 263 insertions(+), 3 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 3320fb7..b1e8900 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1413,12 +1413,14 @@ async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPag remaining = limit - len(collected) # Safe oversampling: only when we can rewind integer-offset cursors. - can_overfetch = isinstance(next_page.after, (int, type(None))) + # IMPORTANT: `after` can be many shapes (str/dict/etc) and may start as None. + # We only enable overfetch when `after` is already an int offset. + can_overfetch = isinstance(next_page.after, int) request_limit = max(1, remaining) if can_overfetch and self.overfetch_factor > 1: request_limit = max(1, remaining * self.overfetch_factor) - start_after = 0 if next_page.after is None else int(next_page.after) + start_after: Optional[int] = int(next_page.after) if can_overfetch else None method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) if not isinstance(method_result, FeedResultClient): @@ -1475,7 +1477,7 @@ async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPag # If we oversampled with a simple integer cursor, rewind to the point we actually consumed. # This prevents skipping un-inspected items that were fetched but not needed. - if can_overfetch and request_limit > remaining: + if can_overfetch and request_limit > remaining and start_after is not None: end_after = next_page.after if isinstance(end_after, int) and end_after == start_after + len(method_result.data): next_page.after = start_after + consumed_in_batch diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 8347e3d..94d122a 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -28,6 +28,100 @@ async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unuse return _method +def make_string_after_paged_method(items, *, max_per_call=None, after_field="created_at"): + """A subfeed method whose cursor is a string (e.g. timestamp). + + Cursor semantics: `after` is the last returned `created_at` value (monotonic). + """ + + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + effective_limit = limit + if isinstance(max_per_call, int) and max_per_call > 0: + effective_limit = min(effective_limit, max_per_call) + + after = next_page.after + start_idx = 0 + if isinstance(after, str) and after: + # Find first item with created_at > after + for i, item in enumerate(items): + if str(item[after_field]) > after: + start_idx = i + break + else: + start_idx = len(items) + + result_data = items[start_idx : start_idx + effective_limit] + has_next_page = (start_idx + len(result_data)) < len(items) + + if result_data: + next_page.after = str(result_data[-1][after_field]) + next_page.page += 1 + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) + + return _method + + +def make_profile_dict_after_method( + profiles_to_items, + *, + max_per_call=None, + after_key="after", +): + """A subfeed method whose cursor is a dict of per-profile offsets. + + Example shape: after = {"p1": 0, "p2": 0} + Cursor semantics: each profile offset increments as items are *read*. + """ + + profile_ids = list(profiles_to_items.keys()) + + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + effective_limit = limit + if isinstance(max_per_call, int) and max_per_call > 0: + effective_limit = min(effective_limit, max_per_call) + + after = next_page.after + if not isinstance(after, dict): + after = {pid: 0 for pid in profile_ids} + else: + after = dict(after) + for pid in profile_ids: + after.setdefault(pid, 0) + + result = [] + has_next_page = False + + # Build a cyclic iteration over profiles. + active_profiles = [pid for pid in profile_ids] + + i = 0 + while active_profiles and len(result) < effective_limit: + pid = active_profiles[i % len(active_profiles)] + idx = after.get(pid, 0) + items = profiles_to_items.get(pid, []) + + if idx >= len(items): + # This profile is exhausted. + active_profiles.remove(pid) + continue + + result.append(items[idx]) + after[pid] = idx + 1 + i += 1 + + # Determine if any profile still has unread items. + for pid in profile_ids: + if after.get(pid, 0) < len(profiles_to_items.get(pid, [])): + has_next_page = True + break + + next_page.after = after + next_page.page += 1 + return FeedResultClient(data=result, next_page=next_page, has_next_page=has_next_page) + + return _method + + def _assert_cursor_monotonic_if_present(res_1, res_2, keys): for key in keys: if key not in res_1.next_page.data: @@ -719,6 +813,170 @@ async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf _assert_pages_no_overlap(res_1, res_2) +@pytest.mark.asyncio +async def test_dedup_refill_loops_advance_dict_after_cursor_not_just_page() -> None: + """Dedup refill loops must correctly advance dict-shaped `after` cursors.""" + + # A produces ids 1,2. + a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] + + # B produces ids 1.. in round-robin across profiles; cursor is per-profile offsets. + b_profiles = { + "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], + "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], + } + + methods_dict = { + "a": make_offset_paged_method(a_items), + "b": make_profile_dict_after_method(b_profiles), + } + + # Use a percentage merger so B is asked for a small limit (2 items for limit=4). + # This forces refill loops when B's first batch is all duplicates. + config = { + "merger_id": "dedup_dict_after", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "max_refill_loops": 50, + "data": { + "merger_id": "pct_mix", + "type": "merger_percentage", + "shuffle": False, + "items": [ + { + "percentage": 50, + "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a", "dedup_priority": 100}, + }, + { + "percentage": 50, + "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b", "dedup_priority": 0}, + }, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=4, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) == 4 + _assert_no_dupes_in_page(res.data) + assert set(_ids(res.data)) == {1, 2, 3, 4} + assert "sf_b" in res.next_page.data + assert isinstance(res.next_page.data["sf_b"].after, dict) + + # B contributed 2 items (3,4) but must have *read* 4 items (1..4) to skip duplicates. + b_after = res.next_page.data["sf_b"].after + read_count = sum(int(v) for v in b_after.values()) + assert read_count == 4 + + +@pytest.mark.asyncio +async def test_dedup_overfetch_does_not_overadvance_non_int_after_cursor() -> None: + """overfetch_factor must not cause over-advancement for non-rewindable cursors.""" + + # Single subfeed with dict after cursor; no dedup skips should happen. + profiles = { + "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], + "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], + } + + methods_dict = { + "b": make_profile_dict_after_method(profiles), + } + + config = { + "merger_id": "dedup_nonint_overfetch", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 5, + "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, + } + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=4, next_page=FeedResultNextPage(data={})) + + assert len(res.data) == 4 + after = res.next_page.data["sf_b"].after + assert isinstance(after, dict) + # If overfetch were incorrectly applied, we'd see more than 4 reads. + assert sum(int(v) for v in after.values()) == 4 + + +@pytest.mark.asyncio +async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplicates() -> None: + """Overfetch should be safe: when we oversample, we must rewind offset cursors. + + Scenario: + - A (high priority) returns ids 1..5 + - B (low priority) initially returns only duplicates (1..5) + - On the next refill loop, B overfetches but must rewind `after` to inspected count + so it doesn't skip items. + """ + + items_a = [{"id": i, "src": "A"} for i in range(1, 300)] + items_b = [{"id": i, "src": "B"} for i in range(1, 300)] + + methods_dict = { + "a": make_offset_paged_method(items_a), + "b": make_offset_paged_method(items_b), + } + + config = { + "merger_id": "dedup_overfetch_rewind", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "overfetch_factor": 3, + "max_refill_loops": 20, + "data": { + "merger_id": "pct_mix", + "type": "merger_percentage", + "shuffle": False, + "items": [ + { + "percentage": 50, + "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a", "dedup_priority": 100}, + }, + { + "percentage": 50, + "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b", "dedup_priority": 0}, + }, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=10, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) == 10 + _assert_no_dupes_in_page(res.data) + + # A provides 1..5, B must provide 6..10. + winning = {item["id"]: item["src"] for item in res.data} + assert all(winning[i] == "A" for i in range(1, 6)) + assert all(winning[i] == "B" for i in range(6, 11)) + + # Cursor rewind check: + # - First loop for B reads 5 duplicates -> after becomes 5 + # - Second loop overfetches, but must rewind to inspected 5 more -> after should end at 10 + assert res.next_page.data["sf_b"].after == 10 + + @pytest.mark.asyncio async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill() -> None: """Distribute: duplicates skipped per-leaf and page slices don't overlap.""" From c6184dbc3e74ba3dfa98fa79ea8763ecbaf7a33c Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 23:46:58 +0000 Subject: [PATCH 07/41] Minor refactor. --- smartfeed/schemas.py | 465 +++++++++++++++++++++++++++++-------------- 1 file changed, 316 insertions(+), 149 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index b1e8900..16fe0cc 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,4 +1,5 @@ import base64 +from dataclasses import dataclass import inspect import json import logging @@ -9,10 +10,99 @@ from typing import Annotated, Any, Callable, Dict, List, Literal, Optional, Union, no_type_check import redis -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, PrivateAttr, model_validator from redis.asyncio import Redis as AsyncRedis from redis.asyncio import RedisCluster as AsyncRedisCluster + +def _pydantic_deep_copy(model: Any) -> Any: + """Deep copy helper compatible with Pydantic v1 and v2.""" + + if hasattr(model, "model_copy"): + return model.model_copy(deep=True) # type: ignore[attr-defined] + return model.copy(deep=True) + + +class _DedupState(ABC): + @abstractmethod + def should_accept(self, key: str, priority: int) -> bool: + raise NotImplementedError + + @abstractmethod + def record(self, key: str, priority: int) -> None: + raise NotImplementedError + + async def prefetch(self, keys: List[str]) -> None: + return + + +@dataclass +class _CursorDedupState(_DedupState): + seen_priority_map: Dict[str, int] + seen_updates_in_order: List[tuple[str, int]] + seen_request_set: set[str] + + def should_accept(self, key: str, priority: int) -> bool: + if key in self.seen_request_set: + return False + existing_priority = self.seen_priority_map.get(key) + if existing_priority is not None and priority <= existing_priority: + return False + return True + + def record(self, key: str, priority: int) -> None: + self.seen_priority_map[key] = priority + self.seen_updates_in_order.append((key, priority)) + self.seen_request_set.add(key) + + +@dataclass +class _RedisDedupState(_DedupState): + redis_client: Union[redis.Redis, AsyncRedis] + redis_state_key: str + redis_seen_cache: Dict[str, Optional[int]] + redis_new_scores: Dict[str, int] + seen_request_set: set[str] + zmscore: Callable[[Union[redis.Redis, AsyncRedis], str, List[str]], Any] + + async def prefetch(self, keys: List[str]) -> None: + if not keys: + return + unique: List[str] = [] + seen: set[str] = set() + for k in keys: + if k in self.seen_request_set: + continue + if k in self.redis_seen_cache: + continue + if k in seen: + continue + seen.add(k) + unique.append(k) + + if not unique: + return + + scores = self.zmscore(self.redis_client, self.redis_state_key, unique) + if inspect.iscoroutine(scores): + scores = await scores + + for k, s in zip(unique, scores): + self.redis_seen_cache[k] = None if s is None else int(s) + + def should_accept(self, key: str, priority: int) -> bool: + if key in self.seen_request_set: + return False + existing_priority = self.redis_seen_cache.get(key) + if existing_priority is not None and priority <= existing_priority: + return False + return True + + def record(self, key: str, priority: int) -> None: + self.seen_request_set.add(key) + self.redis_seen_cache[key] = priority + self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) + FeedTypes = Annotated[ Union[ "MergerDeduplication", @@ -1154,6 +1244,8 @@ class MergerDeduplication(BaseFeedConfigModel): max_refill_loops: int = 20 + _descendant_cursor_keys_cache: Optional[set[str]] = PrivateAttr(default=None) + @model_validator(mode="after") def validate_merger_deduplication(self) -> "MergerDeduplication": if self.overfetch_factor < 1: @@ -1199,6 +1291,18 @@ def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str] return keys + def _get_descendant_cursor_keys_cached(self) -> set[str]: + cached = self._descendant_cursor_keys_cache + if cached is None: + cached = self._collect_descendant_cursor_keys(self.data) + self._descendant_cursor_keys_cache = cached + return cached + + def _reset_descendant_cursors(self, next_page: FeedResultNextPage) -> None: + descendant_keys = self._get_descendant_cursor_keys_cached() + for key in descendant_keys: + next_page.data.pop(key, None) + def _normalize_key(self, value: Any) -> str: if isinstance(value, (str, int)): return str(value) @@ -1221,6 +1325,161 @@ def _extract_dedup_value(self, item: Any) -> Any: ) return value + def _get_entity_key(self, entity: Any) -> Optional[str]: + """Return normalized dedup key for entity, or None if entity should be skipped.""" + + raw_value = self._extract_dedup_value(entity) + if raw_value is None: + if self.missing_key_policy == "drop": + return None + if self.missing_key_policy == "keep": + raw_value = ("__missing__", id(entity)) + return self._normalize_key(raw_value) + + def _compute_overfetch_params(self, *, remaining: int, next_after: Any) -> tuple[bool, int, Optional[int]]: + """Compute safe overfetch params. + + Overfetch is only safe when `after` is an integer offset (so we can rewind). + + Returns: (can_overfetch, request_limit, start_after) + """ + + can_overfetch = isinstance(next_after, int) + request_limit = max(1, remaining) + if can_overfetch and self.overfetch_factor > 1: + request_limit = max(1, remaining * self.overfetch_factor) + start_after: Optional[int] = int(next_after) if can_overfetch else None + return can_overfetch, request_limit, start_after + + def _iter_subfeeds(self, feed: BaseFeedConfigModel): + if isinstance(feed, SubFeed): + yield feed + return + + for attr_name in ("data", "positional", "default"): + inner = getattr(feed, attr_name, None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + for attr_name in ("item_from", "item_to"): + wrapper = getattr(feed, attr_name, None) + inner = getattr(wrapper, "data", None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + items = getattr(feed, "items", None) + if isinstance(items, list): + for item in items: + if isinstance(item, BaseFeedConfigModel): + yield from self._iter_subfeeds(item) + continue + inner = getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + def _register_wrapped_subfeed_method( + self, + *, + subfeed: "SubFeed", + original_methods_dict: Dict[str, Callable], + rewritten_methods_dict: Dict[str, Callable], + dedup_state: "_DedupState", + ) -> None: + original_name = subfeed.method_name + original_method = original_methods_dict[original_name] + unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" + + # Idempotency: if the same subfeed id appears multiple times, don't re-wrap. + if unique_name in rewritten_methods_dict: + subfeed.method_name = unique_name + return + + subfeed.method_name = unique_name + leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) + + wrapped = self._make_wrapped_leaf_method( + original_method=original_method, + dedup_state=dedup_state, + leaf_priority=leaf_priority, + ) + setattr(wrapped, "_smartfeed_original", original_method) + rewritten_methods_dict[unique_name] = wrapped + + def _make_wrapped_leaf_method( + self, + *, + original_method: Callable, + dedup_state: "_DedupState", + leaf_priority: int, + ) -> Callable: + async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPageInside, **kw: Any): + collected: List[Any] = [] + upstream_has_next_page = False + + loops = 0 + while len(collected) < limit and loops < self.max_refill_loops: + loops += 1 + before_len = len(collected) + + remaining = limit - len(collected) + can_overfetch, request_limit, start_after = self._compute_overfetch_params( + remaining=remaining, + next_after=next_page.after, + ) + + method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) + if not isinstance(method_result, FeedResultClient): + raise TypeError('SubFeed function must return "FeedResultClient" instance.') + + upstream_has_next_page = upstream_has_next_page or method_result.has_next_page + + inspected_count = 0 + + # Backend-specific optimization: Redis batches zmscore. + # For cursor backend, prefetch is a no-op and we avoid the extra pass entirely. + keys_by_index: Optional[List[Optional[str]]] = None + if isinstance(dedup_state, _RedisDedupState): + keys_by_index = [] + batch_keys: List[str] = [] + for entity in method_result.data: + key = self._get_entity_key(entity) + keys_by_index.append(key) + if key is not None: + batch_keys.append(key) + await dedup_state.prefetch(batch_keys) + + for idx, entity in enumerate(method_result.data, start=1): + inspected_count = idx + + key = keys_by_index[idx - 1] if keys_by_index is not None else self._get_entity_key(entity) + if key is None: + continue + + if not dedup_state.should_accept(key, leaf_priority): + continue + + collected.append(entity) + dedup_state.record(key, leaf_priority) + + if len(collected) >= limit: + break + + if len(collected) == before_len: + # No progress this loop. Stop if upstream is exhausted. + if not method_result.has_next_page: + break + + # If we oversampled with a simple integer cursor, rewind to the point we actually consumed. + # This prevents skipping un-inspected items that were fetched but not needed. + if can_overfetch and request_limit > remaining and start_after is not None: + end_after = next_page.after + if isinstance(end_after, int) and end_after == start_after + len(method_result.data): + next_page.after = start_after + inspected_count + + return FeedResultClient(data=collected, next_page=next_page, has_next_page=upstream_has_next_page) + + return _wrapped_method + def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> Dict[str, int]: entry = next_page.data.get(self.merger_id) if not entry or entry.after is None: @@ -1268,11 +1527,34 @@ def _encode_seen_for_cursor(self, seen_updates_in_order: List[tuple[str, int]]) "z": base64.urlsafe_b64encode(compressed).decode(), } - async def _redis_zscore(self, redis_client: Union[redis.Redis, AsyncRedis], key: str, member: str) -> Optional[float]: - res = redis_client.zscore(key, member) + async def _redis_zmscore( + self, + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + members: List[str], + ) -> List[Optional[float]]: + """Batch zscore for multiple members. + + Falls back to pipelined zscore when zmscore isn't available. + """ + + if not members: + return [] + + if hasattr(redis_client, "zmscore"): + res = redis_client.zmscore(key, members) # type: ignore[attr-defined] + if inspect.iscoroutine(res): + res = await res + # redis-py returns list[Optional[float]] + return [None if v is None else float(v) for v in list(res)] + + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + res = pipe.execute() if inspect.iscoroutine(res): res = await res - return None if res is None else float(res) + return [None if v is None else float(v) for v in list(res)] async def _redis_zadd_and_expire( self, @@ -1315,16 +1597,11 @@ async def get_data( if self.state_backend == "redis" and not redis_client: raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") - if hasattr(next_page, "model_copy"): - working_next_page = next_page.model_copy(deep=True) # type: ignore[attr-defined] - else: - working_next_page = next_page.copy(deep=True) + working_next_page = _pydantic_deep_copy(next_page) if is_fresh_session: # Reset cursors for all descendants under this merger so upstream nodes also restart. - descendant_keys = self._collect_descendant_cursor_keys(self.data) - for key in descendant_keys: - working_next_page.data.pop(key, None) + self._reset_descendant_cursors(working_next_page) # Shared dedup state (cross-page) seen_priority_map: Dict[str, int] = {} @@ -1336,6 +1613,8 @@ async def get_data( seen_request_set: set[str] = set(seen_priority_map.keys()) redis_state_key = "" + redis_new_scores: Dict[str, int] = {} + redis_seen_cache: Dict[str, Optional[int]] = {} if self.state_backend == "redis" and redis_client: redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) if is_fresh_session: @@ -1344,7 +1623,23 @@ async def get_data( if inspect.iscoroutine(deleted): await deleted - redis_new_scores: Dict[str, int] = {} + # Create a single state helper shared across all leaf wrappers. + if self.state_backend == "cursor": + dedup_state: _DedupState = _CursorDedupState( + seen_priority_map=seen_priority_map, + seen_updates_in_order=seen_updates_in_order, + seen_request_set=seen_request_set, + ) + else: + assert redis_client is not None + dedup_state = _RedisDedupState( + redis_client=redis_client, + redis_state_key=redis_state_key, + redis_seen_cache=redis_seen_cache, + redis_new_scores=redis_new_scores, + seen_request_set=seen_request_set, + zmscore=self._redis_zmscore, + ) # Preserve inner merger ordering/mixing semantics by deduplicating at the leaf method level # with a shared seen-set. @@ -1353,142 +1648,17 @@ async def get_data( # Create a deep copy of the child tree and rewrite each SubFeed to call a unique wrapper # so we can associate a dedup_priority with each leaf. child = self.data - if hasattr(child, "model_copy"): - child = child.model_copy(deep=True) # type: ignore[attr-defined] - else: - child = child.copy(deep=True) - - def iter_subfeeds(feed: BaseFeedConfigModel) -> List["SubFeed"]: - found: List[SubFeed] = [] - - if isinstance(feed, SubFeed): - found.append(feed) - return found - - for attr_name in ("data", "positional", "default"): - inner = getattr(feed, attr_name, None) - if isinstance(inner, BaseFeedConfigModel): - found.extend(iter_subfeeds(inner)) - - for attr_name in ("item_from", "item_to"): - wrapper = getattr(feed, attr_name, None) - inner = getattr(wrapper, "data", None) - if isinstance(inner, BaseFeedConfigModel): - found.extend(iter_subfeeds(inner)) - - items = getattr(feed, "items", None) - if isinstance(items, list): - for item in items: - if isinstance(item, BaseFeedConfigModel): - found.extend(iter_subfeeds(item)) - continue - inner = getattr(item, "data", None) - if isinstance(inner, BaseFeedConfigModel): - found.extend(iter_subfeeds(inner)) - - return found + child = _pydantic_deep_copy(child) rewritten_methods_dict = dict(original_methods_dict) - def wrap_leaf_method(*, subfeed: "SubFeed") -> None: - original_name = subfeed.method_name - original_method = original_methods_dict[original_name] - unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" - # Idempotency: if the same subfeed id appears multiple times, don't re-wrap. - if unique_name in rewritten_methods_dict: - subfeed.method_name = unique_name - return - subfeed.method_name = unique_name - leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) - - async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPageInside, **kw: Any): - collected: List[Any] = [] - local_seen: set[str] = set() - any_has_next_page = False - - loops = 0 - while len(collected) < limit and loops < self.max_refill_loops: - loops += 1 - before_len = len(collected) - - remaining = limit - len(collected) - # Safe oversampling: only when we can rewind integer-offset cursors. - # IMPORTANT: `after` can be many shapes (str/dict/etc) and may start as None. - # We only enable overfetch when `after` is already an int offset. - can_overfetch = isinstance(next_page.after, int) - request_limit = max(1, remaining) - if can_overfetch and self.overfetch_factor > 1: - request_limit = max(1, remaining * self.overfetch_factor) - - start_after: Optional[int] = int(next_page.after) if can_overfetch else None - - method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) - if not isinstance(method_result, FeedResultClient): - raise TypeError('SubFeed function must return "FeedResultClient" instance.') - - any_has_next_page = any_has_next_page or method_result.has_next_page - - consumed_in_batch = 0 - - for entity in method_result.data: - consumed_in_batch += 1 - raw_value = self._extract_dedup_value(entity) - if raw_value is None: - if self.missing_key_policy == "drop": - continue - if self.missing_key_policy == "keep": - raw_value = ("__missing__", id(entity)) - - key = self._normalize_key(raw_value) - if key in local_seen: - continue - - if key in seen_request_set: - continue - - if self.state_backend == "cursor": - existing_priority = seen_priority_map.get(key) - if existing_priority is not None and leaf_priority <= existing_priority: - continue - else: - assert redis_client is not None - existing_score = await self._redis_zscore(redis_client, redis_state_key, key) - if existing_score is not None and leaf_priority <= int(existing_score): - continue - - local_seen.add(key) - collected.append(entity) - - seen_request_set.add(key) - - if self.state_backend == "cursor": - seen_priority_map[key] = leaf_priority - seen_updates_in_order.append((key, leaf_priority)) - else: - redis_new_scores[key] = max(redis_new_scores.get(key, 0), leaf_priority) - - if len(collected) >= limit: - break - - if len(collected) == before_len: - # No progress this loop. Stop if upstream is exhausted. - if not method_result.has_next_page: - break - - # If we oversampled with a simple integer cursor, rewind to the point we actually consumed. - # This prevents skipping un-inspected items that were fetched but not needed. - if can_overfetch and request_limit > remaining and start_after is not None: - end_after = next_page.after - if isinstance(end_after, int) and end_after == start_after + len(method_result.data): - next_page.after = start_after + consumed_in_batch - - return FeedResultClient(data=collected, next_page=next_page, has_next_page=any_has_next_page) - - setattr(_wrapped_method, "_smartfeed_original", original_method) - rewritten_methods_dict[unique_name] = _wrapped_method - - for sf in iter_subfeeds(child): - wrap_leaf_method(subfeed=sf) + for sf in self._iter_subfeeds(child): + self._register_wrapped_subfeed_method( + subfeed=sf, + original_methods_dict=original_methods_dict, + rewritten_methods_dict=rewritten_methods_dict, + dedup_state=dedup_state, + ) child_result = await child.get_data( methods_dict=rewritten_methods_dict, @@ -1508,10 +1678,7 @@ async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPag if self.state_backend == "cursor": merger_after = self._encode_seen_for_cursor(seen_updates_in_order) - if hasattr(child_result.next_page, "model_copy"): - result_next_page = child_result.next_page.model_copy(deep=True) # type: ignore[attr-defined] - else: - result_next_page = child_result.next_page.copy(deep=True) + result_next_page = _pydantic_deep_copy(child_result.next_page) result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) return FeedResult(data=child_result.data, next_page=result_next_page, has_next_page=child_result.has_next_page) From b0fa3ea4d198ce807fd4df98a4838cf340df1c18 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Tue, 16 Dec 2025 23:54:27 +0000 Subject: [PATCH 08/41] Readme updated. --- README.md | 112 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/README.md b/README.md index fe96da4..a86e8a4 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,9 @@ Python-package для формирования ленты (Feed) из клиен - [Использование](#использование) - [Установка](#установка) - [Формирование конфигурации](#формирование-конфигурации) + - [MergerDeduplication (дедупликация)](#mergerdeduplication-дедупликация) + - [Параметры MergerDeduplication](#параметры-mergerdeduplication) + - [Важные нюансы (сброс, cursor/redis, overfetch)](#важные-нюансы-сброс-cursorredis-overfetch) - [Требования к клиентскому методу](#требования-к-клиентскому-методу) - [Запуск](#запуск) @@ -68,6 +71,115 @@ poetry add git+ssh://git@github.com:epoch8/looky-timeline.git }, ``` +### MergerDeduplication (дедупликация) + +MergerDeduplication — обёртка над одним дочерним узлом (merger или subfeed), которая удаляет дубли по ключу. + +Ключевые свойства реализации: + +- Дедупликация выполняется на уровне листьев (SubFeed), а не пост-обработкой результата мерджера. + Это важно: вложенные мерджеры (positional/percentage/gradient/append/distribute) сохраняют свои правила смешивания. + Если элемент удалён как дубль, MergerDeduplication «дозапросит» следующий элемент из того же источника. +- Состояние «уже видели» может храниться: + - в курсоре (state_backend="cursor") — удобно без Redis, но курсор может расти; + - в Redis (state_backend="redis") — удобно для большого состояния. + +Пример: обернуть существующую конфигурацию фида дедупликацией: + +```json +{ + "version": "1", + "feed": { + "merger_id": "dedup_main", + "type": "merger_deduplication", + "dedup_key": "id", + "missing_key_policy": "error", + "state_backend": "cursor", + "cursor_compress": true, + "cursor_max_keys": 2000, + "overfetch_factor": 2, + "max_refill_loops": 20, + "data": { + "merger_id": "merger_percent", + "type": "merger_percentage", + "items": [ + { + "percentage": 60, + "data": { + "subfeed_id": "sf_posts", + "type": "subfeed", + "method_name": "posts", + "dedup_priority": 10 + } + }, + { + "percentage": 40, + "data": { + "subfeed_id": "sf_ads", + "type": "subfeed", + "method_name": "ads", + "dedup_priority": 0 + } + } + ] + } + } +} +``` + +В примере выше, если `posts` и `ads` отдают объекты с одинаковым `id`, то «побеждает» источник с большим `dedup_priority`. + +### Параметры MergerDeduplication + +Обязательные поля: + +- `merger_id: str` — уникальный ID мерджера. +- `type: "merger_deduplication"` +- `data` — ровно один дочерний узел (subfeed или merger). + +Поля дедупликации: + +- `dedup_key: str | null` — имя ключа/атрибута для поиска дублей. + - если `null`, ключом считается сам объект (подходит, когда объекты уже hashable/строковые). +- `missing_key_policy: "error" | "keep" | "drop"` (default: `"error"`) + - `error`: выбросить ошибку, если у элемента нет `dedup_key`; + - `keep`: сохранить элемент, даже если ключа нет; + - `drop`: выкинуть элемент без ключа. + +Состояние seen (межстраничная дедупликация): + +- `state_backend: "cursor" | "redis"` (default: `"cursor"`) +- `state_ttl_seconds: int` (default: `3600`) — TTL для Redis состояния (только для backend=`redis`). +- `cursor_compress: bool` (default: `true`) — сжимать seen-состояние в cursor backend. +- `cursor_max_keys: int | null` — ограничить размер seen-состояния в cursor backend (полезно для контроля размера курсора). + +Производительность/поведение: + +- `overfetch_factor: int` (default: `1`) — «перезапрос» внутри листьев, чтобы быстрее добрать `limit` без множества рефиллов. +- `max_refill_loops: int` (default: `20`) — верхняя граница количества дозапросов на один лист. + +### Важные нюансы (сброс, cursor/redis, overfetch) + +- Сброс состояния при `page <= 0` или отсутствии курсора для `merger_id`. + - MergerDeduplication воспринимает это как «fresh session» и очищает курсоры всех дочерних узлов. + - Для backend=`redis` дополнительно удаляет ключ состояния в Redis. + +- Если `state_backend="redis"`, нужно передать `redis_client` в `FeedManager`. + - Ключ состояния в Redis строится как `dedup:{merger_id}:{user_id}`. + - Можно добавить суффикс через параметр запроса `custom_deduplication_key` (или `custom_view_session_key`), + чтобы разделять состояния для разных режимов выдачи. + +- Приоритет (`dedup_priority`) — это приоритет победы при конфликте дублей, а не порядок вывода. + - Больше `dedup_priority` → элемент «побеждает» и будет считаться seen с этим приоритетом. + - Это поле доступно у всех узлов (merger/subfeed) и используется MergerDeduplication при дедупликации. + +- overfetch работает безопасно только для «перематываемых» курсоров. + - Сейчас overfetch включается только если `next_page.after` у листа — целочисленный offset. + - Если `after` — строка/словарь/любой другой объект, он считается непрозрачным и overfetch не применяется. + +- Главный реальный bottleneck в дедупликации — не обёртки/копии, а рефиллы. + - Если дублей много и upstream-методы дорогие, стоит аккуратно подобрать `overfetch_factor` и `max_refill_loops`. + ### Требования к клиентскому методу Клиентский метод для получения данных должен обязательно включать в себя следующие параметры: From 8c0091536109e8b22db649ef743f441acabe089a Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Wed, 17 Dec 2025 00:14:40 +0000 Subject: [PATCH 09/41] Linter fixes. --- smartfeed/examples/example_client.py | 5 ++- smartfeed/manager.py | 5 ++- smartfeed/schemas.py | 64 +++++++++++++++++----------- 3 files changed, 46 insertions(+), 28 deletions(-) diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index 8b9d6a7..b1a55e3 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -25,8 +25,9 @@ class TestClientRequest(BaseModel): def validate_next_page(cls, value: Union[str, FeedResultNextPage]) -> Union[str, FeedResultNextPage]: if isinstance(value, str): payload = json.loads(base64.urlsafe_b64decode(value)) - if hasattr(FeedResultNextPage, "model_validate"): - return FeedResultNextPage.model_validate(payload) # type: ignore[attr-defined] + validate = getattr(FeedResultNextPage, "model_validate", None) + if validate is not None: + return validate(payload) return FeedResultNextPage.parse_obj(payload) return value diff --git a/smartfeed/manager.py b/smartfeed/manager.py index 7ac06f9..5ef6eb1 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -20,8 +20,9 @@ def __init__(self, config: Dict, methods_dict: Dict, redis_client: Optional[Unio :param redis_client: объект клиента Redis (для конфигурации с view_session = True). """ - if hasattr(FeedConfig, "model_validate"): - self.feed_config = FeedConfig.model_validate(config) # type: ignore[attr-defined] + validate = getattr(FeedConfig, "model_validate", None) + if validate is not None: + self.feed_config = validate(config) else: self.feed_config = FeedConfig.parse_obj(config) self.methods_dict = methods_dict diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 16fe0cc..f818b1b 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from collections import defaultdict, deque from random import shuffle -from typing import Annotated, Any, Callable, Dict, List, Literal, Optional, Union, no_type_check +from typing import Annotated, Any, Awaitable, Callable, Dict, Iterator, List, Literal, Optional, Union, cast, no_type_check import redis from pydantic import BaseModel, Field, PrivateAttr, model_validator @@ -19,7 +19,7 @@ def _pydantic_deep_copy(model: Any) -> Any: """Deep copy helper compatible with Pydantic v1 and v2.""" if hasattr(model, "model_copy"): - return model.model_copy(deep=True) # type: ignore[attr-defined] + return model.model_copy(deep=True) return model.copy(deep=True) @@ -63,7 +63,10 @@ class _RedisDedupState(_DedupState): redis_seen_cache: Dict[str, Optional[int]] redis_new_scores: Dict[str, int] seen_request_set: set[str] - zmscore: Callable[[Union[redis.Redis, AsyncRedis], str, List[str]], Any] + zmscore: Callable[ + [Union[redis.Redis, AsyncRedis], str, List[str]], + Union[Awaitable[List[Optional[float]]], List[Optional[float]]], + ] async def prefetch(self, keys: List[str]) -> None: if not keys: @@ -83,9 +86,11 @@ async def prefetch(self, keys: List[str]) -> None: if not unique: return - scores = self.zmscore(self.redis_client, self.redis_state_key, unique) - if inspect.iscoroutine(scores): - scores = await scores + scores_result = self.zmscore(self.redis_client, self.redis_state_key, unique) + if inspect.iscoroutine(scores_result): + scores = await cast(Awaitable[List[Optional[float]]], scores_result) + else: + scores = cast(List[Optional[float]], scores_result) for k, s in zip(unique, scores): self.redis_seen_cache[k] = None if s is None else int(s) @@ -847,8 +852,8 @@ async def get_data( dedup_active = bool(params.pop("_sf_dedup_active", False)) - items_data: List = [None] * len(self.items) - results: List[Optional[FeedResult]] = [None] * len(self.items) + items_data: List[List[Any]] = [[] for _ in self.items] + results: List[Optional[FeedResult]] = [None for _ in self.items] indexed_items = list(enumerate(self.items)) fetch_order = indexed_items @@ -860,7 +865,9 @@ async def get_data( ) for idx, item in fetch_order: - item_result = await item.data.get_data( + item_result = cast( + FeedResult, + await item.data.get_data( methods_dict=methods_dict, user_id=user_id, limit=limit * item.percentage // 100, @@ -868,17 +875,18 @@ async def get_data( redis_client=redis_client, _sf_dedup_active=dedup_active, **params, + ), ) results[idx] = item_result - for idx, item_result in enumerate(results): - assert item_result is not None - items_data[idx] = item_result.data + for idx, result_item in enumerate(results): + assert result_item is not None + items_data[idx] = result_item.data - if not result.has_next_page and item_result.has_next_page: + if not result.has_next_page and result_item.has_next_page: result.has_next_page = True - result.next_page.data.update(item_result.next_page.data) + result.next_page.data.update(result_item.next_page.data) # Добавляем данные позиции к общему результату процентного мерджера. result.data = await self._merge_items_data(items_data=items_data) @@ -1351,7 +1359,7 @@ def _compute_overfetch_params(self, *, remaining: int, next_after: Any) -> tuple start_after: Optional[int] = int(next_after) if can_overfetch else None return can_overfetch, request_limit, start_after - def _iter_subfeeds(self, feed: BaseFeedConfigModel): + def _iter_subfeeds(self, feed: BaseFeedConfigModel) -> Iterator["SubFeed"]: if isinstance(feed, SubFeed): yield feed return @@ -1412,7 +1420,12 @@ def _make_wrapped_leaf_method( dedup_state: "_DedupState", leaf_priority: int, ) -> Callable: - async def _wrapped_method(user_id: Any, limit: int, next_page: FeedResultNextPageInside, **kw: Any): + async def _wrapped_method( + user_id: Any, + limit: int, + next_page: FeedResultNextPageInside, + **kw: Any, + ) -> FeedResultClient: collected: List[Any] = [] upstream_has_next_page = False @@ -1541,8 +1554,9 @@ async def _redis_zmscore( if not members: return [] - if hasattr(redis_client, "zmscore"): - res = redis_client.zmscore(key, members) # type: ignore[attr-defined] + zmscore_fn = getattr(redis_client, "zmscore", None) + if zmscore_fn is not None: + res = zmscore_fn(key, members) if inspect.iscoroutine(res): res = await res # redis-py returns list[Optional[float]] @@ -1567,9 +1581,10 @@ async def _redis_zadd_and_expire( res = redis_client.zadd(key, mapping={m: float(s) for m, s in member_scores.items()}) if inspect.iscoroutine(res): await res - await redis_client.expire(key, self.state_ttl_seconds) - else: - redis_client.expire(key, self.state_ttl_seconds) + + expire_res = redis_client.expire(key, self.state_ttl_seconds) + if inspect.iscoroutine(expire_res): + await expire_res def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") @@ -1591,7 +1606,8 @@ async def get_data( # Treat an explicit "page 0" (or missing cursor for this merger) as a fresh session. # This allows clients to restart the feed (e.g., full reload) without carrying over seen state. - requested_page = next_page.data.get(self.merger_id).page if self.merger_id in next_page.data else None + entry = next_page.data.get(self.merger_id) + requested_page = entry.page if entry is not None else None is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) if self.state_backend == "redis" and not redis_client: @@ -1792,9 +1808,9 @@ class FeedConfig(BaseModel): # Update Forward Refs def _rebuild_model(model: Any) -> None: if hasattr(model, "model_rebuild"): - model.model_rebuild() # type: ignore[attr-defined] + model.model_rebuild() else: - model.update_forward_refs() # type: ignore[attr-defined] + model.update_forward_refs() _rebuild_model(MergerPositional) From c70c6c4bb24b145e276a1ae114c09793d6475ae3 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Wed, 17 Dec 2025 00:18:21 +0000 Subject: [PATCH 10/41] black fixes. --- smartfeed/schemas.py | 38 +++++++++++------ tests/test_merger_deduplication.py | 35 ++++++++++------ tests/test_parsing_config.py | 2 +- tests/test_redis_live.py | 67 ++++++++++++++++-------------- tests/utils.py | 1 - 5 files changed, 83 insertions(+), 60 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index f818b1b..dc45d0f 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,13 +1,26 @@ import base64 -from dataclasses import dataclass import inspect import json import logging import zlib from abc import ABC, abstractmethod from collections import defaultdict, deque +from dataclasses import dataclass from random import shuffle -from typing import Annotated, Any, Awaitable, Callable, Dict, Iterator, List, Literal, Optional, Union, cast, no_type_check +from typing import ( + Annotated, + Any, + Awaitable, + Callable, + Dict, + Iterator, + List, + Literal, + Optional, + Union, + cast, + no_type_check, +) import redis from pydantic import BaseModel, Field, PrivateAttr, model_validator @@ -108,6 +121,7 @@ def record(self, key: str, priority: int) -> None: self.redis_seen_cache[key] = priority self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) + FeedTypes = Annotated[ Union[ "MergerDeduplication", @@ -868,13 +882,13 @@ async def get_data( item_result = cast( FeedResult, await item.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit * item.percentage // 100, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, + methods_dict=methods_dict, + user_id=user_id, + limit=limit * item.percentage // 100, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, ), ) @@ -1328,9 +1342,7 @@ def _extract_dedup_value(self, item: Any) -> Any: value = getattr(item, self.dedup_key, None) if value is None and self.missing_key_policy == "error": - raise AssertionError( - f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}" - ) + raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}") return value def _get_entity_key(self, entity: Any) -> Optional[str]: @@ -1821,4 +1833,4 @@ def _rebuild_model(model: Any) -> None: _rebuild_model(MergerAppendDistribute) _rebuild_model(MergerPercentageGradient) _rebuild_model(MergerViewSession) -_rebuild_model(MergerDeduplication) \ No newline at end of file +_rebuild_model(MergerDeduplication) diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 94d122a..9bda09f 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -2,13 +2,7 @@ import pytest -from smartfeed.schemas import ( - FeedResultClient, - FeedResultNextPage, - FeedResultNextPageInside, - MergerDeduplication, -) - +from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication from tests.fixtures.redis import redis_client # noqa: F401 from tests.utils import parse_model @@ -1009,7 +1003,9 @@ async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_ref } merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) + res_1 = await merger.get_data( + methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={}) + ) res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) assert len(res_1.data) == 10 @@ -1026,7 +1022,7 @@ async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_ref @pytest.mark.asyncio async def test_dedup_percentage_gradient_cursor_backend_across_pages() -> None: a_items = [{"id": i, "src": "A"} for i in range(1, 300)] - b_items = ([{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)]) + b_items = [{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)] methods_dict = { "a": make_offset_paged_method(a_items), @@ -1052,7 +1048,9 @@ async def test_dedup_percentage_gradient_cursor_backend_across_pages() -> None: } merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) + res_1 = await merger.get_data( + methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={}) + ) res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) _assert_no_dupes_in_page(res_1.data) @@ -1148,8 +1146,19 @@ async def test_dedup_wrapper_with_view_session_merger(redis_client) -> None: "type": "merger_percentage", "shuffle": False, "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}}, - {"percentage": 50, "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high", "dedup_priority": 100}}, + { + "percentage": 50, + "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}, + }, + { + "percentage": 50, + "data": { + "subfeed_id": "sf_high", + "type": "subfeed", + "method_name": "high", + "dedup_priority": 100, + }, + }, ], }, }, @@ -1290,7 +1299,7 @@ async def test_dedup_percentage_gradient_slot_ownership_cursor_backend() -> None a_items = [{"id": i, "src": "A"} for i in range(1, 300)] # Start with duplicates, then provide unique tail. - b_items = ([{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)]) + b_items = [{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)] methods_dict = { "a": make_offset_paged_method(a_items), diff --git a/tests/test_parsing_config.py b/tests/test_parsing_config.py index 66c2387..bf61007 100644 --- a/tests/test_parsing_config.py +++ b/tests/test_parsing_config.py @@ -2,9 +2,9 @@ from smartfeed.manager import FeedManager from smartfeed.schemas import ( - MergerDeduplication, FeedConfig, MergerAppend, + MergerDeduplication, MergerPercentage, MergerPercentageGradient, MergerPercentageItem, diff --git a/tests/test_redis_live.py b/tests/test_redis_live.py index 85effa2..e1dc13d 100644 --- a/tests/test_redis_live.py +++ b/tests/test_redis_live.py @@ -1,9 +1,10 @@ import asyncio import json +import time + import pytest import redis from redis.asyncio import Redis as AsyncRedis -import time from smartfeed.schemas import FeedResultNextPage, MergerViewSession from tests.fixtures.configs import METHODS_DICT @@ -15,21 +16,22 @@ class RedisReplicationSimulator: """ Симулятор задержки репликации Redis для тестирования проблемы кластера. """ + def __init__(self, real_client): self.real_client = real_client self.write_delay = 0.1 # Задержка для имитации репликации self.pending_writes = {} # Ключи которые только что записали - + def exists(self, cache_key): return self.real_client.exists(cache_key) - + def set(self, name, value, ex=None): # Записываем в реальный Redis result = self.real_client.set(name, value, ex=ex) # Помечаем что этот ключ только что записан (имитация репликации) self.pending_writes[name] = time.time() return result - + def get(self, name): # Если ключ только что записан (в течение write_delay секунд), возвращаем None if name in self.pending_writes: @@ -39,7 +41,7 @@ def get(self, name): else: # Задержка прошла, можно удалить из pending del self.pending_writes[name] - + # Обычное чтение из Redis return self.real_client.get(name) @@ -50,24 +52,24 @@ async def test_redis_replication_delay_problem(): Тест для воспроизведения проблемы репликации Redis с использованием RedisReplicationSimulator для имитации задержки. """ - + # Подключаемся к Redis (должен быть запущен локально) try: - real_client = redis.Redis(host='localhost', port=6379, db=0) + real_client = redis.Redis(host="localhost", port=6379, db=0) real_client.ping() # Проверяем соединение except (redis.ConnectionError, redis.ResponseError): pytest.skip("Redis not available for live testing") - + # Очищаем тестовый ключ test_key = "test_merger_view_session_test_user" real_client.delete(test_key) - + # Используем симулятор задержки репликации redis_client = RedisReplicationSimulator(real_client) merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - + print("\n=== Демонстрация проблемы с задержкой репликации ===") - + try: # Этот вызов должен воспроизвести проблему с оригинальным кодом result = await merger_vs.get_data( @@ -77,17 +79,17 @@ async def test_redis_replication_delay_problem(): user_id="test_user", redis_client=redis_client, ) - + print("✅ Исправление работает! Получили результат без ошибки:") print(f" Данные: {result.data[:5]}... (показаны первые 5)") print(f" Размер: {len(result.data)}") print(f" Есть следующая страница: {result.has_next_page}") - + # Проверяем что получили валидные данные assert len(result.data) == 10 assert result.data[0] == "test_user_1" assert result.has_next_page is True - + except TypeError as e: if "the JSON object must be str, bytes or bytearray, not NoneType" in str(e): print("❌ Проблема НЕ исправлена! Все еще получаем TypeError") @@ -95,33 +97,33 @@ async def test_redis_replication_delay_problem(): else: print(f"❓ Неожиданная ошибка: {e}") raise - + finally: # Очистка real_client.delete(test_key) real_client.close() -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_redis_multiple_requests(): """ Тест множественных запросов для проверки стабильности исправления. """ - + try: - real_client = redis.Redis(host='localhost', port=6379, db=0) + real_client = redis.Redis(host="localhost", port=6379, db=0) real_client.ping() except (redis.ConnectionError, redis.ResponseError): pytest.skip("Redis not available for live testing") - + test_key = "test_merger_multiple_test_user" real_client.delete(test_key) - + redis_client = RedisReplicationSimulator(real_client) merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - + print("\n=== Тест множественных запросов ===") - + try: # Первый запрос - создает кэш result1 = await merger_vs.get_data( @@ -131,33 +133,34 @@ async def test_redis_multiple_requests(): user_id="test_user", redis_client=redis_client, ) - + print(f"Первый запрос: получили {len(result1.data)} элементов") - + # Ждем чтобы задержка репликации прошла await asyncio.sleep(0.2) - - # Второй запрос - должен использовать кэш + + # Второй запрос - должен использовать кэш from smartfeed.schemas import FeedResultNextPageInside + result2 = await merger_vs.get_data( methods_dict=METHODS_DICT, - limit=5, + limit=5, next_page=FeedResultNextPage( data={"merger_view_session_example": FeedResultNextPageInside(page=2, after=None)} ), user_id="test_user", redis_client=redis_client, ) - + print(f"Второй запрос: получили {len(result2.data)} элементов") print(f"Данные второй страницы: {result2.data}") - + # Проверяем что получили разные данные (пагинация работает) assert result1.data != result2.data assert len(result2.data) == 5 - + print("✅ Множественные запросы работают корректно!") - + finally: real_client.delete(test_key) real_client.close() @@ -165,4 +168,4 @@ async def test_redis_multiple_requests(): if __name__ == "__main__": # Для запуска напрямую без pytest - asyncio.run(test_redis_replication_delay_problem()) \ No newline at end of file + asyncio.run(test_redis_replication_delay_problem()) diff --git a/tests/utils.py b/tests/utils.py index f38689f..331f1e7 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,7 +2,6 @@ from typing import Any, Dict, Type, TypeVar - T = TypeVar("T") From 43b06c7705b56e01314f4cd743af07666e19c3e1 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 5 Feb 2026 16:29:59 +0000 Subject: [PATCH 11/41] Minor redis fix. --- smartfeed/schemas.py | 53 +++++++++++++++++++++++++++++++------------- 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index dc45d0f..ca8fff6 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,3 +1,4 @@ +import asyncio import base64 import inspect import json @@ -28,6 +29,23 @@ from redis.asyncio import RedisCluster as AsyncRedisCluster +def _is_async_redis_client(client: Any) -> bool: + return isinstance(client, (AsyncRedis, AsyncRedisCluster)) + + +async def _redis_call(client: Any, method_name: str, *args: Any, **kwargs: Any) -> Any: + """Call a Redis method without blocking the event loop. + + - For `redis.asyncio` clients, calls are awaited directly. + - For sync `redis.Redis`, calls are offloaded via `asyncio.to_thread()`. + """ + + method = getattr(client, method_name) + if _is_async_redis_client(client): + return await method(*args, **kwargs) + return await asyncio.to_thread(method, *args, **kwargs) + + def _pydantic_deep_copy(model: Any) -> Any: """Deep copy helper compatible with Pydantic v1 and v2.""" @@ -287,7 +305,7 @@ async def _set_cache( self, methods_dict: Dict[str, Callable], user_id: Any, - redis_client: redis.Redis, + redis_client: Union[redis.Redis, AsyncRedis], cache_key: str, **params: Any, ) -> List[Any]: @@ -313,7 +331,7 @@ async def _set_cache( data = result.data if self.deduplicate: data = self._dedup_data(data) - redis_client.set(name=cache_key, value=json.dumps(data), ex=self.session_live_time) + await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) return data async def _set_cache_async( @@ -356,7 +374,7 @@ async def _get_cache( user_id: Any, limit: int, next_page: FeedResultNextPage, - redis_client: redis.Redis, + redis_client: Union[redis.Redis, AsyncRedis], **params: Any, ) -> FeedResult: """ @@ -380,7 +398,8 @@ async def _get_cache( logging.info("MergerViewSession cache request for %s", cache_key) # Если кэш не найден или передан пустой курсор пагинации на мерджер, обновляем данные и записываем в кэш. - if not redis_client.exists(cache_key) or self.merger_id not in next_page.data: + cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) + if not cache_exists or self.merger_id not in next_page.data: logging.info("Cache miss or new session - generating fresh data for %s", cache_key) # Получаем свежие данные и используем их напрямую (избегаем чтение из кэша) session_data = await self._set_cache( @@ -389,7 +408,7 @@ async def _get_cache( else: logging.info("Cache exists - attempting read from Redis for %s", cache_key) # Читаем из кэша только если он уже существовал - cached_data = redis_client.get(name=cache_key) + cached_data = await _redis_call(redis_client, "get", cache_key) if cached_data is None: # Fallback: если кэш пропал, получаем свежие данные logging.info( @@ -1574,6 +1593,17 @@ async def _redis_zmscore( # redis-py returns list[Optional[float]] return [None if v is None else float(v) for v in list(res)] + # Fallback: pipelined zscore. For sync clients, run the whole pipeline in a thread. + if not _is_async_redis_client(redis_client): + def _sync_pipeline_execute() -> Any: + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + return pipe.execute() + + res = await asyncio.to_thread(_sync_pipeline_execute) + return [None if v is None else float(v) for v in list(res)] + pipe = redis_client.pipeline() for m in members: pipe.zscore(key, m) @@ -1590,13 +1620,8 @@ async def _redis_zadd_and_expire( ) -> None: if not member_scores: return - res = redis_client.zadd(key, mapping={m: float(s) for m, s in member_scores.items()}) - if inspect.iscoroutine(res): - await res - - expire_res = redis_client.expire(key, self.state_ttl_seconds) - if inspect.iscoroutine(expire_res): - await expire_res + await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) + await _redis_call(redis_client, "expire", key, self.state_ttl_seconds) def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") @@ -1647,9 +1672,7 @@ async def get_data( redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) if is_fresh_session: # Drop state for a full restart. - deleted = redis_client.delete(redis_state_key) - if inspect.iscoroutine(deleted): - await deleted + await _redis_call(redis_client, "delete", redis_state_key) # Create a single state helper shared across all leaf wrappers. if self.state_backend == "cursor": From 0f11803d121b4d54979b94782d0e7f7743340e0a Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 5 Feb 2026 21:48:53 +0000 Subject: [PATCH 12/41] Minor parrallelism fix. --- smartfeed/schemas.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index ca8fff6..939a504 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -582,19 +582,35 @@ async def get_data( if dedup_active: indexed_items = list(enumerate(self.items)) - fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) fetched: Dict[int, FeedResult] = {} - for idx, item in fetch_order: - fetched[idx] = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) + # Always-on parallelism in dedup mode: preserve dedup semantics by ensuring + # higher priority is recorded first; fetch same-priority children concurrently. + groups: Dict[int, List[tuple[int, FeedTypes]]] = defaultdict(list) + for idx, item in indexed_items: + prio = int(getattr(item, "dedup_priority", 0)) + groups[prio].append((idx, item)) + + for prio in sorted(groups.keys(), reverse=True): + group = groups[prio] + coros: List[Awaitable[FeedResult]] = [] + order: List[int] = [] + for idx, item in group: + order.append(idx) + coros.append( + item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=_pydantic_deep_copy(next_page), + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + ) + group_results = await asyncio.gather(*coros) + for idx, r in zip(order, group_results): + fetched[idx] = cast(FeedResult, r) for idx, _item in indexed_items: item_result = fetched[idx] From 2086519a1461e55d23d03d2f37f8759b5650d83e Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 5 Feb 2026 21:54:30 +0000 Subject: [PATCH 13/41] Switched to orjson. --- pyproject.toml | 1 + smartfeed/examples/example_client.py | 2 +- smartfeed/jsonlib.py | 41 ++++++++++++++++++++++++++++ smartfeed/schemas.py | 2 +- 4 files changed, 44 insertions(+), 2 deletions(-) create mode 100644 smartfeed/jsonlib.py diff --git a/pyproject.toml b/pyproject.toml index 108a4c7..2c4d669 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ packages = [ python = ">=3.9" pydantic = ">=1.10.7" redis = ">=4.5.5" +orjson = ">=3.9.0" [tool.poetry.group.dev.dependencies] isort = "^5.12.0" diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index b1a55e3..27ca2fa 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -1,5 +1,5 @@ import base64 -import json +from smartfeed import jsonlib as json from typing import Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_validator diff --git a/smartfeed/jsonlib.py b/smartfeed/jsonlib.py new file mode 100644 index 0000000..401e6b4 --- /dev/null +++ b/smartfeed/jsonlib.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Any, Callable, Optional + +import orjson + + +DefaultFn = Callable[[Any], Any] + + +def dumps( + obj: Any, + *, + default: Optional[DefaultFn] = None, + sort_keys: bool = False, +) -> str: + """Serialize *obj* to JSON text using orjson. + + This is a small compatibility layer meant to cover the subset of the stdlib + `json.dumps` API used inside this package. + + Key differences vs `orjson.dumps`: + - Returns `str` (UTF-8) instead of `bytes`. + - Supports `default=` and `sort_keys=`. + """ + + option = 0 + if sort_keys: + option |= orjson.OPT_SORT_KEYS + + return orjson.dumps(obj, default=default, option=option).decode("utf-8") + + +def loads(data: Any) -> Any: + """Deserialize JSON from *data* using orjson. + + Accepts `str`, `bytes`, `bytearray`, or `memoryview` (same spirit as + stdlib `json.loads`). + """ + + return orjson.loads(data) diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 939a504..ed04203 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,7 +1,7 @@ import asyncio import base64 import inspect -import json +from . import jsonlib as json import logging import zlib from abc import ABC, abstractmethod From 795c5748f9bb6ec1fac3934b906a3d9e542e7977 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 5 Feb 2026 22:34:39 +0000 Subject: [PATCH 14/41] Refactor into separate merger modules. --- smartfeed/examples/example_client.py | 2 +- smartfeed/feed_models.py | 161 ++ smartfeed/jsonlib.py | 1 - smartfeed/mergers/__init__.py | 24 + smartfeed/mergers/append.py | 103 ++ smartfeed/mergers/append_distribute.py | 111 ++ smartfeed/mergers/deduplication.py | 548 ++++++ smartfeed/mergers/percentage.py | 111 ++ smartfeed/mergers/percentage_gradient.py | 164 ++ smartfeed/mergers/positional.py | 135 ++ smartfeed/mergers/view_session.py | 205 +++ smartfeed/schemas.py | 1926 +--------------------- 12 files changed, 1636 insertions(+), 1855 deletions(-) create mode 100644 smartfeed/feed_models.py create mode 100644 smartfeed/mergers/__init__.py create mode 100644 smartfeed/mergers/append.py create mode 100644 smartfeed/mergers/append_distribute.py create mode 100644 smartfeed/mergers/deduplication.py create mode 100644 smartfeed/mergers/percentage.py create mode 100644 smartfeed/mergers/percentage_gradient.py create mode 100644 smartfeed/mergers/positional.py create mode 100644 smartfeed/mergers/view_session.py diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index 27ca2fa..9b48842 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -1,9 +1,9 @@ import base64 -from smartfeed import jsonlib as json from typing import Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_validator +from smartfeed import jsonlib as json from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py new file mode 100644 index 0000000..b98f11f --- /dev/null +++ b/smartfeed/feed_models.py @@ -0,0 +1,161 @@ +import asyncio +import inspect +from abc import ABC, abstractmethod +from dataclasses import dataclass +from random import shuffle +from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast + +import redis +from pydantic import BaseModel +from redis.asyncio import Redis as AsyncRedis +from redis.asyncio import RedisCluster as AsyncRedisCluster + + +def _is_async_redis_client(client: Any) -> bool: + return isinstance(client, (AsyncRedis, AsyncRedisCluster)) + + +async def _redis_call(client: Any, method_name: str, *args: Any, **kwargs: Any) -> Any: + """Call a Redis method without blocking the event loop. + + - For `redis.asyncio` clients, calls are awaited directly. + - For sync `redis.Redis`, calls are offloaded via `asyncio.to_thread()`. + """ + + method = getattr(client, method_name) + if _is_async_redis_client(client): + return await method(*args, **kwargs) + return await asyncio.to_thread(method, *args, **kwargs) + + +def _pydantic_deep_copy(model: Any) -> Any: + """Deep copy helper compatible with Pydantic v1 and v2.""" + + if hasattr(model, "model_copy"): + return model.model_copy(deep=True) + return model.copy(deep=True) + + +class FeedResultNextPageInside(BaseModel): + """Cursor model for one feed node.""" + + page: int = 1 + after: Any = None + + +class FeedResultNextPage(BaseModel): + """Cursor model for a whole feed traversal.""" + + data: Dict[str, FeedResultNextPageInside] + + +class FeedResult(BaseModel): + """Normalized output of any feed node `get_data()`.""" + + data: List + next_page: FeedResultNextPage + has_next_page: bool + + +class FeedResultClient(BaseModel): + """Result returned by client subfeed methods.""" + + data: List + next_page: FeedResultNextPageInside + has_next_page: bool + + +class BaseFeedConfigModel(ABC, BaseModel): + """Base class for merger/subfeed config models.""" + + # Higher value means the item should "win" deduplication when duplicates exist. + # This is primarily used by MergerDeduplication and by mergers when a dedup wrapper is active. + dedup_priority: int = 0 + + @abstractmethod + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + """Fetch data according to this node config.""" + + +@dataclass +class _SubFeedMethodSpec: + method: Callable + args: List[str] + + +class SubFeed(BaseFeedConfigModel): + """Leaf node pointing at a client method.""" + + subfeed_id: str + type: Literal["subfeed"] + method_name: str + subfeed_params: Dict[str, Any] = {} + raise_error: Optional[bool] = True + shuffle: bool = False + + def _get_method_spec(self, methods_dict: Dict[str, Callable]) -> _SubFeedMethodSpec: + method = methods_dict[self.method_name] + method_spec = getattr(method, "_smartfeed_original", method) + method_args = inspect.getfullargspec(method_spec).args + return _SubFeedMethodSpec(method=method, args=list(method_args)) + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + subfeed_next_page = FeedResultNextPageInside( + page=next_page.data[self.subfeed_id].page if self.subfeed_id in next_page.data else 1, + after=next_page.data[self.subfeed_id].after if self.subfeed_id in next_page.data else None, + ) + + method_spec = self._get_method_spec(methods_dict) + + method_params: Dict[str, Any] = {} + for arg in method_spec.args: + if arg in params: + method_params[arg] = params[arg] + + try: + method_result = await method_spec.method( + user_id=user_id, + limit=limit, + next_page=subfeed_next_page, + **method_params, + **self.subfeed_params, + ) + except Exception: + if self.raise_error: + raise + + method_result = FeedResultClient( + data=[], + next_page=subfeed_next_page, + has_next_page=False, + ) + + if not isinstance(method_result, FeedResultClient): + raise TypeError('SubFeed function must return "FeedResultClient" instance.') + + if self.shuffle: + shuffle(method_result.data) + + return FeedResult( + data=method_result.data, + next_page=FeedResultNextPage( + data={self.subfeed_id: cast(FeedResultNextPageInside, method_result.next_page)} + ), + has_next_page=bool(method_result.has_next_page), + ) diff --git a/smartfeed/jsonlib.py b/smartfeed/jsonlib.py index 401e6b4..9ab650f 100644 --- a/smartfeed/jsonlib.py +++ b/smartfeed/jsonlib.py @@ -4,7 +4,6 @@ import orjson - DefaultFn = Callable[[Any], Any] diff --git a/smartfeed/mergers/__init__.py b/smartfeed/mergers/__init__.py new file mode 100644 index 0000000..24b1c8e --- /dev/null +++ b/smartfeed/mergers/__init__.py @@ -0,0 +1,24 @@ +"""Merger implementations. + +Each merger schema lives in its own module. +`smartfeed.schemas` re-exports these classes for backwards compatibility. +""" + +from .append import MergerAppend +from .append_distribute import MergerAppendDistribute +from .deduplication import MergerDeduplication +from .percentage import MergerPercentage, MergerPercentageItem +from .percentage_gradient import MergerPercentageGradient +from .positional import MergerPositional +from .view_session import MergerViewSession + +__all__ = [ + "MergerAppend", + "MergerAppendDistribute", + "MergerDeduplication", + "MergerPercentage", + "MergerPercentageItem", + "MergerPercentageGradient", + "MergerPositional", + "MergerViewSession", +] diff --git a/smartfeed/mergers/append.py b/smartfeed/mergers/append.py new file mode 100644 index 0000000..bfe0718 --- /dev/null +++ b/smartfeed/mergers/append.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict +from random import shuffle +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast + +import redis +from redis.asyncio import Redis as AsyncRedis + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, _pydantic_deep_copy + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class MergerAppend(BaseFeedConfigModel): + """Append merger.""" + + merger_id: str + type: Literal["merger_append"] + items: List[FeedTypes] + shuffle: bool = False + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) + + if dedup_active: + indexed_items = list(enumerate(self.items)) + fetched: Dict[int, FeedResult] = {} + + groups: Dict[int, List[tuple[int, "FeedTypes"]]] = defaultdict(list) + for idx, item in indexed_items: + prio = int(getattr(item, "dedup_priority", 0)) + groups[prio].append((idx, item)) + + for prio in sorted(groups.keys(), reverse=True): + group = groups[prio] + coros: List[Awaitable[FeedResult]] = [] + order: List[int] = [] + for idx, item in group: + order.append(idx) + coros.append( + item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=_pydantic_deep_copy(next_page), + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + ) + group_results = await asyncio.gather(*coros) + for idx, r in zip(order, group_results): + fetched[idx] = cast(FeedResult, r) + + for idx, _item in indexed_items: + item_result = fetched[idx] + result.data.extend(item_result.data) + result.next_page.data.update(item_result.next_page.data) + if item_result.has_next_page: + result.has_next_page = True + + if len(result.data) > limit: + result.data = result.data[:limit] + else: + result_limit = limit + for item in self.items: + item_result = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=result_limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) + + result.data.extend(item_result.data) + result_limit -= len(item_result.data) + + if not result.has_next_page and item_result.has_next_page: + result.has_next_page = True + + result.next_page.data.update(item_result.next_page.data) + + if result_limit <= 0: + break + + if self.shuffle: + shuffle(result.data) + + return result diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py new file mode 100644 index 0000000..5f07d66 --- /dev/null +++ b/smartfeed/mergers/append_distribute.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from collections import defaultdict, deque +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union + +import redis +from redis.asyncio import Redis as AsyncRedis +from typing_extensions import no_type_check + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class MergerAppendDistribute(BaseFeedConfigModel): + """Merger that uniformly distributes items by a key.""" + + merger_id: str + type: Literal["merger_distribute"] + items: List["FeedTypes"] + distribution_key: str + sorting_key: Optional[str] = None + sorting_desc: bool = False + + @no_type_check + async def _uniform_distribute(self, data: list) -> list: + if self.sorting_key: + data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) + + grouped_entries = defaultdict(deque) + for entry in data: + grouped_entries[entry[self.distribution_key]].append(entry) + result = [] + prev_profile_id = None + while any(grouped_entries.values()): + for profile_id in list(grouped_entries.keys()): + if grouped_entries[profile_id]: + if profile_id != prev_profile_id or len(grouped_entries) == 1: + result.append(grouped_entries[profile_id].popleft()) + prev_profile_id = profile_id + if not grouped_entries[profile_id]: + del grouped_entries[profile_id] + else: + del grouped_entries[profile_id] + + return result + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) + + if dedup_active: + indexed_items = list(enumerate(self.items)) + fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) + fetched: Dict[int, FeedResult] = {} + + for idx, item in fetch_order: + fetched[idx] = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + + for idx, _item in indexed_items: + item_result = fetched[idx] + result.data.extend(item_result.data) + result.next_page.data.update(item_result.next_page.data) + if item_result.has_next_page: + result.has_next_page = True + + if len(result.data) > limit: + result.data = result.data[:limit] + else: + result_limit = limit + for item in self.items: + item_result = await item.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=result_limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) + + result.data.extend(item_result.data) + result_limit -= len(item_result.data) + + if not result.has_next_page and item_result.has_next_page: + result.has_next_page = True + + result.next_page.data.update(item_result.next_page.data) + + if result_limit <= 0: + break + + result.data = await self._uniform_distribute(result.data) + return result diff --git a/smartfeed/mergers/deduplication.py b/smartfeed/mergers/deduplication.py new file mode 100644 index 0000000..a9902a5 --- /dev/null +++ b/smartfeed/mergers/deduplication.py @@ -0,0 +1,548 @@ +from __future__ import annotations + +import asyncio +import base64 +import inspect +import zlib +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterator, List, Literal, Optional, Union, cast + +import redis +from pydantic import PrivateAttr, model_validator +from redis.asyncio import Redis as AsyncRedis + +from .. import jsonlib as json +from ..feed_models import ( + BaseFeedConfigModel, + FeedResult, + FeedResultClient, + FeedResultNextPage, + FeedResultNextPageInside, + SubFeed, + _is_async_redis_client, + _pydantic_deep_copy, + _redis_call, +) + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class _DedupState(ABC): + @abstractmethod + def should_accept(self, key: str, priority: int) -> bool: + raise NotImplementedError + + @abstractmethod + def record(self, key: str, priority: int) -> None: + raise NotImplementedError + + async def prefetch(self, keys: List[str]) -> None: + return + + +@dataclass +class _CursorDedupState(_DedupState): + seen_priority_map: Dict[str, int] + seen_updates_in_order: List[tuple[str, int]] + seen_request_set: set[str] + + def should_accept(self, key: str, priority: int) -> bool: + if key in self.seen_request_set: + return False + existing_priority = self.seen_priority_map.get(key) + if existing_priority is not None and priority <= existing_priority: + return False + return True + + def record(self, key: str, priority: int) -> None: + self.seen_priority_map[key] = priority + self.seen_updates_in_order.append((key, priority)) + self.seen_request_set.add(key) + + +@dataclass +class _RedisDedupState(_DedupState): + redis_client: Union[redis.Redis, AsyncRedis] + redis_state_key: str + redis_seen_cache: Dict[str, Optional[int]] + redis_new_scores: Dict[str, int] + seen_request_set: set[str] + zmscore: Callable[ + [Union[redis.Redis, AsyncRedis], str, List[str]], + Union[Awaitable[List[Optional[float]]], List[Optional[float]]], + ] + + async def prefetch(self, keys: List[str]) -> None: + if not keys: + return + unique: List[str] = [] + seen: set[str] = set() + for k in keys: + if k in self.seen_request_set: + continue + if k in self.redis_seen_cache: + continue + if k in seen: + continue + seen.add(k) + unique.append(k) + + if not unique: + return + + scores_result = self.zmscore(self.redis_client, self.redis_state_key, unique) + if inspect.iscoroutine(scores_result): + scores = await cast(Awaitable[List[Optional[float]]], scores_result) + else: + scores = cast(List[Optional[float]], scores_result) + + for k, s in zip(unique, scores): + self.redis_seen_cache[k] = None if s is None else int(s) + + def should_accept(self, key: str, priority: int) -> bool: + if key in self.seen_request_set: + return False + existing_priority = self.redis_seen_cache.get(key) + if existing_priority is not None and priority <= existing_priority: + return False + return True + + def record(self, key: str, priority: int) -> None: + self.seen_request_set.add(key) + self.redis_seen_cache[key] = priority + self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) + + +class MergerDeduplication(BaseFeedConfigModel): + """Merger that deduplicates while preserving child mixing/position semantics.""" + + merger_id: str + type: Literal["merger_deduplication"] + data: "FeedTypes" + + dedup_key: Optional[str] = None + missing_key_policy: Literal["error", "keep", "drop"] = "error" + + state_backend: Literal["cursor", "redis"] = "cursor" + state_ttl_seconds: int = 3600 + cursor_compress: bool = True + cursor_max_keys: Optional[int] = None + + overfetch_factor: int = 1 + + max_refill_loops: int = 20 + + _descendant_cursor_keys_cache: Optional[set[str]] = PrivateAttr(default=None) + + @model_validator(mode="after") + def validate_merger_deduplication(self) -> "MergerDeduplication": + if self.overfetch_factor < 1: + raise ValueError('"overfetch_factor" must be >= 1') + if self.max_refill_loops < 1: + raise ValueError('"max_refill_loops" must be >= 1') + return self + + def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: + keys: set[str] = set() + + subfeed_id = getattr(feed, "subfeed_id", None) + if isinstance(subfeed_id, str) and subfeed_id: + keys.add(subfeed_id) + + merger_id = getattr(feed, "merger_id", None) + if isinstance(merger_id, str) and merger_id: + keys.add(merger_id) + + child: Any + for attr_name in ("data", "positional", "default"): + child = getattr(feed, attr_name, None) + if isinstance(child, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(child)) + + for attr_name in ("item_from", "item_to"): + child = getattr(feed, attr_name, None) + inner = getattr(child, "data", None) + if isinstance(inner, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(inner)) + + items = getattr(feed, "items", None) + if isinstance(items, list): + for item in items: + if isinstance(item, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(item)) + continue + + inner = getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + keys.update(self._collect_descendant_cursor_keys(inner)) + + return keys + + def _get_descendant_cursor_keys_cached(self) -> set[str]: + cached = self._descendant_cursor_keys_cache + if cached is None: + cached = self._collect_descendant_cursor_keys(self.data) + self._descendant_cursor_keys_cache = cached + return cached + + def _reset_descendant_cursors(self, next_page: FeedResultNextPage) -> None: + descendant_keys = self._get_descendant_cursor_keys_cached() + for key in descendant_keys: + next_page.data.pop(key, None) + + def _normalize_key(self, value: Any) -> str: + if isinstance(value, (str, int)): + return str(value) + if isinstance(value, (dict, list)): + return json.dumps(value, sort_keys=True, default=str) + return str(value) + + def _extract_dedup_value(self, item: Any) -> Any: + if not self.dedup_key: + return item + + try: + value = item.get(self.dedup_key) + except AttributeError: + value = getattr(item, self.dedup_key, None) + + if value is None and self.missing_key_policy == "error": + raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}") + return value + + def _get_entity_key(self, entity: Any) -> Optional[str]: + raw_value = self._extract_dedup_value(entity) + if raw_value is None: + if self.missing_key_policy == "drop": + return None + if self.missing_key_policy == "keep": + raw_value = ("__missing__", id(entity)) + return self._normalize_key(raw_value) + + def _compute_overfetch_params(self, *, remaining: int, next_after: Any) -> tuple[bool, int, Optional[int]]: + can_overfetch = isinstance(next_after, int) + request_limit = max(1, remaining) + if can_overfetch and self.overfetch_factor > 1: + request_limit = max(1, remaining * self.overfetch_factor) + start_after: Optional[int] = int(next_after) if can_overfetch else None + return can_overfetch, request_limit, start_after + + def _iter_subfeeds(self, feed: BaseFeedConfigModel) -> Iterator[SubFeed]: + if isinstance(feed, SubFeed): + yield feed + return + + for attr_name in ("data", "positional", "default"): + inner = getattr(feed, attr_name, None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + for attr_name in ("item_from", "item_to"): + wrapper = getattr(feed, attr_name, None) + inner = getattr(wrapper, "data", None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + items = getattr(feed, "items", None) + if isinstance(items, list): + for item in items: + if isinstance(item, BaseFeedConfigModel): + yield from self._iter_subfeeds(item) + continue + inner = getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + yield from self._iter_subfeeds(inner) + + def _register_wrapped_subfeed_method( + self, + *, + subfeed: SubFeed, + original_methods_dict: Dict[str, Callable], + rewritten_methods_dict: Dict[str, Callable], + dedup_state: _DedupState, + ) -> None: + original_name = subfeed.method_name + original_method = original_methods_dict[original_name] + unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" + + if unique_name in rewritten_methods_dict: + subfeed.method_name = unique_name + return + + subfeed.method_name = unique_name + leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) + + wrapped = self._make_wrapped_leaf_method( + original_method=original_method, + dedup_state=dedup_state, + leaf_priority=leaf_priority, + ) + setattr(wrapped, "_smartfeed_original", original_method) + rewritten_methods_dict[unique_name] = wrapped + + def _make_wrapped_leaf_method( + self, + *, + original_method: Callable, + dedup_state: _DedupState, + leaf_priority: int, + ) -> Callable: + async def _wrapped_method( + user_id: Any, + limit: int, + next_page: FeedResultNextPageInside, + **kw: Any, + ) -> FeedResultClient: + collected: List[Any] = [] + upstream_has_next_page = False + + loops = 0 + while len(collected) < limit and loops < self.max_refill_loops: + loops += 1 + before_len = len(collected) + + remaining = limit - len(collected) + can_overfetch, request_limit, start_after = self._compute_overfetch_params( + remaining=remaining, + next_after=next_page.after, + ) + + method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) + if not isinstance(method_result, FeedResultClient): + raise TypeError('SubFeed function must return "FeedResultClient" instance.') + + upstream_has_next_page = upstream_has_next_page or method_result.has_next_page + + inspected_count = 0 + + keys_by_index: Optional[List[Optional[str]]] = None + if isinstance(dedup_state, _RedisDedupState): + keys_by_index = [] + batch_keys: List[str] = [] + for entity in method_result.data: + key = self._get_entity_key(entity) + keys_by_index.append(key) + if key is not None: + batch_keys.append(key) + await dedup_state.prefetch(batch_keys) + + for idx, entity in enumerate(method_result.data, start=1): + inspected_count = idx + + key = keys_by_index[idx - 1] if keys_by_index is not None else self._get_entity_key(entity) + if key is None: + continue + + if not dedup_state.should_accept(key, leaf_priority): + continue + + collected.append(entity) + dedup_state.record(key, leaf_priority) + + if len(collected) >= limit: + break + + if len(collected) == before_len: + if not method_result.has_next_page: + break + + if can_overfetch and request_limit > remaining and start_after is not None: + end_after = next_page.after + if isinstance(end_after, int) and end_after == start_after + len(method_result.data): + next_page.after = start_after + inspected_count + + return FeedResultClient(data=collected, next_page=next_page, has_next_page=upstream_has_next_page) + + return _wrapped_method + + def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> Dict[str, int]: + entry = next_page.data.get(self.merger_id) + if not entry or entry.after is None: + return {} + + after = entry.after + if isinstance(after, dict) and "z" in after: + payload = base64.urlsafe_b64decode(after["z"].encode()) + raw = zlib.decompress(payload).decode() + decoded = json.loads(raw) + if isinstance(decoded, dict): + return {str(k): int(v) for k, v in decoded.items()} + if isinstance(decoded, list): + seen_map: Dict[str, int] = {} + for entry_item in decoded: + if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: + seen_map[str(entry_item[0])] = int(entry_item[1]) + else: + seen_map[str(entry_item)] = 0 + return seen_map + return {} + if isinstance(after, dict) and "seen" in after: + return {str(k): 0 for k in list(after["seen"])} + if isinstance(after, list): + return {str(k): 0 for k in list(after)} + if isinstance(after, dict): + return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} + return {} + + def _encode_seen_for_cursor(self, seen_updates_in_order: List[tuple[str, int]]) -> Any: + if self.cursor_max_keys is not None: + seen_updates_in_order = seen_updates_in_order[-self.cursor_max_keys :] + + if not self.cursor_compress: + return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} + + raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() + compressed = zlib.compress(raw) + return { + "v": 2, + "c": "zlib+base64", + "n": len(seen_updates_in_order), + "z": base64.urlsafe_b64encode(compressed).decode(), + } + + async def _redis_zmscore( + self, + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + members: List[str], + ) -> List[Optional[float]]: + if not members: + return [] + + zmscore_fn = getattr(redis_client, "zmscore", None) + if zmscore_fn is not None: + res = zmscore_fn(key, members) + if inspect.iscoroutine(res): + res = await res + return [None if v is None else float(v) for v in list(res)] + + if not _is_async_redis_client(redis_client): + + def _sync_pipeline_execute() -> Any: + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + return pipe.execute() + + res = await asyncio.to_thread(_sync_pipeline_execute) + return [None if v is None else float(v) for v in list(res)] + + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + res = pipe.execute() + if inspect.iscoroutine(res): + res = await res + return [None if v is None else float(v) for v in list(res)] + + async def _redis_zadd_and_expire( + self, + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + member_scores: Dict[str, int], + ) -> None: + if not member_scores: + return + await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) + await _redis_call(redis_client, "expire", key, self.state_ttl_seconds) + + def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: + suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") + if suffix: + return f"dedup:{self.merger_id}:{user_id}:{suffix}" + return f"dedup:{self.merger_id}:{user_id}" + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + if limit <= 0: + return FeedResult(data=[], next_page=next_page, has_next_page=False) + + entry = next_page.data.get(self.merger_id) + requested_page = entry.page if entry is not None else None + is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) + + if self.state_backend == "redis" and not redis_client: + raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") + + working_next_page = _pydantic_deep_copy(next_page) + + if is_fresh_session: + self._reset_descendant_cursors(working_next_page) + + seen_priority_map: Dict[str, int] = {} + seen_updates_in_order: List[tuple[str, int]] = [] + if self.state_backend == "cursor" and not is_fresh_session: + seen_priority_map = self._decode_seen_from_cursor(next_page) + + seen_request_set: set[str] = set(seen_priority_map.keys()) + + redis_state_key = "" + redis_new_scores: Dict[str, int] = {} + redis_seen_cache: Dict[str, Optional[int]] = {} + if self.state_backend == "redis" and redis_client: + redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) + if is_fresh_session: + await _redis_call(redis_client, "delete", redis_state_key) + + if self.state_backend == "cursor": + dedup_state: _DedupState = _CursorDedupState( + seen_priority_map=seen_priority_map, + seen_updates_in_order=seen_updates_in_order, + seen_request_set=seen_request_set, + ) + else: + assert redis_client is not None + dedup_state = _RedisDedupState( + redis_client=redis_client, + redis_state_key=redis_state_key, + redis_seen_cache=redis_seen_cache, + redis_new_scores=redis_new_scores, + seen_request_set=seen_request_set, + zmscore=self._redis_zmscore, + ) + + original_methods_dict = methods_dict + + child = _pydantic_deep_copy(self.data) + + rewritten_methods_dict = dict(original_methods_dict) + + for sf in self._iter_subfeeds(child): + self._register_wrapped_subfeed_method( + subfeed=sf, + original_methods_dict=original_methods_dict, + rewritten_methods_dict=rewritten_methods_dict, + dedup_state=dedup_state, + ) + + child_result = await child.get_data( + methods_dict=rewritten_methods_dict, + user_id=user_id, + limit=limit, + next_page=working_next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + + if self.state_backend == "redis" and redis_client: + await self._redis_zadd_and_expire(redis_client, redis_state_key, redis_new_scores) + + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + merger_after: Any = None + if self.state_backend == "cursor": + merger_after = self._encode_seen_for_cursor(seen_updates_in_order) + + result_next_page = _pydantic_deep_copy(child_result.next_page) + result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) + + return FeedResult(data=child_result.data, next_page=result_next_page, has_next_page=child_result.has_next_page) diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py new file mode 100644 index 0000000..76982fb --- /dev/null +++ b/smartfeed/mergers/percentage.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from random import shuffle +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast + +import redis +from pydantic import BaseModel +from redis.asyncio import Redis as AsyncRedis + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class MergerPercentageItem(BaseModel): + """One percentage slot.""" + + percentage: int + data: FeedTypes + + +class MergerPercentage(BaseFeedConfigModel): + """Percentage-based mixing merger.""" + + merger_id: str + type: Literal["merger_percentage"] + items: List[MergerPercentageItem] + shuffle: bool = False + + @staticmethod + async def _merge_items_data(items_data: List[List]) -> List: + result: List = [] + cursor: List[Dict] = [] + + min_length = min(len(item_data) for item_data in items_data) or 1 + for item_data in items_data: + cursor.append( + { + "items": item_data, + "current": 0, + "size": round(len(item_data) / min_length), + } + ) + + full_length = sum(len(item_data) for item_data in items_data) + while len(result) < full_length: + for item_cursor in cursor: + items = item_cursor["items"] + start = item_cursor["current"] + end = start + item_cursor["size"] if start + item_cursor["size"] < len(items) else len(items) + result.extend(items[start:end]) + item_cursor["current"] = end + + return result + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) + + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + items_data: List[List[Any]] = [[] for _ in self.items] + results: List[Optional[FeedResult]] = [None for _ in self.items] + + indexed_items = list(enumerate(self.items)) + fetch_order = indexed_items + if dedup_active: + fetch_order = sorted( + indexed_items, + key=lambda p: (getattr(p[1].data, "dedup_priority", 0), -p[0]), + reverse=True, + ) + + for idx, item in fetch_order: + item_result = cast( + FeedResult, + await item.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit * item.percentage // 100, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ), + ) + + results[idx] = item_result + + for idx, result_item in enumerate(results): + assert result_item is not None + items_data[idx] = result_item.data + + if not result.has_next_page and result_item.has_next_page: + result.has_next_page = True + result.next_page.data.update(result_item.next_page.data) + + result.data = await self._merge_items_data(items_data=items_data) + + if self.shuffle: + shuffle(result.data) + + return result diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py new file mode 100644 index 0000000..17ed530 --- /dev/null +++ b/smartfeed/mergers/percentage_gradient.py @@ -0,0 +1,164 @@ +from random import shuffle +from typing import Any, Callable, Dict, Literal, Optional, Union + +import redis +from pydantic import model_validator +from redis.asyncio import Redis as AsyncRedis + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside +from .percentage import MergerPercentageItem + + +class MergerPercentageGradient(BaseFeedConfigModel): + """Percentage-gradient merger.""" + + merger_id: str + type: Literal["merger_percentage_gradient"] + item_from: MergerPercentageItem + item_to: MergerPercentageItem + step: int + size_to_step: int + shuffle: bool = False + + @model_validator(mode="after") + def validate_merger_percentage_gradient(self) -> "MergerPercentageGradient": + if self.step < 1 or self.step > 100: + raise ValueError('"step" must be in range from 1 to 100') + if self.size_to_step < 1: + raise ValueError('"size_to_step" must be bigger than 1') + return self + + async def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: + result: Dict = { + "limit_from": 0, + "limit_to": 0, + "percentages": [], + } + + percentage_from = self.item_from.percentage + percentage_to = self.item_to.percentage + start_position = limit * (page - 1) + first_iter = True + + for i in range(self.size_to_step, limit * page + self.size_to_step, self.size_to_step): + if not first_iter and percentage_to < 100: + percentage_from -= self.step + percentage_to += self.step + + if percentage_to > 100 or percentage_from < 0: + percentage_from = 0 + percentage_to = 100 + + if i > start_position: + iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) + start_position = i + + if result["percentages"] and result["percentages"][-1]["to"] >= 100: + result["limit_to"] += iter_limit + result["percentages"][-1]["limit"] += iter_limit + else: + result["limit_from"] += iter_limit * percentage_from // 100 + result["limit_to"] += iter_limit * percentage_to // 100 + iter_result = {"limit": iter_limit, "from": percentage_from, "to": percentage_to} + result["percentages"].append(iter_result) + + if first_iter: + first_iter = False + + return result + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + result = FeedResult( + data=[], + next_page=FeedResultNextPage( + data={ + self.merger_id: FeedResultNextPageInside( + page=next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1, + after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, + ) + }, + ), + has_next_page=False, + ) + + limits_and_percents = await self._calculate_limits_and_percents( + page=result.next_page.data[self.merger_id].page, + limit=limit, + ) + + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + from_priority = getattr(self.item_from.data, "dedup_priority", 0) + to_priority = getattr(self.item_to.data, "dedup_priority", 0) + + if dedup_active and to_priority > from_priority: + item_to = await self.item_to.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_to"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + item_from = await self.item_from.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_from"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + else: + item_from = await self.item_from.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_from"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + item_to = await self.item_to.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limits_and_percents["limit_to"], + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + + from_start_index = 0 + to_start_index = 0 + for lp_data in limits_and_percents["percentages"]: + from_end_index = (lp_data["limit"] * lp_data["from"] // 100) + from_start_index + to_end_index = (lp_data["limit"] * lp_data["to"] // 100) + to_start_index + + result.data.extend(item_from.data[from_start_index:from_end_index]) + result.data.extend(item_to.data[to_start_index:to_end_index]) + + from_start_index = from_end_index + to_start_index = to_end_index + + result.next_page.data.update(item_from.next_page.data) + result.next_page.data.update(item_to.next_page.data) + + if any([item_from.has_next_page, item_to.has_next_page]): + result.has_next_page = True + + if self.shuffle: + shuffle(result.data) + + result.next_page.data[self.merger_id].page += 1 + + return result diff --git a/smartfeed/mergers/positional.py b/smartfeed/mergers/positional.py new file mode 100644 index 0000000..a543ddb --- /dev/null +++ b/smartfeed/mergers/positional.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union + +import redis +from pydantic import model_validator +from redis.asyncio import Redis as AsyncRedis + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class MergerPositional(BaseFeedConfigModel): + """Positional merger.""" + + merger_id: str + type: Literal["merger_positional"] + positions: List[int] = [] + start: Optional[int] = None + end: Optional[int] = None + step: Optional[int] = None + positional: FeedTypes + default: FeedTypes + + @model_validator(mode="after") + def validate_merger_positional(self) -> "MergerPositional": + if not self.positions and not all((self.start, self.end, self.step)): + raise ValueError('Either "positions" or "start", "end", and "step" must be provided') + if self.start and self.positions: + if isinstance(self.start, int) and self.start <= max(self.positions): + raise ValueError('"start" must be bigger than maximum value of "positions"') + if isinstance(self.start, int) and isinstance(self.end, int): + if self.end <= self.start: + raise ValueError('"end" must be bigger than "start"') + return self + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + dedup_active = bool(params.pop("_sf_dedup_active", False)) + + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + + positional_has_next_page = True + page_positions: List[int] = [] + available_positions = range((page - 1) * limit, (page * limit) + 1) + for position in self.positions: + if position in available_positions: + page_positions.append(available_positions.index(position)) + + if max(available_positions) >= max(self.positions, default=0): + positional_has_next_page = False + + if self.start is not None and self.end is not None and self.step is not None: + positional_has_next_page = not max(available_positions) >= self.end + + for position in range(self.start, self.end, self.step): + if position in available_positions: + page_positions.append(available_positions.index(position)) + + if dedup_active and getattr(self.positional, "dedup_priority", 0) > getattr(self.default, "dedup_priority", 0): + pos_res = await self.positional.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=len(page_positions), + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + default_res = await self.default.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=True, + **params, + ) + else: + default_res = await self.default.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + pos_res = await self.positional.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=len(page_positions), + next_page=next_page, + redis_client=redis_client, + _sf_dedup_active=dedup_active, + **params, + ) + + result = FeedResult( + data=default_res.data, + next_page=FeedResultNextPage( + data={ + self.merger_id: FeedResultNextPageInside( + page=page, + after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, + ) + }, + ), + has_next_page=default_res.has_next_page, + ) + + if not result.has_next_page and all([positional_has_next_page, pos_res.has_next_page]): + result.has_next_page = True + + result.next_page.data.update(default_res.next_page.data) + result.next_page.data.update(pos_res.next_page.data) + + for i, post in enumerate(pos_res.data): + result.data = result.data[: page_positions[i] - 1] + [post] + result.data[page_positions[i] - 1 :] + + if len(result.data) > limit: + result.data = result.data[:limit] + + result.next_page.data[self.merger_id].page += 1 + + return result diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py new file mode 100644 index 0000000..cd3c2f8 --- /dev/null +++ b/smartfeed/mergers/view_session.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import logging +from random import shuffle +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union + +import redis +from redis.asyncio import Redis as AsyncRedis +from redis.asyncio import RedisCluster as AsyncRedisCluster + +from .. import jsonlib as json +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside, _redis_call + +if TYPE_CHECKING: + from ..schemas import FeedTypes + + +class MergerViewSession(BaseFeedConfigModel): + """Merger with view-session caching.""" + + merger_id: str + type: Literal["merger_view_session"] + session_size: int + session_live_time: int + data: "FeedTypes" + deduplicate: bool = False + dedup_key: str = None # type: ignore + shuffle: bool = False + + def _get_dedup_key_or_attr(self, item: Any) -> str: + if not self.dedup_key: + return item + + try: + dedup_value = item.get(self.dedup_key) + except AttributeError: + dedup_value = getattr(item, self.dedup_key, None) + + assert dedup_value is not None, f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}" + return dedup_value + + def _dedup_data(self, data: List[Any]) -> List[Any]: + deduplicated_data = {self._get_dedup_key_or_attr(item): item for item in data} + return list(deduplicated_data.values()) + + async def _set_cache( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + redis_client: Union[redis.Redis, AsyncRedis], + cache_key: str, + **params: Any, + ) -> List[Any]: + result = await self.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=self.session_size, + next_page=FeedResultNextPage(data={}), + **params, + ) + + data = result.data + if self.deduplicate: + data = self._dedup_data(data) + await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) + return data + + async def _set_cache_async( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + redis_client: AsyncRedis, + cache_key: str, + **params: Any, + ) -> List[Any]: + result = await self.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=self.session_size, + next_page=FeedResultNextPage(data={}), + **params, + ) + + data = result.data + if self.deduplicate: + data = self._dedup_data(data) + await redis_client.set(cache_key, json.dumps(data)) + await redis_client.expire(cache_key, self.session_live_time) + return data + + async def _get_cache( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Union[redis.Redis, AsyncRedis], + **params: Any, + ) -> FeedResult: + if session_cache_key := params.get("custom_view_session_key", None): + cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" + else: + cache_key = f"{self.merger_id}_{user_id}" + + logging.info("MergerViewSession cache request for %s", cache_key) + cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) + if not cache_exists or self.merger_id not in next_page.data: + logging.info("Cache miss or new session - generating fresh data for %s", cache_key) + session_data = await self._set_cache( + methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + ) + else: + logging.info("Cache exists - attempting read from Redis for %s", cache_key) + cached_data = await _redis_call(redis_client, "get", cache_key) + if cached_data is None: + logging.info( + "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key + ) + session_data = await self._set_cache( + methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + ) + else: + logging.info("Successfully read cached data for %s", cache_key) + session_data = json.loads(cached_data) + + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + return FeedResult( + data=session_data[(page - 1) * limit :][:limit], + next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), + has_next_page=bool(len(session_data) > limit * page), + ) + + async def _get_cache_async( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: AsyncRedis, + **params: Any, + ) -> FeedResult: + if session_cache_key := params.get("custom_view_session_key", None): + cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" + else: + cache_key = f"{self.merger_id}_{user_id}" + + if not await redis_client.exists(cache_key) or self.merger_id not in next_page.data: + session_data = await self._set_cache_async( + methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + ) + else: + cached_data = await redis_client.get(cache_key) + if cached_data is None: + logging.info( + "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key + ) + session_data = await self._set_cache_async( + methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + ) + else: + logging.info("Successfully read cached data for %s", cache_key) + session_data = json.loads(cached_data) + + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + return FeedResult( + data=session_data[(page - 1) * limit :][:limit], + next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), + has_next_page=bool(len(session_data) > limit * page), + ) + + async def get_data( + self, + methods_dict: Dict[str, Callable], + user_id: Any, + limit: int, + next_page: FeedResultNextPage, + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + **params: Any, + ) -> FeedResult: + if not redis_client: + raise ValueError("Redis client must be provided if using Merger View Session") + + if isinstance(redis_client, (AsyncRedis, AsyncRedisCluster)): + result = await self._get_cache_async( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) + else: + result = await self._get_cache( + methods_dict=methods_dict, + user_id=user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + **params, + ) + + if self.shuffle: + shuffle(result.data) + + return result diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index ed04203..9863bb4 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -1,1875 +1,95 @@ -import asyncio -import base64 -import inspect -from . import jsonlib as json -import logging -import zlib -from abc import ABC, abstractmethod -from collections import defaultdict, deque -from dataclasses import dataclass -from random import shuffle -from typing import ( - Annotated, - Any, - Awaitable, - Callable, - Dict, - Iterator, - List, - Literal, - Optional, - Union, - cast, - no_type_check, -) - -import redis -from pydantic import BaseModel, Field, PrivateAttr, model_validator -from redis.asyncio import Redis as AsyncRedis -from redis.asyncio import RedisCluster as AsyncRedisCluster - - -def _is_async_redis_client(client: Any) -> bool: - return isinstance(client, (AsyncRedis, AsyncRedisCluster)) - - -async def _redis_call(client: Any, method_name: str, *args: Any, **kwargs: Any) -> Any: - """Call a Redis method without blocking the event loop. - - - For `redis.asyncio` clients, calls are awaited directly. - - For sync `redis.Redis`, calls are offloaded via `asyncio.to_thread()`. - """ - - method = getattr(client, method_name) - if _is_async_redis_client(client): - return await method(*args, **kwargs) - return await asyncio.to_thread(method, *args, **kwargs) - - -def _pydantic_deep_copy(model: Any) -> Any: - """Deep copy helper compatible with Pydantic v1 and v2.""" - - if hasattr(model, "model_copy"): - return model.model_copy(deep=True) - return model.copy(deep=True) - - -class _DedupState(ABC): - @abstractmethod - def should_accept(self, key: str, priority: int) -> bool: - raise NotImplementedError - - @abstractmethod - def record(self, key: str, priority: int) -> None: - raise NotImplementedError - - async def prefetch(self, keys: List[str]) -> None: - return - - -@dataclass -class _CursorDedupState(_DedupState): - seen_priority_map: Dict[str, int] - seen_updates_in_order: List[tuple[str, int]] - seen_request_set: set[str] - - def should_accept(self, key: str, priority: int) -> bool: - if key in self.seen_request_set: - return False - existing_priority = self.seen_priority_map.get(key) - if existing_priority is not None and priority <= existing_priority: - return False - return True - - def record(self, key: str, priority: int) -> None: - self.seen_priority_map[key] = priority - self.seen_updates_in_order.append((key, priority)) - self.seen_request_set.add(key) +"""Public schema surface. +This module keeps the public import path (`smartfeed.schemas`) stable while +moving merger implementations into `smartfeed.mergers.*`. +""" -@dataclass -class _RedisDedupState(_DedupState): - redis_client: Union[redis.Redis, AsyncRedis] - redis_state_key: str - redis_seen_cache: Dict[str, Optional[int]] - redis_new_scores: Dict[str, int] - seen_request_set: set[str] - zmscore: Callable[ - [Union[redis.Redis, AsyncRedis], str, List[str]], - Union[Awaitable[List[Optional[float]]], List[Optional[float]]], - ] +from __future__ import annotations - async def prefetch(self, keys: List[str]) -> None: - if not keys: - return - unique: List[str] = [] - seen: set[str] = set() - for k in keys: - if k in self.seen_request_set: - continue - if k in self.redis_seen_cache: - continue - if k in seen: - continue - seen.add(k) - unique.append(k) +from typing import Annotated, Any, Dict, Union - if not unique: - return - - scores_result = self.zmscore(self.redis_client, self.redis_state_key, unique) - if inspect.iscoroutine(scores_result): - scores = await cast(Awaitable[List[Optional[float]]], scores_result) - else: - scores = cast(List[Optional[float]], scores_result) - - for k, s in zip(unique, scores): - self.redis_seen_cache[k] = None if s is None else int(s) - - def should_accept(self, key: str, priority: int) -> bool: - if key in self.seen_request_set: - return False - existing_priority = self.redis_seen_cache.get(key) - if existing_priority is not None and priority <= existing_priority: - return False - return True - - def record(self, key: str, priority: int) -> None: - self.seen_request_set.add(key) - self.redis_seen_cache[key] = priority - self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) +from pydantic import BaseModel, Field +from .feed_models import ( + BaseFeedConfigModel, + FeedResult, + FeedResultClient, + FeedResultNextPage, + FeedResultNextPageInside, + SubFeed, +) +from .mergers import ( + MergerAppend, + MergerAppendDistribute, + MergerDeduplication, + MergerPercentage, + MergerPercentageGradient, + MergerPercentageItem, + MergerPositional, + MergerViewSession, +) FeedTypes = Annotated[ Union[ - "MergerDeduplication", - "MergerAppend", - "MergerAppendDistribute", - "MergerPositional", - "MergerPercentage", - "MergerPercentageGradient", - "MergerViewSession", - "SubFeed", + MergerDeduplication, + MergerAppend, + MergerAppendDistribute, + MergerPositional, + MergerPercentage, + MergerPercentageGradient, + MergerViewSession, + SubFeed, ], Field(discriminator="type"), ] -class FeedResultNextPageInside(BaseModel): - """ - Модель данных курсора пагинации конкретной позиции. - - Attributes: - page порядковый номер страницы. - after данные для пагинации клиентского метода. - """ - - page: int = 1 - after: Any = None - - -class FeedResultNextPage(BaseModel): - """ - Модель курсора пагинации. - - Attributes: - data словарь вида "ключ: данные по пагинации", где ключ - subfeed_id или merger_id. - """ - - data: Dict[str, FeedResultNextPageInside] - - -class FeedResult(BaseModel): - """ - Модель результата метода get_data() любой позиции / целого фида. - - Attributes: - data список данных, возвращенных мерджером / субфидом. - next_page курсор пагинации. - has_next_page флаг наличия следующей страницы данных. - """ - - data: List - next_page: FeedResultNextPage - has_next_page: bool - - -class FeedResultClient(BaseModel): - """ - Модель результата клиентского метода субфида. - - Attributes: - data список данных, возвращенных мерджером / субфидом. - next_page курсор пагинации клиентского метода. - has_next_page флаг наличия следующей страницы данных. - """ - - data: List - next_page: FeedResultNextPageInside - has_next_page: bool - - -class BaseFeedConfigModel(ABC, BaseModel): - """ - Абстрактный класс для мерджера / субфида конфигурации. - """ - - # Higher value means the item should "win" deduplication when duplicates exist. - # This is primarily used by MergerDeduplication and by mergers when a dedup wrapper is active. - dedup_priority: int = 0 - - @abstractmethod - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: параметры для метода. - :return: список данных. - """ - - -class MergerViewSession(BaseFeedConfigModel): - """ - Модель мерджера с кэшированием. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_view_session". - view_session флаг использования механизма расчета всего фида сразу и сохранения в кэш. - session_size размер кэшируемого фида (limit получения данных для сохранения в кэш). - session_live_time срок хранения в кэше для кэшируемого фида (в секундах). - data мерджер или субфид. - deduplicate флаг дедупликации (удаления дублей из сессии). - dedup_key название ключа или атрибута, по которому логика дедпликации найдет дубли. - shuffle флаг для перемешивания полученных данных мерджера. - """ - - merger_id: str - type: Literal["merger_view_session"] - session_size: int - session_live_time: int - data: FeedTypes - deduplicate: bool = False - dedup_key: str = None # type: ignore - shuffle: bool = False - - def _get_dedup_key_or_attr(self, item: Any) -> str: - """ - Метод для получения ключа объекта кешируемой сессии. - - Если указанное в конфиге сессии название ключа имеет значение None, - в качестве ключа вернется сам объект. - Если название ключа не None, и для одного из объектов ни найден ни ключ, ни атрибут, - метод выбросит AssertionError. - - :param item: объект, для которого нужен ключ. - :return: ключ объекта. - """ - - if not self.dedup_key: - return item - - try: - dedup_value = item.get(self.dedup_key) - except AttributeError: - dedup_value = getattr(item, self.dedup_key, None) - - assert dedup_value is not None, f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}" - return dedup_value - - def _dedup_data(self, data: List[Any]) -> List[Any]: - """ - Метод для удаления дублей в списке data с сохранением последовательности. - - :param data: список, в котором нужно удалить дубли. - :return: результат удаления дублей. - """ - - deduplicated_data = {self._get_dedup_key_or_attr(item): item for item in data} - return list(deduplicated_data.values()) - - async def _set_cache( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - redis_client: Union[redis.Redis, AsyncRedis], - cache_key: str, - **params: Any, - ) -> List[Any]: - """ - Метод для кэширования данных Merger View Session. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param redis_client: объект клиента Redis. - :param cache_key: ключ для кэширования. - :param params: любые внешние параметры, передаваемые в исполняемую функцию на клиентской стороне. - :return: обработанные данные, которые были записаны в кэш. - """ - - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - **params, - ) - - data = result.data - if self.deduplicate: - data = self._dedup_data(data) - await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) - return data - - async def _set_cache_async( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - redis_client: AsyncRedis, - cache_key: str, - **params: Any, - ) -> List[Any]: - """ - Метод для кэширования данных Merger View Session. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param redis_client: объект клиента Redis. - :param cache_key: ключ для кэширования. - :param params: любые внешние параметры, передаваемые в исполняемую функцию на клиентской стороне. - :return: обработанные данные, которые были записаны в кэш. - """ - - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - **params, - ) - - data = result.data - if self.deduplicate: - data = self._dedup_data(data) - await redis_client.set(cache_key, json.dumps(data)) - await redis_client.expire(cache_key, self.session_live_time) - return data - - async def _get_cache( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Union[redis.Redis, AsyncRedis], - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных Merger View Session из кэша Redis. - При отсутствии данных в кэше - получить и сохранить. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: лимит на выдачу данных. - :param next_page: курсор для пагинации в формате SmartFeedResultNextPage. - :param redis_client: объект клиента Redis. - :param params: любые внешние параметры, передаваемые в исполняемую функцию на клиентской стороне. - :return: результат получения данных согласно конфигурации фида. - """ - - # Формируем ключ для кэширования данных мерджера. - if session_cache_key := params.get("custom_view_session_key", None): - cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" - else: - cache_key = f"{self.merger_id}_{user_id}" - - logging.info("MergerViewSession cache request for %s", cache_key) - # Если кэш не найден или передан пустой курсор пагинации на мерджер, обновляем данные и записываем в кэш. - cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) - if not cache_exists or self.merger_id not in next_page.data: - logging.info("Cache miss or new session - generating fresh data for %s", cache_key) - # Получаем свежие данные и используем их напрямую (избегаем чтение из кэша) - session_data = await self._set_cache( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params - ) - else: - logging.info("Cache exists - attempting read from Redis for %s", cache_key) - # Читаем из кэша только если он уже существовал - cached_data = await _redis_call(redis_client, "get", cache_key) - if cached_data is None: - # Fallback: если кэш пропал, получаем свежие данные - logging.info( - "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key - ) - session_data = await self._set_cache( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params - ) - else: - logging.info("Successfully read cached data for %s", cache_key) - session_data = json.loads(cached_data) - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - result = FeedResult( - data=session_data[(page - 1) * limit :][:limit], - next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - has_next_page=bool(len(session_data) > limit * page), - ) - return result - - async def _get_cache_async( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: AsyncRedis, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных Merger View Session из кэша Redis. - При отсутствии данных в кэше - получить и сохранить. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: лимит на выдачу данных. - :param next_page: курсор для пагинации в формате SmartFeedResultNextPage. - :param redis_client: объект клиента Redis. - :param params: любые внешние параметры, передаваемые в исполняемую функцию на клиентской стороне. - :return: результат получения данных согласно конфигурации фида. - """ - - # Формируем ключ для кэширования данных мерджера. - if session_cache_key := params.get("custom_view_session_key", None): - cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" - else: - cache_key = f"{self.merger_id}_{user_id}" - - # Если кэш не найден или передан пустой курсор пагинации на мерджер, обновляем данные и записываем в кэш. - if not await redis_client.exists(cache_key) or self.merger_id not in next_page.data: - # Получаем свежие данные и используем их напрямую (избегаем чтение из кэша) - session_data = await self._set_cache_async( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params - ) - else: - # Читаем из кэша только если он уже существовал - cached_data = await redis_client.get(cache_key) - if cached_data is None: - # Fallback: если кэш пропал, получаем свежие данные - logging.info( - "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key - ) - session_data = await self._set_cache_async( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params - ) - else: - logging.info("Successfully read cached data for %s", cache_key) - session_data = json.loads(cached_data) - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - result = FeedResult( - data=session_data[(page - 1) * limit :][:limit], - next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - has_next_page=bool(len(session_data) > limit * page), - ) - return result - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных методом append. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных методом append. - """ - - # Проверяем наличие клиента Redis в конфигурации фида. - if not redis_client: - raise ValueError("Redis client must be provided if using Merger View Session") - - # Формируем результат view session мерджера. - if isinstance(redis_client, (AsyncRedis, AsyncRedisCluster)): - result = await self._get_cache_async( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - else: - result = await self._get_cache( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - - # Если в конфигурации указано "смешать" данные. - if self.shuffle: - shuffle(result.data) - - return result - - -class MergerAppend(BaseFeedConfigModel): - """ - Модель append мерджера. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_append". - items позиции мерджера. - shuffle флаг для перемешивания полученных данных мерджера. - """ - - merger_id: str - type: Literal["merger_append"] - items: List[FeedTypes] - shuffle: bool = False - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных методом append. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных методом append. - """ - - # When a MergerDeduplication wrapper is active, we may need to respect dedup_priority - # across children without changing the append output order. In that mode we fetch in - # priority order, then concatenate in the configured order and trim to `limit`. - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - - if dedup_active: - indexed_items = list(enumerate(self.items)) - fetched: Dict[int, FeedResult] = {} - - # Always-on parallelism in dedup mode: preserve dedup semantics by ensuring - # higher priority is recorded first; fetch same-priority children concurrently. - groups: Dict[int, List[tuple[int, FeedTypes]]] = defaultdict(list) - for idx, item in indexed_items: - prio = int(getattr(item, "dedup_priority", 0)) - groups[prio].append((idx, item)) - - for prio in sorted(groups.keys(), reverse=True): - group = groups[prio] - coros: List[Awaitable[FeedResult]] = [] - order: List[int] = [] - for idx, item in group: - order.append(idx) - coros.append( - item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=_pydantic_deep_copy(next_page), - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - ) - group_results = await asyncio.gather(*coros) - for idx, r in zip(order, group_results): - fetched[idx] = cast(FeedResult, r) - - for idx, _item in indexed_items: - item_result = fetched[idx] - result.data.extend(item_result.data) - result.next_page.data.update(item_result.next_page.data) - if item_result.has_next_page: - result.has_next_page = True - - if len(result.data) > limit: - result.data = result.data[:limit] - else: - result_limit = limit - for item in self.items: - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - - result.data.extend(item_result.data) - result_limit -= len(item_result.data) - - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True - - result.next_page.data.update(item_result.next_page.data) - - if result_limit <= 0: - break - - # Если в конфигурации указано "смешать" данные. - if self.shuffle: - shuffle(result.data) - - return result - - -class MergerPositional(BaseFeedConfigModel): - """ - Модель позиционного мерджера. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_positional". - positions позиции для вставки из мерджера / субфида "positional" [обязателен, если нет start, end, step]. - start начальная позиция [обязателен, если нет positions]. - end завершающая позиция [обязателен, если нет positions]. - step шаг позиций между "start" и "end" [обязателен, если нет positions]. - positional мерджер / субфид из которого берутся позиционные данные. - default мерджер / субфид из которого берутся остальные данные. - """ - - merger_id: str - type: Literal["merger_positional"] - positions: List[int] = [] - start: Optional[int] = None - end: Optional[int] = None - step: Optional[int] = None - positional: FeedTypes - default: FeedTypes - - @model_validator(mode="after") - def validate_merger_positional(self) -> "MergerPositional": - if not self.positions and not all((self.start, self.end, self.step)): - raise ValueError('Either "positions" or "start", "end", and "step" must be provided') - if self.start and self.positions: - if isinstance(self.start, int) and self.start <= max(self.positions): - raise ValueError('"start" must be bigger than maximum value of "positions"') - if isinstance(self.start, int) and isinstance(self.end, int): - if self.end <= self.start: - raise ValueError('"end" must be bigger than "start"') - return self - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных в позиционном соотношении из данных позиций. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных в процентном соотношении. - """ - - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - # Determine the merger page first (independent of children). - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - - positional_has_next_page = True - page_positions: List[int] = [] - available_positions = range((page - 1) * limit, (page * limit) + 1) - for position in self.positions: - if position in available_positions: - page_positions.append(available_positions.index(position)) - - # Если конечная позиция текущей страницы больше или равна MAX позиции в конфигурации, то has_next_page = False - if max(available_positions) >= max(self.positions, default=0): - positional_has_next_page = False - - if self.start is not None and self.end is not None and self.step is not None: - # Если конечная позиция текущей страницы больше или равна конечной шаговой позиции, то has_next_page = False - positional_has_next_page = not max(available_positions) >= self.end - - for position in range(self.start, self.end, self.step): - if position in available_positions: - page_positions.append(available_positions.index(position)) - - default_res: FeedResult - pos_res: FeedResult - - if dedup_active and getattr(self.positional, "dedup_priority", 0) > getattr(self.default, "dedup_priority", 0): - pos_res = await self.positional.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=len(page_positions), - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - default_res = await self.default.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - else: - default_res = await self.default.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - pos_res = await self.positional.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=len(page_positions), - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - - result = FeedResult( - data=default_res.data, - next_page=FeedResultNextPage( - data={ - self.merger_id: FeedResultNextPageInside( - page=page, - after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, - ) - }, - ), - has_next_page=default_res.has_next_page, - ) - - # Если has_next_page = False, то проверяем has_next_page у позиции и, если необходимо, обновляем. - if not result.has_next_page and all([positional_has_next_page, pos_res.has_next_page]): - result.has_next_page = True - - # Обновляем next_page. - result.next_page.data.update(default_res.next_page.data) - result.next_page.data.update(pos_res.next_page.data) - - # Формируем общие данные позиционного мерджера. - for i, post in enumerate(pos_res.data): - result.data = result.data[: page_positions[i] - 1] + [post] + result.data[page_positions[i] - 1 :] - - # Проверка на возврат данных в количестве не более limit. - if len(result.data) > limit: - result.data = result.data[:limit] - - # Обновляем страницу для курсора пагинации мерджера. - result.next_page.data[self.merger_id].page += 1 - - return result - - -class MergerPercentageItem(BaseModel): - """ - Модель позиции процентного мерджера. - - Attributes: - percentage процент позиции в мерджере. - data мерджер / субфид. - """ - - percentage: int - data: FeedTypes - - -class MergerPercentage(BaseFeedConfigModel): - """ - Модель процентного мерджера. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_percentage". - shuffle флаг для перемешивания полученных данных мерджера. - items позиции мерджера. - """ - - merger_id: str - type: Literal["merger_percentage"] - items: List[MergerPercentageItem] - shuffle: bool = False - - @staticmethod - async def _merge_items_data(items_data: List[List]) -> List: - """ - Метод для получения максимально равномерно распределенных данных позиций процентного мерджера. - - :param items_data: список со списками данных из каждой позиции. - :return: максимально равномерно распределенные данные позиций процентного мерджера. - """ - - # Формируем возвращаемый результат и список курсоров для списка каждой позиции. - result: List = [] - cursor: List[Dict] = [] - - # Получаем длину самого маленького списка и формируем курсор для каждого списка. - min_length = min(len(item_data) for item_data in items_data) or 1 - for item_data in items_data: - cursor.append( - { - "items": item_data, - "current": 0, - "size": round(len(item_data) / min_length), - } - ) - - # Получаем общий размер всех элементов всех списков и пока не получаем результат такого же размера - # производим операции по распределению элементов. - full_length = sum(len(item_data) for item_data in items_data) - while len(result) < full_length: - for item_cursor in cursor: - items = item_cursor["items"] - start = item_cursor["current"] - end = start + item_cursor["size"] if start + item_cursor["size"] < len(items) else len(items) - result.extend(items[start:end]) - item_cursor["current"] = end - - return result - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных в процентном соотношении из данных позиций. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных в процентном соотношении. - """ - - # Формируем результат процентного мерджера. - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - items_data: List[List[Any]] = [[] for _ in self.items] - results: List[Optional[FeedResult]] = [None for _ in self.items] - - indexed_items = list(enumerate(self.items)) - fetch_order = indexed_items - if dedup_active: - fetch_order = sorted( - indexed_items, - key=lambda p: (getattr(p[1].data, "dedup_priority", 0), -p[0]), - reverse=True, - ) - - for idx, item in fetch_order: - item_result = cast( - FeedResult, - await item.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit * item.percentage // 100, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ), - ) - - results[idx] = item_result - - for idx, result_item in enumerate(results): - assert result_item is not None - items_data[idx] = result_item.data - - if not result.has_next_page and result_item.has_next_page: - result.has_next_page = True - result.next_page.data.update(result_item.next_page.data) - - # Добавляем данные позиции к общему результату процентного мерджера. - result.data = await self._merge_items_data(items_data=items_data) - - # Если в конфигурации указано "смешать" данные. - if self.shuffle: - shuffle(result.data) - - return result - - -class MergerPercentageGradient(BaseFeedConfigModel): - """ - Модель процентного мерджера с градиентном. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_percentage_gradient". - item_from мерджер / субфид из которого начинается "перетекание" градиента. - item_to мерджер / субфид в который "перетекает" градиент. - step изменение в % соотношения из item_from в item_to. - size_to_step шаг для применения изменений % соотношения (например, через каждые 30 позиций). - shuffle флаг для перемешивания полученных данных мерджера. - """ - - merger_id: str - type: Literal["merger_percentage_gradient"] - item_from: MergerPercentageItem - item_to: MergerPercentageItem - step: int - size_to_step: int - shuffle: bool = False - - @model_validator(mode="after") - def validate_merger_percentage_gradient(self) -> "MergerPercentageGradient": - if self.step < 1 or self.step > 100: - raise ValueError('"step" must be in range from 1 to 100') - if self.size_to_step < 1: - raise ValueError('"size_to_step" must be bigger than 1') - return self - - async def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: - """ - Метод для получения списка лимитов данных с процентным соотношением позиций item_from & item_to, - учитывая градиентное изменение соотношений. - - :param page: порядковый номер страницы. - :param limit: общий лимит данных для страницы. - :return: список лимитов данных с процентным соотношением позиций item_from & item_to. - """ - - result: Dict = { - "limit_from": 0, - "limit_to": 0, - "percentages": [], - } - - percentage_from = self.item_from.percentage - percentage_to = self.item_to.percentage - start_position = limit * (page - 1) - first_iter = True - - for i in range(self.size_to_step, limit * page + self.size_to_step, self.size_to_step): - # При первой итерации и percentage_to >= 100 не меняем соотношение % между позициями. - if not first_iter and percentage_to < 100: - # Меняем процентное соотношение позиций на "шаг", указанный в конфигурации. - percentage_from -= self.step - percentage_to += self.step - - # Если процентное соотношение вышло за 100+, то устанавливаем предельные значения. - if percentage_to > 100 or percentage_from < 0: - percentage_from = 0 - percentage_to = 100 - - # Если индекс итерации по величине больше стартовой позиции согласно переданной странице, - # то начинаем обработку. - if i > start_position: - # Рассчитываем лимит получения данных для конкретной итерации. - iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) - start_position = i - - # Формируем результат для каждой итерации и добавляем в возвращаемый список, но если процентное - # соотношение у последней итерации 0 - 100, то добавляем лимит к ней. - if result["percentages"] and result["percentages"][-1]["to"] >= 100: - result["limit_to"] += iter_limit - result["percentages"][-1]["limit"] += iter_limit - else: - result["limit_from"] += iter_limit * percentage_from // 100 - result["limit_to"] += iter_limit * percentage_to // 100 - iter_result = {"limit": iter_limit, "from": percentage_from, "to": percentage_to} - result["percentages"].append(iter_result) - - # Если первая итерация цикла - if first_iter: - first_iter = False - - return result - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных в процентном соотношении с градиентом из данных позиций. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных в процентном соотношении. - """ - - # Формируем результат процентного мерджера с градиентом. - result = FeedResult( - data=[], - next_page=FeedResultNextPage( - data={ - self.merger_id: FeedResultNextPageInside( - page=next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1, - after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, - ) - }, - ), - has_next_page=False, - ) - - # Получаем список лимитов данных и соотношений согласно странице и градиенту. - limits_and_percents = await self._calculate_limits_and_percents( - page=result.next_page.data[self.merger_id].page, - limit=limit, - ) - - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - from_priority = getattr(self.item_from.data, "dedup_priority", 0) - to_priority = getattr(self.item_to.data, "dedup_priority", 0) - - if dedup_active and to_priority > from_priority: - item_to = await self.item_to.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_to"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - item_from = await self.item_from.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_from"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - else: - item_from = await self.item_from.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_from"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - item_to = await self.item_to.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_to"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - - from_start_index = 0 - to_start_index = 0 - for lp_data in limits_and_percents["percentages"]: - # Высчитываем лимиты для каждой позиции исходя из процентного соотношения. - from_end_index = (lp_data["limit"] * lp_data["from"] // 100) + from_start_index - to_end_index = (lp_data["limit"] * lp_data["to"] // 100) + to_start_index - - # Добавляем данные позиции к общему результату процентного мерджера с градиентом. - result.data.extend(item_from.data[from_start_index:from_end_index]) - result.data.extend(item_to.data[to_start_index:to_end_index]) - - # Обновляем стартовые индексы. - from_start_index = from_end_index - to_start_index = to_end_index - - # Обновляем next_page. - result.next_page.data.update(item_from.next_page.data) - result.next_page.data.update(item_to.next_page.data) - - # Если has_next_page = False, то проверяем has_next_page у позиций и, если необходимо, обновляем. - if any([item_from.has_next_page, item_to.has_next_page]): - result.has_next_page = True - - # Если в конфигурации указано "смешать" данные. - if self.shuffle: - shuffle(result.data) - - # Обновляем страницу для курсора пагинации мерджера. - result.next_page.data[self.merger_id].page += 1 - - return result - - -class MergerAppendDistribute(BaseFeedConfigModel): - """ - Модель мерджера, равномерно распределяющего данные по ключу. - - Attributes: - merger_id уникальный ID мерджера. - type тип объекта - всегда "merger_distribute". - items позиции мерджера. - distribution_key ключ для распределения данных мерджера. - sorting_key ключ сортировки. - sorting_desc флаг сортировки по убыванию. - """ - - merger_id: str - type: Literal["merger_distribute"] - items: List[FeedTypes] - distribution_key: str - sorting_key: Optional[str] = None - sorting_desc: bool = False - - @no_type_check - async def _uniform_distribute(self, data: list) -> list: - # Сортируем записи глобально по `created_at` в порядке убывания - if self.sorting_key: - data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) - - # Группируем записи по `distribution_key` - grouped_entries = defaultdict(deque) - for entry in data: - grouped_entries[entry[self.distribution_key]].append(entry) - result = [] - prev_profile_id = None - while any(grouped_entries.values()): - for profile_id in list(grouped_entries.keys()): - if grouped_entries[profile_id]: - # Если текущий `distribution_key` отличается от предыдущего или он последний, берем его - if profile_id != prev_profile_id or len(grouped_entries) == 1: - result.append(grouped_entries[profile_id].popleft()) - prev_profile_id = profile_id - if not grouped_entries[profile_id]: # Если записи закончились, удаляем ключ из группы - del grouped_entries[profile_id] - else: - del grouped_entries[profile_id] - - return result - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных методом append. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: для метода класса. - :return: список данных методом append. - """ - - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - - if dedup_active: - indexed_items = list(enumerate(self.items)) - fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) - fetched: Dict[int, FeedResult] = {} - - for idx, item in fetch_order: - fetched[idx] = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - - for idx, _item in indexed_items: - item_result = fetched[idx] - result.data.extend(item_result.data) - result.next_page.data.update(item_result.next_page.data) - if item_result.has_next_page: - result.has_next_page = True - - if len(result.data) > limit: - result.data = result.data[:limit] - else: - result_limit = limit - for item in self.items: - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - - result.data.extend(item_result.data) - result_limit -= len(item_result.data) - - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True - - result.next_page.data.update(item_result.next_page.data) - - if result_limit <= 0: - break - - # Распределяем данные равномерно по ключу. - result.data = await self._uniform_distribute(result.data) - return result - - -class MergerDeduplication(BaseFeedConfigModel): - """Merger that deduplicates while preserving child mixing/position semantics. - - This merger acts as a wrapper around exactly one child feed node. - Deduplication is applied at the leaf SubFeed method level with a shared seen-set. - This lets nested mergers (positional/percentage/gradient/etc.) keep their slot rules: - duplicates are skipped by fetching additional items from the *same* leaf source. - """ - - merger_id: str - type: Literal["merger_deduplication"] - data: FeedTypes - - dedup_key: Optional[str] = None - missing_key_policy: Literal["error", "keep", "drop"] = "error" - - state_backend: Literal["cursor", "redis"] = "cursor" - state_ttl_seconds: int = 3600 - cursor_compress: bool = True - cursor_max_keys: Optional[int] = None - - overfetch_factor: int = 1 - - max_refill_loops: int = 20 - - _descendant_cursor_keys_cache: Optional[set[str]] = PrivateAttr(default=None) - - @model_validator(mode="after") - def validate_merger_deduplication(self) -> "MergerDeduplication": - if self.overfetch_factor < 1: - raise ValueError('"overfetch_factor" must be >= 1') - if self.max_refill_loops < 1: - raise ValueError('"max_refill_loops" must be >= 1') - return self - - def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: - keys: set[str] = set() - - subfeed_id = getattr(feed, "subfeed_id", None) - if isinstance(subfeed_id, str) and subfeed_id: - keys.add(subfeed_id) - - merger_id = getattr(feed, "merger_id", None) - if isinstance(merger_id, str) and merger_id: - keys.add(merger_id) - - # Recurse into known child containers across existing feed types. - child: Any - for attr_name in ("data", "positional", "default"): - child = getattr(feed, attr_name, None) - if isinstance(child, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(child)) - - for attr_name in ("item_from", "item_to"): - child = getattr(feed, attr_name, None) - inner = getattr(child, "data", None) - if isinstance(inner, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(inner)) - - items = getattr(feed, "items", None) - if isinstance(items, list): - for item in items: - if isinstance(item, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(item)) - continue - - inner = getattr(item, "data", None) - if isinstance(inner, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(inner)) - - return keys - - def _get_descendant_cursor_keys_cached(self) -> set[str]: - cached = self._descendant_cursor_keys_cache - if cached is None: - cached = self._collect_descendant_cursor_keys(self.data) - self._descendant_cursor_keys_cache = cached - return cached - - def _reset_descendant_cursors(self, next_page: FeedResultNextPage) -> None: - descendant_keys = self._get_descendant_cursor_keys_cached() - for key in descendant_keys: - next_page.data.pop(key, None) - - def _normalize_key(self, value: Any) -> str: - if isinstance(value, (str, int)): - return str(value) - if isinstance(value, (dict, list)): - return json.dumps(value, sort_keys=True, default=str) - return str(value) - - def _extract_dedup_value(self, item: Any) -> Any: - if not self.dedup_key: - return item - - try: - value = item.get(self.dedup_key) - except AttributeError: - value = getattr(item, self.dedup_key, None) - - if value is None and self.missing_key_policy == "error": - raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}") - return value - - def _get_entity_key(self, entity: Any) -> Optional[str]: - """Return normalized dedup key for entity, or None if entity should be skipped.""" - - raw_value = self._extract_dedup_value(entity) - if raw_value is None: - if self.missing_key_policy == "drop": - return None - if self.missing_key_policy == "keep": - raw_value = ("__missing__", id(entity)) - return self._normalize_key(raw_value) - - def _compute_overfetch_params(self, *, remaining: int, next_after: Any) -> tuple[bool, int, Optional[int]]: - """Compute safe overfetch params. - - Overfetch is only safe when `after` is an integer offset (so we can rewind). - - Returns: (can_overfetch, request_limit, start_after) - """ - - can_overfetch = isinstance(next_after, int) - request_limit = max(1, remaining) - if can_overfetch and self.overfetch_factor > 1: - request_limit = max(1, remaining * self.overfetch_factor) - start_after: Optional[int] = int(next_after) if can_overfetch else None - return can_overfetch, request_limit, start_after - - def _iter_subfeeds(self, feed: BaseFeedConfigModel) -> Iterator["SubFeed"]: - if isinstance(feed, SubFeed): - yield feed - return - - for attr_name in ("data", "positional", "default"): - inner = getattr(feed, attr_name, None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - for attr_name in ("item_from", "item_to"): - wrapper = getattr(feed, attr_name, None) - inner = getattr(wrapper, "data", None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - items = getattr(feed, "items", None) - if isinstance(items, list): - for item in items: - if isinstance(item, BaseFeedConfigModel): - yield from self._iter_subfeeds(item) - continue - inner = getattr(item, "data", None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - def _register_wrapped_subfeed_method( - self, - *, - subfeed: "SubFeed", - original_methods_dict: Dict[str, Callable], - rewritten_methods_dict: Dict[str, Callable], - dedup_state: "_DedupState", - ) -> None: - original_name = subfeed.method_name - original_method = original_methods_dict[original_name] - unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" - - # Idempotency: if the same subfeed id appears multiple times, don't re-wrap. - if unique_name in rewritten_methods_dict: - subfeed.method_name = unique_name - return - - subfeed.method_name = unique_name - leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) - - wrapped = self._make_wrapped_leaf_method( - original_method=original_method, - dedup_state=dedup_state, - leaf_priority=leaf_priority, - ) - setattr(wrapped, "_smartfeed_original", original_method) - rewritten_methods_dict[unique_name] = wrapped - - def _make_wrapped_leaf_method( - self, - *, - original_method: Callable, - dedup_state: "_DedupState", - leaf_priority: int, - ) -> Callable: - async def _wrapped_method( - user_id: Any, - limit: int, - next_page: FeedResultNextPageInside, - **kw: Any, - ) -> FeedResultClient: - collected: List[Any] = [] - upstream_has_next_page = False - - loops = 0 - while len(collected) < limit and loops < self.max_refill_loops: - loops += 1 - before_len = len(collected) - - remaining = limit - len(collected) - can_overfetch, request_limit, start_after = self._compute_overfetch_params( - remaining=remaining, - next_after=next_page.after, - ) - - method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) - if not isinstance(method_result, FeedResultClient): - raise TypeError('SubFeed function must return "FeedResultClient" instance.') - - upstream_has_next_page = upstream_has_next_page or method_result.has_next_page - - inspected_count = 0 - - # Backend-specific optimization: Redis batches zmscore. - # For cursor backend, prefetch is a no-op and we avoid the extra pass entirely. - keys_by_index: Optional[List[Optional[str]]] = None - if isinstance(dedup_state, _RedisDedupState): - keys_by_index = [] - batch_keys: List[str] = [] - for entity in method_result.data: - key = self._get_entity_key(entity) - keys_by_index.append(key) - if key is not None: - batch_keys.append(key) - await dedup_state.prefetch(batch_keys) - - for idx, entity in enumerate(method_result.data, start=1): - inspected_count = idx - - key = keys_by_index[idx - 1] if keys_by_index is not None else self._get_entity_key(entity) - if key is None: - continue - - if not dedup_state.should_accept(key, leaf_priority): - continue - - collected.append(entity) - dedup_state.record(key, leaf_priority) - - if len(collected) >= limit: - break - - if len(collected) == before_len: - # No progress this loop. Stop if upstream is exhausted. - if not method_result.has_next_page: - break - - # If we oversampled with a simple integer cursor, rewind to the point we actually consumed. - # This prevents skipping un-inspected items that were fetched but not needed. - if can_overfetch and request_limit > remaining and start_after is not None: - end_after = next_page.after - if isinstance(end_after, int) and end_after == start_after + len(method_result.data): - next_page.after = start_after + inspected_count - - return FeedResultClient(data=collected, next_page=next_page, has_next_page=upstream_has_next_page) - - return _wrapped_method - - def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> Dict[str, int]: - entry = next_page.data.get(self.merger_id) - if not entry or entry.after is None: - return {} - - after = entry.after - if isinstance(after, dict) and "z" in after: - payload = base64.urlsafe_b64decode(after["z"].encode()) - raw = zlib.decompress(payload).decode() - decoded = json.loads(raw) - if isinstance(decoded, dict): - return {str(k): int(v) for k, v in decoded.items()} - if isinstance(decoded, list): - # v2: list of [key, priority] entries - seen_map: Dict[str, int] = {} - for entry_item in decoded: - if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: - seen_map[str(entry_item[0])] = int(entry_item[1]) - else: - seen_map[str(entry_item)] = 0 - return seen_map - return {} - if isinstance(after, dict) and "seen" in after: - return {str(k): 0 for k in list(after["seen"])} - if isinstance(after, list): - return {str(k): 0 for k in list(after)} - if isinstance(after, dict): - # v2 uncompressed map - return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} - return {} - - def _encode_seen_for_cursor(self, seen_updates_in_order: List[tuple[str, int]]) -> Any: - if self.cursor_max_keys is not None: - seen_updates_in_order = seen_updates_in_order[-self.cursor_max_keys :] - - if not self.cursor_compress: - return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} - - raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() - compressed = zlib.compress(raw) - return { - "v": 2, - "c": "zlib+base64", - "n": len(seen_updates_in_order), - "z": base64.urlsafe_b64encode(compressed).decode(), - } - - async def _redis_zmscore( - self, - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - members: List[str], - ) -> List[Optional[float]]: - """Batch zscore for multiple members. - - Falls back to pipelined zscore when zmscore isn't available. - """ - - if not members: - return [] - - zmscore_fn = getattr(redis_client, "zmscore", None) - if zmscore_fn is not None: - res = zmscore_fn(key, members) - if inspect.iscoroutine(res): - res = await res - # redis-py returns list[Optional[float]] - return [None if v is None else float(v) for v in list(res)] - - # Fallback: pipelined zscore. For sync clients, run the whole pipeline in a thread. - if not _is_async_redis_client(redis_client): - def _sync_pipeline_execute() -> Any: - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - return pipe.execute() - - res = await asyncio.to_thread(_sync_pipeline_execute) - return [None if v is None else float(v) for v in list(res)] - - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - res = pipe.execute() - if inspect.iscoroutine(res): - res = await res - return [None if v is None else float(v) for v in list(res)] - - async def _redis_zadd_and_expire( - self, - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - member_scores: Dict[str, int], - ) -> None: - if not member_scores: - return - await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) - await _redis_call(redis_client, "expire", key, self.state_ttl_seconds) - - def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: - suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") - if suffix: - return f"dedup:{self.merger_id}:{user_id}:{suffix}" - return f"dedup:{self.merger_id}:{user_id}" - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - if limit <= 0: - return FeedResult(data=[], next_page=next_page, has_next_page=False) - - # Treat an explicit "page 0" (or missing cursor for this merger) as a fresh session. - # This allows clients to restart the feed (e.g., full reload) without carrying over seen state. - entry = next_page.data.get(self.merger_id) - requested_page = entry.page if entry is not None else None - is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) - - if self.state_backend == "redis" and not redis_client: - raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") - - working_next_page = _pydantic_deep_copy(next_page) - - if is_fresh_session: - # Reset cursors for all descendants under this merger so upstream nodes also restart. - self._reset_descendant_cursors(working_next_page) - - # Shared dedup state (cross-page) - seen_priority_map: Dict[str, int] = {} - seen_updates_in_order: List[tuple[str, int]] = [] - if self.state_backend == "cursor" and not is_fresh_session: - seen_priority_map = self._decode_seen_from_cursor(next_page) - - # Always maintain a per-request seen set to prevent duplicates within a single get_data() call. - seen_request_set: set[str] = set(seen_priority_map.keys()) - - redis_state_key = "" - redis_new_scores: Dict[str, int] = {} - redis_seen_cache: Dict[str, Optional[int]] = {} - if self.state_backend == "redis" and redis_client: - redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) - if is_fresh_session: - # Drop state for a full restart. - await _redis_call(redis_client, "delete", redis_state_key) - - # Create a single state helper shared across all leaf wrappers. - if self.state_backend == "cursor": - dedup_state: _DedupState = _CursorDedupState( - seen_priority_map=seen_priority_map, - seen_updates_in_order=seen_updates_in_order, - seen_request_set=seen_request_set, - ) - else: - assert redis_client is not None - dedup_state = _RedisDedupState( - redis_client=redis_client, - redis_state_key=redis_state_key, - redis_seen_cache=redis_seen_cache, - redis_new_scores=redis_new_scores, - seen_request_set=seen_request_set, - zmscore=self._redis_zmscore, - ) - - # Preserve inner merger ordering/mixing semantics by deduplicating at the leaf method level - # with a shared seen-set. - original_methods_dict = methods_dict - - # Create a deep copy of the child tree and rewrite each SubFeed to call a unique wrapper - # so we can associate a dedup_priority with each leaf. - child = self.data - child = _pydantic_deep_copy(child) - - rewritten_methods_dict = dict(original_methods_dict) - - for sf in self._iter_subfeeds(child): - self._register_wrapped_subfeed_method( - subfeed=sf, - original_methods_dict=original_methods_dict, - rewritten_methods_dict=rewritten_methods_dict, - dedup_state=dedup_state, - ) - - child_result = await child.get_data( - methods_dict=rewritten_methods_dict, - user_id=user_id, - limit=limit, - next_page=working_next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - - if self.state_backend == "redis" and redis_client: - await self._redis_zadd_and_expire(redis_client, redis_state_key, redis_new_scores) - - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - merger_after: Any = None - if self.state_backend == "cursor": - merger_after = self._encode_seen_for_cursor(seen_updates_in_order) - - result_next_page = _pydantic_deep_copy(child_result.next_page) - result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) - - return FeedResult(data=child_result.data, next_page=result_next_page, has_next_page=child_result.has_next_page) - - -class SubFeed(BaseFeedConfigModel): - """ - Модель субфида. - - Attributes: - subfeed_id уникальный ID субфида. - type тип объекта - всегда "subfeed". - method_name название клиентского метода для получения данных субфида. - subfeed_params статичные параметры для метода субфида. - shuffle флаг для перемешивания полученных данных мерджера. - """ - - subfeed_id: str - type: Literal["subfeed"] - method_name: str - subfeed_params: Dict[str, Any] = {} - raise_error: Optional[bool] = True - shuffle: bool = False - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - **params: Any, - ) -> FeedResult: - """ - Метод для получения данных из метода субфида. - - :param methods_dict: словарь с используемыми методами. - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param redis_client: объект клиента Redis (для конфигурации с view_session мерджером). - :param params: параметры для метода. - :return: список данных. - """ - - # Формируем next_page конкретного субфида. - subfeed_next_page = FeedResultNextPageInside( - page=next_page.data[self.subfeed_id].page if self.subfeed_id in next_page.data else 1, - after=next_page.data[self.subfeed_id].after if self.subfeed_id in next_page.data else None, - ) - - # Формируем params для функции субфида. - method = methods_dict[self.method_name] - method_spec = getattr(method, "_smartfeed_original", method) - method_args = inspect.getfullargspec(method_spec).args - method_params: Dict[str, Any] = {} - for arg in method_args: - if arg in params: - method_params[arg] = params[arg] - - # Получаем результат функции клиента в формате SubFeedResult. - try: - method_result = await methods_dict[self.method_name]( - user_id=user_id, - limit=limit, - next_page=subfeed_next_page, - **method_params, - **self.subfeed_params, - ) - except (Exception,) as _: - if self.raise_error: - raise - - method_result = FeedResultClient( - data=[], - next_page=subfeed_next_page, - has_next_page=False, - ) - - if not isinstance(method_result, FeedResultClient): - raise TypeError('SubFeed function must return "FeedResultClient" instance.') - - # Если в конфигурации указано "смешать" данные. - if self.shuffle: - shuffle(method_result.data) - - result = FeedResult( - data=method_result.data, - next_page=FeedResultNextPage(data={self.subfeed_id: method_result.next_page}), - has_next_page=method_result.has_next_page, - ) - return result - - class FeedConfig(BaseModel): - """ - Модель конфигурации фида. - - Attributes: - version версия конфигурации. - view_session флаг использования механизма расчета всего фида сразу и сохранения в кэш. - session_size размер кэшируемого фида (limit получения данных для сохранения в кэш). - session_live_time срок хранения в кэше для кэшируемого фида (в секундах). - feed мерджер или субфид. - """ + """Top-level feed config model.""" version: str feed: FeedTypes -# Update Forward Refs def _rebuild_model(model: Any) -> None: + """Resolve forward refs across modules (Pydantic v1/v2 compatible).""" + if hasattr(model, "model_rebuild"): - model.model_rebuild() + model.model_rebuild(force=True, _types_namespace={"FeedTypes": FeedTypes}) else: - model.update_forward_refs() - - -_rebuild_model(MergerPositional) -_rebuild_model(MergerPercentage) -_rebuild_model(SubFeed) -_rebuild_model(MergerPercentageItem) -_rebuild_model(MergerAppend) -_rebuild_model(MergerAppendDistribute) -_rebuild_model(MergerPercentageGradient) -_rebuild_model(MergerViewSession) -_rebuild_model(MergerDeduplication) + model.update_forward_refs(FeedTypes=FeedTypes) + + +for _m in ( + MergerPositional, + MergerPercentage, + MergerPercentageItem, + MergerAppend, + MergerAppendDistribute, + MergerPercentageGradient, + MergerViewSession, + MergerDeduplication, + SubFeed, + FeedConfig, +): + _rebuild_model(_m) + + +__all__ = [ + "BaseFeedConfigModel", + "FeedResult", + "FeedResultClient", + "FeedResultNextPage", + "FeedResultNextPageInside", + "SubFeed", + "MergerAppend", + "MergerAppendDistribute", + "MergerDeduplication", + "MergerPercentage", + "MergerPercentageGradient", + "MergerPercentageItem", + "MergerPositional", + "MergerViewSession", + "FeedTypes", + "FeedConfig", +] From e3fdefbc9938e24364b00a244774450317a3de5d Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sat, 7 Feb 2026 13:05:35 +0000 Subject: [PATCH 15/41] Refactor continues. --- smartfeed/execution/context.py | 38 ++ smartfeed/execution/cursors.py | 63 +++ smartfeed/execution/executor.py | 616 +++++++++++++++++++++++ smartfeed/execution/plans.py | 60 +++ smartfeed/feed_models.py | 12 +- smartfeed/manager.py | 13 +- smartfeed/mergers/append.py | 110 ++-- smartfeed/mergers/append_distribute.py | 92 ++-- smartfeed/mergers/deduplication.py | 509 ++++--------------- smartfeed/mergers/percentage.py | 96 ++-- smartfeed/mergers/percentage_gradient.py | 169 +++---- smartfeed/mergers/positional.py | 137 ++--- smartfeed/mergers/view_session.py | 147 ++++-- smartfeed/policies/dedup.py | 220 ++++++++ smartfeed/policies/dedup_utils.py | 112 +++++ smartfeed/policies/seen_store.py | 158 ++++++ tests/fixtures/redis.py | 4 +- tests/test_manager_params.py | 40 ++ tests/test_merger_deduplication.py | 74 ++- 19 files changed, 1875 insertions(+), 795 deletions(-) create mode 100644 smartfeed/execution/context.py create mode 100644 smartfeed/execution/cursors.py create mode 100644 smartfeed/execution/executor.py create mode 100644 smartfeed/execution/plans.py create mode 100644 smartfeed/policies/dedup.py create mode 100644 smartfeed/policies/dedup_utils.py create mode 100644 smartfeed/policies/seen_store.py create mode 100644 tests/test_manager_params.py diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py new file mode 100644 index 0000000..4a118ea --- /dev/null +++ b/smartfeed/execution/context.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Union + +import redis +from redis.asyncio import Redis as AsyncRedis + + +@dataclass +class ExecutionContext: + """Execution context propagated through the feed tree. + + Keeps internal state (policies, backends) out of user params. + """ + + methods_dict: Dict[str, Callable] + user_id: Any + redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None + + # Assigned by the caller (FeedManager / tests) to avoid circular imports. + executor: Any = None + + # Policies (optional) + dedup: Optional[object] = None + + # Execution settings (optional) + refill_settings: Optional["RefillExecutionSettings"] = None + dedup_settings: Optional["DedupExecutionSettings"] = None + + +@dataclass(frozen=True) +class RefillExecutionSettings: + overfetch_factor: int = 1 + max_refill_loops: int = 20 + + +DedupExecutionSettings = RefillExecutionSettings diff --git a/smartfeed/execution/cursors.py b/smartfeed/execution/cursors.py new file mode 100644 index 0000000..eef2fd0 --- /dev/null +++ b/smartfeed/execution/cursors.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable + +from ..feed_models import BaseFeedConfigModel, FeedResultNextPage + + +@dataclass +class CursorMap: + next_page: FeedResultNextPage + + def merge_delta(self, *, base_next_page: FeedResultNextPage, owner_next_page: FeedResultNextPage) -> None: + """Merge only the cursor keys that actually changed.""" + + for key, value in owner_next_page.data.items(): + base_value = base_next_page.data.get(key) + if base_value == value: + continue + self.next_page.data[key] = value + + def reset_keys(self, keys: Iterable[str]) -> None: + for key in keys: + self.next_page.data.pop(key, None) + + @staticmethod + def can_overfetch(*, node: BaseFeedConfigModel, base_next_page: FeedResultNextPage) -> bool: + sub_id = getattr(node, "subfeed_id", None) + if not isinstance(sub_id, str): + return False + entry = base_next_page.data.get(sub_id) + if entry is None: + return False + return isinstance(entry.after, int) + + @staticmethod + def rewind_overfetch( + *, + node: BaseFeedConfigModel, + base_next_page: FeedResultNextPage, + result_next_page: FeedResultNextPage, + inspected_count: int, + batch_size: int, + ) -> None: + sub_id = getattr(node, "subfeed_id", None) + if not isinstance(sub_id, str): + return + if sub_id not in result_next_page.data: + return + + entry = result_next_page.data[sub_id] + end_after = entry.after + if not isinstance(end_after, int): + return + + base_entry = base_next_page.data.get(sub_id) + prev_after = base_entry.after if base_entry is not None else None + if not isinstance(prev_after, int): + return + + expected_end = prev_after + batch_size + if end_after == expected_end: + entry.after = prev_after + inspected_count diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py new file mode 100644 index 0000000..20e3cbc --- /dev/null +++ b/smartfeed/execution/executor.py @@ -0,0 +1,616 @@ +from __future__ import annotations + +import asyncio +import inspect +from typing import Any, Dict, List, Optional, Tuple + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, _pydantic_deep_copy +from .context import ExecutionContext +from .cursors import CursorMap +from .plans import CallablePlan, Plan, SlotSpec, SlotsPlan + + +class Executor: + """Shared execution engine. + + Owns recursion and concurrency. Nodes can optionally expose `build_plan(...)`. + """ + + async def run( + self, + node: BaseFeedConfigModel, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> FeedResult: + result, plan = await self._run_node_raw(node, ctx, limit, next_page, params) + + dedup = getattr(ctx, "dedup", None) + if dedup is None: + return result + + if isinstance(plan, SlotsPlan): + return result + + return await self._run_node_with_dedup_refill(node, ctx, limit, next_page, params, result) + + async def execute_plan(self, plan: Plan) -> FeedResult: + """Interpret and execute a declarative plan. + + Plans must not perform execution themselves; they are data structures. + """ + + if isinstance(plan, SlotsPlan): + return await self._execute_slots_plan(plan) + if isinstance(plan, CallablePlan): + return await plan.fn(self) + raise TypeError(f"Unknown plan type: {type(plan)!r}") + + async def gather(self, *coros: Any) -> List[Any]: + """Execute coroutines concurrently. + + Centralizes concurrency in the executor layer. + """ + + return list(await asyncio.gather(*coros)) + + async def _maybe_await(self, value: Any) -> Any: + if inspect.isawaitable(value): + return await value + return value + + async def _run_node_raw( + self, + node: BaseFeedConfigModel, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + params: Dict[str, Any], + ) -> Tuple[FeedResult, Optional[Plan]]: + build_plan = getattr(node, "build_plan", None) + if callable(build_plan): + plan: Plan = build_plan(ctx=ctx, limit=limit, next_page=next_page, **params) + result = await self.execute_plan(plan) + return result, plan + + result = await node.get_data( + methods_dict=ctx.methods_dict, + user_id=ctx.user_id, + limit=limit, + next_page=next_page, + redis_client=ctx.redis_client, + ctx=ctx, + **params, + ) + return result, None + + async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: + if plan.limit <= 0: + assembled = await self._maybe_await(plan.assemble([], plan.next_page, {})) + return assembled + + working_next_page = _pydantic_deep_copy(plan.next_page) + cursor = CursorMap(working_next_page) + owners, owner_index = self._collect_plan_owners(plan) + dedup_policy = getattr(plan.ctx, "dedup", None) + refill_settings = getattr(plan.ctx, "refill_settings", None) or getattr(plan.ctx, "dedup_settings", None) + dedup_active = dedup_policy is not None + + owner_max_demand = self._owner_slot_demand(plan) + owner_buffers, owner_results = await self._run_plan_owners( + plan=plan, + owners=owners, + owner_max_demand=owner_max_demand, + dedup_active=dedup_active, + cursor=cursor, + ) + + if dedup_policy is not None: + owner_buffers, owner_results = await self._arbitrate_owner_buffers( + owners=owners, + owner_index=owner_index, + owner_buffers=owner_buffers, + owner_results=owner_results, + dedup_policy=dedup_policy, + ) + + deficits = self._compute_slot_deficits( + plan=plan, + owner_buffers=owner_buffers, + ) + if deficits: + await self._refill_deficits( + plan=plan, + deficits=deficits, + owners=owners, + owner_index=owner_index, + owner_buffers=owner_buffers, + owner_results=owner_results, + dedup_policy=dedup_policy, + refill_settings=refill_settings, + cursor=cursor, + ) + + output = self._consume_slots(plan=plan, owner_buffers=owner_buffers) + assembled = await self._maybe_await(plan.assemble(output, cursor.next_page, owner_results)) + return assembled + + def _owner_slot_demand(self, plan: SlotsPlan) -> Dict[int, int]: + """Compute a per-owner maximum demand based on the slot schedule.""" + + demand: Dict[int, int] = {} + for slot in plan.slots: + owner_id = id(slot.owner) + demand[owner_id] = demand.get(owner_id, 0) + int(slot.max_count) + return demand + + def _collect_plan_owners(self, plan: SlotsPlan) -> tuple[List[Any], Dict[int, int]]: + owners: List[Any] = [] + owner_index: Dict[int, int] = {} + for slot in plan.slots: + owner_id = id(slot.owner) + if owner_id in owner_index: + continue + owner_index[owner_id] = len(owners) + owners.append(slot.owner) + return owners, owner_index + + async def _run_owner( + self, + *, + plan: SlotsPlan, + owner: Any, + demand: int, + base_next_page: FeedResultNextPage, + dedup_active: bool, + ) -> FeedResult: + isolated_next_page = _pydantic_deep_copy(base_next_page) + owner_ctx = plan.ctx + if dedup_active: + owner_ctx = ExecutionContext( + methods_dict=plan.ctx.methods_dict, + user_id=plan.ctx.user_id, + redis_client=plan.ctx.redis_client, + executor=plan.ctx.executor, + dedup=None, + refill_settings=None, + dedup_settings=None, + ) + return await self.run(owner, owner_ctx, demand, isolated_next_page, **plan.params) + + async def _run_plan_owners( + self, + *, + plan: SlotsPlan, + owners: List[Any], + owner_max_demand: Dict[int, int], + dedup_active: bool, + cursor: CursorMap, + ) -> tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: + owner_buffers: Dict[int, List[Any]] = {id(o): [] for o in owners} + owner_results: Dict[int, FeedResult] = {} + + ops: List[tuple[Any, int]] = [] + for owner in owners: + if plan.owner_fetch_limits is not None and id(owner) in plan.owner_fetch_limits: + demand = int(plan.owner_fetch_limits[id(owner)]) + else: + demand = min(plan.limit, int(owner_max_demand.get(id(owner), 0))) + if demand > 0: + ops.append((owner, demand)) + + if not ops: + return owner_buffers, owner_results + + results = await self.gather( + *[ + self._run_owner( + plan=plan, + owner=owner, + demand=demand, + base_next_page=plan.next_page, + dedup_active=dedup_active, + ) + for owner, demand in ops + ] + ) + for (owner, _demand), owner_result in zip(ops, results): + owner_results[id(owner)] = owner_result + owner_buffers[id(owner)] = list(owner_result.data) + cursor.merge_delta( + base_next_page=plan.next_page, + owner_next_page=owner_result.next_page, + ) + + return owner_buffers, owner_results + + async def _arbitrate_owner_buffers( + self, + *, + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + dedup_policy: Any, + ) -> tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: + owner_buffers = await dedup_policy.arbitrate_owner_buffers( + owners=owners, + owner_buffers=owner_buffers, + owner_rank=owner_index, + ) + + for owner in owners: + owner_id = id(owner) + if owner_id not in owner_results: + continue + old = owner_results[owner_id] + owner_results[owner_id] = FeedResult( + data=list(owner_buffers.get(owner_id, [])), + next_page=old.next_page, + has_next_page=old.has_next_page, + ) + + return owner_buffers, owner_results + + def _compute_slot_deficits( + self, + *, + plan: SlotsPlan, + owner_buffers: Dict[int, List[Any]], + ) -> Dict[int, int]: + total_max = sum(int(s.max_count) for s in plan.slots) + quota_schedule = total_max <= int(plan.limit) + + consumed: Dict[int, int] = {} + remaining = int(plan.limit) + deficit_slots: List[int] = [] + + for slot in plan.slots: + if remaining <= 0: + break + + owner_id = id(slot.owner) + want = min(int(slot.max_count), remaining) + if want <= 0: + continue + + have_total = len(owner_buffers.get(owner_id, [])) + already = int(consumed.get(owner_id, 0)) + available = max(0, have_total - already) + take = min(want, available) + if take < want: + deficit_slots.append(owner_id) + consumed[owner_id] = already + take + remaining -= take + + page_underfilled = remaining > 0 + + if not quota_schedule and not page_underfilled: + return {} + + deficits: Dict[int, int] = {} + + if quota_schedule: + return self._compute_quota_deficits(plan=plan, owner_buffers=owner_buffers) + + return self._compute_fill_deficits( + plan=plan, + remaining=remaining, + deficit_slots=deficit_slots, + ) + + def _compute_quota_deficits( + self, + *, + plan: SlotsPlan, + owner_buffers: Dict[int, List[Any]], + ) -> Dict[int, int]: + deficits: Dict[int, int] = {} + remaining = int(plan.limit) + consumed: Dict[int, int] = {} + for slot in plan.slots: + if remaining <= 0: + break + + owner_id = id(slot.owner) + want = min(int(slot.max_count), remaining) + if want <= 0: + continue + + have_total = len(owner_buffers.get(owner_id, [])) + already = int(consumed.get(owner_id, 0)) + available = max(0, have_total - already) + take = min(want, available) + missing = max(0, want - take) + if missing: + deficits[owner_id] = deficits.get(owner_id, 0) + missing + consumed[owner_id] = already + take + remaining -= take + + return deficits + + def _compute_fill_deficits( + self, + *, + plan: SlotsPlan, + remaining: int, + deficit_slots: List[int], + ) -> Dict[int, int]: + deficits: Dict[int, int] = {} + to_fill = int(remaining) + if to_fill <= 0: + return deficits + + if deficit_slots: + deficits[deficit_slots[-1]] = deficits.get(deficit_slots[-1], 0) + to_fill + return deficits + + if plan.slots: + last_owner_id = id(plan.slots[-1].owner) + deficits[last_owner_id] = deficits.get(last_owner_id, 0) + to_fill + return deficits + + return deficits + + async def _refill_deficits( + self, + *, + plan: SlotsPlan, + deficits: Dict[int, int], + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + dedup_policy: Any, + refill_settings: Any, + cursor: CursorMap, + ) -> None: + overfetch_factor = max(1, int(getattr(refill_settings, "overfetch_factor", 1))) + max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) + + deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] + deficit_owners = sorted( + deficit_owners, + key=lambda o: ( + int(getattr(o, "dedup_priority", 0)), + owner_index.get(id(o), 0), + ), + ) + + state: Dict[int, Dict[str, Any]] = {} + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + missing_total = int(deficits.get(refill_owner_id, 0)) + if missing_total <= 0: + continue + + base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page + state[refill_owner_id] = { + "owner": refill_owner, + "missing_total": missing_total, + "remaining": int(missing_total), + "accepted": [], + "loops": 0, + "current_next_page": base_np, + "has_next_page": True, + "last_result": None, + "last_request_limit": 0, + "last_can_overfetch": False, + "last_base_next_page": base_np, + } + + if not state: + return + + while True: + wave_ops: List[Tuple[Any, int, FeedResultNextPage, int, bool]] = [] + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + if owner_state["remaining"] <= 0: + continue + if not owner_state["has_next_page"]: + continue + if owner_state["loops"] >= max_refill_loops: + continue + + base_np = owner_state["current_next_page"] + request_limit = max(1, int(owner_state["remaining"])) + can_overfetch = CursorMap.can_overfetch(node=refill_owner, base_next_page=base_np) + if can_overfetch and overfetch_factor > 1: + request_limit = max(1, int(owner_state["remaining"]) * overfetch_factor) + + wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit, can_overfetch)) + + if not wave_ops: + break + + results = await self.gather( + *[ + self._run_owner( + plan=plan, + owner=owner, + demand=request_limit, + base_next_page=base_np, + dedup_active=True, + ) + for owner, _owner_id, base_np, request_limit, _can_overfetch in wave_ops + ] + ) + + for (owner, owner_id, base_np, request_limit, can_overfetch), result in zip(wave_ops, results): + owner_state = state[owner_id] + owner_state["last_result"] = result + owner_state["last_request_limit"] = request_limit + owner_state["last_can_overfetch"] = can_overfetch + owner_state["last_base_next_page"] = base_np + owner_state["current_next_page"] = result.next_page + owner_state["has_next_page"] = bool(result.has_next_page) + + cursor.merge_delta( + base_next_page=plan.next_page, + owner_next_page=result.next_page, + ) + + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + if owner_state["remaining"] <= 0: + continue + last_result = owner_state["last_result"] + if last_result is None: + continue + + refill_prio = int(getattr(refill_owner, "dedup_priority", 0)) + wave_accepted, inspected_count = await dedup_policy.accept_batch( + items=list(last_result.data), + priority=refill_prio, + limit=int(owner_state["remaining"]), + ) + + if owner_state["last_can_overfetch"] and owner_state["last_request_limit"] > owner_state["remaining"]: + CursorMap.rewind_overfetch( + node=refill_owner, + base_next_page=owner_state["last_base_next_page"], + result_next_page=owner_state["current_next_page"], + inspected_count=inspected_count, + batch_size=len(last_result.data), + ) + + if wave_accepted: + owner_state["accepted"].extend(wave_accepted) + owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) + + if owner_state["remaining"] > 0 and owner_state["has_next_page"]: + owner_state["loops"] += 1 + + owner_state["last_result"] = None + + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + + accepted = owner_state["accepted"] + if accepted: + owner_buffers.setdefault(refill_owner_id, []) + owner_buffers[refill_owner_id].extend(accepted) + + owner_results[refill_owner_id] = FeedResult( + data=list(owner_buffers.get(refill_owner_id, [])), + next_page=owner_state["current_next_page"], + has_next_page=owner_state["has_next_page"], + ) + + def _consume_slots(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> List[Any]: + output: List[Any] = [] + for slot in plan.slots: + if len(output) >= plan.limit: + break + + remaining = plan.limit - len(output) + take = min(int(slot.max_count), remaining) + if take <= 0: + continue + + owner_buffer = owner_buffers.get(id(slot.owner), []) + if not owner_buffer: + continue + + chunk = owner_buffer[:take] + del owner_buffer[: len(chunk)] + output.extend(chunk) + + return output + + async def _run_node_with_dedup_refill( + self, + node: BaseFeedConfigModel, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + params: Dict[str, Any], + initial_result: FeedResult, + ) -> FeedResult: + dedup = getattr(ctx, "dedup", None) + if dedup is None: + return initial_result + + settings = getattr(ctx, "refill_settings", None) or getattr(ctx, "dedup_settings", None) + overfetch_factor = max(1, int(getattr(settings, "overfetch_factor", 1))) + max_refill_loops = max(1, int(getattr(settings, "max_refill_loops", 20))) + priority = int(getattr(node, "dedup_priority", 0)) + + collected: List[Any] = [] + remaining = int(limit) + loops = 0 + + current_result = initial_result + current_next_page = current_result.next_page + current_request_limit = max(1, remaining) + has_next_page = bool(current_result.has_next_page) + base_next_page = next_page + + while remaining > 0: + can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) + + accepted, inspected_count = await dedup.accept_batch( + items=list(current_result.data), + priority=priority, + limit=remaining, + ) + + if can_overfetch and current_request_limit > remaining: + CursorMap.rewind_overfetch( + node=node, + base_next_page=base_next_page, + result_next_page=current_next_page, + inspected_count=inspected_count, + batch_size=len(current_result.data), + ) + + if accepted: + collected.extend(accepted) + remaining = limit - len(collected) + + if remaining <= 0 or not has_next_page or loops >= max_refill_loops: + break + loops += 1 + + base_next_page = current_next_page + next_request_limit = max(1, remaining) + can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) + if can_overfetch and overfetch_factor > 1: + next_request_limit = max(1, remaining * overfetch_factor) + + current_result, _plan = await self._run_node_raw( + node, + ctx, + next_request_limit, + base_next_page, + params, + ) + current_next_page = current_result.next_page + current_request_limit = next_request_limit + has_next_page = bool(current_result.has_next_page) + + return FeedResult( + data=collected, + next_page=current_next_page, + has_next_page=has_next_page, + ) + + +__all__ = [ + "Executor", + "Plan", + "CallablePlan", + "SlotSpec", + "SlotsPlan", +] diff --git a/smartfeed/execution/plans.py b/smartfeed/execution/plans.py new file mode 100644 index 0000000..8a23a29 --- /dev/null +++ b/smartfeed/execution/plans.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage +from .context import ExecutionContext + + +class Plan(Protocol): + """Declarative execution plan. + + Plans describe what to run; the `Executor` is responsible for interpreting + and executing them. + """ + + +@dataclass(frozen=True) +class CallablePlan: + """A plan implemented as an async callable. + + Useful for mergers whose child limits depend on previous child results. + """ + + fn: Callable[["Executor"], Awaitable[FeedResult]] + + +@dataclass(frozen=True) +class SlotSpec: + """A slot segment owned by a child node. + + Output order is defined by the sequence of SlotSpecs. + """ + + owner: BaseFeedConfigModel + max_count: int + + +@dataclass(frozen=True) +class SlotsPlan: + """Plan expressed as slot ownership + an assembly function. + + The executor will fetch children (possibly in priority order) and then assemble + results in the slot schedule order. + """ + + ctx: ExecutionContext + limit: int + next_page: FeedResultNextPage + params: Dict[str, Any] + slots: List[SlotSpec] + assemble: Callable[[List[Any], FeedResultNextPage, Dict[int, FeedResult]], Any] + owner_fetch_limits: Optional[Dict[int, int]] = None + + +# NOTE: `Executor` is imported only for typing to avoid an import cycle. +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .executor import Executor diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py index b98f11f..50abcbe 100644 --- a/smartfeed/feed_models.py +++ b/smartfeed/feed_models.py @@ -3,13 +3,16 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from random import shuffle -from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast import redis from pydantic import BaseModel from redis.asyncio import Redis as AsyncRedis from redis.asyncio import RedisCluster as AsyncRedisCluster +if TYPE_CHECKING: + from .execution.context import ExecutionContext + def _is_async_redis_client(client: Any) -> bool: return isinstance(client, (AsyncRedis, AsyncRedisCluster)) @@ -80,6 +83,7 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional["ExecutionContext"] = None, **params: Any, ) -> FeedResult: """Fetch data according to this node config.""" @@ -114,8 +118,14 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional["ExecutionContext"] = None, **params: Any, ) -> FeedResult: + if ctx is None: + from .execution.context import ExecutionContext as _ExecutionContext + + ctx = _ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + subfeed_next_page = FeedResultNextPageInside( page=next_page.data[self.subfeed_id].page if self.subfeed_id in next_page.data else 1, after=next_page.data[self.subfeed_id].after if self.subfeed_id in next_page.data else None, diff --git a/smartfeed/manager.py b/smartfeed/manager.py index 5ef6eb1..f42e95f 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -3,6 +3,8 @@ import redis from redis.asyncio import Redis as AsyncRedis +from .execution.context import ExecutionContext +from .execution.executor import Executor from .schemas import FeedConfig, FeedResult, FeedResultNextPage @@ -39,12 +41,7 @@ async def get_data(self, user_id: Any, limit: int, next_page: FeedResultNextPage :return: результат получения данных согласно конфигурации фида. """ - result = await self.feed_config.feed.get_data( - methods_dict=self.methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=self.redis_client, - **params, - ) + ctx = ExecutionContext(methods_dict=self.methods_dict, user_id=user_id, redis_client=self.redis_client) + ctx.executor = Executor() + result = await ctx.executor.run(self.feed_config.feed, ctx, limit, next_page, **params) return result diff --git a/smartfeed/mergers/append.py b/smartfeed/mergers/append.py index bfe0718..e4fc609 100644 --- a/smartfeed/mergers/append.py +++ b/smartfeed/mergers/append.py @@ -1,14 +1,14 @@ from __future__ import annotations -import asyncio -from collections import defaultdict from random import shuffle -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast import redis from redis.asyncio import Redis as AsyncRedis -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, _pydantic_deep_copy +from ..execution.context import ExecutionContext +from ..execution.executor import SlotSpec, SlotsPlan +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage if TYPE_CHECKING: from ..schemas import FeedTypes @@ -22,6 +22,34 @@ class MergerAppend(BaseFeedConfigModel): items: List[FeedTypes] shuffle: bool = False + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> SlotsPlan: + slots = [SlotSpec(owner=cast(BaseFeedConfigModel, item), max_count=limit) for item in self.items] + + def _assemble( + output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] + ) -> FeedResult: + has_next_page = any(r.has_next_page for r in owner_results.values()) + result = FeedResult(data=output, next_page=merged_next_page, has_next_page=has_next_page) + if self.shuffle: + shuffle(result.data) + return result + + return SlotsPlan( + ctx=ctx, + limit=limit, + next_page=next_page, + params=dict(params), + slots=slots, + assemble=_assemble, + ) + async def get_data( self, methods_dict: Dict[str, Callable], @@ -29,75 +57,15 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - - if dedup_active: - indexed_items = list(enumerate(self.items)) - fetched: Dict[int, FeedResult] = {} - - groups: Dict[int, List[tuple[int, "FeedTypes"]]] = defaultdict(list) - for idx, item in indexed_items: - prio = int(getattr(item, "dedup_priority", 0)) - groups[prio].append((idx, item)) - - for prio in sorted(groups.keys(), reverse=True): - group = groups[prio] - coros: List[Awaitable[FeedResult]] = [] - order: List[int] = [] - for idx, item in group: - order.append(idx) - coros.append( - item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=_pydantic_deep_copy(next_page), - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - ) - group_results = await asyncio.gather(*coros) - for idx, r in zip(order, group_results): - fetched[idx] = cast(FeedResult, r) - - for idx, _item in indexed_items: - item_result = fetched[idx] - result.data.extend(item_result.data) - result.next_page.data.update(item_result.next_page.data) - if item_result.has_next_page: - result.has_next_page = True - - if len(result.data) > limit: - result.data = result.data[:limit] - else: - result_limit = limit - for item in self.items: - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - - result.data.extend(item_result.data) - result_limit -= len(item_result.data) - - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True - - result.next_page.data.update(item_result.next_page.data) + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - if result_limit <= 0: - break + if ctx.executor is None: + from ..execution.executor import Executor - if self.shuffle: - shuffle(result.data) + ctx.executor = Executor() - return result + return await ctx.executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py index 5f07d66..442ee00 100644 --- a/smartfeed/mergers/append_distribute.py +++ b/smartfeed/mergers/append_distribute.py @@ -7,6 +7,8 @@ from redis.asyncio import Redis as AsyncRedis from typing_extensions import no_type_check +from ..execution.context import ExecutionContext +from ..execution.executor import SlotSpec, SlotsPlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage if TYPE_CHECKING: @@ -46,6 +48,32 @@ async def _uniform_distribute(self, data: list) -> list: return result + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> SlotsPlan: + slots = [SlotSpec(owner=item, max_count=limit) for item in self.items] + + async def _assemble( + output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] + ) -> FeedResult: + has_next_page = any(r.has_next_page for r in owner_results.values()) + distributed = await self._uniform_distribute(output) + return FeedResult(data=distributed, next_page=merged_next_page, has_next_page=has_next_page) + + return SlotsPlan( + ctx=ctx, + limit=limit, + next_page=next_page, + params=dict(params), + slots=slots, + assemble=_assemble, + ) + async def get_data( self, methods_dict: Dict[str, Callable], @@ -53,59 +81,15 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) - - if dedup_active: - indexed_items = list(enumerate(self.items)) - fetch_order = sorted(indexed_items, key=lambda p: (getattr(p[1], "dedup_priority", 0), -p[0]), reverse=True) - fetched: Dict[int, FeedResult] = {} - - for idx, item in fetch_order: - fetched[idx] = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - - for idx, _item in indexed_items: - item_result = fetched[idx] - result.data.extend(item_result.data) - result.next_page.data.update(item_result.next_page.data) - if item_result.has_next_page: - result.has_next_page = True - - if len(result.data) > limit: - result.data = result.data[:limit] - else: - result_limit = limit - for item in self.items: - item_result = await item.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=result_limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - - result.data.extend(item_result.data) - result_limit -= len(item_result.data) - - if not result.has_next_page and item_result.has_next_page: - result.has_next_page = True - - result.next_page.data.update(item_result.next_page.data) - - if result_limit <= 0: - break - - result.data = await self._uniform_distribute(result.data) - return result + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + + if ctx.executor is None: + from ..execution.executor import Executor + + ctx.executor = Executor() + + return await ctx.executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/deduplication.py b/smartfeed/mergers/deduplication.py index a9902a5..94c9841 100644 --- a/smartfeed/mergers/deduplication.py +++ b/smartfeed/mergers/deduplication.py @@ -1,120 +1,28 @@ from __future__ import annotations -import asyncio -import base64 -import inspect -import zlib -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterator, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Union import redis from pydantic import PrivateAttr, model_validator from redis.asyncio import Redis as AsyncRedis -from .. import jsonlib as json +from ..execution.context import ExecutionContext, RefillExecutionSettings +from ..execution.cursors import CursorMap +from ..execution.executor import CallablePlan from ..feed_models import ( BaseFeedConfigModel, FeedResult, - FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, - SubFeed, - _is_async_redis_client, _pydantic_deep_copy, - _redis_call, ) +from ..policies.dedup import DeduplicationPolicy +from ..policies.seen_store import CursorSeenStore, RedisSeenStore, SeenStore if TYPE_CHECKING: from ..schemas import FeedTypes -class _DedupState(ABC): - @abstractmethod - def should_accept(self, key: str, priority: int) -> bool: - raise NotImplementedError - - @abstractmethod - def record(self, key: str, priority: int) -> None: - raise NotImplementedError - - async def prefetch(self, keys: List[str]) -> None: - return - - -@dataclass -class _CursorDedupState(_DedupState): - seen_priority_map: Dict[str, int] - seen_updates_in_order: List[tuple[str, int]] - seen_request_set: set[str] - - def should_accept(self, key: str, priority: int) -> bool: - if key in self.seen_request_set: - return False - existing_priority = self.seen_priority_map.get(key) - if existing_priority is not None and priority <= existing_priority: - return False - return True - - def record(self, key: str, priority: int) -> None: - self.seen_priority_map[key] = priority - self.seen_updates_in_order.append((key, priority)) - self.seen_request_set.add(key) - - -@dataclass -class _RedisDedupState(_DedupState): - redis_client: Union[redis.Redis, AsyncRedis] - redis_state_key: str - redis_seen_cache: Dict[str, Optional[int]] - redis_new_scores: Dict[str, int] - seen_request_set: set[str] - zmscore: Callable[ - [Union[redis.Redis, AsyncRedis], str, List[str]], - Union[Awaitable[List[Optional[float]]], List[Optional[float]]], - ] - - async def prefetch(self, keys: List[str]) -> None: - if not keys: - return - unique: List[str] = [] - seen: set[str] = set() - for k in keys: - if k in self.seen_request_set: - continue - if k in self.redis_seen_cache: - continue - if k in seen: - continue - seen.add(k) - unique.append(k) - - if not unique: - return - - scores_result = self.zmscore(self.redis_client, self.redis_state_key, unique) - if inspect.iscoroutine(scores_result): - scores = await cast(Awaitable[List[Optional[float]]], scores_result) - else: - scores = cast(List[Optional[float]], scores_result) - - for k, s in zip(unique, scores): - self.redis_seen_cache[k] = None if s is None else int(s) - - def should_accept(self, key: str, priority: int) -> bool: - if key in self.seen_request_set: - return False - existing_priority = self.redis_seen_cache.get(key) - if existing_priority is not None and priority <= existing_priority: - return False - return True - - def record(self, key: str, priority: int) -> None: - self.seen_request_set.add(key) - self.redis_seen_cache[key] = priority - self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) - - class MergerDeduplication(BaseFeedConfigModel): """Merger that deduplicates while preserving child mixing/position semantics.""" @@ -189,264 +97,7 @@ def _get_descendant_cursor_keys_cached(self) -> set[str]: def _reset_descendant_cursors(self, next_page: FeedResultNextPage) -> None: descendant_keys = self._get_descendant_cursor_keys_cached() - for key in descendant_keys: - next_page.data.pop(key, None) - - def _normalize_key(self, value: Any) -> str: - if isinstance(value, (str, int)): - return str(value) - if isinstance(value, (dict, list)): - return json.dumps(value, sort_keys=True, default=str) - return str(value) - - def _extract_dedup_value(self, item: Any) -> Any: - if not self.dedup_key: - return item - - try: - value = item.get(self.dedup_key) - except AttributeError: - value = getattr(item, self.dedup_key, None) - - if value is None and self.missing_key_policy == "error": - raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}") - return value - - def _get_entity_key(self, entity: Any) -> Optional[str]: - raw_value = self._extract_dedup_value(entity) - if raw_value is None: - if self.missing_key_policy == "drop": - return None - if self.missing_key_policy == "keep": - raw_value = ("__missing__", id(entity)) - return self._normalize_key(raw_value) - - def _compute_overfetch_params(self, *, remaining: int, next_after: Any) -> tuple[bool, int, Optional[int]]: - can_overfetch = isinstance(next_after, int) - request_limit = max(1, remaining) - if can_overfetch and self.overfetch_factor > 1: - request_limit = max(1, remaining * self.overfetch_factor) - start_after: Optional[int] = int(next_after) if can_overfetch else None - return can_overfetch, request_limit, start_after - - def _iter_subfeeds(self, feed: BaseFeedConfigModel) -> Iterator[SubFeed]: - if isinstance(feed, SubFeed): - yield feed - return - - for attr_name in ("data", "positional", "default"): - inner = getattr(feed, attr_name, None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - for attr_name in ("item_from", "item_to"): - wrapper = getattr(feed, attr_name, None) - inner = getattr(wrapper, "data", None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - items = getattr(feed, "items", None) - if isinstance(items, list): - for item in items: - if isinstance(item, BaseFeedConfigModel): - yield from self._iter_subfeeds(item) - continue - inner = getattr(item, "data", None) - if isinstance(inner, BaseFeedConfigModel): - yield from self._iter_subfeeds(inner) - - def _register_wrapped_subfeed_method( - self, - *, - subfeed: SubFeed, - original_methods_dict: Dict[str, Callable], - rewritten_methods_dict: Dict[str, Callable], - dedup_state: _DedupState, - ) -> None: - original_name = subfeed.method_name - original_method = original_methods_dict[original_name] - unique_name = f"__dedup__{self.merger_id}__{subfeed.subfeed_id}" - - if unique_name in rewritten_methods_dict: - subfeed.method_name = unique_name - return - - subfeed.method_name = unique_name - leaf_priority = int(getattr(subfeed, "dedup_priority", 0)) - - wrapped = self._make_wrapped_leaf_method( - original_method=original_method, - dedup_state=dedup_state, - leaf_priority=leaf_priority, - ) - setattr(wrapped, "_smartfeed_original", original_method) - rewritten_methods_dict[unique_name] = wrapped - - def _make_wrapped_leaf_method( - self, - *, - original_method: Callable, - dedup_state: _DedupState, - leaf_priority: int, - ) -> Callable: - async def _wrapped_method( - user_id: Any, - limit: int, - next_page: FeedResultNextPageInside, - **kw: Any, - ) -> FeedResultClient: - collected: List[Any] = [] - upstream_has_next_page = False - - loops = 0 - while len(collected) < limit and loops < self.max_refill_loops: - loops += 1 - before_len = len(collected) - - remaining = limit - len(collected) - can_overfetch, request_limit, start_after = self._compute_overfetch_params( - remaining=remaining, - next_after=next_page.after, - ) - - method_result = await original_method(user_id=user_id, limit=request_limit, next_page=next_page, **kw) - if not isinstance(method_result, FeedResultClient): - raise TypeError('SubFeed function must return "FeedResultClient" instance.') - - upstream_has_next_page = upstream_has_next_page or method_result.has_next_page - - inspected_count = 0 - - keys_by_index: Optional[List[Optional[str]]] = None - if isinstance(dedup_state, _RedisDedupState): - keys_by_index = [] - batch_keys: List[str] = [] - for entity in method_result.data: - key = self._get_entity_key(entity) - keys_by_index.append(key) - if key is not None: - batch_keys.append(key) - await dedup_state.prefetch(batch_keys) - - for idx, entity in enumerate(method_result.data, start=1): - inspected_count = idx - - key = keys_by_index[idx - 1] if keys_by_index is not None else self._get_entity_key(entity) - if key is None: - continue - - if not dedup_state.should_accept(key, leaf_priority): - continue - - collected.append(entity) - dedup_state.record(key, leaf_priority) - - if len(collected) >= limit: - break - - if len(collected) == before_len: - if not method_result.has_next_page: - break - - if can_overfetch and request_limit > remaining and start_after is not None: - end_after = next_page.after - if isinstance(end_after, int) and end_after == start_after + len(method_result.data): - next_page.after = start_after + inspected_count - - return FeedResultClient(data=collected, next_page=next_page, has_next_page=upstream_has_next_page) - - return _wrapped_method - - def _decode_seen_from_cursor(self, next_page: FeedResultNextPage) -> Dict[str, int]: - entry = next_page.data.get(self.merger_id) - if not entry or entry.after is None: - return {} - - after = entry.after - if isinstance(after, dict) and "z" in after: - payload = base64.urlsafe_b64decode(after["z"].encode()) - raw = zlib.decompress(payload).decode() - decoded = json.loads(raw) - if isinstance(decoded, dict): - return {str(k): int(v) for k, v in decoded.items()} - if isinstance(decoded, list): - seen_map: Dict[str, int] = {} - for entry_item in decoded: - if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: - seen_map[str(entry_item[0])] = int(entry_item[1]) - else: - seen_map[str(entry_item)] = 0 - return seen_map - return {} - if isinstance(after, dict) and "seen" in after: - return {str(k): 0 for k in list(after["seen"])} - if isinstance(after, list): - return {str(k): 0 for k in list(after)} - if isinstance(after, dict): - return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} - return {} - - def _encode_seen_for_cursor(self, seen_updates_in_order: List[tuple[str, int]]) -> Any: - if self.cursor_max_keys is not None: - seen_updates_in_order = seen_updates_in_order[-self.cursor_max_keys :] - - if not self.cursor_compress: - return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} - - raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() - compressed = zlib.compress(raw) - return { - "v": 2, - "c": "zlib+base64", - "n": len(seen_updates_in_order), - "z": base64.urlsafe_b64encode(compressed).decode(), - } - - async def _redis_zmscore( - self, - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - members: List[str], - ) -> List[Optional[float]]: - if not members: - return [] - - zmscore_fn = getattr(redis_client, "zmscore", None) - if zmscore_fn is not None: - res = zmscore_fn(key, members) - if inspect.iscoroutine(res): - res = await res - return [None if v is None else float(v) for v in list(res)] - - if not _is_async_redis_client(redis_client): - - def _sync_pipeline_execute() -> Any: - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - return pipe.execute() - - res = await asyncio.to_thread(_sync_pipeline_execute) - return [None if v is None else float(v) for v in list(res)] - - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - res = pipe.execute() - if inspect.iscoroutine(res): - res = await res - return [None if v is None else float(v) for v in list(res)] - - async def _redis_zadd_and_expire( - self, - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - member_scores: Dict[str, int], - ) -> None: - if not member_scores: - return - await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) - await _redis_call(redis_client, "expire", key, self.state_ttl_seconds) + CursorMap(next_page).reset_keys(descendant_keys) def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") @@ -461,88 +112,102 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - if limit <= 0: - return FeedResult(data=[], next_page=next_page, has_next_page=False) + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + elif ctx.redis_client is None and redis_client is not None: + ctx.redis_client = redis_client + + if ctx.executor is None: + from ..execution.executor import Executor - entry = next_page.data.get(self.merger_id) - requested_page = entry.page if entry is not None else None - is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) + ctx.executor = Executor() - if self.state_backend == "redis" and not redis_client: - raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") + return await ctx.executor.run(self, ctx, limit, next_page, **params) - working_next_page = _pydantic_deep_copy(next_page) + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> CallablePlan: + async def _run(executor: Any) -> FeedResult: + if limit <= 0: + return FeedResult(data=[], next_page=next_page, has_next_page=False) - if is_fresh_session: - self._reset_descendant_cursors(working_next_page) + if ctx.executor is None: + ctx.executor = executor - seen_priority_map: Dict[str, int] = {} - seen_updates_in_order: List[tuple[str, int]] = [] - if self.state_backend == "cursor" and not is_fresh_session: - seen_priority_map = self._decode_seen_from_cursor(next_page) + entry = next_page.data.get(self.merger_id) + requested_page = entry.page if entry is not None else None + is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) - seen_request_set: set[str] = set(seen_priority_map.keys()) + redis_client = ctx.redis_client + if self.state_backend == "redis" and not redis_client: + raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") - redis_state_key = "" - redis_new_scores: Dict[str, int] = {} - redis_seen_cache: Dict[str, Optional[int]] = {} - if self.state_backend == "redis" and redis_client: - redis_state_key = self._build_redis_state_key(user_id=user_id, params=params) + working_next_page = _pydantic_deep_copy(next_page) if is_fresh_session: - await _redis_call(redis_client, "delete", redis_state_key) + self._reset_descendant_cursors(working_next_page) + + seen_request_set: set[str] = set() + store: SeenStore + if self.state_backend == "cursor": + cursor_entry = next_page.data.get(self.merger_id) + store = CursorSeenStore.from_after( + cursor_entry.after if (cursor_entry is not None and not is_fresh_session) else None, + cursor_compress=self.cursor_compress, + cursor_max_keys=self.cursor_max_keys, + ) + else: + assert redis_client is not None + redis_state_key = self._build_redis_state_key(user_id=ctx.user_id, params=params) + store = RedisSeenStore.create( + redis_client=redis_client, + redis_key=redis_state_key, + ttl_seconds=self.state_ttl_seconds, + ) + if is_fresh_session: + await store.reset() - if self.state_backend == "cursor": - dedup_state: _DedupState = _CursorDedupState( - seen_priority_map=seen_priority_map, - seen_updates_in_order=seen_updates_in_order, + policy = DeduplicationPolicy( + dedup_key=self.dedup_key, + missing_key_policy=self.missing_key_policy, + store=store, seen_request_set=seen_request_set, ) - else: - assert redis_client is not None - dedup_state = _RedisDedupState( - redis_client=redis_client, - redis_state_key=redis_state_key, - redis_seen_cache=redis_seen_cache, - redis_new_scores=redis_new_scores, - seen_request_set=seen_request_set, - zmscore=self._redis_zmscore, - ) - - original_methods_dict = methods_dict - child = _pydantic_deep_copy(self.data) - - rewritten_methods_dict = dict(original_methods_dict) - - for sf in self._iter_subfeeds(child): - self._register_wrapped_subfeed_method( - subfeed=sf, - original_methods_dict=original_methods_dict, - rewritten_methods_dict=rewritten_methods_dict, - dedup_state=dedup_state, + refill_settings = RefillExecutionSettings( + overfetch_factor=self.overfetch_factor, + max_refill_loops=self.max_refill_loops, + ) + child_ctx = ExecutionContext( + methods_dict=ctx.methods_dict, + user_id=ctx.user_id, + redis_client=ctx.redis_client, + executor=ctx.executor, + dedup=policy, + refill_settings=refill_settings, + dedup_settings=refill_settings, ) - child_result = await child.get_data( - methods_dict=rewritten_methods_dict, - user_id=user_id, - limit=limit, - next_page=working_next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - - if self.state_backend == "redis" and redis_client: - await self._redis_zadd_and_expire(redis_client, redis_state_key, redis_new_scores) + child = _pydantic_deep_copy(self.data) + child_result = await executor.run(child, child_ctx, limit, working_next_page, **params) - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - merger_after: Any = None - if self.state_backend == "cursor": - merger_after = self._encode_seen_for_cursor(seen_updates_in_order) + commit_result: Any = await store.commit() + merger_after: Any = commit_result if self.state_backend == "cursor" else None - result_next_page = _pydantic_deep_copy(child_result.next_page) - result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + result_next_page = _pydantic_deep_copy(child_result.next_page) + result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) + return FeedResult( + data=child_result.data, + next_page=result_next_page, + has_next_page=child_result.has_next_page, + ) - return FeedResult(data=child_result.data, next_page=result_next_page, has_next_page=child_result.has_next_page) + return CallablePlan(fn=_run) diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py index 76982fb..1da51ea 100644 --- a/smartfeed/mergers/percentage.py +++ b/smartfeed/mergers/percentage.py @@ -7,6 +7,8 @@ from pydantic import BaseModel from redis.asyncio import Redis as AsyncRedis +from ..execution.context import ExecutionContext +from ..execution.executor import SlotSpec, SlotsPlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage if TYPE_CHECKING: @@ -61,51 +63,61 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - result = FeedResult(data=[], next_page=FeedResultNextPage(data={}), has_next_page=False) + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - dedup_active = bool(params.pop("_sf_dedup_active", False)) + if ctx.executor is None: + from ..execution.executor import Executor - items_data: List[List[Any]] = [[] for _ in self.items] - results: List[Optional[FeedResult]] = [None for _ in self.items] + ctx.executor = Executor() - indexed_items = list(enumerate(self.items)) - fetch_order = indexed_items - if dedup_active: - fetch_order = sorted( - indexed_items, - key=lambda p: (getattr(p[1].data, "dedup_priority", 0), -p[0]), - reverse=True, - ) - - for idx, item in fetch_order: - item_result = cast( - FeedResult, - await item.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit * item.percentage // 100, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ), - ) - - results[idx] = item_result - - for idx, result_item in enumerate(results): - assert result_item is not None - items_data[idx] = result_item.data - - if not result.has_next_page and result_item.has_next_page: - result.has_next_page = True - result.next_page.data.update(result_item.next_page.data) - - result.data = await self._merge_items_data(items_data=items_data) + return await ctx.executor.run(self, ctx, limit, next_page, **params) - if self.shuffle: - shuffle(result.data) - - return result + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> SlotsPlan: + owners: List[BaseFeedConfigModel] = [cast(BaseFeedConfigModel, item.data) for item in self.items] + + slots: List[SlotSpec] = [] + for item, owner in zip(self.items, owners): + child_limit = limit * int(item.percentage) // 100 + slots.append(SlotSpec(owner=owner, max_count=max(0, child_limit))) + + async def _assemble( + output: List[Any], + merged_next_page: FeedResultNextPage, + owner_results: Dict[int, FeedResult], + ) -> FeedResult: + items_data: List[List[Any]] = [] + has_next_page = False + + for owner in owners: + child_res = owner_results.get(id(owner)) + if child_res is None: + items_data.append([]) + continue + items_data.append(list(child_res.data)) + has_next_page = has_next_page or bool(child_res.has_next_page) + + data = await self._merge_items_data(items_data=items_data) + if self.shuffle: + shuffle(data) + + return FeedResult(data=data, next_page=merged_next_page, has_next_page=has_next_page) + + return SlotsPlan( + ctx=ctx, + limit=limit, + next_page=next_page, + params=dict(params), + slots=slots, + assemble=_assemble, + ) diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py index 17ed530..460e9f0 100644 --- a/smartfeed/mergers/percentage_gradient.py +++ b/smartfeed/mergers/percentage_gradient.py @@ -1,10 +1,12 @@ from random import shuffle -from typing import Any, Callable, Dict, Literal, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast import redis from pydantic import model_validator from redis.asyncio import Redis as AsyncRedis +from ..execution.context import ExecutionContext +from ..execution.executor import SlotSpec, SlotsPlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside from .percentage import MergerPercentageItem @@ -28,7 +30,7 @@ def validate_merger_percentage_gradient(self) -> "MergerPercentageGradient": raise ValueError('"size_to_step" must be bigger than 1') return self - async def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: + def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: result: Dict = { "limit_from": 0, "limit_to": 0, @@ -74,91 +76,90 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - result = FeedResult( - data=[], - next_page=FeedResultNextPage( - data={ - self.merger_id: FeedResultNextPageInside( - page=next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1, - after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, - ) - }, - ), - has_next_page=False, - ) + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - limits_and_percents = await self._calculate_limits_and_percents( - page=result.next_page.data[self.merger_id].page, - limit=limit, + if ctx.executor is None: + from ..execution.executor import Executor + + ctx.executor = Executor() + + return await ctx.executor.run(self, ctx, limit, next_page, **params) + + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> SlotsPlan: + start_page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + start_after = next_page.data[self.merger_id].after if self.merger_id in next_page.data else None + + plan_next_page = FeedResultNextPage( + data={ + **next_page.data, + self.merger_id: FeedResultNextPageInside(page=start_page, after=start_after), + } ) - dedup_active = bool(params.pop("_sf_dedup_active", False)) - - from_priority = getattr(self.item_from.data, "dedup_priority", 0) - to_priority = getattr(self.item_to.data, "dedup_priority", 0) - - if dedup_active and to_priority > from_priority: - item_to = await self.item_to.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_to"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - item_from = await self.item_from.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_from"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - else: - item_from = await self.item_from.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_from"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - item_to = await self.item_to.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limits_and_percents["limit_to"], - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - - from_start_index = 0 - to_start_index = 0 - for lp_data in limits_and_percents["percentages"]: - from_end_index = (lp_data["limit"] * lp_data["from"] // 100) + from_start_index - to_end_index = (lp_data["limit"] * lp_data["to"] // 100) + to_start_index - - result.data.extend(item_from.data[from_start_index:from_end_index]) - result.data.extend(item_to.data[to_start_index:to_end_index]) - - from_start_index = from_end_index - to_start_index = to_end_index - - result.next_page.data.update(item_from.next_page.data) - result.next_page.data.update(item_to.next_page.data) - - if any([item_from.has_next_page, item_to.has_next_page]): - result.has_next_page = True - - if self.shuffle: - shuffle(result.data) - - result.next_page.data[self.merger_id].page += 1 + limits_and_percents = self._calculate_limits_and_percents(page=start_page, limit=limit) - return result + owner_from = cast(BaseFeedConfigModel, self.item_from.data) + owner_to = cast(BaseFeedConfigModel, self.item_to.data) + + slots = [ + SlotSpec(owner=owner_from, max_count=int(limits_and_percents["limit_from"])), + SlotSpec(owner=owner_to, max_count=int(limits_and_percents["limit_to"])), + ] + + async def _assemble( + output: List[Any], + merged_next_page: FeedResultNextPage, + owner_results: Dict[int, FeedResult], + ) -> FeedResult: + from_res = owner_results.get(id(owner_from)) + to_res = owner_results.get(id(owner_to)) + + from_data = list(from_res.data) if from_res is not None else [] + to_data = list(to_res.data) if to_res is not None else [] + + data: List[Any] = [] + from_start_index = 0 + to_start_index = 0 + for lp_data in limits_and_percents["percentages"]: + from_end_index = (lp_data["limit"] * lp_data["from"] // 100) + from_start_index + to_end_index = (lp_data["limit"] * lp_data["to"] // 100) + to_start_index + + data.extend(from_data[from_start_index:from_end_index]) + data.extend(to_data[to_start_index:to_end_index]) + + from_start_index = from_end_index + to_start_index = to_end_index + + has_next_page = False + if from_res is not None and from_res.has_next_page: + has_next_page = True + if to_res is not None and to_res.has_next_page: + has_next_page = True + + if self.shuffle: + shuffle(data) + + if self.merger_id in merged_next_page.data: + merged_next_page.data[self.merger_id].page += 1 + + return FeedResult(data=data, next_page=merged_next_page, has_next_page=has_next_page) + + return SlotsPlan( + ctx=ctx, + limit=limit, + next_page=plan_next_page, + params=dict(params), + slots=slots, + assemble=_assemble, + ) diff --git a/smartfeed/mergers/positional.py b/smartfeed/mergers/positional.py index a543ddb..bc3a095 100644 --- a/smartfeed/mergers/positional.py +++ b/smartfeed/mergers/positional.py @@ -6,6 +6,8 @@ from pydantic import model_validator from redis.asyncio import Redis as AsyncRedis +from ..execution.context import ExecutionContext +from ..execution.executor import SlotSpec, SlotsPlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside if TYPE_CHECKING: @@ -43,10 +45,27 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - dedup_active = bool(params.pop("_sf_dedup_active", False)) + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + if ctx.executor is None: + from ..execution.executor import Executor + + ctx.executor = Executor() + + return await ctx.executor.run(self, ctx, limit, next_page, **params) + + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> SlotsPlan: page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 positional_has_next_page = True @@ -66,70 +85,54 @@ async def get_data( if position in available_positions: page_positions.append(available_positions.index(position)) - if dedup_active and getattr(self.positional, "dedup_priority", 0) > getattr(self.default, "dedup_priority", 0): - pos_res = await self.positional.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=len(page_positions), - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - default_res = await self.default.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=True, - **params, - ) - else: - default_res = await self.default.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - pos_res = await self.positional.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=len(page_positions), - next_page=next_page, - redis_client=redis_client, - _sf_dedup_active=dedup_active, - **params, - ) - - result = FeedResult( - data=default_res.data, - next_page=FeedResultNextPage( - data={ - self.merger_id: FeedResultNextPageInside( - page=page, - after=next_page.data[self.merger_id].after if self.merger_id in next_page.data else None, - ) - }, - ), - has_next_page=default_res.has_next_page, + pos_limit = len(page_positions) + + # Build a slot ownership schedule by applying the same sequential insert + # semantics as the legacy assembly logic. + schedule: List[BaseFeedConfigModel] = [self.default for _ in range(limit)] + for insert_index in [p - 1 for p in page_positions[:pos_limit]]: + schedule.insert(insert_index, self.positional) + schedule = schedule[:limit] + + # Compress the schedule into contiguous segments. + slots: List[SlotSpec] = [] + if schedule: + current_owner = schedule[0] + count = 1 + for owner in schedule[1:]: + if owner is current_owner: + count += 1 + continue + slots.append(SlotSpec(owner=current_owner, max_count=count)) + current_owner = owner + count = 1 + slots.append(SlotSpec(owner=current_owner, max_count=count)) + + after = next_page.data[self.merger_id].after if self.merger_id in next_page.data else None + + def _assemble( + output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] + ) -> FeedResult: + default_res = owner_results.get(id(self.default)) + pos_res = owner_results.get(id(self.positional)) + + has_next_page = bool(default_res.has_next_page) if default_res is not None else False + if not has_next_page and positional_has_next_page and pos_res is not None and pos_res.has_next_page: + has_next_page = True + + result_next_page = merged_next_page + result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=after) + return FeedResult(data=output, next_page=result_next_page, has_next_page=has_next_page) + + return SlotsPlan( + ctx=ctx, + limit=limit, + next_page=next_page, + params=dict(params), + slots=slots, + owner_fetch_limits={ + id(self.default): limit, + id(self.positional): pos_limit, + }, + assemble=_assemble, ) - - if not result.has_next_page and all([positional_has_next_page, pos_res.has_next_page]): - result.has_next_page = True - - result.next_page.data.update(default_res.next_page.data) - result.next_page.data.update(pos_res.next_page.data) - - for i, post in enumerate(pos_res.data): - result.data = result.data[: page_positions[i] - 1] + [post] + result.data[page_positions[i] - 1 :] - - if len(result.data) > limit: - result.data = result.data[:limit] - - result.next_page.data[self.merger_id].page += 1 - - return result diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index cd3c2f8..b32fdf3 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -9,6 +9,8 @@ from redis.asyncio import RedisCluster as AsyncRedisCluster from .. import jsonlib as json +from ..execution.context import ExecutionContext +from ..execution.executor import CallablePlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside, _redis_call if TYPE_CHECKING: @@ -49,15 +51,21 @@ async def _set_cache( user_id: Any, redis_client: Union[redis.Redis, AsyncRedis], cache_key: str, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> List[Any]: - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - **params, - ) + if ctx is not None and ctx.executor is not None: + result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) + else: + result = await self.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=self.session_size, + next_page=FeedResultNextPage(data={}), + redis_client=ctx.redis_client if ctx is not None else None, + ctx=ctx, + **params, + ) data = result.data if self.deduplicate: @@ -71,15 +79,21 @@ async def _set_cache_async( user_id: Any, redis_client: AsyncRedis, cache_key: str, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> List[Any]: - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - **params, - ) + if ctx is not None and ctx.executor is not None: + result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) + else: + result = await self.data.get_data( + methods_dict=methods_dict, + user_id=user_id, + limit=self.session_size, + next_page=FeedResultNextPage(data={}), + redis_client=ctx.redis_client if ctx is not None else None, + ctx=ctx, + **params, + ) data = result.data if self.deduplicate: @@ -95,6 +109,7 @@ async def _get_cache( limit: int, next_page: FeedResultNextPage, redis_client: Union[redis.Redis, AsyncRedis], + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: if session_cache_key := params.get("custom_view_session_key", None): @@ -107,7 +122,12 @@ async def _get_cache( if not cache_exists or self.merger_id not in next_page.data: logging.info("Cache miss or new session - generating fresh data for %s", cache_key) session_data = await self._set_cache( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + methods_dict=methods_dict, + user_id=user_id, + redis_client=redis_client, + cache_key=cache_key, + ctx=ctx, + **params, ) else: logging.info("Cache exists - attempting read from Redis for %s", cache_key) @@ -117,7 +137,12 @@ async def _get_cache( "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key ) session_data = await self._set_cache( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + methods_dict=methods_dict, + user_id=user_id, + redis_client=redis_client, + cache_key=cache_key, + ctx=ctx, + **params, ) else: logging.info("Successfully read cached data for %s", cache_key) @@ -137,6 +162,7 @@ async def _get_cache_async( limit: int, next_page: FeedResultNextPage, redis_client: AsyncRedis, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: if session_cache_key := params.get("custom_view_session_key", None): @@ -146,7 +172,12 @@ async def _get_cache_async( if not await redis_client.exists(cache_key) or self.merger_id not in next_page.data: session_data = await self._set_cache_async( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + methods_dict=methods_dict, + user_id=user_id, + redis_client=redis_client, + cache_key=cache_key, + ctx=ctx, + **params, ) else: cached_data = await redis_client.get(cache_key) @@ -155,7 +186,12 @@ async def _get_cache_async( "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key ) session_data = await self._set_cache_async( - methods_dict=methods_dict, user_id=user_id, redis_client=redis_client, cache_key=cache_key, **params + methods_dict=methods_dict, + user_id=user_id, + redis_client=redis_client, + cache_key=cache_key, + ctx=ctx, + **params, ) else: logging.info("Successfully read cached data for %s", cache_key) @@ -175,31 +211,60 @@ async def get_data( limit: int, next_page: FeedResultNextPage, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: - if not redis_client: - raise ValueError("Redis client must be provided if using Merger View Session") + if ctx is None: + ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + elif ctx.redis_client is None and redis_client is not None: + ctx.redis_client = redis_client - if isinstance(redis_client, (AsyncRedis, AsyncRedisCluster)): - result = await self._get_cache_async( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) - else: - result = await self._get_cache( - methods_dict=methods_dict, - user_id=user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - **params, - ) + if ctx.executor is None: + from ..execution.executor import Executor + + ctx.executor = Executor() + + return await ctx.executor.run(self, ctx, limit, next_page, **params) + + def build_plan( + self, + *, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + **params: Any, + ) -> CallablePlan: + async def _run(executor: Any) -> FeedResult: + if ctx.redis_client is None: + raise ValueError("Redis client must be provided if using Merger View Session") + + if ctx.executor is None: + ctx.executor = executor + + redis_client = ctx.redis_client + if isinstance(redis_client, (AsyncRedis, AsyncRedisCluster)): + result = await self._get_cache_async( + methods_dict=ctx.methods_dict, + user_id=ctx.user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + ctx=ctx, + **params, + ) + else: + result = await self._get_cache( + methods_dict=ctx.methods_dict, + user_id=ctx.user_id, + limit=limit, + next_page=next_page, + redis_client=redis_client, + ctx=ctx, + **params, + ) - if self.shuffle: - shuffle(result.data) + if self.shuffle: + shuffle(result.data) + return result - return result + return CallablePlan(fn=_run) diff --git a/smartfeed/policies/dedup.py b/smartfeed/policies/dedup.py new file mode 100644 index 0000000..b040c47 --- /dev/null +++ b/smartfeed/policies/dedup.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple + +from .. import jsonlib as json +from .seen_store import SeenStore + +MissingKeyPolicy = Literal["error", "keep", "drop"] + + +def normalize_key(value: Any) -> str: + if isinstance(value, (str, int)): + return str(value) + if isinstance(value, (dict, list)): + return json.dumps(value, sort_keys=True, default=str) + return str(value) + + +def extract_dedup_value(item: Any, dedup_key: Optional[str], missing_key_policy: MissingKeyPolicy) -> Any: + if not dedup_key: + return item + + try: + value = item.get(dedup_key) + except AttributeError: + value = getattr(item, dedup_key, None) + + if value is None and missing_key_policy == "error": + raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {dedup_key}") + return value + + +def entity_key(item: Any, dedup_key: Optional[str], missing_key_policy: MissingKeyPolicy) -> Optional[str]: + raw_value = extract_dedup_value(item, dedup_key, missing_key_policy) + if raw_value is None: + if missing_key_policy == "drop": + return None + if missing_key_policy == "keep": + raw_value = ("__missing__", id(item)) + return normalize_key(raw_value) + + +@dataclass +class DeduplicationPolicy: + """Deduplication policy applied during execution. + + This keeps dedup logic out of merger implementations and plan interpreters. + """ + + dedup_key: Optional[str] + missing_key_policy: MissingKeyPolicy + + # Store backend (cursor or redis) + store: SeenStore + + # Keys encountered/accepted within this request. Prevents duplicates inside one response. + seen_request_set: set[str] + + def key_for(self, item: Any) -> Optional[str]: + return entity_key(item, self.dedup_key, self.missing_key_policy) + + async def prefetch_keys(self, keys: List[str]) -> None: + if not keys: + return + + filtered: List[str] = [] + seen: set[str] = set() + for k in keys: + if k in self.seen_request_set: + continue + if k in seen: + continue + seen.add(k) + filtered.append(k) + + if not filtered: + return + + await self.store.prefetch(filtered) + + def should_accept(self, key: str, priority: int) -> bool: + if key in self.seen_request_set: + return False + + existing_priority = self.store.get(key) + if existing_priority is not None and priority <= existing_priority: + return False + return True + + def record(self, key: str, priority: int) -> None: + self.seen_request_set.add(key) + self.store.set_max(key, priority) + + async def arbitrate_owner_buffers( + self, + *, + owners: List[Any], + owner_buffers: Dict[int, List[Any]], + owner_rank: Dict[int, int], + ) -> Dict[int, List[Any]]: + """Arbitrate winners across multiple owners. + + - Deterministic tie-break: (-priority, owner_rank, item_rank) + - Records winners into the store + - Returns per-owner buffers containing only accepted winners + """ + + keys_to_prefetch: List[str] = [] + keys_seen_local: set[str] = set() + for owner in owners: + owner_id = id(owner) + for item in owner_buffers.get(owner_id, []): + key = self.key_for(item) + if key is None: + continue + if key in keys_seen_local: + continue + keys_seen_local.add(key) + keys_to_prefetch.append(key) + + if keys_to_prefetch: + await self.prefetch_keys(keys_to_prefetch) + + winners: Dict[str, int] = {} + winner_prio: Dict[str, int] = {} + winner_tie: Dict[str, Tuple[int, int, int]] = {} + + for owner in owners: + owner_id = id(owner) + prio = int(getattr(owner, "dedup_priority", 0)) + rank = int(owner_rank.get(owner_id, 0)) + for item_rank, item in enumerate(owner_buffers.get(owner_id, [])): + key = self.key_for(item) + if key is None: + continue + tie = (-prio, rank, item_rank) + existing = winner_tie.get(key) + if existing is None or tie < existing: + winners[key] = owner_id + winner_prio[key] = prio + winner_tie[key] = tie + + for key, _tie in sorted(winner_tie.items(), key=lambda kv: kv[1]): + winner_owner_id = winners.get(key) + if winner_owner_id is None: + continue + prio = int(winner_prio.get(key, 0)) + if not self.should_accept(key, prio): + continue + self.record(key, prio) + + request_set = self.seen_request_set + per_owner_accepted: Dict[int, List[Any]] = {id(o): [] for o in owners} + for owner in owners: + owner_id = id(owner) + accepted: List[Any] = [] + for item in owner_buffers.get(owner_id, []): + key = self.key_for(item) + if key is None: + continue + if winners.get(key) != owner_id: + continue + if key not in request_set: + continue + accepted.append(item) + per_owner_accepted[owner_id] = accepted + + return per_owner_accepted + + async def accept_batch( + self, + *, + items: List[Any], + priority: int, + limit: Optional[int] = None, + ) -> Tuple[List[Any], int]: + """Accept items from a single stream in order. + + Returns accepted items and the number of inspected items. + """ + + if not items: + return [], 0 + + keys_to_prefetch: List[str] = [] + keys_seen_local: set[str] = set() + for item in items: + key = self.key_for(item) + if key is None: + continue + if key in self.seen_request_set: + continue + if key in keys_seen_local: + continue + keys_seen_local.add(key) + keys_to_prefetch.append(key) + + if keys_to_prefetch: + await self.prefetch_keys(keys_to_prefetch) + + accepted: List[Any] = [] + inspected_count = 0 + max_accept = int(limit) if limit is not None else len(items) + + for idx, item in enumerate(items, start=1): + inspected_count = idx + if len(accepted) >= max_accept: + break + key = self.key_for(item) + if key is None: + continue + if not self.should_accept(key, priority): + continue + accepted.append(item) + self.record(key, priority) + if len(accepted) >= max_accept: + break + + return accepted, inspected_count diff --git a/smartfeed/policies/dedup_utils.py b/smartfeed/policies/dedup_utils.py new file mode 100644 index 0000000..e12aed1 --- /dev/null +++ b/smartfeed/policies/dedup_utils.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import asyncio +import base64 +import inspect +import zlib +from typing import Any, Dict, List, Optional, Tuple, Union + +import redis +from redis.asyncio import Redis as AsyncRedis + +from .. import jsonlib as json +from ..feed_models import _is_async_redis_client, _redis_call + + +def decode_seen_from_cursor(after: Any) -> Dict[str, int]: + if after is None: + return {} + + if isinstance(after, dict) and "z" in after: + payload = base64.urlsafe_b64decode(str(after["z"]).encode()) + raw = zlib.decompress(payload).decode() + decoded = json.loads(raw) + if isinstance(decoded, dict): + return {str(k): int(v) for k, v in decoded.items()} + if isinstance(decoded, list): + seen_map: Dict[str, int] = {} + for entry_item in decoded: + if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: + seen_map[str(entry_item[0])] = int(entry_item[1]) + else: + seen_map[str(entry_item)] = 0 + return seen_map + return {} + + if isinstance(after, dict) and "seen" in after: + return {str(k): 0 for k in list(after["seen"])} + if isinstance(after, list): + return {str(k): 0 for k in list(after)} + if isinstance(after, dict): + return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} + return {} + + +def encode_seen_for_cursor( + seen_updates_in_order: List[Tuple[str, int]], + *, + cursor_compress: bool, + cursor_max_keys: Optional[int], +) -> Any: + if cursor_max_keys is not None: + seen_updates_in_order = seen_updates_in_order[-cursor_max_keys:] + + if not cursor_compress: + return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} + + raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() + compressed = zlib.compress(raw) + return { + "v": 2, + "c": "zlib+base64", + "n": len(seen_updates_in_order), + "z": base64.urlsafe_b64encode(compressed).decode(), + } + + +async def redis_zmscore( + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + members: List[str], +) -> List[Optional[float]]: + if not members: + return [] + + zmscore_fn = getattr(redis_client, "zmscore", None) + if zmscore_fn is not None: + res = zmscore_fn(key, members) + if inspect.iscoroutine(res): + res = await res + return [None if v is None else float(v) for v in list(res)] + + if not _is_async_redis_client(redis_client): + + def _sync_pipeline_execute() -> Any: + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + return pipe.execute() + + res = await asyncio.to_thread(_sync_pipeline_execute) + return [None if v is None else float(v) for v in list(res)] + + pipe = redis_client.pipeline() + for m in members: + pipe.zscore(key, m) + res = pipe.execute() + if inspect.iscoroutine(res): + res = await res + return [None if v is None else float(v) for v in list(res)] + + +async def redis_zadd_and_expire( + redis_client: Union[redis.Redis, AsyncRedis], + key: str, + member_scores: Dict[str, int], + *, + ttl_seconds: int, +) -> None: + if not member_scores: + return + await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) + await _redis_call(redis_client, "expire", key, ttl_seconds) diff --git a/smartfeed/policies/seen_store.py b/smartfeed/policies/seen_store.py new file mode 100644 index 0000000..a867346 --- /dev/null +++ b/smartfeed/policies/seen_store.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, cast + +import redis +from redis.asyncio import Redis as AsyncRedis + +from ..feed_models import _redis_call +from .dedup_utils import decode_seen_from_cursor, encode_seen_for_cursor, redis_zadd_and_expire, redis_zmscore + + +class SeenStore(Protocol): + async def prefetch(self, keys: List[str]) -> None: + raise NotImplementedError + + def get(self, key: str) -> Optional[int]: + raise NotImplementedError + + def set_max(self, key: str, priority: int) -> None: + raise NotImplementedError + + async def reset(self) -> None: + raise NotImplementedError + + async def commit(self) -> Any: + raise NotImplementedError + + +@dataclass +class CursorSeenStore: + """Seen-store that persists in the cursor (`after`).""" + + cursor_compress: bool + cursor_max_keys: Optional[int] + + seen_priority_map: Dict[str, int] + seen_updates_in_order: List[Tuple[str, int]] + + @classmethod + def from_after( + cls, + after: Any, + *, + cursor_compress: bool, + cursor_max_keys: Optional[int], + ) -> "CursorSeenStore": + seen_priority_map = decode_seen_from_cursor(after) + return cls( + cursor_compress=cursor_compress, + cursor_max_keys=cursor_max_keys, + seen_priority_map=seen_priority_map, + seen_updates_in_order=[], + ) + + async def prefetch(self, keys: List[str]) -> None: + return None + + def get(self, key: str) -> Optional[int]: + return self.seen_priority_map.get(key) + + def set_max(self, key: str, priority: int) -> None: + existing = self.seen_priority_map.get(key) + if existing is not None and priority <= existing: + return + self.seen_priority_map[key] = priority + self.seen_updates_in_order.append((key, priority)) + + async def reset(self) -> None: + self.seen_priority_map.clear() + self.seen_updates_in_order.clear() + + async def commit(self) -> Any: + return encode_seen_for_cursor( + self.seen_updates_in_order, + cursor_compress=self.cursor_compress, + cursor_max_keys=self.cursor_max_keys, + ) + + +@dataclass +class RedisSeenStore: + """Seen-store backed by a Redis zset (member=key, score=priority).""" + + redis_client: Union[redis.Redis, AsyncRedis] + redis_key: str + ttl_seconds: int + + redis_seen_cache: Dict[str, Optional[int]] + redis_new_scores: Dict[str, int] + + @classmethod + def create( + cls, + *, + redis_client: Union[redis.Redis, AsyncRedis], + redis_key: str, + ttl_seconds: int, + ) -> "RedisSeenStore": + return cls( + redis_client=redis_client, + redis_key=redis_key, + ttl_seconds=ttl_seconds, + redis_seen_cache={}, + redis_new_scores={}, + ) + + async def prefetch(self, keys: List[str]) -> None: + if not keys: + return + + unique: List[str] = [] + seen: set[str] = set() + for k in keys: + if k in self.redis_seen_cache: + continue + if k in seen: + continue + seen.add(k) + unique.append(k) + + if not unique: + return + + scores_result = redis_zmscore(self.redis_client, self.redis_key, unique) + if inspect.iscoroutine(scores_result): + scores = await cast(Any, scores_result) + else: + scores = scores_result + + for k, s in zip(unique, scores): + self.redis_seen_cache[k] = None if s is None else int(s) + + def get(self, key: str) -> Optional[int]: + return self.redis_seen_cache.get(key) + + def set_max(self, key: str, priority: int) -> None: + existing = self.redis_seen_cache.get(key) + if existing is not None and priority <= existing: + return + self.redis_seen_cache[key] = priority + self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) + + async def reset(self) -> None: + await _redis_call(self.redis_client, "delete", self.redis_key) + self.redis_seen_cache.clear() + self.redis_new_scores.clear() + + async def commit(self) -> Any: + await redis_zadd_and_expire( + self.redis_client, + self.redis_key, + self.redis_new_scores, + ttl_seconds=self.ttl_seconds, + ) + self.redis_new_scores.clear() + return None diff --git a/tests/fixtures/redis.py b/tests/fixtures/redis.py index 2c07678..5c9af72 100644 --- a/tests/fixtures/redis.py +++ b/tests/fixtures/redis.py @@ -15,7 +15,7 @@ async def redis_client(request): client = AsyncRedis(host="localhost", port=6379) try: await client.ping() - except Exception: # pragma: no cover + except Exception: pytest.skip("Redis is not available on localhost:6379") yield client await client.aclose() @@ -24,7 +24,7 @@ async def redis_client(request): client = redis.Redis(host="localhost", port=6379, db=0) try: client.ping() - except Exception: # pragma: no cover + except Exception: pytest.skip("Redis is not available on localhost:6379") yield client client.close() diff --git a/tests/test_manager_params.py b/tests/test_manager_params.py new file mode 100644 index 0000000..a62d160 --- /dev/null +++ b/tests/test_manager_params.py @@ -0,0 +1,40 @@ +import pytest + +from smartfeed.manager import FeedManager +from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside + + +async def meta_method( + user_id: str, + limit: int, + next_page: FeedResultNextPageInside, + meta: dict, +) -> FeedResultClient: + assert meta["tag"] == "alpha" + take = int(meta.get("take", limit)) + data = [f"{user_id}:{meta['tag']}"] * min(limit, take) + next_page.after = None + next_page.page += 1 + return FeedResultClient(data=data, next_page=next_page, has_next_page=False) + + +@pytest.mark.asyncio +async def test_manager_passes_params_to_subfeed() -> None: + config = { + "version": "1", + "feed": { + "subfeed_id": "sf_meta", + "type": "subfeed", + "method_name": "meta_method", + }, + } + + manager = FeedManager(config=config, methods_dict={"meta_method": meta_method}) + result = await manager.get_data( + user_id="u1", + limit=5, + next_page=FeedResultNextPage(data={}), + meta={"tag": "alpha", "take": 2}, + ) + + assert result.data == ["u1:alpha", "u1:alpha"] diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 9bda09f..2ef210a 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -1,3 +1,4 @@ +import asyncio import inspect import pytest @@ -789,11 +790,12 @@ async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf ) assert _ids(res_1.data) == [1, 2, 3, 4, 5] + assert res_1.next_page.data["dedup_append_pages"].page == 2 - # In dedup-active append mode, each child is requested with the full page limit (5). - # B must therefore collect 5 unique items while skipping 2 duplicates -> scan ids 1..7. - assert res_1.next_page.data["sf_b"].after == 7 + # In default arbitrate mode, B only needs to scan far enough to fill the remaining + # portion of the page after arbitration (here: 3 items: ids 3..5). + assert res_1.next_page.data["sf_b"].after == 5 res_2 = await merger.get_data( methods_dict=methods_dict, @@ -807,6 +809,72 @@ async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf _assert_pages_no_overlap(res_1, res_2) +@pytest.mark.asyncio +async def test_dedup_arbitrate_mode_runs_parallel_prefetch_and_arbitrates_winners() -> None: + started_a = asyncio.Event() + started_b = asyncio.Event() + release = asyncio.Event() + + items_a = [{"id": i, "src": "A"} for i in range(1, 200)] + items_b = [{"id": i, "src": "B"} for i in range(1, 200)] + + async def method_a(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + started_a.set() + await release.wait() + offset = int(next_page.after or 0) + data = items_a[offset : offset + limit] + next_page.after = offset + len(data) + next_page.page += 1 + return FeedResultClient(data=data, next_page=next_page, has_next_page=True) + + async def method_b(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + started_b.set() + await release.wait() + offset = int(next_page.after or 0) + data = items_b[offset : offset + limit] + next_page.after = offset + len(data) + next_page.page += 1 + return FeedResultClient(data=data, next_page=next_page, has_next_page=True) + + methods_dict = {"a": method_a, "b": method_b} + + config = { + "merger_id": "dedup_arbitrate", + "type": "merger_deduplication", + "dedup_key": "id", + "state_backend": "cursor", + "cursor_compress": True, + "data": { + "merger_id": "pct", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, + {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, + ], + }, + } + + merger = parse_model(MergerDeduplication, config) + + task = asyncio.create_task( + merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) + ) + + # If execution is sequential, one of these would time out. + await asyncio.wait_for(started_a.wait(), timeout=1) + await asyncio.wait_for(started_b.wait(), timeout=1) + release.set() + + res = await asyncio.wait_for(task, timeout=2) + + assert len(res.data) == 10 + _assert_no_dupes_in_page(res.data) + # With equal priorities, stable tie-breaker should pick A (first branch) for overlapping keys. + winning = {item["id"]: item["src"] for item in res.data} + assert all(winning[i] == "A" for i in range(1, 6)) + + @pytest.mark.asyncio async def test_dedup_refill_loops_advance_dict_after_cursor_not_just_page() -> None: """Dedup refill loops must correctly advance dict-shaped `after` cursors.""" From 076c1017f584500592d69bf6518fc98ff99c0f3f Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sat, 7 Feb 2026 14:40:54 +0000 Subject: [PATCH 16/41] Test cleanup. --- smartfeed/examples/example_client.py | 4 +- smartfeed/manager.py | 3 +- tests/fixtures/dedup_helpers.py | 441 ++++++++ tests/test_merger_deduplication.py | 1396 +++++++------------------- 4 files changed, 782 insertions(+), 1062 deletions(-) create mode 100644 tests/fixtures/dedup_helpers.py diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index 9b48842..d11e00a 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -5,6 +5,7 @@ from smartfeed import jsonlib as json from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside +from tests.utils import parse_model class TestClientRequest(BaseModel): @@ -28,7 +29,8 @@ def validate_next_page(cls, value: Union[str, FeedResultNextPage]) -> Union[str, validate = getattr(FeedResultNextPage, "model_validate", None) if validate is not None: return validate(payload) - return FeedResultNextPage.parse_obj(payload) + return parse_model(FeedResultNextPage, payload) # type: ignore + return value diff --git a/smartfeed/manager.py b/smartfeed/manager.py index f42e95f..3c0ea05 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -6,6 +6,7 @@ from .execution.context import ExecutionContext from .execution.executor import Executor from .schemas import FeedConfig, FeedResult, FeedResultNextPage +from tests.utils import parse_model class FeedManager: @@ -26,7 +27,7 @@ def __init__(self, config: Dict, methods_dict: Dict, redis_client: Optional[Unio if validate is not None: self.feed_config = validate(config) else: - self.feed_config = FeedConfig.parse_obj(config) + self.feed_config = parse_model(FeedConfig, config) # type: ignore self.methods_dict = methods_dict self.redis_client = redis_client diff --git a/tests/fixtures/dedup_helpers.py b/tests/fixtures/dedup_helpers.py new file mode 100644 index 0000000..eb3fbdc --- /dev/null +++ b/tests/fixtures/dedup_helpers.py @@ -0,0 +1,441 @@ +from smartfeed.schemas import FeedResultClient, FeedResultNextPage + + +def _effective_limit(limit, max_per_call): + effective_limit = limit + if isinstance(max_per_call, int) and max_per_call > 0: + effective_limit = min(effective_limit, max_per_call) + return effective_limit + + +def make_offset_paged_method(items, *, max_per_call=None): + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + offset = int(next_page.after or 0) + effective_limit = _effective_limit(limit, max_per_call) + result_data = items[offset : offset + effective_limit] + next_page.after = offset + len(result_data) + next_page.page += 1 + has_next_page = (offset + len(result_data)) < len(items) + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) + + return _method + + +def make_string_after_paged_method(items, *, max_per_call=None, after_field="created_at"): + """A subfeed method whose cursor is a string (e.g. timestamp). + + Cursor semantics: `after` is the last returned `created_at` value (monotonic). + """ + + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + effective_limit = _effective_limit(limit, max_per_call) + + after = next_page.after + start_idx = 0 + if isinstance(after, str) and after: + # Find first item with created_at > after + for i, item in enumerate(items): + if str(item[after_field]) > after: + start_idx = i + break + else: + start_idx = len(items) + + result_data = items[start_idx : start_idx + effective_limit] + has_next_page = (start_idx + len(result_data)) < len(items) + + if result_data: + next_page.after = str(result_data[-1][after_field]) + next_page.page += 1 + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) + + return _method + + +def make_profile_dict_after_method( + profiles_to_items, + *, + max_per_call=None, + after_key="after", +): + """A subfeed method whose cursor is a dict of per-profile offsets. + + Example shape: after = {"p1": 0, "p2": 0} + Cursor semantics: each profile offset increments as items are *read*. + """ + + profile_ids = list(profiles_to_items.keys()) + + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + effective_limit = _effective_limit(limit, max_per_call) + + after = next_page.after + if not isinstance(after, dict): + after = {pid: 0 for pid in profile_ids} + else: + after = dict(after) + for pid in profile_ids: + after.setdefault(pid, 0) + + result = [] + has_next_page = False + + # Build a cyclic iteration over profiles. + active_profiles = [pid for pid in profile_ids] + + i = 0 + while active_profiles and len(result) < effective_limit: + pid = active_profiles[i % len(active_profiles)] + idx = after.get(pid, 0) + items = profiles_to_items.get(pid, []) + + if idx >= len(items): + # This profile is exhausted. + active_profiles.remove(pid) + continue + + result.append(items[idx]) + after[pid] = idx + 1 + i += 1 + + # Determine if any profile still has unread items. + for pid in profile_ids: + if after.get(pid, 0) < len(profiles_to_items.get(pid, [])): + has_next_page = True + break + + next_page.after = after + next_page.page += 1 + return FeedResultClient(data=result, next_page=next_page, has_next_page=has_next_page) + + return _method + + +def _assert_cursor_monotonic_if_present(res_1, res_2, keys): + for key in keys: + if key not in res_1.next_page.data: + continue + + assert key in res_2.next_page.data + + after_1 = res_1.next_page.data[key].after + after_2 = res_2.next_page.data[key].after + + if after_1 is None or after_2 is None: + continue + + if isinstance(after_1, int) and isinstance(after_2, int): + assert after_2 >= after_1 + continue + + if isinstance(after_1, dict) and isinstance(after_2, dict): + continue + + try: + assert after_2 >= after_1 + except TypeError: + pass + + +def _sources(data): + return [x.get("src") for x in data] + + +def _ids(data): + return [x.get("id") for x in data] + + +def _assert_no_dupes_in_page(data): + ids = _ids(data) + assert len(ids) == len(set(ids)) + + +def _assert_pages_no_overlap(res_1, res_2): + assert not (set(_ids(res_1.data)) & set(_ids(res_2.data))) + + +def _assert_two_pages_no_dupes(res_1, res_2): + _assert_no_dupes_in_page(res_1.data) + _assert_no_dupes_in_page(res_2.data) + _assert_pages_no_overlap(res_1, res_2) + + +def _assert_sources_at_positions(data, positions, expected_src): + sources = _sources(data) + for pos in positions: + assert sources[pos - 1] == expected_src + + +def make_items(src, start, end, *, user_id_mod=None, id_offset=0, extra=None): + items = [] + for i in range(start, end): + item_id = id_offset + i + item = {"id": item_id, "src": src} + if user_id_mod is not None: + item["user_id"] = f"u{item_id % user_id_mod}" + if extra: + item.update(extra) + items.append(item) + return items + + +def _subfeed(subfeed_id, method_name, *, dedup_priority=None): + data = {"subfeed_id": subfeed_id, "type": "subfeed", "method_name": method_name} + if dedup_priority is not None: + data["dedup_priority"] = dedup_priority + return data + + +def _dedup_config(merger_id, data, *, dedup_key="id", state_backend="cursor", cursor_compress=True, **kwargs): + config = { + "merger_id": merger_id, + "type": "merger_deduplication", + "dedup_key": dedup_key, + "state_backend": state_backend, + "cursor_compress": cursor_compress, + "data": data, + } + config.update(kwargs) + return config + + +def _percentage_config(merger_id, items, *, shuffle=False): + return {"merger_id": merger_id, "type": "merger_percentage", "shuffle": shuffle, "items": items} + + +def _append_config(merger_id, items, *, shuffle=False): + return {"merger_id": merger_id, "type": "merger_append", "shuffle": shuffle, "items": items} + + +def _distribute_config(merger_id, items, *, distribution_key="user_id"): + return { + "merger_id": merger_id, + "type": "merger_distribute", + "distribution_key": distribution_key, + "items": items, + } + + +def _positional_config(merger_id, *, positions, positional, default): + return { + "merger_id": merger_id, + "type": "merger_positional", + "positions": positions, + "positional": positional, + "default": default, + } + + +def _gradient_config( + merger_id, + *, + item_from, + item_to, + step, + size_to_step, + shuffle=False, +): + return { + "merger_id": merger_id, + "type": "merger_percentage_gradient", + "item_from": item_from, + "item_to": item_to, + "step": step, + "size_to_step": size_to_step, + "shuffle": shuffle, + } + + +async def _run_two_pages(merger, methods_dict, limit, *, next_page=None, **kwargs): + if next_page is None: + next_page = FeedResultNextPage(data={}) + res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=limit, next_page=next_page, **kwargs) + res_2 = await merger.get_data( + methods_dict=methods_dict, user_id="u", limit=limit, next_page=res_1.next_page, **kwargs + ) + return res_1, res_2 + + +def _percentage_items(first, second, *, first_pct=50, second_pct=50): + return [ + {"percentage": first_pct, "data": first}, + {"percentage": second_pct, "data": second}, + ] + + +def _two_subfeed_spec(*, name="a", subfeed_id="sf_a", max_per_call=None, dedup_priority=None): + return { + "name": name, + "subfeed_id": subfeed_id, + "max_per_call": max_per_call, + "dedup_priority": dedup_priority, + } + + +def _build_two_subfeed_methods(items_a, items_b, *, spec_a=None, spec_b=None): + if spec_a is None: + spec_a = _two_subfeed_spec() + if spec_b is None: + spec_b = _two_subfeed_spec(name="b", subfeed_id="sf_b") + + methods_dict = { + spec_a["name"]: make_offset_paged_method(items_a, max_per_call=spec_a["max_per_call"]), + spec_b["name"]: make_offset_paged_method(items_b, max_per_call=spec_b["max_per_call"]), + } + subfeed_a = _subfeed(spec_a["subfeed_id"], spec_a["name"], dedup_priority=spec_a["dedup_priority"]) + subfeed_b = _subfeed(spec_b["subfeed_id"], spec_b["name"], dedup_priority=spec_b["dedup_priority"]) + return methods_dict, subfeed_a, subfeed_b + + +def _build_two_subfeed_dedup_merger( + *, + items_a, + items_b, + child_builder, + merger_id, + spec_a=None, + spec_b=None, + dedup_kwargs=None, +): + methods_dict, subfeed_a, subfeed_b = _build_two_subfeed_methods( + items_a, + items_b, + spec_a=spec_a, + spec_b=spec_b, + ) + config = _dedup_config(merger_id, child_builder(subfeed_a, subfeed_b), **(dedup_kwargs or {})) + return config, methods_dict, subfeed_a, subfeed_b + + +def _build_deep_positional_pct_dedup_merger( + *, + items_p, + items_d1, + items_d2, + dedup_merger_id, + pos_merger_id, + pct_merger_id, + positions, + overfetch_factor=None, + max_refill_loops=None, +): + methods_dict = { + "p": make_offset_paged_method(items_p), + "d1": make_offset_paged_method(items_d1), + "d2": make_offset_paged_method(items_d2), + } + + dedup_kwargs = {} + if overfetch_factor is not None: + dedup_kwargs["overfetch_factor"] = overfetch_factor + if max_refill_loops is not None: + dedup_kwargs["max_refill_loops"] = max_refill_loops + + config = _dedup_config( + dedup_merger_id, + _positional_config( + pos_merger_id, + positions=positions, + positional=_subfeed("sf_p", "p"), + default=_percentage_config( + pct_merger_id, + items=_percentage_items(_subfeed("sf_d1", "d1"), _subfeed("sf_d2", "d2")), + ), + ), + **dedup_kwargs, + ) + return config, methods_dict + + +def _inner_append_config(*, merger_id: str, subfeed_id: str, method_name: str, dedup_priority: int): + return { + "merger_id": merger_id, + "type": "merger_append", + # Important: dedup deletion priority must be visible at this node so parent mergers + # can fetch higher-priority subtrees first when a dedup wrapper is active. + "dedup_priority": dedup_priority, + "shuffle": False, + "items": [ + { + "subfeed_id": subfeed_id, + "type": "subfeed", + "method_name": method_name, + "dedup_priority": dedup_priority, + } + ], + } + + +def _build_deep_priority_tree_for_merger_type(*, merger_type: str): + """Return a deep tree config where low/high leaves overlap by id. + + The inner leaves are wrapped into an append merger to ensure a "deep" tree even + when the outer merger is flat. + """ + + low = _inner_append_config(merger_id="inner_low", subfeed_id="sf_low", method_name="low", dedup_priority=0) + high = _inner_append_config(merger_id="inner_high", subfeed_id="sf_high", method_name="high", dedup_priority=100) + + if merger_type == "merger_append": + return { + "merger_id": "outer_append", + "type": "merger_append", + "shuffle": False, + # Put low first intentionally; priority must still make high win for overlapping ids. + "items": [low, high], + } + + if merger_type == "merger_distribute": + return { + "merger_id": "outer_dist", + "type": "merger_distribute", + "distribution_key": "user_id", + # Put low first intentionally. + "items": [low, high], + } + + if merger_type == "merger_percentage": + return { + "merger_id": "outer_pct", + "type": "merger_percentage", + "shuffle": False, + "items": [ + {"percentage": 50, "data": low}, + {"percentage": 50, "data": high}, + ], + } + + if merger_type == "merger_percentage_gradient": + return { + "merger_id": "outer_grad", + "type": "merger_percentage_gradient", + "item_from": {"percentage": 60, "data": low}, + "item_to": {"percentage": 40, "data": high}, + "step": 20, + "size_to_step": 5, + "shuffle": False, + } + + if merger_type == "merger_positional": + # High priority on positional branch so it must win duplicates. + high_pos = _inner_append_config( + merger_id="inner_pos_high", + subfeed_id="sf_high", + method_name="high", + dedup_priority=100, + ) + low_def = _inner_append_config( + merger_id="inner_def_low", + subfeed_id="sf_low", + method_name="low", + dedup_priority=0, + ) + return { + "merger_id": "outer_pos", + "type": "merger_positional", + "positions": [1, 3, 5, 7, 9, 11], + "positional": high_pos, + "default": low_def, + } + + raise AssertionError(f"Unknown merger_type: {merger_type}") diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 2ef210a..ca83885 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -4,254 +4,24 @@ import pytest from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication +from tests.fixtures import dedup_helpers as dh from tests.fixtures.redis import redis_client # noqa: F401 from tests.utils import parse_model -def make_offset_paged_method(items, *, max_per_call=None): - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - offset = int(next_page.after or 0) - effective_limit = limit - if isinstance(max_per_call, int) and max_per_call > 0: - effective_limit = min(effective_limit, max_per_call) - result_data = items[offset : offset + effective_limit] - next_page.after = offset + len(result_data) - next_page.page += 1 - has_next_page = (offset + len(result_data)) < len(items) - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) +PROFILES_B_1_TO_8 = { + "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], + "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], +} - return _method +def _assert_winning_src_for_ids(data, ids, expected_src: str) -> None: + winning = {item["id"]: item["src"] for item in data} + assert all(winning[i] == expected_src for i in ids if i in winning) -def make_string_after_paged_method(items, *, max_per_call=None, after_field="created_at"): - """A subfeed method whose cursor is a string (e.g. timestamp). - Cursor semantics: `after` is the last returned `created_at` value (monotonic). - """ - - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - effective_limit = limit - if isinstance(max_per_call, int) and max_per_call > 0: - effective_limit = min(effective_limit, max_per_call) - - after = next_page.after - start_idx = 0 - if isinstance(after, str) and after: - # Find first item with created_at > after - for i, item in enumerate(items): - if str(item[after_field]) > after: - start_idx = i - break - else: - start_idx = len(items) - - result_data = items[start_idx : start_idx + effective_limit] - has_next_page = (start_idx + len(result_data)) < len(items) - - if result_data: - next_page.after = str(result_data[-1][after_field]) - next_page.page += 1 - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) - - return _method - - -def make_profile_dict_after_method( - profiles_to_items, - *, - max_per_call=None, - after_key="after", -): - """A subfeed method whose cursor is a dict of per-profile offsets. - - Example shape: after = {"p1": 0, "p2": 0} - Cursor semantics: each profile offset increments as items are *read*. - """ - - profile_ids = list(profiles_to_items.keys()) - - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - effective_limit = limit - if isinstance(max_per_call, int) and max_per_call > 0: - effective_limit = min(effective_limit, max_per_call) - - after = next_page.after - if not isinstance(after, dict): - after = {pid: 0 for pid in profile_ids} - else: - after = dict(after) - for pid in profile_ids: - after.setdefault(pid, 0) - - result = [] - has_next_page = False - - # Build a cyclic iteration over profiles. - active_profiles = [pid for pid in profile_ids] - - i = 0 - while active_profiles and len(result) < effective_limit: - pid = active_profiles[i % len(active_profiles)] - idx = after.get(pid, 0) - items = profiles_to_items.get(pid, []) - - if idx >= len(items): - # This profile is exhausted. - active_profiles.remove(pid) - continue - - result.append(items[idx]) - after[pid] = idx + 1 - i += 1 - - # Determine if any profile still has unread items. - for pid in profile_ids: - if after.get(pid, 0) < len(profiles_to_items.get(pid, [])): - has_next_page = True - break - - next_page.after = after - next_page.page += 1 - return FeedResultClient(data=result, next_page=next_page, has_next_page=has_next_page) - - return _method - - -def _assert_cursor_monotonic_if_present(res_1, res_2, keys): - for key in keys: - if key not in res_1.next_page.data: - continue - - assert key in res_2.next_page.data - - after_1 = res_1.next_page.data[key].after - after_2 = res_2.next_page.data[key].after - - if after_1 is None or after_2 is None: - continue - - if isinstance(after_1, int) and isinstance(after_2, int): - assert after_2 >= after_1 - continue - - if isinstance(after_1, dict) and isinstance(after_2, dict): - continue - - try: - assert after_2 >= after_1 - except TypeError: - pass - - -def _sources(data): - return [x.get("src") for x in data] - - -def _ids(data): - return [x.get("id") for x in data] - - -def _assert_no_dupes_in_page(data): - ids = _ids(data) - assert len(ids) == len(set(ids)) - - -def _assert_pages_no_overlap(res_1, res_2): - assert not (set(_ids(res_1.data)) & set(_ids(res_2.data))) - - -def _inner_append_config(*, merger_id: str, subfeed_id: str, method_name: str, dedup_priority: int): - return { - "merger_id": merger_id, - "type": "merger_append", - # Important: dedup deletion priority must be visible at this node so parent mergers - # can fetch higher-priority subtrees first when a dedup wrapper is active. - "dedup_priority": dedup_priority, - "shuffle": False, - "items": [ - { - "subfeed_id": subfeed_id, - "type": "subfeed", - "method_name": method_name, - "dedup_priority": dedup_priority, - } - ], - } - - -def _build_deep_priority_tree_for_merger_type(*, merger_type: str): - """Return a deep tree config where low/high leaves overlap by id. - - The inner leaves are wrapped into an append merger to ensure a "deep" tree even - when the outer merger is flat. - """ - - low = _inner_append_config(merger_id="inner_low", subfeed_id="sf_low", method_name="low", dedup_priority=0) - high = _inner_append_config(merger_id="inner_high", subfeed_id="sf_high", method_name="high", dedup_priority=100) - - if merger_type == "merger_append": - return { - "merger_id": "outer_append", - "type": "merger_append", - "shuffle": False, - # Put low first intentionally; priority must still make high win for overlapping ids. - "items": [low, high], - } - - if merger_type == "merger_distribute": - return { - "merger_id": "outer_dist", - "type": "merger_distribute", - "distribution_key": "user_id", - # Put low first intentionally. - "items": [low, high], - } - - if merger_type == "merger_percentage": - return { - "merger_id": "outer_pct", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": low}, - {"percentage": 50, "data": high}, - ], - } - - if merger_type == "merger_percentage_gradient": - return { - "merger_id": "outer_grad", - "type": "merger_percentage_gradient", - "item_from": {"percentage": 60, "data": low}, - "item_to": {"percentage": 40, "data": high}, - "step": 20, - "size_to_step": 5, - "shuffle": False, - } - - if merger_type == "merger_positional": - # High priority on positional branch so it must win duplicates. - high_pos = _inner_append_config( - merger_id="inner_pos_high", - subfeed_id="sf_high", - method_name="high", - dedup_priority=100, - ) - low_def = _inner_append_config( - merger_id="inner_def_low", - subfeed_id="sf_low", - method_name="low", - dedup_priority=0, - ) - return { - "merger_id": "outer_pos", - "type": "merger_positional", - "positions": [1, 3, 5, 7, 9, 11], - "positional": high_pos, - "default": low_def, - } - - raise AssertionError(f"Unknown merger_type: {merger_type}") +def _make_items_by_ids(src: str, ids, *, user_id_mod: int): + return [{"id": i, "user_id": f"u{i % user_id_mod}", "src": src} for i in ids] @pytest.mark.asyncio @@ -263,68 +33,45 @@ async def test_dedup_positional_slot_ownership_cursor_backend() -> None: """ # Default branch has early ids 1..3, which will be seen first. - default_items = [{"id": i, "src": "default"} for i in range(1, 300)] + default_items = dh.make_items("default", 1, 300) # Positional branch starts with duplicates 1..3; it must skip them and fetch 4.. instead. - positional_items = [{"id": i, "src": "pos"} for i in range(1, 300)] + positional_items = dh.make_items("pos", 1, 300) methods_dict = { - "default": make_offset_paged_method(default_items), - "pos": make_offset_paged_method(positional_items), + "default": dh.make_offset_paged_method(default_items), + "pos": dh.make_offset_paged_method(positional_items), } - config = { - "merger_id": "dedup_wrapper", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 20, - "data": { - "merger_id": "positional_mix", - "type": "merger_positional", + config = dh._dedup_config( + "dedup_wrapper", + dh._positional_config( + "positional_mix", # Ensure positional inserts exist on both pages for limit=6: # page1 uses (1,3,5), page2 uses (7,9,11) which map to the same in-page slots. - "positions": [1, 3, 5, 7, 9, 11], - "positional": {"subfeed_id": "sf_pos", "type": "subfeed", "method_name": "pos"}, - "default": {"subfeed_id": "sf_default", "type": "subfeed", "method_name": "default"}, - }, - } + positions=[1, 3, 5, 7, 9, 11], + positional=dh._subfeed("sf_pos", "pos"), + default=dh._subfeed("sf_default", "default"), + ), + max_refill_loops=20, + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=6, - next_page=FeedResultNextPage(data={}), - ) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 6) assert len(res_1.data) == 6 - _assert_no_dupes_in_page(res_1.data) + dh._assert_no_dupes_in_page(res_1.data) # Slot ownership: configured positions [1,3,5] are the positional branch. - assert _sources(res_1.data)[0] == "pos" - assert _sources(res_1.data)[2] == "pos" - assert _sources(res_1.data)[4] == "pos" + dh._assert_sources_at_positions(res_1.data, [1, 3, 5], "pos") # Next page: still no overlap across pages, and positional slots remain owned. - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=6, - next_page=res_1.next_page, - ) - assert len(res_2.data) == 6 - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) + dh._assert_sources_at_positions(res_2.data, [1, 3, 5], "pos") - assert _sources(res_2.data)[0] == "pos" - assert _sources(res_2.data)[2] == "pos" - assert _sources(res_2.data)[4] == "pos" - - _assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"]) + dh._assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"]) @pytest.mark.asyncio @@ -332,69 +79,37 @@ async def test_dedup_percentage_slot_ownership_cursor_backend() -> None: """Percentage mixing order must be preserved even with duplicates across sources.""" # A is called first by the percentage merger; its ids will be seen before B. - a_items = [{"id": i, "src": "A"} for i in range(1, 300)] + a_items = dh.make_items("A", 1, 300) # B starts with duplicates 1..3; it must skip them and fetch unique tail items. # Same IDs as A to force cross-source duplicates. - b_items = [{"id": i, "src": "B"} for i in range(1, 300)] - - methods_dict = { - "a": make_offset_paged_method(a_items), - "b": make_offset_paged_method(b_items), - } - - config = { - "merger_id": "dedup_wrapper_pct", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "pct_mix", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, - ], - }, - } - + b_items = dh.make_items("B", 1, 300) + + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=a_items, + items_b=b_items, + merger_id="dedup_wrapper_pct", + child_builder=lambda sf_a, sf_b: dh._percentage_config( + "pct_mix", + items=dh._percentage_items(sf_a, sf_b), + ), + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - ) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) assert len(res_1.data) == 10 - _assert_no_dupes_in_page(res_1.data) + dh._assert_no_dupes_in_page(res_1.data) # Slot ownership: percentage merge alternates when list sizes are equal. - sources_1 = _sources(res_1.data) - assert sources_1[0] == "A" - assert sources_1[1] == "B" - assert sources_1[2] == "A" - assert sources_1[3] == "B" - - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=res_1.next_page, - ) + assert dh._sources(res_1.data)[:4] == ["A", "B", "A", "B"] assert len(res_2.data) == 10 - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) - sources_2 = _sources(res_2.data) - assert sources_2[0] == "A" - assert sources_2[1] == "B" + assert dh._sources(res_2.data)[:2] == ["A", "B"] - _assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_a", "sf_b", "pct_mix", "dedup_wrapper_pct"]) + dh._assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_a", "sf_b", "pct_mix", "dedup_wrapper_pct"]) @pytest.mark.asyncio @@ -402,70 +117,35 @@ async def test_dedup_deep_tree_cursor_backend() -> None: """Dedup must work through deep merger trees (wrapping leaf methods).""" # Leaf sources: intentionally overlapping ids across different leaves. - p_items = [{"id": i, "src": "P"} for i in range(1, 30)] - d1_items = [{"id": i, "src": "D1"} for i in range(1, 30)] # overlaps P - d2_items = [{"id": 100 + i, "src": "D2"} for i in range(1, 30)] - - methods_dict = { - "p": make_offset_paged_method(p_items), - "d1": make_offset_paged_method(d1_items), - "d2": make_offset_paged_method(d2_items), - } - - # Deep tree: Dedup -> Positional(default=Percentage(D1,D2), positional=SubFeed(P)) - config = { - "merger_id": "dedup_deep", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "pos_deep", - "type": "merger_positional", - # Ensure positional positions exist on both page 1 (1,4) and page 2 (9,12) for limit=8. - "positions": [1, 4, 9, 12], - "positional": {"subfeed_id": "sf_p", "type": "subfeed", "method_name": "p"}, - "default": { - "merger_id": "pct_deep", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_d1", "type": "subfeed", "method_name": "d1"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_d2", "type": "subfeed", "method_name": "d2"}}, - ], - }, - }, - } + p_items = dh.make_items("P", 1, 30) + d1_items = dh.make_items("D1", 1, 30) # overlaps P + d2_items = dh.make_items("D2", 1, 30, id_offset=100) + + config, methods_dict = dh._build_deep_positional_pct_dedup_merger( + items_p=p_items, + items_d1=d1_items, + items_d2=d2_items, + dedup_merger_id="dedup_deep", + pos_merger_id="pos_deep", + pct_merger_id="pct_deep", + # Ensure positional positions exist on both page 1 (1,4) and page 2 (9,12) for limit=8. + positions=[1, 4, 9, 12], + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=8, - next_page=FeedResultNextPage(data={}), - ) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 8) assert len(res_1.data) == 8 - _assert_no_dupes_in_page(res_1.data) + dh._assert_no_dupes_in_page(res_1.data) # Positional ownership must hold even with deep defaults. - assert _sources(res_1.data)[0] == "P" # position 1 - assert _sources(res_1.data)[3] == "P" # position 4 - - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=8, - next_page=res_1.next_page, - ) + dh._assert_sources_at_positions(res_1.data, [1, 4], "P") assert len(res_2.data) == 8 - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) - assert _sources(res_2.data)[0] == "P" - assert _sources(res_2.data)[3] == "P" + dh._assert_sources_at_positions(res_2.data, [1, 4], "P") @pytest.mark.parametrize( @@ -492,32 +172,19 @@ async def test_dedup_deletion_priority_works_for_deep_trees_all_merger_types(mer # "priority" is unobservable because earlier branches can fill the page). We do that by # making the low branch short and duplicate-heavy. if merger_type in {"merger_append", "merger_distribute"}: - low_items = [ - {"id": 1, "user_id": "u0", "src": "low"}, - {"id": 2, "user_id": "u1", "src": "low"}, - {"id": 3, "user_id": "u2", "src": "low"}, - {"id": 1000, "user_id": "u0", "src": "low"}, - {"id": 1001, "user_id": "u1", "src": "low"}, - ] - high_items = [{"id": i, "user_id": f"u{i%3}", "src": "high"} for i in range(1, 200)] + low_items = _make_items_by_ids("low", [1, 2, 3, 1000, 1001], user_id_mod=3) + high_items = dh.make_items("high", 1, 200, user_id_mod=3) else: - low_items = [{"id": i, "user_id": f"u{i%3}", "src": "low"} for i in range(1, 200)] - high_items = [{"id": i, "user_id": f"u{i%3}", "src": "high"} for i in range(1, 200)] + low_items = dh.make_items("low", 1, 200, user_id_mod=3) + high_items = dh.make_items("high", 1, 200, user_id_mod=3) methods_dict = { - "low": make_offset_paged_method(low_items), - "high": make_offset_paged_method(high_items), + "low": dh.make_offset_paged_method(low_items), + "high": dh.make_offset_paged_method(high_items), } - deep_tree = _build_deep_priority_tree_for_merger_type(merger_type=merger_type) - config = { - "merger_id": f"dedup_priority_deep_{merger_type}", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": deep_tree, - } + deep_tree = dh._build_deep_priority_tree_for_merger_type(merger_type=merger_type) + config = dh._dedup_config(f"dedup_priority_deep_{merger_type}", deep_tree) merger = parse_model(MergerDeduplication, config) res = await merger.get_data( @@ -527,15 +194,20 @@ async def test_dedup_deletion_priority_works_for_deep_trees_all_merger_types(mer next_page=FeedResultNextPage(data={}), ) - _assert_no_dupes_in_page(res.data) + dh._assert_no_dupes_in_page(res.data) + + # For append/distribute, priority is only observable if both branches contribute something. + if merger_type in {"merger_append", "merger_distribute"}: + sources = set(dh._sources(res.data)) + assert "high" in sources + assert "low" in sources # Priority is about which source wins overlapping ids (not about output order). - winning = {item["id"]: item["src"] for item in res.data} - assert all(winning[i] == "high" for i in range(1, 6) if i in winning) + _assert_winning_src_for_ids(res.data, range(1, 6), "high") # Placement invariant for positional: positional slots must still be owned by positional branch. if merger_type == "merger_positional": - sources = _sources(res.data) + sources = dh._sources(res.data) assert sources[0] == "high" assert sources[2] == "high" assert sources[4] == "high" @@ -550,52 +222,28 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu un-inspected items. In a deep tree, this must hold for all descendant SubFeeds. """ - p_items = [{"id": 1000 + i, "src": "P"} for i in range(1, 200)] - d1_items = [{"id": i, "src": "D1"} for i in range(1, 200)] - d2_items = [{"id": 500 + i, "src": "D2"} for i in range(1, 200)] - - methods_dict = { - "p": make_offset_paged_method(p_items), - "d1": make_offset_paged_method(d1_items), - "d2": make_offset_paged_method(d2_items), - } - - config = { - "merger_id": "dedup_overfetch", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "overfetch_factor": 3, - "data": { - "merger_id": "pos_overfetch", - "type": "merger_positional", - "positions": [1, 4, 9, 12], - "positional": {"subfeed_id": "sf_p", "type": "subfeed", "method_name": "p"}, - "default": { - "merger_id": "pct_overfetch", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_d1", "type": "subfeed", "method_name": "d1"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_d2", "type": "subfeed", "method_name": "d2"}}, - ], - }, - }, - } + p_items = dh.make_items("P", 1, 200, id_offset=1000) + d1_items = dh.make_items("D1", 1, 200) + d2_items = dh.make_items("D2", 1, 200, id_offset=500) + + config, methods_dict = dh._build_deep_positional_pct_dedup_merger( + items_p=p_items, + items_d1=d1_items, + items_d2=d2_items, + dedup_merger_id="dedup_overfetch", + pos_merger_id="pos_overfetch", + pct_merger_id="pct_overfetch", + positions=[1, 4, 9, 12], + overfetch_factor=3, + ) merger = parse_model(MergerDeduplication, config) - # Page 1 - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=8, - next_page=FeedResultNextPage(data={}), - ) + # Page 1/2 + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 8) assert len(res_1.data) == 8 - _assert_no_dupes_in_page(res_1.data) + dh._assert_no_dupes_in_page(res_1.data) # Dedup merger cursor must exist and advance page. assert "dedup_overfetch" in res_1.next_page.data @@ -613,16 +261,8 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu assert res_1.next_page.data["sf_d2"].after == 4 # Page 2 (monotonic advancement, still no over-advancement) - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=8, - next_page=res_1.next_page, - ) - assert len(res_2.data) == 8 - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) assert res_2.next_page.data["dedup_overfetch"].page == 3 assert res_2.next_page.data["pos_overfetch"].page == 3 @@ -631,7 +271,7 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu assert res_2.next_page.data["sf_d1"].after == 8 assert res_2.next_page.data["sf_d2"].after == 8 - _assert_cursor_monotonic_if_present( + dh._assert_cursor_monotonic_if_present( res_1, res_2, keys=["sf_p", "sf_d1", "sf_d2", "pos_overfetch", "dedup_overfetch"], @@ -640,17 +280,10 @@ async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cu @pytest.mark.asyncio async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: - items = [{"id": i, "src": "S"} for i in range(1, 50)] - methods_dict = {"s": make_offset_paged_method(items)} - - config = { - "merger_id": "dedup_reset", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": {"subfeed_id": "sf_stream", "type": "subfeed", "method_name": "s"}, - } + items = dh.make_items("S", 1, 50) + methods_dict = {"s": dh.make_offset_paged_method(items)} + + config = dh._dedup_config("dedup_reset", dh._subfeed("sf_stream", "s")) merger = parse_model(MergerDeduplication, config) @@ -660,7 +293,7 @@ async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: limit=5, next_page=FeedResultNextPage(data={}), ) - assert _ids(res_1.data) == [1, 2, 3, 4, 5] + assert dh._ids(res_1.data) == [1, 2, 3, 4, 5] # Simulate client "full reload": page=0 for the dedup merger. # Also include the stale descendant cursor; dedup should clear it. @@ -671,79 +304,17 @@ async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: next_page=FeedResultNextPage( data={ "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after), - "sf_stream": res_1.next_page.data["sf_stream"], + # Use a deliberately bogus descendant cursor; the dedup wrapper must ignore/reset it. + "sf_stream": FeedResultNextPageInside(page=99, after=999), } ), ) # Must restart from the beginning. - assert _ids(res_2.data) == [1, 2, 3, 4, 5] - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_redis_backend_cross_page(redis_client) -> None: - items_a = [{"id": i, "src": "A"} for i in range(1, 300)] - # Same IDs as A to force cross-source duplicates. - items_b = [{"id": i, "src": "B"} for i in range(1, 300)] - - methods_dict = { - "a": make_offset_paged_method(items_a), - "b": make_offset_paged_method(items_b), - } - - config = { - "merger_id": "dedup_redis", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "data": { - "merger_id": "pct_mix", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, - ], - }, - } - - merger = parse_model(MergerDeduplication, config) - - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - redis_client=redis_client, - custom_deduplication_key="t1", - ) - - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=res_1.next_page, - redis_client=redis_client, - custom_deduplication_key="t1", - ) - - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) - - # Redis backend should not store seen ids in cursor after. - assert "dedup_redis" in res_2.next_page.data - assert res_2.next_page.data["dedup_redis"].after is None - - # Ensure state is persisted in Redis. - key = "dedup:dedup_redis:u:t1" - members = redis_client.zrange(key, 0, -1) - if inspect.iscoroutine(members): - members = await members - assert len(members) >= len(set(_ids(res_1.data) + _ids(res_2.data))) - + assert dh._ids(res_2.data) == [1, 2, 3, 4, 5] + # And must not propagate the bogus descendant cursor. + assert res_2.next_page.data["sf_stream"].after == 5 + assert res_2.next_page.data["sf_stream"].page == 2 @pytest.mark.asyncio async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf_cursor_exactly() -> None: @@ -754,59 +325,38 @@ async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf """ a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] - b_items = [{"id": i, "src": "B"} for i in range(1, 50)] - - methods_dict = { - "a": make_offset_paged_method(a_items), - # Force multiple internal calls. - "b": make_offset_paged_method(b_items, max_per_call=1), - } - - config = { - "merger_id": "dedup_append_pages", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 50, - "data": { - "merger_id": "append_mix", - "type": "merger_append", - "shuffle": False, - "items": [ - {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - ], - }, - } - + b_items = dh.make_items("B", 1, 50) + + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=a_items, + items_b=b_items, + merger_id="dedup_append_pages", + child_builder=lambda sf_a, sf_b: dh._append_config("append_mix", [sf_a, sf_b]), + spec_b=dh._two_subfeed_spec(name="b", subfeed_id="sf_b", max_per_call=1), + dedup_kwargs={"max_refill_loops": 50}, + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage(data={}), - ) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 5) - assert _ids(res_1.data) == [1, 2, 3, 4, 5] + assert dh._ids(res_1.data) == [1, 2, 3, 4, 5] + assert dh._sources(res_1.data) == ["A", "A", "B", "B", "B"] assert res_1.next_page.data["dedup_append_pages"].page == 2 # In default arbitrate mode, B only needs to scan far enough to fill the remaining # portion of the page after arbitration (here: 3 items: ids 3..5). assert res_1.next_page.data["sf_b"].after == 5 + b_contributed = sum(1 for x in res_1.data if x.get("src") == "B") + assert res_1.next_page.data["sf_b"].after > b_contributed - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=res_1.next_page, - ) + # A is exhausted after 2 reads; ensure cursor reflects that. + assert res_1.next_page.data["sf_a"].after == 2 assert len(res_2.data) == 5 - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) + # Across two pages, B should have advanced exactly 5 more items. + assert res_2.next_page.data["sf_b"].after == 10 @pytest.mark.asyncio @@ -815,46 +365,34 @@ async def test_dedup_arbitrate_mode_runs_parallel_prefetch_and_arbitrates_winner started_b = asyncio.Event() release = asyncio.Event() - items_a = [{"id": i, "src": "A"} for i in range(1, 200)] - items_b = [{"id": i, "src": "B"} for i in range(1, 200)] - - async def method_a(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - started_a.set() - await release.wait() - offset = int(next_page.after or 0) - data = items_a[offset : offset + limit] - next_page.after = offset + len(data) - next_page.page += 1 - return FeedResultClient(data=data, next_page=next_page, has_next_page=True) - - async def method_b(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - started_b.set() - await release.wait() - offset = int(next_page.after or 0) - data = items_b[offset : offset + limit] - next_page.after = offset + len(data) - next_page.page += 1 - return FeedResultClient(data=data, next_page=next_page, has_next_page=True) - - methods_dict = {"a": method_a, "b": method_b} - - config = { - "merger_id": "dedup_arbitrate", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "pct", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, - {"percentage": 50, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, - ], - }, + items_a = dh.make_items("A", 1, 200) + items_b = dh.make_items("B", 1, 200) + + def make_method(*, items, started_event): + async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + started_event.set() + await release.wait() + offset = int(next_page.after or 0) + data = items[offset : offset + limit] + next_page.after = offset + len(data) + next_page.page += 1 + return FeedResultClient(data=data, next_page=next_page, has_next_page=True) + + return _method + + methods_dict = { + "a": make_method(items=items_a, started_event=started_a), + "b": make_method(items=items_b, started_event=started_b), } + config = dh._dedup_config( + "dedup_arbitrate", + dh._percentage_config( + "pct", + items=dh._percentage_items(dh._subfeed("sf_a", "a"), dh._subfeed("sf_b", "b")), + ), + ) + merger = parse_model(MergerDeduplication, config) task = asyncio.create_task( @@ -869,10 +407,9 @@ async def method_b(user_id, limit, next_page, **kwargs): # pylint: disable=unus res = await asyncio.wait_for(task, timeout=2) assert len(res.data) == 10 - _assert_no_dupes_in_page(res.data) + dh._assert_no_dupes_in_page(res.data) # With equal priorities, stable tie-breaker should pick A (first branch) for overlapping keys. - winning = {item["id"]: item["src"] for item in res.data} - assert all(winning[i] == "A" for i in range(1, 6)) + _assert_winning_src_for_ids(res.data, range(1, 6), "A") @pytest.mark.asyncio @@ -883,41 +420,26 @@ async def test_dedup_refill_loops_advance_dict_after_cursor_not_just_page() -> N a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] # B produces ids 1.. in round-robin across profiles; cursor is per-profile offsets. - b_profiles = { - "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], - "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], - } + b_profiles = PROFILES_B_1_TO_8 methods_dict = { - "a": make_offset_paged_method(a_items), - "b": make_profile_dict_after_method(b_profiles), + "a": dh.make_offset_paged_method(a_items), + "b": dh.make_profile_dict_after_method(b_profiles), } # Use a percentage merger so B is asked for a small limit (2 items for limit=4). # This forces refill loops when B's first batch is all duplicates. - config = { - "merger_id": "dedup_dict_after", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 50, - "data": { - "merger_id": "pct_mix", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 50, - "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a", "dedup_priority": 100}, - }, - { - "percentage": 50, - "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b", "dedup_priority": 0}, - }, - ], - }, - } + config = dh._dedup_config( + "dedup_dict_after", + dh._percentage_config( + "pct_mix", + items=dh._percentage_items( + dh._subfeed("sf_a", "a", dedup_priority=100), + dh._subfeed("sf_b", "b", dedup_priority=0), + ), + ), + max_refill_loops=50, + ) merger = parse_model(MergerDeduplication, config) res = await merger.get_data( @@ -928,8 +450,8 @@ async def test_dedup_refill_loops_advance_dict_after_cursor_not_just_page() -> N ) assert len(res.data) == 4 - _assert_no_dupes_in_page(res.data) - assert set(_ids(res.data)) == {1, 2, 3, 4} + dh._assert_no_dupes_in_page(res.data) + assert set(dh._ids(res.data)) == {1, 2, 3, 4} assert "sf_b" in res.next_page.data assert isinstance(res.next_page.data["sf_b"].after, dict) @@ -944,24 +466,17 @@ async def test_dedup_overfetch_does_not_overadvance_non_int_after_cursor() -> No """overfetch_factor must not cause over-advancement for non-rewindable cursors.""" # Single subfeed with dict after cursor; no dedup skips should happen. - profiles = { - "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], - "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], - } + profiles = PROFILES_B_1_TO_8 methods_dict = { - "b": make_profile_dict_after_method(profiles), + "b": dh.make_profile_dict_after_method(profiles), } - config = { - "merger_id": "dedup_nonint_overfetch", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "overfetch_factor": 5, - "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - } + config = dh._dedup_config( + "dedup_nonint_overfetch", + dh._subfeed("sf_b", "b"), + overfetch_factor=5, + ) merger = parse_model(MergerDeduplication, config) res = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=4, next_page=FeedResultNextPage(data={})) @@ -984,39 +499,21 @@ async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplic so it doesn't skip items. """ - items_a = [{"id": i, "src": "A"} for i in range(1, 300)] - items_b = [{"id": i, "src": "B"} for i in range(1, 300)] - - methods_dict = { - "a": make_offset_paged_method(items_a), - "b": make_offset_paged_method(items_b), - } - - config = { - "merger_id": "dedup_overfetch_rewind", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "overfetch_factor": 3, - "max_refill_loops": 20, - "data": { - "merger_id": "pct_mix", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 50, - "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a", "dedup_priority": 100}, - }, - { - "percentage": 50, - "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b", "dedup_priority": 0}, - }, - ], - }, - } + items_a = dh.make_items("A", 1, 300) + items_b = dh.make_items("B", 1, 300) + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=items_a, + items_b=items_b, + merger_id="dedup_overfetch_rewind", + child_builder=lambda sf_a, sf_b: dh._percentage_config( + "pct_mix", + items=dh._percentage_items(sf_a, sf_b), + ), + spec_a=dh._two_subfeed_spec(dedup_priority=100), + spec_b=dh._two_subfeed_spec(name="b", subfeed_id="sf_b", dedup_priority=0), + dedup_kwargs={"overfetch_factor": 3, "max_refill_loops": 20}, + ) merger = parse_model(MergerDeduplication, config) res = await merger.get_data( methods_dict=methods_dict, @@ -1026,12 +523,11 @@ async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplic ) assert len(res.data) == 10 - _assert_no_dupes_in_page(res.data) + dh._assert_no_dupes_in_page(res.data) # A provides 1..5, B must provide 6..10. - winning = {item["id"]: item["src"] for item in res.data} - assert all(winning[i] == "A" for i in range(1, 6)) - assert all(winning[i] == "B" for i in range(6, 11)) + _assert_winning_src_for_ids(res.data, range(1, 6), "A") + _assert_winning_src_for_ids(res.data, range(6, 11), "B") # Cursor rewind check: # - First loop for B reads 5 duplicates -> after becomes 5 @@ -1039,147 +535,132 @@ async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplic assert res.next_page.data["sf_b"].after == 10 +@pytest.mark.parametrize( + "items_a,items_b,min_b_id", + [ + (dh.make_items("A", 1, 4, user_id_mod=2), dh.make_items("B", 1, 200, user_id_mod=2), 4), + (dh.make_items("A", 1, 200, user_id_mod=3), dh.make_items("B", 1, 200, user_id_mod=3), None), + ], +) @pytest.mark.asyncio -async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill() -> None: +async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill( + items_a, items_b, min_b_id +) -> None: """Distribute: duplicates skipped per-leaf and page slices don't overlap.""" - # A is short so B must contribute. - items_a = [{"id": i, "user_id": f"u{i%2}", "src": "A"} for i in range(1, 4)] - # B overlaps A by id and continues. - items_b = [{"id": i, "user_id": f"u{i%2}", "src": "B"} for i in range(1, 200)] - - methods_dict = { - "a": make_offset_paged_method(items_a), - "b": make_offset_paged_method(items_b), - } - - config = { - "merger_id": "dedup_dist_pages", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "dist", - "type": "merger_distribute", - "distribution_key": "user_id", - "items": [ - {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - ], - }, - } - - merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={}) + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=items_a, + items_b=items_b, + merger_id="dedup_dist_pages", + child_builder=lambda sf_a, sf_b: dh._distribute_config("dist", [sf_a, sf_b]), ) - res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) + merger = parse_model(MergerDeduplication, config) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) assert len(res_1.data) == 10 assert len(res_2.data) == 10 - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) # Placement/refill: B must skip duplicate ids 1..3 and still fill the page. - b_ids_1 = [x["id"] for x in res_1.data if x.get("src") == "B"] - assert b_ids_1 and min(b_ids_1) >= 4 + if min_b_id is not None: + b_ids_1 = [x["id"] for x in res_1.data if x.get("src") == "B"] + assert b_ids_1 and min(b_ids_1) >= min_b_id @pytest.mark.asyncio async def test_dedup_percentage_gradient_cursor_backend_across_pages() -> None: - a_items = [{"id": i, "src": "A"} for i in range(1, 300)] - b_items = [{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)] + a_items = dh.make_items("A", 1, 300) + b_items = dh.make_items("B", 1, 30) + dh.make_items("B", 1, 300, id_offset=1000) methods_dict = { - "a": make_offset_paged_method(a_items), - "b": make_offset_paged_method(b_items), - } - - config = { - "merger_id": "dedup_grad_pages", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 50, - "data": { - "merger_id": "grad_mix", - "type": "merger_percentage_gradient", - "item_from": {"percentage": 60, "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}}, - "item_to": {"percentage": 40, "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}}, - "step": 20, - "size_to_step": 5, - "shuffle": False, - }, - } + "a": dh.make_offset_paged_method(a_items), + "b": dh.make_offset_paged_method(b_items), + } + + config = dh._dedup_config( + "dedup_grad_pages", + dh._gradient_config( + "grad_mix", + item_from={"percentage": 60, "data": dh._subfeed("sf_a", "a")}, + item_to={"percentage": 40, "data": dh._subfeed("sf_b", "b")}, + step=20, + size_to_step=5, + shuffle=False, + ), + max_refill_loops=50, + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={}) - ) - res_2 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=res_1.next_page) + res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) + + sources = dh._sources(res_1.data) + assert sources == ["A", "A", "A", "B", "B", "A", "A", "B", "B", "B"] # Gradient merger cursor should exist and advance. assert res_1.next_page.data["grad_mix"].page == 2 assert res_2.next_page.data["grad_mix"].page == 3 +@pytest.mark.parametrize( + "merger_id,custom_deduplication_key,items_a,items_b,child_builder", + [ + ( + "dedup_redis", + "t1", + dh.make_items("A", 1, 300), + dh.make_items("B", 1, 300), # Same IDs as A to force cross-source duplicates. + lambda sf_a, sf_b: dh._percentage_config("pct_mix", items=dh._percentage_items(sf_a, sf_b)), + ), + ( + "dedup_redis_append", + "t2", + dh.make_items("A", 1, 20), + dh.make_items("B", 1, 300), + lambda sf_a, sf_b: dh._append_config("append_mix", [sf_a, sf_b]), + ), + ], +) @pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio -async def test_dedup_redis_backend_cross_page_append(redis_client) -> None: - items_a = [{"id": i, "src": "A"} for i in range(1, 20)] - items_b = [{"id": i, "src": "B"} for i in range(1, 300)] - - methods_dict = { - "a": make_offset_paged_method(items_a), - "b": make_offset_paged_method(items_b), - } - - config = { - "merger_id": "dedup_redis_append", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "redis", - "state_ttl_seconds": 60, - "data": { - "merger_id": "append_mix", - "type": "merger_append", - "shuffle": False, - "items": [ - {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - ], - }, - } - - merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - redis_client=redis_client, - custom_deduplication_key="t2", +async def test_dedup_redis_backend_cross_page( + redis_client, + merger_id, + custom_deduplication_key, + items_a, + items_b, + child_builder, +) -> None: + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=items_a, + items_b=items_b, + merger_id=merger_id, + child_builder=child_builder, + dedup_kwargs={"state_backend": "redis", "state_ttl_seconds": 60}, ) - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=res_1.next_page, + merger = parse_model(MergerDeduplication, config) + + res_1, res_2 = await dh._run_two_pages( + merger, + methods_dict, + 10, redis_client=redis_client, - custom_deduplication_key="t2", + custom_deduplication_key=custom_deduplication_key, ) - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) - assert res_2.next_page.data["dedup_redis_append"].after is None + dh._assert_two_pages_no_dupes(res_1, res_2) + + # Redis backend should not store seen ids in cursor after. + assert merger_id in res_2.next_page.data + assert res_2.next_page.data[merger_id].after is None + + # Ensure state is persisted in Redis. + key = f"dedup:{merger_id}:u:{custom_deduplication_key}" + members = redis_client.zrange(key, 0, -1) + if inspect.iscoroutine(members): + members = await members + assert len(members) >= len(set(dh._ids(res_1.data) + dh._ids(res_2.data))) @pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @@ -1188,116 +669,46 @@ async def test_dedup_wrapper_with_view_session_merger(redis_client) -> None: """Dedup wrapper must work when the child is a view_session merger.""" # Two leaves with overlapping ids; view_session computes a session once. - items_low = [{"id": i, "src": "low"} for i in range(1, 100)] - items_high = [{"id": i, "src": "high"} for i in range(1, 100)] - - methods_dict = { - "low": make_offset_paged_method(items_low), - "high": make_offset_paged_method(items_high), - } + items_low = dh.make_items("low", 1, 100) + items_high = dh.make_items("high", 1, 100) + + methods_dict, subfeed_low, subfeed_high = dh._build_two_subfeed_methods( + items_low, + items_high, + spec_a=dh._two_subfeed_spec(name="low", subfeed_id="sf_low", dedup_priority=0), + spec_b=dh._two_subfeed_spec(name="high", subfeed_id="sf_high", dedup_priority=100), + ) - config = { - "merger_id": "dedup_vs", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { + config = dh._dedup_config( + "dedup_vs", + { "merger_id": "vs", "type": "merger_view_session", "session_size": 30, "session_live_time": 60, "deduplicate": False, "shuffle": False, - "data": { - "merger_id": "pct", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 50, - "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}, - }, - { - "percentage": 50, - "data": { - "subfeed_id": "sf_high", - "type": "subfeed", - "method_name": "high", - "dedup_priority": 100, - }, - }, - ], - }, + "data": dh._percentage_config( + "pct", + items=dh._percentage_items(subfeed_low, subfeed_high), + ), }, - } + ) merger = parse_model(MergerDeduplication, config) - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), + res_1, res_2 = await dh._run_two_pages( + merger, + methods_dict, + 10, redis_client=redis_client, custom_view_session_key="vs1", ) - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=res_1.next_page, - redis_client=redis_client, - custom_view_session_key="vs1", - ) - - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) + dh._assert_two_pages_no_dupes(res_1, res_2) # Deletion priority: for the overlapping early ids, the winning entity must be from high. - winning = {item["id"]: item["src"] for item in (res_1.data + res_2.data)} - assert all(winning[i] == "high" for i in range(1, 11) if i in winning) - - -@pytest.mark.asyncio -async def test_dedup_append_distribute_cursor_backend_no_dupes() -> None: - items_a = [{"id": i, "user_id": f"u{i%3}", "src": "A"} for i in range(1, 200)] - items_b = [{"id": i, "user_id": f"u{i%3}", "src": "B"} for i in range(1, 200)] - - methods_dict = { - "a": make_offset_paged_method(items_a), - "b": make_offset_paged_method(items_b), - } - - config = { - "merger_id": "dedup_dist", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "dist", - "type": "merger_distribute", - "distribution_key": "user_id", - "items": [ - {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - ], - }, - } - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=30, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 30 - _assert_no_dupes_in_page(res.data) + _assert_winning_src_for_ids(res_1.data + res_2.data, range(1, 11), "high") @pytest.mark.asyncio @@ -1308,37 +719,20 @@ async def test_dedup_in_page_deletion_priority_keeps_high_priority_even_if_confi The "high" branch is second in config, but has higher dedup_priority. """ - low_items = [{"id": i, "src": "low"} for i in range(1, 200)] - high_items = [{"id": i, "src": "high"} for i in range(1, 200)] - - methods_dict = { - "low": make_offset_paged_method(low_items), - "high": make_offset_paged_method(high_items), - } - - config = { - "merger_id": "dedup_priority", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "pct", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 50, - "data": {"subfeed_id": "sf_low", "type": "subfeed", "method_name": "low", "dedup_priority": 0}, - }, - { - "percentage": 50, - "data": {"subfeed_id": "sf_high", "type": "subfeed", "method_name": "high", "dedup_priority": 100}, - }, - ], - }, - } + low_items = dh.make_items("low", 1, 200) + high_items = dh.make_items("high", 1, 200) + config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( + items_a=low_items, + items_b=high_items, + merger_id="dedup_priority", + child_builder=lambda sf_a, sf_b: dh._percentage_config( + "pct", + items=dh._percentage_items(sf_a, sf_b), + ), + spec_a=dh._two_subfeed_spec(name="low", subfeed_id="sf_low", dedup_priority=0), + spec_b=dh._two_subfeed_spec(name="high", subfeed_id="sf_high", dedup_priority=100), + ) merger = parse_model(MergerDeduplication, config) res = await merger.get_data( methods_dict=methods_dict, @@ -1347,127 +741,9 @@ async def test_dedup_in_page_deletion_priority_keeps_high_priority_even_if_confi next_page=FeedResultNextPage(data={}), ) - _assert_no_dupes_in_page(res.data) + dh._assert_no_dupes_in_page(res.data) # Priority is about which source "wins" for a given dedup_key, not about output order. # With 50/50 limits, the high-priority branch should supply ids 1..5, while the low-priority # branch will be advanced to avoid duplicates. - winning = {item["id"]: item["src"] for item in res.data} - assert all(winning[i] == "high" for i in range(1, 6)) - - -@pytest.mark.asyncio -async def test_dedup_percentage_gradient_slot_ownership_cursor_backend() -> None: - """Dedup must preserve gradient chunking semantics. + _assert_winning_src_for_ids(res.data, range(1, 6), "high") - For limit=10, size_to_step=5, from/to percentages should yield chunks: - - first 5: 3 from A, 2 from B - - next 5: 2 from A, 3 from B - Dedup must refill within each leaf so these chunk sizes remain true. - """ - - a_items = [{"id": i, "src": "A"} for i in range(1, 300)] - # Start with duplicates, then provide unique tail. - b_items = [{"id": i, "src": "B"} for i in range(1, 30)] + [{"id": 1000 + i, "src": "B"} for i in range(1, 300)] - - methods_dict = { - "a": make_offset_paged_method(a_items), - "b": make_offset_paged_method(b_items), - } - - config = { - "merger_id": "dedup_gradient", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 50, - "data": { - "merger_id": "grad_mix", - "type": "merger_percentage_gradient", - "item_from": { - "percentage": 60, - "data": {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - }, - "item_to": { - "percentage": 40, - "data": {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - }, - "step": 20, - "size_to_step": 5, - "shuffle": False, - }, - } - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 10 - _assert_no_dupes_in_page(res.data) - - sources = _sources(res.data) - assert sources[:3] == ["A", "A", "A"] - assert sources[3:5] == ["B", "B"] - assert sources[5:7] == ["A", "A"] - assert sources[7:10] == ["B", "B", "B"] - - -@pytest.mark.asyncio -async def test_dedup_preserves_append_priority_and_advances_cursors_cursor_backend() -> None: - """Append order is the priority signal; dedup must not let later sources win duplicates. - - Also asserts that a leaf cursor advances even when items are skipped as duplicates. - """ - - a_items = [ - {"id": 1, "src": "A"}, - {"id": 2, "src": "A"}, - ] - # B repeats A's ids first, then continues with unique ids. - b_items = [{"id": i, "src": "B"} for i in range(1, 50)] - - methods_dict = { - "a": make_offset_paged_method(a_items), - "b": make_offset_paged_method(b_items), - } - - config = { - "merger_id": "dedup_append", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "max_refill_loops": 20, - "data": { - "merger_id": "append_mix", - "type": "merger_append", - "shuffle": False, - "items": [ - {"subfeed_id": "sf_a", "type": "subfeed", "method_name": "a"}, - {"subfeed_id": "sf_b", "type": "subfeed", "method_name": "b"}, - ], - }, - } - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage(data={}), - ) - - assert _ids(res.data) == [1, 2, 3, 4, 5] - assert _sources(res.data)[:2] == ["A", "A"] - assert _sources(res.data)[2:] == ["B", "B", "B"] - - # B had to scan past duplicated ids 1 and 2, so its cursor should advance - # farther than the number of items it contributed to the final page. - assert "sf_b" in res.next_page.data - assert isinstance(res.next_page.data["sf_b"].after, int) - b_contributed = sum(1 for x in res.data if x.get("src") == "B") - assert res.next_page.data["sf_b"].after > b_contributed From 70e0006cf4e1ae9f8964d0d27d33ec0d34b50943 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sat, 7 Feb 2026 15:18:27 +0000 Subject: [PATCH 17/41] Minor cleanup. --- smartfeed/examples/example_client.py | 85 ++---------- smartfeed/execution/context.py | 16 ++- smartfeed/execution/executor.py | 75 +++-------- smartfeed/manager.py | 14 +- smartfeed/mergers/append.py | 10 +- smartfeed/mergers/append_distribute.py | 16 +-- smartfeed/mergers/deduplication.py | 69 +++++----- smartfeed/mergers/percentage.py | 16 +-- smartfeed/mergers/percentage_gradient.py | 12 +- smartfeed/mergers/positional.py | 10 +- smartfeed/mergers/view_session.py | 156 ++++------------------- smartfeed/policies/dedup_utils.py | 14 +- smartfeed/policies/seen_store.py | 9 +- smartfeed/pydantic_compat.py | 16 +++ tests/test_merger_deduplication.py | 17 +-- tests/test_merger_view_session.py | 34 ++--- tests/utils.py | 15 +-- 17 files changed, 165 insertions(+), 419 deletions(-) create mode 100644 smartfeed/pydantic_compat.py diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index d11e00a..a24e130 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -4,14 +4,12 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator from smartfeed import jsonlib as json +from smartfeed.pydantic_compat import parse_model from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside -from tests.utils import parse_model class TestClientRequest(BaseModel): - """ - Пример модели клиентского входящего запроса. - """ + """Example client request model.""" profile_id: str = Field(...) limit: int = Field(...) @@ -26,18 +24,13 @@ class TestClientRequest(BaseModel): def validate_next_page(cls, value: Union[str, FeedResultNextPage]) -> Union[str, FeedResultNextPage]: if isinstance(value, str): payload = json.loads(base64.urlsafe_b64decode(value)) - validate = getattr(FeedResultNextPage, "model_validate", None) - if validate is not None: - return validate(payload) - return parse_model(FeedResultNextPage, payload) # type: ignore + return parse_model(FeedResultNextPage, payload) return value class ClientMixerClass: - """ - Пример клиентского класса ClientMixer. - """ + """Example client methods for SmartFeed.""" @staticmethod async def example_method( @@ -46,16 +39,6 @@ async def example_method( next_page: FeedResultNextPageInside, limit_to_return: Optional[int] = None, ) -> FeedResultClient: - """ - Пример клиентского метода. - - :param user_id: ID профиля. - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param limit_to_return: ограничить кол-во результата. - :return: массив букв "profile_id" в количестве "limit" штук. - """ - data = [f"{user_id}_{i}" for i in range(1, 1000)] from_index = (data.index(next_page.after) + 1) if next_page.after else 0 @@ -68,9 +51,7 @@ async def example_method( next_page.after = result_data[-1] if result_data else None next_page.page += 1 - - result = FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) - return result + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) @staticmethod async def empty_method( @@ -79,21 +60,9 @@ async def empty_method( next_page: FeedResultNextPageInside, limit_to_return: Optional[int] = None, # pylint: disable=W0613 ) -> FeedResultClient: - """ - Пример клиентского метода, возвращающего пустые данные. - - :param user_id: ID профиля. - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param limit_to_return: ограничить кол-во результата. - :return: массив букв "profile_id" в количестве "limit" штук. - """ - next_page.after = None next_page.page += 1 - - result = FeedResultClient(data=[], next_page=next_page, has_next_page=False) - return result + return FeedResultClient(data=[], next_page=next_page, has_next_page=False) @staticmethod async def error_method( @@ -102,21 +71,9 @@ async def error_method( next_page: FeedResultNextPageInside, limit_to_return: Optional[int] = None, # pylint: disable=W0613 ) -> FeedResultClient: - """ - Пример клиентского метода, возвращающего пустые данные. - - :param user_id: ID профиля. - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param limit_to_return: ограничить кол-во результата. - :return: массив букв "profile_id" в количестве "limit" штук. - """ - next_page.after = None next_page.page = int(10 / 0) - - result = FeedResultClient(data=[], next_page=next_page, has_next_page=False) - return result + return FeedResultClient(data=[], next_page=next_page, has_next_page=False) @staticmethod async def doubles_method( @@ -125,23 +82,11 @@ async def doubles_method( next_page: FeedResultNextPageInside, limit_to_return: Optional[int] = None, # pylint: disable=W0613 ) -> FeedResultClient: - """ - Пример клиентского метода, возвращающего данные с дублями. - - :param user_id: ID профиля. - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param limit_to_return: ограничить кол-во результата. - :return: массив целых чисел, равный [i for i in range(1, 11)] после удаления дублей. - """ - data = [1, 2, 3, 4, 3, 2, 5, 6, 4, 4, 7, 8, 9, 10, 9, 9, 9] next_page.after = None next_page.page += 1 - - result = FeedResultClient(data=data, next_page=next_page, has_next_page=False) - return result + return FeedResultClient(data=data, next_page=next_page, has_next_page=False) @staticmethod async def keys_method( @@ -150,16 +95,6 @@ async def keys_method( next_page: FeedResultNextPageInside, limit_to_return: Optional[int] = None, ) -> FeedResultClient: - """ - Пример клиентского метода. - - :param user_id: ID профиля. - :param limit: кол-во элементов. - :param next_page: курсор пагинации. - :param limit_to_return: ограничить кол-во результата. - :return: массив букв "profile_id" в количестве "limit" штук. - """ - data = [{"user_id": f"{user_id}_{i%10}", "value": i} for i in range(1, 1000)] from_index = (data.index(next_page.after) + 1) if next_page.after else 0 @@ -172,6 +107,4 @@ async def keys_method( next_page.after = result_data[-1] if result_data else None next_page.page += 1 - - result = FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) - return result + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py index 4a118ea..134c261 100644 --- a/smartfeed/execution/context.py +++ b/smartfeed/execution/context.py @@ -26,13 +26,21 @@ class ExecutionContext: # Execution settings (optional) refill_settings: Optional["RefillExecutionSettings"] = None - dedup_settings: Optional["DedupExecutionSettings"] = None + dedup_settings: Optional["RefillExecutionSettings"] = None + + def ensure_redis_client(self, redis_client: Optional[Union[redis.Redis, AsyncRedis]]) -> None: + if self.redis_client is None and redis_client is not None: + self.redis_client = redis_client + + def ensure_executor(self) -> Any: + if self.executor is None: + from .executor import Executor + + self.executor = Executor() + return self.executor @dataclass(frozen=True) class RefillExecutionSettings: overfetch_factor: int = 1 max_refill_loops: int = 20 - - -DedupExecutionSettings = RefillExecutionSettings diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index 20e3cbc..cd3ba4c 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -286,19 +286,11 @@ def _compute_slot_deficits( page_underfilled = remaining > 0 - if not quota_schedule and not page_underfilled: - return {} - - deficits: Dict[int, int] = {} - if quota_schedule: return self._compute_quota_deficits(plan=plan, owner_buffers=owner_buffers) - - return self._compute_fill_deficits( - plan=plan, - remaining=remaining, - deficit_slots=deficit_slots, - ) + if not page_underfilled: + return {} + return self._compute_fill_deficits(plan=plan, remaining=remaining, deficit_slots=deficit_slots) def _compute_quota_deficits( self, @@ -337,21 +329,12 @@ def _compute_fill_deficits( remaining: int, deficit_slots: List[int], ) -> Dict[int, int]: - deficits: Dict[int, int] = {} to_fill = int(remaining) if to_fill <= 0: - return deficits - - if deficit_slots: - deficits[deficit_slots[-1]] = deficits.get(deficit_slots[-1], 0) + to_fill - return deficits - - if plan.slots: - last_owner_id = id(plan.slots[-1].owner) - deficits[last_owner_id] = deficits.get(last_owner_id, 0) + to_fill - return deficits + return {} - return deficits + owner_id = deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) + return {owner_id: to_fill} if owner_id is not None else {} async def _refill_deficits( self, @@ -418,10 +401,11 @@ async def _refill_deficits( continue base_np = owner_state["current_next_page"] - request_limit = max(1, int(owner_state["remaining"])) + remaining_before = max(1, int(owner_state["remaining"])) + request_limit = remaining_before can_overfetch = CursorMap.can_overfetch(node=refill_owner, base_next_page=base_np) if can_overfetch and overfetch_factor > 1: - request_limit = max(1, int(owner_state["remaining"]) * overfetch_factor) + request_limit = max(1, remaining_before * overfetch_factor) wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit, can_overfetch)) @@ -443,43 +427,26 @@ async def _refill_deficits( for (owner, owner_id, base_np, request_limit, can_overfetch), result in zip(wave_ops, results): owner_state = state[owner_id] - owner_state["last_result"] = result - owner_state["last_request_limit"] = request_limit - owner_state["last_can_overfetch"] = can_overfetch - owner_state["last_base_next_page"] = base_np + remaining_before = int(owner_state["remaining"]) + owner_state["current_next_page"] = result.next_page owner_state["has_next_page"] = bool(result.has_next_page) + cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) - cursor.merge_delta( - base_next_page=plan.next_page, - owner_next_page=result.next_page, - ) - - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - if owner_state["remaining"] <= 0: - continue - last_result = owner_state["last_result"] - if last_result is None: - continue - - refill_prio = int(getattr(refill_owner, "dedup_priority", 0)) + refill_prio = int(getattr(owner, "dedup_priority", 0)) wave_accepted, inspected_count = await dedup_policy.accept_batch( - items=list(last_result.data), + items=list(result.data), priority=refill_prio, - limit=int(owner_state["remaining"]), + limit=max(0, remaining_before), ) - if owner_state["last_can_overfetch"] and owner_state["last_request_limit"] > owner_state["remaining"]: + if can_overfetch and request_limit > remaining_before: CursorMap.rewind_overfetch( - node=refill_owner, - base_next_page=owner_state["last_base_next_page"], - result_next_page=owner_state["current_next_page"], + node=owner, + base_next_page=base_np, + result_next_page=result.next_page, inspected_count=inspected_count, - batch_size=len(last_result.data), + batch_size=len(result.data), ) if wave_accepted: @@ -489,8 +456,6 @@ async def _refill_deficits( if owner_state["remaining"] > 0 and owner_state["has_next_page"]: owner_state["loops"] += 1 - owner_state["last_result"] = None - for refill_owner in deficit_owners: refill_owner_id = id(refill_owner) owner_state = state.get(refill_owner_id) diff --git a/smartfeed/manager.py b/smartfeed/manager.py index 3c0ea05..d6c7a76 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -4,9 +4,8 @@ from redis.asyncio import Redis as AsyncRedis from .execution.context import ExecutionContext -from .execution.executor import Executor +from .pydantic_compat import parse_model from .schemas import FeedConfig, FeedResult, FeedResultNextPage -from tests.utils import parse_model class FeedManager: @@ -23,11 +22,7 @@ def __init__(self, config: Dict, methods_dict: Dict, redis_client: Optional[Unio :param redis_client: объект клиента Redis (для конфигурации с view_session = True). """ - validate = getattr(FeedConfig, "model_validate", None) - if validate is not None: - self.feed_config = validate(config) - else: - self.feed_config = parse_model(FeedConfig, config) # type: ignore + self.feed_config = parse_model(FeedConfig, config) self.methods_dict = methods_dict self.redis_client = redis_client @@ -43,6 +38,5 @@ async def get_data(self, user_id: Any, limit: int, next_page: FeedResultNextPage """ ctx = ExecutionContext(methods_dict=self.methods_dict, user_id=user_id, redis_client=self.redis_client) - ctx.executor = Executor() - result = await ctx.executor.run(self.feed_config.feed, ctx, limit, next_page, **params) - return result + executor = ctx.ensure_executor() + return await executor.run(self.feed_config.feed, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/append.py b/smartfeed/mergers/append.py index e4fc609..a415159 100644 --- a/smartfeed/mergers/append.py +++ b/smartfeed/mergers/append.py @@ -62,10 +62,8 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() - - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py index 442ee00..3e0a8e1 100644 --- a/smartfeed/mergers/append_distribute.py +++ b/smartfeed/mergers/append_distribute.py @@ -26,7 +26,7 @@ class MergerAppendDistribute(BaseFeedConfigModel): sorting_desc: bool = False @no_type_check - async def _uniform_distribute(self, data: list) -> list: + def _uniform_distribute(self, data: list) -> list: if self.sorting_key: data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) @@ -58,11 +58,11 @@ def build_plan( ) -> SlotsPlan: slots = [SlotSpec(owner=item, max_count=limit) for item in self.items] - async def _assemble( + def _assemble( output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] ) -> FeedResult: has_next_page = any(r.has_next_page for r in owner_results.values()) - distributed = await self._uniform_distribute(output) + distributed = self._uniform_distribute(output) return FeedResult(data=distributed, next_page=merged_next_page, has_next_page=has_next_page) return SlotsPlan( @@ -86,10 +86,8 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() - - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/deduplication.py b/smartfeed/mergers/deduplication.py index 94c9841..2686ac9 100644 --- a/smartfeed/mergers/deduplication.py +++ b/smartfeed/mergers/deduplication.py @@ -54,37 +54,34 @@ def validate_merger_deduplication(self) -> "MergerDeduplication": def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: keys: set[str] = set() - - subfeed_id = getattr(feed, "subfeed_id", None) - if isinstance(subfeed_id, str) and subfeed_id: - keys.add(subfeed_id) - - merger_id = getattr(feed, "merger_id", None) - if isinstance(merger_id, str) and merger_id: - keys.add(merger_id) - - child: Any - for attr_name in ("data", "positional", "default"): - child = getattr(feed, attr_name, None) - if isinstance(child, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(child)) - - for attr_name in ("item_from", "item_to"): - child = getattr(feed, attr_name, None) - inner = getattr(child, "data", None) - if isinstance(inner, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(inner)) - - items = getattr(feed, "items", None) - if isinstance(items, list): - for item in items: - if isinstance(item, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(item)) - continue - - inner = getattr(item, "data", None) + stack = [feed] + while stack: + node = stack.pop() + + for attr in ("subfeed_id", "merger_id"): + value = getattr(node, attr, None) + if isinstance(value, str) and value: + keys.add(value) + + for child in ( + getattr(node, "data", None), + getattr(node, "positional", None), + getattr(node, "default", None), + ): + if isinstance(child, BaseFeedConfigModel): + stack.append(child) + + for wrapper in (getattr(node, "item_from", None), getattr(node, "item_to", None)): + inner = getattr(wrapper, "data", None) if isinstance(inner, BaseFeedConfigModel): - keys.update(self._collect_descendant_cursor_keys(inner)) + stack.append(inner) + + items = getattr(node, "items", None) + if isinstance(items, list): + for item in items: + inner = item if isinstance(item, BaseFeedConfigModel) else getattr(item, "data", None) + if isinstance(inner, BaseFeedConfigModel): + stack.append(inner) return keys @@ -117,15 +114,11 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - elif ctx.redis_client is None and redis_client is not None: - ctx.redis_client = redis_client - - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() + else: + ctx.ensure_redis_client(redis_client) - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) def build_plan( self, diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py index 1da51ea..1b034b4 100644 --- a/smartfeed/mergers/percentage.py +++ b/smartfeed/mergers/percentage.py @@ -31,7 +31,7 @@ class MergerPercentage(BaseFeedConfigModel): shuffle: bool = False @staticmethod - async def _merge_items_data(items_data: List[List]) -> List: + def _merge_items_data(items_data: List[List]) -> List: result: List = [] cursor: List[Dict] = [] @@ -68,13 +68,11 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() - - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) def build_plan( self, @@ -91,7 +89,7 @@ def build_plan( child_limit = limit * int(item.percentage) // 100 slots.append(SlotSpec(owner=owner, max_count=max(0, child_limit))) - async def _assemble( + def _assemble( output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult], @@ -107,7 +105,7 @@ async def _assemble( items_data.append(list(child_res.data)) has_next_page = has_next_page or bool(child_res.has_next_page) - data = await self._merge_items_data(items_data=items_data) + data = self._merge_items_data(items_data=items_data) if self.shuffle: shuffle(data) diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py index 460e9f0..b1ff586 100644 --- a/smartfeed/mergers/percentage_gradient.py +++ b/smartfeed/mergers/percentage_gradient.py @@ -81,13 +81,11 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() - - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) def build_plan( self, @@ -117,7 +115,7 @@ def build_plan( SlotSpec(owner=owner_to, max_count=int(limits_and_percents["limit_to"])), ] - async def _assemble( + def _assemble( output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult], diff --git a/smartfeed/mergers/positional.py b/smartfeed/mergers/positional.py index bc3a095..3ac9b32 100644 --- a/smartfeed/mergers/positional.py +++ b/smartfeed/mergers/positional.py @@ -50,13 +50,11 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() - - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) def build_plan( self, diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index b32fdf3..8adf814 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -6,7 +6,6 @@ import redis from redis.asyncio import Redis as AsyncRedis -from redis.asyncio import RedisCluster as AsyncRedisCluster from .. import jsonlib as json from ..execution.context import ExecutionContext @@ -47,83 +46,41 @@ def _dedup_data(self, data: List[Any]) -> List[Any]: async def _set_cache( self, - methods_dict: Dict[str, Callable], - user_id: Any, redis_client: Union[redis.Redis, AsyncRedis], cache_key: str, - ctx: Optional[ExecutionContext] = None, + ctx: ExecutionContext, **params: Any, ) -> List[Any]: - if ctx is not None and ctx.executor is not None: - result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) - else: - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - redis_client=ctx.redis_client if ctx is not None else None, - ctx=ctx, - **params, - ) + if ctx.executor is None: + raise ValueError("Executor must be initialized for MergerViewSession") - data = result.data - if self.deduplicate: - data = self._dedup_data(data) - await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) - return data - - async def _set_cache_async( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - redis_client: AsyncRedis, - cache_key: str, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> List[Any]: - if ctx is not None and ctx.executor is not None: - result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) - else: - result = await self.data.get_data( - methods_dict=methods_dict, - user_id=user_id, - limit=self.session_size, - next_page=FeedResultNextPage(data={}), - redis_client=ctx.redis_client if ctx is not None else None, - ctx=ctx, - **params, - ) + result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) data = result.data if self.deduplicate: data = self._dedup_data(data) - await redis_client.set(cache_key, json.dumps(data)) - await redis_client.expire(cache_key, self.session_live_time) + await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) return data async def _get_cache( self, - methods_dict: Dict[str, Callable], - user_id: Any, limit: int, next_page: FeedResultNextPage, redis_client: Union[redis.Redis, AsyncRedis], - ctx: Optional[ExecutionContext] = None, + ctx: ExecutionContext, **params: Any, ) -> FeedResult: - if session_cache_key := params.get("custom_view_session_key", None): - cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" - else: - cache_key = f"{self.merger_id}_{user_id}" + cache_key = ( + f"{self.merger_id}_{ctx.user_id}_{session_cache_key}" + if (session_cache_key := params.get("custom_view_session_key")) + else f"{self.merger_id}_{ctx.user_id}" + ) logging.info("MergerViewSession cache request for %s", cache_key) cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) if not cache_exists or self.merger_id not in next_page.data: logging.info("Cache miss or new session - generating fresh data for %s", cache_key) session_data = await self._set_cache( - methods_dict=methods_dict, - user_id=user_id, redis_client=redis_client, cache_key=cache_key, ctx=ctx, @@ -137,57 +94,6 @@ async def _get_cache( "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key ) session_data = await self._set_cache( - methods_dict=methods_dict, - user_id=user_id, - redis_client=redis_client, - cache_key=cache_key, - ctx=ctx, - **params, - ) - else: - logging.info("Successfully read cached data for %s", cache_key) - session_data = json.loads(cached_data) - - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - return FeedResult( - data=session_data[(page - 1) * limit :][:limit], - next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - has_next_page=bool(len(session_data) > limit * page), - ) - - async def _get_cache_async( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: AsyncRedis, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if session_cache_key := params.get("custom_view_session_key", None): - cache_key = f"{self.merger_id}_{user_id}_{session_cache_key}" - else: - cache_key = f"{self.merger_id}_{user_id}" - - if not await redis_client.exists(cache_key) or self.merger_id not in next_page.data: - session_data = await self._set_cache_async( - methods_dict=methods_dict, - user_id=user_id, - redis_client=redis_client, - cache_key=cache_key, - ctx=ctx, - **params, - ) - else: - cached_data = await redis_client.get(cache_key) - if cached_data is None: - logging.info( - "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key - ) - session_data = await self._set_cache_async( - methods_dict=methods_dict, - user_id=user_id, redis_client=redis_client, cache_key=cache_key, ctx=ctx, @@ -216,15 +122,11 @@ async def get_data( ) -> FeedResult: if ctx is None: ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - elif ctx.redis_client is None and redis_client is not None: - ctx.redis_client = redis_client - - if ctx.executor is None: - from ..execution.executor import Executor - - ctx.executor = Executor() + else: + ctx.ensure_redis_client(redis_client) - return await ctx.executor.run(self, ctx, limit, next_page, **params) + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) def build_plan( self, @@ -241,27 +143,13 @@ async def _run(executor: Any) -> FeedResult: if ctx.executor is None: ctx.executor = executor - redis_client = ctx.redis_client - if isinstance(redis_client, (AsyncRedis, AsyncRedisCluster)): - result = await self._get_cache_async( - methods_dict=ctx.methods_dict, - user_id=ctx.user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - ctx=ctx, - **params, - ) - else: - result = await self._get_cache( - methods_dict=ctx.methods_dict, - user_id=ctx.user_id, - limit=limit, - next_page=next_page, - redis_client=redis_client, - ctx=ctx, - **params, - ) + result = await self._get_cache( + limit=limit, + next_page=next_page, + redis_client=ctx.redis_client, + ctx=ctx, + **params, + ) if self.shuffle: shuffle(result.data) diff --git a/smartfeed/policies/dedup_utils.py b/smartfeed/policies/dedup_utils.py index e12aed1..2559e0f 100644 --- a/smartfeed/policies/dedup_utils.py +++ b/smartfeed/policies/dedup_utils.py @@ -2,9 +2,8 @@ import asyncio import base64 -import inspect import zlib -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union, cast import redis from redis.asyncio import Redis as AsyncRedis @@ -72,11 +71,8 @@ async def redis_zmscore( if not members: return [] - zmscore_fn = getattr(redis_client, "zmscore", None) - if zmscore_fn is not None: - res = zmscore_fn(key, members) - if inspect.iscoroutine(res): - res = await res + if getattr(redis_client, "zmscore", None) is not None: + res = await _redis_call(redis_client, "zmscore", key, members) return [None if v is None else float(v) for v in list(res)] if not _is_async_redis_client(redis_client): @@ -93,9 +89,7 @@ def _sync_pipeline_execute() -> Any: pipe = redis_client.pipeline() for m in members: pipe.zscore(key, m) - res = pipe.execute() - if inspect.iscoroutine(res): - res = await res + res = await cast(Awaitable[Any], pipe.execute()) return [None if v is None else float(v) for v in list(res)] diff --git a/smartfeed/policies/seen_store.py b/smartfeed/policies/seen_store.py index a867346..d0a9258 100644 --- a/smartfeed/policies/seen_store.py +++ b/smartfeed/policies/seen_store.py @@ -1,8 +1,7 @@ from __future__ import annotations -import inspect from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, cast +from typing import Any, Dict, List, Optional, Protocol, Tuple, Union import redis from redis.asyncio import Redis as AsyncRedis @@ -123,11 +122,7 @@ async def prefetch(self, keys: List[str]) -> None: if not unique: return - scores_result = redis_zmscore(self.redis_client, self.redis_key, unique) - if inspect.iscoroutine(scores_result): - scores = await cast(Any, scores_result) - else: - scores = scores_result + scores = await redis_zmscore(self.redis_client, self.redis_key, unique) for k, s in zip(unique, scores): self.redis_seen_cache[k] = None if s is None else int(s) diff --git a/smartfeed/pydantic_compat.py b/smartfeed/pydantic_compat.py new file mode 100644 index 0000000..185e004 --- /dev/null +++ b/smartfeed/pydantic_compat.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from typing import Any, Mapping, Type, TypeVar + +T = TypeVar("T") + + +def parse_model(model_cls: Type[T], obj: Mapping[str, Any]) -> T: + """Parse a mapping into a Pydantic model. + + Uses Pydantic v2 `model_validate` when available, otherwise falls back to v1 `parse_obj`. + """ + + if hasattr(model_cls, "model_validate"): + return model_cls.model_validate(obj) # type: ignore[attr-defined] + return model_cls.parse_obj(obj) # type: ignore[attr-defined] diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index ca83885..ab9838f 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -1,14 +1,13 @@ import asyncio -import inspect import pytest +from smartfeed.feed_models import _redis_call from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication from tests.fixtures import dedup_helpers as dh from tests.fixtures.redis import redis_client # noqa: F401 from tests.utils import parse_model - PROFILES_B_1_TO_8 = { "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], @@ -71,7 +70,9 @@ async def test_dedup_positional_slot_ownership_cursor_backend() -> None: dh._assert_two_pages_no_dupes(res_1, res_2) dh._assert_sources_at_positions(res_2.data, [1, 3, 5], "pos") - dh._assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"]) + dh._assert_cursor_monotonic_if_present( + res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"] + ) @pytest.mark.asyncio @@ -316,6 +317,7 @@ async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: assert res_2.next_page.data["sf_stream"].after == 5 assert res_2.next_page.data["sf_stream"].page == 2 + @pytest.mark.asyncio async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf_cursor_exactly() -> None: """Append: across pages there is no overlap; refill advances cursors correctly. @@ -543,9 +545,7 @@ async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplic ], ) @pytest.mark.asyncio -async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill( - items_a, items_b, min_b_id -) -> None: +async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill(items_a, items_b, min_b_id) -> None: """Distribute: duplicates skipped per-leaf and page slices don't overlap.""" config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( @@ -657,9 +657,7 @@ async def test_dedup_redis_backend_cross_page( # Ensure state is persisted in Redis. key = f"dedup:{merger_id}:u:{custom_deduplication_key}" - members = redis_client.zrange(key, 0, -1) - if inspect.iscoroutine(members): - members = await members + members = await _redis_call(redis_client, "zrange", key, 0, -1) assert len(members) >= len(set(dh._ids(res_1.data) + dh._ids(res_2.data))) @@ -746,4 +744,3 @@ async def test_dedup_in_page_deletion_priority_keeps_high_priority_even_if_confi # With 50/50 limits, the high-priority branch should supply ids 1..5, while the low-priority # branch will be advanced to avoid duplicates. _assert_winning_src_for_ids(res.data, range(1, 6), "high") - diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py index bd02bfe..f62a096 100644 --- a/tests/test_merger_view_session.py +++ b/tests/test_merger_view_session.py @@ -1,8 +1,8 @@ -import inspect import json import pytest +from smartfeed.feed_models import _redis_call from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerViewSession from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG, MERGER_VIEW_SESSION_DUPS_CONFIG @@ -10,6 +10,10 @@ from tests.utils import parse_model +async def _get_cache_json(redis_client, key: str): + return json.loads(await _redis_call(redis_client, "get", key)) + + @pytest.mark.asyncio async def test_merger_view_session_no_redis() -> None: """ @@ -41,12 +45,7 @@ async def test_merger_view_session(redis_client) -> None: user_id="x", redis_client=redis_client, ) - merger_vs_cache = redis_client.get(name="merger_view_session_example_x") - # Для использования синхронной и асинхронной фикстуры в одном тесте проверяем метод get - if inspect.iscoroutine(merger_vs_cache): - merger_vs_cache = json.loads(await merger_vs_cache) - else: - merger_vs_cache = json.loads(merger_vs_cache) + merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") assert merger_vs_res.data == ["x_1", "x_2", "x_3", "x_4", "x_5", "x_6", "x_7", "x_8", "x_9", "x_10"] assert len(merger_vs_cache) == merger_vs.session_size @@ -70,12 +69,7 @@ async def test_merger_view_session_custom_key(redis_client) -> None: redis_client=redis_client, custom_view_session_key="foo", ) - merger_vs_cache = redis_client.get(name="merger_view_session_example_x_foo") - # Для использования синхронной и асинхронной фикстуры в одном тесте проверяем метод get - if inspect.iscoroutine(merger_vs_cache): - merger_vs_cache = json.loads(await merger_vs_cache) - else: - merger_vs_cache = json.loads(merger_vs_cache) + merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x_foo") assert merger_vs_res.data == ["x_1", "x_2", "x_3", "x_4", "x_5", "x_6", "x_7", "x_8", "x_9", "x_10"] assert len(merger_vs_cache) == merger_vs.session_size @@ -99,12 +93,7 @@ async def test_merger_view_session_next_page(redis_client) -> None: user_id="x", redis_client=redis_client, ) - merger_vs_cache = redis_client.get(name="merger_view_session_example_x") - # Для использования синхронной и асинхронной фикстуры в одном тесте проверяем метод get - if inspect.iscoroutine(merger_vs_cache): - merger_vs_cache = json.loads(await merger_vs_cache) - else: - merger_vs_cache = json.loads(merger_vs_cache) + merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") assert merger_vs_res.data == ["x_11", "x_12", "x_13", "x_14", "x_15", "x_16", "x_17", "x_18", "x_19", "x_20"] assert len(merger_vs_cache) == merger_vs.session_size @@ -122,12 +111,7 @@ async def test_merger_view_session_deduplication(redis_client) -> None: user_id="x", redis_client=redis_client, ) - merger_vs_cache = redis_client.get(name="merger_view_session_example_x") - # Для использования синхронной и асинхронной фикстуры в одном тесте проверяем метод get - if inspect.iscoroutine(merger_vs_cache): - merger_vs_cache = json.loads(await merger_vs_cache) - else: - merger_vs_cache = json.loads(merger_vs_cache) + merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") assert merger_vs_res.data == [i for i in range(1, 11)] assert len(merger_vs_cache) == merger_vs.session_size diff --git a/tests/utils.py b/tests/utils.py index 331f1e7..af29848 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,16 +1,5 @@ from __future__ import annotations -from typing import Any, Dict, Type, TypeVar +from smartfeed.pydantic_compat import parse_model -T = TypeVar("T") - - -def parse_model(model_cls: Type[T], obj: Dict[str, Any]) -> T: - """Parse a dict into a Pydantic model. - - Uses Pydantic v2 `model_validate` when available, otherwise falls back to v1 `parse_obj`. - """ - - if hasattr(model_cls, "model_validate"): - return model_cls.model_validate(obj) # type: ignore[attr-defined] - return model_cls.parse_obj(obj) # type: ignore[attr-defined] +__all__ = ["parse_model"] From 8032ce3982bfa3c091e63be20613266f15f1beda Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sat, 7 Feb 2026 15:23:01 +0000 Subject: [PATCH 18/41] If subfeed is sync - throw it into thread. --- smartfeed/feed_models.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py index 50abcbe..6d29d67 100644 --- a/smartfeed/feed_models.py +++ b/smartfeed/feed_models.py @@ -138,14 +138,27 @@ async def get_data( if arg in params: method_params[arg] = params[arg] + method = method_spec.method + is_async = inspect.iscoroutinefunction(method) or inspect.iscoroutinefunction(getattr(method, "__call__", None)) + try: - method_result = await method_spec.method( - user_id=user_id, - limit=limit, - next_page=subfeed_next_page, - **method_params, - **self.subfeed_params, - ) + if is_async: + method_result = await method( + user_id=user_id, + limit=limit, + next_page=subfeed_next_page, + **method_params, + **self.subfeed_params, + ) + else: + method_result = await asyncio.to_thread( + method, + user_id=user_id, + limit=limit, + next_page=subfeed_next_page, + **method_params, + **self.subfeed_params, + ) except Exception: if self.raise_error: raise From efece1c94d4a537a9b64e93f26719665e026e368 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 00:54:33 +0000 Subject: [PATCH 19/41] Dedup runtime separated. --- smartfeed/execution/dedup_runtime.py | 354 ++++++++++++++++++++++++++ smartfeed/execution/executor.py | 364 ++------------------------- 2 files changed, 374 insertions(+), 344 deletions(-) create mode 100644 smartfeed/execution/dedup_runtime.py diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py new file mode 100644 index 0000000..15773a1 --- /dev/null +++ b/smartfeed/execution/dedup_runtime.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage +from .context import ExecutionContext +from .cursors import CursorMap +from .plans import SlotsPlan + +if TYPE_CHECKING: + from .executor import Executor + + +class DedupRuntime: + """Dedup/refill orchestration. + + This owns the control flow (refill loops, slot deficit refills, etc.) + while `DeduplicationPolicy` stays focused on acceptance/arbitration decisions. + """ + + def __init__(self, executor: "Executor") -> None: + self._executor = executor + + def _get_refill_settings(self, ctx: ExecutionContext) -> Any: + return getattr(ctx, "refill_settings", None) or getattr(ctx, "dedup_settings", None) + + async def run_node_with_dedup_refill( + self, + *, + node: BaseFeedConfigModel, + ctx: ExecutionContext, + limit: int, + next_page: FeedResultNextPage, + params: Dict[str, Any], + initial_result: FeedResult, + ) -> FeedResult: + dedup = getattr(ctx, "dedup", None) + if dedup is None: + return initial_result + + settings = self._get_refill_settings(ctx) + overfetch_factor = max(1, int(getattr(settings, "overfetch_factor", 1))) + max_refill_loops = max(1, int(getattr(settings, "max_refill_loops", 20))) + priority = int(getattr(node, "dedup_priority", 0)) + + collected: List[Any] = [] + remaining = int(limit) + loops = 0 + + current_result = initial_result + current_next_page = current_result.next_page + current_request_limit = max(1, remaining) + has_next_page = bool(current_result.has_next_page) + base_next_page = next_page + + # NOTE: Refill loops are inherently sequential for a single node because + # each subsequent request depends on the previous cursor. + while remaining > 0: + can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) + + accepted, inspected_count = await dedup.accept_batch( + items=list(current_result.data), + priority=priority, + limit=remaining, + ) + + if can_overfetch and current_request_limit > remaining: + CursorMap.rewind_overfetch( + node=node, + base_next_page=base_next_page, + result_next_page=current_next_page, + inspected_count=inspected_count, + batch_size=len(current_result.data), + ) + + if accepted: + collected.extend(accepted) + remaining = limit - len(collected) + + if remaining <= 0 or not has_next_page or loops >= max_refill_loops: + break + loops += 1 + + base_next_page = current_next_page + next_request_limit = max(1, remaining) + can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) + if can_overfetch and overfetch_factor > 1: + next_request_limit = max(1, remaining * overfetch_factor) + + current_result, _plan = await self._executor._run_node_raw( + node, + ctx, + next_request_limit, + base_next_page, + params, + ) + current_next_page = current_result.next_page + current_request_limit = next_request_limit + has_next_page = bool(current_result.has_next_page) + + return FeedResult( + data=collected, + next_page=current_next_page, + has_next_page=has_next_page, + ) + + async def apply_slots_plan_dedup( + self, + *, + plan: SlotsPlan, + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + dedup_policy: Any, + refill_settings: Any, + cursor: CursorMap, + ) -> Tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: + owner_buffers = await dedup_policy.arbitrate_owner_buffers( + owners=owners, + owner_buffers=owner_buffers, + owner_rank=owner_index, + ) + + for owner in owners: + owner_id = id(owner) + if owner_id not in owner_results: + continue + old = owner_results[owner_id] + owner_results[owner_id] = FeedResult( + data=list(owner_buffers.get(owner_id, [])), + next_page=old.next_page, + has_next_page=old.has_next_page, + ) + + deficits = self._compute_slot_deficits(plan=plan, owner_buffers=owner_buffers) + if deficits: + await self._refill_deficits( + plan=plan, + deficits=deficits, + owners=owners, + owner_index=owner_index, + owner_buffers=owner_buffers, + owner_results=owner_results, + dedup_policy=dedup_policy, + refill_settings=refill_settings, + cursor=cursor, + ) + + return owner_buffers, owner_results + + def _compute_slot_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: + total_max = sum(int(s.max_count) for s in plan.slots) + quota_schedule = total_max <= int(plan.limit) + + consumed: Dict[int, int] = {} + remaining = int(plan.limit) + deficit_slots: List[int] = [] + + for slot in plan.slots: + if remaining <= 0: + break + + owner_id = id(slot.owner) + want = min(int(slot.max_count), remaining) + if want <= 0: + continue + + have_total = len(owner_buffers.get(owner_id, [])) + already = int(consumed.get(owner_id, 0)) + available = max(0, have_total - already) + take = min(want, available) + if take < want: + deficit_slots.append(owner_id) + consumed[owner_id] = already + take + remaining -= take + + page_underfilled = remaining > 0 + + if quota_schedule: + return self._compute_quota_deficits(plan=plan, owner_buffers=owner_buffers) + if not page_underfilled: + return {} + return self._compute_fill_deficits(plan=plan, remaining=remaining, deficit_slots=deficit_slots) + + def _compute_quota_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: + deficits: Dict[int, int] = {} + remaining = int(plan.limit) + consumed: Dict[int, int] = {} + for slot in plan.slots: + if remaining <= 0: + break + + owner_id = id(slot.owner) + want = min(int(slot.max_count), remaining) + if want <= 0: + continue + + have_total = len(owner_buffers.get(owner_id, [])) + already = int(consumed.get(owner_id, 0)) + available = max(0, have_total - already) + take = min(want, available) + missing = max(0, want - take) + if missing: + deficits[owner_id] = deficits.get(owner_id, 0) + missing + consumed[owner_id] = already + take + remaining -= take + + return deficits + + def _compute_fill_deficits(self, *, plan: SlotsPlan, remaining: int, deficit_slots: List[int]) -> Dict[int, int]: + to_fill = int(remaining) + if to_fill <= 0: + return {} + + owner_id = deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) + return {owner_id: to_fill} if owner_id is not None else {} + + async def _refill_deficits( + self, + *, + plan: SlotsPlan, + deficits: Dict[int, int], + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + dedup_policy: Any, + refill_settings: Any, + cursor: CursorMap, + ) -> None: + overfetch_factor = max(1, int(getattr(refill_settings, "overfetch_factor", 1))) + max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) + + deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] + deficit_owners = sorted( + deficit_owners, + key=lambda o: ( + int(getattr(o, "dedup_priority", 0)), + owner_index.get(id(o), 0), + ), + ) + + state: Dict[int, Dict[str, Any]] = {} + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + missing_total = int(deficits.get(refill_owner_id, 0)) + if missing_total <= 0: + continue + + base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page + state[refill_owner_id] = { + "owner": refill_owner, + "missing_total": missing_total, + "remaining": int(missing_total), + "accepted": [], + "loops": 0, + "current_next_page": base_np, + "has_next_page": True, + "last_result": None, + "last_request_limit": 0, + "last_can_overfetch": False, + "last_base_next_page": base_np, + } + + if not state: + return + + while True: + wave_ops: List[Tuple[Any, int, FeedResultNextPage, int, bool]] = [] + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + if owner_state["remaining"] <= 0: + continue + if not owner_state["has_next_page"]: + continue + if owner_state["loops"] >= max_refill_loops: + continue + + base_np = owner_state["current_next_page"] + remaining_before = max(1, int(owner_state["remaining"])) + request_limit = remaining_before + can_overfetch = CursorMap.can_overfetch(node=refill_owner, base_next_page=base_np) + if can_overfetch and overfetch_factor > 1: + request_limit = max(1, remaining_before * overfetch_factor) + + wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit, can_overfetch)) + + if not wave_ops: + break + + results = await self._executor.gather( + *[ + self._executor._run_owner( + plan=plan, + owner=owner, + demand=request_limit, + base_next_page=base_np, + dedup_active=True, + ) + for owner, _owner_id, base_np, request_limit, _can_overfetch in wave_ops + ] + ) + + for (owner, owner_id, base_np, request_limit, can_overfetch), result in zip(wave_ops, results): + owner_state = state[owner_id] + remaining_before = int(owner_state["remaining"]) + + owner_state["current_next_page"] = result.next_page + owner_state["has_next_page"] = bool(result.has_next_page) + cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) + + refill_prio = int(getattr(owner, "dedup_priority", 0)) + wave_accepted, inspected_count = await dedup_policy.accept_batch( + items=list(result.data), + priority=refill_prio, + limit=max(0, remaining_before), + ) + + if can_overfetch and request_limit > remaining_before: + CursorMap.rewind_overfetch( + node=owner, + base_next_page=base_np, + result_next_page=result.next_page, + inspected_count=inspected_count, + batch_size=len(result.data), + ) + + if wave_accepted: + owner_state["accepted"].extend(wave_accepted) + owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) + + if owner_state["remaining"] > 0 and owner_state["has_next_page"]: + owner_state["loops"] += 1 + + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + + accepted = owner_state["accepted"] + if accepted: + owner_buffers.setdefault(refill_owner_id, []) + owner_buffers[refill_owner_id].extend(accepted) + + owner_results[refill_owner_id] = FeedResult( + data=list(owner_buffers.get(refill_owner_id, [])), + next_page=owner_state["current_next_page"], + has_next_page=owner_state["has_next_page"], + ) diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index cd3ba4c..e4fb3ca 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -7,6 +7,7 @@ from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, _pydantic_deep_copy from .context import ExecutionContext from .cursors import CursorMap +from .dedup_runtime import DedupRuntime from .plans import CallablePlan, Plan, SlotSpec, SlotsPlan @@ -33,7 +34,21 @@ async def run( if isinstance(plan, SlotsPlan): return result - return await self._run_node_with_dedup_refill(node, ctx, limit, next_page, params, result) + return await self._dedup_runtime().run_node_with_dedup_refill( + node=node, + ctx=ctx, + limit=limit, + next_page=next_page, + params=params, + initial_result=result, + ) + + def _dedup_runtime(self) -> DedupRuntime: + runtime = getattr(self, "_dedup_runtime_instance", None) + if runtime is None: + runtime = DedupRuntime(self) + setattr(self, "_dedup_runtime_instance", runtime) + return runtime async def execute_plan(self, plan: Plan) -> FeedResult: """Interpret and execute a declarative plan. @@ -107,31 +122,17 @@ async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: ) if dedup_policy is not None: - owner_buffers, owner_results = await self._arbitrate_owner_buffers( + owner_buffers, owner_results = await self._dedup_runtime().apply_slots_plan_dedup( + plan=plan, owners=owners, owner_index=owner_index, owner_buffers=owner_buffers, owner_results=owner_results, dedup_policy=dedup_policy, + refill_settings=refill_settings, + cursor=cursor, ) - deficits = self._compute_slot_deficits( - plan=plan, - owner_buffers=owner_buffers, - ) - if deficits: - await self._refill_deficits( - plan=plan, - deficits=deficits, - owners=owners, - owner_index=owner_index, - owner_buffers=owner_buffers, - owner_results=owner_results, - dedup_policy=dedup_policy, - refill_settings=refill_settings, - cursor=cursor, - ) - output = self._consume_slots(plan=plan, owner_buffers=owner_buffers) assembled = await self._maybe_await(plan.assemble(output, cursor.next_page, owner_results)) return assembled @@ -225,254 +226,6 @@ async def _run_plan_owners( return owner_buffers, owner_results - async def _arbitrate_owner_buffers( - self, - *, - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - dedup_policy: Any, - ) -> tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: - owner_buffers = await dedup_policy.arbitrate_owner_buffers( - owners=owners, - owner_buffers=owner_buffers, - owner_rank=owner_index, - ) - - for owner in owners: - owner_id = id(owner) - if owner_id not in owner_results: - continue - old = owner_results[owner_id] - owner_results[owner_id] = FeedResult( - data=list(owner_buffers.get(owner_id, [])), - next_page=old.next_page, - has_next_page=old.has_next_page, - ) - - return owner_buffers, owner_results - - def _compute_slot_deficits( - self, - *, - plan: SlotsPlan, - owner_buffers: Dict[int, List[Any]], - ) -> Dict[int, int]: - total_max = sum(int(s.max_count) for s in plan.slots) - quota_schedule = total_max <= int(plan.limit) - - consumed: Dict[int, int] = {} - remaining = int(plan.limit) - deficit_slots: List[int] = [] - - for slot in plan.slots: - if remaining <= 0: - break - - owner_id = id(slot.owner) - want = min(int(slot.max_count), remaining) - if want <= 0: - continue - - have_total = len(owner_buffers.get(owner_id, [])) - already = int(consumed.get(owner_id, 0)) - available = max(0, have_total - already) - take = min(want, available) - if take < want: - deficit_slots.append(owner_id) - consumed[owner_id] = already + take - remaining -= take - - page_underfilled = remaining > 0 - - if quota_schedule: - return self._compute_quota_deficits(plan=plan, owner_buffers=owner_buffers) - if not page_underfilled: - return {} - return self._compute_fill_deficits(plan=plan, remaining=remaining, deficit_slots=deficit_slots) - - def _compute_quota_deficits( - self, - *, - plan: SlotsPlan, - owner_buffers: Dict[int, List[Any]], - ) -> Dict[int, int]: - deficits: Dict[int, int] = {} - remaining = int(plan.limit) - consumed: Dict[int, int] = {} - for slot in plan.slots: - if remaining <= 0: - break - - owner_id = id(slot.owner) - want = min(int(slot.max_count), remaining) - if want <= 0: - continue - - have_total = len(owner_buffers.get(owner_id, [])) - already = int(consumed.get(owner_id, 0)) - available = max(0, have_total - already) - take = min(want, available) - missing = max(0, want - take) - if missing: - deficits[owner_id] = deficits.get(owner_id, 0) + missing - consumed[owner_id] = already + take - remaining -= take - - return deficits - - def _compute_fill_deficits( - self, - *, - plan: SlotsPlan, - remaining: int, - deficit_slots: List[int], - ) -> Dict[int, int]: - to_fill = int(remaining) - if to_fill <= 0: - return {} - - owner_id = deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) - return {owner_id: to_fill} if owner_id is not None else {} - - async def _refill_deficits( - self, - *, - plan: SlotsPlan, - deficits: Dict[int, int], - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - dedup_policy: Any, - refill_settings: Any, - cursor: CursorMap, - ) -> None: - overfetch_factor = max(1, int(getattr(refill_settings, "overfetch_factor", 1))) - max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) - - deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] - deficit_owners = sorted( - deficit_owners, - key=lambda o: ( - int(getattr(o, "dedup_priority", 0)), - owner_index.get(id(o), 0), - ), - ) - - state: Dict[int, Dict[str, Any]] = {} - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - missing_total = int(deficits.get(refill_owner_id, 0)) - if missing_total <= 0: - continue - - base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page - state[refill_owner_id] = { - "owner": refill_owner, - "missing_total": missing_total, - "remaining": int(missing_total), - "accepted": [], - "loops": 0, - "current_next_page": base_np, - "has_next_page": True, - "last_result": None, - "last_request_limit": 0, - "last_can_overfetch": False, - "last_base_next_page": base_np, - } - - if not state: - return - - while True: - wave_ops: List[Tuple[Any, int, FeedResultNextPage, int, bool]] = [] - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - if owner_state["remaining"] <= 0: - continue - if not owner_state["has_next_page"]: - continue - if owner_state["loops"] >= max_refill_loops: - continue - - base_np = owner_state["current_next_page"] - remaining_before = max(1, int(owner_state["remaining"])) - request_limit = remaining_before - can_overfetch = CursorMap.can_overfetch(node=refill_owner, base_next_page=base_np) - if can_overfetch and overfetch_factor > 1: - request_limit = max(1, remaining_before * overfetch_factor) - - wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit, can_overfetch)) - - if not wave_ops: - break - - results = await self.gather( - *[ - self._run_owner( - plan=plan, - owner=owner, - demand=request_limit, - base_next_page=base_np, - dedup_active=True, - ) - for owner, _owner_id, base_np, request_limit, _can_overfetch in wave_ops - ] - ) - - for (owner, owner_id, base_np, request_limit, can_overfetch), result in zip(wave_ops, results): - owner_state = state[owner_id] - remaining_before = int(owner_state["remaining"]) - - owner_state["current_next_page"] = result.next_page - owner_state["has_next_page"] = bool(result.has_next_page) - cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) - - refill_prio = int(getattr(owner, "dedup_priority", 0)) - wave_accepted, inspected_count = await dedup_policy.accept_batch( - items=list(result.data), - priority=refill_prio, - limit=max(0, remaining_before), - ) - - if can_overfetch and request_limit > remaining_before: - CursorMap.rewind_overfetch( - node=owner, - base_next_page=base_np, - result_next_page=result.next_page, - inspected_count=inspected_count, - batch_size=len(result.data), - ) - - if wave_accepted: - owner_state["accepted"].extend(wave_accepted) - owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) - - if owner_state["remaining"] > 0 and owner_state["has_next_page"]: - owner_state["loops"] += 1 - - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - - accepted = owner_state["accepted"] - if accepted: - owner_buffers.setdefault(refill_owner_id, []) - owner_buffers[refill_owner_id].extend(accepted) - - owner_results[refill_owner_id] = FeedResult( - data=list(owner_buffers.get(refill_owner_id, [])), - next_page=owner_state["current_next_page"], - has_next_page=owner_state["has_next_page"], - ) - def _consume_slots(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> List[Any]: output: List[Any] = [] for slot in plan.slots: @@ -494,83 +247,6 @@ def _consume_slots(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]] return output - async def _run_node_with_dedup_refill( - self, - node: BaseFeedConfigModel, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - params: Dict[str, Any], - initial_result: FeedResult, - ) -> FeedResult: - dedup = getattr(ctx, "dedup", None) - if dedup is None: - return initial_result - - settings = getattr(ctx, "refill_settings", None) or getattr(ctx, "dedup_settings", None) - overfetch_factor = max(1, int(getattr(settings, "overfetch_factor", 1))) - max_refill_loops = max(1, int(getattr(settings, "max_refill_loops", 20))) - priority = int(getattr(node, "dedup_priority", 0)) - - collected: List[Any] = [] - remaining = int(limit) - loops = 0 - - current_result = initial_result - current_next_page = current_result.next_page - current_request_limit = max(1, remaining) - has_next_page = bool(current_result.has_next_page) - base_next_page = next_page - - while remaining > 0: - can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) - - accepted, inspected_count = await dedup.accept_batch( - items=list(current_result.data), - priority=priority, - limit=remaining, - ) - - if can_overfetch and current_request_limit > remaining: - CursorMap.rewind_overfetch( - node=node, - base_next_page=base_next_page, - result_next_page=current_next_page, - inspected_count=inspected_count, - batch_size=len(current_result.data), - ) - - if accepted: - collected.extend(accepted) - remaining = limit - len(collected) - - if remaining <= 0 or not has_next_page or loops >= max_refill_loops: - break - loops += 1 - - base_next_page = current_next_page - next_request_limit = max(1, remaining) - can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) - if can_overfetch and overfetch_factor > 1: - next_request_limit = max(1, remaining * overfetch_factor) - - current_result, _plan = await self._run_node_raw( - node, - ctx, - next_request_limit, - base_next_page, - params, - ) - current_next_page = current_result.next_page - current_request_limit = next_request_limit - has_next_page = bool(current_result.has_next_page) - - return FeedResult( - data=collected, - next_page=current_next_page, - has_next_page=has_next_page, - ) - __all__ = [ "Executor", From fd2201b487004f103fc56c0b6363c0cb1dd8613b Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 01:20:19 +0000 Subject: [PATCH 20/41] More test coverage. --- smartfeed/policies/dedup_utils.py | 43 +++++++++++------ tests/test_dedup_policy_unit.py | 74 +++++++++++++++++++++++++++++ tests/test_dedup_utils.py | 78 +++++++++++++++++++++++++++++++ tests/test_seen_store_unit.py | 52 +++++++++++++++++++++ tests/test_view_session_unit.py | 54 +++++++++++++++++++++ 5 files changed, 287 insertions(+), 14 deletions(-) create mode 100644 tests/test_dedup_policy_unit.py create mode 100644 tests/test_dedup_utils.py create mode 100644 tests/test_seen_store_unit.py create mode 100644 tests/test_view_session_unit.py diff --git a/smartfeed/policies/dedup_utils.py b/smartfeed/policies/dedup_utils.py index 2559e0f..4483619 100644 --- a/smartfeed/policies/dedup_utils.py +++ b/smartfeed/policies/dedup_utils.py @@ -12,32 +12,47 @@ from ..feed_models import _is_async_redis_client, _redis_call +def _seen_entries_to_map(entries: Any) -> Dict[str, int]: + """Coerce a legacy cursor "seen" list into a {key: priority} map. + + Supports: + - ["k1", "k2", ...] (implies priority 0) + - [["k1", 10], ["k2", 3], ...] (explicit priorities) + """ + + seen_map: Dict[str, int] = {} + if not isinstance(entries, list): + return seen_map + + for entry_item in entries: + if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: + seen_map[str(entry_item[0])] = int(entry_item[1]) + else: + seen_map[str(entry_item)] = 0 + return seen_map + + def decode_seen_from_cursor(after: Any) -> Dict[str, int]: if after is None: return {} if isinstance(after, dict) and "z" in after: + if after.get("v") != 2: + return {} + if after.get("c") != "zlib+base64": + return {} payload = base64.urlsafe_b64decode(str(after["z"]).encode()) raw = zlib.decompress(payload).decode() decoded = json.loads(raw) - if isinstance(decoded, dict): - return {str(k): int(v) for k, v in decoded.items()} if isinstance(decoded, list): - seen_map: Dict[str, int] = {} - for entry_item in decoded: - if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: - seen_map[str(entry_item[0])] = int(entry_item[1]) - else: - seen_map[str(entry_item)] = 0 - return seen_map + return _seen_entries_to_map(decoded) return {} if isinstance(after, dict) and "seen" in after: - return {str(k): 0 for k in list(after["seen"])} - if isinstance(after, list): - return {str(k): 0 for k in list(after)} - if isinstance(after, dict): - return {str(k): int(v) for k, v in after.items() if k not in {"v", "c", "n"}} + if after.get("v") != 2: + return {} + return _seen_entries_to_map(list(after["seen"])) + return {} diff --git a/tests/test_dedup_policy_unit.py b/tests/test_dedup_policy_unit.py new file mode 100644 index 0000000..b846f6b --- /dev/null +++ b/tests/test_dedup_policy_unit.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import pytest + +from smartfeed.policies.dedup import DeduplicationPolicy +from smartfeed.policies.seen_store import CursorSeenStore + + +def _policy(*, dedup_key: Optional[str], missing_key_policy: str, after: Any = None) -> DeduplicationPolicy: + store = CursorSeenStore.from_after(after=after, cursor_compress=False, cursor_max_keys=None) + return DeduplicationPolicy( + dedup_key=dedup_key, + missing_key_policy=missing_key_policy, # type: ignore[arg-type] + store=store, + seen_request_set=set(), + ) + + +@pytest.mark.asyncio +async def test_accept_batch_dedups_within_request_and_respects_existing_priority() -> None: + # Existing seen key with high priority should block lower/equal priority + existing_after = {"v": 2, "seen": [["1", 10]]} + policy = _policy(dedup_key="id", missing_key_policy="keep", after=existing_after) + + items = [{"id": 1}, {"id": 1}, {"id": 2}] + accepted, inspected = await policy.accept_batch(items=items, priority=5) + + assert inspected == 3 + assert accepted == [{"id": 2}] + + +@pytest.mark.asyncio +async def test_accept_batch_missing_key_policies() -> None: + items = [{"id": 1}, {"nope": 2}] + + policy_drop = _policy(dedup_key="id", missing_key_policy="drop") + accepted_drop, _ = await policy_drop.accept_batch(items=items, priority=0) + assert accepted_drop == [{"id": 1}] + + policy_error = _policy(dedup_key="id", missing_key_policy="error") + with pytest.raises(AssertionError): + await policy_error.accept_batch(items=items, priority=0) + + +@dataclass +class _Owner: + dedup_priority: int + + +@pytest.mark.asyncio +async def test_arbitrate_owner_buffers_prefers_higher_priority_owner() -> None: + policy = _policy(dedup_key="id", missing_key_policy="keep") + + owner_low = _Owner(dedup_priority=1) + owner_high = _Owner(dedup_priority=2) + + owners = [owner_low, owner_high] + owner_rank: Dict[int, int] = {id(owner_low): 0, id(owner_high): 1} + + shared = {"id": "same"} + owner_buffers: Dict[int, List[Any]] = { + id(owner_low): [shared, {"id": "low_only"}], + id(owner_high): [shared, {"id": "high_only"}], + } + + per_owner = await policy.arbitrate_owner_buffers( + owners=owners, + owner_buffers=owner_buffers, + owner_rank=owner_rank, + ) + + assert per_owner[id(owner_high)] == [shared, {"id": "high_only"}] + assert per_owner[id(owner_low)] == [{"id": "low_only"}] diff --git a/tests/test_dedup_utils.py b/tests/test_dedup_utils.py new file mode 100644 index 0000000..9556af3 --- /dev/null +++ b/tests/test_dedup_utils.py @@ -0,0 +1,78 @@ +from typing import Any, List + +import pytest + +from smartfeed.feed_models import _redis_call +from smartfeed.policies.dedup_utils import ( + decode_seen_from_cursor, + encode_seen_for_cursor, + redis_zmscore, +) + +from tests.fixtures.redis import redis_client + + +class _RedisNoZmscore: + """Wrapper around a real sync redis client to force the pipeline fallback. + + This keeps the backend "real Redis" while exercising the no-zmscore branch. + """ + + zmscore = None # type: ignore[assignment] + + def __init__(self, client: Any) -> None: + self._client = client + + def pipeline(self) -> Any: + return self._client.pipeline() + + +def test_encode_decode_seen_cursor_compressed_roundtrip_and_truncation() -> None: + seen_updates = [("a", 1), ("b", 2), ("c", 3)] + + encoded = encode_seen_for_cursor(seen_updates, cursor_compress=True, cursor_max_keys=None) + decoded = decode_seen_from_cursor(encoded) + assert decoded == {"a": 1, "b": 2, "c": 3} + + encoded_trunc = encode_seen_for_cursor(seen_updates, cursor_compress=True, cursor_max_keys=2) + decoded_trunc = decode_seen_from_cursor(encoded_trunc) + assert decoded_trunc == {"b": 2, "c": 3} + + +def test_decode_seen_cursor_v2_only() -> None: + assert decode_seen_from_cursor(None) == {} + + # Supported v2 uncompressed format + assert decode_seen_from_cursor({"v": 2, "seen": [["a", 9], ["b", 1]]}) == {"a": 9, "b": 1} + + # Legacy/unknown shapes are intentionally rejected + assert decode_seen_from_cursor(["x", "y"]) == {} + assert decode_seen_from_cursor({"a": 1}) == {} + assert decode_seen_from_cursor({"v": 1, "seen": [["a", 1]]}) == {} + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_redis_zmscore_native(redis_client) -> None: + key = "test_zmscore_native" + await _redis_call(redis_client, "delete", key) + await _redis_call(redis_client, "zadd", key, mapping={"a": 1.0, "b": 2.0}) + + res = await redis_zmscore(redis_client, key, ["a", "missing", "b"]) + assert res == [1.0, None, 2.0] + + await _redis_call(redis_client, "delete", key) + + +@pytest.mark.parametrize("redis_client", ["sync"], indirect=True) +@pytest.mark.asyncio +async def test_redis_zmscore_pipeline_fallback_for_sync_client_without_zmscore(redis_client) -> None: + key = "test_zmscore_fallback" + await _redis_call(redis_client, "delete", key) + await _redis_call(redis_client, "zadd", key, mapping={"a": 1.0, "b": 2.0}) + + wrapped = _RedisNoZmscore(redis_client) + res = await redis_zmscore(wrapped, key, ["a", "missing", "b"]) + assert res == [1.0, None, 2.0] + + await _redis_call(redis_client, "delete", key) diff --git a/tests/test_seen_store_unit.py b/tests/test_seen_store_unit.py new file mode 100644 index 0000000..50dab93 --- /dev/null +++ b/tests/test_seen_store_unit.py @@ -0,0 +1,52 @@ +from typing import Any + +import pytest + +from smartfeed.feed_models import _redis_call +from smartfeed.policies.dedup_utils import decode_seen_from_cursor +from smartfeed.policies.seen_store import CursorSeenStore, RedisSeenStore + +from tests.fixtures.redis import redis_client + + +@pytest.mark.asyncio +async def test_cursor_seen_store_set_max_and_commit_roundtrip() -> None: + store = CursorSeenStore.from_after(after=None, cursor_compress=True, cursor_max_keys=None) + store.set_max("a", 1) + store.set_max("a", 1) # no-op + store.set_max("a", 0) # no-op (lower) + store.set_max("b", 2) + + after = await store.commit() + decoded = decode_seen_from_cursor(after) + assert decoded == {"a": 1, "b": 2} + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_redis_seen_store_prefetch_set_max_commit_and_reset(redis_client) -> None: + key = "test_seen_store" + await _redis_call(redis_client, "delete", key) + # Pre-seed zset state + await _redis_call(redis_client, "zadd", key, mapping={"a": 5.0}) + + store = RedisSeenStore.create(redis_client=redis_client, redis_key=key, ttl_seconds=60) + + await store.prefetch(["a", "a", "b"]) # duplicates + assert store.get("a") == 5 + assert store.get("b") is None + + store.set_max("a", 3) # should not reduce existing + store.set_max("b", 2) + + await store.commit() + + # New state should be present in redis + scores = list(await _redis_call(redis_client, "zmscore", key, ["a", "b"])) + assert scores == [5.0, 2.0] + + await store.reset() + scores_after_reset = list(await _redis_call(redis_client, "zmscore", key, ["a", "b"])) + assert scores_after_reset == [None, None] + + await _redis_call(redis_client, "delete", key) diff --git a/tests/test_view_session_unit.py b/tests/test_view_session_unit.py new file mode 100644 index 0000000..a847eea --- /dev/null +++ b/tests/test_view_session_unit.py @@ -0,0 +1,54 @@ +from dataclasses import dataclass +from typing import Any + +import pytest + +from smartfeed.feed_models import _redis_call +from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerViewSession +from tests.fixtures.configs import METHODS_DICT +from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG +from tests.fixtures.redis import redis_client +from tests.utils import parse_model + + +@dataclass +class _ItemWithAttr: + id: str + + +def test_get_dedup_key_supports_dict_and_attr_and_raises_on_missing() -> None: + cfg = dict(MERGER_VIEW_SESSION_CONFIG) + cfg.update({"deduplicate": True, "dedup_key": "id"}) + merger = parse_model(MergerViewSession, cfg) + + assert merger._get_dedup_key_or_attr({"id": "x"}) == "x" + assert merger._get_dedup_key_or_attr(_ItemWithAttr(id="y")) == "y" + + with pytest.raises(AssertionError): + merger._get_dedup_key_or_attr({"nope": 1}) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_view_session_shuffle_applies_to_result(redis_client, monkeypatch) -> None: + import smartfeed.mergers.view_session as vs_mod + + cfg = dict(MERGER_VIEW_SESSION_CONFIG) + cfg.update({"shuffle": True}) + merger = parse_model(MergerViewSession, cfg) + cache_key = f"{merger.merger_id}_x" + await _redis_call(redis_client, "delete", cache_key) + + # Make shuffle deterministic: reverse in-place + monkeypatch.setattr(vs_mod, "shuffle", lambda data: data.reverse()) + + res = await merger.get_data( + methods_dict=METHODS_DICT, + limit=5, + next_page=FeedResultNextPage(data={}), + user_id="x", + redis_client=redis_client, + ) + + assert res.data == ["x_5", "x_4", "x_3", "x_2", "x_1"] + await _redis_call(redis_client, "delete", cache_key) From 65601b947cbe9f1ab5debb0e12b27f433f3a0919 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 01:26:41 +0000 Subject: [PATCH 21/41] Even more test coverage. --- tests/test_cursor_and_refill_edges.py | 92 ++++++++ tests/test_executor_slots_plan_invariants.py | 208 +++++++++++++++++++ 2 files changed, 300 insertions(+) create mode 100644 tests/test_cursor_and_refill_edges.py create mode 100644 tests/test_executor_slots_plan_invariants.py diff --git a/tests/test_cursor_and_refill_edges.py b/tests/test_cursor_and_refill_edges.py new file mode 100644 index 0000000..27b5b72 --- /dev/null +++ b/tests/test_cursor_and_refill_edges.py @@ -0,0 +1,92 @@ +import base64 +import zlib + +import pytest + +from smartfeed.execution.context import ExecutionContext, RefillExecutionSettings +from smartfeed.execution.executor import Executor +from smartfeed.feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage +from smartfeed.policies.dedup import DeduplicationPolicy +from smartfeed.policies.dedup_utils import decode_seen_from_cursor +from smartfeed.policies.seen_store import CursorSeenStore + + +class _DuplicateOnlyNode(BaseFeedConfigModel): + """A node that always returns the same duplicate item and never ends.""" + + type: str = "test_node" # satisfies pydantic + + def __init__(self, **data): + super().__init__(**data) + object.__setattr__(self, "calls", 0) + + async def get_data( # type: ignore[override] + self, + methods_dict, + user_id, + limit, + next_page, + redis_client=None, + ctx=None, + **params, + ) -> FeedResult: + object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) + return FeedResult( + data=[{"id": "dup"} for _ in range(int(limit) or 1)], + next_page=FeedResultNextPage(data={}), + has_next_page=True, + ) + + +def _policy_with_seen_dup(*, max_refill_loops: int) -> tuple[DeduplicationPolicy, RefillExecutionSettings]: + store = CursorSeenStore.from_after( + after={"v": 2, "seen": [["dup", 10]]}, + cursor_compress=False, + cursor_max_keys=None, + ) + policy = DeduplicationPolicy( + dedup_key="id", + missing_key_policy="keep", + store=store, + seen_request_set=set(), + ) + settings = RefillExecutionSettings(overfetch_factor=1, max_refill_loops=max_refill_loops) + return policy, settings + + +@pytest.mark.asyncio +async def test_dedup_refill_stops_at_max_loops_when_only_duplicates() -> None: + node = _DuplicateOnlyNode() + executor = Executor() + + dedup_policy, settings = _policy_with_seen_dup(max_refill_loops=2) + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + ctx.dedup = dedup_policy + ctx.refill_settings = settings + + res = await executor.run(node, ctx, limit=3, next_page=FeedResultNextPage(data={})) + + assert res.data == [] + # initial call + 2 refill loops + assert getattr(node, "calls") == 3 + + +def test_decode_seen_from_cursor_raises_on_corrupt_compressed_payload() -> None: + # invalid base64 + with pytest.raises(Exception): + decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": "not-base64"}) + + # base64 ok, zlib invalid + bad_zlib = base64.urlsafe_b64encode(b"not-a-zlib-stream").decode("ascii") + with pytest.raises(Exception): + decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": bad_zlib}) + + # zlib ok, json invalid + bad_json = base64.urlsafe_b64encode(zlib.compress(b"not json")).decode("ascii") + with pytest.raises(Exception): + decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": bad_json}) + + +def test_decode_seen_from_cursor_rejects_wrong_version_or_codec() -> None: + assert decode_seen_from_cursor({"v": 1, "c": "zlib+base64", "z": ""}) == {} + assert decode_seen_from_cursor({"v": 2, "c": "other", "z": ""}) == {} diff --git a/tests/test_executor_slots_plan_invariants.py b/tests/test_executor_slots_plan_invariants.py new file mode 100644 index 0000000..8eb5175 --- /dev/null +++ b/tests/test_executor_slots_plan_invariants.py @@ -0,0 +1,208 @@ +import pytest + +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution.executor import Executor +from smartfeed.execution.plans import SlotSpec, SlotsPlan +from smartfeed.feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside +from smartfeed.policies.dedup import DeduplicationPolicy +from smartfeed.policies.seen_store import CursorSeenStore + + +class _Owner(BaseFeedConfigModel): + type: str = "test_owner" + + def __init__(self, *, name: str, **data): + super().__init__(**data) + object.__setattr__(self, "name", name) + object.__setattr__(self, "last_limit", None) + object.__setattr__(self, "calls", 0) + + async def get_data( # type: ignore[override] + self, + methods_dict, + user_id, + limit, + next_page, + redis_client=None, + ctx=None, + **params, + ) -> FeedResult: + object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) + object.__setattr__(self, "last_limit", int(limit)) + return FeedResult(data=[self.name] * int(limit), next_page=FeedResultNextPage(data={}), has_next_page=False) + + +class _PagedOwner(BaseFeedConfigModel): + type: str = "test_paged_owner" + subfeed_id: str + total: int = 10 + + def __init__(self, *, subfeed_id: str, total: int = 10, **data): + super().__init__(subfeed_id=subfeed_id, total=total, **data) + object.__setattr__(self, "calls", 0) + object.__setattr__(self, "limits", []) + + async def get_data( # type: ignore[override] + self, + methods_dict, + user_id, + limit, + next_page, + redis_client=None, + ctx=None, + **params, + ) -> FeedResult: + object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) + limits = list(getattr(self, "limits", [])) + limits.append(int(limit)) + object.__setattr__(self, "limits", limits) + + entry = next_page.data.get(self.subfeed_id) + offset = int(entry.after) if (entry is not None and isinstance(entry.after, int)) else 0 + + take = max(0, min(int(limit), int(self.total) - offset)) + data = [{"id": f"{self.subfeed_id}_{i}"} for i in range(offset + 1, offset + take + 1)] + new_after = offset + take + + next_page.data[self.subfeed_id] = FeedResultNextPageInside( + page=(entry.page + 1 if entry is not None else 2), + after=new_after, + ) + return FeedResult( + data=data, + next_page=next_page, + has_next_page=bool(new_after < int(self.total)), + ) + + +def _dedup_policy() -> DeduplicationPolicy: + store = CursorSeenStore.from_after(after=None, cursor_compress=False, cursor_max_keys=None) + return DeduplicationPolicy( + dedup_key="id", + missing_key_policy="keep", # type: ignore[arg-type] + store=store, + seen_request_set=set(), + ) + + +@pytest.mark.asyncio +async def test_slots_plan_limit_le_zero_calls_assemble_only() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + + owner = _Owner(name="x") + + called = {"assemble": 0} + + def assemble(output, next_page, owner_results): + called["assemble"] += 1 + assert output == [] + assert owner_results == {} + return FeedResult(data=output, next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=0, + next_page=FeedResultNextPage(data={}), + params={}, + slots=[SlotSpec(owner=owner, max_count=10)], + assemble=assemble, + ) + + res = await executor.execute_plan(plan) + assert res.data == [] + assert called["assemble"] == 1 + assert getattr(owner, "calls") == 0 + + +@pytest.mark.asyncio +async def test_slots_plan_owner_fetch_limits_overrides_demand() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + + owner = _Owner(name="x") + + def assemble(output, next_page, owner_results): + return FeedResult(data=output, next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=5, + next_page=FeedResultNextPage(data={}), + params={}, + slots=[SlotSpec(owner=owner, max_count=5)], + assemble=assemble, + owner_fetch_limits={id(owner): 1}, + ) + + res = await executor.execute_plan(plan) + assert getattr(owner, "calls") == 1 + assert getattr(owner, "last_limit") == 1 + assert res.data == ["x"] + + +@pytest.mark.asyncio +async def test_slots_plan_no_ops_path_assemble_still_runs() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + + owner = _Owner(name="x") + + called = {"assemble": 0} + + def assemble(output, next_page, owner_results): + called["assemble"] += 1 + assert output == [] + assert owner_results == {} + return FeedResult(data=[], next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=5, + next_page=FeedResultNextPage(data={}), + params={}, + slots=[SlotSpec(owner=owner, max_count=0)], + assemble=assemble, + ) + + res = await executor.execute_plan(plan) + assert res.data == [] + assert called["assemble"] == 1 + assert getattr(owner, "calls") == 0 + + +@pytest.mark.asyncio +async def test_slots_plan_quota_deficit_triggers_refill_wave() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + ctx.dedup = _dedup_policy() + + a = _PagedOwner(subfeed_id="a", total=10) + b = _PagedOwner(subfeed_id="b", total=10) + + def assemble(output, next_page, owner_results): + return FeedResult(data=output, next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=6, + next_page=FeedResultNextPage(data={ + "a": FeedResultNextPageInside(page=1, after=0), + "b": FeedResultNextPageInside(page=1, after=0), + }), + params={}, + slots=[ + SlotSpec(owner=a, max_count=3), + SlotSpec(owner=b, max_count=3), + ], + assemble=assemble, + # Force an initial under-fetch for owner a (quota deficit). + owner_fetch_limits={id(a): 1}, + ) + + res = await executor.execute_plan(plan) + + # a should be refilled from 1 -> 3 items + assert getattr(a, "calls") >= 2 + assert res.data[:3] == [{"id": "a_1"}, {"id": "a_2"}, {"id": "a_3"}] + assert res.data[3:] == [{"id": "b_1"}, {"id": "b_2"}, {"id": "b_3"}] From 86be2a457ac5ff55ca1d5de909fa1f70891b9232 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 15:07:11 +0000 Subject: [PATCH 22/41] Tests for when one subfeed is empty. --- tests/test_executor_slots_plan_invariants.py | 48 ++++++++++++++++++++ tests/test_merger_append.py | 22 +++++++++ tests/test_merger_percentage.py | 26 +++++++++++ 3 files changed, 96 insertions(+) diff --git a/tests/test_executor_slots_plan_invariants.py b/tests/test_executor_slots_plan_invariants.py index 8eb5175..74f8540 100644 --- a/tests/test_executor_slots_plan_invariants.py +++ b/tests/test_executor_slots_plan_invariants.py @@ -206,3 +206,51 @@ def assemble(output, next_page, owner_results): assert getattr(a, "calls") >= 2 assert res.data[:3] == [{"id": "a_1"}, {"id": "a_2"}, {"id": "a_3"}] assert res.data[3:] == [{"id": "b_1"}, {"id": "b_2"}, {"id": "b_3"}] + + +@pytest.mark.asyncio +async def test_slots_plan_quota_deficit_stops_refill_when_owner_exhausts() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + ctx.dedup = _dedup_policy() + + # Owner a can never satisfy its full slot quota. + a = _PagedOwner(subfeed_id="a", total=2) + b = _PagedOwner(subfeed_id="b", total=10) + + def assemble(output, next_page, owner_results): + return FeedResult(data=output, next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=6, + next_page=FeedResultNextPage( + data={ + "a": FeedResultNextPageInside(page=1, after=0), + "b": FeedResultNextPageInside(page=1, after=0), + } + ), + params={}, + slots=[ + SlotSpec(owner=a, max_count=3), + SlotSpec(owner=b, max_count=3), + ], + assemble=assemble, + # Force an initial under-fetch to create a quota deficit for a. + owner_fetch_limits={id(a): 1}, + ) + + res = await executor.execute_plan(plan) + + # a is exhausted after returning 2 total items; refill should stop. + assert getattr(a, "calls") == 2 + assert getattr(a, "limits") == [1, 2] + assert getattr(b, "calls") == 1 + + assert res.data == [ + {"id": "a_1"}, + {"id": "a_2"}, + {"id": "b_1"}, + {"id": "b_2"}, + {"id": "b_3"}, + ] diff --git a/tests/test_merger_append.py b/tests/test_merger_append.py index 309ea82..290dd04 100644 --- a/tests/test_merger_append.py +++ b/tests/test_merger_append.py @@ -1,3 +1,5 @@ +import copy + import pytest from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerAppend @@ -42,3 +44,23 @@ async def test_merger_append_with_item_1_page_2() -> None: assert merger_append_res.data == ["x_6", "x_7", "x_8", "x_9", "x_10", "x_1", "x_2", "x_3", "x_4", "x_5", "x_6"] assert merger_append_res.next_page.data["subfeed_merger_append_example"].page == 3 assert merger_append_res.next_page.data["subfeed_merger_append_example"].after == "x_10" + + +@pytest.mark.asyncio +async def test_merger_append_when_one_leaf_is_empty() -> None: + config = copy.deepcopy(MERGER_APPEND_CONFIG) + # Make the second leaf return no data + has_next_page=False. + config["items"][1]["method_name"] = "empty" + + merger_append = parse_model(MergerAppend, config) + res = await merger_append.get_data( + methods_dict=METHODS_DICT, + limit=11, + next_page=FeedResultNextPage(data={}), + user_id="x", + ) + + # Only the first subfeed contributes (it is capped to 5 by config). + assert res.data == ["x_1", "x_2", "x_3", "x_4", "x_5"] + # First subfeed's example method still reports more pages. + assert res.has_next_page is True diff --git a/tests/test_merger_percentage.py b/tests/test_merger_percentage.py index 89e225e..ae57f00 100644 --- a/tests/test_merger_percentage.py +++ b/tests/test_merger_percentage.py @@ -1,3 +1,5 @@ +import copy + import pytest from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPercentage @@ -26,3 +28,27 @@ async def test_merger_percentage() -> None: ) assert merger_percentage_res.data == ["x_4", "x_21", "x_22", "x_5", "x_23", "x_24", "x_6", "x_25", "x_26", "x_7"] + + +@pytest.mark.asyncio +async def test_merger_percentage_when_one_leaf_is_empty() -> None: + config = copy.deepcopy(MERGER_PERCENTAGE_CONFIG) + # Make the second leaf return no data + has_next_page=False. + config["items"][1]["data"]["method_name"] = "empty" + + merger_percentage = parse_model(MergerPercentage, config) + res = await merger_percentage.get_data( + methods_dict=METHODS_DICT, + limit=10, + next_page=FeedResultNextPage( + data={ + "subfeed_merger_percentage_example": FeedResultNextPageInside(page=2, after="x_3"), + "subfeed_2_merger_percentage_example": FeedResultNextPageInside(page=3, after="x_20"), + } + ), + user_id="x", + ) + + assert res.data == ["x_4", "x_5", "x_6", "x_7"] + # The non-empty leaf still reports more pages. + assert res.has_next_page is True From 218e927a5c555eb4f9095bcd2d9c1b34642c8a4c Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 21:50:23 +0000 Subject: [PATCH 23/41] Minor cleanup --- Makefile | 14 + smartfeed/execution/dedup_runtime.py | 89 ++--- smartfeed/execution/executor.py | 25 +- tests/test_async_loop_blocks_trace.py | 453 ++++++++++++++++++++++++++ 4 files changed, 498 insertions(+), 83 deletions(-) create mode 100644 tests/test_async_loop_blocks_trace.py diff --git a/Makefile b/Makefile index 7e9caf0..682148e 100644 --- a/Makefile +++ b/Makefile @@ -11,3 +11,17 @@ test: test_cache: pytest -s -vv -k "test_merger_view_session" + +.PHONY: test_async_chart charting + +# Runs only the async loop block + Chrome trace test. +# Writes trace.json next to this Makefile (project root). +test_async_chart: + rm -f ./trace.json + SMARTFEED_CHROME_TRACE=./trace.json pytest -q tests/test_async_loop_blocks_trace.py + @echo "\nWrote trace: $(CURDIR)/trace.json" + @echo "Open Chrome -> chrome://tracing -> Load -> select trace.json" + +# Convenience target: generate the trace + try to open chrome://tracing. +charting: test_async_chart + -@open -a "Google Chrome" "chrome://tracing" 2>/dev/null || true diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py index 15773a1..3181150 100644 --- a/smartfeed/execution/dedup_runtime.py +++ b/smartfeed/execution/dedup_runtime.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Tuple from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage from .context import ExecutionContext @@ -47,11 +47,9 @@ async def run_node_with_dedup_refill( remaining = int(limit) loops = 0 - current_result = initial_result - current_next_page = current_result.next_page - current_request_limit = max(1, remaining) - has_next_page = bool(current_result.has_next_page) base_next_page = next_page + current_result = initial_result + request_limit = max(1, remaining) # NOTE: Refill loops are inherently sequential for a single node because # each subsequent request depends on the previous cursor. @@ -64,11 +62,11 @@ async def run_node_with_dedup_refill( limit=remaining, ) - if can_overfetch and current_request_limit > remaining: + if can_overfetch and request_limit > remaining: CursorMap.rewind_overfetch( node=node, base_next_page=base_next_page, - result_next_page=current_next_page, + result_next_page=current_result.next_page, inspected_count=inspected_count, batch_size=len(current_result.data), ) @@ -77,31 +75,27 @@ async def run_node_with_dedup_refill( collected.extend(accepted) remaining = limit - len(collected) - if remaining <= 0 or not has_next_page or loops >= max_refill_loops: + if remaining <= 0 or not current_result.has_next_page or loops >= max_refill_loops: break loops += 1 - base_next_page = current_next_page - next_request_limit = max(1, remaining) - can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) - if can_overfetch and overfetch_factor > 1: - next_request_limit = max(1, remaining * overfetch_factor) + base_next_page = current_result.next_page + request_limit = max(1, remaining) + if CursorMap.can_overfetch(node=node, base_next_page=base_next_page) and overfetch_factor > 1: + request_limit = max(1, remaining * overfetch_factor) current_result, _plan = await self._executor._run_node_raw( node, ctx, - next_request_limit, + request_limit, base_next_page, params, ) - current_next_page = current_result.next_page - current_request_limit = next_request_limit - has_next_page = bool(current_result.has_next_page) return FeedResult( data=collected, - next_page=current_next_page, - has_next_page=has_next_page, + next_page=current_result.next_page, + has_next_page=bool(current_result.has_next_page), ) async def apply_slots_plan_dedup( @@ -150,43 +144,12 @@ async def apply_slots_plan_dedup( return owner_buffers, owner_results def _compute_slot_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: - total_max = sum(int(s.max_count) for s in plan.slots) - quota_schedule = total_max <= int(plan.limit) - + quota_schedule = sum(int(s.max_count) for s in plan.slots) <= int(plan.limit) + deficits: Dict[int, int] = {} consumed: Dict[int, int] = {} remaining = int(plan.limit) deficit_slots: List[int] = [] - for slot in plan.slots: - if remaining <= 0: - break - - owner_id = id(slot.owner) - want = min(int(slot.max_count), remaining) - if want <= 0: - continue - - have_total = len(owner_buffers.get(owner_id, [])) - already = int(consumed.get(owner_id, 0)) - available = max(0, have_total - already) - take = min(want, available) - if take < want: - deficit_slots.append(owner_id) - consumed[owner_id] = already + take - remaining -= take - - page_underfilled = remaining > 0 - - if quota_schedule: - return self._compute_quota_deficits(plan=plan, owner_buffers=owner_buffers) - if not page_underfilled: - return {} - return self._compute_fill_deficits(plan=plan, remaining=remaining, deficit_slots=deficit_slots) - - def _compute_quota_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: - deficits: Dict[int, int] = {} - remaining = int(plan.limit) - consumed: Dict[int, int] = {} for slot in plan.slots: if remaining <= 0: break @@ -202,19 +165,18 @@ def _compute_quota_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, L take = min(want, available) missing = max(0, want - take) if missing: - deficits[owner_id] = deficits.get(owner_id, 0) + missing + deficit_slots.append(owner_id) + if quota_schedule: + deficits[owner_id] = deficits.get(owner_id, 0) + missing consumed[owner_id] = already + take remaining -= take - return deficits - - def _compute_fill_deficits(self, *, plan: SlotsPlan, remaining: int, deficit_slots: List[int]) -> Dict[int, int]: - to_fill = int(remaining) - if to_fill <= 0: + if quota_schedule: + return deficits + if remaining <= 0: return {} - owner_id = deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) - return {owner_id: to_fill} if owner_id is not None else {} + return {owner_id: remaining} if owner_id is not None else {} async def _refill_deficits( self, @@ -250,17 +212,12 @@ async def _refill_deficits( base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page state[refill_owner_id] = { - "owner": refill_owner, "missing_total": missing_total, - "remaining": int(missing_total), + "remaining": missing_total, "accepted": [], "loops": 0, "current_next_page": base_np, "has_next_page": True, - "last_result": None, - "last_request_limit": 0, - "last_can_overfetch": False, - "last_base_next_page": base_np, } if not state: diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index e4fb3ca..22b6a8c 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -107,12 +107,11 @@ async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: working_next_page = _pydantic_deep_copy(plan.next_page) cursor = CursorMap(working_next_page) - owners, owner_index = self._collect_plan_owners(plan) + owners, owner_index, owner_max_demand = self._collect_plan_owners(plan) dedup_policy = getattr(plan.ctx, "dedup", None) refill_settings = getattr(plan.ctx, "refill_settings", None) or getattr(plan.ctx, "dedup_settings", None) dedup_active = dedup_policy is not None - owner_max_demand = self._owner_slot_demand(plan) owner_buffers, owner_results = await self._run_plan_owners( plan=plan, owners=owners, @@ -137,25 +136,17 @@ async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: assembled = await self._maybe_await(plan.assemble(output, cursor.next_page, owner_results)) return assembled - def _owner_slot_demand(self, plan: SlotsPlan) -> Dict[int, int]: - """Compute a per-owner maximum demand based on the slot schedule.""" - - demand: Dict[int, int] = {} - for slot in plan.slots: - owner_id = id(slot.owner) - demand[owner_id] = demand.get(owner_id, 0) + int(slot.max_count) - return demand - - def _collect_plan_owners(self, plan: SlotsPlan) -> tuple[List[Any], Dict[int, int]]: + def _collect_plan_owners(self, plan: SlotsPlan) -> tuple[List[Any], Dict[int, int], Dict[int, int]]: owners: List[Any] = [] owner_index: Dict[int, int] = {} + owner_demand: Dict[int, int] = {} for slot in plan.slots: owner_id = id(slot.owner) - if owner_id in owner_index: - continue - owner_index[owner_id] = len(owners) - owners.append(slot.owner) - return owners, owner_index + if owner_id not in owner_index: + owner_index[owner_id] = len(owners) + owners.append(slot.owner) + owner_demand[owner_id] = owner_demand.get(owner_id, 0) + int(slot.max_count) + return owners, owner_index, owner_demand async def _run_owner( self, diff --git a/tests/test_async_loop_blocks_trace.py b/tests/test_async_loop_blocks_trace.py new file mode 100644 index 0000000..2dc8def --- /dev/null +++ b/tests/test_async_loop_blocks_trace.py @@ -0,0 +1,453 @@ +import asyncio +import json +import os +import time +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable, Dict, List, Optional + +import pytest + +from smartfeed.schemas import FeedResultNextPage, MergerDeduplication +from tests.fixtures import dedup_helpers as dh +from tests.fixtures.redis import redis_client # noqa: F401 +from tests.utils import parse_model + + +def _now_us() -> int: + return time.perf_counter_ns() // 1000 + + +@dataclass +class ChromeTraceRecorder: + """Writes Chrome Trace Events JSON for chrome://tracing. + + This is intentionally tiny and test-only: no production dependencies. + """ + + pid: int = 1 + events: List[Dict[str, Any]] = field(default_factory=list) + + def _emit(self, event: Dict[str, Any]) -> None: + self.events.append(event) + + def begin(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: + self._emit( + { + "name": name, + "ph": "B", + "ts": int(_now_us() if ts_us is None else ts_us), + "pid": int(self.pid), + "tid": int(tid), + "args": args or {}, + } + ) + + def end(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: + self._emit( + { + "name": name, + "ph": "E", + "ts": int(_now_us() if ts_us is None else ts_us), + "pid": int(self.pid), + "tid": int(tid), + "args": args or {}, + } + ) + + def instant(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: + self._emit( + { + "name": name, + "ph": "i", + "s": "t", + "ts": int(_now_us() if ts_us is None else ts_us), + "pid": int(self.pid), + "tid": int(tid), + "args": args or {}, + } + ) + + def write(self, path: str) -> None: + payload = {"traceEvents": self.events} + with open(path, "w", encoding="utf-8") as f: + json.dump(payload, f) + + +class LoopBlockMonitor: + """Detects event-loop blocking by measuring scheduling lag. + + If the event loop is blocked by long sync work, a periodic sleeper will wake + up late; we track the maximum observed lag. + """ + + def __init__(self, *, sample_interval_s: float = 0.01, block_threshold_s: float = 0.25) -> None: + self.sample_interval_s = float(sample_interval_s) + self.block_threshold_s = float(block_threshold_s) + self.max_lag_s: float = 0.0 + self.block_events: List[float] = [] + self._task: Optional[asyncio.Task[None]] = None + self._stop = asyncio.Event() + + async def __aenter__(self) -> "LoopBlockMonitor": + self._stop.clear() + self._task = asyncio.create_task(self._run()) + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: # type: ignore[override] + self._stop.set() + if self._task is not None: + await self._task + + async def _run(self) -> None: + loop = asyncio.get_running_loop() + expected = loop.time() + self.sample_interval_s + while not self._stop.is_set(): + await asyncio.sleep(self.sample_interval_s) + now = loop.time() + lag = max(0.0, now - expected) + expected = now + self.sample_interval_s + self.max_lag_s = max(self.max_lag_s, lag) + if lag >= self.block_threshold_s: + self.block_events.append(lag) + + +@dataclass +class LeafConcurrencyTracker: + """Tracks how many leaf calls are in-flight concurrently.""" + + current: int = 0 + peak: int = 0 + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + async def enter(self) -> int: + async with self._lock: + self.current += 1 + if self.current > self.peak: + self.peak = self.current + return self.current + + async def exit(self) -> int: + async with self._lock: + self.current = max(0, self.current - 1) + return self.current + + +def _trace_wrap_awaitable(rec: ChromeTraceRecorder, name: str, awaitable: Awaitable[Any], *, args: Dict[str, Any]) -> Awaitable[Any]: + async def _wrapped() -> Any: + task = asyncio.current_task() + tid = id(task) if task is not None else 0 + rec.begin(name, tid=tid, args=args) + try: + return await awaitable + finally: + rec.end(name, tid=tid) + + return _wrapped() + + +def _wrap_method_latency(method: Callable[..., Awaitable[Any]], *, latency_s: float) -> Callable[..., Awaitable[Any]]: + async def _wrapped(*args: Any, **kwargs: Any) -> Any: + await asyncio.sleep(latency_s) + return await method(*args, **kwargs) + + return _wrapped + + +def _wrap_leaf_method_traced( + *, + rec: ChromeTraceRecorder, + key: str, + method: Callable[..., Awaitable[Any]], + latency_s: float, + concurrency: LeafConcurrencyTracker, +) -> Callable[..., Awaitable[Any]]: + async def _wrapped(user_id: Any, limit: int, next_page: Any, **kwargs: Any) -> Any: + task = asyncio.current_task() + tid = id(task) if task is not None else 0 + + page = getattr(next_page, "page", None) + after = getattr(next_page, "after", None) + after_type = type(after).__name__ + + if after is None: + after_preview = None + else: + after_preview = str(after) + if len(after_preview) > 120: + after_preview = after_preview[:117] + "..." + + span = f"leaf.{key}" + + in_flight = await concurrency.enter() + rec.begin( + span, + tid=tid, + args={ + "key": key, + "limit": int(limit), + "page": page, + "after_type": after_type, + "after_preview": after_preview, + "in_flight": int(in_flight), + }, + ) + try: + if latency_s > 0: + await asyncio.sleep(float(latency_s)) + return await method(user_id, limit, next_page, **kwargs) + finally: + rec.end(span, tid=tid) + await concurrency.exit() + + return _wrapped + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_async_loop_blocks_and_trace_for_deep_tree_all_mergers(redis_client, monkeypatch, tmp_path) -> None: + """A smoke-test for detecting async loop blocks + visualizing concurrency. + + - Builds one deep tree that includes ALL merger types. + - Simulates 2 sequential requests (fresh + next page). + - Forces refills by creating lots of cross-branch duplicates. + - Records loop scheduling lag (blocks/hangs) and optionally exports a Chrome trace. + + Set `SMARTFEED_CHROME_TRACE=/path/to/trace.json` to write a trace. + Open it in Chrome via chrome://tracing. + """ + + # Keep IDs disjoint across sources so "no dupes" is stable. + # Refill waves are forced via max_per_call limits (under-fetch), not via dedup collisions. + items_a = dh.make_items("A", 1, 400, user_id_mod=5, id_offset=1_000) + items_b = dh.make_items("B", 1, 400, user_id_mod=5, id_offset=10_000) + + # Distribute branch: needs distribution_key present (user_id). + items_posted_1 = dh.make_items("posted_1", 1, 80, user_id_mod=3, id_offset=20_000) + items_posted_2 = dh.make_items("posted_2", 1, 120, user_id_mod=3, id_offset=21_000) + + # Gradient branch: overlapping ids again. + items_g1 = dh.make_items("G1", 1, 250, user_id_mod=7, id_offset=30_000) + items_g2 = dh.make_items("G2", 1, 250, user_id_mod=7, id_offset=40_000) + + # View-session leaf. + items_vs = dh.make_items("VS", 1, 160, user_id_mod=11, id_offset=50_000) + + # Positional leaf that intentionally under-fetches to force refill waves. + items_pos_leaf = dh.make_items("POS", 1, 500, user_id_mod=13, id_offset=60_000) + + # --- tracing (test-only monkeypatch) --- + rec = ChromeTraceRecorder() + leaf_concurrency = LeafConcurrencyTracker() + + # Leaf method tracing: wrap the *actual* subfeed method calls. + # These spans are what you want to inspect for "are leaf calls parallel?". + leaf_latency_s = 0.02 + methods_dict = { + "a": _wrap_leaf_method_traced( + rec=rec, + key="a", + method=dh.make_offset_paged_method(items_a), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "b": _wrap_leaf_method_traced( + rec=rec, + key="b", + method=dh.make_offset_paged_method(items_b), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "posted_1": _wrap_leaf_method_traced( + rec=rec, + key="posted_1", + method=dh.make_offset_paged_method(items_posted_1), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "posted_2": _wrap_leaf_method_traced( + rec=rec, + key="posted_2", + method=dh.make_offset_paged_method(items_posted_2), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "g1": _wrap_leaf_method_traced( + rec=rec, + key="g1", + method=dh.make_offset_paged_method(items_g1), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "g2": _wrap_leaf_method_traced( + rec=rec, + key="g2", + method=dh.make_offset_paged_method(items_g2), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + "vs": _wrap_leaf_method_traced( + rec=rec, + key="vs", + method=dh.make_offset_paged_method(items_vs), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + # Fetch only 1 item per call even if demand is higher -> triggers refill loops. + "pos_leaf": _wrap_leaf_method_traced( + rec=rec, + key="pos_leaf", + method=dh.make_offset_paged_method(items_pos_leaf, max_per_call=1), + latency_s=leaf_latency_s, + concurrency=leaf_concurrency, + ), + } + + view_session_cfg = { + "merger_id": "vs_all", + "type": "merger_view_session", + "session_size": 100, + "session_live_time": 60, + "deduplicate": True, + "dedup_key": "id", + "data": dh._subfeed("sf_vs", "vs"), + } + + pct_cfg = dh._percentage_config( + "pct_all", + items=dh._percentage_items(dh._subfeed("sf_a", "a"), dh._subfeed("sf_b", "b"), first_pct=50, second_pct=50), + ) + + pos_cfg = dh._positional_config( + "pos_all", + # Ensure positional inserts appear across pages for limit~12. + # Use even positions so the schedule starts with the default branch; + # this keeps ordering deterministic. + positions=[2, 4, 6, 8, 10, 12, 14, 16, 18], + positional=dh._subfeed("sf_pos_leaf", "pos_leaf"), + default=pct_cfg, + ) + + dist_cfg = dh._distribute_config( + "dist_all", + items=[dh._subfeed("sf_posted_1", "posted_1"), dh._subfeed("sf_posted_2", "posted_2")], + distribution_key="user_id", + ) + + grad_cfg = dh._gradient_config( + "grad_all", + item_from={"percentage": 70, "data": dh._subfeed("sf_g1", "g1")}, + item_to={"percentage": 30, "data": dh._subfeed("sf_g2", "g2")}, + step=10, + size_to_step=5, + shuffle=False, + ) + + # Include all merger types as siblings so they are executed (and visible in trace), + # while keeping the main output driven by the first branch. + deep_tree = dh._append_config("append_all", [pos_cfg, view_session_cfg, dist_cfg, grad_cfg]) + config = dh._dedup_config( + "dedup_all", + deep_tree, + dedup_key="id", + state_backend="cursor", + overfetch_factor=3, + max_refill_loops=50, + ) + merger = parse_model(MergerDeduplication, config) + + # Patch Executor.gather to wrap each awaitable for Chrome trace. + from smartfeed.execution.executor import Executor # local import for monkeypatch + + original_gather = Executor.gather + + async def _gather_traced(self: Any, *coros: Any) -> List[Any]: + wrapped = [ + _trace_wrap_awaitable(rec, "executor.gather.op", c, args={"idx": i, "total": len(coros)}) + for i, c in enumerate(coros) + ] + task = asyncio.current_task() + tid = id(task) if task is not None else 0 + rec.begin("executor.gather", tid=tid, args={"n": len(coros)}) + try: + return await original_gather(self, *wrapped) + finally: + rec.end("executor.gather", tid=tid) + + monkeypatch.setattr(Executor, "gather", _gather_traced) + + # Patch Executor.run to show sequential refill loops vs plan execution. + original_run = Executor.run + + async def _run_traced( + self: Any, + node: Any, + ctx: Any, + limit: int, + next_page: Any, + **params: Any, + ) -> Any: + task = asyncio.current_task() + tid = id(task) if task is not None else 0 + node_type = getattr(node, "type", node.__class__.__name__) + node_id = getattr(node, "merger_id", getattr(node, "subfeed_id", None)) + rec.begin( + "executor.run_node", + tid=tid, + args={"node_type": node_type, "node_id": node_id, "limit": int(limit)}, + ) + try: + return await original_run(self, node, ctx, limit, next_page, **params) + finally: + rec.end("executor.run_node", tid=tid) + + monkeypatch.setattr(Executor, "run", _run_traced) + + # --- run: fresh request + next_page --- + limit = 12 + np0 = FeedResultNextPage(data={}) + + async with LoopBlockMonitor(sample_interval_s=0.01, block_threshold_s=0.05) as monitor: + res1 = await asyncio.wait_for( + merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=limit, + next_page=np0, + redis_client=redis_client, + ), + timeout=15, + ) + res2 = await asyncio.wait_for( + merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=limit, + next_page=res1.next_page, + redis_client=redis_client, + ), + timeout=15, + ) + + # Sanity: we should fill the page and maintain dedup invariants. + assert len(res1.data) == limit + assert len({x["id"] for x in res1.data}) == limit + assert len(res2.data) == limit + assert len({x["id"] for x in res2.data}) == limit + + # Hard assertion: leaf calls must overlap (async concurrency), not serialize. + assert leaf_concurrency.peak > 1 + + # Primary signal: event-loop should remain responsive under load. + assert monitor.max_lag_s < 0.1 + + out = os.environ.get("SMARTFEED_CHROME_TRACE") + if out: + # Allow writing to an explicit file path, or to a directory. + out_path = out + if os.path.isdir(out_path): + out_path = os.path.join(out_path, "smartfeed_trace.json") + rec.instant("loop.max_lag", tid=0, args={"max_lag_s": monitor.max_lag_s, "blocks": len(monitor.block_events)}) + rec.write(out_path) + + # Keep references so this test remains useful in local debugging. + _ = tmp_path From 4e6bfefa61daf20e0046b34ec875f552415466c6 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 22:32:30 +0000 Subject: [PATCH 24/41] Remove some boilerplate from mergers. --- smartfeed/execution/context.py | 1 - smartfeed/execution/dedup_runtime.py | 10 +++-- smartfeed/execution/executor.py | 3 +- smartfeed/feed_models.py | 21 ++++++++-- smartfeed/mergers/append.py | 23 +---------- smartfeed/mergers/append_distribute.py | 23 +---------- smartfeed/mergers/deduplication.py | 23 +---------- smartfeed/mergers/percentage.py | 22 +---------- smartfeed/mergers/percentage_gradient.py | 22 +---------- smartfeed/mergers/positional.py | 22 +---------- smartfeed/mergers/view_session.py | 50 +++++++++--------------- smartfeed/policies/seen_store.py | 16 +++++--- tests/test_merger_deduplication.py | 44 +++++++++++++++++++++ tests/test_seen_store_unit.py | 14 +++++++ 14 files changed, 117 insertions(+), 177 deletions(-) diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py index 134c261..18fc769 100644 --- a/smartfeed/execution/context.py +++ b/smartfeed/execution/context.py @@ -26,7 +26,6 @@ class ExecutionContext: # Execution settings (optional) refill_settings: Optional["RefillExecutionSettings"] = None - dedup_settings: Optional["RefillExecutionSettings"] = None def ensure_redis_client(self, redis_client: Optional[Union[redis.Redis, AsyncRedis]]) -> None: if self.redis_client is None and redis_client is not None: diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py index 3181150..90dadc7 100644 --- a/smartfeed/execution/dedup_runtime.py +++ b/smartfeed/execution/dedup_runtime.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, List, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage from .context import ExecutionContext @@ -22,7 +22,7 @@ def __init__(self, executor: "Executor") -> None: self._executor = executor def _get_refill_settings(self, ctx: ExecutionContext) -> Any: - return getattr(ctx, "refill_settings", None) or getattr(ctx, "dedup_settings", None) + return getattr(ctx, "refill_settings", None) async def run_node_with_dedup_refill( self, @@ -175,8 +175,10 @@ def _compute_slot_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, Li return deficits if remaining <= 0: return {} - owner_id = deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) - return {owner_id: remaining} if owner_id is not None else {} + fallback_owner_id: Optional[int] = ( + deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) + ) + return {fallback_owner_id: remaining} if fallback_owner_id is not None else {} async def _refill_deficits( self, diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index 22b6a8c..11eb551 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -109,7 +109,7 @@ async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: cursor = CursorMap(working_next_page) owners, owner_index, owner_max_demand = self._collect_plan_owners(plan) dedup_policy = getattr(plan.ctx, "dedup", None) - refill_settings = getattr(plan.ctx, "refill_settings", None) or getattr(plan.ctx, "dedup_settings", None) + refill_settings = getattr(plan.ctx, "refill_settings", None) dedup_active = dedup_policy is not None owner_buffers, owner_results = await self._run_plan_owners( @@ -167,7 +167,6 @@ async def _run_owner( executor=plan.ctx.executor, dedup=None, refill_settings=None, - dedup_settings=None, ) return await self.run(owner, owner_ctx, demand, isolated_next_page, **plan.params) diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py index 6d29d67..38d26d6 100644 --- a/smartfeed/feed_models.py +++ b/smartfeed/feed_models.py @@ -1,6 +1,5 @@ import asyncio import inspect -from abc import ABC, abstractmethod from dataclasses import dataclass from random import shuffle from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast @@ -68,14 +67,13 @@ class FeedResultClient(BaseModel): has_next_page: bool -class BaseFeedConfigModel(ABC, BaseModel): +class BaseFeedConfigModel(BaseModel): """Base class for merger/subfeed config models.""" # Higher value means the item should "win" deduplication when duplicates exist. # This is primarily used by MergerDeduplication and by mergers when a dedup wrapper is active. dedup_priority: int = 0 - @abstractmethod async def get_data( self, methods_dict: Dict[str, Callable], @@ -86,7 +84,22 @@ async def get_data( ctx: Optional["ExecutionContext"] = None, **params: Any, ) -> FeedResult: - """Fetch data according to this node config.""" + """Default merger execution path via the shared executor.""" + + if not callable(getattr(self, "build_plan", None)): + raise NotImplementedError( + f"{self.__class__.__name__} must implement build_plan(...) or override get_data(...)." + ) + + if ctx is None: + from .execution.context import ExecutionContext as _ExecutionContext + + ctx = _ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) + else: + ctx.ensure_redis_client(redis_client) + + executor = ctx.ensure_executor() + return await executor.run(self, ctx, limit, next_page, **params) @dataclass diff --git a/smartfeed/mergers/append.py b/smartfeed/mergers/append.py index a415159..9c5c5c6 100644 --- a/smartfeed/mergers/append.py +++ b/smartfeed/mergers/append.py @@ -1,10 +1,7 @@ from __future__ import annotations from random import shuffle -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast - -import redis -from redis.asyncio import Redis as AsyncRedis +from typing import TYPE_CHECKING, Any, Dict, List, Literal, cast from ..execution.context import ExecutionContext from ..execution.executor import SlotSpec, SlotsPlan @@ -49,21 +46,3 @@ def _assemble( slots=slots, assemble=_assemble, ) - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py index 3e0a8e1..3ef891d 100644 --- a/smartfeed/mergers/append_distribute.py +++ b/smartfeed/mergers/append_distribute.py @@ -1,10 +1,7 @@ from __future__ import annotations from collections import defaultdict, deque -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union - -import redis -from redis.asyncio import Redis as AsyncRedis +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional from typing_extensions import no_type_check from ..execution.context import ExecutionContext @@ -73,21 +70,3 @@ def _assemble( slots=slots, assemble=_assemble, ) - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) diff --git a/smartfeed/mergers/deduplication.py b/smartfeed/mergers/deduplication.py index 2686ac9..1b35a7b 100644 --- a/smartfeed/mergers/deduplication.py +++ b/smartfeed/mergers/deduplication.py @@ -1,10 +1,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional -import redis from pydantic import PrivateAttr, model_validator -from redis.asyncio import Redis as AsyncRedis from ..execution.context import ExecutionContext, RefillExecutionSettings from ..execution.cursors import CursorMap @@ -102,24 +100,6 @@ def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: return f"dedup:{self.merger_id}:{user_id}:{suffix}" return f"dedup:{self.merger_id}:{user_id}" - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - def build_plan( self, *, @@ -185,7 +165,6 @@ async def _run(executor: Any) -> FeedResult: executor=ctx.executor, dedup=policy, refill_settings=refill_settings, - dedup_settings=refill_settings, ) child = _pydantic_deep_copy(self.data) diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py index 1b034b4..a21fbd3 100644 --- a/smartfeed/mergers/percentage.py +++ b/smartfeed/mergers/percentage.py @@ -1,11 +1,9 @@ from __future__ import annotations from random import shuffle -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, cast -import redis from pydantic import BaseModel -from redis.asyncio import Redis as AsyncRedis from ..execution.context import ExecutionContext from ..execution.executor import SlotSpec, SlotsPlan @@ -56,24 +54,6 @@ def _merge_items_data(items_data: List[List]) -> List: return result - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - def build_plan( self, *, diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py index b1ff586..1aea051 100644 --- a/smartfeed/mergers/percentage_gradient.py +++ b/smartfeed/mergers/percentage_gradient.py @@ -1,9 +1,7 @@ from random import shuffle -from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast +from typing import Any, Dict, List, Literal, cast -import redis from pydantic import model_validator -from redis.asyncio import Redis as AsyncRedis from ..execution.context import ExecutionContext from ..execution.executor import SlotSpec, SlotsPlan @@ -69,24 +67,6 @@ def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: return result - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - def build_plan( self, *, diff --git a/smartfeed/mergers/positional.py b/smartfeed/mergers/positional.py index 3ac9b32..6023aae 100644 --- a/smartfeed/mergers/positional.py +++ b/smartfeed/mergers/positional.py @@ -1,10 +1,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional -import redis from pydantic import model_validator -from redis.asyncio import Redis as AsyncRedis from ..execution.context import ExecutionContext from ..execution.executor import SlotSpec, SlotsPlan @@ -38,24 +36,6 @@ def validate_merger_positional(self) -> "MergerPositional": raise ValueError('"end" must be bigger than "start"') return self - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - def build_plan( self, *, diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index 8adf814..dad3314 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -2,7 +2,7 @@ import logging from random import shuffle -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union import redis from redis.asyncio import Redis as AsyncRedis @@ -11,6 +11,7 @@ from ..execution.context import ExecutionContext from ..execution.executor import CallablePlan from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside, _redis_call +from ..policies.dedup import entity_key if TYPE_CHECKING: from ..schemas import FeedTypes @@ -25,24 +26,27 @@ class MergerViewSession(BaseFeedConfigModel): session_live_time: int data: "FeedTypes" deduplicate: bool = False - dedup_key: str = None # type: ignore + dedup_key: Optional[str] = None + missing_key_policy: Literal["error", "keep", "drop"] = "error" shuffle: bool = False def _get_dedup_key_or_attr(self, item: Any) -> str: - if not self.dedup_key: - return item - - try: - dedup_value = item.get(self.dedup_key) - except AttributeError: - dedup_value = getattr(item, self.dedup_key, None) - - assert dedup_value is not None, f"Deduplication failed: entity {item} has no key or attr {self.dedup_key}" - return dedup_value + key = entity_key(item, self.dedup_key, self.missing_key_policy) + assert key is not None, "Deduplication key is missing and item was dropped by missing_key_policy='drop'" + return key def _dedup_data(self, data: List[Any]) -> List[Any]: - deduplicated_data = {self._get_dedup_key_or_attr(item): item for item in data} - return list(deduplicated_data.values()) + deduplicated: List[Any] = [] + seen: set[str] = set() + for item in data: + key = entity_key(item, self.dedup_key, self.missing_key_policy) + if key is None: + continue + if key in seen: + continue + seen.add(key) + deduplicated.append(item) + return deduplicated async def _set_cache( self, @@ -110,24 +114,6 @@ async def _get_cache( has_next_page=bool(len(session_data) > limit * page), ) - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional[ExecutionContext] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - ctx = ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - def build_plan( self, *, diff --git a/smartfeed/policies/seen_store.py b/smartfeed/policies/seen_store.py index d0a9258..c4712c3 100644 --- a/smartfeed/policies/seen_store.py +++ b/smartfeed/policies/seen_store.py @@ -35,7 +35,7 @@ class CursorSeenStore: cursor_max_keys: Optional[int] seen_priority_map: Dict[str, int] - seen_updates_in_order: List[Tuple[str, int]] + seen_order: List[str] @classmethod def from_after( @@ -50,7 +50,7 @@ def from_after( cursor_compress=cursor_compress, cursor_max_keys=cursor_max_keys, seen_priority_map=seen_priority_map, - seen_updates_in_order=[], + seen_order=list(seen_priority_map.keys()), ) async def prefetch(self, keys: List[str]) -> None: @@ -64,15 +64,21 @@ def set_max(self, key: str, priority: int) -> None: if existing is not None and priority <= existing: return self.seen_priority_map[key] = priority - self.seen_updates_in_order.append((key, priority)) + if key in self.seen_order: + self.seen_order.remove(key) + self.seen_order.append(key) async def reset(self) -> None: self.seen_priority_map.clear() - self.seen_updates_in_order.clear() + self.seen_order.clear() async def commit(self) -> Any: + # Persist the full snapshot so dedup state survives beyond 2 pages. + seen_snapshot_in_order: List[Tuple[str, int]] = [ + (key, self.seen_priority_map[key]) for key in self.seen_order if key in self.seen_priority_map + ] return encode_seen_for_cursor( - self.seen_updates_in_order, + seen_snapshot_in_order, cursor_compress=self.cursor_compress, cursor_max_keys=self.cursor_max_keys, ) diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index ab9838f..35fb3d1 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -318,6 +318,50 @@ async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: assert res_2.next_page.data["sf_stream"].page == 2 +@pytest.mark.asyncio +async def test_dedup_cursor_backend_persists_seen_state_beyond_two_pages() -> None: + # First 2 pages are unique, then page 3 starts with duplicates from page 1. + items = dh.make_items("S", 1, 11) + dh.make_items("S", 1, 4) + dh.make_items("S", 11, 31) + methods_dict = {"s": dh.make_offset_paged_method(items)} + + config = dh._dedup_config( + "dedup_cursor_3p", + dh._subfeed("sf_stream", "s"), + state_backend="cursor", + ) + merger = parse_model(MergerDeduplication, config) + + res_1 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=FeedResultNextPage(data={}), + ) + res_2 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=res_1.next_page, + ) + res_3 = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=5, + next_page=res_2.next_page, + ) + + ids_1 = dh._ids(res_1.data) + ids_2 = dh._ids(res_2.data) + ids_3 = dh._ids(res_3.data) + + assert ids_1 == [1, 2, 3, 4, 5] + assert ids_2 == [6, 7, 8, 9, 10] + assert ids_3 == [11, 12, 13, 14, 15] + assert not (set(ids_1) & set(ids_2)) + assert not (set(ids_1) & set(ids_3)) + assert not (set(ids_2) & set(ids_3)) + + @pytest.mark.asyncio async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf_cursor_exactly() -> None: """Append: across pages there is no overlap; refill advances cursors correctly. diff --git a/tests/test_seen_store_unit.py b/tests/test_seen_store_unit.py index 50dab93..f70661f 100644 --- a/tests/test_seen_store_unit.py +++ b/tests/test_seen_store_unit.py @@ -22,6 +22,20 @@ async def test_cursor_seen_store_set_max_and_commit_roundtrip() -> None: assert decoded == {"a": 1, "b": 2} +@pytest.mark.asyncio +async def test_cursor_seen_store_commit_keeps_previous_cursor_state() -> None: + store = CursorSeenStore.from_after( + after={"v": 2, "seen": [["a", 1], ["b", 2]]}, + cursor_compress=False, + cursor_max_keys=None, + ) + store.set_max("c", 3) + + after = await store.commit() + decoded = decode_seen_from_cursor(after) + assert decoded == {"a": 1, "b": 2, "c": 3} + + @pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) @pytest.mark.asyncio async def test_redis_seen_store_prefetch_set_max_commit_and_reset(redis_client) -> None: From e3310f06e3291286ad1fdca8897c2b2b58c227f4 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Sun, 8 Feb 2026 23:38:07 +0000 Subject: [PATCH 25/41] Patch for positional leak when underfetched. --- smartfeed/execution/dedup_runtime.py | 135 +++++++++++++++++++++++++++ smartfeed/execution/executor.py | 14 ++- tests/test_merger_deduplication.py | 62 ++++++++++++ 3 files changed, 210 insertions(+), 1 deletion(-) diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py index 90dadc7..1b9b26c 100644 --- a/smartfeed/execution/dedup_runtime.py +++ b/smartfeed/execution/dedup_runtime.py @@ -143,6 +143,32 @@ async def apply_slots_plan_dedup( return owner_buffers, owner_results + async def apply_slots_plan_refill( + self, + *, + plan: SlotsPlan, + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + refill_settings: Any, + cursor: CursorMap, + ) -> Tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: + deficits = self._compute_slot_deficits(plan=plan, owner_buffers=owner_buffers) + if deficits: + await self._refill_deficits_without_dedup( + plan=plan, + deficits=deficits, + owners=owners, + owner_index=owner_index, + owner_buffers=owner_buffers, + owner_results=owner_results, + refill_settings=refill_settings, + cursor=cursor, + ) + + return owner_buffers, owner_results + def _compute_slot_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: quota_schedule = sum(int(s.max_count) for s in plan.slots) <= int(plan.limit) deficits: Dict[int, int] = {} @@ -311,3 +337,112 @@ async def _refill_deficits( next_page=owner_state["current_next_page"], has_next_page=owner_state["has_next_page"], ) + + async def _refill_deficits_without_dedup( + self, + *, + plan: SlotsPlan, + deficits: Dict[int, int], + owners: List[Any], + owner_index: Dict[int, int], + owner_buffers: Dict[int, List[Any]], + owner_results: Dict[int, FeedResult], + refill_settings: Any, + cursor: CursorMap, + ) -> None: + max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) + + deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] + deficit_owners = sorted( + deficit_owners, + key=lambda o: ( + int(getattr(o, "dedup_priority", 0)), + owner_index.get(id(o), 0), + ), + ) + + state: Dict[int, Dict[str, Any]] = {} + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + missing_total = int(deficits.get(refill_owner_id, 0)) + if missing_total <= 0: + continue + + base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page + state[refill_owner_id] = { + "missing_total": missing_total, + "remaining": missing_total, + "accepted": [], + "loops": 0, + "current_next_page": base_np, + "has_next_page": True, + } + + if not state: + return + + while True: + wave_ops: List[Tuple[Any, int, FeedResultNextPage, int]] = [] + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + if owner_state["remaining"] <= 0: + continue + if not owner_state["has_next_page"]: + continue + if owner_state["loops"] >= max_refill_loops: + continue + + base_np = owner_state["current_next_page"] + request_limit = max(1, int(owner_state["remaining"])) + wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit)) + + if not wave_ops: + break + + results = await self._executor.gather( + *[ + self._executor._run_owner( + plan=plan, + owner=owner, + demand=request_limit, + base_next_page=base_np, + dedup_active=False, + ) + for owner, _owner_id, base_np, request_limit in wave_ops + ] + ) + + for (_owner, owner_id, _base_np, _request_limit), result in zip(wave_ops, results): + owner_state = state[owner_id] + remaining_before = int(owner_state["remaining"]) + + owner_state["current_next_page"] = result.next_page + owner_state["has_next_page"] = bool(result.has_next_page) + cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) + + if remaining_before > 0: + owner_state["accepted"].extend(list(result.data)[:remaining_before]) + owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) + + if owner_state["remaining"] > 0 and owner_state["has_next_page"]: + owner_state["loops"] += 1 + + for refill_owner in deficit_owners: + refill_owner_id = id(refill_owner) + owner_state = state.get(refill_owner_id) + if owner_state is None: + continue + + accepted = owner_state["accepted"] + if accepted: + owner_buffers.setdefault(refill_owner_id, []) + owner_buffers[refill_owner_id].extend(accepted) + + owner_results[refill_owner_id] = FeedResult( + data=list(owner_buffers.get(refill_owner_id, [])), + next_page=owner_state["current_next_page"], + has_next_page=owner_state["has_next_page"], + ) diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index 11eb551..d6a92e8 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -131,6 +131,16 @@ async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: refill_settings=refill_settings, cursor=cursor, ) + elif refill_settings is not None: + owner_buffers, owner_results = await self._dedup_runtime().apply_slots_plan_refill( + plan=plan, + owners=owners, + owner_index=owner_index, + owner_buffers=owner_buffers, + owner_results=owner_results, + refill_settings=refill_settings, + cursor=cursor, + ) output = self._consume_slots(plan=plan, owner_buffers=owner_buffers) assembled = await self._maybe_await(plan.assemble(output, cursor.next_page, owner_results)) @@ -166,7 +176,9 @@ async def _run_owner( redis_client=plan.ctx.redis_client, executor=plan.ctx.executor, dedup=None, - refill_settings=None, + # Keep refill settings so nested slots plans can still compensate + # owner deficits while top-level dedup arbitration remains centralized. + refill_settings=plan.ctx.refill_settings, ) return await self.run(owner, owner_ctx, demand, isolated_next_page, **plan.params) diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 35fb3d1..bbeb011 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -149,6 +149,68 @@ async def test_dedup_deep_tree_cursor_backend() -> None: dh._assert_sources_at_positions(res_2.data, [1, 4], "P") +@pytest.mark.asyncio +async def test_dedup_nested_positional_refill_not_masked_by_parent_append() -> None: + """Nested positional refills must run even when parent append can fill the page. + + Regression: + - parent dedup wrapper executes append owners with dedup disabled in owner ctx + - positional child under-fetches (`max_per_call=1`) and needs internal slot refills + - if those refills are skipped, append sibling backfills the page and positional slots are lost + """ + + items_default = dh.make_items("D", 1, 400, id_offset=1_000) + items_pos = dh.make_items("P", 1, 400, id_offset=10_000) + items_fill = dh.make_items("F", 1, 400, id_offset=20_000) + + pos_calls = {"count": 0} + pos_base = dh.make_offset_paged_method(items_pos, max_per_call=1) + + async def _pos_method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + pos_calls["count"] += 1 + return await pos_base(user_id, limit, next_page, **kwargs) + + methods_dict = { + "default": dh.make_offset_paged_method(items_default), + "pos": _pos_method, + "fill": dh.make_offset_paged_method(items_fill), + } + + config = dh._dedup_config( + "dedup_nested_refill", + dh._append_config( + "append_nested_refill", + [ + dh._positional_config( + "pos_nested_refill", + positions=[2, 4, 6, 8, 10, 12], + positional=dh._subfeed("sf_pos_nested", "pos"), + default=dh._subfeed("sf_default_nested", "default"), + ), + dh._subfeed("sf_fill_nested", "fill"), + ], + ), + dedup_key="id", + state_backend="cursor", + overfetch_factor=3, + max_refill_loops=50, + ) + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=12, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) == 12 + dh._assert_no_dupes_in_page(res.data) + dh._assert_sources_at_positions(res.data, [2, 4, 6, 8, 10, 12], "P") + assert "F" not in set(dh._sources(res.data)) + assert pos_calls["count"] > 1 + + @pytest.mark.parametrize( "merger_type", [ From 051373a3b4213b59c4aba8a789b3ad67919672fc Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Mon, 9 Feb 2026 00:56:27 +0000 Subject: [PATCH 26/41] Minor bugfixes and formatting. --- smartfeed/mergers/append_distribute.py | 1 + smartfeed/mergers/percentage.py | 27 +++++++-- smartfeed/mergers/percentage_gradient.py | 22 ++++++-- tests/test_async_loop_blocks_trace.py | 22 ++++++-- tests/test_dedup_utils.py | 7 +-- tests/test_executor_slots_plan_invariants.py | 51 +++++++++++++++-- tests/test_merger_deduplication.py | 59 ++++++++++++++++++++ tests/test_merger_percentage.py | 18 ++++++ tests/test_merger_percentage_gradient.py | 18 ++++++ tests/test_seen_store_unit.py | 1 - 10 files changed, 201 insertions(+), 25 deletions(-) diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py index 3ef891d..220e3e3 100644 --- a/smartfeed/mergers/append_distribute.py +++ b/smartfeed/mergers/append_distribute.py @@ -2,6 +2,7 @@ from collections import defaultdict, deque from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + from typing_extensions import no_type_check from ..execution.context import ExecutionContext diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py index a21fbd3..9ea2096 100644 --- a/smartfeed/mergers/percentage.py +++ b/smartfeed/mergers/percentage.py @@ -64,10 +64,29 @@ def build_plan( ) -> SlotsPlan: owners: List[BaseFeedConfigModel] = [cast(BaseFeedConfigModel, item.data) for item in self.items] - slots: List[SlotSpec] = [] - for item, owner in zip(self.items, owners): - child_limit = limit * int(item.percentage) // 100 - slots.append(SlotSpec(owner=owner, max_count=max(0, child_limit))) + slot_limits: List[int] = [] + remainders: List[tuple[int, int]] = [] + total_percentage = sum(int(item.percentage) for item in self.items) + + for idx, item in enumerate(self.items): + raw = int(limit) * int(item.percentage) + child_limit = raw // 100 + slot_limits.append(max(0, child_limit)) + remainders.append((raw % 100, idx)) + + # avoid underfilling for the common "percentages sum to 100" case + if total_percentage == 100: + missing = max(0, int(limit) - sum(slot_limits)) + if missing > 0: + for _rem, idx in sorted(remainders, key=lambda x: (-x[0], x[1])): + if missing <= 0: + break + slot_limits[idx] += 1 + missing -= 1 + + slots: List[SlotSpec] = [ + SlotSpec(owner=owner, max_count=max(0, int(slot_limits[idx]))) for idx, owner in enumerate(owners) + ] def _assemble( output: List[Any], diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py index 1aea051..fb70891 100644 --- a/smartfeed/mergers/percentage_gradient.py +++ b/smartfeed/mergers/percentage_gradient.py @@ -56,10 +56,19 @@ def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: if result["percentages"] and result["percentages"][-1]["to"] >= 100: result["limit_to"] += iter_limit result["percentages"][-1]["limit"] += iter_limit + result["percentages"][-1]["to_take"] += iter_limit else: - result["limit_from"] += iter_limit * percentage_from // 100 - result["limit_to"] += iter_limit * percentage_to // 100 - iter_result = {"limit": iter_limit, "from": percentage_from, "to": percentage_to} + from_take = iter_limit * percentage_from // 100 + to_take = iter_limit - from_take + result["limit_from"] += from_take + result["limit_to"] += to_take + iter_result = { + "limit": iter_limit, + "from": percentage_from, + "to": percentage_to, + "from_take": from_take, + "to_take": to_take, + } result["percentages"].append(iter_result) if first_iter: @@ -110,8 +119,11 @@ def _assemble( from_start_index = 0 to_start_index = 0 for lp_data in limits_and_percents["percentages"]: - from_end_index = (lp_data["limit"] * lp_data["from"] // 100) + from_start_index - to_end_index = (lp_data["limit"] * lp_data["to"] // 100) + to_start_index + from_take = int(lp_data.get("from_take", lp_data["limit"] * lp_data["from"] // 100)) + to_take = int(lp_data.get("to_take", lp_data["limit"] - from_take)) + + from_end_index = from_start_index + from_take + to_end_index = to_start_index + to_take data.extend(from_data[from_start_index:from_end_index]) data.extend(to_data[to_start_index:to_end_index]) diff --git a/tests/test_async_loop_blocks_trace.py b/tests/test_async_loop_blocks_trace.py index 2dc8def..a1bafd9 100644 --- a/tests/test_async_loop_blocks_trace.py +++ b/tests/test_async_loop_blocks_trace.py @@ -54,7 +54,9 @@ def end(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optiona } ) - def instant(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: + def instant( + self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None + ) -> None: self._emit( { "name": name, @@ -132,7 +134,9 @@ async def exit(self) -> int: return self.current -def _trace_wrap_awaitable(rec: ChromeTraceRecorder, name: str, awaitable: Awaitable[Any], *, args: Dict[str, Any]) -> Awaitable[Any]: +def _trace_wrap_awaitable( + rec: ChromeTraceRecorder, name: str, awaitable: Awaitable[Any], *, args: Dict[str, Any] +) -> Awaitable[Any]: async def _wrapped() -> Any: task = asyncio.current_task() tid = id(task) if task is not None else 0 @@ -209,7 +213,7 @@ async def test_async_loop_blocks_and_trace_for_deep_tree_all_mergers(redis_clien - Builds one deep tree that includes ALL merger types. - Simulates 2 sequential requests (fresh + next page). - - Forces refills by creating lots of cross-branch duplicates. + - Forces refills via positional under-fetch (`max_per_call=1`). - Records loop scheduling lag (blocks/hangs) and optionally exports a Chrome trace. Set `SMARTFEED_CHROME_TRACE=/path/to/trace.json` to write a trace. @@ -238,6 +242,13 @@ async def test_async_loop_blocks_and_trace_for_deep_tree_all_mergers(redis_clien # --- tracing (test-only monkeypatch) --- rec = ChromeTraceRecorder() leaf_concurrency = LeafConcurrencyTracker() + pos_leaf_calls = {"count": 0} + + pos_leaf_base = dh.make_offset_paged_method(items_pos_leaf, max_per_call=1) + + async def _pos_leaf_counted(user_id: Any, limit: int, next_page: Any, **kwargs: Any) -> Any: + pos_leaf_calls["count"] += 1 + return await pos_leaf_base(user_id, limit, next_page, **kwargs) # Leaf method tracing: wrap the *actual* subfeed method calls. # These spans are what you want to inspect for "are leaf calls parallel?". @@ -296,7 +307,7 @@ async def test_async_loop_blocks_and_trace_for_deep_tree_all_mergers(redis_clien "pos_leaf": _wrap_leaf_method_traced( rec=rec, key="pos_leaf", - method=dh.make_offset_paged_method(items_pos_leaf, max_per_call=1), + method=_pos_leaf_counted, latency_s=leaf_latency_s, concurrency=leaf_concurrency, ), @@ -436,6 +447,9 @@ async def _run_traced( # Hard assertion: leaf calls must overlap (async concurrency), not serialize. assert leaf_concurrency.peak > 1 + # Refill signal: with max_per_call=1, two page requests should trigger + # multiple extra positional calls to satisfy positional slots. + assert pos_leaf_calls["count"] > 2 # Primary signal: event-loop should remain responsive under load. assert monitor.max_lag_s < 0.1 diff --git a/tests/test_dedup_utils.py b/tests/test_dedup_utils.py index 9556af3..06150ad 100644 --- a/tests/test_dedup_utils.py +++ b/tests/test_dedup_utils.py @@ -3,12 +3,7 @@ import pytest from smartfeed.feed_models import _redis_call -from smartfeed.policies.dedup_utils import ( - decode_seen_from_cursor, - encode_seen_for_cursor, - redis_zmscore, -) - +from smartfeed.policies.dedup_utils import decode_seen_from_cursor, encode_seen_for_cursor, redis_zmscore from tests.fixtures.redis import redis_client diff --git a/tests/test_executor_slots_plan_invariants.py b/tests/test_executor_slots_plan_invariants.py index 74f8540..4b774b4 100644 --- a/tests/test_executor_slots_plan_invariants.py +++ b/tests/test_executor_slots_plan_invariants.py @@ -1,6 +1,6 @@ import pytest -from smartfeed.execution.context import ExecutionContext +from smartfeed.execution.context import ExecutionContext, RefillExecutionSettings from smartfeed.execution.executor import Executor from smartfeed.execution.plans import SlotSpec, SlotsPlan from smartfeed.feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside @@ -186,10 +186,12 @@ def assemble(output, next_page, owner_results): plan = SlotsPlan( ctx=ctx, limit=6, - next_page=FeedResultNextPage(data={ - "a": FeedResultNextPageInside(page=1, after=0), - "b": FeedResultNextPageInside(page=1, after=0), - }), + next_page=FeedResultNextPage( + data={ + "a": FeedResultNextPageInside(page=1, after=0), + "b": FeedResultNextPageInside(page=1, after=0), + } + ), params={}, slots=[ SlotSpec(owner=a, max_count=3), @@ -254,3 +256,42 @@ def assemble(output, next_page, owner_results): {"id": "b_2"}, {"id": "b_3"}, ] + + +@pytest.mark.asyncio +async def test_slots_plan_quota_deficit_refills_without_dedup_when_refill_settings_present() -> None: + executor = Executor() + ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) + ctx.refill_settings = RefillExecutionSettings(overfetch_factor=3, max_refill_loops=10) + + a = _PagedOwner(subfeed_id="a", total=10) + b = _PagedOwner(subfeed_id="b", total=10) + + def assemble(output, next_page, owner_results): + return FeedResult(data=output, next_page=next_page, has_next_page=False) + + plan = SlotsPlan( + ctx=ctx, + limit=6, + next_page=FeedResultNextPage( + data={ + "a": FeedResultNextPageInside(page=1, after=0), + "b": FeedResultNextPageInside(page=1, after=0), + } + ), + params={}, + slots=[ + SlotSpec(owner=a, max_count=3), + SlotSpec(owner=b, max_count=3), + ], + # force an initial under-fetch for owner a. + owner_fetch_limits={id(a): 1}, + assemble=assemble, + ) + + res = await executor.execute_plan(plan) + + # refill must still happen even when dedup policy is absent. + assert getattr(a, "calls") >= 2 + assert res.data[:3] == [{"id": "a_1"}, {"id": "a_2"}, {"id": "a_3"}] + assert res.data[3:] == [{"id": "b_1"}, {"id": "b_2"}, {"id": "b_3"}] diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index bbeb011..101eaac 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -211,6 +211,65 @@ async def _pos_method(user_id, limit, next_page, **kwargs): # pylint: disable=u assert pos_calls["count"] > 1 +@pytest.mark.asyncio +async def test_dedup_nested_percentage_refill_not_masked_by_parent_append() -> None: + """Nested percentage refills must run even when parent append can fill.""" + + items_a = dh.make_items("A", 1, 400, id_offset=1_000) + items_b = dh.make_items("B", 1, 400, id_offset=10_000) + items_fill = dh.make_items("F", 1, 400, id_offset=20_000) + + b_calls = {"count": 0} + b_base = dh.make_offset_paged_method(items_b, max_per_call=1) + + async def _b_method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument + b_calls["count"] += 1 + return await b_base(user_id, limit, next_page, **kwargs) + + methods_dict = { + "a": dh.make_offset_paged_method(items_a), + "b": _b_method, + "fill": dh.make_offset_paged_method(items_fill), + } + + percentage_cfg = dh._percentage_config( + "pct_nested_refill", + items=dh._percentage_items( + dh._subfeed("sf_a_nested", "a"), + dh._subfeed("sf_b_nested", "b"), + first_pct=50, + second_pct=50, + ), + ) + + config = dh._dedup_config( + "dedup_nested_pct_refill", + dh._append_config( + "append_nested_pct_refill", + [percentage_cfg, dh._subfeed("sf_fill_nested_pct", "fill")], + ), + dedup_key="id", + state_backend="cursor", + overfetch_factor=3, + max_refill_loops=50, + ) + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=12, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) == 12 + dh._assert_no_dupes_in_page(res.data) + assert "F" not in set(dh._sources(res.data)) + assert dh._sources(res.data).count("A") == 6 + assert dh._sources(res.data).count("B") == 6 + assert b_calls["count"] > 1 + + @pytest.mark.parametrize( "merger_type", [ diff --git a/tests/test_merger_percentage.py b/tests/test_merger_percentage.py index ae57f00..328dc39 100644 --- a/tests/test_merger_percentage.py +++ b/tests/test_merger_percentage.py @@ -52,3 +52,21 @@ async def test_merger_percentage_when_one_leaf_is_empty() -> None: assert res.data == ["x_4", "x_5", "x_6", "x_7"] # The non-empty leaf still reports more pages. assert res.has_next_page is True + + +@pytest.mark.asyncio +async def test_merger_percentage_odd_limit_fills_page_when_sources_have_data() -> None: + merger_percentage = parse_model(MergerPercentage, MERGER_PERCENTAGE_CONFIG) + res = await merger_percentage.get_data( + methods_dict=METHODS_DICT, + limit=11, + next_page=FeedResultNextPage( + data={ + "subfeed_merger_percentage_example": FeedResultNextPageInside(page=2, after="x_3"), + "subfeed_2_merger_percentage_example": FeedResultNextPageInside(page=3, after="x_20"), + } + ), + user_id="x", + ) + + assert len(res.data) == 11 diff --git a/tests/test_merger_percentage_gradient.py b/tests/test_merger_percentage_gradient.py index 73bc769..e2c6607 100644 --- a/tests/test_merger_percentage_gradient.py +++ b/tests/test_merger_percentage_gradient.py @@ -71,3 +71,21 @@ async def test_merger_percentage_gradient_next_page() -> None: "x_22", "x_23", ] + + +@pytest.mark.asyncio +async def test_merger_percentage_gradient_odd_limit_fills_page_when_sources_have_data() -> None: + merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) + res = await merger_percentage_gradient.get_data( + methods_dict=METHODS_DICT, + limit=11, + next_page=FeedResultNextPage( + data={ + "subfeed_from_merger_percentage_gradient_example": FeedResultNextPageInside(page=2, after="x_3"), + "subfeed_to_merger_percentage_gradient_example": FeedResultNextPageInside(page=3, after="x_20"), + } + ), + user_id="x", + ) + + assert len(res.data) == 11 diff --git a/tests/test_seen_store_unit.py b/tests/test_seen_store_unit.py index f70661f..fc7f6e7 100644 --- a/tests/test_seen_store_unit.py +++ b/tests/test_seen_store_unit.py @@ -5,7 +5,6 @@ from smartfeed.feed_models import _redis_call from smartfeed.policies.dedup_utils import decode_seen_from_cursor from smartfeed.policies.seen_store import CursorSeenStore, RedisSeenStore - from tests.fixtures.redis import redis_client From ec63967f46713bf8c3602e9122d586e1228f9737 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Mon, 9 Feb 2026 11:11:26 +0000 Subject: [PATCH 27/41] Readme added. --- ARCHITECTURE.md | 238 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 238 insertions(+) create mode 100644 ARCHITECTURE.md diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 0000000..a9ba69c --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,238 @@ +# SmartFeed Architecture (medium-brief) + +## 1) What SmartFeed does + +SmartFeed builds one paginated feed from multiple client-provided sources (“subfeeds”) using a declarative tree config: + +- **Leaf**: `SubFeed` (calls one client method) +- **Mergers**: compose children (`append`, `distribute`, `positional`, `percentage`, `percentage_gradient`, `view_session`) +- **Wrapper**: `MergerDeduplication` (changes execution semantics around one child) + +Core runtime: + +- parse config -> create request `ExecutionContext` -> run tree via shared `Executor` -> return `FeedResult` + `next_page`. + + +## 2) Public surfaces and core data types + +### Public entrypoint + +- `FeedManager(config, methods_dict, redis_client=None)` + - `get_data(user_id, limit, next_page, **params) -> FeedResult` + +`methods_dict` maps config `method_name` strings to host-app callables. + +### Config schema surface + +`smartfeed.schemas` keeps stable imports for: + +- `FeedConfig`: top-level model (`version`, `feed`) +- `FeedTypes`: discriminated union by `type` + +### Cursor / pagination models + +- `FeedResultNextPageInside`: one node cursor (`page`, `after`) +- `FeedResultNextPage`: full-tree cursor map (`data: {node_id -> FeedResultNextPageInside}`) + +### Result models + +- `FeedResultClient`: required return type of client subfeed methods +- `FeedResult`: normalized return type of any SmartFeed node + + +## 3) Node interface contract + +All nodes inherit `BaseFeedConfigModel` and are executed through: + +- `get_data(methods_dict, user_id, limit, next_page, redis_client=None, ctx=None, **params) -> FeedResult` + +Important notes: + +- If a node implements `build_plan(...)`, executor uses the plan path. +- Base `get_data(...)` delegates back to executor and expects `build_plan(...)` to exist. +- Every node has `dedup_priority: int` (used by dedup arbitration/refill ordering). + + +## 4) ExecutionContext + +`ExecutionContext` is per-request state propagated through the tree: + +- `methods_dict`, `user_id`, `redis_client` +- `executor` (lazy via `ensure_executor()`) +- optional policy/settings: + - `dedup`: `DeduplicationPolicy` when dedup wrapper is active + - `refill_settings`: `RefillExecutionSettings(overfetch_factor, max_refill_loops)` + +Responsibilities: + +- centralize shared plumbing (executor + redis client) +- keep execution policies out of user params + + +## 5) Executor (runtime engine) + +Primary entry: + +- `Executor.run(node, ctx, limit, next_page, **params) -> FeedResult` + +Execution strategy: + +1. **Plan-first** + - `build_plan(...)` -> execute returned `Plan` + - otherwise call node `get_data(...)` +2. **Centralized concurrency** + - child runs use executor-managed `asyncio.gather(...)` +3. **Dedup/refill hooks** + - for non-slot nodes with `ctx.dedup`, run `DedupRuntime.run_node_with_dedup_refill(...)` + - for `SlotsPlan`, dedup/refill is handled inside slot execution + +`SlotsPlan` execution highlights: + +1. collect unique owners + demand per owner +2. fetch owners concurrently (with optional `owner_fetch_limits` overrides) +3. merge only changed cursor keys (`CursorMap.merge_delta`) +4. apply: + - dedup arbitration + refill (`apply_slots_plan_dedup`) when `ctx.dedup` exists + - refill-only deficits (`apply_slots_plan_refill`) when only `ctx.refill_settings` exists +5. consume slot schedule and call `assemble(...)` + +When dedup is active for a slots plan, owners are executed with `dedup=None` in owner context so global arbitration stays centralized. + + +## 6) Plans: declarative execution + +Plans separate “what to run” from “how to run it”. + +- `CallablePlan(fn)` + - node-provided async function with custom flow, still executed by executor + +- `SlotsPlan(ctx, limit, next_page, params, slots, assemble, owner_fetch_limits=None)` + - `slots`: ordered `SlotSpec(owner, max_count)` schedule + - `assemble(output, merged_next_page, owner_results)`: builds final `FeedResult` + + +## 7) Mergers and leaf responsibilities + +### SubFeed (leaf) + +- derives its local cursor from `next_page.data[subfeed_id]` (defaults page=1/after=None) +- calls `methods_dict[method_name]` +- passes only params present in method signature + `subfeed_params` +- async methods are awaited; sync methods run via `asyncio.to_thread(...)` +- `raise_error=False` converts method failure into empty `FeedResultClient` +- optional `shuffle` then normalizes to `FeedResult` + +### Slot-based mergers + +These build `SlotsPlan`: + +- `MergerAppend`: concatenation (optional shuffle) +- `MergerAppendDistribute` (`type="merger_distribute"`): append then redistribute by `distribution_key` +- `MergerPositional`: page-local slot ownership for `positional` vs `default`, keeps its own merger cursor +- `MergerPercentage`: integer allocation by percentages; when total is exactly 100, remainder is distributed to avoid underfill +- `MergerPercentageGradient`: two-owner percentage curve across the page, then advances merger page cursor + +### MergerViewSession (Redis-backed session cache) + +Goal: cache a session-sized list and serve slices. + +Flow: + +1. build cache key: `{merger_id}_{user_id}` + optional suffix from `custom_view_session_key` +2. check Redis `exists`; if no cache or no merger cursor in request -> regenerate session +3. on hit, `get`; if Redis returns `None` unexpectedly, regenerate +4. on generation: execute child once for `session_size`, optional dedup, store JSON with TTL +5. return page slice and increment merger cursor page +6. optional `shuffle` is applied to returned page slice (cache payload is not reshuffled) + +### MergerDeduplication (single-child wrapper) + +Goal: deduplicate while keeping child mix/slot semantics. + +Key behavior: + +- fresh session when merger cursor is absent or `page <= 0` + - reset descendant cursors + - for Redis backend, reset Redis seen-state key +- seen-state backend: + - `cursor`: encoded into merger cursor `after` + - `redis`: ZSET `dedup:{merger_id}:{user_id}` (+ optional custom suffix) +- builds `DeduplicationPolicy` + child `ExecutionContext(dedup=..., refill_settings=...)` +- executes child via shared executor, commits store, writes merger cursor (`page+1`, `after` for cursor backend) + +Refill/overfetch behavior: + +- duplicates trigger bounded refill loops (`max_refill_loops`) +- overfetch (`overfetch_factor`) is applied only for rewindable integer-offset cursors +- when overfetch is used, leaf cursor is rewound to inspected-count to avoid skipping unseen items + + +## 8) Dedup policy + seen stores + +### DeduplicationPolicy + +Owns key extraction + acceptance rules: + +- entity key from `dedup_key` + `missing_key_policy` +- reject duplicates already seen in current response (`seen_request_set`) +- compare candidate priority vs persisted seen priority + +Capabilities: + +- batched prefetch from store +- per-owner arbitration with deterministic tie-break: `(-dedup_priority, owner_rank, item_rank)` +- ordered single-stream acceptance (`accept_batch`) returning accepted items + inspected count + +### Seen stores + +- `CursorSeenStore` + - in-cursor map of `{key -> max_priority}` + - optional compression + max-key trimming at commit + +- `RedisSeenStore` + - cached reads via `redis_zmscore(...)` + - buffered writes via `redis_zadd_and_expire(...)` + + +## 9) Redis/JSON helpers + +- `_redis_call(client, method_name, *args, **kwargs)` + - async redis client: direct await + - sync redis client: `asyncio.to_thread(...)` + +Other helpers: + +- `jsonlib`: thin `orjson` wrapper compatible with package usage (`dumps`/`loads`) +- `dedup_utils`: cursor encode/decode + Redis ZSET helper fallbacks (`zmscore` / pipeline) + + +## 10) End-to-end call flows + +### A) Standard request (no view session, no dedup) + +1. `FeedManager.get_data(...)` builds `ExecutionContext` +2. `Executor.run(root, ctx, limit, next_page)` +3. recursive execution via plans or direct `get_data(...)` +4. returns `FeedResult(data, next_page, has_next_page)` + +### B) Slot-based merger request + +1. merger returns `SlotsPlan` +2. executor fetches owners concurrently +3. optional arbitration/refill runs +4. slots are consumed in schedule order +5. `assemble(...)` builds final result + +### C) Dedup wrapper request + +1. wrapper creates store + policy and child context +2. child executes under dedup/refill control +3. executor performs acceptance/arbitration + bounded refills +4. store commits; wrapper writes merger cursor state + +### D) View-session request + +1. wrapper resolves cache key +2. cache miss/new session -> regenerate and cache +3. cache hit -> load session list from Redis +4. return requested slice + advanced merger page From 9ceabad96c35d086ee53ceb5589cccf4614a899a Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Thu, 12 Mar 2026 20:24:55 +0300 Subject: [PATCH 28/41] bugfix --- smartfeed/mergers/view_session.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index dad3314..3bab829 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -58,7 +58,17 @@ async def _set_cache( if ctx.executor is None: raise ValueError("Executor must be initialized for MergerViewSession") - result = await ctx.executor.run(self.data, ctx, self.session_size, FeedResultNextPage(data={}), **params) + # Strip dedup from context to avoid marking items as "seen" during cache build. + # Items will be deduped later when returned page-by-page through the parent merger. + cache_ctx = ExecutionContext( + methods_dict=ctx.methods_dict, + user_id=ctx.user_id, + redis_client=ctx.redis_client, + executor=ctx.executor, + dedup=None, + refill_settings=None, + ) + result = await ctx.executor.run(self.data, cache_ctx, self.session_size, FeedResultNextPage(data={}), **params) data = result.data if self.deduplicate: From 92a628ad837f06bd479a7ececc239ddf9ff3bb2f Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Wed, 18 Mar 2026 12:49:08 +0000 Subject: [PATCH 29/41] State fix dedup+view_session stacked. --- smartfeed/mergers/view_session.py | 11 +++++------ smartfeed/policies/dedup.py | 21 +++++++++++++++++++++ tests/fixtures/mergers.py | 17 +++++++++++++++++ tests/test_merger_deduplication.py | 1 + tests/test_merger_view_session.py | 25 +++++++++++++++++++++++-- 5 files changed, 67 insertions(+), 8 deletions(-) diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index 3bab829..857e69a 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -58,17 +58,16 @@ async def _set_cache( if ctx.executor is None: raise ValueError("Executor must be initialized for MergerViewSession") - # Strip dedup from context to avoid marking items as "seen" during cache build. - # Items will be deduped later when returned page-by-page through the parent merger. - cache_ctx = ExecutionContext( + inner_dedup = ctx.dedup.create_isolated() if ctx.dedup is not None else None + inner_ctx = ExecutionContext( methods_dict=ctx.methods_dict, user_id=ctx.user_id, redis_client=ctx.redis_client, executor=ctx.executor, - dedup=None, - refill_settings=None, + dedup=inner_dedup, + refill_settings=ctx.refill_settings if inner_dedup is not None else None, ) - result = await ctx.executor.run(self.data, cache_ctx, self.session_size, FeedResultNextPage(data={}), **params) + result = await ctx.executor.run(self.data, inner_ctx, self.session_size, FeedResultNextPage(data={}), **params) data = result.data if self.deduplicate: diff --git a/smartfeed/policies/dedup.py b/smartfeed/policies/dedup.py index b040c47..e744931 100644 --- a/smartfeed/policies/dedup.py +++ b/smartfeed/policies/dedup.py @@ -92,6 +92,27 @@ def record(self, key: str, priority: int) -> None: self.seen_request_set.add(key) self.store.set_max(key, priority) + def create_isolated(self) -> "DeduplicationPolicy": + """Create an isolated copy for inner execution contexts. + + Uses a fresh in-memory store and empty seen_request_set so that + inner priority arbitration works without contaminating the caller's + dedup state. + """ + from .seen_store import CursorSeenStore + + return DeduplicationPolicy( + dedup_key=self.dedup_key, + missing_key_policy=self.missing_key_policy, + store=CursorSeenStore( + cursor_compress=True, + cursor_max_keys=None, + seen_priority_map={}, + seen_order=[], + ), + seen_request_set=set(), + ) + async def arbitrate_owner_buffers( self, *, diff --git a/tests/fixtures/mergers.py b/tests/fixtures/mergers.py index c6bbf50..41c2150 100644 --- a/tests/fixtures/mergers.py +++ b/tests/fixtures/mergers.py @@ -136,3 +136,20 @@ "method_name": "doubles", }, } + +MERGER_DEDUP_VIEW_SESSION_CONFIG = { + "merger_id": "dedup_over_session", + "type": "merger_deduplication", + "dedup_key": None, + "data": { + "merger_id": "inner_session", + "type": "merger_view_session", + "session_size": 60, + "session_live_time": 300, + "data": { + "subfeed_id": "subfeed_dedup_vs", + "type": "subfeed", + "method_name": "followings", + }, + }, +} diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py index 101eaac..f0e57b7 100644 --- a/tests/test_merger_deduplication.py +++ b/tests/test_merger_deduplication.py @@ -868,6 +868,7 @@ async def test_dedup_wrapper_with_view_session_merger(redis_client) -> None: custom_view_session_key="vs1", ) + assert len(res_1.data) == 10 dh._assert_two_pages_no_dupes(res_1, res_2) # Deletion priority: for the overlapping early ids, the winning entity must be from high. diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py index f62a096..4e41536 100644 --- a/tests/test_merger_view_session.py +++ b/tests/test_merger_view_session.py @@ -3,9 +3,13 @@ import pytest from smartfeed.feed_models import _redis_call -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerViewSession +from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication, MergerViewSession from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG, MERGER_VIEW_SESSION_DUPS_CONFIG +from tests.fixtures.mergers import ( + MERGER_DEDUP_VIEW_SESSION_CONFIG, + MERGER_VIEW_SESSION_CONFIG, + MERGER_VIEW_SESSION_DUPS_CONFIG, +) from tests.fixtures.redis import redis_client from tests.utils import parse_model @@ -116,3 +120,20 @@ async def test_merger_view_session_deduplication(redis_client) -> None: assert merger_vs_res.data == [i for i in range(1, 11)] assert len(merger_vs_cache) == merger_vs.session_size assert merger_vs_cache[:10] == merger_vs_res.data + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_dedup_directly_over_view_session(redis_client) -> None: + """Regression: dedup context must not leak into view session cache generation.""" + merger = parse_model(MergerDeduplication, MERGER_DEDUP_VIEW_SESSION_CONFIG) + result = await merger.get_data( + methods_dict=METHODS_DICT, + limit=20, + next_page=FeedResultNextPage(data={}), + user_id="x", + redis_client=redis_client, + ) + # Currently returns 0 items — all rejected as "seen" during cache build + assert len(result.data) == 20 + assert result.data == [f"x_{i}" for i in range(1, 21)] From 64de66963e08001e8e0a01e414dd1c54a3aedd4e Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Wed, 18 Mar 2026 13:29:05 +0000 Subject: [PATCH 30/41] Bugfix. --- smartfeed/mergers/view_session.py | 47 +++++++++++++++++++++++++++---- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index 857e69a..9ae0e01 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -53,8 +53,9 @@ async def _set_cache( redis_client: Union[redis.Redis, AsyncRedis], cache_key: str, ctx: ExecutionContext, + child_next_page: Optional[FeedResultNextPage] = None, **params: Any, - ) -> List[Any]: + ) -> tuple[List[Any], bool, Optional[FeedResultNextPage]]: if ctx.executor is None: raise ValueError("Executor must be initialized for MergerViewSession") @@ -67,13 +68,22 @@ async def _set_cache( dedup=inner_dedup, refill_settings=ctx.refill_settings if inner_dedup is not None else None, ) - result = await ctx.executor.run(self.data, inner_ctx, self.session_size, FeedResultNextPage(data={}), **params) + start_cursor = child_next_page if child_next_page is not None else FeedResultNextPage(data={}) + result = await ctx.executor.run(self.data, inner_ctx, self.session_size, start_cursor, **params) data = result.data if self.deduplicate: data = self._dedup_data(data) await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) - return data + + meta = { + "child_has_next": result.has_next_page, + "child_cursor": result.next_page.model_dump(), + } + await _redis_call(redis_client, "set", f"{cache_key}:meta", json.dumps(meta), ex=self.session_live_time) + + child_cursor = result.next_page if result.has_next_page else None + return data, result.has_next_page, child_cursor async def _get_cache( self, @@ -90,10 +100,14 @@ async def _get_cache( ) logging.info("MergerViewSession cache request for %s", cache_key) + + child_has_next = False + child_cursor: Optional[FeedResultNextPage] = None + cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) if not cache_exists or self.merger_id not in next_page.data: logging.info("Cache miss or new session - generating fresh data for %s", cache_key) - session_data = await self._set_cache( + session_data, child_has_next, child_cursor = await self._set_cache( redis_client=redis_client, cache_key=cache_key, ctx=ctx, @@ -106,7 +120,7 @@ async def _get_cache( logging.info( "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key ) - session_data = await self._set_cache( + session_data, child_has_next, child_cursor = await self._set_cache( redis_client=redis_client, cache_key=cache_key, ctx=ctx, @@ -116,11 +130,32 @@ async def _get_cache( logging.info("Successfully read cached data for %s", cache_key) session_data = json.loads(cached_data) + meta_raw = await _redis_call(redis_client, "get", f"{cache_key}:meta") + if meta_raw: + meta = json.loads(meta_raw) + child_has_next = meta.get("child_has_next", False) + child_cursor_data = meta.get("child_cursor") + if child_cursor_data: + child_cursor = FeedResultNextPage.model_validate(child_cursor_data) + page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 + + # Session exhausted but child has more data -- rebuild from continuation cursor + if (page - 1) * limit >= len(session_data) and child_has_next and child_cursor is not None: + logging.info("Session exhausted at page %d for %s - rebuilding with child cursor", page, cache_key) + session_data, child_has_next, child_cursor = await self._set_cache( + redis_client=redis_client, + cache_key=cache_key, + ctx=ctx, + child_next_page=child_cursor, + **params, + ) + page = 1 + return FeedResult( data=session_data[(page - 1) * limit :][:limit], next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - has_next_page=bool(len(session_data) > limit * page), + has_next_page=bool(len(session_data) > limit * page or child_has_next), ) def build_plan( From 951aff2fb294d9e26fc9e85c0c1c0b816b2efd10 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Wed, 18 Mar 2026 13:53:21 +0000 Subject: [PATCH 31/41] Lint fixes. --- pyproject.toml | 3 +++ smartfeed/execution/context.py | 9 +++++++-- smartfeed/execution/plans.py | 5 +---- smartfeed/feed_models.py | 2 +- smartfeed/mergers/view_session.py | 2 +- smartfeed/schemas.py | 2 +- tests/test_dedup_utils.py | 6 +++--- tests/test_executor_slots_plan_invariants.py | 10 +++++----- tests/test_merger_append.py | 2 +- tests/test_merger_percentage.py | 2 +- tests/test_merger_positional.py | 2 +- tests/test_merger_view_session.py | 2 +- tests/test_redis_live.py | 2 -- tests/test_seen_store_unit.py | 3 +-- tests/test_view_session_unit.py | 5 ++--- 15 files changed, 29 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2c4d669..f474b19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,9 @@ strict_optional = true disallow_any_expr = false python_version = "3.10" +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["F811"] + [tool.pytest.ini_options] testpaths = ["tests"] python_files = ["tests.py", "test_*.py", "*_test.py"] diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py index 18fc769..20d213f 100644 --- a/smartfeed/execution/context.py +++ b/smartfeed/execution/context.py @@ -1,11 +1,16 @@ from __future__ import annotations +from __future__ import annotations + from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union import redis from redis.asyncio import Redis as AsyncRedis +if TYPE_CHECKING: + from ..policies.dedup import DeduplicationPolicy + @dataclass class ExecutionContext: @@ -22,7 +27,7 @@ class ExecutionContext: executor: Any = None # Policies (optional) - dedup: Optional[object] = None + dedup: Optional[DeduplicationPolicy] = None # Execution settings (optional) refill_settings: Optional["RefillExecutionSettings"] = None diff --git a/smartfeed/execution/plans.py b/smartfeed/execution/plans.py index 8a23a29..9824c86 100644 --- a/smartfeed/execution/plans.py +++ b/smartfeed/execution/plans.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional, Protocol from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage from .context import ExecutionContext @@ -53,8 +53,5 @@ class SlotsPlan: owner_fetch_limits: Optional[Dict[int, int]] = None -# NOTE: `Executor` is imported only for typing to avoid an import cycle. -from typing import TYPE_CHECKING - if TYPE_CHECKING: from .executor import Executor diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py index 38d26d6..4227253 100644 --- a/smartfeed/feed_models.py +++ b/smartfeed/feed_models.py @@ -2,7 +2,7 @@ import inspect from dataclasses import dataclass from random import shuffle -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast import redis from pydantic import BaseModel diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index 9ae0e01..e29ed22 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -2,7 +2,7 @@ import logging from random import shuffle -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union import redis from redis.asyncio import Redis as AsyncRedis diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py index 9863bb4..298cde7 100644 --- a/smartfeed/schemas.py +++ b/smartfeed/schemas.py @@ -6,7 +6,7 @@ from __future__ import annotations -from typing import Annotated, Any, Dict, Union +from typing import Annotated, Any, Union from pydantic import BaseModel, Field diff --git a/tests/test_dedup_utils.py b/tests/test_dedup_utils.py index 06150ad..1805516 100644 --- a/tests/test_dedup_utils.py +++ b/tests/test_dedup_utils.py @@ -1,10 +1,10 @@ -from typing import Any, List +from typing import Any import pytest from smartfeed.feed_models import _redis_call from smartfeed.policies.dedup_utils import decode_seen_from_cursor, encode_seen_for_cursor, redis_zmscore -from tests.fixtures.redis import redis_client +from tests.fixtures.redis import redis_client # noqa: F401 class _RedisNoZmscore: @@ -67,7 +67,7 @@ async def test_redis_zmscore_pipeline_fallback_for_sync_client_without_zmscore(r await _redis_call(redis_client, "zadd", key, mapping={"a": 1.0, "b": 2.0}) wrapped = _RedisNoZmscore(redis_client) - res = await redis_zmscore(wrapped, key, ["a", "missing", "b"]) + res = await redis_zmscore(wrapped, key, ["a", "missing", "b"]) # type: ignore[arg-type] assert res == [1.0, None, 2.0] await _redis_call(redis_client, "delete", key) diff --git a/tests/test_executor_slots_plan_invariants.py b/tests/test_executor_slots_plan_invariants.py index 4b774b4..596f10d 100644 --- a/tests/test_executor_slots_plan_invariants.py +++ b/tests/test_executor_slots_plan_invariants.py @@ -10,10 +10,10 @@ class _Owner(BaseFeedConfigModel): type: str = "test_owner" + name: str = "" - def __init__(self, *, name: str, **data): - super().__init__(**data) - object.__setattr__(self, "name", name) + def __init__(self, *, name: str, **data): # type: ignore[override] + super().__init__(name=name, **data) # type: ignore[call-arg] object.__setattr__(self, "last_limit", None) object.__setattr__(self, "calls", 0) @@ -37,8 +37,8 @@ class _PagedOwner(BaseFeedConfigModel): subfeed_id: str total: int = 10 - def __init__(self, *, subfeed_id: str, total: int = 10, **data): - super().__init__(subfeed_id=subfeed_id, total=total, **data) + def __init__(self, *, subfeed_id: str, total: int = 10, **data): # type: ignore[override] + super().__init__(subfeed_id=subfeed_id, total=total, **data) # type: ignore[call-arg] object.__setattr__(self, "calls", 0) object.__setattr__(self, "limits", []) diff --git a/tests/test_merger_append.py b/tests/test_merger_append.py index 290dd04..c6f67fb 100644 --- a/tests/test_merger_append.py +++ b/tests/test_merger_append.py @@ -48,7 +48,7 @@ async def test_merger_append_with_item_1_page_2() -> None: @pytest.mark.asyncio async def test_merger_append_when_one_leaf_is_empty() -> None: - config = copy.deepcopy(MERGER_APPEND_CONFIG) + config: dict = copy.deepcopy(MERGER_APPEND_CONFIG) # Make the second leaf return no data + has_next_page=False. config["items"][1]["method_name"] = "empty" diff --git a/tests/test_merger_percentage.py b/tests/test_merger_percentage.py index 328dc39..10da578 100644 --- a/tests/test_merger_percentage.py +++ b/tests/test_merger_percentage.py @@ -32,7 +32,7 @@ async def test_merger_percentage() -> None: @pytest.mark.asyncio async def test_merger_percentage_when_one_leaf_is_empty() -> None: - config = copy.deepcopy(MERGER_PERCENTAGE_CONFIG) + config: dict = copy.deepcopy(MERGER_PERCENTAGE_CONFIG) # Make the second leaf return no data + has_next_page=False. config["items"][1]["data"]["method_name"] = "empty" diff --git a/tests/test_merger_positional.py b/tests/test_merger_positional.py index 370f770..8653d80 100644 --- a/tests/test_merger_positional.py +++ b/tests/test_merger_positional.py @@ -58,7 +58,7 @@ async def test_merger_positional_with_empty_default() -> None: """ merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) - merger_positional.default.method_name = "empty" + merger_positional.default.method_name = "empty" # type: ignore[union-attr] merger_positional_res = await merger_positional.get_data( methods_dict=METHODS_DICT, limit=10, diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py index 4e41536..aef0a25 100644 --- a/tests/test_merger_view_session.py +++ b/tests/test_merger_view_session.py @@ -10,7 +10,7 @@ MERGER_VIEW_SESSION_CONFIG, MERGER_VIEW_SESSION_DUPS_CONFIG, ) -from tests.fixtures.redis import redis_client +from tests.fixtures.redis import redis_client # noqa: F401 from tests.utils import parse_model diff --git a/tests/test_redis_live.py b/tests/test_redis_live.py index e1dc13d..32f59a7 100644 --- a/tests/test_redis_live.py +++ b/tests/test_redis_live.py @@ -1,10 +1,8 @@ import asyncio -import json import time import pytest import redis -from redis.asyncio import Redis as AsyncRedis from smartfeed.schemas import FeedResultNextPage, MergerViewSession from tests.fixtures.configs import METHODS_DICT diff --git a/tests/test_seen_store_unit.py b/tests/test_seen_store_unit.py index fc7f6e7..2ae77c3 100644 --- a/tests/test_seen_store_unit.py +++ b/tests/test_seen_store_unit.py @@ -1,11 +1,10 @@ -from typing import Any import pytest from smartfeed.feed_models import _redis_call from smartfeed.policies.dedup_utils import decode_seen_from_cursor from smartfeed.policies.seen_store import CursorSeenStore, RedisSeenStore -from tests.fixtures.redis import redis_client +from tests.fixtures.redis import redis_client # noqa: F401 @pytest.mark.asyncio diff --git a/tests/test_view_session_unit.py b/tests/test_view_session_unit.py index a847eea..d237185 100644 --- a/tests/test_view_session_unit.py +++ b/tests/test_view_session_unit.py @@ -1,13 +1,12 @@ from dataclasses import dataclass -from typing import Any import pytest from smartfeed.feed_models import _redis_call -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerViewSession +from smartfeed.schemas import FeedResultNextPage, MergerViewSession from tests.fixtures.configs import METHODS_DICT from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG -from tests.fixtures.redis import redis_client +from tests.fixtures.redis import redis_client # noqa: F401 from tests.utils import parse_model From acc66dc3485080169d864b2819087a4336c8f932 Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Thu, 19 Mar 2026 12:46:04 +0300 Subject: [PATCH 32/41] more tests for large feeds --- smartfeed/examples/example_client.py | 22 +++ smartfeed/mergers/view_session.py | 7 +- tests/fixtures/configs.py | 1 + tests/test_merger_view_session.py | 225 +++++++++++++++++++++++++++ 4 files changed, 254 insertions(+), 1 deletion(-) diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py index a24e130..485c4b5 100644 --- a/smartfeed/examples/example_client.py +++ b/smartfeed/examples/example_client.py @@ -53,6 +53,28 @@ async def example_method( next_page.page += 1 return FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) + @staticmethod + async def large_method( + user_id: str, + limit: int, + next_page: FeedResultNextPageInside, + limit_to_return: Optional[int] = None, + ) -> FeedResultClient: + data = [f"{user_id}_{i}" for i in range(1, 5001)] + + from_index = (data.index(next_page.after) + 1) if next_page.after else 0 + to_index = from_index + limit + + result_data = data[from_index:to_index] + + if isinstance(limit_to_return, int) and limit_to_return > 0: + result_data = result_data[:limit_to_return] + + has_next = to_index < len(data) + next_page.after = result_data[-1] if result_data else None + next_page.page += 1 + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next) + @staticmethod async def empty_method( user_id: str, # pylint: disable=W0613 diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py index e29ed22..d7126e2 100644 --- a/smartfeed/mergers/view_session.py +++ b/smartfeed/mergers/view_session.py @@ -152,10 +152,15 @@ async def _get_cache( ) page = 1 + has_more_in_session = len(session_data) > limit * page + can_rebuild = child_has_next and child_cursor is not None + return FeedResult( data=session_data[(page - 1) * limit :][:limit], next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - has_next_page=bool(len(session_data) > limit * page or child_has_next), + # True while the current session still has pages OR the child + # can provide a new session (triggers rebuild on the next call). + has_next_page=bool(has_more_in_session or can_rebuild), ) def build_plan( diff --git a/tests/fixtures/configs.py b/tests/fixtures/configs.py index 6aff5cb..642d169 100644 --- a/tests/fixtures/configs.py +++ b/tests/fixtures/configs.py @@ -5,6 +5,7 @@ METHODS_DICT: Dict[str, Callable] = { "ads": ClientMixerClass().example_method, "followings": ClientMixerClass().example_method, + "large": ClientMixerClass().large_method, "empty": ClientMixerClass().empty_method, "error": ClientMixerClass().error_method, "doubles": ClientMixerClass().doubles_method, diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py index aef0a25..3269e6f 100644 --- a/tests/test_merger_view_session.py +++ b/tests/test_merger_view_session.py @@ -137,3 +137,228 @@ async def test_dedup_directly_over_view_session(redis_client) -> None: # Currently returns 0 items — all rejected as "seen" during cache build assert len(result.data) == 20 assert result.data == [f"x_{i}" for i in range(1, 21)] + + +MERGER_VIEW_SESSION_SMALL_CONFIG = { + "merger_id": "small_session", + "type": "merger_view_session", + "session_size": 20, + "session_live_time": 300, + "data": { + "subfeed_id": "subfeed_small_session", + "type": "subfeed", + "method_name": "followings", + }, +} + +MERGER_DEDUP_SMALL_SESSION_CONFIG = { + "merger_id": "dedup_small", + "type": "merger_deduplication", + "dedup_key": None, + "max_refill_loops": 3, + "data": { + "merger_id": "small_session_dedup", + "type": "merger_view_session", + "session_size": 30, + "session_live_time": 300, + "data": { + "subfeed_id": "subfeed_dedup_small", + "type": "subfeed", + "method_name": "followings", + }, + }, +} + +MERGER_DEDUP_LARGE_POOL_CONFIG = { + "merger_id": "dedup_large", + "type": "merger_deduplication", + "dedup_key": None, + "max_refill_loops": 3, + "data": { + "merger_id": "large_pool_session", + "type": "merger_view_session", + "session_size": 300, + "session_live_time": 300, + "data": { + "subfeed_id": "subfeed_large_pool", + "type": "subfeed", + "method_name": "large", + }, + }, +} + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_session_rebuild_continues_pagination(redis_client) -> None: + """Pagination continues beyond session_size via session rebuild. + + session_size=20, limit=10. After 2 pages (20 items), session exhausts. + has_next_page should remain True because child has more data. + Page 3 triggers rebuild and returns fresh items. + """ + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_SMALL_CONFIG) + user_id = "rebuild_test" + + all_items = [] + next_page = FeedResultNextPage(data={}) + num_pages = 5 # 50 items, crosses 2+ sessions of 20 + + for page_num in range(1, num_pages + 1): + result = await merger_vs.get_data( + methods_dict=METHODS_DICT, + limit=10, + next_page=next_page, + user_id=user_id, + redis_client=redis_client, + ) + assert len(result.data) == 10, f"Page {page_num}: expected 10 items, got {len(result.data)}" + all_items.extend(result.data) + next_page = result.next_page + + if page_num < num_pages: + assert result.has_next_page, ( + f"Page {page_num}: has_next_page=False but expected True " + f"(session_size=20, should rebuild)" + ) + + # 50 unique items across 2+ sessions + assert len(all_items) == 50 + assert all_items == [f"{user_id}_{i}" for i in range(1, 51)] + # No duplicates + assert len(set(all_items)) == 50 + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_session_has_next_page_true_at_boundary(redis_client) -> None: + """has_next_page=True when session exactly exhausted but child has more. + + session_size=20, limit=10. Page 2 returns items 11-20 (session exhausted). + has_next_page must be True because child can rebuild. + """ + merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_SMALL_CONFIG) + user_id = "boundary_test" + + # Page 1 + result1 = await merger_vs.get_data( + methods_dict=METHODS_DICT, + limit=10, + next_page=FeedResultNextPage(data={}), + user_id=user_id, + redis_client=redis_client, + ) + assert result1.has_next_page is True + assert result1.data == [f"{user_id}_{i}" for i in range(1, 11)] + + # Page 2 -- session exactly exhausted (items 11-20, session_size=20) + result2 = await merger_vs.get_data( + methods_dict=METHODS_DICT, + limit=10, + next_page=result1.next_page, + user_id=user_id, + redis_client=redis_client, + ) + assert result2.data == [f"{user_id}_{i}" for i in range(11, 21)] + # KEY ASSERTION: has_next_page=True despite session exhausted + assert result2.has_next_page is True, ( + "has_next_page should be True when session exhausted but child has more data" + ) + + # Page 3 -- rebuild triggers, fresh items 21-30 + result3 = await merger_vs.get_data( + methods_dict=METHODS_DICT, + limit=10, + next_page=result2.next_page, + user_id=user_id, + redis_client=redis_client, + ) + assert len(result3.data) == 10 + assert result3.data == [f"{user_id}_{i}" for i in range(21, 31)] + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_dedup_over_session_no_duplicates_across_rebuilds(redis_client) -> None: + """dedup + session_size=30: пагинация до 900 туров, 0 дупликатов. + + session_size=30 при limit=10 дает 3 страницы на сессию. + 90 страниц = 30 пересборок сессии. + Все 900 туров должны быть уникальными (dedup трекает виденные). + (example_method генерирует 999 элементов, берем 900 чтобы не упереться в лимит) + """ + merger = parse_model(MergerDeduplication, MERGER_DEDUP_SMALL_SESSION_CONFIG) + user_id = "dedup_900" + limit = 10 + num_pages = 90 + + all_items = [] + next_page = FeedResultNextPage(data={}) + + for page_num in range(1, num_pages + 1): + result = await merger.get_data( + methods_dict=METHODS_DICT, + limit=limit, + next_page=next_page, + user_id=user_id, + redis_client=redis_client, + ) + assert len(result.data) == limit, ( + f"Page {page_num}: expected {limit} items, got {len(result.data)}" + ) + all_items.extend(result.data) + next_page = result.next_page + + if page_num < num_pages: + assert result.has_next_page, ( + f"Page {page_num}: has_next_page=False, got only {len(all_items)} items" + ) + + assert len(all_items) == 900 + unique = set(all_items) + assert len(unique) == 900, ( + f"Found {900 - len(unique)} duplicates among 900 items" + ) + + +@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) +@pytest.mark.asyncio +async def test_dedup_over_session_1500_unique_tours(redis_client) -> None: + """dedup + session_size=300, pool=5000: 1500 unique tours, 0 duplicates. + + session_size=300 with limit=150 gives 2 pages per session. + 10 pages = 5 session rebuilds = 1500 tours. + All must be unique (dedup tracks seen items across sessions). + """ + merger = parse_model(MergerDeduplication, MERGER_DEDUP_LARGE_POOL_CONFIG) + user_id = "dedup_1500" + limit = 150 + num_pages = 10 + + all_items = [] + next_page = FeedResultNextPage(data={}) + + for page_num in range(1, num_pages + 1): + result = await merger.get_data( + methods_dict=METHODS_DICT, + limit=limit, + next_page=next_page, + user_id=user_id, + redis_client=redis_client, + ) + assert len(result.data) == limit, ( + f"Page {page_num}: expected {limit} items, got {len(result.data)}" + ) + all_items.extend(result.data) + next_page = result.next_page + + if page_num < num_pages: + assert result.has_next_page, ( + f"Page {page_num}: has_next_page=False after only {len(all_items)} items" + ) + + assert len(all_items) == 1500 + unique = set(all_items) + assert len(unique) == 1500, ( + f"Found {1500 - len(unique)} duplicates among 1500 items" + ) From ec499b01e2e9c749c958fa3ced226b9703190d20 Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Wed, 25 Mar 2026 18:01:12 +0300 Subject: [PATCH 33/41] dedup fix --- smartfeed/execution/dedup_runtime.py | 12 +- ...test_positional_no_refill_high_priority.py | 128 ++++++++++++++++++ 2 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tests/test_positional_no_refill_high_priority.py diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py index 1b9b26c..c706f9d 100644 --- a/smartfeed/execution/dedup_runtime.py +++ b/smartfeed/execution/dedup_runtime.py @@ -238,14 +238,16 @@ async def _refill_deficits( if missing_total <= 0: continue - base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page + owner_res = owner_results.get(refill_owner_id) + base_np = owner_res.next_page if owner_res is not None else plan.next_page + initial_has_next = bool(owner_res.has_next_page) if owner_res is not None else True state[refill_owner_id] = { "missing_total": missing_total, "remaining": missing_total, "accepted": [], "loops": 0, "current_next_page": base_np, - "has_next_page": True, + "has_next_page": initial_has_next, } if not state: @@ -368,14 +370,16 @@ async def _refill_deficits_without_dedup( if missing_total <= 0: continue - base_np = owner_results[refill_owner_id].next_page if refill_owner_id in owner_results else plan.next_page + owner_res = owner_results.get(refill_owner_id) + base_np = owner_res.next_page if owner_res is not None else plan.next_page + initial_has_next = bool(owner_res.has_next_page) if owner_res is not None else True state[refill_owner_id] = { "missing_total": missing_total, "remaining": missing_total, "accepted": [], "loops": 0, "current_next_page": base_np, - "has_next_page": True, + "has_next_page": initial_has_next, } if not state: diff --git a/tests/test_positional_no_refill_high_priority.py b/tests/test_positional_no_refill_high_priority.py new file mode 100644 index 0000000..8b8fd81 --- /dev/null +++ b/tests/test_positional_no_refill_high_priority.py @@ -0,0 +1,128 @@ +"""Positional subfeed with highest dedup_priority must NOT be refilled +when it returns fewer items than requested slots. + +Reproduces production bug: TopSort (promo) returns fewer ads than +positional slots → _refill_deficits hardcodes has_next_page=True +in initial state → pointless refill calls even though the subfeed +already signalled has_next_page=False. +""" + +import pytest + +from smartfeed.schemas import FeedResultNextPage, MergerDeduplication +from tests.fixtures import dedup_helpers as dh +from tests.utils import parse_model + + +def _make_counting_method(items): + """Offset-paged method that also counts how many times it was called.""" + call_count = {"value": 0} + + async def _method(user_id, limit, next_page, **kwargs): + call_count["value"] += 1 + from smartfeed.schemas import FeedResultClient + + offset = int(next_page.after or 0) + result_data = items[offset : offset + limit] + next_page.after = offset + len(result_data) + next_page.page += 1 + has_next_page = (offset + len(result_data)) < len(items) + return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) + + return _method, call_count + + +@pytest.mark.asyncio +async def test_positional_high_priority_no_refill_when_exhausted(): + """When the positional subfeed (dedup_priority=1, highest) returns fewer + items than requested AND has_next_page=False, it must NOT be refilled. + + Setup: + - positional (promo): only 2 items, dedup_priority=1 + - default: 100 items, dedup_priority=0 + - positions=[1,3,5,7] → needs 4 promo items + - promo returns 2 out of 4 → has_next_page=False + - Expected: promo called exactly ONCE (no refill) + """ + promo_items = dh.make_items("promo", 1001, 1003) # only 2 items + default_items = dh.make_items("default", 1, 101) # 100 items, no overlap + + promo_method, promo_calls = _make_counting_method(promo_items) + default_method, default_calls = _make_counting_method(default_items) + + methods_dict = { + "promo": promo_method, + "default": default_method, + } + + config = dh._dedup_config( + "dedup_wrapper", + dh._positional_config( + "pos_mix", + positions=[1, 3, 5, 7], + positional=dh._subfeed("sf_promo", "promo", dedup_priority=1), + default=dh._subfeed("sf_default", "default", dedup_priority=0), + ), + max_refill_loops=5, + overfetch_factor=1, + ) + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=20, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) > 0 + + # The critical assertion: promo must be called only once. + # Before the fix, it was called 1 + max_refill_loops times + # because _refill_deficits hardcoded has_next_page=True. + assert promo_calls["value"] == 1, ( + f"Promo was called {promo_calls['value']} times, expected 1. " + f"Refill should not retry a subfeed that returned has_next_page=False." + ) + + +@pytest.mark.asyncio +async def test_positional_high_priority_no_refill_even_with_overfetch(): + """Same as above but with overfetch_factor > 1 to ensure the fix + works regardless of overfetch settings.""" + promo_items = dh.make_items("promo", 1001, 1003) # only 2 items + default_items = dh.make_items("default", 1, 101) + + promo_method, promo_calls = _make_counting_method(promo_items) + default_method, _ = _make_counting_method(default_items) + + methods_dict = { + "promo": promo_method, + "default": default_method, + } + + config = dh._dedup_config( + "dedup_wrapper", + dh._positional_config( + "pos_mix", + positions=[1, 3, 5, 7], + positional=dh._subfeed("sf_promo", "promo", dedup_priority=1), + default=dh._subfeed("sf_default", "default", dedup_priority=0), + ), + max_refill_loops=5, + overfetch_factor=5, + ) + + merger = parse_model(MergerDeduplication, config) + res = await merger.get_data( + methods_dict=methods_dict, + user_id="u", + limit=20, + next_page=FeedResultNextPage(data={}), + ) + + assert len(res.data) > 0 + assert promo_calls["value"] == 1, ( + f"Promo was called {promo_calls['value']} times with overfetch_factor=5. " + f"Expected 1 — exhausted subfeed must not be refilled." + ) From bb438e5fe04b1ded5c96490b9bbe2731cdf6969c Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Wed, 8 Apr 2026 18:31:52 +0300 Subject: [PATCH 34/41] NEW ARCH --- pyproject.toml | 7 +- smartfeed/__init__.py | 2 + smartfeed/examples/__init__.py | 0 smartfeed/examples/example_client.py | 132 --- smartfeed/execution/context.py | 43 +- smartfeed/execution/cursors.py | 63 -- smartfeed/execution/dedup_runtime.py | 452 --------- smartfeed/execution/executor.py | 269 +----- smartfeed/execution/plans.py | 60 +- smartfeed/execution/redis_lock.py | 36 + smartfeed/feed_models.py | 197 ---- smartfeed/jsonlib.py | 40 - smartfeed/manager.py | 67 +- smartfeed/mergers/__init__.py | 24 - smartfeed/mergers/append.py | 48 - smartfeed/mergers/append_distribute.py | 73 -- smartfeed/mergers/deduplication.py | 185 ---- smartfeed/mergers/percentage.py | 120 --- smartfeed/mergers/percentage_gradient.py | 155 --- smartfeed/mergers/positional.py | 116 --- smartfeed/mergers/view_session.py | 193 ---- smartfeed/models/__init__.py | 83 ++ smartfeed/models/base.py | 70 ++ smartfeed/models/mixers.py | 392 ++++++++ smartfeed/models/subfeed.py | 53 + smartfeed/models/wrapper.py | 553 +++++++++++ smartfeed/policies/dedup.py | 241 ----- smartfeed/policies/dedup_utils.py | 121 --- smartfeed/policies/seen_store.py | 159 --- smartfeed/pydantic_compat.py | 16 - smartfeed/schemas.py | 95 -- tests/conftest.py | 27 + tests/fixtures/configs.py | 124 --- tests/fixtures/dedup_helpers.py | 441 --------- tests/fixtures/mergers.py | 155 --- tests/fixtures/redis.py | 30 - tests/fixtures/subfeeds.py | 27 - tests/test_async_concurrency.py | 283 ++++++ tests/test_async_loop_blocks_trace.py | 467 --------- tests/test_config_parsing.py | 60 ++ tests/test_cursor_and_refill_edges.py | 92 -- tests/test_cursors.py | 54 ++ tests/test_dedup_policy_unit.py | 74 -- tests/test_dedup_priority.py | 59 ++ tests/test_dedup_utils.py | 73 -- tests/test_executor_slots_plan_invariants.py | 297 ------ tests/test_high_priority.py | 550 +++++++++++ tests/test_manager_params.py | 40 - tests/test_medium_priority.py | 418 ++++++++ tests/test_merger_append.py | 66 -- tests/test_merger_append_distribute.py | 57 -- tests/test_merger_deduplication.py | 912 ------------------ tests/test_merger_percentage.py | 72 -- tests/test_merger_percentage_gradient.py | 91 -- tests/test_merger_positional.py | 75 -- tests/test_merger_view_session.py | 364 ------- tests/test_mixers.py | 71 ++ tests/test_parsing_config.py | 58 -- ...test_positional_no_refill_high_priority.py | 128 --- tests/test_redis_live.py | 169 ---- tests/test_redis_lock.py | 33 + tests/test_resilience.py | 250 +++++ tests/test_seen_store_unit.py | 64 -- tests/test_sub_feed.py | 78 -- tests/test_subfeed.py | 51 + tests/test_view_session_unit.py | 53 - tests/test_wrapper_cache.py | 77 ++ tests/test_wrapper_dedup.py | 86 ++ tests/test_wrapper_full_pipeline.py | 72 ++ tests/test_wrapper_rerank.py | 83 ++ tests/test_wrapper_shared_cache.py | 96 ++ tests/utils.py | 5 - 72 files changed, 3537 insertions(+), 6810 deletions(-) delete mode 100644 smartfeed/examples/__init__.py delete mode 100644 smartfeed/examples/example_client.py delete mode 100644 smartfeed/execution/cursors.py delete mode 100644 smartfeed/execution/dedup_runtime.py create mode 100644 smartfeed/execution/redis_lock.py delete mode 100644 smartfeed/feed_models.py delete mode 100644 smartfeed/jsonlib.py delete mode 100644 smartfeed/mergers/__init__.py delete mode 100644 smartfeed/mergers/append.py delete mode 100644 smartfeed/mergers/append_distribute.py delete mode 100644 smartfeed/mergers/deduplication.py delete mode 100644 smartfeed/mergers/percentage.py delete mode 100644 smartfeed/mergers/percentage_gradient.py delete mode 100644 smartfeed/mergers/positional.py delete mode 100644 smartfeed/mergers/view_session.py create mode 100644 smartfeed/models/__init__.py create mode 100644 smartfeed/models/base.py create mode 100644 smartfeed/models/mixers.py create mode 100644 smartfeed/models/subfeed.py create mode 100644 smartfeed/models/wrapper.py delete mode 100644 smartfeed/policies/dedup.py delete mode 100644 smartfeed/policies/dedup_utils.py delete mode 100644 smartfeed/policies/seen_store.py delete mode 100644 smartfeed/pydantic_compat.py delete mode 100644 smartfeed/schemas.py create mode 100644 tests/conftest.py delete mode 100644 tests/fixtures/configs.py delete mode 100644 tests/fixtures/dedup_helpers.py delete mode 100644 tests/fixtures/mergers.py delete mode 100644 tests/fixtures/redis.py delete mode 100644 tests/fixtures/subfeeds.py create mode 100644 tests/test_async_concurrency.py delete mode 100644 tests/test_async_loop_blocks_trace.py create mode 100644 tests/test_config_parsing.py delete mode 100644 tests/test_cursor_and_refill_edges.py create mode 100644 tests/test_cursors.py delete mode 100644 tests/test_dedup_policy_unit.py create mode 100644 tests/test_dedup_priority.py delete mode 100644 tests/test_dedup_utils.py delete mode 100644 tests/test_executor_slots_plan_invariants.py create mode 100644 tests/test_high_priority.py delete mode 100644 tests/test_manager_params.py create mode 100644 tests/test_medium_priority.py delete mode 100644 tests/test_merger_append.py delete mode 100644 tests/test_merger_append_distribute.py delete mode 100644 tests/test_merger_deduplication.py delete mode 100644 tests/test_merger_percentage.py delete mode 100644 tests/test_merger_percentage_gradient.py delete mode 100644 tests/test_merger_positional.py delete mode 100644 tests/test_merger_view_session.py create mode 100644 tests/test_mixers.py delete mode 100644 tests/test_parsing_config.py delete mode 100644 tests/test_positional_no_refill_high_priority.py delete mode 100644 tests/test_redis_live.py create mode 100644 tests/test_redis_lock.py create mode 100644 tests/test_resilience.py delete mode 100644 tests/test_seen_store_unit.py delete mode 100644 tests/test_sub_feed.py create mode 100644 tests/test_subfeed.py delete mode 100644 tests/test_view_session_unit.py create mode 100644 tests/test_wrapper_cache.py create mode 100644 tests/test_wrapper_dedup.py create mode 100644 tests/test_wrapper_full_pipeline.py create mode 100644 tests/test_wrapper_rerank.py create mode 100644 tests/test_wrapper_shared_cache.py delete mode 100644 tests/utils.py diff --git a/pyproject.toml b/pyproject.toml index f474b19..f6e40c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "epoch8-smartfeed" -version = "0.2.0" +version = "0.3.0" description = "" authors = [ "Epoch8 Team " @@ -12,13 +12,14 @@ packages = [ [tool.poetry.dependencies] python = ">=3.9" +orjson = ">=3.9.0" pydantic = ">=1.10.7" redis = ">=4.5.5" -orjson = ">=3.9.0" [tool.poetry.group.dev.dependencies] -isort = "^5.12.0" black = "^23.3.0" +fakeredis = "^2.34.1" +isort = "^5.12.0" mypy = "^1.3.0" pytest = "^7.3.1" pytest-asyncio = "^0.21.0" diff --git a/smartfeed/__init__.py b/smartfeed/__init__.py index e69de29..cb612a7 100644 --- a/smartfeed/__init__.py +++ b/smartfeed/__init__.py @@ -0,0 +1,2 @@ +from .models import * # noqa +from .manager import FeedManager diff --git a/smartfeed/examples/__init__.py b/smartfeed/examples/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/smartfeed/examples/example_client.py b/smartfeed/examples/example_client.py deleted file mode 100644 index 485c4b5..0000000 --- a/smartfeed/examples/example_client.py +++ /dev/null @@ -1,132 +0,0 @@ -import base64 -from typing import Optional, Union - -from pydantic import BaseModel, ConfigDict, Field, field_validator - -from smartfeed import jsonlib as json -from smartfeed.pydantic_compat import parse_model -from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside - - -class TestClientRequest(BaseModel): - """Example client request model.""" - - profile_id: str = Field(...) - limit: int = Field(...) - next_page: Union[str, FeedResultNextPage] = Field( - base64.urlsafe_b64encode(json.dumps({"data": {}}).encode()).decode() - ) - - model_config = ConfigDict(validate_default=True) - - @field_validator("next_page") - @classmethod - def validate_next_page(cls, value: Union[str, FeedResultNextPage]) -> Union[str, FeedResultNextPage]: - if isinstance(value, str): - payload = json.loads(base64.urlsafe_b64decode(value)) - return parse_model(FeedResultNextPage, payload) - - return value - - -class ClientMixerClass: - """Example client methods for SmartFeed.""" - - @staticmethod - async def example_method( - user_id: str, - limit: int, - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, - ) -> FeedResultClient: - data = [f"{user_id}_{i}" for i in range(1, 1000)] - - from_index = (data.index(next_page.after) + 1) if next_page.after else 0 - to_index = from_index + limit - - result_data = data[from_index:to_index] - - if isinstance(limit_to_return, int) and limit_to_return > 0: - result_data = result_data[:limit_to_return] - - next_page.after = result_data[-1] if result_data else None - next_page.page += 1 - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) - - @staticmethod - async def large_method( - user_id: str, - limit: int, - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, - ) -> FeedResultClient: - data = [f"{user_id}_{i}" for i in range(1, 5001)] - - from_index = (data.index(next_page.after) + 1) if next_page.after else 0 - to_index = from_index + limit - - result_data = data[from_index:to_index] - - if isinstance(limit_to_return, int) and limit_to_return > 0: - result_data = result_data[:limit_to_return] - - has_next = to_index < len(data) - next_page.after = result_data[-1] if result_data else None - next_page.page += 1 - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next) - - @staticmethod - async def empty_method( - user_id: str, # pylint: disable=W0613 - limit: int, # pylint: disable=W0613 - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, # pylint: disable=W0613 - ) -> FeedResultClient: - next_page.after = None - next_page.page += 1 - return FeedResultClient(data=[], next_page=next_page, has_next_page=False) - - @staticmethod - async def error_method( - user_id: str, # pylint: disable=W0613 - limit: int, # pylint: disable=W0613 - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, # pylint: disable=W0613 - ) -> FeedResultClient: - next_page.after = None - next_page.page = int(10 / 0) - return FeedResultClient(data=[], next_page=next_page, has_next_page=False) - - @staticmethod - async def doubles_method( - user_id: str, # pylint: disable=W0613 - limit: int, # pylint: disable=W0613 - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, # pylint: disable=W0613 - ) -> FeedResultClient: - data = [1, 2, 3, 4, 3, 2, 5, 6, 4, 4, 7, 8, 9, 10, 9, 9, 9] - - next_page.after = None - next_page.page += 1 - return FeedResultClient(data=data, next_page=next_page, has_next_page=False) - - @staticmethod - async def keys_method( - user_id: str, - limit: int, - next_page: FeedResultNextPageInside, - limit_to_return: Optional[int] = None, - ) -> FeedResultClient: - data = [{"user_id": f"{user_id}_{i%10}", "value": i} for i in range(1, 1000)] - - from_index = (data.index(next_page.after) + 1) if next_page.after else 0 - to_index = from_index + limit - - result_data = data[from_index:to_index] - - if isinstance(limit_to_return, int) and limit_to_return > 0: - result_data = result_data[:limit_to_return] - - next_page.after = result_data[-1] if result_data else None - next_page.page += 1 - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=True) diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py index 20d213f..0f2bd27 100644 --- a/smartfeed/execution/context.py +++ b/smartfeed/execution/context.py @@ -1,50 +1,13 @@ from __future__ import annotations -from __future__ import annotations - from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional -import redis from redis.asyncio import Redis as AsyncRedis -if TYPE_CHECKING: - from ..policies.dedup import DeduplicationPolicy - @dataclass class ExecutionContext: - """Execution context propagated through the feed tree. - - Keeps internal state (policies, backends) out of user params. - """ - + session_id: str methods_dict: Dict[str, Callable] - user_id: Any - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None - - # Assigned by the caller (FeedManager / tests) to avoid circular imports. - executor: Any = None - - # Policies (optional) - dedup: Optional[DeduplicationPolicy] = None - - # Execution settings (optional) - refill_settings: Optional["RefillExecutionSettings"] = None - - def ensure_redis_client(self, redis_client: Optional[Union[redis.Redis, AsyncRedis]]) -> None: - if self.redis_client is None and redis_client is not None: - self.redis_client = redis_client - - def ensure_executor(self) -> Any: - if self.executor is None: - from .executor import Executor - - self.executor = Executor() - return self.executor - - -@dataclass(frozen=True) -class RefillExecutionSettings: - overfetch_factor: int = 1 - max_refill_loops: int = 20 + redis: Optional[AsyncRedis] = None diff --git a/smartfeed/execution/cursors.py b/smartfeed/execution/cursors.py deleted file mode 100644 index eef2fd0..0000000 --- a/smartfeed/execution/cursors.py +++ /dev/null @@ -1,63 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Iterable - -from ..feed_models import BaseFeedConfigModel, FeedResultNextPage - - -@dataclass -class CursorMap: - next_page: FeedResultNextPage - - def merge_delta(self, *, base_next_page: FeedResultNextPage, owner_next_page: FeedResultNextPage) -> None: - """Merge only the cursor keys that actually changed.""" - - for key, value in owner_next_page.data.items(): - base_value = base_next_page.data.get(key) - if base_value == value: - continue - self.next_page.data[key] = value - - def reset_keys(self, keys: Iterable[str]) -> None: - for key in keys: - self.next_page.data.pop(key, None) - - @staticmethod - def can_overfetch(*, node: BaseFeedConfigModel, base_next_page: FeedResultNextPage) -> bool: - sub_id = getattr(node, "subfeed_id", None) - if not isinstance(sub_id, str): - return False - entry = base_next_page.data.get(sub_id) - if entry is None: - return False - return isinstance(entry.after, int) - - @staticmethod - def rewind_overfetch( - *, - node: BaseFeedConfigModel, - base_next_page: FeedResultNextPage, - result_next_page: FeedResultNextPage, - inspected_count: int, - batch_size: int, - ) -> None: - sub_id = getattr(node, "subfeed_id", None) - if not isinstance(sub_id, str): - return - if sub_id not in result_next_page.data: - return - - entry = result_next_page.data[sub_id] - end_after = entry.after - if not isinstance(end_after, int): - return - - base_entry = base_next_page.data.get(sub_id) - prev_after = base_entry.after if base_entry is not None else None - if not isinstance(prev_after, int): - return - - expected_end = prev_after + batch_size - if end_after == expected_end: - entry.after = prev_after + inspected_count diff --git a/smartfeed/execution/dedup_runtime.py b/smartfeed/execution/dedup_runtime.py deleted file mode 100644 index c706f9d..0000000 --- a/smartfeed/execution/dedup_runtime.py +++ /dev/null @@ -1,452 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple - -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage -from .context import ExecutionContext -from .cursors import CursorMap -from .plans import SlotsPlan - -if TYPE_CHECKING: - from .executor import Executor - - -class DedupRuntime: - """Dedup/refill orchestration. - - This owns the control flow (refill loops, slot deficit refills, etc.) - while `DeduplicationPolicy` stays focused on acceptance/arbitration decisions. - """ - - def __init__(self, executor: "Executor") -> None: - self._executor = executor - - def _get_refill_settings(self, ctx: ExecutionContext) -> Any: - return getattr(ctx, "refill_settings", None) - - async def run_node_with_dedup_refill( - self, - *, - node: BaseFeedConfigModel, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - params: Dict[str, Any], - initial_result: FeedResult, - ) -> FeedResult: - dedup = getattr(ctx, "dedup", None) - if dedup is None: - return initial_result - - settings = self._get_refill_settings(ctx) - overfetch_factor = max(1, int(getattr(settings, "overfetch_factor", 1))) - max_refill_loops = max(1, int(getattr(settings, "max_refill_loops", 20))) - priority = int(getattr(node, "dedup_priority", 0)) - - collected: List[Any] = [] - remaining = int(limit) - loops = 0 - - base_next_page = next_page - current_result = initial_result - request_limit = max(1, remaining) - - # NOTE: Refill loops are inherently sequential for a single node because - # each subsequent request depends on the previous cursor. - while remaining > 0: - can_overfetch = CursorMap.can_overfetch(node=node, base_next_page=base_next_page) - - accepted, inspected_count = await dedup.accept_batch( - items=list(current_result.data), - priority=priority, - limit=remaining, - ) - - if can_overfetch and request_limit > remaining: - CursorMap.rewind_overfetch( - node=node, - base_next_page=base_next_page, - result_next_page=current_result.next_page, - inspected_count=inspected_count, - batch_size=len(current_result.data), - ) - - if accepted: - collected.extend(accepted) - remaining = limit - len(collected) - - if remaining <= 0 or not current_result.has_next_page or loops >= max_refill_loops: - break - loops += 1 - - base_next_page = current_result.next_page - request_limit = max(1, remaining) - if CursorMap.can_overfetch(node=node, base_next_page=base_next_page) and overfetch_factor > 1: - request_limit = max(1, remaining * overfetch_factor) - - current_result, _plan = await self._executor._run_node_raw( - node, - ctx, - request_limit, - base_next_page, - params, - ) - - return FeedResult( - data=collected, - next_page=current_result.next_page, - has_next_page=bool(current_result.has_next_page), - ) - - async def apply_slots_plan_dedup( - self, - *, - plan: SlotsPlan, - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - dedup_policy: Any, - refill_settings: Any, - cursor: CursorMap, - ) -> Tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: - owner_buffers = await dedup_policy.arbitrate_owner_buffers( - owners=owners, - owner_buffers=owner_buffers, - owner_rank=owner_index, - ) - - for owner in owners: - owner_id = id(owner) - if owner_id not in owner_results: - continue - old = owner_results[owner_id] - owner_results[owner_id] = FeedResult( - data=list(owner_buffers.get(owner_id, [])), - next_page=old.next_page, - has_next_page=old.has_next_page, - ) - - deficits = self._compute_slot_deficits(plan=plan, owner_buffers=owner_buffers) - if deficits: - await self._refill_deficits( - plan=plan, - deficits=deficits, - owners=owners, - owner_index=owner_index, - owner_buffers=owner_buffers, - owner_results=owner_results, - dedup_policy=dedup_policy, - refill_settings=refill_settings, - cursor=cursor, - ) - - return owner_buffers, owner_results - - async def apply_slots_plan_refill( - self, - *, - plan: SlotsPlan, - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - refill_settings: Any, - cursor: CursorMap, - ) -> Tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: - deficits = self._compute_slot_deficits(plan=plan, owner_buffers=owner_buffers) - if deficits: - await self._refill_deficits_without_dedup( - plan=plan, - deficits=deficits, - owners=owners, - owner_index=owner_index, - owner_buffers=owner_buffers, - owner_results=owner_results, - refill_settings=refill_settings, - cursor=cursor, - ) - - return owner_buffers, owner_results - - def _compute_slot_deficits(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> Dict[int, int]: - quota_schedule = sum(int(s.max_count) for s in plan.slots) <= int(plan.limit) - deficits: Dict[int, int] = {} - consumed: Dict[int, int] = {} - remaining = int(plan.limit) - deficit_slots: List[int] = [] - - for slot in plan.slots: - if remaining <= 0: - break - - owner_id = id(slot.owner) - want = min(int(slot.max_count), remaining) - if want <= 0: - continue - - have_total = len(owner_buffers.get(owner_id, [])) - already = int(consumed.get(owner_id, 0)) - available = max(0, have_total - already) - take = min(want, available) - missing = max(0, want - take) - if missing: - deficit_slots.append(owner_id) - if quota_schedule: - deficits[owner_id] = deficits.get(owner_id, 0) + missing - consumed[owner_id] = already + take - remaining -= take - - if quota_schedule: - return deficits - if remaining <= 0: - return {} - fallback_owner_id: Optional[int] = ( - deficit_slots[-1] if deficit_slots else (id(plan.slots[-1].owner) if plan.slots else None) - ) - return {fallback_owner_id: remaining} if fallback_owner_id is not None else {} - - async def _refill_deficits( - self, - *, - plan: SlotsPlan, - deficits: Dict[int, int], - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - dedup_policy: Any, - refill_settings: Any, - cursor: CursorMap, - ) -> None: - overfetch_factor = max(1, int(getattr(refill_settings, "overfetch_factor", 1))) - max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) - - deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] - deficit_owners = sorted( - deficit_owners, - key=lambda o: ( - int(getattr(o, "dedup_priority", 0)), - owner_index.get(id(o), 0), - ), - ) - - state: Dict[int, Dict[str, Any]] = {} - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - missing_total = int(deficits.get(refill_owner_id, 0)) - if missing_total <= 0: - continue - - owner_res = owner_results.get(refill_owner_id) - base_np = owner_res.next_page if owner_res is not None else plan.next_page - initial_has_next = bool(owner_res.has_next_page) if owner_res is not None else True - state[refill_owner_id] = { - "missing_total": missing_total, - "remaining": missing_total, - "accepted": [], - "loops": 0, - "current_next_page": base_np, - "has_next_page": initial_has_next, - } - - if not state: - return - - while True: - wave_ops: List[Tuple[Any, int, FeedResultNextPage, int, bool]] = [] - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - if owner_state["remaining"] <= 0: - continue - if not owner_state["has_next_page"]: - continue - if owner_state["loops"] >= max_refill_loops: - continue - - base_np = owner_state["current_next_page"] - remaining_before = max(1, int(owner_state["remaining"])) - request_limit = remaining_before - can_overfetch = CursorMap.can_overfetch(node=refill_owner, base_next_page=base_np) - if can_overfetch and overfetch_factor > 1: - request_limit = max(1, remaining_before * overfetch_factor) - - wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit, can_overfetch)) - - if not wave_ops: - break - - results = await self._executor.gather( - *[ - self._executor._run_owner( - plan=plan, - owner=owner, - demand=request_limit, - base_next_page=base_np, - dedup_active=True, - ) - for owner, _owner_id, base_np, request_limit, _can_overfetch in wave_ops - ] - ) - - for (owner, owner_id, base_np, request_limit, can_overfetch), result in zip(wave_ops, results): - owner_state = state[owner_id] - remaining_before = int(owner_state["remaining"]) - - owner_state["current_next_page"] = result.next_page - owner_state["has_next_page"] = bool(result.has_next_page) - cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) - - refill_prio = int(getattr(owner, "dedup_priority", 0)) - wave_accepted, inspected_count = await dedup_policy.accept_batch( - items=list(result.data), - priority=refill_prio, - limit=max(0, remaining_before), - ) - - if can_overfetch and request_limit > remaining_before: - CursorMap.rewind_overfetch( - node=owner, - base_next_page=base_np, - result_next_page=result.next_page, - inspected_count=inspected_count, - batch_size=len(result.data), - ) - - if wave_accepted: - owner_state["accepted"].extend(wave_accepted) - owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) - - if owner_state["remaining"] > 0 and owner_state["has_next_page"]: - owner_state["loops"] += 1 - - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - - accepted = owner_state["accepted"] - if accepted: - owner_buffers.setdefault(refill_owner_id, []) - owner_buffers[refill_owner_id].extend(accepted) - - owner_results[refill_owner_id] = FeedResult( - data=list(owner_buffers.get(refill_owner_id, [])), - next_page=owner_state["current_next_page"], - has_next_page=owner_state["has_next_page"], - ) - - async def _refill_deficits_without_dedup( - self, - *, - plan: SlotsPlan, - deficits: Dict[int, int], - owners: List[Any], - owner_index: Dict[int, int], - owner_buffers: Dict[int, List[Any]], - owner_results: Dict[int, FeedResult], - refill_settings: Any, - cursor: CursorMap, - ) -> None: - max_refill_loops = max(1, int(getattr(refill_settings, "max_refill_loops", 20))) - - deficit_owners: List[Any] = [o for o in owners if id(o) in deficits] - deficit_owners = sorted( - deficit_owners, - key=lambda o: ( - int(getattr(o, "dedup_priority", 0)), - owner_index.get(id(o), 0), - ), - ) - - state: Dict[int, Dict[str, Any]] = {} - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - missing_total = int(deficits.get(refill_owner_id, 0)) - if missing_total <= 0: - continue - - owner_res = owner_results.get(refill_owner_id) - base_np = owner_res.next_page if owner_res is not None else plan.next_page - initial_has_next = bool(owner_res.has_next_page) if owner_res is not None else True - state[refill_owner_id] = { - "missing_total": missing_total, - "remaining": missing_total, - "accepted": [], - "loops": 0, - "current_next_page": base_np, - "has_next_page": initial_has_next, - } - - if not state: - return - - while True: - wave_ops: List[Tuple[Any, int, FeedResultNextPage, int]] = [] - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - if owner_state["remaining"] <= 0: - continue - if not owner_state["has_next_page"]: - continue - if owner_state["loops"] >= max_refill_loops: - continue - - base_np = owner_state["current_next_page"] - request_limit = max(1, int(owner_state["remaining"])) - wave_ops.append((refill_owner, refill_owner_id, base_np, request_limit)) - - if not wave_ops: - break - - results = await self._executor.gather( - *[ - self._executor._run_owner( - plan=plan, - owner=owner, - demand=request_limit, - base_next_page=base_np, - dedup_active=False, - ) - for owner, _owner_id, base_np, request_limit in wave_ops - ] - ) - - for (_owner, owner_id, _base_np, _request_limit), result in zip(wave_ops, results): - owner_state = state[owner_id] - remaining_before = int(owner_state["remaining"]) - - owner_state["current_next_page"] = result.next_page - owner_state["has_next_page"] = bool(result.has_next_page) - cursor.merge_delta(base_next_page=plan.next_page, owner_next_page=result.next_page) - - if remaining_before > 0: - owner_state["accepted"].extend(list(result.data)[:remaining_before]) - owner_state["remaining"] = int(owner_state["missing_total"]) - len(owner_state["accepted"]) - - if owner_state["remaining"] > 0 and owner_state["has_next_page"]: - owner_state["loops"] += 1 - - for refill_owner in deficit_owners: - refill_owner_id = id(refill_owner) - owner_state = state.get(refill_owner_id) - if owner_state is None: - continue - - accepted = owner_state["accepted"] - if accepted: - owner_buffers.setdefault(refill_owner_id, []) - owner_buffers[refill_owner_id].extend(accepted) - - owner_results[refill_owner_id] = FeedResult( - data=list(owner_buffers.get(refill_owner_id, [])), - next_page=owner_state["current_next_page"], - has_next_page=owner_state["has_next_page"], - ) diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index d6a92e8..dace07e 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -1,259 +1,44 @@ from __future__ import annotations import asyncio -import inspect -from typing import Any, Dict, List, Optional, Tuple +from typing import Any -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, _pydantic_deep_copy +from smartfeed.models.base import BaseNode, FeedResult from .context import ExecutionContext -from .cursors import CursorMap -from .dedup_runtime import DedupRuntime -from .plans import CallablePlan, Plan, SlotSpec, SlotsPlan +from .plans import MixPlan -class Executor: - """Shared execution engine. - - Owns recursion and concurrency. Nodes can optionally expose `build_plan(...)`. - """ - - async def run( - self, - node: BaseFeedConfigModel, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> FeedResult: - result, plan = await self._run_node_raw(node, ctx, limit, next_page, params) - - dedup = getattr(ctx, "dedup", None) - if dedup is None: - return result - - if isinstance(plan, SlotsPlan): - return result - - return await self._dedup_runtime().run_node_with_dedup_refill( - node=node, - ctx=ctx, - limit=limit, - next_page=next_page, - params=params, - initial_result=result, - ) - - def _dedup_runtime(self) -> DedupRuntime: - runtime = getattr(self, "_dedup_runtime_instance", None) - if runtime is None: - runtime = DedupRuntime(self) - setattr(self, "_dedup_runtime_instance", runtime) - return runtime - - async def execute_plan(self, plan: Plan) -> FeedResult: - """Interpret and execute a declarative plan. - - Plans must not perform execution themselves; they are data structures. - """ - - if isinstance(plan, SlotsPlan): - return await self._execute_slots_plan(plan) - if isinstance(plan, CallablePlan): - return await plan.fn(self) - raise TypeError(f"Unknown plan type: {type(plan)!r}") - - async def gather(self, *coros: Any) -> List[Any]: - """Execute coroutines concurrently. - - Centralizes concurrency in the executor layer. - """ - - return list(await asyncio.gather(*coros)) - - async def _maybe_await(self, value: Any) -> Any: - if inspect.isawaitable(value): - return await value - return value - - async def _run_node_raw( - self, - node: BaseFeedConfigModel, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - params: Dict[str, Any], - ) -> Tuple[FeedResult, Optional[Plan]]: - build_plan = getattr(node, "build_plan", None) - if callable(build_plan): - plan: Plan = build_plan(ctx=ctx, limit=limit, next_page=next_page, **params) - result = await self.execute_plan(plan) - return result, plan - - result = await node.get_data( +async def run(node: BaseNode, ctx: ExecutionContext, limit: int, cursor: dict) -> FeedResult: + """Execute a feed node. Dispatches to node.execute() or build_mix_plan().""" + if hasattr(node, "execute"): + return await node.execute( methods_dict=ctx.methods_dict, - user_id=ctx.user_id, + session_id=ctx.session_id, limit=limit, - next_page=next_page, - redis_client=ctx.redis_client, - ctx=ctx, - **params, - ) - return result, None - - async def _execute_slots_plan(self, plan: SlotsPlan) -> FeedResult: - if plan.limit <= 0: - assembled = await self._maybe_await(plan.assemble([], plan.next_page, {})) - return assembled - - working_next_page = _pydantic_deep_copy(plan.next_page) - cursor = CursorMap(working_next_page) - owners, owner_index, owner_max_demand = self._collect_plan_owners(plan) - dedup_policy = getattr(plan.ctx, "dedup", None) - refill_settings = getattr(plan.ctx, "refill_settings", None) - dedup_active = dedup_policy is not None - - owner_buffers, owner_results = await self._run_plan_owners( - plan=plan, - owners=owners, - owner_max_demand=owner_max_demand, - dedup_active=dedup_active, cursor=cursor, + ctx=ctx, ) - if dedup_policy is not None: - owner_buffers, owner_results = await self._dedup_runtime().apply_slots_plan_dedup( - plan=plan, - owners=owners, - owner_index=owner_index, - owner_buffers=owner_buffers, - owner_results=owner_results, - dedup_policy=dedup_policy, - refill_settings=refill_settings, - cursor=cursor, - ) - elif refill_settings is not None: - owner_buffers, owner_results = await self._dedup_runtime().apply_slots_plan_refill( - plan=plan, - owners=owners, - owner_index=owner_index, - owner_buffers=owner_buffers, - owner_results=owner_results, - refill_settings=refill_settings, - cursor=cursor, - ) - - output = self._consume_slots(plan=plan, owner_buffers=owner_buffers) - assembled = await self._maybe_await(plan.assemble(output, cursor.next_page, owner_results)) - return assembled - - def _collect_plan_owners(self, plan: SlotsPlan) -> tuple[List[Any], Dict[int, int], Dict[int, int]]: - owners: List[Any] = [] - owner_index: Dict[int, int] = {} - owner_demand: Dict[int, int] = {} - for slot in plan.slots: - owner_id = id(slot.owner) - if owner_id not in owner_index: - owner_index[owner_id] = len(owners) - owners.append(slot.owner) - owner_demand[owner_id] = owner_demand.get(owner_id, 0) + int(slot.max_count) - return owners, owner_index, owner_demand - - async def _run_owner( - self, - *, - plan: SlotsPlan, - owner: Any, - demand: int, - base_next_page: FeedResultNextPage, - dedup_active: bool, - ) -> FeedResult: - isolated_next_page = _pydantic_deep_copy(base_next_page) - owner_ctx = plan.ctx - if dedup_active: - owner_ctx = ExecutionContext( - methods_dict=plan.ctx.methods_dict, - user_id=plan.ctx.user_id, - redis_client=plan.ctx.redis_client, - executor=plan.ctx.executor, - dedup=None, - # Keep refill settings so nested slots plans can still compensate - # owner deficits while top-level dedup arbitration remains centralized. - refill_settings=plan.ctx.refill_settings, - ) - return await self.run(owner, owner_ctx, demand, isolated_next_page, **plan.params) - - async def _run_plan_owners( - self, - *, - plan: SlotsPlan, - owners: List[Any], - owner_max_demand: Dict[int, int], - dedup_active: bool, - cursor: CursorMap, - ) -> tuple[Dict[int, List[Any]], Dict[int, FeedResult]]: - owner_buffers: Dict[int, List[Any]] = {id(o): [] for o in owners} - owner_results: Dict[int, FeedResult] = {} - - ops: List[tuple[Any, int]] = [] - for owner in owners: - if plan.owner_fetch_limits is not None and id(owner) in plan.owner_fetch_limits: - demand = int(plan.owner_fetch_limits[id(owner)]) - else: - demand = min(plan.limit, int(owner_max_demand.get(id(owner), 0))) - if demand > 0: - ops.append((owner, demand)) - - if not ops: - return owner_buffers, owner_results - - results = await self.gather( - *[ - self._run_owner( - plan=plan, - owner=owner, - demand=demand, - base_next_page=plan.next_page, - dedup_active=dedup_active, - ) - for owner, demand in ops - ] - ) - for (owner, _demand), owner_result in zip(ops, results): - owner_results[id(owner)] = owner_result - owner_buffers[id(owner)] = list(owner_result.data) - cursor.merge_delta( - base_next_page=plan.next_page, - owner_next_page=owner_result.next_page, - ) - - return owner_buffers, owner_results - - def _consume_slots(self, *, plan: SlotsPlan, owner_buffers: Dict[int, List[Any]]) -> List[Any]: - output: List[Any] = [] - for slot in plan.slots: - if len(output) >= plan.limit: - break - - remaining = plan.limit - len(output) - take = min(int(slot.max_count), remaining) - if take <= 0: - continue + plan: MixPlan = node.build_mix_plan(ctx=ctx, limit=limit, cursor=cursor) + return await _execute_mix(plan, ctx, cursor) - owner_buffer = owner_buffers.get(id(slot.owner), []) - if not owner_buffer: - continue - chunk = owner_buffer[:take] - del owner_buffer[: len(chunk)] - output.extend(chunk) +async def _execute_mix(plan: MixPlan, ctx: ExecutionContext, cursor: dict) -> FeedResult: + """Execute a MixPlan: run children in parallel, assemble results.""" + if not plan.children: + merged, merged_cursor = plan.assemble({}, {}) + return FeedResult(data=merged, next_page=merged_cursor, has_next_page=False) - return output + tasks = [run(c.node, ctx, c.demand, cursor) for c in plan.children] + results = await asyncio.gather(*tasks) + buffers = {} + child_cursors = {} + any_has_next = False + for child, result in zip(plan.children, results): + buffers[child.node_id] = result.data + child_cursors[child.node_id] = result.next_page + any_has_next = any_has_next or result.has_next_page -__all__ = [ - "Executor", - "Plan", - "CallablePlan", - "SlotSpec", - "SlotsPlan", -] + merged, merged_cursor = plan.assemble(buffers, child_cursors) + return FeedResult(data=merged, next_page=merged_cursor, has_next_page=any_has_next) diff --git a/smartfeed/execution/plans.py b/smartfeed/execution/plans.py index 9824c86..3bd0cfc 100644 --- a/smartfeed/execution/plans.py +++ b/smartfeed/execution/plans.py @@ -1,57 +1,17 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional, Protocol +from typing import Any, Callable, List -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage -from .context import ExecutionContext +@dataclass +class MixChild: + node_id: str + node: Any # BaseNode + demand: int -class Plan(Protocol): - """Declarative execution plan. - Plans describe what to run; the `Executor` is responsible for interpreting - and executing them. - """ - - -@dataclass(frozen=True) -class CallablePlan: - """A plan implemented as an async callable. - - Useful for mergers whose child limits depend on previous child results. - """ - - fn: Callable[["Executor"], Awaitable[FeedResult]] - - -@dataclass(frozen=True) -class SlotSpec: - """A slot segment owned by a child node. - - Output order is defined by the sequence of SlotSpecs. - """ - - owner: BaseFeedConfigModel - max_count: int - - -@dataclass(frozen=True) -class SlotsPlan: - """Plan expressed as slot ownership + an assembly function. - - The executor will fetch children (possibly in priority order) and then assemble - results in the slot schedule order. - """ - - ctx: ExecutionContext - limit: int - next_page: FeedResultNextPage - params: Dict[str, Any] - slots: List[SlotSpec] - assemble: Callable[[List[Any], FeedResultNextPage, Dict[int, FeedResult]], Any] - owner_fetch_limits: Optional[Dict[int, int]] = None - - -if TYPE_CHECKING: - from .executor import Executor +@dataclass +class MixPlan: + children: List[MixChild] + assemble: Callable # (buffers: dict[str,list], cursors: dict[str,dict]) -> (list, dict) diff --git a/smartfeed/execution/redis_lock.py b/smartfeed/execution/redis_lock.py new file mode 100644 index 0000000..86b6daf --- /dev/null +++ b/smartfeed/execution/redis_lock.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import secrets + +from redis.asyncio import Redis as AsyncRedis + + +class RedisLock: + """Distributed lock via SETNX with unique token. + + async with RedisLock(redis, "key") as acquired: + if acquired: + ... # we hold the lock + else: + ... # someone else holds it + """ + + def __init__(self, redis: AsyncRedis, key: str, ttl: int = 10): + self._redis = redis + self._key = key + self._ttl = ttl + self._token = secrets.token_hex(8) + self._owned = False + + async def __aenter__(self) -> bool: + self._owned = bool( + await self._redis.set(self._key, self._token, nx=True, ex=self._ttl) + ) + return self._owned + + async def __aexit__(self, *exc): + if not self._owned: + return + val = await self._redis.get(self._key) + if val and val.decode() == self._token: + await self._redis.delete(self._key) diff --git a/smartfeed/feed_models.py b/smartfeed/feed_models.py deleted file mode 100644 index 4227253..0000000 --- a/smartfeed/feed_models.py +++ /dev/null @@ -1,197 +0,0 @@ -import asyncio -import inspect -from dataclasses import dataclass -from random import shuffle -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union, cast - -import redis -from pydantic import BaseModel -from redis.asyncio import Redis as AsyncRedis -from redis.asyncio import RedisCluster as AsyncRedisCluster - -if TYPE_CHECKING: - from .execution.context import ExecutionContext - - -def _is_async_redis_client(client: Any) -> bool: - return isinstance(client, (AsyncRedis, AsyncRedisCluster)) - - -async def _redis_call(client: Any, method_name: str, *args: Any, **kwargs: Any) -> Any: - """Call a Redis method without blocking the event loop. - - - For `redis.asyncio` clients, calls are awaited directly. - - For sync `redis.Redis`, calls are offloaded via `asyncio.to_thread()`. - """ - - method = getattr(client, method_name) - if _is_async_redis_client(client): - return await method(*args, **kwargs) - return await asyncio.to_thread(method, *args, **kwargs) - - -def _pydantic_deep_copy(model: Any) -> Any: - """Deep copy helper compatible with Pydantic v1 and v2.""" - - if hasattr(model, "model_copy"): - return model.model_copy(deep=True) - return model.copy(deep=True) - - -class FeedResultNextPageInside(BaseModel): - """Cursor model for one feed node.""" - - page: int = 1 - after: Any = None - - -class FeedResultNextPage(BaseModel): - """Cursor model for a whole feed traversal.""" - - data: Dict[str, FeedResultNextPageInside] - - -class FeedResult(BaseModel): - """Normalized output of any feed node `get_data()`.""" - - data: List - next_page: FeedResultNextPage - has_next_page: bool - - -class FeedResultClient(BaseModel): - """Result returned by client subfeed methods.""" - - data: List - next_page: FeedResultNextPageInside - has_next_page: bool - - -class BaseFeedConfigModel(BaseModel): - """Base class for merger/subfeed config models.""" - - # Higher value means the item should "win" deduplication when duplicates exist. - # This is primarily used by MergerDeduplication and by mergers when a dedup wrapper is active. - dedup_priority: int = 0 - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional["ExecutionContext"] = None, - **params: Any, - ) -> FeedResult: - """Default merger execution path via the shared executor.""" - - if not callable(getattr(self, "build_plan", None)): - raise NotImplementedError( - f"{self.__class__.__name__} must implement build_plan(...) or override get_data(...)." - ) - - if ctx is None: - from .execution.context import ExecutionContext as _ExecutionContext - - ctx = _ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - else: - ctx.ensure_redis_client(redis_client) - - executor = ctx.ensure_executor() - return await executor.run(self, ctx, limit, next_page, **params) - - -@dataclass -class _SubFeedMethodSpec: - method: Callable - args: List[str] - - -class SubFeed(BaseFeedConfigModel): - """Leaf node pointing at a client method.""" - - subfeed_id: str - type: Literal["subfeed"] - method_name: str - subfeed_params: Dict[str, Any] = {} - raise_error: Optional[bool] = True - shuffle: bool = False - - def _get_method_spec(self, methods_dict: Dict[str, Callable]) -> _SubFeedMethodSpec: - method = methods_dict[self.method_name] - method_spec = getattr(method, "_smartfeed_original", method) - method_args = inspect.getfullargspec(method_spec).args - return _SubFeedMethodSpec(method=method, args=list(method_args)) - - async def get_data( - self, - methods_dict: Dict[str, Callable], - user_id: Any, - limit: int, - next_page: FeedResultNextPage, - redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None, - ctx: Optional["ExecutionContext"] = None, - **params: Any, - ) -> FeedResult: - if ctx is None: - from .execution.context import ExecutionContext as _ExecutionContext - - ctx = _ExecutionContext(methods_dict=methods_dict, user_id=user_id, redis_client=redis_client) - - subfeed_next_page = FeedResultNextPageInside( - page=next_page.data[self.subfeed_id].page if self.subfeed_id in next_page.data else 1, - after=next_page.data[self.subfeed_id].after if self.subfeed_id in next_page.data else None, - ) - - method_spec = self._get_method_spec(methods_dict) - - method_params: Dict[str, Any] = {} - for arg in method_spec.args: - if arg in params: - method_params[arg] = params[arg] - - method = method_spec.method - is_async = inspect.iscoroutinefunction(method) or inspect.iscoroutinefunction(getattr(method, "__call__", None)) - - try: - if is_async: - method_result = await method( - user_id=user_id, - limit=limit, - next_page=subfeed_next_page, - **method_params, - **self.subfeed_params, - ) - else: - method_result = await asyncio.to_thread( - method, - user_id=user_id, - limit=limit, - next_page=subfeed_next_page, - **method_params, - **self.subfeed_params, - ) - except Exception: - if self.raise_error: - raise - - method_result = FeedResultClient( - data=[], - next_page=subfeed_next_page, - has_next_page=False, - ) - - if not isinstance(method_result, FeedResultClient): - raise TypeError('SubFeed function must return "FeedResultClient" instance.') - - if self.shuffle: - shuffle(method_result.data) - - return FeedResult( - data=method_result.data, - next_page=FeedResultNextPage( - data={self.subfeed_id: cast(FeedResultNextPageInside, method_result.next_page)} - ), - has_next_page=bool(method_result.has_next_page), - ) diff --git a/smartfeed/jsonlib.py b/smartfeed/jsonlib.py deleted file mode 100644 index 9ab650f..0000000 --- a/smartfeed/jsonlib.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from typing import Any, Callable, Optional - -import orjson - -DefaultFn = Callable[[Any], Any] - - -def dumps( - obj: Any, - *, - default: Optional[DefaultFn] = None, - sort_keys: bool = False, -) -> str: - """Serialize *obj* to JSON text using orjson. - - This is a small compatibility layer meant to cover the subset of the stdlib - `json.dumps` API used inside this package. - - Key differences vs `orjson.dumps`: - - Returns `str` (UTF-8) instead of `bytes`. - - Supports `default=` and `sort_keys=`. - """ - - option = 0 - if sort_keys: - option |= orjson.OPT_SORT_KEYS - - return orjson.dumps(obj, default=default, option=option).decode("utf-8") - - -def loads(data: Any) -> Any: - """Deserialize JSON from *data* using orjson. - - Accepts `str`, `bytes`, `bytearray`, or `memoryview` (same spirit as - stdlib `json.loads`). - """ - - return orjson.loads(data) diff --git a/smartfeed/manager.py b/smartfeed/manager.py index d6c7a76..d4ca18c 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -1,42 +1,43 @@ -from typing import Any, Dict, Optional, Union +from __future__ import annotations -import redis -from redis.asyncio import Redis as AsyncRedis +from typing import Any, Dict, Optional from .execution.context import ExecutionContext -from .pydantic_compat import parse_model -from .schemas import FeedConfig, FeedResult, FeedResultNextPage +from .execution import executor as _executor +from .models import FeedConfig +from .models.base import FeedResult class FeedManager: - """ - Класс FeedManager. - """ - - def __init__(self, config: Dict, methods_dict: Dict, redis_client: Optional[Union[redis.Redis, AsyncRedis]] = None): - """ - Инициализация класса FeedManager. - - :param config: конфигурация. - :param methods_dict: словарь с используемыми методами. - :param redis_client: объект клиента Redis (для конфигурации с view_session = True). - """ - - self.feed_config = parse_model(FeedConfig, config) + def __init__( + self, + config: Dict, + methods_dict: Dict, + redis_client: Optional[Any] = None, + ) -> None: + if hasattr(FeedConfig, "model_validate"): + # Pydantic v2 + self.config = FeedConfig.model_validate(config) + else: + # Pydantic v1 + self.config = FeedConfig.parse_obj(config) self.methods_dict = methods_dict self.redis_client = redis_client - async def get_data(self, user_id: Any, limit: int, next_page: FeedResultNextPage, **params: Any) -> FeedResult: - """ - Метод для получения данных согласно конфигурации. - - :param user_id: ID объекта для получения данных (например, ID пользователя). - :param limit: лимит на выдачу данных. - :param next_page: курсор для пагинации в формате SmartFeedResultNextPage. - :param params: любые внешние параметры, передаваемые в исполняемую функцию на клиентской стороне. - :return: результат получения данных согласно конфигурации фида. - """ - - ctx = ExecutionContext(methods_dict=self.methods_dict, user_id=user_id, redis_client=self.redis_client) - executor = ctx.ensure_executor() - return await executor.run(self.feed_config.feed, ctx, limit, next_page, **params) + async def get_feed( + self, + session_id: str, + limit: int, + cursor: Optional[Dict] = None, + ) -> FeedResult: + cursor = cursor or {} + ctx = ExecutionContext( + session_id=session_id, + methods_dict=self.methods_dict, + redis=self.redis_client, + ) + result = await _executor.run(self.config.feed, ctx, limit, cursor) + for i, item in enumerate(result.data): + if isinstance(item, dict): + item.setdefault("_smartfeed_debug_info", {})["smartfeed_position"] = i + return result diff --git a/smartfeed/mergers/__init__.py b/smartfeed/mergers/__init__.py deleted file mode 100644 index 24b1c8e..0000000 --- a/smartfeed/mergers/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Merger implementations. - -Each merger schema lives in its own module. -`smartfeed.schemas` re-exports these classes for backwards compatibility. -""" - -from .append import MergerAppend -from .append_distribute import MergerAppendDistribute -from .deduplication import MergerDeduplication -from .percentage import MergerPercentage, MergerPercentageItem -from .percentage_gradient import MergerPercentageGradient -from .positional import MergerPositional -from .view_session import MergerViewSession - -__all__ = [ - "MergerAppend", - "MergerAppendDistribute", - "MergerDeduplication", - "MergerPercentage", - "MergerPercentageItem", - "MergerPercentageGradient", - "MergerPositional", - "MergerViewSession", -] diff --git a/smartfeed/mergers/append.py b/smartfeed/mergers/append.py deleted file mode 100644 index 9c5c5c6..0000000 --- a/smartfeed/mergers/append.py +++ /dev/null @@ -1,48 +0,0 @@ -from __future__ import annotations - -from random import shuffle -from typing import TYPE_CHECKING, Any, Dict, List, Literal, cast - -from ..execution.context import ExecutionContext -from ..execution.executor import SlotSpec, SlotsPlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerAppend(BaseFeedConfigModel): - """Append merger.""" - - merger_id: str - type: Literal["merger_append"] - items: List[FeedTypes] - shuffle: bool = False - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> SlotsPlan: - slots = [SlotSpec(owner=cast(BaseFeedConfigModel, item), max_count=limit) for item in self.items] - - def _assemble( - output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] - ) -> FeedResult: - has_next_page = any(r.has_next_page for r in owner_results.values()) - result = FeedResult(data=output, next_page=merged_next_page, has_next_page=has_next_page) - if self.shuffle: - shuffle(result.data) - return result - - return SlotsPlan( - ctx=ctx, - limit=limit, - next_page=next_page, - params=dict(params), - slots=slots, - assemble=_assemble, - ) diff --git a/smartfeed/mergers/append_distribute.py b/smartfeed/mergers/append_distribute.py deleted file mode 100644 index 220e3e3..0000000 --- a/smartfeed/mergers/append_distribute.py +++ /dev/null @@ -1,73 +0,0 @@ -from __future__ import annotations - -from collections import defaultdict, deque -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional - -from typing_extensions import no_type_check - -from ..execution.context import ExecutionContext -from ..execution.executor import SlotSpec, SlotsPlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerAppendDistribute(BaseFeedConfigModel): - """Merger that uniformly distributes items by a key.""" - - merger_id: str - type: Literal["merger_distribute"] - items: List["FeedTypes"] - distribution_key: str - sorting_key: Optional[str] = None - sorting_desc: bool = False - - @no_type_check - def _uniform_distribute(self, data: list) -> list: - if self.sorting_key: - data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) - - grouped_entries = defaultdict(deque) - for entry in data: - grouped_entries[entry[self.distribution_key]].append(entry) - result = [] - prev_profile_id = None - while any(grouped_entries.values()): - for profile_id in list(grouped_entries.keys()): - if grouped_entries[profile_id]: - if profile_id != prev_profile_id or len(grouped_entries) == 1: - result.append(grouped_entries[profile_id].popleft()) - prev_profile_id = profile_id - if not grouped_entries[profile_id]: - del grouped_entries[profile_id] - else: - del grouped_entries[profile_id] - - return result - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> SlotsPlan: - slots = [SlotSpec(owner=item, max_count=limit) for item in self.items] - - def _assemble( - output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] - ) -> FeedResult: - has_next_page = any(r.has_next_page for r in owner_results.values()) - distributed = self._uniform_distribute(output) - return FeedResult(data=distributed, next_page=merged_next_page, has_next_page=has_next_page) - - return SlotsPlan( - ctx=ctx, - limit=limit, - next_page=next_page, - params=dict(params), - slots=slots, - assemble=_assemble, - ) diff --git a/smartfeed/mergers/deduplication.py b/smartfeed/mergers/deduplication.py deleted file mode 100644 index 1b35a7b..0000000 --- a/smartfeed/mergers/deduplication.py +++ /dev/null @@ -1,185 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Dict, Literal, Optional - -from pydantic import PrivateAttr, model_validator - -from ..execution.context import ExecutionContext, RefillExecutionSettings -from ..execution.cursors import CursorMap -from ..execution.executor import CallablePlan -from ..feed_models import ( - BaseFeedConfigModel, - FeedResult, - FeedResultNextPage, - FeedResultNextPageInside, - _pydantic_deep_copy, -) -from ..policies.dedup import DeduplicationPolicy -from ..policies.seen_store import CursorSeenStore, RedisSeenStore, SeenStore - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerDeduplication(BaseFeedConfigModel): - """Merger that deduplicates while preserving child mixing/position semantics.""" - - merger_id: str - type: Literal["merger_deduplication"] - data: "FeedTypes" - - dedup_key: Optional[str] = None - missing_key_policy: Literal["error", "keep", "drop"] = "error" - - state_backend: Literal["cursor", "redis"] = "cursor" - state_ttl_seconds: int = 3600 - cursor_compress: bool = True - cursor_max_keys: Optional[int] = None - - overfetch_factor: int = 1 - - max_refill_loops: int = 20 - - _descendant_cursor_keys_cache: Optional[set[str]] = PrivateAttr(default=None) - - @model_validator(mode="after") - def validate_merger_deduplication(self) -> "MergerDeduplication": - if self.overfetch_factor < 1: - raise ValueError('"overfetch_factor" must be >= 1') - if self.max_refill_loops < 1: - raise ValueError('"max_refill_loops" must be >= 1') - return self - - def _collect_descendant_cursor_keys(self, feed: BaseFeedConfigModel) -> set[str]: - keys: set[str] = set() - stack = [feed] - while stack: - node = stack.pop() - - for attr in ("subfeed_id", "merger_id"): - value = getattr(node, attr, None) - if isinstance(value, str) and value: - keys.add(value) - - for child in ( - getattr(node, "data", None), - getattr(node, "positional", None), - getattr(node, "default", None), - ): - if isinstance(child, BaseFeedConfigModel): - stack.append(child) - - for wrapper in (getattr(node, "item_from", None), getattr(node, "item_to", None)): - inner = getattr(wrapper, "data", None) - if isinstance(inner, BaseFeedConfigModel): - stack.append(inner) - - items = getattr(node, "items", None) - if isinstance(items, list): - for item in items: - inner = item if isinstance(item, BaseFeedConfigModel) else getattr(item, "data", None) - if isinstance(inner, BaseFeedConfigModel): - stack.append(inner) - - return keys - - def _get_descendant_cursor_keys_cached(self) -> set[str]: - cached = self._descendant_cursor_keys_cache - if cached is None: - cached = self._collect_descendant_cursor_keys(self.data) - self._descendant_cursor_keys_cache = cached - return cached - - def _reset_descendant_cursors(self, next_page: FeedResultNextPage) -> None: - descendant_keys = self._get_descendant_cursor_keys_cached() - CursorMap(next_page).reset_keys(descendant_keys) - - def _build_redis_state_key(self, user_id: Any, params: Dict[str, Any]) -> str: - suffix = params.get("custom_deduplication_key") or params.get("custom_view_session_key") - if suffix: - return f"dedup:{self.merger_id}:{user_id}:{suffix}" - return f"dedup:{self.merger_id}:{user_id}" - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> CallablePlan: - async def _run(executor: Any) -> FeedResult: - if limit <= 0: - return FeedResult(data=[], next_page=next_page, has_next_page=False) - - if ctx.executor is None: - ctx.executor = executor - - entry = next_page.data.get(self.merger_id) - requested_page = entry.page if entry is not None else None - is_fresh_session = requested_page is None or (isinstance(requested_page, int) and requested_page <= 0) - - redis_client = ctx.redis_client - if self.state_backend == "redis" and not redis_client: - raise ValueError("Redis client must be provided if using MergerDeduplication with state_backend=redis") - - working_next_page = _pydantic_deep_copy(next_page) - if is_fresh_session: - self._reset_descendant_cursors(working_next_page) - - seen_request_set: set[str] = set() - store: SeenStore - if self.state_backend == "cursor": - cursor_entry = next_page.data.get(self.merger_id) - store = CursorSeenStore.from_after( - cursor_entry.after if (cursor_entry is not None and not is_fresh_session) else None, - cursor_compress=self.cursor_compress, - cursor_max_keys=self.cursor_max_keys, - ) - else: - assert redis_client is not None - redis_state_key = self._build_redis_state_key(user_id=ctx.user_id, params=params) - store = RedisSeenStore.create( - redis_client=redis_client, - redis_key=redis_state_key, - ttl_seconds=self.state_ttl_seconds, - ) - if is_fresh_session: - await store.reset() - - policy = DeduplicationPolicy( - dedup_key=self.dedup_key, - missing_key_policy=self.missing_key_policy, - store=store, - seen_request_set=seen_request_set, - ) - - refill_settings = RefillExecutionSettings( - overfetch_factor=self.overfetch_factor, - max_refill_loops=self.max_refill_loops, - ) - child_ctx = ExecutionContext( - methods_dict=ctx.methods_dict, - user_id=ctx.user_id, - redis_client=ctx.redis_client, - executor=ctx.executor, - dedup=policy, - refill_settings=refill_settings, - ) - - child = _pydantic_deep_copy(self.data) - child_result = await executor.run(child, child_ctx, limit, working_next_page, **params) - - commit_result: Any = await store.commit() - merger_after: Any = commit_result if self.state_backend == "cursor" else None - - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - result_next_page = _pydantic_deep_copy(child_result.next_page) - result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=merger_after) - return FeedResult( - data=child_result.data, - next_page=result_next_page, - has_next_page=child_result.has_next_page, - ) - - return CallablePlan(fn=_run) diff --git a/smartfeed/mergers/percentage.py b/smartfeed/mergers/percentage.py deleted file mode 100644 index 9ea2096..0000000 --- a/smartfeed/mergers/percentage.py +++ /dev/null @@ -1,120 +0,0 @@ -from __future__ import annotations - -from random import shuffle -from typing import TYPE_CHECKING, Any, Dict, List, Literal, cast - -from pydantic import BaseModel - -from ..execution.context import ExecutionContext -from ..execution.executor import SlotSpec, SlotsPlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerPercentageItem(BaseModel): - """One percentage slot.""" - - percentage: int - data: FeedTypes - - -class MergerPercentage(BaseFeedConfigModel): - """Percentage-based mixing merger.""" - - merger_id: str - type: Literal["merger_percentage"] - items: List[MergerPercentageItem] - shuffle: bool = False - - @staticmethod - def _merge_items_data(items_data: List[List]) -> List: - result: List = [] - cursor: List[Dict] = [] - - min_length = min(len(item_data) for item_data in items_data) or 1 - for item_data in items_data: - cursor.append( - { - "items": item_data, - "current": 0, - "size": round(len(item_data) / min_length), - } - ) - - full_length = sum(len(item_data) for item_data in items_data) - while len(result) < full_length: - for item_cursor in cursor: - items = item_cursor["items"] - start = item_cursor["current"] - end = start + item_cursor["size"] if start + item_cursor["size"] < len(items) else len(items) - result.extend(items[start:end]) - item_cursor["current"] = end - - return result - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> SlotsPlan: - owners: List[BaseFeedConfigModel] = [cast(BaseFeedConfigModel, item.data) for item in self.items] - - slot_limits: List[int] = [] - remainders: List[tuple[int, int]] = [] - total_percentage = sum(int(item.percentage) for item in self.items) - - for idx, item in enumerate(self.items): - raw = int(limit) * int(item.percentage) - child_limit = raw // 100 - slot_limits.append(max(0, child_limit)) - remainders.append((raw % 100, idx)) - - # avoid underfilling for the common "percentages sum to 100" case - if total_percentage == 100: - missing = max(0, int(limit) - sum(slot_limits)) - if missing > 0: - for _rem, idx in sorted(remainders, key=lambda x: (-x[0], x[1])): - if missing <= 0: - break - slot_limits[idx] += 1 - missing -= 1 - - slots: List[SlotSpec] = [ - SlotSpec(owner=owner, max_count=max(0, int(slot_limits[idx]))) for idx, owner in enumerate(owners) - ] - - def _assemble( - output: List[Any], - merged_next_page: FeedResultNextPage, - owner_results: Dict[int, FeedResult], - ) -> FeedResult: - items_data: List[List[Any]] = [] - has_next_page = False - - for owner in owners: - child_res = owner_results.get(id(owner)) - if child_res is None: - items_data.append([]) - continue - items_data.append(list(child_res.data)) - has_next_page = has_next_page or bool(child_res.has_next_page) - - data = self._merge_items_data(items_data=items_data) - if self.shuffle: - shuffle(data) - - return FeedResult(data=data, next_page=merged_next_page, has_next_page=has_next_page) - - return SlotsPlan( - ctx=ctx, - limit=limit, - next_page=next_page, - params=dict(params), - slots=slots, - assemble=_assemble, - ) diff --git a/smartfeed/mergers/percentage_gradient.py b/smartfeed/mergers/percentage_gradient.py deleted file mode 100644 index fb70891..0000000 --- a/smartfeed/mergers/percentage_gradient.py +++ /dev/null @@ -1,155 +0,0 @@ -from random import shuffle -from typing import Any, Dict, List, Literal, cast - -from pydantic import model_validator - -from ..execution.context import ExecutionContext -from ..execution.executor import SlotSpec, SlotsPlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside -from .percentage import MergerPercentageItem - - -class MergerPercentageGradient(BaseFeedConfigModel): - """Percentage-gradient merger.""" - - merger_id: str - type: Literal["merger_percentage_gradient"] - item_from: MergerPercentageItem - item_to: MergerPercentageItem - step: int - size_to_step: int - shuffle: bool = False - - @model_validator(mode="after") - def validate_merger_percentage_gradient(self) -> "MergerPercentageGradient": - if self.step < 1 or self.step > 100: - raise ValueError('"step" must be in range from 1 to 100') - if self.size_to_step < 1: - raise ValueError('"size_to_step" must be bigger than 1') - return self - - def _calculate_limits_and_percents(self, page: int, limit: int) -> Dict: - result: Dict = { - "limit_from": 0, - "limit_to": 0, - "percentages": [], - } - - percentage_from = self.item_from.percentage - percentage_to = self.item_to.percentage - start_position = limit * (page - 1) - first_iter = True - - for i in range(self.size_to_step, limit * page + self.size_to_step, self.size_to_step): - if not first_iter and percentage_to < 100: - percentage_from -= self.step - percentage_to += self.step - - if percentage_to > 100 or percentage_from < 0: - percentage_from = 0 - percentage_to = 100 - - if i > start_position: - iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) - start_position = i - - if result["percentages"] and result["percentages"][-1]["to"] >= 100: - result["limit_to"] += iter_limit - result["percentages"][-1]["limit"] += iter_limit - result["percentages"][-1]["to_take"] += iter_limit - else: - from_take = iter_limit * percentage_from // 100 - to_take = iter_limit - from_take - result["limit_from"] += from_take - result["limit_to"] += to_take - iter_result = { - "limit": iter_limit, - "from": percentage_from, - "to": percentage_to, - "from_take": from_take, - "to_take": to_take, - } - result["percentages"].append(iter_result) - - if first_iter: - first_iter = False - - return result - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> SlotsPlan: - start_page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - start_after = next_page.data[self.merger_id].after if self.merger_id in next_page.data else None - - plan_next_page = FeedResultNextPage( - data={ - **next_page.data, - self.merger_id: FeedResultNextPageInside(page=start_page, after=start_after), - } - ) - - limits_and_percents = self._calculate_limits_and_percents(page=start_page, limit=limit) - - owner_from = cast(BaseFeedConfigModel, self.item_from.data) - owner_to = cast(BaseFeedConfigModel, self.item_to.data) - - slots = [ - SlotSpec(owner=owner_from, max_count=int(limits_and_percents["limit_from"])), - SlotSpec(owner=owner_to, max_count=int(limits_and_percents["limit_to"])), - ] - - def _assemble( - output: List[Any], - merged_next_page: FeedResultNextPage, - owner_results: Dict[int, FeedResult], - ) -> FeedResult: - from_res = owner_results.get(id(owner_from)) - to_res = owner_results.get(id(owner_to)) - - from_data = list(from_res.data) if from_res is not None else [] - to_data = list(to_res.data) if to_res is not None else [] - - data: List[Any] = [] - from_start_index = 0 - to_start_index = 0 - for lp_data in limits_and_percents["percentages"]: - from_take = int(lp_data.get("from_take", lp_data["limit"] * lp_data["from"] // 100)) - to_take = int(lp_data.get("to_take", lp_data["limit"] - from_take)) - - from_end_index = from_start_index + from_take - to_end_index = to_start_index + to_take - - data.extend(from_data[from_start_index:from_end_index]) - data.extend(to_data[to_start_index:to_end_index]) - - from_start_index = from_end_index - to_start_index = to_end_index - - has_next_page = False - if from_res is not None and from_res.has_next_page: - has_next_page = True - if to_res is not None and to_res.has_next_page: - has_next_page = True - - if self.shuffle: - shuffle(data) - - if self.merger_id in merged_next_page.data: - merged_next_page.data[self.merger_id].page += 1 - - return FeedResult(data=data, next_page=merged_next_page, has_next_page=has_next_page) - - return SlotsPlan( - ctx=ctx, - limit=limit, - next_page=plan_next_page, - params=dict(params), - slots=slots, - assemble=_assemble, - ) diff --git a/smartfeed/mergers/positional.py b/smartfeed/mergers/positional.py deleted file mode 100644 index 6023aae..0000000 --- a/smartfeed/mergers/positional.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional - -from pydantic import model_validator - -from ..execution.context import ExecutionContext -from ..execution.executor import SlotSpec, SlotsPlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerPositional(BaseFeedConfigModel): - """Positional merger.""" - - merger_id: str - type: Literal["merger_positional"] - positions: List[int] = [] - start: Optional[int] = None - end: Optional[int] = None - step: Optional[int] = None - positional: FeedTypes - default: FeedTypes - - @model_validator(mode="after") - def validate_merger_positional(self) -> "MergerPositional": - if not self.positions and not all((self.start, self.end, self.step)): - raise ValueError('Either "positions" or "start", "end", and "step" must be provided') - if self.start and self.positions: - if isinstance(self.start, int) and self.start <= max(self.positions): - raise ValueError('"start" must be bigger than maximum value of "positions"') - if isinstance(self.start, int) and isinstance(self.end, int): - if self.end <= self.start: - raise ValueError('"end" must be bigger than "start"') - return self - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> SlotsPlan: - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - - positional_has_next_page = True - page_positions: List[int] = [] - available_positions = range((page - 1) * limit, (page * limit) + 1) - for position in self.positions: - if position in available_positions: - page_positions.append(available_positions.index(position)) - - if max(available_positions) >= max(self.positions, default=0): - positional_has_next_page = False - - if self.start is not None and self.end is not None and self.step is not None: - positional_has_next_page = not max(available_positions) >= self.end - - for position in range(self.start, self.end, self.step): - if position in available_positions: - page_positions.append(available_positions.index(position)) - - pos_limit = len(page_positions) - - # Build a slot ownership schedule by applying the same sequential insert - # semantics as the legacy assembly logic. - schedule: List[BaseFeedConfigModel] = [self.default for _ in range(limit)] - for insert_index in [p - 1 for p in page_positions[:pos_limit]]: - schedule.insert(insert_index, self.positional) - schedule = schedule[:limit] - - # Compress the schedule into contiguous segments. - slots: List[SlotSpec] = [] - if schedule: - current_owner = schedule[0] - count = 1 - for owner in schedule[1:]: - if owner is current_owner: - count += 1 - continue - slots.append(SlotSpec(owner=current_owner, max_count=count)) - current_owner = owner - count = 1 - slots.append(SlotSpec(owner=current_owner, max_count=count)) - - after = next_page.data[self.merger_id].after if self.merger_id in next_page.data else None - - def _assemble( - output: List[Any], merged_next_page: FeedResultNextPage, owner_results: Dict[int, FeedResult] - ) -> FeedResult: - default_res = owner_results.get(id(self.default)) - pos_res = owner_results.get(id(self.positional)) - - has_next_page = bool(default_res.has_next_page) if default_res is not None else False - if not has_next_page and positional_has_next_page and pos_res is not None and pos_res.has_next_page: - has_next_page = True - - result_next_page = merged_next_page - result_next_page.data[self.merger_id] = FeedResultNextPageInside(page=page + 1, after=after) - return FeedResult(data=output, next_page=result_next_page, has_next_page=has_next_page) - - return SlotsPlan( - ctx=ctx, - limit=limit, - next_page=next_page, - params=dict(params), - slots=slots, - owner_fetch_limits={ - id(self.default): limit, - id(self.positional): pos_limit, - }, - assemble=_assemble, - ) diff --git a/smartfeed/mergers/view_session.py b/smartfeed/mergers/view_session.py deleted file mode 100644 index d7126e2..0000000 --- a/smartfeed/mergers/view_session.py +++ /dev/null @@ -1,193 +0,0 @@ -from __future__ import annotations - -import logging -from random import shuffle -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union - -import redis -from redis.asyncio import Redis as AsyncRedis - -from .. import jsonlib as json -from ..execution.context import ExecutionContext -from ..execution.executor import CallablePlan -from ..feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside, _redis_call -from ..policies.dedup import entity_key - -if TYPE_CHECKING: - from ..schemas import FeedTypes - - -class MergerViewSession(BaseFeedConfigModel): - """Merger with view-session caching.""" - - merger_id: str - type: Literal["merger_view_session"] - session_size: int - session_live_time: int - data: "FeedTypes" - deduplicate: bool = False - dedup_key: Optional[str] = None - missing_key_policy: Literal["error", "keep", "drop"] = "error" - shuffle: bool = False - - def _get_dedup_key_or_attr(self, item: Any) -> str: - key = entity_key(item, self.dedup_key, self.missing_key_policy) - assert key is not None, "Deduplication key is missing and item was dropped by missing_key_policy='drop'" - return key - - def _dedup_data(self, data: List[Any]) -> List[Any]: - deduplicated: List[Any] = [] - seen: set[str] = set() - for item in data: - key = entity_key(item, self.dedup_key, self.missing_key_policy) - if key is None: - continue - if key in seen: - continue - seen.add(key) - deduplicated.append(item) - return deduplicated - - async def _set_cache( - self, - redis_client: Union[redis.Redis, AsyncRedis], - cache_key: str, - ctx: ExecutionContext, - child_next_page: Optional[FeedResultNextPage] = None, - **params: Any, - ) -> tuple[List[Any], bool, Optional[FeedResultNextPage]]: - if ctx.executor is None: - raise ValueError("Executor must be initialized for MergerViewSession") - - inner_dedup = ctx.dedup.create_isolated() if ctx.dedup is not None else None - inner_ctx = ExecutionContext( - methods_dict=ctx.methods_dict, - user_id=ctx.user_id, - redis_client=ctx.redis_client, - executor=ctx.executor, - dedup=inner_dedup, - refill_settings=ctx.refill_settings if inner_dedup is not None else None, - ) - start_cursor = child_next_page if child_next_page is not None else FeedResultNextPage(data={}) - result = await ctx.executor.run(self.data, inner_ctx, self.session_size, start_cursor, **params) - - data = result.data - if self.deduplicate: - data = self._dedup_data(data) - await _redis_call(redis_client, "set", cache_key, json.dumps(data), ex=self.session_live_time) - - meta = { - "child_has_next": result.has_next_page, - "child_cursor": result.next_page.model_dump(), - } - await _redis_call(redis_client, "set", f"{cache_key}:meta", json.dumps(meta), ex=self.session_live_time) - - child_cursor = result.next_page if result.has_next_page else None - return data, result.has_next_page, child_cursor - - async def _get_cache( - self, - limit: int, - next_page: FeedResultNextPage, - redis_client: Union[redis.Redis, AsyncRedis], - ctx: ExecutionContext, - **params: Any, - ) -> FeedResult: - cache_key = ( - f"{self.merger_id}_{ctx.user_id}_{session_cache_key}" - if (session_cache_key := params.get("custom_view_session_key")) - else f"{self.merger_id}_{ctx.user_id}" - ) - - logging.info("MergerViewSession cache request for %s", cache_key) - - child_has_next = False - child_cursor: Optional[FeedResultNextPage] = None - - cache_exists = bool(await _redis_call(redis_client, "exists", cache_key)) - if not cache_exists or self.merger_id not in next_page.data: - logging.info("Cache miss or new session - generating fresh data for %s", cache_key) - session_data, child_has_next, child_cursor = await self._set_cache( - redis_client=redis_client, - cache_key=cache_key, - ctx=ctx, - **params, - ) - else: - logging.info("Cache exists - attempting read from Redis for %s", cache_key) - cached_data = await _redis_call(redis_client, "get", cache_key) - if cached_data is None: - logging.info( - "Redis returned None for %s - falling back to fresh data (cluster replication issue)", cache_key - ) - session_data, child_has_next, child_cursor = await self._set_cache( - redis_client=redis_client, - cache_key=cache_key, - ctx=ctx, - **params, - ) - else: - logging.info("Successfully read cached data for %s", cache_key) - session_data = json.loads(cached_data) - - meta_raw = await _redis_call(redis_client, "get", f"{cache_key}:meta") - if meta_raw: - meta = json.loads(meta_raw) - child_has_next = meta.get("child_has_next", False) - child_cursor_data = meta.get("child_cursor") - if child_cursor_data: - child_cursor = FeedResultNextPage.model_validate(child_cursor_data) - - page = next_page.data[self.merger_id].page if self.merger_id in next_page.data else 1 - - # Session exhausted but child has more data -- rebuild from continuation cursor - if (page - 1) * limit >= len(session_data) and child_has_next and child_cursor is not None: - logging.info("Session exhausted at page %d for %s - rebuilding with child cursor", page, cache_key) - session_data, child_has_next, child_cursor = await self._set_cache( - redis_client=redis_client, - cache_key=cache_key, - ctx=ctx, - child_next_page=child_cursor, - **params, - ) - page = 1 - - has_more_in_session = len(session_data) > limit * page - can_rebuild = child_has_next and child_cursor is not None - - return FeedResult( - data=session_data[(page - 1) * limit :][:limit], - next_page=FeedResultNextPage(data={self.merger_id: FeedResultNextPageInside(page=page + 1, after=None)}), - # True while the current session still has pages OR the child - # can provide a new session (triggers rebuild on the next call). - has_next_page=bool(has_more_in_session or can_rebuild), - ) - - def build_plan( - self, - *, - ctx: ExecutionContext, - limit: int, - next_page: FeedResultNextPage, - **params: Any, - ) -> CallablePlan: - async def _run(executor: Any) -> FeedResult: - if ctx.redis_client is None: - raise ValueError("Redis client must be provided if using Merger View Session") - - if ctx.executor is None: - ctx.executor = executor - - result = await self._get_cache( - limit=limit, - next_page=next_page, - redis_client=ctx.redis_client, - ctx=ctx, - **params, - ) - - if self.shuffle: - shuffle(result.data) - return result - - return CallablePlan(fn=_run) diff --git a/smartfeed/models/__init__.py b/smartfeed/models/__init__.py new file mode 100644 index 0000000..e50a531 --- /dev/null +++ b/smartfeed/models/__init__.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from typing import Annotated, Union + +from pydantic import Field + +from .base import BaseNode, FeedResult, SmartFeedDebugInfo, FeedItem +from .subfeed import SubFeed +from .mixers import ( + MergerPercentage, + MergerPercentageItem, + MergerAppend, + MergerPositional, + MergerPercentageGradient, + MergerAppendDistribute, +) +from .wrapper import Wrapper, WrapperCache, WrapperRerank, WrapperDedup + +# --------------------------------------------------------------------------- +# FeedNode: discriminated union of all node types +# --------------------------------------------------------------------------- + +FeedNode = Annotated[ + Union[ + SubFeed, + Wrapper, + MergerAppend, + MergerPercentage, + MergerPercentageGradient, + MergerPositional, + MergerAppendDistribute, + ], + Field(discriminator="type"), +] + +# --------------------------------------------------------------------------- +# FeedConfig: top-level config model +# --------------------------------------------------------------------------- + +from pydantic import BaseModel + + +class FeedConfig(BaseModel): + version: str + feed: FeedNode # type: ignore[valid-type] + + +# Rebuild forward refs so that nested FeedNode fields resolve correctly +_types_ns = {"FeedNode": FeedNode} +for _model in ( + SubFeed, + Wrapper, + WrapperCache, + MergerAppend, + MergerPercentage, + MergerPercentageItem, + MergerPositional, + MergerPercentageGradient, + MergerAppendDistribute, + FeedConfig, +): + _model.model_rebuild(force=True, _types_namespace=_types_ns) + +__all__ = [ + "BaseNode", + "FeedResult", + "SmartFeedDebugInfo", + "FeedItem", + + "SubFeed", + "MergerPercentage", + "MergerPercentageItem", + "MergerAppend", + "MergerPositional", + "MergerPercentageGradient", + "MergerAppendDistribute", + "Wrapper", + "WrapperCache", + "WrapperRerank", + "WrapperDedup", + "FeedNode", + "FeedConfig", +] diff --git a/smartfeed/models/base.py b/smartfeed/models/base.py new file mode 100644 index 0000000..9df9e0f --- /dev/null +++ b/smartfeed/models/base.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import hashlib +import json +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class SmartFeedDebugInfo(BaseModel): + """Metadata stamped by SmartFeed on every item.""" + + # Always present (stamped by SubFeed) + source: str # subfeed_id (regular_tours, recommended_tours, promo_tours) + + # Stamped by FeedManager (top-level position in final feed) + smartfeed_position: Optional[int] = None # 0-based position in page + + # Stamped by subfeed methods (optional, source-specific) + strategy: Optional[str] = None # ML model strategy (model_hot_users, model_cold_users) + + # Stamped by rerank callable (optional, present only when rerank is configured) + rerank_position: Optional[int] = None # position after rerank (1-based) + rrf_score: Optional[float] = None # RRF score + feature_score: Optional[float] = None # feature score from ES coefficients + feature_position: Optional[int] = None # rank by feature score (1-based) + total_reranked: Optional[int] = None # total items in rerank batch + raw_params: Optional[Dict[str, Any]] = None # raw tour params from Redis + + model_config = ConfigDict(extra="allow") + + +class FeedItem(BaseModel): + """One item in the feed output. Wraps tour data + SmartFeed metadata.""" + + id: Any + smartfeed_debug_info: Optional[SmartFeedDebugInfo] = Field( + default=None, alias="_smartfeed_debug_info" + ) + + model_config = ConfigDict(extra="allow", populate_by_name=True) + + +class BaseNode(BaseModel): + dedup_priority: int = 0 + + def config_hash(self) -> str: + # Use json.dumps with sort_keys for deterministic hashing + raw = json.dumps(self.model_dump(), sort_keys=True, default=str) + return hashlib.md5(raw.encode()).hexdigest()[:8] + + +class FeedResult(BaseModel): + data: List + next_page: Dict[str, Any] + has_next_page: bool + + +def coerce_feed_node(value: Any) -> Any: + """Coerce a dict to the appropriate FeedNode model instance. + + Uses a lazy import of FeedNode to avoid circular imports. + Non-dict values are returned unchanged. + """ + if not isinstance(value, dict): + return value + # Lazy import to break circular dependency + from smartfeed.models import FeedNode # noqa: PLC0415 + from pydantic import TypeAdapter + return TypeAdapter(FeedNode).validate_python(value) diff --git a/smartfeed/models/mixers.py b/smartfeed/models/mixers.py new file mode 100644 index 0000000..3fb4241 --- /dev/null +++ b/smartfeed/models/mixers.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +from collections import defaultdict, deque +from typing import Any, Dict, List, Literal, Optional + +from pydantic import BaseModel, model_validator + +from .base import BaseNode, coerce_feed_node +from ..execution.plans import MixChild, MixPlan + + +def _merge_cursor(child_cursors: Dict[str, dict]) -> dict: + merged: dict = {} + for c in child_cursors.values(): + merged.update(c) + return merged + + +# --------------------------------------------------------------------------- +# MergerPercentageItem +# --------------------------------------------------------------------------- + +class MergerPercentageItem(BaseModel): + percentage: int + data: Any # BaseNode subclass + + @model_validator(mode="before") + @classmethod + def _coerce_data(cls, values: Any) -> Any: + if isinstance(values, dict) and isinstance(values.get("data"), dict): + values["data"] = coerce_feed_node(values["data"]) + return values + + +# --------------------------------------------------------------------------- +# MergerPercentage +# --------------------------------------------------------------------------- + +class MergerPercentage(BaseNode): + type: Literal["merger_percentage"] = "merger_percentage" + node_id: str + items: List[MergerPercentageItem] + dedup_priority: int = 0 + + def build_mix_plan( + self, + *, + ctx: Any, + limit: int, + cursor: dict, + ) -> MixPlan: + total_pct = sum(item.percentage for item in self.items) + + # Compute per-child demand with remainder distribution + demands: List[int] = [] + remainders: List[tuple] = [] + for idx, item in enumerate(self.items): + raw = limit * item.percentage + child_limit = raw // 100 + demands.append(max(0, child_limit)) + remainders.append((raw % 100, idx)) + + if total_pct == 100: + missing = max(0, limit - sum(demands)) + if missing > 0: + for _rem, idx in sorted(remainders, key=lambda x: (-x[0], x[1])): + if missing <= 0: + break + demands[idx] += 1 + missing -= 1 + + children = [ + MixChild( + node_id=f"{self.node_id}_{idx}", + node=item.data, + demand=demands[idx], + ) + for idx, item in enumerate(self.items) + ] + + def assemble( + buffers: Dict[str, list], + child_cursors: Dict[str, dict], + ): + # Simple concatenation: demand already ensures correct proportions + merged_data: List[Any] = [] + for child in children: + merged_data.extend(buffers.get(child.node_id, [])) + return merged_data, _merge_cursor(child_cursors) + + return MixPlan(children=children, assemble=assemble) + + +# --------------------------------------------------------------------------- +# MergerAppend +# --------------------------------------------------------------------------- + +class MergerAppend(BaseNode): + type: Literal["merger_append"] = "merger_append" + node_id: str + items: List[Any] # list of BaseNode subclasses + dedup_priority: int = 0 + + @model_validator(mode="before") + @classmethod + def _coerce_items(cls, values: Any) -> Any: + if isinstance(values, dict): + raw_items = values.get("items") + if isinstance(raw_items, list): + values["items"] = [ + coerce_feed_node(item) if isinstance(item, dict) else item + for item in raw_items + ] + return values + + def build_mix_plan( + self, + *, + ctx: Any, + limit: int, + cursor: dict, + ) -> MixPlan: + # Each child gets equal demand share; leftover goes to first children + n = len(self.items) + if n == 0: + def assemble_empty(buffers, child_cursors): + return [], _merge_cursor(child_cursors) + return MixPlan(children=[], assemble=assemble_empty) + + base_demand = limit // n + extra = limit - base_demand * n + + children = [ + MixChild( + node_id=f"{self.node_id}_{idx}", + node=item, + demand=base_demand + (1 if idx < extra else 0), + ) + for idx, item in enumerate(self.items) + ] + + def assemble( + buffers: Dict[str, list], + child_cursors: Dict[str, dict], + ): + merged_data: List[Any] = [] + for child in children: + merged_data.extend(buffers.get(child.node_id, [])) + # Trim to limit in case of overfill + return merged_data[:limit], _merge_cursor(child_cursors) + + return MixPlan(children=children, assemble=assemble) + + +# --------------------------------------------------------------------------- +# MergerPositional +# --------------------------------------------------------------------------- + +class MergerPositional(BaseNode): + type: Literal["merger_positional"] = "merger_positional" + node_id: str + positions: List[int] = [] + positional: Any # BaseNode subclass + default: Any # BaseNode subclass + dedup_priority: int = 0 + + @model_validator(mode="before") + @classmethod + def _coerce_children(cls, values: Any) -> Any: + if isinstance(values, dict): + for field in ("positional", "default"): + if isinstance(values.get(field), dict): + values[field] = coerce_feed_node(values[field]) + return values + + def build_mix_plan( + self, + *, + ctx: Any, + limit: int, + cursor: dict, + ) -> MixPlan: + # positions are 1-indexed; determine how many positional slots fall within [1..limit] + pos_slots = [p for p in self.positions if 1 <= p <= limit] + pos_count = len(pos_slots) + default_count = limit - pos_count + + positional_child = MixChild( + node_id=f"{self.node_id}_positional", + node=self.positional, + demand=pos_count, + ) + default_child = MixChild( + node_id=f"{self.node_id}_default", + node=self.default, + demand=default_count, + ) + + children = [positional_child, default_child] + + def assemble( + buffers: Dict[str, list], + child_cursors: Dict[str, dict], + ): + pos_items = deque(buffers.get(positional_child.node_id, [])) + def_items = deque(buffers.get(default_child.node_id, [])) + + result: List[Any] = [] + pos_set = set(pos_slots) + for position in range(1, limit + 1): + if position in pos_set and pos_items: + result.append(pos_items.popleft()) + elif def_items: + result.append(def_items.popleft()) + elif pos_items: + result.append(pos_items.popleft()) + + return result, _merge_cursor(child_cursors) + + return MixPlan(children=children, assemble=assemble) + + +# --------------------------------------------------------------------------- +# MergerPercentageGradient +# --------------------------------------------------------------------------- + +class MergerPercentageGradient(BaseNode): + """Percentage-based merger that shifts the ratio over pages.""" + + type: Literal["merger_percentage_gradient"] = "merger_percentage_gradient" + node_id: str + item_from: MergerPercentageItem + item_to: MergerPercentageItem + step: int + size_to_step: int + dedup_priority: int = 0 + + def _calculate_demands(self, page: int, limit: int) -> tuple: + percentage_from = self.item_from.percentage + percentage_to = self.item_to.percentage + start_position = limit * (page - 1) + first_iter = True + limit_from = 0 + limit_to = 0 + segments: List[Dict] = [] + + for i in range(self.size_to_step, limit * page + self.size_to_step, self.size_to_step): + if not first_iter and percentage_to < 100: + percentage_from -= self.step + percentage_to += self.step + if percentage_to > 100 or percentage_from < 0: + percentage_from = 0 + percentage_to = 100 + + if i > start_position: + iter_limit = ( + (limit * page - start_position) + if i > limit * page + else (i - start_position) + ) + start_position = i + from_take = iter_limit * percentage_from // 100 + to_take = iter_limit - from_take + limit_from += from_take + limit_to += to_take + segments.append({"limit": iter_limit, "from_take": from_take, "to_take": to_take}) + + if first_iter: + first_iter = False + + return limit_from, limit_to, segments + + def build_mix_plan( + self, + *, + ctx: Any, + limit: int, + cursor: dict, + ) -> MixPlan: + page = cursor.get(self.node_id, {}).get("page", 1) + limit_from, limit_to, segments = self._calculate_demands(page, limit) + + from_child = MixChild( + node_id=f"{self.node_id}_from", + node=self.item_from.data, + demand=max(0, limit_from), + ) + to_child = MixChild( + node_id=f"{self.node_id}_to", + node=self.item_to.data, + demand=max(0, limit_to), + ) + children = [from_child, to_child] + + def assemble( + buffers: Dict[str, list], + child_cursors: Dict[str, dict], + ): + from_data = list(buffers.get(from_child.node_id, [])) + to_data = list(buffers.get(to_child.node_id, [])) + result: List[Any] = [] + fi = ti = 0 + for seg in segments: + ft = int(seg["from_take"]) + tt = int(seg["to_take"]) + result.extend(from_data[fi: fi + ft]) + result.extend(to_data[ti: ti + tt]) + fi += ft + ti += tt + merged_cur = _merge_cursor(child_cursors) + merged_cur[self.node_id] = {"page": page + 1} + return result, merged_cur + + return MixPlan(children=children, assemble=assemble) + + +# --------------------------------------------------------------------------- +# MergerAppendDistribute +# --------------------------------------------------------------------------- + +class MergerAppendDistribute(BaseNode): + """Append merger that round-robins items by a distribution key.""" + + type: Literal["merger_distribute"] = "merger_distribute" + node_id: str + items: List[Any] # list of BaseNode subclasses + distribution_key: str + sorting_key: Optional[str] = None + sorting_desc: bool = False + dedup_priority: int = 0 + + @model_validator(mode="before") + @classmethod + def _coerce_items(cls, values: Any) -> Any: + if isinstance(values, dict): + raw_items = values.get("items") + if isinstance(raw_items, list): + values["items"] = [ + coerce_feed_node(item) if isinstance(item, dict) else item + for item in raw_items + ] + return values + + def _uniform_distribute(self, data: list) -> list: + if self.sorting_key: + data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) + + grouped_entries: Dict[Any, deque] = defaultdict(deque) + for entry in data: + grouped_entries[entry[self.distribution_key]].append(entry) + + result: List[Any] = [] + prev_key = None + while any(grouped_entries.values()): + for key in list(grouped_entries.keys()): + if grouped_entries[key]: + if key != prev_key or len(grouped_entries) == 1: + result.append(grouped_entries[key].popleft()) + prev_key = key + if not grouped_entries[key]: + del grouped_entries[key] + else: + del grouped_entries[key] + return result + + def build_mix_plan( + self, + *, + ctx: Any, + limit: int, + cursor: dict, + ) -> MixPlan: + children = [ + MixChild( + node_id=f"{self.node_id}_{idx}", + node=item, + demand=limit, + ) + for idx, item in enumerate(self.items) + ] + + def assemble( + buffers: Dict[str, list], + child_cursors: Dict[str, dict], + ): + all_items: List[Any] = [] + for child in children: + all_items.extend(buffers.get(child.node_id, [])) + distributed = self._uniform_distribute(all_items) + return distributed[:limit], _merge_cursor(child_cursors) + + return MixPlan(children=children, assemble=assemble) diff --git a/smartfeed/models/subfeed.py b/smartfeed/models/subfeed.py new file mode 100644 index 0000000..7d8fa72 --- /dev/null +++ b/smartfeed/models/subfeed.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from random import shuffle +from typing import Any, Dict, List, Literal + +from .base import BaseNode, FeedResult + + +class SubFeed(BaseNode): + type: Literal["subfeed"] = "subfeed" + subfeed_id: str + method_name: str + subfeed_params: Dict[str, Any] = {} + raise_error: bool = True + shuffle: bool = False + + async def execute( + self, + methods_dict: dict, + session_id: str, + limit: int, + cursor: dict, + **params: Any, + ) -> FeedResult: + method = methods_dict[self.method_name] + subfeed_cursor = cursor.get(self.subfeed_id, {}) + + try: + result: FeedResult = await method( + user_id=session_id, + limit=limit, + next_page=subfeed_cursor, + **params, + **self.subfeed_params, + ) + except Exception: + if self.raise_error: + raise + return FeedResult(data=[], next_page=cursor, has_next_page=False) + + if self.shuffle: + shuffle(result.data) + + # Stamp source on every dict item (merge, don't overwrite) + for item in result.data: + if isinstance(item, dict): + item.setdefault("_smartfeed_debug_info", {})["source"] = self.subfeed_id + + return FeedResult( + data=result.data, + next_page={self.subfeed_id: result.next_page}, + has_next_page=result.has_next_page, + ) diff --git a/smartfeed/models/wrapper.py b/smartfeed/models/wrapper.py new file mode 100644 index 0000000..808ea57 --- /dev/null +++ b/smartfeed/models/wrapper.py @@ -0,0 +1,553 @@ +from __future__ import annotations + +import asyncio +import secrets +from typing import Any, Dict, List, Literal, Optional + +from pydantic import BaseModel, model_validator + +import orjson +from smartfeed.execution import executor as _executor +from .base import BaseNode, FeedResult, coerce_feed_node + + +class WrapperCache(BaseModel): + session_size: int = 300 + session_ttl: int = 300 + + +class WrapperRerank(BaseModel): + method_name: str + raise_error: bool = True # True = crash if rerank fails, False = keep original order + + +class WrapperDedup(BaseModel): + dedup_key: str + missing_key_policy: Literal["error", "keep", "drop"] = "error" + overfetch_factor: int = 4 + max_refill_loops: int = 2 + state_ttl: int = 300 # TTL for Redis seen-set (seconds) + + +class Wrapper(BaseNode): + type: Literal["wrapper"] = "wrapper" + node_id: str + cache: Optional[WrapperCache] = None + rerank: Optional[WrapperRerank] = None + dedup: Optional[WrapperDedup] = None + data: Any # BaseNode subclass; typed as Any to support discriminated union deserialization + cache_key: Optional[str] = None + + @model_validator(mode="before") + @classmethod + def _coerce_data(cls, values: Any) -> Any: + if isinstance(values, dict) and isinstance(values.get("data"), dict): + values["data"] = coerce_feed_node(values["data"]) + return values + + async def execute( + self, + methods_dict: dict, + session_id: str, + limit: int, + cursor: dict, + ctx: Any = None, + **params: Any, + ) -> FeedResult: + if self.cache is None or ctx is None or ctx.redis is None: + return await self._passthrough(methods_dict, session_id, limit, cursor, ctx) + + return await self._execute_with_cache( + methods_dict, session_id, limit, cursor, ctx + ) + + # -- dedup ----------------------------------------------------------------- + + def _resolve_dedup_priorities(self) -> Dict[str, int]: + """Walk the config subtree to build {subfeed_id: priority} map.""" + result: Dict[str, int] = {} + self._collect_priorities(self.data, override_priority=0, out=result) + return result + + @staticmethod + def _collect_priorities( + node: BaseNode, override_priority: int, out: Dict[str, int] + ) -> None: + from .subfeed import SubFeed + + # Determine the effective priority for this subtree. + # A node with dedup_priority != 0 overrides all children. + effective = node.dedup_priority if node.dedup_priority != 0 else override_priority + + if isinstance(node, SubFeed): + out[node.subfeed_id] = effective + return + + # Walk children: look in known child-bearing attributes + for attr in ("items", "data", "positional", "default", + "item_from", "item_to"): + child = getattr(node, attr, None) + if child is None: + continue + if isinstance(child, list): + for item in child: + child_node = getattr(item, "data", item) + if isinstance(child_node, BaseNode): + Wrapper._collect_priorities(child_node, effective, out) + elif isinstance(child, BaseModel): + # MergerPercentageItem has .data + child_node = getattr(child, "data", child) + if isinstance(child_node, BaseNode): + Wrapper._collect_priorities(child_node, effective, out) + elif isinstance(child, BaseNode): + Wrapper._collect_priorities(child, effective, out) + + def _dedup(self, data: List, seen: Optional[set] = None) -> List: + """Deduplicate items by dedup_key with priority arbitration. + + Args: + data: items to deduplicate + seen: optional external seen-set (for cross-page dedup). Mutated in-place. + """ + if not self.dedup: + return data + + priorities = self._resolve_dedup_priorities() + key_field = self.dedup.dedup_key + policy = self.dedup.missing_key_policy + if seen is None: + seen = set() + + # {key_val: (priority, index_in_result)} for in-batch priority arbitration + batch: Dict[Any, tuple] = {} + result: List = [] + + for item in data: + if not isinstance(item, dict): + result.append(item) + continue + + if key_field not in item: + if policy == "error": + raise KeyError(f"Dedup key '{key_field}' missing from item: {item}") + elif policy == "keep": + result.append(item) + elif policy == "drop": + pass + continue + + key_val = str(item[key_field]) + + # Cross-page: already seen on previous pages + if key_val in seen: + continue + + # In-batch priority arbitration + source = (item.get("_smartfeed_debug_info") or {}).get("source") + item_priority = priorities.get(source, 0) if source else 0 + + if key_val not in batch: + idx = len(result) + batch[key_val] = (item_priority, idx) + result.append(item) + else: + existing_priority, existing_idx = batch[key_val] + if item_priority > existing_priority: + result[existing_idx] = item + batch[key_val] = (item_priority, existing_idx) + + seen.add(key_val) + + return result + + # -- rerank ---------------------------------------------------------------- + + async def _apply_rerank( + self, data: list, methods_dict: dict, session_id: str + ) -> list: + if not self.rerank: + return data + rerank_fn = methods_dict[self.rerank.method_name] + original_len = len(data) + try: + result = await rerank_fn(data, session_id) + except Exception: + if self.rerank.raise_error: + raise + # Fail-soft: keep original order + return data + if len(result) != original_len: + raise ValueError( + f"Rerank '{self.rerank.method_name}' must return exactly " + f"{original_len} items, got {len(result)}" + ) + return result + + # -- debug stamping -------------------------------------------------------- + + def _stamp_pre_rerank(self, data: list) -> None: + """Stamp position before rerank (order from tree/dedup).""" + for i, item in enumerate(data): + if isinstance(item, dict): + item.setdefault("_smartfeed_debug_info", {})[self.node_id] = { + "smartfeed_position": i, + } + + def _stamp_post_rerank(self, data: list) -> None: + """Stamp position after rerank.""" + for i, item in enumerate(data): + if isinstance(item, dict): + item.setdefault("_smartfeed_debug_info", {}).setdefault(self.node_id, {})["rerank_position"] = i + + # -- passthrough (no cache) ----------------------------------------------- + + async def _passthrough( + self, + methods_dict: dict, + session_id: str, + limit: int, + cursor: dict, + ctx: Any, + ) -> FeedResult: + if self.dedup: + # Load persisted seen-set from Redis (cross-page dedup) + seen_keys = await self._load_seen_set(ctx, session_id) + + overfetch = self.dedup.overfetch_factor + max_loops = self.dedup.max_refill_loops + fetch_size = limit * overfetch + current_cursor = dict(cursor) + + result = await _executor.run(self.data, ctx, fetch_size, current_cursor) + data = self._dedup(result.data, seen_keys) + current_cursor = result.next_page + has_next = result.has_next_page + + loop = 0 + while len(data) < limit and has_next and loop < max_loops: + deficit = limit - len(data) + result = await _executor.run(self.data, ctx, deficit * overfetch, current_cursor) + data.extend(self._dedup(result.data, seen_keys)) + current_cursor = result.next_page + has_next = result.has_next_page + loop += 1 + + data = data[:limit] + + # Persist seen-set to Redis for next page + new_keys = set() + key_field = self.dedup.dedup_key + for item in data: + if isinstance(item, dict) and key_field in item: + new_keys.add(str(item[key_field])) + await self._save_seen_set(ctx, session_id, new_keys) + else: + result = await _executor.run(self.data, ctx, limit, cursor) + data = result.data + current_cursor = result.next_page + has_next = result.has_next_page + + self._stamp_pre_rerank(data) + if self.rerank: + data = await self._apply_rerank(data, methods_dict, session_id) + self._stamp_post_rerank(data) + return FeedResult( + data=data, + next_page=current_cursor, + has_next_page=has_next, + ) + + def _seen_set_key(self, session_id: str) -> str: + return f"sf:{session_id}:{self.node_id}:{self.config_hash()}:seen" + + async def _load_seen_set(self, ctx: Any, session_id: str) -> set: + if not ctx or not ctx.redis: + return set() + key = self._seen_set_key(session_id) + members = await ctx.redis.smembers(key) + return {m.decode() if isinstance(m, bytes) else m for m in members} if members else set() + + async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None: + if not ctx or not ctx.redis or not new_keys: + return + key = self._seen_set_key(session_id) + await ctx.redis.sadd(key, *new_keys) + ttl = self.dedup.state_ttl + await ctx.redis.expire(key, ttl) + + # -- cache helpers --------------------------------------------------------- + + def _base_key(self, session_id: str) -> str: + key_part = self.cache_key or self.node_id + return f"sf:{session_id}:{key_part}:{self.config_hash()}" + + async def _read_cache(self, ctx: Any, session_id: str) -> Optional[List]: + base = self._base_key(session_id) + raw = await ctx.redis.get(base) + if raw is None: + return None + return orjson.loads(raw) + + async def _read_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: + base = self._base_key(session_id) + raw = await ctx.redis.get(f"{base}:meta") + if raw is None: + return None + return orjson.loads(raw) + + async def _write_cache( + self, + ctx: Any, + session_id: str, + data: List, + gen: str, + child_cursor: Dict, + child_has_next: bool = False, + ) -> None: + base = self._base_key(session_id) + ttl = self.cache.session_ttl + + pipe = ctx.redis.pipeline() + pipe.set(base, orjson.dumps(data), ex=ttl) + pipe.set( + f"{base}:meta", + orjson.dumps({"gen": gen, "child_cursor": child_cursor, "child_has_next": child_has_next}), + ex=ttl, + ) + await pipe.execute() + + async def _touch_ttl(self, ctx: Any, session_id: str) -> None: + if not self.cache: + return + base = self._base_key(session_id) + ttl = self.cache.session_ttl + pipe = ctx.redis.pipeline() + pipe.expire(base, ttl) + pipe.expire(f"{base}:meta", ttl) + await pipe.execute() + + # -- pagination ------------------------------------------------------------ + + def _paginate( + self, data: List, limit: int, page: int, gen: str, + child_has_next: bool = False, + ) -> FeedResult: + start = (page - 1) * limit + end = start + limit + page_data = data[start:end] + has_next = end < len(data) or child_has_next + + next_cursor = { + self.node_id: { + "page": page + 1, + "gen": gen, + } + } + + return FeedResult( + data=page_data, + next_page=next_cursor, + has_next_page=has_next, + ) + + # -- main cached flow ------------------------------------------------------ + + async def _execute_with_cache( + self, + methods_dict: dict, + session_id: str, + limit: int, + cursor: dict, + ctx: Any, + ) -> FeedResult: + my_cursor = cursor.get(self.node_id, {}) + cursor_gen = my_cursor.get("gen") + cursor_page = my_cursor.get("page", 1) + + # Warm path: try to read existing cache + if cursor_gen: + meta = await self._read_meta(ctx, session_id) + if meta and meta.get("gen") == cursor_gen: + cached_data = await self._read_cache(ctx, session_id) + if cached_data is not None: + # Check if we have enough data for this page + start = (cursor_page - 1) * limit + if start < len(cached_data): + await self._touch_ttl(ctx, session_id) + child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) + return self._paginate( + cached_data, limit, cursor_page, cursor_gen, + child_has_next=child_has_next, + ) + # Cache exhausted: rebuild with continuation cursor + return await self._cold_build( + methods_dict, + session_id, + limit, + ctx, + child_cursor=meta.get("child_cursor", {}), + ) + + # Cold path: no gen or stale gen -> fresh build + return await self._cold_build( + methods_dict, session_id, limit, ctx, child_cursor={} + ) + + # -- shared base cache helpers ---------------------------------------------- + + def _base_shared_key(self, session_id: str) -> str: + """Key for the shared base cache (deduped, NOT reranked), keyed by cache_key.""" + return f"sf:{session_id}:{self.cache_key}:{self.data.config_hash()}" + + async def _read_shared_base(self, ctx: Any, session_id: str) -> Optional[List]: + key = self._base_shared_key(session_id) + raw = await ctx.redis.get(key) + if raw is None: + return None + return orjson.loads(raw) + + async def _write_shared_base( + self, ctx: Any, session_id: str, data: List, child_cursor: Dict, + child_has_next: bool = False, + ) -> None: + key = self._base_shared_key(session_id) + ttl = self.cache.session_ttl + pipe = ctx.redis.pipeline() + pipe.set(key, orjson.dumps(data), ex=ttl) + pipe.set( + f"{key}:meta", + orjson.dumps({"child_cursor": child_cursor, "child_has_next": child_has_next}), + ex=ttl, + ) + await pipe.execute() + + async def _read_shared_base_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: + key = self._base_shared_key(session_id) + raw = await ctx.redis.get(f"{key}:meta") + if raw is None: + return None + return orjson.loads(raw) + + async def _fetch_and_dedup( + self, + ctx: Any, + target: int, + child_cursor: Dict, + ) -> tuple: + """Fetch target items from child with overfetch + refill loop for dedup. + + Returns (data, child_cursor, child_has_next). + """ + overfetch = self.dedup.overfetch_factor if self.dedup else 1 + max_loops = self.dedup.max_refill_loops if self.dedup else 0 + fetch_size = target * overfetch + cursor = dict(child_cursor) + + # First fetch + result = await _executor.run(self.data, ctx, fetch_size, cursor) + data = self._dedup(result.data) if self.dedup else result.data + cursor = result.next_page + has_next = result.has_next_page + + # Refill loop: keep fetching if dedup ate too many items. + # Accumulate seen keys so refill batches are deduped against all prior items. + loop = 0 + while len(data) < target and has_next and loop < max_loops: + deficit = target - len(data) + refill_size = deficit * overfetch + result = await _executor.run(self.data, ctx, refill_size, cursor) + # Dedup refill against ALL accumulated data + combined = data + result.data + combined = self._dedup(combined) if self.dedup else combined + data = combined + cursor = result.next_page + has_next = result.has_next_page + loop += 1 + + return data[:target], cursor, has_next + + async def _cold_build( + self, + methods_dict: dict, + session_id: str, + limit: int, + ctx: Any, + child_cursor: Dict, + ) -> FeedResult: + session_size = self.cache.session_size + + if self.cache_key is not None: + # Shared cache path: base data (fetch + dedup) is shared across wrappers + # with the same cache_key; rerank is applied per-wrapper after. + base_data = await self._build_shared_base( + methods_dict, session_id, ctx, child_cursor + ) + # Read child_cursor from shared base meta for continuation + shared_meta = await self._read_shared_base_meta(ctx, session_id) + child_cursor_out = shared_meta.get("child_cursor", {}) if shared_meta else {} + child_has_next = shared_meta.get("child_has_next", False) if shared_meta else False + # Per-wrapper: stamp pre-rerank, apply rerank, stamp post-rerank + data = list(base_data) + self._stamp_pre_rerank(data) + data = await self._apply_rerank(data, methods_dict, session_id) + if self.rerank: + self._stamp_post_rerank(data) + else: + # Standard path: fetch -> dedup (with overfetch/refill) -> rerank -> stamp -> cache + data, child_cursor_out, child_has_next = await self._fetch_and_dedup( + ctx, session_size, child_cursor + ) + self._stamp_pre_rerank(data) + data = await self._apply_rerank(data, methods_dict, session_id) + if self.rerank: + self._stamp_post_rerank(data) + + gen = secrets.token_hex(4) + await self._write_cache(ctx, session_id, data, gen, child_cursor_out, child_has_next) + return self._paginate(data, limit, 1, gen, child_has_next=child_has_next) + + async def _build_shared_base( + self, + methods_dict: dict, + session_id: str, + ctx: Any, + child_cursor: Dict, + ) -> List: + """Fetch and dedup the shared base data, using a distributed lock to ensure + only one caller fetches from child on a cold build. Others wait and read.""" + from smartfeed.execution.redis_lock import RedisLock + + session_size = self.cache.session_size + lock_key = f"{self._base_shared_key(session_id)}:lock" + + # Fast path: base already cached + existing = await self._read_shared_base(ctx, session_id) + if existing is not None: + return existing + + # Try to acquire the lock + async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: + if acquired: + # Double-check after acquiring lock + existing = await self._read_shared_base(ctx, session_id) + if existing is not None: + return existing + + base_data, child_cursor_out, child_has_next = await self._fetch_and_dedup( + ctx, session_size, child_cursor + ) + await self._write_shared_base(ctx, session_id, base_data, child_cursor_out, child_has_next) + return base_data + else: + # Another coroutine holds the lock -- poll until base cache appears + for _ in range(50): + await asyncio.sleep(0.1) + existing = await self._read_shared_base(ctx, session_id) + if existing is not None: + return existing + # Fallback: fetch ourselves if lock holder never wrote + existing = await self._read_shared_base(ctx, session_id) + if existing is not None: + return existing + data, _, _ = await self._fetch_and_dedup(ctx, session_size, child_cursor) + return data diff --git a/smartfeed/policies/dedup.py b/smartfeed/policies/dedup.py deleted file mode 100644 index e744931..0000000 --- a/smartfeed/policies/dedup.py +++ /dev/null @@ -1,241 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Literal, Optional, Tuple - -from .. import jsonlib as json -from .seen_store import SeenStore - -MissingKeyPolicy = Literal["error", "keep", "drop"] - - -def normalize_key(value: Any) -> str: - if isinstance(value, (str, int)): - return str(value) - if isinstance(value, (dict, list)): - return json.dumps(value, sort_keys=True, default=str) - return str(value) - - -def extract_dedup_value(item: Any, dedup_key: Optional[str], missing_key_policy: MissingKeyPolicy) -> Any: - if not dedup_key: - return item - - try: - value = item.get(dedup_key) - except AttributeError: - value = getattr(item, dedup_key, None) - - if value is None and missing_key_policy == "error": - raise AssertionError(f"Deduplication failed: entity {item} has no key or attr {dedup_key}") - return value - - -def entity_key(item: Any, dedup_key: Optional[str], missing_key_policy: MissingKeyPolicy) -> Optional[str]: - raw_value = extract_dedup_value(item, dedup_key, missing_key_policy) - if raw_value is None: - if missing_key_policy == "drop": - return None - if missing_key_policy == "keep": - raw_value = ("__missing__", id(item)) - return normalize_key(raw_value) - - -@dataclass -class DeduplicationPolicy: - """Deduplication policy applied during execution. - - This keeps dedup logic out of merger implementations and plan interpreters. - """ - - dedup_key: Optional[str] - missing_key_policy: MissingKeyPolicy - - # Store backend (cursor or redis) - store: SeenStore - - # Keys encountered/accepted within this request. Prevents duplicates inside one response. - seen_request_set: set[str] - - def key_for(self, item: Any) -> Optional[str]: - return entity_key(item, self.dedup_key, self.missing_key_policy) - - async def prefetch_keys(self, keys: List[str]) -> None: - if not keys: - return - - filtered: List[str] = [] - seen: set[str] = set() - for k in keys: - if k in self.seen_request_set: - continue - if k in seen: - continue - seen.add(k) - filtered.append(k) - - if not filtered: - return - - await self.store.prefetch(filtered) - - def should_accept(self, key: str, priority: int) -> bool: - if key in self.seen_request_set: - return False - - existing_priority = self.store.get(key) - if existing_priority is not None and priority <= existing_priority: - return False - return True - - def record(self, key: str, priority: int) -> None: - self.seen_request_set.add(key) - self.store.set_max(key, priority) - - def create_isolated(self) -> "DeduplicationPolicy": - """Create an isolated copy for inner execution contexts. - - Uses a fresh in-memory store and empty seen_request_set so that - inner priority arbitration works without contaminating the caller's - dedup state. - """ - from .seen_store import CursorSeenStore - - return DeduplicationPolicy( - dedup_key=self.dedup_key, - missing_key_policy=self.missing_key_policy, - store=CursorSeenStore( - cursor_compress=True, - cursor_max_keys=None, - seen_priority_map={}, - seen_order=[], - ), - seen_request_set=set(), - ) - - async def arbitrate_owner_buffers( - self, - *, - owners: List[Any], - owner_buffers: Dict[int, List[Any]], - owner_rank: Dict[int, int], - ) -> Dict[int, List[Any]]: - """Arbitrate winners across multiple owners. - - - Deterministic tie-break: (-priority, owner_rank, item_rank) - - Records winners into the store - - Returns per-owner buffers containing only accepted winners - """ - - keys_to_prefetch: List[str] = [] - keys_seen_local: set[str] = set() - for owner in owners: - owner_id = id(owner) - for item in owner_buffers.get(owner_id, []): - key = self.key_for(item) - if key is None: - continue - if key in keys_seen_local: - continue - keys_seen_local.add(key) - keys_to_prefetch.append(key) - - if keys_to_prefetch: - await self.prefetch_keys(keys_to_prefetch) - - winners: Dict[str, int] = {} - winner_prio: Dict[str, int] = {} - winner_tie: Dict[str, Tuple[int, int, int]] = {} - - for owner in owners: - owner_id = id(owner) - prio = int(getattr(owner, "dedup_priority", 0)) - rank = int(owner_rank.get(owner_id, 0)) - for item_rank, item in enumerate(owner_buffers.get(owner_id, [])): - key = self.key_for(item) - if key is None: - continue - tie = (-prio, rank, item_rank) - existing = winner_tie.get(key) - if existing is None or tie < existing: - winners[key] = owner_id - winner_prio[key] = prio - winner_tie[key] = tie - - for key, _tie in sorted(winner_tie.items(), key=lambda kv: kv[1]): - winner_owner_id = winners.get(key) - if winner_owner_id is None: - continue - prio = int(winner_prio.get(key, 0)) - if not self.should_accept(key, prio): - continue - self.record(key, prio) - - request_set = self.seen_request_set - per_owner_accepted: Dict[int, List[Any]] = {id(o): [] for o in owners} - for owner in owners: - owner_id = id(owner) - accepted: List[Any] = [] - for item in owner_buffers.get(owner_id, []): - key = self.key_for(item) - if key is None: - continue - if winners.get(key) != owner_id: - continue - if key not in request_set: - continue - accepted.append(item) - per_owner_accepted[owner_id] = accepted - - return per_owner_accepted - - async def accept_batch( - self, - *, - items: List[Any], - priority: int, - limit: Optional[int] = None, - ) -> Tuple[List[Any], int]: - """Accept items from a single stream in order. - - Returns accepted items and the number of inspected items. - """ - - if not items: - return [], 0 - - keys_to_prefetch: List[str] = [] - keys_seen_local: set[str] = set() - for item in items: - key = self.key_for(item) - if key is None: - continue - if key in self.seen_request_set: - continue - if key in keys_seen_local: - continue - keys_seen_local.add(key) - keys_to_prefetch.append(key) - - if keys_to_prefetch: - await self.prefetch_keys(keys_to_prefetch) - - accepted: List[Any] = [] - inspected_count = 0 - max_accept = int(limit) if limit is not None else len(items) - - for idx, item in enumerate(items, start=1): - inspected_count = idx - if len(accepted) >= max_accept: - break - key = self.key_for(item) - if key is None: - continue - if not self.should_accept(key, priority): - continue - accepted.append(item) - self.record(key, priority) - if len(accepted) >= max_accept: - break - - return accepted, inspected_count diff --git a/smartfeed/policies/dedup_utils.py b/smartfeed/policies/dedup_utils.py deleted file mode 100644 index 4483619..0000000 --- a/smartfeed/policies/dedup_utils.py +++ /dev/null @@ -1,121 +0,0 @@ -from __future__ import annotations - -import asyncio -import base64 -import zlib -from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union, cast - -import redis -from redis.asyncio import Redis as AsyncRedis - -from .. import jsonlib as json -from ..feed_models import _is_async_redis_client, _redis_call - - -def _seen_entries_to_map(entries: Any) -> Dict[str, int]: - """Coerce a legacy cursor "seen" list into a {key: priority} map. - - Supports: - - ["k1", "k2", ...] (implies priority 0) - - [["k1", 10], ["k2", 3], ...] (explicit priorities) - """ - - seen_map: Dict[str, int] = {} - if not isinstance(entries, list): - return seen_map - - for entry_item in entries: - if isinstance(entry_item, (list, tuple)) and len(entry_item) == 2: - seen_map[str(entry_item[0])] = int(entry_item[1]) - else: - seen_map[str(entry_item)] = 0 - return seen_map - - -def decode_seen_from_cursor(after: Any) -> Dict[str, int]: - if after is None: - return {} - - if isinstance(after, dict) and "z" in after: - if after.get("v") != 2: - return {} - if after.get("c") != "zlib+base64": - return {} - payload = base64.urlsafe_b64decode(str(after["z"]).encode()) - raw = zlib.decompress(payload).decode() - decoded = json.loads(raw) - if isinstance(decoded, list): - return _seen_entries_to_map(decoded) - return {} - - if isinstance(after, dict) and "seen" in after: - if after.get("v") != 2: - return {} - return _seen_entries_to_map(list(after["seen"])) - - return {} - - -def encode_seen_for_cursor( - seen_updates_in_order: List[Tuple[str, int]], - *, - cursor_compress: bool, - cursor_max_keys: Optional[int], -) -> Any: - if cursor_max_keys is not None: - seen_updates_in_order = seen_updates_in_order[-cursor_max_keys:] - - if not cursor_compress: - return {"v": 2, "seen": [[k, p] for k, p in seen_updates_in_order]} - - raw = json.dumps([[k, p] for k, p in seen_updates_in_order]).encode() - compressed = zlib.compress(raw) - return { - "v": 2, - "c": "zlib+base64", - "n": len(seen_updates_in_order), - "z": base64.urlsafe_b64encode(compressed).decode(), - } - - -async def redis_zmscore( - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - members: List[str], -) -> List[Optional[float]]: - if not members: - return [] - - if getattr(redis_client, "zmscore", None) is not None: - res = await _redis_call(redis_client, "zmscore", key, members) - return [None if v is None else float(v) for v in list(res)] - - if not _is_async_redis_client(redis_client): - - def _sync_pipeline_execute() -> Any: - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - return pipe.execute() - - res = await asyncio.to_thread(_sync_pipeline_execute) - return [None if v is None else float(v) for v in list(res)] - - pipe = redis_client.pipeline() - for m in members: - pipe.zscore(key, m) - res = await cast(Awaitable[Any], pipe.execute()) - return [None if v is None else float(v) for v in list(res)] - - -async def redis_zadd_and_expire( - redis_client: Union[redis.Redis, AsyncRedis], - key: str, - member_scores: Dict[str, int], - *, - ttl_seconds: int, -) -> None: - if not member_scores: - return - await _redis_call(redis_client, "zadd", key, mapping={m: float(s) for m, s in member_scores.items()}) - await _redis_call(redis_client, "expire", key, ttl_seconds) diff --git a/smartfeed/policies/seen_store.py b/smartfeed/policies/seen_store.py deleted file mode 100644 index c4712c3..0000000 --- a/smartfeed/policies/seen_store.py +++ /dev/null @@ -1,159 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Protocol, Tuple, Union - -import redis -from redis.asyncio import Redis as AsyncRedis - -from ..feed_models import _redis_call -from .dedup_utils import decode_seen_from_cursor, encode_seen_for_cursor, redis_zadd_and_expire, redis_zmscore - - -class SeenStore(Protocol): - async def prefetch(self, keys: List[str]) -> None: - raise NotImplementedError - - def get(self, key: str) -> Optional[int]: - raise NotImplementedError - - def set_max(self, key: str, priority: int) -> None: - raise NotImplementedError - - async def reset(self) -> None: - raise NotImplementedError - - async def commit(self) -> Any: - raise NotImplementedError - - -@dataclass -class CursorSeenStore: - """Seen-store that persists in the cursor (`after`).""" - - cursor_compress: bool - cursor_max_keys: Optional[int] - - seen_priority_map: Dict[str, int] - seen_order: List[str] - - @classmethod - def from_after( - cls, - after: Any, - *, - cursor_compress: bool, - cursor_max_keys: Optional[int], - ) -> "CursorSeenStore": - seen_priority_map = decode_seen_from_cursor(after) - return cls( - cursor_compress=cursor_compress, - cursor_max_keys=cursor_max_keys, - seen_priority_map=seen_priority_map, - seen_order=list(seen_priority_map.keys()), - ) - - async def prefetch(self, keys: List[str]) -> None: - return None - - def get(self, key: str) -> Optional[int]: - return self.seen_priority_map.get(key) - - def set_max(self, key: str, priority: int) -> None: - existing = self.seen_priority_map.get(key) - if existing is not None and priority <= existing: - return - self.seen_priority_map[key] = priority - if key in self.seen_order: - self.seen_order.remove(key) - self.seen_order.append(key) - - async def reset(self) -> None: - self.seen_priority_map.clear() - self.seen_order.clear() - - async def commit(self) -> Any: - # Persist the full snapshot so dedup state survives beyond 2 pages. - seen_snapshot_in_order: List[Tuple[str, int]] = [ - (key, self.seen_priority_map[key]) for key in self.seen_order if key in self.seen_priority_map - ] - return encode_seen_for_cursor( - seen_snapshot_in_order, - cursor_compress=self.cursor_compress, - cursor_max_keys=self.cursor_max_keys, - ) - - -@dataclass -class RedisSeenStore: - """Seen-store backed by a Redis zset (member=key, score=priority).""" - - redis_client: Union[redis.Redis, AsyncRedis] - redis_key: str - ttl_seconds: int - - redis_seen_cache: Dict[str, Optional[int]] - redis_new_scores: Dict[str, int] - - @classmethod - def create( - cls, - *, - redis_client: Union[redis.Redis, AsyncRedis], - redis_key: str, - ttl_seconds: int, - ) -> "RedisSeenStore": - return cls( - redis_client=redis_client, - redis_key=redis_key, - ttl_seconds=ttl_seconds, - redis_seen_cache={}, - redis_new_scores={}, - ) - - async def prefetch(self, keys: List[str]) -> None: - if not keys: - return - - unique: List[str] = [] - seen: set[str] = set() - for k in keys: - if k in self.redis_seen_cache: - continue - if k in seen: - continue - seen.add(k) - unique.append(k) - - if not unique: - return - - scores = await redis_zmscore(self.redis_client, self.redis_key, unique) - - for k, s in zip(unique, scores): - self.redis_seen_cache[k] = None if s is None else int(s) - - def get(self, key: str) -> Optional[int]: - return self.redis_seen_cache.get(key) - - def set_max(self, key: str, priority: int) -> None: - existing = self.redis_seen_cache.get(key) - if existing is not None and priority <= existing: - return - self.redis_seen_cache[key] = priority - self.redis_new_scores[key] = max(self.redis_new_scores.get(key, 0), priority) - - async def reset(self) -> None: - await _redis_call(self.redis_client, "delete", self.redis_key) - self.redis_seen_cache.clear() - self.redis_new_scores.clear() - - async def commit(self) -> Any: - await redis_zadd_and_expire( - self.redis_client, - self.redis_key, - self.redis_new_scores, - ttl_seconds=self.ttl_seconds, - ) - self.redis_new_scores.clear() - return None diff --git a/smartfeed/pydantic_compat.py b/smartfeed/pydantic_compat.py deleted file mode 100644 index 185e004..0000000 --- a/smartfeed/pydantic_compat.py +++ /dev/null @@ -1,16 +0,0 @@ -from __future__ import annotations - -from typing import Any, Mapping, Type, TypeVar - -T = TypeVar("T") - - -def parse_model(model_cls: Type[T], obj: Mapping[str, Any]) -> T: - """Parse a mapping into a Pydantic model. - - Uses Pydantic v2 `model_validate` when available, otherwise falls back to v1 `parse_obj`. - """ - - if hasattr(model_cls, "model_validate"): - return model_cls.model_validate(obj) # type: ignore[attr-defined] - return model_cls.parse_obj(obj) # type: ignore[attr-defined] diff --git a/smartfeed/schemas.py b/smartfeed/schemas.py deleted file mode 100644 index 298cde7..0000000 --- a/smartfeed/schemas.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Public schema surface. - -This module keeps the public import path (`smartfeed.schemas`) stable while -moving merger implementations into `smartfeed.mergers.*`. -""" - -from __future__ import annotations - -from typing import Annotated, Any, Union - -from pydantic import BaseModel, Field - -from .feed_models import ( - BaseFeedConfigModel, - FeedResult, - FeedResultClient, - FeedResultNextPage, - FeedResultNextPageInside, - SubFeed, -) -from .mergers import ( - MergerAppend, - MergerAppendDistribute, - MergerDeduplication, - MergerPercentage, - MergerPercentageGradient, - MergerPercentageItem, - MergerPositional, - MergerViewSession, -) - -FeedTypes = Annotated[ - Union[ - MergerDeduplication, - MergerAppend, - MergerAppendDistribute, - MergerPositional, - MergerPercentage, - MergerPercentageGradient, - MergerViewSession, - SubFeed, - ], - Field(discriminator="type"), -] - - -class FeedConfig(BaseModel): - """Top-level feed config model.""" - - version: str - feed: FeedTypes - - -def _rebuild_model(model: Any) -> None: - """Resolve forward refs across modules (Pydantic v1/v2 compatible).""" - - if hasattr(model, "model_rebuild"): - model.model_rebuild(force=True, _types_namespace={"FeedTypes": FeedTypes}) - else: - model.update_forward_refs(FeedTypes=FeedTypes) - - -for _m in ( - MergerPositional, - MergerPercentage, - MergerPercentageItem, - MergerAppend, - MergerAppendDistribute, - MergerPercentageGradient, - MergerViewSession, - MergerDeduplication, - SubFeed, - FeedConfig, -): - _rebuild_model(_m) - - -__all__ = [ - "BaseFeedConfigModel", - "FeedResult", - "FeedResultClient", - "FeedResultNextPage", - "FeedResultNextPageInside", - "SubFeed", - "MergerAppend", - "MergerAppendDistribute", - "MergerDeduplication", - "MergerPercentage", - "MergerPercentageGradient", - "MergerPercentageItem", - "MergerPositional", - "MergerViewSession", - "FeedTypes", - "FeedConfig", -] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..5c79e4e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,27 @@ +from smartfeed.models.base import FeedResult + + +async def make_items(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": start + i, "title": f"item_{start + i}"} for i in range(limit)] + return FeedResult( + data=data, + next_page={"page": page + 1}, + has_next_page=True, + ) + + +async def make_empty(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: + return FeedResult(data=[], next_page=next_page, has_next_page=False) + + +async def make_error(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: + raise RuntimeError("subfeed error") + + +METHODS = { + "items": make_items, + "empty": make_empty, + "error": make_error, +} diff --git a/tests/fixtures/configs.py b/tests/fixtures/configs.py deleted file mode 100644 index 642d169..0000000 --- a/tests/fixtures/configs.py +++ /dev/null @@ -1,124 +0,0 @@ -from typing import Callable, Dict - -from smartfeed.examples.example_client import ClientMixerClass - -METHODS_DICT: Dict[str, Callable] = { - "ads": ClientMixerClass().example_method, - "followings": ClientMixerClass().example_method, - "large": ClientMixerClass().large_method, - "empty": ClientMixerClass().empty_method, - "error": ClientMixerClass().error_method, - "doubles": ClientMixerClass().doubles_method, - "posted": ClientMixerClass().keys_method, -} - -PARSING_CONFIG_FIXTURE = { - "version": "1", - "feed": { - "merger_id": "merger_positional_parsing_example", - "type": "merger_positional", - "positions": [1, 3, 15], - "start": 17, - "end": 200, - "step": 2, - "positional": { - "merger_id": "merger_append_parsing_example", - "type": "merger_append", - "items": [ - { - "subfeed_id": "subfeed_merger_append_parsing_example", - "type": "subfeed", - "method_name": "ads", - "subfeed_params": { - "limit_to_return": 10, - }, - }, - { - "merger_id": "merger_percentage_gradient_parsing_example", - "type": "merger_percentage_gradient", - "item_from": { - "percentage": 80, - "data": { - "subfeed_id": "subfeed_from_merger_percentage_gradient_parsing_example", - "type": "subfeed", - "method_name": "followings", - }, - }, - "item_to": { - "percentage": 20, - "data": { - "subfeed_id": "subfeed_to_merger_percentage_gradient_parsing_example", - "type": "subfeed", - "method_name": "followings", - }, - }, - "step": 10, - "size_to_step": 30, - "shuffle": False, - }, - { - "merger_id": "merger_view_session_parsing_example", - "type": "merger_view_session", - "session_size": 800, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_merger_view_session_parsing_example", - "type": "subfeed", - "method_name": "followings", - }, - }, - ], - }, - "default": { - "merger_id": "merger_percentage_parsing_example", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 100, - "data": { - "subfeed_id": "subfeed_merger_percentage_parsing_example", - "type": "subfeed", - "method_name": "followings", - "raise_error": False, - }, - }, - ], - }, - }, -} - - -PARSING_DEDUP_CONFIG_FIXTURE = { - "version": "1", - "feed": { - "merger_id": "merger_deduplication_parsing_example", - "type": "merger_deduplication", - "dedup_key": "id", - "state_backend": "cursor", - "cursor_compress": True, - "data": { - "merger_id": "merger_percentage_inside_dedup_parsing_example", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 50, - "data": { - "subfeed_id": "subfeed_dedup_a", - "type": "subfeed", - "method_name": "posted", - }, - }, - { - "percentage": 50, - "data": { - "subfeed_id": "subfeed_dedup_b", - "type": "subfeed", - "method_name": "posted", - }, - }, - ], - }, - }, -} diff --git a/tests/fixtures/dedup_helpers.py b/tests/fixtures/dedup_helpers.py deleted file mode 100644 index eb3fbdc..0000000 --- a/tests/fixtures/dedup_helpers.py +++ /dev/null @@ -1,441 +0,0 @@ -from smartfeed.schemas import FeedResultClient, FeedResultNextPage - - -def _effective_limit(limit, max_per_call): - effective_limit = limit - if isinstance(max_per_call, int) and max_per_call > 0: - effective_limit = min(effective_limit, max_per_call) - return effective_limit - - -def make_offset_paged_method(items, *, max_per_call=None): - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - offset = int(next_page.after or 0) - effective_limit = _effective_limit(limit, max_per_call) - result_data = items[offset : offset + effective_limit] - next_page.after = offset + len(result_data) - next_page.page += 1 - has_next_page = (offset + len(result_data)) < len(items) - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) - - return _method - - -def make_string_after_paged_method(items, *, max_per_call=None, after_field="created_at"): - """A subfeed method whose cursor is a string (e.g. timestamp). - - Cursor semantics: `after` is the last returned `created_at` value (monotonic). - """ - - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - effective_limit = _effective_limit(limit, max_per_call) - - after = next_page.after - start_idx = 0 - if isinstance(after, str) and after: - # Find first item with created_at > after - for i, item in enumerate(items): - if str(item[after_field]) > after: - start_idx = i - break - else: - start_idx = len(items) - - result_data = items[start_idx : start_idx + effective_limit] - has_next_page = (start_idx + len(result_data)) < len(items) - - if result_data: - next_page.after = str(result_data[-1][after_field]) - next_page.page += 1 - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) - - return _method - - -def make_profile_dict_after_method( - profiles_to_items, - *, - max_per_call=None, - after_key="after", -): - """A subfeed method whose cursor is a dict of per-profile offsets. - - Example shape: after = {"p1": 0, "p2": 0} - Cursor semantics: each profile offset increments as items are *read*. - """ - - profile_ids = list(profiles_to_items.keys()) - - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - effective_limit = _effective_limit(limit, max_per_call) - - after = next_page.after - if not isinstance(after, dict): - after = {pid: 0 for pid in profile_ids} - else: - after = dict(after) - for pid in profile_ids: - after.setdefault(pid, 0) - - result = [] - has_next_page = False - - # Build a cyclic iteration over profiles. - active_profiles = [pid for pid in profile_ids] - - i = 0 - while active_profiles and len(result) < effective_limit: - pid = active_profiles[i % len(active_profiles)] - idx = after.get(pid, 0) - items = profiles_to_items.get(pid, []) - - if idx >= len(items): - # This profile is exhausted. - active_profiles.remove(pid) - continue - - result.append(items[idx]) - after[pid] = idx + 1 - i += 1 - - # Determine if any profile still has unread items. - for pid in profile_ids: - if after.get(pid, 0) < len(profiles_to_items.get(pid, [])): - has_next_page = True - break - - next_page.after = after - next_page.page += 1 - return FeedResultClient(data=result, next_page=next_page, has_next_page=has_next_page) - - return _method - - -def _assert_cursor_monotonic_if_present(res_1, res_2, keys): - for key in keys: - if key not in res_1.next_page.data: - continue - - assert key in res_2.next_page.data - - after_1 = res_1.next_page.data[key].after - after_2 = res_2.next_page.data[key].after - - if after_1 is None or after_2 is None: - continue - - if isinstance(after_1, int) and isinstance(after_2, int): - assert after_2 >= after_1 - continue - - if isinstance(after_1, dict) and isinstance(after_2, dict): - continue - - try: - assert after_2 >= after_1 - except TypeError: - pass - - -def _sources(data): - return [x.get("src") for x in data] - - -def _ids(data): - return [x.get("id") for x in data] - - -def _assert_no_dupes_in_page(data): - ids = _ids(data) - assert len(ids) == len(set(ids)) - - -def _assert_pages_no_overlap(res_1, res_2): - assert not (set(_ids(res_1.data)) & set(_ids(res_2.data))) - - -def _assert_two_pages_no_dupes(res_1, res_2): - _assert_no_dupes_in_page(res_1.data) - _assert_no_dupes_in_page(res_2.data) - _assert_pages_no_overlap(res_1, res_2) - - -def _assert_sources_at_positions(data, positions, expected_src): - sources = _sources(data) - for pos in positions: - assert sources[pos - 1] == expected_src - - -def make_items(src, start, end, *, user_id_mod=None, id_offset=0, extra=None): - items = [] - for i in range(start, end): - item_id = id_offset + i - item = {"id": item_id, "src": src} - if user_id_mod is not None: - item["user_id"] = f"u{item_id % user_id_mod}" - if extra: - item.update(extra) - items.append(item) - return items - - -def _subfeed(subfeed_id, method_name, *, dedup_priority=None): - data = {"subfeed_id": subfeed_id, "type": "subfeed", "method_name": method_name} - if dedup_priority is not None: - data["dedup_priority"] = dedup_priority - return data - - -def _dedup_config(merger_id, data, *, dedup_key="id", state_backend="cursor", cursor_compress=True, **kwargs): - config = { - "merger_id": merger_id, - "type": "merger_deduplication", - "dedup_key": dedup_key, - "state_backend": state_backend, - "cursor_compress": cursor_compress, - "data": data, - } - config.update(kwargs) - return config - - -def _percentage_config(merger_id, items, *, shuffle=False): - return {"merger_id": merger_id, "type": "merger_percentage", "shuffle": shuffle, "items": items} - - -def _append_config(merger_id, items, *, shuffle=False): - return {"merger_id": merger_id, "type": "merger_append", "shuffle": shuffle, "items": items} - - -def _distribute_config(merger_id, items, *, distribution_key="user_id"): - return { - "merger_id": merger_id, - "type": "merger_distribute", - "distribution_key": distribution_key, - "items": items, - } - - -def _positional_config(merger_id, *, positions, positional, default): - return { - "merger_id": merger_id, - "type": "merger_positional", - "positions": positions, - "positional": positional, - "default": default, - } - - -def _gradient_config( - merger_id, - *, - item_from, - item_to, - step, - size_to_step, - shuffle=False, -): - return { - "merger_id": merger_id, - "type": "merger_percentage_gradient", - "item_from": item_from, - "item_to": item_to, - "step": step, - "size_to_step": size_to_step, - "shuffle": shuffle, - } - - -async def _run_two_pages(merger, methods_dict, limit, *, next_page=None, **kwargs): - if next_page is None: - next_page = FeedResultNextPage(data={}) - res_1 = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=limit, next_page=next_page, **kwargs) - res_2 = await merger.get_data( - methods_dict=methods_dict, user_id="u", limit=limit, next_page=res_1.next_page, **kwargs - ) - return res_1, res_2 - - -def _percentage_items(first, second, *, first_pct=50, second_pct=50): - return [ - {"percentage": first_pct, "data": first}, - {"percentage": second_pct, "data": second}, - ] - - -def _two_subfeed_spec(*, name="a", subfeed_id="sf_a", max_per_call=None, dedup_priority=None): - return { - "name": name, - "subfeed_id": subfeed_id, - "max_per_call": max_per_call, - "dedup_priority": dedup_priority, - } - - -def _build_two_subfeed_methods(items_a, items_b, *, spec_a=None, spec_b=None): - if spec_a is None: - spec_a = _two_subfeed_spec() - if spec_b is None: - spec_b = _two_subfeed_spec(name="b", subfeed_id="sf_b") - - methods_dict = { - spec_a["name"]: make_offset_paged_method(items_a, max_per_call=spec_a["max_per_call"]), - spec_b["name"]: make_offset_paged_method(items_b, max_per_call=spec_b["max_per_call"]), - } - subfeed_a = _subfeed(spec_a["subfeed_id"], spec_a["name"], dedup_priority=spec_a["dedup_priority"]) - subfeed_b = _subfeed(spec_b["subfeed_id"], spec_b["name"], dedup_priority=spec_b["dedup_priority"]) - return methods_dict, subfeed_a, subfeed_b - - -def _build_two_subfeed_dedup_merger( - *, - items_a, - items_b, - child_builder, - merger_id, - spec_a=None, - spec_b=None, - dedup_kwargs=None, -): - methods_dict, subfeed_a, subfeed_b = _build_two_subfeed_methods( - items_a, - items_b, - spec_a=spec_a, - spec_b=spec_b, - ) - config = _dedup_config(merger_id, child_builder(subfeed_a, subfeed_b), **(dedup_kwargs or {})) - return config, methods_dict, subfeed_a, subfeed_b - - -def _build_deep_positional_pct_dedup_merger( - *, - items_p, - items_d1, - items_d2, - dedup_merger_id, - pos_merger_id, - pct_merger_id, - positions, - overfetch_factor=None, - max_refill_loops=None, -): - methods_dict = { - "p": make_offset_paged_method(items_p), - "d1": make_offset_paged_method(items_d1), - "d2": make_offset_paged_method(items_d2), - } - - dedup_kwargs = {} - if overfetch_factor is not None: - dedup_kwargs["overfetch_factor"] = overfetch_factor - if max_refill_loops is not None: - dedup_kwargs["max_refill_loops"] = max_refill_loops - - config = _dedup_config( - dedup_merger_id, - _positional_config( - pos_merger_id, - positions=positions, - positional=_subfeed("sf_p", "p"), - default=_percentage_config( - pct_merger_id, - items=_percentage_items(_subfeed("sf_d1", "d1"), _subfeed("sf_d2", "d2")), - ), - ), - **dedup_kwargs, - ) - return config, methods_dict - - -def _inner_append_config(*, merger_id: str, subfeed_id: str, method_name: str, dedup_priority: int): - return { - "merger_id": merger_id, - "type": "merger_append", - # Important: dedup deletion priority must be visible at this node so parent mergers - # can fetch higher-priority subtrees first when a dedup wrapper is active. - "dedup_priority": dedup_priority, - "shuffle": False, - "items": [ - { - "subfeed_id": subfeed_id, - "type": "subfeed", - "method_name": method_name, - "dedup_priority": dedup_priority, - } - ], - } - - -def _build_deep_priority_tree_for_merger_type(*, merger_type: str): - """Return a deep tree config where low/high leaves overlap by id. - - The inner leaves are wrapped into an append merger to ensure a "deep" tree even - when the outer merger is flat. - """ - - low = _inner_append_config(merger_id="inner_low", subfeed_id="sf_low", method_name="low", dedup_priority=0) - high = _inner_append_config(merger_id="inner_high", subfeed_id="sf_high", method_name="high", dedup_priority=100) - - if merger_type == "merger_append": - return { - "merger_id": "outer_append", - "type": "merger_append", - "shuffle": False, - # Put low first intentionally; priority must still make high win for overlapping ids. - "items": [low, high], - } - - if merger_type == "merger_distribute": - return { - "merger_id": "outer_dist", - "type": "merger_distribute", - "distribution_key": "user_id", - # Put low first intentionally. - "items": [low, high], - } - - if merger_type == "merger_percentage": - return { - "merger_id": "outer_pct", - "type": "merger_percentage", - "shuffle": False, - "items": [ - {"percentage": 50, "data": low}, - {"percentage": 50, "data": high}, - ], - } - - if merger_type == "merger_percentage_gradient": - return { - "merger_id": "outer_grad", - "type": "merger_percentage_gradient", - "item_from": {"percentage": 60, "data": low}, - "item_to": {"percentage": 40, "data": high}, - "step": 20, - "size_to_step": 5, - "shuffle": False, - } - - if merger_type == "merger_positional": - # High priority on positional branch so it must win duplicates. - high_pos = _inner_append_config( - merger_id="inner_pos_high", - subfeed_id="sf_high", - method_name="high", - dedup_priority=100, - ) - low_def = _inner_append_config( - merger_id="inner_def_low", - subfeed_id="sf_low", - method_name="low", - dedup_priority=0, - ) - return { - "merger_id": "outer_pos", - "type": "merger_positional", - "positions": [1, 3, 5, 7, 9, 11], - "positional": high_pos, - "default": low_def, - } - - raise AssertionError(f"Unknown merger_type: {merger_type}") diff --git a/tests/fixtures/mergers.py b/tests/fixtures/mergers.py deleted file mode 100644 index 41c2150..0000000 --- a/tests/fixtures/mergers.py +++ /dev/null @@ -1,155 +0,0 @@ -MERGER_APPEND_CONFIG = { - "merger_id": "merger_append_example", - "type": "merger_append", - "items": [ - { - "subfeed_id": "subfeed_merger_append_example", - "type": "subfeed", - "method_name": "ads", - "subfeed_params": { - "limit_to_return": 5, - }, - }, - { - "subfeed_id": "subfeed_2_merger_append_example", - "type": "subfeed", - "method_name": "followings", - }, - ], -} - -MERGER_APPEND_DISTRIBUTE_CONFIG = { - "merger_id": "merger_distribute_example", - "type": "merger_distribute", - "distribution_key": "user_id", - "items": [ - { - "subfeed_id": "subfeed_merger_distribute_example", - "type": "subfeed", - "method_name": "posted", - "subfeed_params": { - "limit_to_return": 25, - }, - }, - { - "subfeed_id": "subfeed_2_merger_distribute_example", - "type": "subfeed", - "method_name": "posted", - "subfeed_params": { - "limit_to_return": 45, - }, - }, - ], -} - -MERGER_PERCENTAGE_CONFIG = { - "merger_id": "merger_percentage_example", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 40, - "data": { - "subfeed_id": "subfeed_merger_percentage_example", - "type": "subfeed", - "method_name": "followings", - }, - }, - { - "percentage": 60, - "data": { - "subfeed_id": "subfeed_2_merger_percentage_example", - "type": "subfeed", - "method_name": "ads", - }, - }, - ], -} - -MERGER_POSITIONAL_CONFIG = { - "merger_id": "merger_positional_example", - "type": "merger_positional", - "positions": [1, 3, 15], - "start": 17, - "end": 200, - "step": 2, - "positional": { - "subfeed_id": "subfeed_positional_merger_positional_example", - "type": "subfeed", - "method_name": "ads", - "subfeed_params": { - "limit_to_return": 10, - }, - }, - "default": { - "subfeed_id": "subfeed_default_merger_positional_example", - "type": "subfeed", - "method_name": "followings", - }, -} - -MERGER_PERCENTAGE_GRADIENT_CONFIG = { - "merger_id": "merger_percentage_gradient_example", - "type": "merger_percentage_gradient", - "item_from": { - "percentage": 80, - "data": { - "subfeed_id": "subfeed_from_merger_percentage_gradient_example", - "type": "subfeed", - "method_name": "ads", - }, - }, - "item_to": { - "percentage": 20, - "data": { - "subfeed_id": "subfeed_to_merger_percentage_gradient_example", - "type": "subfeed", - "method_name": "followings", - }, - }, - "step": 10, - "size_to_step": 10, - "shuffle": False, -} - -MERGER_VIEW_SESSION_CONFIG = { - "merger_id": "merger_view_session_example", - "type": "merger_view_session", - "session_size": 800, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_merger_view_session_example", - "type": "subfeed", - "method_name": "followings", - }, -} - -MERGER_VIEW_SESSION_DUPS_CONFIG = { - "merger_id": "merger_view_session_example", - "type": "merger_view_session", - "deduplicate": True, - "session_size": 10, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_merger_view_session_example", - "type": "subfeed", - "method_name": "doubles", - }, -} - -MERGER_DEDUP_VIEW_SESSION_CONFIG = { - "merger_id": "dedup_over_session", - "type": "merger_deduplication", - "dedup_key": None, - "data": { - "merger_id": "inner_session", - "type": "merger_view_session", - "session_size": 60, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_dedup_vs", - "type": "subfeed", - "method_name": "followings", - }, - }, -} diff --git a/tests/fixtures/redis.py b/tests/fixtures/redis.py deleted file mode 100644 index 5c9af72..0000000 --- a/tests/fixtures/redis.py +++ /dev/null @@ -1,30 +0,0 @@ -import pytest -import pytest_asyncio -import redis -from redis.asyncio import Redis as AsyncRedis - - -@pytest_asyncio.fixture(scope="function") -async def redis_client(request): - """Provide a Redis client for tests. - - If Redis is not available on localhost:6379, skip tests that depend on it. - """ - - if request.param == "async": - client = AsyncRedis(host="localhost", port=6379) - try: - await client.ping() - except Exception: - pytest.skip("Redis is not available on localhost:6379") - yield client - await client.aclose() - return - - client = redis.Redis(host="localhost", port=6379, db=0) - try: - client.ping() - except Exception: - pytest.skip("Redis is not available on localhost:6379") - yield client - client.close() diff --git a/tests/fixtures/subfeeds.py b/tests/fixtures/subfeeds.py deleted file mode 100644 index 78ee89c..0000000 --- a/tests/fixtures/subfeeds.py +++ /dev/null @@ -1,27 +0,0 @@ -SUBFEED_CONFIG = { - "subfeed_id": "subfeed_example", - "type": "subfeed", - "method_name": "ads", -} - -SUBFEED_CONFIG_RAISE_ERROR = { - "subfeed_id": "subfeed_example", - "type": "subfeed", - "method_name": "error", -} - -SUBFEED_CONFIG_NO_RAISE_ERROR = { - "subfeed_id": "subfeed_example", - "type": "subfeed", - "method_name": "error", - "raise_error": False, -} - -SUBFEED_WITH_PARAMS_CONFIG = { - "subfeed_id": "subfeed_with_params_example", - "type": "subfeed", - "method_name": "ads", - "subfeed_params": { - "limit_to_return": 10, - }, -} diff --git a/tests/test_async_concurrency.py b/tests/test_async_concurrency.py new file mode 100644 index 0000000..65e2cf4 --- /dev/null +++ b/tests/test_async_concurrency.py @@ -0,0 +1,283 @@ +"""Test that subfeeds in parallel branches run concurrently via asyncio.gather, +and that the event loop stays responsive under load.""" + +import asyncio +import time +from dataclasses import dataclass, field +from typing import Any, List, Optional + +import pytest + +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperDedup +from smartfeed.models.mixers import MergerAppend, MergerPercentage, MergerPercentageItem, MergerPositional +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +# --------------------------------------------------------------------------- +# Helpers: concurrency tracker + loop block monitor +# --------------------------------------------------------------------------- + +@dataclass +class ConcurrencyTracker: + """Tracks how many leaf calls are in-flight at the same time.""" + current: int = 0 + peak: int = 0 + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + async def enter(self) -> None: + async with self._lock: + self.current += 1 + self.peak = max(self.peak, self.current) + + async def exit(self) -> None: + async with self._lock: + self.current -= 1 + + +class LoopBlockMonitor: + """Detects event-loop blocking by measuring scheduling lag.""" + + def __init__(self, interval: float = 0.01, threshold: float = 0.1): + self.interval = interval + self.threshold = threshold + self.max_lag: float = 0.0 + self._task: Optional[asyncio.Task] = None + self._stop = asyncio.Event() + + async def __aenter__(self): + self._stop.clear() + self._task = asyncio.create_task(self._run()) + return self + + async def __aexit__(self, *exc): + self._stop.set() + if self._task: + await self._task + + async def _run(self): + loop = asyncio.get_running_loop() + expected = loop.time() + self.interval + while not self._stop.is_set(): + await asyncio.sleep(self.interval) + now = loop.time() + lag = max(0.0, now - expected) + self.max_lag = max(self.max_lag, lag) + expected = now + self.interval + + +# --------------------------------------------------------------------------- +# Slow subfeed factory +# --------------------------------------------------------------------------- + +def make_slow_subfeed_method( + source: str, + latency: float, + tracker: ConcurrencyTracker, + call_counter: dict, +): + """Create an async subfeed method that sleeps (simulating I/O) and tracks concurrency.""" + + async def method(user_id: str, limit: int, next_page: dict, **kw) -> FeedResult: + call_counter[source] = call_counter.get(source, 0) + 1 + page = next_page.get("page", 1) + start = (page - 1) * limit + + await tracker.enter() + try: + await asyncio.sleep(latency) + finally: + await tracker.exit() + + # Each source gets unique id ranges to avoid dedup collisions + offset = hash(source) % 100_000 + data = [{"id": offset + start + i, "val": source} for i in range(limit)] + return FeedResult( + data=data, + next_page={"page": page + 1}, + has_next_page=True, + ) + + return method + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestParallelSubfeeds: + """Verify that sibling subfeeds in a mixer run concurrently.""" + + @pytest.mark.asyncio + async def test_two_subfeeds_run_in_parallel(self): + tracker = ConcurrencyTracker() + calls = {} + latency = 0.05 + + methods = { + "slow_a": make_slow_subfeed_method("a", latency, tracker, calls), + "slow_b": make_slow_subfeed_method("b", latency, tracker, calls), + } + ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + executor = None + + node = MergerAppend( + node_id="par", + items=[ + SubFeed(subfeed_id="a", method_name="slow_a"), + SubFeed(subfeed_id="b", method_name="slow_b"), + ], + ) + + t0 = time.monotonic() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + elapsed = time.monotonic() - t0 + + assert len(result.data) == 10 + # Both called + assert calls.get("a", 0) >= 1 + assert calls.get("b", 0) >= 1 + # Peak concurrency > 1 means they ran in parallel + assert tracker.peak > 1, f"Expected parallel execution, peak={tracker.peak}" + # Total time should be ~latency (parallel), not ~2*latency (serial) + assert elapsed < latency * 1.8, f"Took {elapsed:.3f}s, expected < {latency * 1.8:.3f}s" + + @pytest.mark.asyncio + async def test_three_percentage_subfeeds_parallel(self): + tracker = ConcurrencyTracker() + calls = {} + latency = 0.05 + + methods = { + "slow_a": make_slow_subfeed_method("a", latency, tracker, calls), + "slow_b": make_slow_subfeed_method("b", latency, tracker, calls), + "slow_c": make_slow_subfeed_method("c", latency, tracker, calls), + } + ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + executor = None + + node = MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem(percentage=40, data=SubFeed(subfeed_id="a", method_name="slow_a")), + MergerPercentageItem(percentage=30, data=SubFeed(subfeed_id="b", method_name="slow_b")), + MergerPercentageItem(percentage=30, data=SubFeed(subfeed_id="c", method_name="slow_c")), + ], + ) + + t0 = time.monotonic() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + elapsed = time.monotonic() - t0 + + assert len(result.data) == 10 + assert tracker.peak >= 2, f"Expected parallel, peak={tracker.peak}" + assert elapsed < latency * 2.0 + + @pytest.mark.asyncio + async def test_positional_branches_parallel(self): + """Positional's two branches (positional + default) should run in parallel.""" + tracker = ConcurrencyTracker() + calls = {} + latency = 0.05 + + methods = { + "slow_promo": make_slow_subfeed_method("promo", latency, tracker, calls), + "slow_regular": make_slow_subfeed_method("regular", latency, tracker, calls), + } + ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + executor = None + + node = MergerPositional( + node_id="pos", + positions=[1, 3, 5], + positional=SubFeed(subfeed_id="promo", method_name="slow_promo"), + default=SubFeed(subfeed_id="regular", method_name="slow_regular"), + ) + + t0 = time.monotonic() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + elapsed = time.monotonic() - t0 + + assert len(result.data) == 10 + assert tracker.peak > 1, f"Expected parallel, peak={tracker.peak}" + assert elapsed < latency * 1.8 + + +class TestEventLoopResponsiveness: + """Verify the event loop stays unblocked during feed execution.""" + + @pytest.mark.asyncio + async def test_loop_not_blocked_under_load(self): + tracker = ConcurrencyTracker() + calls = {} + latency = 0.02 + + methods = {} + subfeeds = [] + for i in range(5): + name = f"source_{i}" + methods[name] = make_slow_subfeed_method(name, latency, tracker, calls) + subfeeds.append(SubFeed(subfeed_id=name, method_name=name)) + + ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + executor = None + + node = MergerAppend(node_id="big", items=subfeeds) + + async with LoopBlockMonitor(interval=0.01, threshold=0.05) as monitor: + result = await run_executor.run(node, ctx, limit=20, cursor={}) + + assert len(result.data) == 20 + assert monitor.max_lag < 0.1, f"Event loop blocked: max_lag={monitor.max_lag:.3f}s" + + +class TestNestedParallelism: + """Verify that nested mixer trees still execute leaves in parallel.""" + + @pytest.mark.asyncio + async def test_nested_mixers_parallel(self): + """append(percentage(a, b), positional(c, d)) -- all 4 subfeeds should overlap.""" + tracker = ConcurrencyTracker() + calls = {} + latency = 0.05 + + methods = { + f"slow_{x}": make_slow_subfeed_method(x, latency, tracker, calls) + for x in ("a", "b", "c", "d") + } + ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + executor = None + + node = MergerAppend( + node_id="outer", + items=[ + MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem(percentage=50, data=SubFeed(subfeed_id="a", method_name="slow_a")), + MergerPercentageItem(percentage=50, data=SubFeed(subfeed_id="b", method_name="slow_b")), + ], + ), + MergerPositional( + node_id="pos", + positions=[1, 3], + positional=SubFeed(subfeed_id="c", method_name="slow_c"), + default=SubFeed(subfeed_id="d", method_name="slow_d"), + ), + ], + ) + + t0 = time.monotonic() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + elapsed = time.monotonic() - t0 + + assert len(result.data) == 10 + # All 4 sources should have been called + for x in ("a", "b", "c", "d"): + assert calls.get(x, 0) >= 1 + # At least 2 in-flight at the same time + assert tracker.peak >= 2, f"Expected parallel, peak={tracker.peak}" + # Should take ~1x latency, not ~4x + assert elapsed < latency * 2.5, f"Took {elapsed:.3f}s, suggests serial execution" diff --git a/tests/test_async_loop_blocks_trace.py b/tests/test_async_loop_blocks_trace.py deleted file mode 100644 index a1bafd9..0000000 --- a/tests/test_async_loop_blocks_trace.py +++ /dev/null @@ -1,467 +0,0 @@ -import asyncio -import json -import os -import time -from dataclasses import dataclass, field -from typing import Any, Awaitable, Callable, Dict, List, Optional - -import pytest - -from smartfeed.schemas import FeedResultNextPage, MergerDeduplication -from tests.fixtures import dedup_helpers as dh -from tests.fixtures.redis import redis_client # noqa: F401 -from tests.utils import parse_model - - -def _now_us() -> int: - return time.perf_counter_ns() // 1000 - - -@dataclass -class ChromeTraceRecorder: - """Writes Chrome Trace Events JSON for chrome://tracing. - - This is intentionally tiny and test-only: no production dependencies. - """ - - pid: int = 1 - events: List[Dict[str, Any]] = field(default_factory=list) - - def _emit(self, event: Dict[str, Any]) -> None: - self.events.append(event) - - def begin(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: - self._emit( - { - "name": name, - "ph": "B", - "ts": int(_now_us() if ts_us is None else ts_us), - "pid": int(self.pid), - "tid": int(tid), - "args": args or {}, - } - ) - - def end(self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None) -> None: - self._emit( - { - "name": name, - "ph": "E", - "ts": int(_now_us() if ts_us is None else ts_us), - "pid": int(self.pid), - "tid": int(tid), - "args": args or {}, - } - ) - - def instant( - self, name: str, *, tid: int, ts_us: Optional[int] = None, args: Optional[Dict[str, Any]] = None - ) -> None: - self._emit( - { - "name": name, - "ph": "i", - "s": "t", - "ts": int(_now_us() if ts_us is None else ts_us), - "pid": int(self.pid), - "tid": int(tid), - "args": args or {}, - } - ) - - def write(self, path: str) -> None: - payload = {"traceEvents": self.events} - with open(path, "w", encoding="utf-8") as f: - json.dump(payload, f) - - -class LoopBlockMonitor: - """Detects event-loop blocking by measuring scheduling lag. - - If the event loop is blocked by long sync work, a periodic sleeper will wake - up late; we track the maximum observed lag. - """ - - def __init__(self, *, sample_interval_s: float = 0.01, block_threshold_s: float = 0.25) -> None: - self.sample_interval_s = float(sample_interval_s) - self.block_threshold_s = float(block_threshold_s) - self.max_lag_s: float = 0.0 - self.block_events: List[float] = [] - self._task: Optional[asyncio.Task[None]] = None - self._stop = asyncio.Event() - - async def __aenter__(self) -> "LoopBlockMonitor": - self._stop.clear() - self._task = asyncio.create_task(self._run()) - return self - - async def __aexit__(self, exc_type, exc, tb) -> None: # type: ignore[override] - self._stop.set() - if self._task is not None: - await self._task - - async def _run(self) -> None: - loop = asyncio.get_running_loop() - expected = loop.time() + self.sample_interval_s - while not self._stop.is_set(): - await asyncio.sleep(self.sample_interval_s) - now = loop.time() - lag = max(0.0, now - expected) - expected = now + self.sample_interval_s - self.max_lag_s = max(self.max_lag_s, lag) - if lag >= self.block_threshold_s: - self.block_events.append(lag) - - -@dataclass -class LeafConcurrencyTracker: - """Tracks how many leaf calls are in-flight concurrently.""" - - current: int = 0 - peak: int = 0 - _lock: asyncio.Lock = field(default_factory=asyncio.Lock) - - async def enter(self) -> int: - async with self._lock: - self.current += 1 - if self.current > self.peak: - self.peak = self.current - return self.current - - async def exit(self) -> int: - async with self._lock: - self.current = max(0, self.current - 1) - return self.current - - -def _trace_wrap_awaitable( - rec: ChromeTraceRecorder, name: str, awaitable: Awaitable[Any], *, args: Dict[str, Any] -) -> Awaitable[Any]: - async def _wrapped() -> Any: - task = asyncio.current_task() - tid = id(task) if task is not None else 0 - rec.begin(name, tid=tid, args=args) - try: - return await awaitable - finally: - rec.end(name, tid=tid) - - return _wrapped() - - -def _wrap_method_latency(method: Callable[..., Awaitable[Any]], *, latency_s: float) -> Callable[..., Awaitable[Any]]: - async def _wrapped(*args: Any, **kwargs: Any) -> Any: - await asyncio.sleep(latency_s) - return await method(*args, **kwargs) - - return _wrapped - - -def _wrap_leaf_method_traced( - *, - rec: ChromeTraceRecorder, - key: str, - method: Callable[..., Awaitable[Any]], - latency_s: float, - concurrency: LeafConcurrencyTracker, -) -> Callable[..., Awaitable[Any]]: - async def _wrapped(user_id: Any, limit: int, next_page: Any, **kwargs: Any) -> Any: - task = asyncio.current_task() - tid = id(task) if task is not None else 0 - - page = getattr(next_page, "page", None) - after = getattr(next_page, "after", None) - after_type = type(after).__name__ - - if after is None: - after_preview = None - else: - after_preview = str(after) - if len(after_preview) > 120: - after_preview = after_preview[:117] + "..." - - span = f"leaf.{key}" - - in_flight = await concurrency.enter() - rec.begin( - span, - tid=tid, - args={ - "key": key, - "limit": int(limit), - "page": page, - "after_type": after_type, - "after_preview": after_preview, - "in_flight": int(in_flight), - }, - ) - try: - if latency_s > 0: - await asyncio.sleep(float(latency_s)) - return await method(user_id, limit, next_page, **kwargs) - finally: - rec.end(span, tid=tid) - await concurrency.exit() - - return _wrapped - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_async_loop_blocks_and_trace_for_deep_tree_all_mergers(redis_client, monkeypatch, tmp_path) -> None: - """A smoke-test for detecting async loop blocks + visualizing concurrency. - - - Builds one deep tree that includes ALL merger types. - - Simulates 2 sequential requests (fresh + next page). - - Forces refills via positional under-fetch (`max_per_call=1`). - - Records loop scheduling lag (blocks/hangs) and optionally exports a Chrome trace. - - Set `SMARTFEED_CHROME_TRACE=/path/to/trace.json` to write a trace. - Open it in Chrome via chrome://tracing. - """ - - # Keep IDs disjoint across sources so "no dupes" is stable. - # Refill waves are forced via max_per_call limits (under-fetch), not via dedup collisions. - items_a = dh.make_items("A", 1, 400, user_id_mod=5, id_offset=1_000) - items_b = dh.make_items("B", 1, 400, user_id_mod=5, id_offset=10_000) - - # Distribute branch: needs distribution_key present (user_id). - items_posted_1 = dh.make_items("posted_1", 1, 80, user_id_mod=3, id_offset=20_000) - items_posted_2 = dh.make_items("posted_2", 1, 120, user_id_mod=3, id_offset=21_000) - - # Gradient branch: overlapping ids again. - items_g1 = dh.make_items("G1", 1, 250, user_id_mod=7, id_offset=30_000) - items_g2 = dh.make_items("G2", 1, 250, user_id_mod=7, id_offset=40_000) - - # View-session leaf. - items_vs = dh.make_items("VS", 1, 160, user_id_mod=11, id_offset=50_000) - - # Positional leaf that intentionally under-fetches to force refill waves. - items_pos_leaf = dh.make_items("POS", 1, 500, user_id_mod=13, id_offset=60_000) - - # --- tracing (test-only monkeypatch) --- - rec = ChromeTraceRecorder() - leaf_concurrency = LeafConcurrencyTracker() - pos_leaf_calls = {"count": 0} - - pos_leaf_base = dh.make_offset_paged_method(items_pos_leaf, max_per_call=1) - - async def _pos_leaf_counted(user_id: Any, limit: int, next_page: Any, **kwargs: Any) -> Any: - pos_leaf_calls["count"] += 1 - return await pos_leaf_base(user_id, limit, next_page, **kwargs) - - # Leaf method tracing: wrap the *actual* subfeed method calls. - # These spans are what you want to inspect for "are leaf calls parallel?". - leaf_latency_s = 0.02 - methods_dict = { - "a": _wrap_leaf_method_traced( - rec=rec, - key="a", - method=dh.make_offset_paged_method(items_a), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "b": _wrap_leaf_method_traced( - rec=rec, - key="b", - method=dh.make_offset_paged_method(items_b), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "posted_1": _wrap_leaf_method_traced( - rec=rec, - key="posted_1", - method=dh.make_offset_paged_method(items_posted_1), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "posted_2": _wrap_leaf_method_traced( - rec=rec, - key="posted_2", - method=dh.make_offset_paged_method(items_posted_2), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "g1": _wrap_leaf_method_traced( - rec=rec, - key="g1", - method=dh.make_offset_paged_method(items_g1), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "g2": _wrap_leaf_method_traced( - rec=rec, - key="g2", - method=dh.make_offset_paged_method(items_g2), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - "vs": _wrap_leaf_method_traced( - rec=rec, - key="vs", - method=dh.make_offset_paged_method(items_vs), - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - # Fetch only 1 item per call even if demand is higher -> triggers refill loops. - "pos_leaf": _wrap_leaf_method_traced( - rec=rec, - key="pos_leaf", - method=_pos_leaf_counted, - latency_s=leaf_latency_s, - concurrency=leaf_concurrency, - ), - } - - view_session_cfg = { - "merger_id": "vs_all", - "type": "merger_view_session", - "session_size": 100, - "session_live_time": 60, - "deduplicate": True, - "dedup_key": "id", - "data": dh._subfeed("sf_vs", "vs"), - } - - pct_cfg = dh._percentage_config( - "pct_all", - items=dh._percentage_items(dh._subfeed("sf_a", "a"), dh._subfeed("sf_b", "b"), first_pct=50, second_pct=50), - ) - - pos_cfg = dh._positional_config( - "pos_all", - # Ensure positional inserts appear across pages for limit~12. - # Use even positions so the schedule starts with the default branch; - # this keeps ordering deterministic. - positions=[2, 4, 6, 8, 10, 12, 14, 16, 18], - positional=dh._subfeed("sf_pos_leaf", "pos_leaf"), - default=pct_cfg, - ) - - dist_cfg = dh._distribute_config( - "dist_all", - items=[dh._subfeed("sf_posted_1", "posted_1"), dh._subfeed("sf_posted_2", "posted_2")], - distribution_key="user_id", - ) - - grad_cfg = dh._gradient_config( - "grad_all", - item_from={"percentage": 70, "data": dh._subfeed("sf_g1", "g1")}, - item_to={"percentage": 30, "data": dh._subfeed("sf_g2", "g2")}, - step=10, - size_to_step=5, - shuffle=False, - ) - - # Include all merger types as siblings so they are executed (and visible in trace), - # while keeping the main output driven by the first branch. - deep_tree = dh._append_config("append_all", [pos_cfg, view_session_cfg, dist_cfg, grad_cfg]) - config = dh._dedup_config( - "dedup_all", - deep_tree, - dedup_key="id", - state_backend="cursor", - overfetch_factor=3, - max_refill_loops=50, - ) - merger = parse_model(MergerDeduplication, config) - - # Patch Executor.gather to wrap each awaitable for Chrome trace. - from smartfeed.execution.executor import Executor # local import for monkeypatch - - original_gather = Executor.gather - - async def _gather_traced(self: Any, *coros: Any) -> List[Any]: - wrapped = [ - _trace_wrap_awaitable(rec, "executor.gather.op", c, args={"idx": i, "total": len(coros)}) - for i, c in enumerate(coros) - ] - task = asyncio.current_task() - tid = id(task) if task is not None else 0 - rec.begin("executor.gather", tid=tid, args={"n": len(coros)}) - try: - return await original_gather(self, *wrapped) - finally: - rec.end("executor.gather", tid=tid) - - monkeypatch.setattr(Executor, "gather", _gather_traced) - - # Patch Executor.run to show sequential refill loops vs plan execution. - original_run = Executor.run - - async def _run_traced( - self: Any, - node: Any, - ctx: Any, - limit: int, - next_page: Any, - **params: Any, - ) -> Any: - task = asyncio.current_task() - tid = id(task) if task is not None else 0 - node_type = getattr(node, "type", node.__class__.__name__) - node_id = getattr(node, "merger_id", getattr(node, "subfeed_id", None)) - rec.begin( - "executor.run_node", - tid=tid, - args={"node_type": node_type, "node_id": node_id, "limit": int(limit)}, - ) - try: - return await original_run(self, node, ctx, limit, next_page, **params) - finally: - rec.end("executor.run_node", tid=tid) - - monkeypatch.setattr(Executor, "run", _run_traced) - - # --- run: fresh request + next_page --- - limit = 12 - np0 = FeedResultNextPage(data={}) - - async with LoopBlockMonitor(sample_interval_s=0.01, block_threshold_s=0.05) as monitor: - res1 = await asyncio.wait_for( - merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=limit, - next_page=np0, - redis_client=redis_client, - ), - timeout=15, - ) - res2 = await asyncio.wait_for( - merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=limit, - next_page=res1.next_page, - redis_client=redis_client, - ), - timeout=15, - ) - - # Sanity: we should fill the page and maintain dedup invariants. - assert len(res1.data) == limit - assert len({x["id"] for x in res1.data}) == limit - assert len(res2.data) == limit - assert len({x["id"] for x in res2.data}) == limit - - # Hard assertion: leaf calls must overlap (async concurrency), not serialize. - assert leaf_concurrency.peak > 1 - # Refill signal: with max_per_call=1, two page requests should trigger - # multiple extra positional calls to satisfy positional slots. - assert pos_leaf_calls["count"] > 2 - - # Primary signal: event-loop should remain responsive under load. - assert monitor.max_lag_s < 0.1 - - out = os.environ.get("SMARTFEED_CHROME_TRACE") - if out: - # Allow writing to an explicit file path, or to a directory. - out_path = out - if os.path.isdir(out_path): - out_path = os.path.join(out_path, "smartfeed_trace.json") - rec.instant("loop.max_lag", tid=0, args={"max_lag_s": monitor.max_lag_s, "blocks": len(monitor.block_events)}) - rec.write(out_path) - - # Keep references so this test remains useful in local debugging. - _ = tmp_path diff --git a/tests/test_config_parsing.py b/tests/test_config_parsing.py new file mode 100644 index 0000000..3ebe991 --- /dev/null +++ b/tests/test_config_parsing.py @@ -0,0 +1,60 @@ +import pytest +from smartfeed.models.base import BaseNode +from smartfeed.models import FeedConfig + + +class DummyNode(BaseNode): + type: str = "dummy" + node_id: str = "test" + params: dict = {} + + +def test_config_hash_deterministic(): + a = DummyNode(params={"b": 2, "a": 1}) + b = DummyNode(params={"a": 1, "b": 2}) + assert a.config_hash() == b.config_hash() + assert len(a.config_hash()) == 8 + + +def test_config_hash_changes_with_params(): + a = DummyNode(params={"x": 1}) + b = DummyNode(params={"x": 2}) + assert a.config_hash() != b.config_hash() + + +def test_parse_simple_subfeed(): + cfg = FeedConfig.parse_obj({ + "version": "2", + "feed": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, + }) + assert cfg.feed.subfeed_id == "x" + + +def test_parse_wrapper_with_cache(): + cfg = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "w", + "cache": {"session_size": 100, "session_ttl": 300}, + "data": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, + }, + }) + assert cfg.feed.cache.session_size == 100 + + +def test_parse_nested_config(): + cfg = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "outer", + "dedup": {"dedup_key": "id"}, + "data": { + "type": "merger_percentage", "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "items"}}, + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "items"}}, + ], + }, + }, + }) + assert cfg.feed.dedup.dedup_key == "id" diff --git a/tests/test_cursor_and_refill_edges.py b/tests/test_cursor_and_refill_edges.py deleted file mode 100644 index 27b5b72..0000000 --- a/tests/test_cursor_and_refill_edges.py +++ /dev/null @@ -1,92 +0,0 @@ -import base64 -import zlib - -import pytest - -from smartfeed.execution.context import ExecutionContext, RefillExecutionSettings -from smartfeed.execution.executor import Executor -from smartfeed.feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage -from smartfeed.policies.dedup import DeduplicationPolicy -from smartfeed.policies.dedup_utils import decode_seen_from_cursor -from smartfeed.policies.seen_store import CursorSeenStore - - -class _DuplicateOnlyNode(BaseFeedConfigModel): - """A node that always returns the same duplicate item and never ends.""" - - type: str = "test_node" # satisfies pydantic - - def __init__(self, **data): - super().__init__(**data) - object.__setattr__(self, "calls", 0) - - async def get_data( # type: ignore[override] - self, - methods_dict, - user_id, - limit, - next_page, - redis_client=None, - ctx=None, - **params, - ) -> FeedResult: - object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) - return FeedResult( - data=[{"id": "dup"} for _ in range(int(limit) or 1)], - next_page=FeedResultNextPage(data={}), - has_next_page=True, - ) - - -def _policy_with_seen_dup(*, max_refill_loops: int) -> tuple[DeduplicationPolicy, RefillExecutionSettings]: - store = CursorSeenStore.from_after( - after={"v": 2, "seen": [["dup", 10]]}, - cursor_compress=False, - cursor_max_keys=None, - ) - policy = DeduplicationPolicy( - dedup_key="id", - missing_key_policy="keep", - store=store, - seen_request_set=set(), - ) - settings = RefillExecutionSettings(overfetch_factor=1, max_refill_loops=max_refill_loops) - return policy, settings - - -@pytest.mark.asyncio -async def test_dedup_refill_stops_at_max_loops_when_only_duplicates() -> None: - node = _DuplicateOnlyNode() - executor = Executor() - - dedup_policy, settings = _policy_with_seen_dup(max_refill_loops=2) - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - ctx.dedup = dedup_policy - ctx.refill_settings = settings - - res = await executor.run(node, ctx, limit=3, next_page=FeedResultNextPage(data={})) - - assert res.data == [] - # initial call + 2 refill loops - assert getattr(node, "calls") == 3 - - -def test_decode_seen_from_cursor_raises_on_corrupt_compressed_payload() -> None: - # invalid base64 - with pytest.raises(Exception): - decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": "not-base64"}) - - # base64 ok, zlib invalid - bad_zlib = base64.urlsafe_b64encode(b"not-a-zlib-stream").decode("ascii") - with pytest.raises(Exception): - decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": bad_zlib}) - - # zlib ok, json invalid - bad_json = base64.urlsafe_b64encode(zlib.compress(b"not json")).decode("ascii") - with pytest.raises(Exception): - decode_seen_from_cursor({"v": 2, "c": "zlib+base64", "z": bad_json}) - - -def test_decode_seen_from_cursor_rejects_wrong_version_or_codec() -> None: - assert decode_seen_from_cursor({"v": 1, "c": "zlib+base64", "z": ""}) == {} - assert decode_seen_from_cursor({"v": 2, "c": "other", "z": ""}) == {} diff --git a/tests/test_cursors.py b/tests/test_cursors.py new file mode 100644 index 0000000..0e3bb42 --- /dev/null +++ b/tests/test_cursors.py @@ -0,0 +1,54 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models import FeedConfig +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=METHODS, redis=redis) + + +@pytest.mark.asyncio +async def test_cached_wrapper_hides_child_cursors(ctx): + config = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "cached", + "cache": {"session_size": 50, "session_ttl": 300}, + "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, + }, + }) + + executor = None + result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) + + # Only wrapper cursor, NOT subfeed cursor + assert "cached" in result.next_page + assert "items" not in result.next_page + assert "gen" in result.next_page["cached"] + + +@pytest.mark.asyncio +async def test_uncached_wrapper_exposes_child_cursors(ctx): + config = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "passthrough", + "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, + }, + }) + + executor = None + result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) + + # Child cursor is visible (no cache to hide it) + assert "items" in result.next_page diff --git a/tests/test_dedup_policy_unit.py b/tests/test_dedup_policy_unit.py deleted file mode 100644 index b846f6b..0000000 --- a/tests/test_dedup_policy_unit.py +++ /dev/null @@ -1,74 +0,0 @@ -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -import pytest - -from smartfeed.policies.dedup import DeduplicationPolicy -from smartfeed.policies.seen_store import CursorSeenStore - - -def _policy(*, dedup_key: Optional[str], missing_key_policy: str, after: Any = None) -> DeduplicationPolicy: - store = CursorSeenStore.from_after(after=after, cursor_compress=False, cursor_max_keys=None) - return DeduplicationPolicy( - dedup_key=dedup_key, - missing_key_policy=missing_key_policy, # type: ignore[arg-type] - store=store, - seen_request_set=set(), - ) - - -@pytest.mark.asyncio -async def test_accept_batch_dedups_within_request_and_respects_existing_priority() -> None: - # Existing seen key with high priority should block lower/equal priority - existing_after = {"v": 2, "seen": [["1", 10]]} - policy = _policy(dedup_key="id", missing_key_policy="keep", after=existing_after) - - items = [{"id": 1}, {"id": 1}, {"id": 2}] - accepted, inspected = await policy.accept_batch(items=items, priority=5) - - assert inspected == 3 - assert accepted == [{"id": 2}] - - -@pytest.mark.asyncio -async def test_accept_batch_missing_key_policies() -> None: - items = [{"id": 1}, {"nope": 2}] - - policy_drop = _policy(dedup_key="id", missing_key_policy="drop") - accepted_drop, _ = await policy_drop.accept_batch(items=items, priority=0) - assert accepted_drop == [{"id": 1}] - - policy_error = _policy(dedup_key="id", missing_key_policy="error") - with pytest.raises(AssertionError): - await policy_error.accept_batch(items=items, priority=0) - - -@dataclass -class _Owner: - dedup_priority: int - - -@pytest.mark.asyncio -async def test_arbitrate_owner_buffers_prefers_higher_priority_owner() -> None: - policy = _policy(dedup_key="id", missing_key_policy="keep") - - owner_low = _Owner(dedup_priority=1) - owner_high = _Owner(dedup_priority=2) - - owners = [owner_low, owner_high] - owner_rank: Dict[int, int] = {id(owner_low): 0, id(owner_high): 1} - - shared = {"id": "same"} - owner_buffers: Dict[int, List[Any]] = { - id(owner_low): [shared, {"id": "low_only"}], - id(owner_high): [shared, {"id": "high_only"}], - } - - per_owner = await policy.arbitrate_owner_buffers( - owners=owners, - owner_buffers=owner_buffers, - owner_rank=owner_rank, - ) - - assert per_owner[id(owner_high)] == [shared, {"id": "high_only"}] - assert per_owner[id(owner_low)] == [{"id": "low_only"}] diff --git a/tests/test_dedup_priority.py b/tests/test_dedup_priority.py new file mode 100644 index 0000000..f1be5a1 --- /dev/null +++ b/tests/test_dedup_priority.py @@ -0,0 +1,59 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models import FeedConfig +from smartfeed.models.base import FeedResult +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +async def promo_source(user_id, limit, next_page, **kw): + data = [{"id": i, "val": "promo"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def regular_source(user_id, limit, next_page, **kw): + data = [{"id": i, "val": "regular"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +PRIO_METHODS = {"promo": promo_source, "regular": regular_source} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=PRIO_METHODS, redis=redis) + + +@pytest.mark.asyncio +async def test_wrapper_dedup_priority_override(ctx): + """Inner wrapper with dedup_priority=3 overrides child subfeed priorities.""" + config = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "outer", + "dedup": {"dedup_key": "id"}, + "data": { + "type": "merger_append", "node_id": "top", + "items": [ + {"type": "subfeed", "subfeed_id": "promo", "method_name": "promo", "dedup_priority": 10}, + { + "type": "wrapper", "node_id": "inner", "dedup_priority": 3, + "data": {"type": "subfeed", "subfeed_id": "regular", "method_name": "regular"}, + }, + ], + }, + }, + }) + + executor = None + result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) + + # promo(10) > inner wrapper override(3) -> all promo + for item in result.data: + assert item["val"] == "promo" diff --git a/tests/test_dedup_utils.py b/tests/test_dedup_utils.py deleted file mode 100644 index 1805516..0000000 --- a/tests/test_dedup_utils.py +++ /dev/null @@ -1,73 +0,0 @@ -from typing import Any - -import pytest - -from smartfeed.feed_models import _redis_call -from smartfeed.policies.dedup_utils import decode_seen_from_cursor, encode_seen_for_cursor, redis_zmscore -from tests.fixtures.redis import redis_client # noqa: F401 - - -class _RedisNoZmscore: - """Wrapper around a real sync redis client to force the pipeline fallback. - - This keeps the backend "real Redis" while exercising the no-zmscore branch. - """ - - zmscore = None # type: ignore[assignment] - - def __init__(self, client: Any) -> None: - self._client = client - - def pipeline(self) -> Any: - return self._client.pipeline() - - -def test_encode_decode_seen_cursor_compressed_roundtrip_and_truncation() -> None: - seen_updates = [("a", 1), ("b", 2), ("c", 3)] - - encoded = encode_seen_for_cursor(seen_updates, cursor_compress=True, cursor_max_keys=None) - decoded = decode_seen_from_cursor(encoded) - assert decoded == {"a": 1, "b": 2, "c": 3} - - encoded_trunc = encode_seen_for_cursor(seen_updates, cursor_compress=True, cursor_max_keys=2) - decoded_trunc = decode_seen_from_cursor(encoded_trunc) - assert decoded_trunc == {"b": 2, "c": 3} - - -def test_decode_seen_cursor_v2_only() -> None: - assert decode_seen_from_cursor(None) == {} - - # Supported v2 uncompressed format - assert decode_seen_from_cursor({"v": 2, "seen": [["a", 9], ["b", 1]]}) == {"a": 9, "b": 1} - - # Legacy/unknown shapes are intentionally rejected - assert decode_seen_from_cursor(["x", "y"]) == {} - assert decode_seen_from_cursor({"a": 1}) == {} - assert decode_seen_from_cursor({"v": 1, "seen": [["a", 1]]}) == {} - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_redis_zmscore_native(redis_client) -> None: - key = "test_zmscore_native" - await _redis_call(redis_client, "delete", key) - await _redis_call(redis_client, "zadd", key, mapping={"a": 1.0, "b": 2.0}) - - res = await redis_zmscore(redis_client, key, ["a", "missing", "b"]) - assert res == [1.0, None, 2.0] - - await _redis_call(redis_client, "delete", key) - - -@pytest.mark.parametrize("redis_client", ["sync"], indirect=True) -@pytest.mark.asyncio -async def test_redis_zmscore_pipeline_fallback_for_sync_client_without_zmscore(redis_client) -> None: - key = "test_zmscore_fallback" - await _redis_call(redis_client, "delete", key) - await _redis_call(redis_client, "zadd", key, mapping={"a": 1.0, "b": 2.0}) - - wrapped = _RedisNoZmscore(redis_client) - res = await redis_zmscore(wrapped, key, ["a", "missing", "b"]) # type: ignore[arg-type] - assert res == [1.0, None, 2.0] - - await _redis_call(redis_client, "delete", key) diff --git a/tests/test_executor_slots_plan_invariants.py b/tests/test_executor_slots_plan_invariants.py deleted file mode 100644 index 596f10d..0000000 --- a/tests/test_executor_slots_plan_invariants.py +++ /dev/null @@ -1,297 +0,0 @@ -import pytest - -from smartfeed.execution.context import ExecutionContext, RefillExecutionSettings -from smartfeed.execution.executor import Executor -from smartfeed.execution.plans import SlotSpec, SlotsPlan -from smartfeed.feed_models import BaseFeedConfigModel, FeedResult, FeedResultNextPage, FeedResultNextPageInside -from smartfeed.policies.dedup import DeduplicationPolicy -from smartfeed.policies.seen_store import CursorSeenStore - - -class _Owner(BaseFeedConfigModel): - type: str = "test_owner" - name: str = "" - - def __init__(self, *, name: str, **data): # type: ignore[override] - super().__init__(name=name, **data) # type: ignore[call-arg] - object.__setattr__(self, "last_limit", None) - object.__setattr__(self, "calls", 0) - - async def get_data( # type: ignore[override] - self, - methods_dict, - user_id, - limit, - next_page, - redis_client=None, - ctx=None, - **params, - ) -> FeedResult: - object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) - object.__setattr__(self, "last_limit", int(limit)) - return FeedResult(data=[self.name] * int(limit), next_page=FeedResultNextPage(data={}), has_next_page=False) - - -class _PagedOwner(BaseFeedConfigModel): - type: str = "test_paged_owner" - subfeed_id: str - total: int = 10 - - def __init__(self, *, subfeed_id: str, total: int = 10, **data): # type: ignore[override] - super().__init__(subfeed_id=subfeed_id, total=total, **data) # type: ignore[call-arg] - object.__setattr__(self, "calls", 0) - object.__setattr__(self, "limits", []) - - async def get_data( # type: ignore[override] - self, - methods_dict, - user_id, - limit, - next_page, - redis_client=None, - ctx=None, - **params, - ) -> FeedResult: - object.__setattr__(self, "calls", int(getattr(self, "calls", 0)) + 1) - limits = list(getattr(self, "limits", [])) - limits.append(int(limit)) - object.__setattr__(self, "limits", limits) - - entry = next_page.data.get(self.subfeed_id) - offset = int(entry.after) if (entry is not None and isinstance(entry.after, int)) else 0 - - take = max(0, min(int(limit), int(self.total) - offset)) - data = [{"id": f"{self.subfeed_id}_{i}"} for i in range(offset + 1, offset + take + 1)] - new_after = offset + take - - next_page.data[self.subfeed_id] = FeedResultNextPageInside( - page=(entry.page + 1 if entry is not None else 2), - after=new_after, - ) - return FeedResult( - data=data, - next_page=next_page, - has_next_page=bool(new_after < int(self.total)), - ) - - -def _dedup_policy() -> DeduplicationPolicy: - store = CursorSeenStore.from_after(after=None, cursor_compress=False, cursor_max_keys=None) - return DeduplicationPolicy( - dedup_key="id", - missing_key_policy="keep", # type: ignore[arg-type] - store=store, - seen_request_set=set(), - ) - - -@pytest.mark.asyncio -async def test_slots_plan_limit_le_zero_calls_assemble_only() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - - owner = _Owner(name="x") - - called = {"assemble": 0} - - def assemble(output, next_page, owner_results): - called["assemble"] += 1 - assert output == [] - assert owner_results == {} - return FeedResult(data=output, next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=0, - next_page=FeedResultNextPage(data={}), - params={}, - slots=[SlotSpec(owner=owner, max_count=10)], - assemble=assemble, - ) - - res = await executor.execute_plan(plan) - assert res.data == [] - assert called["assemble"] == 1 - assert getattr(owner, "calls") == 0 - - -@pytest.mark.asyncio -async def test_slots_plan_owner_fetch_limits_overrides_demand() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - - owner = _Owner(name="x") - - def assemble(output, next_page, owner_results): - return FeedResult(data=output, next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=5, - next_page=FeedResultNextPage(data={}), - params={}, - slots=[SlotSpec(owner=owner, max_count=5)], - assemble=assemble, - owner_fetch_limits={id(owner): 1}, - ) - - res = await executor.execute_plan(plan) - assert getattr(owner, "calls") == 1 - assert getattr(owner, "last_limit") == 1 - assert res.data == ["x"] - - -@pytest.mark.asyncio -async def test_slots_plan_no_ops_path_assemble_still_runs() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - - owner = _Owner(name="x") - - called = {"assemble": 0} - - def assemble(output, next_page, owner_results): - called["assemble"] += 1 - assert output == [] - assert owner_results == {} - return FeedResult(data=[], next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=5, - next_page=FeedResultNextPage(data={}), - params={}, - slots=[SlotSpec(owner=owner, max_count=0)], - assemble=assemble, - ) - - res = await executor.execute_plan(plan) - assert res.data == [] - assert called["assemble"] == 1 - assert getattr(owner, "calls") == 0 - - -@pytest.mark.asyncio -async def test_slots_plan_quota_deficit_triggers_refill_wave() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - ctx.dedup = _dedup_policy() - - a = _PagedOwner(subfeed_id="a", total=10) - b = _PagedOwner(subfeed_id="b", total=10) - - def assemble(output, next_page, owner_results): - return FeedResult(data=output, next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=6, - next_page=FeedResultNextPage( - data={ - "a": FeedResultNextPageInside(page=1, after=0), - "b": FeedResultNextPageInside(page=1, after=0), - } - ), - params={}, - slots=[ - SlotSpec(owner=a, max_count=3), - SlotSpec(owner=b, max_count=3), - ], - assemble=assemble, - # Force an initial under-fetch for owner a (quota deficit). - owner_fetch_limits={id(a): 1}, - ) - - res = await executor.execute_plan(plan) - - # a should be refilled from 1 -> 3 items - assert getattr(a, "calls") >= 2 - assert res.data[:3] == [{"id": "a_1"}, {"id": "a_2"}, {"id": "a_3"}] - assert res.data[3:] == [{"id": "b_1"}, {"id": "b_2"}, {"id": "b_3"}] - - -@pytest.mark.asyncio -async def test_slots_plan_quota_deficit_stops_refill_when_owner_exhausts() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - ctx.dedup = _dedup_policy() - - # Owner a can never satisfy its full slot quota. - a = _PagedOwner(subfeed_id="a", total=2) - b = _PagedOwner(subfeed_id="b", total=10) - - def assemble(output, next_page, owner_results): - return FeedResult(data=output, next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=6, - next_page=FeedResultNextPage( - data={ - "a": FeedResultNextPageInside(page=1, after=0), - "b": FeedResultNextPageInside(page=1, after=0), - } - ), - params={}, - slots=[ - SlotSpec(owner=a, max_count=3), - SlotSpec(owner=b, max_count=3), - ], - assemble=assemble, - # Force an initial under-fetch to create a quota deficit for a. - owner_fetch_limits={id(a): 1}, - ) - - res = await executor.execute_plan(plan) - - # a is exhausted after returning 2 total items; refill should stop. - assert getattr(a, "calls") == 2 - assert getattr(a, "limits") == [1, 2] - assert getattr(b, "calls") == 1 - - assert res.data == [ - {"id": "a_1"}, - {"id": "a_2"}, - {"id": "b_1"}, - {"id": "b_2"}, - {"id": "b_3"}, - ] - - -@pytest.mark.asyncio -async def test_slots_plan_quota_deficit_refills_without_dedup_when_refill_settings_present() -> None: - executor = Executor() - ctx = ExecutionContext(methods_dict={}, user_id="u", executor=executor) - ctx.refill_settings = RefillExecutionSettings(overfetch_factor=3, max_refill_loops=10) - - a = _PagedOwner(subfeed_id="a", total=10) - b = _PagedOwner(subfeed_id="b", total=10) - - def assemble(output, next_page, owner_results): - return FeedResult(data=output, next_page=next_page, has_next_page=False) - - plan = SlotsPlan( - ctx=ctx, - limit=6, - next_page=FeedResultNextPage( - data={ - "a": FeedResultNextPageInside(page=1, after=0), - "b": FeedResultNextPageInside(page=1, after=0), - } - ), - params={}, - slots=[ - SlotSpec(owner=a, max_count=3), - SlotSpec(owner=b, max_count=3), - ], - # force an initial under-fetch for owner a. - owner_fetch_limits={id(a): 1}, - assemble=assemble, - ) - - res = await executor.execute_plan(plan) - - # refill must still happen even when dedup policy is absent. - assert getattr(a, "calls") >= 2 - assert res.data[:3] == [{"id": "a_1"}, {"id": "a_2"}, {"id": "a_3"}] - assert res.data[3:] == [{"id": "b_1"}, {"id": "b_2"}, {"id": "b_3"}] diff --git a/tests/test_high_priority.py b/tests/test_high_priority.py new file mode 100644 index 0000000..69e1781 --- /dev/null +++ b/tests/test_high_priority.py @@ -0,0 +1,550 @@ +"""High-priority missing tests for SmartFeed v2. + +Covers: +1. Dedup across 3+ pages (with and without cache) +2. Session rebuild continues pagination +3. has_next_page at session boundary +4. Missing dedup key policies (error / drop / keep) +5. Positional with exhausted high-priority subfeed +""" + +import pytest +import fakeredis.aioredis + +from smartfeed.models import ( + FeedResult, + MergerAppend, + MergerPositional, + SubFeed, + Wrapper, + WrapperCache, + WrapperDedup, +) +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +# --------------------------------------------------------------------------- +# Subfeed helpers +# --------------------------------------------------------------------------- + + +async def make_overlapping_a(user_id, limit, next_page, **kw): + """Produces ids 0, 2, 4, 6, ... (page-aware, even numbers).""" + page = next_page.get("page", 1) + start = (page - 1) * limit * 2 # ids 0,2,4,6... + data = [{"id": start + i * 2, "val": "a"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +async def make_overlapping_b(user_id, limit, next_page, **kw): + """Produces ids 1, 3, 5, 7, ... with some overlap with source a.""" + page = next_page.get("page", 1) + start = (page - 1) * limit * 2 + 1 # ids 1,3,5... + data = [{"id": start + i * 2, "val": "b"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +async def make_overlapping_ab(user_id, limit, next_page, **kw): + """Produces ids that heavily overlap with source a: 0, 2, 4, ...""" + page = next_page.get("page", 1) + start = (page - 1) * limit * 2 + data = [{"id": start + i * 2, "val": "ab_overlap"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +async def make_sequential(user_id, limit, next_page, **kw): + """Produces sequentially numbered items across pages (0,1,2,...) -- no gaps.""" + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": start + i, "val": f"seq_{start + i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +async def make_promo_limited(user_id, limit, next_page, **kw): + """Returns only 2 items total, then empty.""" + page = next_page.get("page", 1) + if page > 1: + return FeedResult(data=[], next_page={"page": page + 1}, has_next_page=False) + data = [{"id": 1000 + i, "val": "promo"} for i in range(2)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=False) + + +async def make_default_items(user_id, limit, next_page, **kw): + """Returns plenty of default items, page-aware.""" + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": start + i, "val": "default"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +OVERLAP_METHODS = { + "overlap_a": make_overlapping_a, + "overlap_b": make_overlapping_b, + "overlap_ab": make_overlapping_ab, +} + +SEQUENTIAL_METHODS = { + "sequential": make_sequential, +} + +POSITIONAL_METHODS = { + "promo_limited": make_promo_limited, + "default_items": make_default_items, +} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +# --------------------------------------------------------------------------- +# 1. Dedup across 3+ pages -- no duplicates +# --------------------------------------------------------------------------- + + +class TestDedupAcrossPages: + """Test dedup behaviour across multiple pages. + + Without cache, dedup is applied per-request (per-page) only because there + is no persisted seen-state between paginate calls. + + With cache, the entire session_size chunk is deduped at build time, so + dedup holds across all pages that come from that cached chunk. + """ + + @pytest.mark.asyncio + async def test_dedup_nocache_within_page_only(self): + """Without cache, dedup removes duplicates WITHIN a single page + but cannot prevent duplicates ACROSS pages because seen-state is not + persisted. + + We use two sources that both return the same fixed set of ids on every + call (ignoring pagination). Within a single page, dedup collapses + duplicates. Across pages, the same ids reappear because there is no + persisted seen-state. + """ + + async def make_fixed_a(user_id, limit, next_page, **kw): + """Always returns ids 0..limit-1 regardless of page.""" + data = [{"id": i, "val": "a"} for i in range(limit)] + return FeedResult( + data=data, next_page={"page": next_page.get("page", 1) + 1}, + has_next_page=True, + ) + + async def make_fixed_b(user_id, limit, next_page, **kw): + """Always returns ids 0..limit-1 (same as a) -- full overlap.""" + data = [{"id": i, "val": "b"} for i in range(limit)] + return FeedResult( + data=data, next_page={"page": next_page.get("page", 1) + 1}, + has_next_page=True, + ) + + methods = {"fixed_a": make_fixed_a, "fixed_b": make_fixed_b} + ctx = ExecutionContext( + session_id="dedup_nocache", + methods_dict=methods, + redis=None, + ) + node = Wrapper( + node_id="w", + dedup=WrapperDedup(dedup_key="id"), + data=MergerAppend( + node_id="mix", + items=[ + SubFeed(subfeed_id="a", method_name="fixed_a"), + SubFeed(subfeed_id="b", method_name="fixed_b"), + ], + ), + ) + + all_ids = [] + cursor = {} + for _ in range(3): + result = await run_executor.run(node, ctx, limit=10, cursor=cursor) + page_ids = [item["id"] for item in result.data] + # Within each page, no duplicates + assert len(page_ids) == len(set(page_ids)), "Duplicates found within a single page" + all_ids.extend(page_ids) + cursor = result.next_page + + # Across pages, duplicates ARE expected (seen-state not persisted). + # This is the known limitation without cache. + # We assert that cross-page duplicates exist to document the behaviour. + assert len(all_ids) > len(set(all_ids)), ( + "Expected cross-page duplicates without cache, but none found. " + "If the implementation now persists dedup state without cache, " + "this test can be updated." + ) + + @pytest.mark.asyncio + async def test_dedup_with_cache_across_pages(self, redis): + """With cache, dedup is applied to the entire session_size chunk at + cold-build time. Pages served from that chunk have zero duplicates.""" + ctx = ExecutionContext( + session_id="dedup_cached", + methods_dict=OVERLAP_METHODS, + redis=redis, + ) + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + dedup=WrapperDedup(dedup_key="id"), + data=MergerAppend( + node_id="mix", + items=[ + SubFeed(subfeed_id="a", method_name="overlap_a"), + SubFeed(subfeed_id="ab", method_name="overlap_ab"), + ], + ), + ) + + all_ids = [] + cursor = {} + for page_num in range(1, 4): + result = await run_executor.run(node, ctx, limit=10, cursor=cursor) + page_ids = [item["id"] for item in result.data] + # Within each page, no duplicates + assert len(page_ids) == len(set(page_ids)), f"Duplicates in page {page_num}" + all_ids.extend(page_ids) + cursor = result.next_page + + # With cache, dedup holds across ALL pages served from the same session + assert len(all_ids) == len(set(all_ids)), ( + "Duplicates found across cached pages -- dedup should cover the entire session" + ) + + +# --------------------------------------------------------------------------- +# 2. Session rebuild continues pagination +# --------------------------------------------------------------------------- + + +class TestSessionRebuildContinuation: + """Wrapper(cache, session_size=20), limit=10. + + Pages 1-2 come from the cache (20 items). Page 3 triggers a cache rebuild + via the continuation cursor stored in :meta. Items on page 3 should + continue where page 2 left off with no gap and no overlap. + """ + + @pytest.mark.asyncio + async def test_page3_continues_after_rebuild(self, redis): + ctx = ExecutionContext( + session_id="rebuild", + methods_dict=SEQUENTIAL_METHODS, + redis=redis, + ) + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=20, session_ttl=300), + data=SubFeed(subfeed_id="seq", method_name="sequential"), + ) + + # Page 1 + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r1.data) == 10 + ids_p1 = [item["id"] for item in r1.data] + + # Page 2 -- still from cache + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + assert len(r2.data) == 10 + ids_p2 = [item["id"] for item in r2.data] + + # Page 3 -- should trigger rebuild (cache exhausted: offset 20 >= 20) + r3 = await run_executor.run(node, ctx, limit=10, cursor=r2.next_page) + assert len(r3.data) == 10, "Page 3 should return items after rebuild" + ids_p3 = [item["id"] for item in r3.data] + + # No overlap between pages + all_ids = ids_p1 + ids_p2 + ids_p3 + assert len(all_ids) == len(set(all_ids)), ( + "Overlap detected between pages after session rebuild" + ) + + # No gap: the maximum id from page 2 + 1 should be the minimum id from page 3 + # (since make_sequential produces strictly sequential ids) + assert min(ids_p3) == max(ids_p2) + 1, ( + f"Gap detected: page 2 ended at {max(ids_p2)}, page 3 starts at {min(ids_p3)}" + ) + + +# --------------------------------------------------------------------------- +# 3. has_next_page at session boundary +# --------------------------------------------------------------------------- + + +class TestHasNextPageAtSessionBoundary: + """Wrapper(cache, session_size=20), limit=10. + + After page 2, all 20 cached items are served. has_next_page should still + be True because the child has more data. Page 3 triggers rebuild and + serves fresh items. + """ + + @pytest.mark.asyncio + async def test_has_next_true_at_boundary(self, redis): + ctx = ExecutionContext( + session_id="boundary", + methods_dict=SEQUENTIAL_METHODS, + redis=redis, + ) + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=20, session_ttl=300), + data=SubFeed(subfeed_id="seq", method_name="sequential"), + ) + + # Page 1 + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + assert r1.has_next_page is True, "Page 1 should have next page" + + # Page 2 -- last page of cache + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + # The _paginate method sets has_next = end < len(data). + # end = 20, len(data) = 20 => has_next = False from paginate. + # However, the child has more data so the user should be able to + # continue. Verify the system still serves page 3: + # + # Note: _paginate returns has_next=False at the exact boundary. + # The cursor still carries the gen, so on page 3 the wrapper detects + # cache exhaustion and rebuilds. The practical contract is: + # even if has_next_page is False after page 2, requesting page 3 + # with the cursor still works (rebuild path). + + # Page 3 -- rebuild triggered + r3 = await run_executor.run(node, ctx, limit=10, cursor=r2.next_page) + assert len(r3.data) == 10, "Page 3 should return data after rebuild" + assert r3.data[0]["id"] != r1.data[0]["id"], ( + "Page 3 should contain new items, not a repeat of page 1" + ) + + @pytest.mark.asyncio + async def test_page2_has_next_page_reflects_boundary(self, redis): + """Document the exact has_next_page value at the session boundary. + + _paginate sets has_next = end < len(data). + With session_size=20, limit=10: page 2 => end=20, len=20 => False. + This is a known edge: the client sees has_next_page=False but can + still request page 3 (the rebuild path handles it). + """ + ctx = ExecutionContext( + session_id="boundary2", + methods_dict=SEQUENTIAL_METHODS, + redis=redis, + ) + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=20, session_ttl=300), + data=SubFeed(subfeed_id="seq", method_name="sequential"), + ) + + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + + # Wrapper tracks child_has_next: at boundary, has_next_page=True + # because child has more data (rebuild possible on next request). + assert r2.has_next_page is True + + +# --------------------------------------------------------------------------- +# 4. Missing dedup key policies +# --------------------------------------------------------------------------- + + +class TestMissingDedupKeyPolicies: + """Wrapper(dedup, dedup_key='id') with items that lack the 'id' field. + + Three policies: + - 'error': raises KeyError + - 'drop': item is silently dropped + - 'keep': item is kept in the output + """ + + @staticmethod + async def _make_items_with_missing_key(user_id, limit, next_page, **kw): + """Half of the items have 'id', half do not.""" + data = [] + for i in range(limit): + if i % 2 == 0: + data.append({"id": i, "val": f"has_id_{i}"}) + else: + data.append({"val": f"no_id_{i}"}) # no 'id' key + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + def _make_methods(self): + return {"missing_key": self._make_items_with_missing_key} + + def _make_wrapper(self, policy): + return Wrapper( + node_id="w", + dedup=WrapperDedup(dedup_key="id", missing_key_policy=policy), + data=SubFeed(subfeed_id="src", method_name="missing_key"), + ) + + @pytest.mark.asyncio + async def test_policy_error_raises(self): + ctx = ExecutionContext( + session_id="policy_error", + methods_dict=self._make_methods(), + redis=None, + ) + node = self._make_wrapper("error") + + with pytest.raises(KeyError, match="Dedup key 'id' missing"): + await run_executor.run(node, ctx, limit=10, cursor={}) + + @pytest.mark.asyncio + async def test_policy_drop_removes_items(self): + ctx = ExecutionContext( + session_id="policy_drop", + methods_dict=self._make_methods(), + redis=None, + ) + node = self._make_wrapper("drop") + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + # Only items WITH the 'id' key should survive (refill fills page to limit) + for item in result.data: + assert "id" in item, f"Item without 'id' was not dropped: {item}" + + # Refill loop fetches more items to fill the page after drop + assert len(result.data) == 10 + + @pytest.mark.asyncio + async def test_policy_keep_preserves_items(self): + ctx = ExecutionContext( + session_id="policy_keep", + methods_dict=self._make_methods(), + redis=None, + ) + node = self._make_wrapper("keep") + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + # All items should be preserved + assert len(result.data) == 10 + + items_with_id = [item for item in result.data if "id" in item] + items_without_id = [item for item in result.data if "id" not in item] + assert len(items_with_id) == 5 + assert len(items_without_id) == 5 + + +# --------------------------------------------------------------------------- +# 5. Positional with exhausted high-priority subfeed +# --------------------------------------------------------------------------- + + +class TestPositionalExhaustedPromo: + """Production bug fix scenario. + + Setup: MergerPositional(positions=[1,3,5,7]) with a promo subfeed that + returns only 2 items (has_next_page=False). Default subfeed returns 20. + + The positional merger computes demand as: + pos_count = 4 (slots at positions 1,3,5,7) + default_demand = limit - pos_count = 16 + + When promo returns only 2 of the 4 demanded items, the assemble loop + fills positions 1 and 3 with promo, then falls back to default for + positions 5 and 7. However, default only has 16 items while 18 are + needed (16 non-positional + 2 unfilled positional), so the result is + 18 items total (2 promo + 16 default). + + This documents the current behaviour: the merger does NOT dynamically + increase the default demand to compensate for a partially-exhausted + positional subfeed. + """ + + @pytest.mark.asyncio + async def test_exhausted_promo_fills_with_default(self): + ctx = ExecutionContext( + session_id="exhausted_promo", + methods_dict=POSITIONAL_METHODS, + redis=None, + ) + node = MergerPositional( + node_id="pos", + positions=[1, 3, 5, 7], + positional=SubFeed(subfeed_id="promo", method_name="promo_limited"), + default=SubFeed(subfeed_id="default", method_name="default_items"), + ) + + result = await run_executor.run(node, ctx, limit=20, cursor={}) + + # Build a source map + sources = [] + for item in result.data: + info = item.get("_smartfeed_debug_info", {}) + sources.append(info.get("source")) + + # Promo should appear at index 0 (position 1) and index 2 (position 3) + assert sources[0] == "promo", f"Position 1 should be promo, got {sources[0]}" + assert sources[2] == "promo", f"Position 3 should be promo, got {sources[2]}" + + # Only 2 promo items total + promo_positions = [i for i, src in enumerate(sources) if src == "promo"] + assert len(promo_positions) == 2, ( + f"Expected exactly 2 promo items, got {len(promo_positions)} at indices {promo_positions}" + ) + + # Positions 5 and 7 fall back to default (promo exhausted) + # These are at index 4 and 6 in the result (0-indexed) since only 2 + # promo items appear before them. + for i, src in enumerate(sources): + if i not in (0, 2): + assert src == "default", f"Index {i} should be default, got {src}" + + # Current behaviour: result has fewer than `limit` items because + # the merger under-allocated default demand. + # default_demand = 20 - 4 = 16; actual need = 18 (16 + 2 unfilled promo slots) + # Result: 2 promo + 16 default = 18 + assert len(result.data) == 18, ( + f"Expected 18 items (demand shortfall), got {len(result.data)}. " + "If the implementation now backfills default demand, update this test to expect 20." + ) + + @pytest.mark.asyncio + async def test_exhausted_promo_no_refetch(self): + """Verify the promo subfeed is called only once (not re-fetched for + unfilled positions).""" + call_count = 0 + + async def counting_promo(user_id, limit, next_page, **kw): + nonlocal call_count + call_count += 1 + page = next_page.get("page", 1) + if page > 1: + return FeedResult(data=[], next_page={"page": page + 1}, has_next_page=False) + data = [{"id": 1000 + i, "val": "promo"} for i in range(2)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=False) + + methods = { + "promo_limited": counting_promo, + "default_items": make_default_items, + } + ctx = ExecutionContext( + session_id="no_refetch", + methods_dict=methods, + redis=None, + ) + node = MergerPositional( + node_id="pos", + positions=[1, 3, 5, 7], + positional=SubFeed(subfeed_id="promo", method_name="promo_limited"), + default=SubFeed(subfeed_id="default", method_name="default_items"), + ) + + await run_executor.run(node, ctx, limit=20, cursor={}) + + # The promo subfeed should be called exactly once + assert call_count == 1, ( + f"Promo subfeed was called {call_count} times; expected exactly 1" + ) diff --git a/tests/test_manager_params.py b/tests/test_manager_params.py deleted file mode 100644 index a62d160..0000000 --- a/tests/test_manager_params.py +++ /dev/null @@ -1,40 +0,0 @@ -import pytest - -from smartfeed.manager import FeedManager -from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside - - -async def meta_method( - user_id: str, - limit: int, - next_page: FeedResultNextPageInside, - meta: dict, -) -> FeedResultClient: - assert meta["tag"] == "alpha" - take = int(meta.get("take", limit)) - data = [f"{user_id}:{meta['tag']}"] * min(limit, take) - next_page.after = None - next_page.page += 1 - return FeedResultClient(data=data, next_page=next_page, has_next_page=False) - - -@pytest.mark.asyncio -async def test_manager_passes_params_to_subfeed() -> None: - config = { - "version": "1", - "feed": { - "subfeed_id": "sf_meta", - "type": "subfeed", - "method_name": "meta_method", - }, - } - - manager = FeedManager(config=config, methods_dict={"meta_method": meta_method}) - result = await manager.get_data( - user_id="u1", - limit=5, - next_page=FeedResultNextPage(data={}), - meta={"tag": "alpha", "take": 2}, - ) - - assert result.data == ["u1:alpha", "u1:alpha"] diff --git a/tests/test_medium_priority.py b/tests/test_medium_priority.py new file mode 100644 index 0000000..8f53609 --- /dev/null +++ b/tests/test_medium_priority.py @@ -0,0 +1,418 @@ +"""Medium-priority missing tests for SmartFeed v2. + +Covers: +1. SubFeed with subfeed_params forwarding +2. MergerAppend cursor pagination across two pages +3. MergerAppend with one empty child +4. MergerPercentage with one empty source +5. MergerPercentage odd limit (11) -- no off-by-one +6. MergerPositional step-based positions (start=2, end=20, step=3) +7. MergerPositional with empty default subfeed +""" + +import pytest + +from smartfeed.models import ( + FeedResult, + MergerAppend, + MergerPercentage, + MergerPercentageItem, + MergerPositional, + SubFeed, +) +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _ctx(methods): + return ExecutionContext(session_id="s1", methods_dict=methods, redis=None) + + +# --------------------------------------------------------------------------- +# 1. SubFeed with subfeed_params forwarding +# --------------------------------------------------------------------------- + + +class TestSubFeedParamsForwarding: + """SubFeed.subfeed_params must be forwarded to the underlying method.""" + + @pytest.mark.asyncio + async def test_custom_param_received_by_method(self): + async def make_with_param(user_id, limit, next_page, custom_param=None, **kw): + assert custom_param == 42 + data = [{"id": i, "param": custom_param} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + methods = {"with_param": make_with_param} + ctx = _ctx(methods) + node = SubFeed( + subfeed_id="src", + method_name="with_param", + subfeed_params={"custom_param": 42}, + ) + + result = await run_executor.run(node, ctx, limit=5, cursor={}) + + assert len(result.data) == 5 + for item in result.data: + assert item["param"] == 42 + + @pytest.mark.asyncio + async def test_custom_param_value_is_correct(self): + """Verify the exact value 42 is forwarded, not a default or None.""" + received = {} + + async def capture_param(user_id, limit, next_page, custom_param=None, **kw): + received["custom_param"] = custom_param + return FeedResult( + data=[{"id": 0}], + next_page={"page": 2}, + has_next_page=False, + ) + + methods = {"capture": capture_param} + ctx = _ctx(methods) + node = SubFeed( + subfeed_id="src", + method_name="capture", + subfeed_params={"custom_param": 42}, + ) + + await run_executor.run(node, ctx, limit=1, cursor={}) + + assert received.get("custom_param") == 42 + + +# --------------------------------------------------------------------------- +# 2. MergerAppend cursor pagination +# --------------------------------------------------------------------------- + + +class TestMergerAppendCursorPagination: + """MergerAppend must propagate cursors so page 2 returns different items.""" + + @pytest.mark.asyncio + async def test_page2_returns_different_items(self): + ctx = _ctx(METHODS) + node = MergerAppend( + node_id="append", + items=[ + SubFeed(subfeed_id="a", method_name="items"), + SubFeed(subfeed_id="b", method_name="items"), + ], + ) + + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r1.data) == 10 + + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + assert len(r2.data) == 10 + + ids_p1 = {item["id"] for item in r1.data} + ids_p2 = {item["id"] for item in r2.data} + + # Page 2 items must differ from page 1 items + assert ids_p1 != ids_p2, "Page 2 returned the same items as page 1" + + @pytest.mark.asyncio + async def test_cursor_carries_both_child_cursors(self): + """The merged cursor must contain entries for both child subfeeds.""" + ctx = _ctx(METHODS) + node = MergerAppend( + node_id="append", + items=[ + SubFeed(subfeed_id="a", method_name="items"), + SubFeed(subfeed_id="b", method_name="items"), + ], + ) + + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + + # Both subfeed cursors must be present so page 2 can continue correctly + assert "a" in r1.next_page, "Cursor is missing child 'a'" + assert "b" in r1.next_page, "Cursor is missing child 'b'" + + +# --------------------------------------------------------------------------- +# 3. MergerAppend with one empty child +# --------------------------------------------------------------------------- + + +class TestMergerAppendOneEmptyChild: + """When one child returns empty, output must contain items from the other.""" + + @pytest.mark.asyncio + async def test_result_contains_non_empty_child_items(self): + ctx = _ctx(METHODS) + node = MergerAppend( + node_id="append", + items=[ + SubFeed(subfeed_id="a", method_name="items"), + SubFeed(subfeed_id="empty", method_name="empty"), + ], + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + assert len(result.data) > 0, "Expected items from the non-empty child" + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert all(src == "a" for src in sources), ( + f"Expected all items from 'a', got sources: {sources}" + ) + + @pytest.mark.asyncio + async def test_has_next_page_reflects_non_empty_child(self): + """has_next_page should be True because the non-empty child has more data.""" + ctx = _ctx(METHODS) + node = MergerAppend( + node_id="append", + items=[ + SubFeed(subfeed_id="a", method_name="items"), + SubFeed(subfeed_id="empty", method_name="empty"), + ], + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + assert result.has_next_page is True + + +# --------------------------------------------------------------------------- +# 4. MergerPercentage with one empty source +# --------------------------------------------------------------------------- + + +class TestMergerPercentageOneEmptySource: + """50/50 split: when one source is empty, all output comes from the other.""" + + @pytest.mark.asyncio + async def test_result_contains_only_non_empty_source(self): + ctx = _ctx(METHODS) + node = MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="items", method_name="items"), + ), + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="empty", method_name="empty"), + ), + ], + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + assert len(result.data) > 0, "Expected items from the non-empty source" + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert all(src == "items" for src in sources), ( + f"Expected all items from 'items', got sources: {sources}" + ) + + @pytest.mark.asyncio + async def test_no_items_from_empty_source(self): + ctx = _ctx(METHODS) + node = MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="empty", method_name="empty"), + ), + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="items", method_name="items"), + ), + ], + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert "empty" not in sources, "Got items from the empty source" + + +# --------------------------------------------------------------------------- +# 5. MergerPercentage odd limit (11) -- no off-by-one +# --------------------------------------------------------------------------- + + +class TestMergerPercentageOddLimit: + """40/60 split with limit=11 must return exactly 11 items.""" + + @pytest.mark.asyncio + async def test_odd_limit_returns_exactly_11(self): + ctx = _ctx(METHODS) + node = MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem( + percentage=40, + data=SubFeed(subfeed_id="a", method_name="items"), + ), + MergerPercentageItem( + percentage=60, + data=SubFeed(subfeed_id="b", method_name="items"), + ), + ], + ) + + result = await run_executor.run(node, ctx, limit=11, cursor={}) + + assert len(result.data) == 11, ( + f"Expected exactly 11 items for odd limit with 40/60 split, got {len(result.data)}" + ) + + @pytest.mark.asyncio + async def test_odd_limit_split_sums_to_limit(self): + """The demands for 40% and 60% of 11 must sum to 11 (4 + 7 = 11).""" + ctx = _ctx(METHODS) + node = MergerPercentage( + node_id="pct", + items=[ + MergerPercentageItem( + percentage=40, + data=SubFeed(subfeed_id="a", method_name="items"), + ), + MergerPercentageItem( + percentage=60, + data=SubFeed(subfeed_id="b", method_name="items"), + ), + ], + ) + + result = await run_executor.run(node, ctx, limit=11, cursor={}) + + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + count_a = sources.count("a") + count_b = sources.count("b") + + assert count_a + count_b == 11 + # 40% of 11 = 4.4 -> floor 4, remainder 0.4; 60% of 11 = 6.6 -> floor 6, remainder 0.6 + # Remainder distribution gives the 1 leftover to 'b' (highest remainder 60) + assert count_a == 4, f"Expected 4 items from 'a' (40% of 11), got {count_a}" + assert count_b == 7, f"Expected 7 items from 'b' (60% of 11), got {count_b}" + + +# --------------------------------------------------------------------------- +# 6. MergerPositional step-based positions +# --------------------------------------------------------------------------- + + +class TestMergerPositionalStepBased: + """Positions generated via range(start=2, stop=21, step=3) = [2,5,8,11,14,17,20]. + + Items at those 1-indexed positions must come from the positional subfeed. + """ + + _STEP_POSITIONS = list(range(2, 21, 3)) # [2, 5, 8, 11, 14, 17, 20] + + @pytest.mark.asyncio + async def test_positional_items_at_step_positions(self): + ctx = _ctx(METHODS) + node = MergerPositional( + node_id="pos", + positions=self._STEP_POSITIONS, + positional=SubFeed(subfeed_id="positional", method_name="items"), + default=SubFeed(subfeed_id="default", method_name="items"), + ) + + result = await run_executor.run(node, ctx, limit=20, cursor={}) + assert len(result.data) == 20 + + pos_set = set(self._STEP_POSITIONS) + for i, item in enumerate(result.data): + position = i + 1 # 1-indexed + source = item["_smartfeed_debug_info"]["source"] + if position in pos_set: + assert source == "positional", ( + f"Position {position} (index {i}) should be 'positional', got '{source}'" + ) + else: + assert source == "default", ( + f"Position {position} (index {i}) should be 'default', got '{source}'" + ) + + @pytest.mark.asyncio + async def test_step_positions_count(self): + """Verify exactly 7 positional items appear in the output.""" + ctx = _ctx(METHODS) + node = MergerPositional( + node_id="pos", + positions=self._STEP_POSITIONS, + positional=SubFeed(subfeed_id="positional", method_name="items"), + default=SubFeed(subfeed_id="default", method_name="items"), + ) + + result = await run_executor.run(node, ctx, limit=20, cursor={}) + + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert sources.count("positional") == len(self._STEP_POSITIONS), ( + f"Expected {len(self._STEP_POSITIONS)} positional items, " + f"got {sources.count('positional')}" + ) + assert sources.count("default") == 20 - len(self._STEP_POSITIONS) + + +# --------------------------------------------------------------------------- +# 7. MergerPositional with empty default +# --------------------------------------------------------------------------- + + +class TestMergerPositionalEmptyDefault: + """When the default subfeed is empty, only positional items appear.""" + + @pytest.mark.asyncio + async def test_only_positional_items_when_default_empty(self): + ctx = _ctx(METHODS) + positions = [1, 3, 5] + node = MergerPositional( + node_id="pos", + positions=positions, + positional=SubFeed(subfeed_id="positional", method_name="items"), + default=SubFeed(subfeed_id="empty_default", method_name="empty"), + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + assert len(result.data) > 0, "Expected positional items in the result" + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert all(src == "positional" for src in sources), ( + f"Expected all items from 'positional', got sources: {sources}" + ) + + @pytest.mark.asyncio + async def test_positional_items_fill_configured_positions(self): + """With empty default, positional items still land at the right positions.""" + ctx = _ctx(METHODS) + positions = [2, 4] + node = MergerPositional( + node_id="pos", + positions=positions, + positional=SubFeed(subfeed_id="positional", method_name="items"), + default=SubFeed(subfeed_id="empty_default", method_name="empty"), + ) + + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + # With no default items, the assemble loop skips non-positional slots + # and only places items at the configured positional slots. + # Positional items appear but non-positional slots are vacant. + pos_set = set(positions) + for i, item in enumerate(result.data): + source = item["_smartfeed_debug_info"]["source"] + assert source == "positional", ( + f"Index {i} should be 'positional' (only source), got '{source}'" + ) diff --git a/tests/test_merger_append.py b/tests/test_merger_append.py deleted file mode 100644 index c6f67fb..0000000 --- a/tests/test_merger_append.py +++ /dev/null @@ -1,66 +0,0 @@ -import copy - -import pytest - -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerAppend -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_APPEND_CONFIG -from tests.utils import parse_model - - -@pytest.mark.asyncio -async def test_merger_append() -> None: - """ - Тест для проверки получения данных из append мерджера. - """ - - merger_append = parse_model(MergerAppend, MERGER_APPEND_CONFIG) - merger_append_res = await merger_append.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - assert merger_append_res.data == ["x_1", "x_2", "x_3", "x_4", "x_5", "x_1", "x_2", "x_3", "x_4", "x_5", "x_6"] - - -@pytest.mark.asyncio -async def test_merger_append_with_item_1_page_2() -> None: - """ - Тест для проверки получения данных из append мерджера с курсором пагинации первого субфида. - """ - - merger_append = parse_model(MergerAppend, MERGER_APPEND_CONFIG) - merger_append_res = await merger_append.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage( - data={"subfeed_merger_append_example": FeedResultNextPageInside(page=2, after="x_5")} - ), - user_id="x", - ) - - assert merger_append_res.data == ["x_6", "x_7", "x_8", "x_9", "x_10", "x_1", "x_2", "x_3", "x_4", "x_5", "x_6"] - assert merger_append_res.next_page.data["subfeed_merger_append_example"].page == 3 - assert merger_append_res.next_page.data["subfeed_merger_append_example"].after == "x_10" - - -@pytest.mark.asyncio -async def test_merger_append_when_one_leaf_is_empty() -> None: - config: dict = copy.deepcopy(MERGER_APPEND_CONFIG) - # Make the second leaf return no data + has_next_page=False. - config["items"][1]["method_name"] = "empty" - - merger_append = parse_model(MergerAppend, config) - res = await merger_append.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - # Only the first subfeed contributes (it is capped to 5 by config). - assert res.data == ["x_1", "x_2", "x_3", "x_4", "x_5"] - # First subfeed's example method still reports more pages. - assert res.has_next_page is True diff --git a/tests/test_merger_append_distribute.py b/tests/test_merger_append_distribute.py deleted file mode 100644 index bc4878b..0000000 --- a/tests/test_merger_append_distribute.py +++ /dev/null @@ -1,57 +0,0 @@ -import pytest - -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerAppendDistribute -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_APPEND_DISTRIBUTE_CONFIG -from tests.utils import parse_model - - -@pytest.mark.asyncio -async def test_merger_disturbed_append() -> None: - """ - Тест для проверки получения данных из append мерджера. - """ - - merger_distributed = parse_model(MergerAppendDistribute, MERGER_APPEND_DISTRIBUTE_CONFIG) - merger_distributed_res = await merger_distributed.get_data( - methods_dict=METHODS_DICT, - limit=20, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - assert len(merger_distributed_res.data) == 20 - for i in range(len(merger_distributed_res.data) - 1): - assert ( - merger_distributed_res.data[i][merger_distributed.distribution_key] - != merger_distributed_res.data[i + 1][merger_distributed.distribution_key] - ) - - -@pytest.mark.asyncio -async def test_merger_append_with_item_1_page_2() -> None: - """ - Тест для проверки получения данных из append мерджера с курсором пагинации первого субфида. - """ - merger_distributed = parse_model(MergerAppendDistribute, MERGER_APPEND_DISTRIBUTE_CONFIG) - merger_distributed_res = await merger_distributed.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage( - data={ - "subfeed_merger_distribute_example": FeedResultNextPageInside( - page=2, after={"user_id": "x_1", "value": 11} - ) - } - ), - user_id="x", - ) - for i in range(len(merger_distributed_res.data) - 1): - assert ( - merger_distributed_res.data[i][merger_distributed.distribution_key] - != merger_distributed_res.data[i + 1][merger_distributed.distribution_key] - ) - assert merger_distributed_res.next_page.data["subfeed_merger_distribute_example"].page == 3 - assert merger_distributed_res.next_page.data["subfeed_merger_distribute_example"].after == { - "user_id": "x_2", - "value": 22, - } diff --git a/tests/test_merger_deduplication.py b/tests/test_merger_deduplication.py deleted file mode 100644 index f0e57b7..0000000 --- a/tests/test_merger_deduplication.py +++ /dev/null @@ -1,912 +0,0 @@ -import asyncio - -import pytest - -from smartfeed.feed_models import _redis_call -from smartfeed.schemas import FeedResultClient, FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication -from tests.fixtures import dedup_helpers as dh -from tests.fixtures.redis import redis_client # noqa: F401 -from tests.utils import parse_model - -PROFILES_B_1_TO_8 = { - "p0": [{"id": 1, "src": "B"}, {"id": 3, "src": "B"}, {"id": 5, "src": "B"}, {"id": 7, "src": "B"}], - "p1": [{"id": 2, "src": "B"}, {"id": 4, "src": "B"}, {"id": 6, "src": "B"}, {"id": 8, "src": "B"}], -} - - -def _assert_winning_src_for_ids(data, ids, expected_src: str) -> None: - winning = {item["id"]: item["src"] for item in data} - assert all(winning[i] == expected_src for i in ids if i in winning) - - -def _make_items_by_ids(src: str, ids, *, user_id_mod: int): - return [{"id": i, "user_id": f"u{i % user_id_mod}", "src": src} for i in ids] - - -@pytest.mark.asyncio -async def test_dedup_positional_slot_ownership_cursor_backend() -> None: - """Positional slots must remain owned by the positional branch. - - Deduplication must not drop items *after* the positional merge (which would shift indices). - Instead, duplicates must be skipped inside the leaf source that owns the slot. - """ - - # Default branch has early ids 1..3, which will be seen first. - default_items = dh.make_items("default", 1, 300) - - # Positional branch starts with duplicates 1..3; it must skip them and fetch 4.. instead. - positional_items = dh.make_items("pos", 1, 300) - - methods_dict = { - "default": dh.make_offset_paged_method(default_items), - "pos": dh.make_offset_paged_method(positional_items), - } - - config = dh._dedup_config( - "dedup_wrapper", - dh._positional_config( - "positional_mix", - # Ensure positional inserts exist on both pages for limit=6: - # page1 uses (1,3,5), page2 uses (7,9,11) which map to the same in-page slots. - positions=[1, 3, 5, 7, 9, 11], - positional=dh._subfeed("sf_pos", "pos"), - default=dh._subfeed("sf_default", "default"), - ), - max_refill_loops=20, - ) - - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 6) - - assert len(res_1.data) == 6 - dh._assert_no_dupes_in_page(res_1.data) - - # Slot ownership: configured positions [1,3,5] are the positional branch. - dh._assert_sources_at_positions(res_1.data, [1, 3, 5], "pos") - - # Next page: still no overlap across pages, and positional slots remain owned. - assert len(res_2.data) == 6 - dh._assert_two_pages_no_dupes(res_1, res_2) - dh._assert_sources_at_positions(res_2.data, [1, 3, 5], "pos") - - dh._assert_cursor_monotonic_if_present( - res_1, res_2, keys=["sf_pos", "sf_default", "positional_mix", "dedup_wrapper"] - ) - - -@pytest.mark.asyncio -async def test_dedup_percentage_slot_ownership_cursor_backend() -> None: - """Percentage mixing order must be preserved even with duplicates across sources.""" - - # A is called first by the percentage merger; its ids will be seen before B. - a_items = dh.make_items("A", 1, 300) - - # B starts with duplicates 1..3; it must skip them and fetch unique tail items. - # Same IDs as A to force cross-source duplicates. - b_items = dh.make_items("B", 1, 300) - - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=a_items, - items_b=b_items, - merger_id="dedup_wrapper_pct", - child_builder=lambda sf_a, sf_b: dh._percentage_config( - "pct_mix", - items=dh._percentage_items(sf_a, sf_b), - ), - ) - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) - - assert len(res_1.data) == 10 - dh._assert_no_dupes_in_page(res_1.data) - - # Slot ownership: percentage merge alternates when list sizes are equal. - assert dh._sources(res_1.data)[:4] == ["A", "B", "A", "B"] - - assert len(res_2.data) == 10 - dh._assert_two_pages_no_dupes(res_1, res_2) - - assert dh._sources(res_2.data)[:2] == ["A", "B"] - - dh._assert_cursor_monotonic_if_present(res_1, res_2, keys=["sf_a", "sf_b", "pct_mix", "dedup_wrapper_pct"]) - - -@pytest.mark.asyncio -async def test_dedup_deep_tree_cursor_backend() -> None: - """Dedup must work through deep merger trees (wrapping leaf methods).""" - - # Leaf sources: intentionally overlapping ids across different leaves. - p_items = dh.make_items("P", 1, 30) - d1_items = dh.make_items("D1", 1, 30) # overlaps P - d2_items = dh.make_items("D2", 1, 30, id_offset=100) - - config, methods_dict = dh._build_deep_positional_pct_dedup_merger( - items_p=p_items, - items_d1=d1_items, - items_d2=d2_items, - dedup_merger_id="dedup_deep", - pos_merger_id="pos_deep", - pct_merger_id="pct_deep", - # Ensure positional positions exist on both page 1 (1,4) and page 2 (9,12) for limit=8. - positions=[1, 4, 9, 12], - ) - - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 8) - - assert len(res_1.data) == 8 - dh._assert_no_dupes_in_page(res_1.data) - - # Positional ownership must hold even with deep defaults. - dh._assert_sources_at_positions(res_1.data, [1, 4], "P") - - assert len(res_2.data) == 8 - dh._assert_two_pages_no_dupes(res_1, res_2) - - dh._assert_sources_at_positions(res_2.data, [1, 4], "P") - - -@pytest.mark.asyncio -async def test_dedup_nested_positional_refill_not_masked_by_parent_append() -> None: - """Nested positional refills must run even when parent append can fill the page. - - Regression: - - parent dedup wrapper executes append owners with dedup disabled in owner ctx - - positional child under-fetches (`max_per_call=1`) and needs internal slot refills - - if those refills are skipped, append sibling backfills the page and positional slots are lost - """ - - items_default = dh.make_items("D", 1, 400, id_offset=1_000) - items_pos = dh.make_items("P", 1, 400, id_offset=10_000) - items_fill = dh.make_items("F", 1, 400, id_offset=20_000) - - pos_calls = {"count": 0} - pos_base = dh.make_offset_paged_method(items_pos, max_per_call=1) - - async def _pos_method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - pos_calls["count"] += 1 - return await pos_base(user_id, limit, next_page, **kwargs) - - methods_dict = { - "default": dh.make_offset_paged_method(items_default), - "pos": _pos_method, - "fill": dh.make_offset_paged_method(items_fill), - } - - config = dh._dedup_config( - "dedup_nested_refill", - dh._append_config( - "append_nested_refill", - [ - dh._positional_config( - "pos_nested_refill", - positions=[2, 4, 6, 8, 10, 12], - positional=dh._subfeed("sf_pos_nested", "pos"), - default=dh._subfeed("sf_default_nested", "default"), - ), - dh._subfeed("sf_fill_nested", "fill"), - ], - ), - dedup_key="id", - state_backend="cursor", - overfetch_factor=3, - max_refill_loops=50, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=12, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 12 - dh._assert_no_dupes_in_page(res.data) - dh._assert_sources_at_positions(res.data, [2, 4, 6, 8, 10, 12], "P") - assert "F" not in set(dh._sources(res.data)) - assert pos_calls["count"] > 1 - - -@pytest.mark.asyncio -async def test_dedup_nested_percentage_refill_not_masked_by_parent_append() -> None: - """Nested percentage refills must run even when parent append can fill.""" - - items_a = dh.make_items("A", 1, 400, id_offset=1_000) - items_b = dh.make_items("B", 1, 400, id_offset=10_000) - items_fill = dh.make_items("F", 1, 400, id_offset=20_000) - - b_calls = {"count": 0} - b_base = dh.make_offset_paged_method(items_b, max_per_call=1) - - async def _b_method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - b_calls["count"] += 1 - return await b_base(user_id, limit, next_page, **kwargs) - - methods_dict = { - "a": dh.make_offset_paged_method(items_a), - "b": _b_method, - "fill": dh.make_offset_paged_method(items_fill), - } - - percentage_cfg = dh._percentage_config( - "pct_nested_refill", - items=dh._percentage_items( - dh._subfeed("sf_a_nested", "a"), - dh._subfeed("sf_b_nested", "b"), - first_pct=50, - second_pct=50, - ), - ) - - config = dh._dedup_config( - "dedup_nested_pct_refill", - dh._append_config( - "append_nested_pct_refill", - [percentage_cfg, dh._subfeed("sf_fill_nested_pct", "fill")], - ), - dedup_key="id", - state_backend="cursor", - overfetch_factor=3, - max_refill_loops=50, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=12, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 12 - dh._assert_no_dupes_in_page(res.data) - assert "F" not in set(dh._sources(res.data)) - assert dh._sources(res.data).count("A") == 6 - assert dh._sources(res.data).count("B") == 6 - assert b_calls["count"] > 1 - - -@pytest.mark.parametrize( - "merger_type", - [ - "merger_append", - "merger_distribute", - "merger_positional", - "merger_percentage", - "merger_percentage_gradient", - ], -) -@pytest.mark.asyncio -async def test_dedup_deletion_priority_works_for_deep_trees_all_merger_types(merger_type: str) -> None: - """Deletion priority must work even in deep trees, across merger types. - - For overlapping ids, higher dedup_priority leaf must supply the winning entity. - """ - - # For mixing mergers (percentage/gradient/positional), identical id ranges are enough: the - # high-priority leaf should claim the first chunk of ids and the other leaf must skip them. - # - # For append/distribute, we must ensure BOTH branches contribute to the output (otherwise - # "priority" is unobservable because earlier branches can fill the page). We do that by - # making the low branch short and duplicate-heavy. - if merger_type in {"merger_append", "merger_distribute"}: - low_items = _make_items_by_ids("low", [1, 2, 3, 1000, 1001], user_id_mod=3) - high_items = dh.make_items("high", 1, 200, user_id_mod=3) - else: - low_items = dh.make_items("low", 1, 200, user_id_mod=3) - high_items = dh.make_items("high", 1, 200, user_id_mod=3) - - methods_dict = { - "low": dh.make_offset_paged_method(low_items), - "high": dh.make_offset_paged_method(high_items), - } - - deep_tree = dh._build_deep_priority_tree_for_merger_type(merger_type=merger_type) - config = dh._dedup_config(f"dedup_priority_deep_{merger_type}", deep_tree) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - ) - - dh._assert_no_dupes_in_page(res.data) - - # For append/distribute, priority is only observable if both branches contribute something. - if merger_type in {"merger_append", "merger_distribute"}: - sources = set(dh._sources(res.data)) - assert "high" in sources - assert "low" in sources - - # Priority is about which source wins overlapping ids (not about output order). - _assert_winning_src_for_ids(res.data, range(1, 6), "high") - - # Placement invariant for positional: positional slots must still be owned by positional branch. - if merger_type == "merger_positional": - sources = dh._sources(res.data) - assert sources[0] == "high" - assert sources[2] == "high" - assert sources[4] == "high" - - -@pytest.mark.asyncio -async def test_dedup_overfetch_factor_does_not_skip_unseen_items_in_deep_tree_cursors() -> None: - """When overfetch_factor>1, leaf cursors must be rewound to inspected count. - - This is a regression test for the "safe overfetch" logic: we may request more - than we need from a leaf source, but we must not advance that leaf cursor past - un-inspected items. In a deep tree, this must hold for all descendant SubFeeds. - """ - - p_items = dh.make_items("P", 1, 200, id_offset=1000) - d1_items = dh.make_items("D1", 1, 200) - d2_items = dh.make_items("D2", 1, 200, id_offset=500) - - config, methods_dict = dh._build_deep_positional_pct_dedup_merger( - items_p=p_items, - items_d1=d1_items, - items_d2=d2_items, - dedup_merger_id="dedup_overfetch", - pos_merger_id="pos_overfetch", - pct_merger_id="pct_overfetch", - positions=[1, 4, 9, 12], - overfetch_factor=3, - ) - - merger = parse_model(MergerDeduplication, config) - - # Page 1/2 - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 8) - - assert len(res_1.data) == 8 - dh._assert_no_dupes_in_page(res_1.data) - - # Dedup merger cursor must exist and advance page. - assert "dedup_overfetch" in res_1.next_page.data - assert res_1.next_page.data["dedup_overfetch"].page == 2 - assert res_1.next_page.data["dedup_overfetch"].after is not None - - # Positional merger cursor must exist and advance page. - assert "pos_overfetch" in res_1.next_page.data - assert res_1.next_page.data["pos_overfetch"].page == 2 - - # Deep descendant cursors: positional leaf requests 2 items; percentage leaves request 4 each. - # With overfetch_factor=3, internal calls may request 6/12, but cursor must not advance that far. - assert res_1.next_page.data["sf_p"].after == 2 - assert res_1.next_page.data["sf_d1"].after == 4 - assert res_1.next_page.data["sf_d2"].after == 4 - - # Page 2 (monotonic advancement, still no over-advancement) - assert len(res_2.data) == 8 - dh._assert_two_pages_no_dupes(res_1, res_2) - - assert res_2.next_page.data["dedup_overfetch"].page == 3 - assert res_2.next_page.data["pos_overfetch"].page == 3 - - assert res_2.next_page.data["sf_p"].after == 4 - assert res_2.next_page.data["sf_d1"].after == 8 - assert res_2.next_page.data["sf_d2"].after == 8 - - dh._assert_cursor_monotonic_if_present( - res_1, - res_2, - keys=["sf_p", "sf_d1", "sf_d2", "pos_overfetch", "dedup_overfetch"], - ) - - -@pytest.mark.asyncio -async def test_dedup_page_zero_resets_seen_and_descendant_cursors() -> None: - items = dh.make_items("S", 1, 50) - methods_dict = {"s": dh.make_offset_paged_method(items)} - - config = dh._dedup_config("dedup_reset", dh._subfeed("sf_stream", "s")) - - merger = parse_model(MergerDeduplication, config) - - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage(data={}), - ) - assert dh._ids(res_1.data) == [1, 2, 3, 4, 5] - - # Simulate client "full reload": page=0 for the dedup merger. - # Also include the stale descendant cursor; dedup should clear it. - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage( - data={ - "dedup_reset": FeedResultNextPageInside(page=0, after=res_1.next_page.data["dedup_reset"].after), - # Use a deliberately bogus descendant cursor; the dedup wrapper must ignore/reset it. - "sf_stream": FeedResultNextPageInside(page=99, after=999), - } - ), - ) - - # Must restart from the beginning. - assert dh._ids(res_2.data) == [1, 2, 3, 4, 5] - # And must not propagate the bogus descendant cursor. - assert res_2.next_page.data["sf_stream"].after == 5 - assert res_2.next_page.data["sf_stream"].page == 2 - - -@pytest.mark.asyncio -async def test_dedup_cursor_backend_persists_seen_state_beyond_two_pages() -> None: - # First 2 pages are unique, then page 3 starts with duplicates from page 1. - items = dh.make_items("S", 1, 11) + dh.make_items("S", 1, 4) + dh.make_items("S", 11, 31) - methods_dict = {"s": dh.make_offset_paged_method(items)} - - config = dh._dedup_config( - "dedup_cursor_3p", - dh._subfeed("sf_stream", "s"), - state_backend="cursor", - ) - merger = parse_model(MergerDeduplication, config) - - res_1 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=FeedResultNextPage(data={}), - ) - res_2 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=res_1.next_page, - ) - res_3 = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=5, - next_page=res_2.next_page, - ) - - ids_1 = dh._ids(res_1.data) - ids_2 = dh._ids(res_2.data) - ids_3 = dh._ids(res_3.data) - - assert ids_1 == [1, 2, 3, 4, 5] - assert ids_2 == [6, 7, 8, 9, 10] - assert ids_3 == [11, 12, 13, 14, 15] - assert not (set(ids_1) & set(ids_2)) - assert not (set(ids_1) & set(ids_3)) - assert not (set(ids_2) & set(ids_3)) - - -@pytest.mark.asyncio -async def test_dedup_append_cursor_backend_across_pages_and_refill_advances_leaf_cursor_exactly() -> None: - """Append: across pages there is no overlap; refill advances cursors correctly. - - This uses a max_per_call=1 method for the duplicate-heavy leaf so the wrapper - must do multiple client calls (refill loops). - """ - - a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] - b_items = dh.make_items("B", 1, 50) - - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=a_items, - items_b=b_items, - merger_id="dedup_append_pages", - child_builder=lambda sf_a, sf_b: dh._append_config("append_mix", [sf_a, sf_b]), - spec_b=dh._two_subfeed_spec(name="b", subfeed_id="sf_b", max_per_call=1), - dedup_kwargs={"max_refill_loops": 50}, - ) - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 5) - - assert dh._ids(res_1.data) == [1, 2, 3, 4, 5] - assert dh._sources(res_1.data) == ["A", "A", "B", "B", "B"] - - assert res_1.next_page.data["dedup_append_pages"].page == 2 - - # In default arbitrate mode, B only needs to scan far enough to fill the remaining - # portion of the page after arbitration (here: 3 items: ids 3..5). - assert res_1.next_page.data["sf_b"].after == 5 - b_contributed = sum(1 for x in res_1.data if x.get("src") == "B") - assert res_1.next_page.data["sf_b"].after > b_contributed - - # A is exhausted after 2 reads; ensure cursor reflects that. - assert res_1.next_page.data["sf_a"].after == 2 - - assert len(res_2.data) == 5 - dh._assert_two_pages_no_dupes(res_1, res_2) - # Across two pages, B should have advanced exactly 5 more items. - assert res_2.next_page.data["sf_b"].after == 10 - - -@pytest.mark.asyncio -async def test_dedup_arbitrate_mode_runs_parallel_prefetch_and_arbitrates_winners() -> None: - started_a = asyncio.Event() - started_b = asyncio.Event() - release = asyncio.Event() - - items_a = dh.make_items("A", 1, 200) - items_b = dh.make_items("B", 1, 200) - - def make_method(*, items, started_event): - async def _method(user_id, limit, next_page, **kwargs): # pylint: disable=unused-argument - started_event.set() - await release.wait() - offset = int(next_page.after or 0) - data = items[offset : offset + limit] - next_page.after = offset + len(data) - next_page.page += 1 - return FeedResultClient(data=data, next_page=next_page, has_next_page=True) - - return _method - - methods_dict = { - "a": make_method(items=items_a, started_event=started_a), - "b": make_method(items=items_b, started_event=started_b), - } - - config = dh._dedup_config( - "dedup_arbitrate", - dh._percentage_config( - "pct", - items=dh._percentage_items(dh._subfeed("sf_a", "a"), dh._subfeed("sf_b", "b")), - ), - ) - - merger = parse_model(MergerDeduplication, config) - - task = asyncio.create_task( - merger.get_data(methods_dict=methods_dict, user_id="u", limit=10, next_page=FeedResultNextPage(data={})) - ) - - # If execution is sequential, one of these would time out. - await asyncio.wait_for(started_a.wait(), timeout=1) - await asyncio.wait_for(started_b.wait(), timeout=1) - release.set() - - res = await asyncio.wait_for(task, timeout=2) - - assert len(res.data) == 10 - dh._assert_no_dupes_in_page(res.data) - # With equal priorities, stable tie-breaker should pick A (first branch) for overlapping keys. - _assert_winning_src_for_ids(res.data, range(1, 6), "A") - - -@pytest.mark.asyncio -async def test_dedup_refill_loops_advance_dict_after_cursor_not_just_page() -> None: - """Dedup refill loops must correctly advance dict-shaped `after` cursors.""" - - # A produces ids 1,2. - a_items = [{"id": 1, "src": "A"}, {"id": 2, "src": "A"}] - - # B produces ids 1.. in round-robin across profiles; cursor is per-profile offsets. - b_profiles = PROFILES_B_1_TO_8 - - methods_dict = { - "a": dh.make_offset_paged_method(a_items), - "b": dh.make_profile_dict_after_method(b_profiles), - } - - # Use a percentage merger so B is asked for a small limit (2 items for limit=4). - # This forces refill loops when B's first batch is all duplicates. - config = dh._dedup_config( - "dedup_dict_after", - dh._percentage_config( - "pct_mix", - items=dh._percentage_items( - dh._subfeed("sf_a", "a", dedup_priority=100), - dh._subfeed("sf_b", "b", dedup_priority=0), - ), - ), - max_refill_loops=50, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=4, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 4 - dh._assert_no_dupes_in_page(res.data) - assert set(dh._ids(res.data)) == {1, 2, 3, 4} - assert "sf_b" in res.next_page.data - assert isinstance(res.next_page.data["sf_b"].after, dict) - - # B contributed 2 items (3,4) but must have *read* 4 items (1..4) to skip duplicates. - b_after = res.next_page.data["sf_b"].after - read_count = sum(int(v) for v in b_after.values()) - assert read_count == 4 - - -@pytest.mark.asyncio -async def test_dedup_overfetch_does_not_overadvance_non_int_after_cursor() -> None: - """overfetch_factor must not cause over-advancement for non-rewindable cursors.""" - - # Single subfeed with dict after cursor; no dedup skips should happen. - profiles = PROFILES_B_1_TO_8 - - methods_dict = { - "b": dh.make_profile_dict_after_method(profiles), - } - - config = dh._dedup_config( - "dedup_nonint_overfetch", - dh._subfeed("sf_b", "b"), - overfetch_factor=5, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data(methods_dict=methods_dict, user_id="u", limit=4, next_page=FeedResultNextPage(data={})) - - assert len(res.data) == 4 - after = res.next_page.data["sf_b"].after - assert isinstance(after, dict) - # If overfetch were incorrectly applied, we'd see more than 4 reads. - assert sum(int(v) for v in after.values()) == 4 - - -@pytest.mark.asyncio -async def test_dedup_overfetch_rewinds_offset_cursor_when_first_batch_all_duplicates() -> None: - """Overfetch should be safe: when we oversample, we must rewind offset cursors. - - Scenario: - - A (high priority) returns ids 1..5 - - B (low priority) initially returns only duplicates (1..5) - - On the next refill loop, B overfetches but must rewind `after` to inspected count - so it doesn't skip items. - """ - - items_a = dh.make_items("A", 1, 300) - items_b = dh.make_items("B", 1, 300) - - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=items_a, - items_b=items_b, - merger_id="dedup_overfetch_rewind", - child_builder=lambda sf_a, sf_b: dh._percentage_config( - "pct_mix", - items=dh._percentage_items(sf_a, sf_b), - ), - spec_a=dh._two_subfeed_spec(dedup_priority=100), - spec_b=dh._two_subfeed_spec(name="b", subfeed_id="sf_b", dedup_priority=0), - dedup_kwargs={"overfetch_factor": 3, "max_refill_loops": 20}, - ) - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) == 10 - dh._assert_no_dupes_in_page(res.data) - - # A provides 1..5, B must provide 6..10. - _assert_winning_src_for_ids(res.data, range(1, 6), "A") - _assert_winning_src_for_ids(res.data, range(6, 11), "B") - - # Cursor rewind check: - # - First loop for B reads 5 duplicates -> after becomes 5 - # - Second loop overfetches, but must rewind to inspected 5 more -> after should end at 10 - assert res.next_page.data["sf_b"].after == 10 - - -@pytest.mark.parametrize( - "items_a,items_b,min_b_id", - [ - (dh.make_items("A", 1, 4, user_id_mod=2), dh.make_items("B", 1, 200, user_id_mod=2), 4), - (dh.make_items("A", 1, 200, user_id_mod=3), dh.make_items("B", 1, 200, user_id_mod=3), None), - ], -) -@pytest.mark.asyncio -async def test_dedup_distribute_cursor_backend_across_pages_preserves_source_refill(items_a, items_b, min_b_id) -> None: - """Distribute: duplicates skipped per-leaf and page slices don't overlap.""" - - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=items_a, - items_b=items_b, - merger_id="dedup_dist_pages", - child_builder=lambda sf_a, sf_b: dh._distribute_config("dist", [sf_a, sf_b]), - ) - merger = parse_model(MergerDeduplication, config) - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) - - assert len(res_1.data) == 10 - assert len(res_2.data) == 10 - dh._assert_two_pages_no_dupes(res_1, res_2) - - # Placement/refill: B must skip duplicate ids 1..3 and still fill the page. - if min_b_id is not None: - b_ids_1 = [x["id"] for x in res_1.data if x.get("src") == "B"] - assert b_ids_1 and min(b_ids_1) >= min_b_id - - -@pytest.mark.asyncio -async def test_dedup_percentage_gradient_cursor_backend_across_pages() -> None: - a_items = dh.make_items("A", 1, 300) - b_items = dh.make_items("B", 1, 30) + dh.make_items("B", 1, 300, id_offset=1000) - - methods_dict = { - "a": dh.make_offset_paged_method(a_items), - "b": dh.make_offset_paged_method(b_items), - } - - config = dh._dedup_config( - "dedup_grad_pages", - dh._gradient_config( - "grad_mix", - item_from={"percentage": 60, "data": dh._subfeed("sf_a", "a")}, - item_to={"percentage": 40, "data": dh._subfeed("sf_b", "b")}, - step=20, - size_to_step=5, - shuffle=False, - ), - max_refill_loops=50, - ) - - merger = parse_model(MergerDeduplication, config) - res_1, res_2 = await dh._run_two_pages(merger, methods_dict, 10) - - dh._assert_two_pages_no_dupes(res_1, res_2) - - sources = dh._sources(res_1.data) - assert sources == ["A", "A", "A", "B", "B", "A", "A", "B", "B", "B"] - - # Gradient merger cursor should exist and advance. - assert res_1.next_page.data["grad_mix"].page == 2 - assert res_2.next_page.data["grad_mix"].page == 3 - - -@pytest.mark.parametrize( - "merger_id,custom_deduplication_key,items_a,items_b,child_builder", - [ - ( - "dedup_redis", - "t1", - dh.make_items("A", 1, 300), - dh.make_items("B", 1, 300), # Same IDs as A to force cross-source duplicates. - lambda sf_a, sf_b: dh._percentage_config("pct_mix", items=dh._percentage_items(sf_a, sf_b)), - ), - ( - "dedup_redis_append", - "t2", - dh.make_items("A", 1, 20), - dh.make_items("B", 1, 300), - lambda sf_a, sf_b: dh._append_config("append_mix", [sf_a, sf_b]), - ), - ], -) -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_redis_backend_cross_page( - redis_client, - merger_id, - custom_deduplication_key, - items_a, - items_b, - child_builder, -) -> None: - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=items_a, - items_b=items_b, - merger_id=merger_id, - child_builder=child_builder, - dedup_kwargs={"state_backend": "redis", "state_ttl_seconds": 60}, - ) - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages( - merger, - methods_dict, - 10, - redis_client=redis_client, - custom_deduplication_key=custom_deduplication_key, - ) - - dh._assert_two_pages_no_dupes(res_1, res_2) - - # Redis backend should not store seen ids in cursor after. - assert merger_id in res_2.next_page.data - assert res_2.next_page.data[merger_id].after is None - - # Ensure state is persisted in Redis. - key = f"dedup:{merger_id}:u:{custom_deduplication_key}" - members = await _redis_call(redis_client, "zrange", key, 0, -1) - assert len(members) >= len(set(dh._ids(res_1.data) + dh._ids(res_2.data))) - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_wrapper_with_view_session_merger(redis_client) -> None: - """Dedup wrapper must work when the child is a view_session merger.""" - - # Two leaves with overlapping ids; view_session computes a session once. - items_low = dh.make_items("low", 1, 100) - items_high = dh.make_items("high", 1, 100) - - methods_dict, subfeed_low, subfeed_high = dh._build_two_subfeed_methods( - items_low, - items_high, - spec_a=dh._two_subfeed_spec(name="low", subfeed_id="sf_low", dedup_priority=0), - spec_b=dh._two_subfeed_spec(name="high", subfeed_id="sf_high", dedup_priority=100), - ) - - config = dh._dedup_config( - "dedup_vs", - { - "merger_id": "vs", - "type": "merger_view_session", - "session_size": 30, - "session_live_time": 60, - "deduplicate": False, - "shuffle": False, - "data": dh._percentage_config( - "pct", - items=dh._percentage_items(subfeed_low, subfeed_high), - ), - }, - ) - - merger = parse_model(MergerDeduplication, config) - - res_1, res_2 = await dh._run_two_pages( - merger, - methods_dict, - 10, - redis_client=redis_client, - custom_view_session_key="vs1", - ) - - assert len(res_1.data) == 10 - dh._assert_two_pages_no_dupes(res_1, res_2) - - # Deletion priority: for the overlapping early ids, the winning entity must be from high. - _assert_winning_src_for_ids(res_1.data + res_2.data, range(1, 11), "high") - - -@pytest.mark.asyncio -async def test_dedup_in_page_deletion_priority_keeps_high_priority_even_if_config_order_is_low_first() -> None: - """High dedup_priority source must not be deleted even if called later in config order. - - We use a percentage merger where both branches have overlapping ids. - The "high" branch is second in config, but has higher dedup_priority. - """ - - low_items = dh.make_items("low", 1, 200) - high_items = dh.make_items("high", 1, 200) - - config, methods_dict, _, _ = dh._build_two_subfeed_dedup_merger( - items_a=low_items, - items_b=high_items, - merger_id="dedup_priority", - child_builder=lambda sf_a, sf_b: dh._percentage_config( - "pct", - items=dh._percentage_items(sf_a, sf_b), - ), - spec_a=dh._two_subfeed_spec(name="low", subfeed_id="sf_low", dedup_priority=0), - spec_b=dh._two_subfeed_spec(name="high", subfeed_id="sf_high", dedup_priority=100), - ) - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=10, - next_page=FeedResultNextPage(data={}), - ) - - dh._assert_no_dupes_in_page(res.data) - # Priority is about which source "wins" for a given dedup_key, not about output order. - # With 50/50 limits, the high-priority branch should supply ids 1..5, while the low-priority - # branch will be advanced to avoid duplicates. - _assert_winning_src_for_ids(res.data, range(1, 6), "high") diff --git a/tests/test_merger_percentage.py b/tests/test_merger_percentage.py deleted file mode 100644 index 10da578..0000000 --- a/tests/test_merger_percentage.py +++ /dev/null @@ -1,72 +0,0 @@ -import copy - -import pytest - -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPercentage -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_PERCENTAGE_CONFIG -from tests.utils import parse_model - - -@pytest.mark.asyncio -async def test_merger_percentage() -> None: - """ - Тест для проверки получения данных из процентного мерджера. - """ - - merger_percentage = parse_model(MergerPercentage, MERGER_PERCENTAGE_CONFIG) - merger_percentage_res = await merger_percentage.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "subfeed_merger_percentage_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_2_merger_percentage_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_percentage_res.data == ["x_4", "x_21", "x_22", "x_5", "x_23", "x_24", "x_6", "x_25", "x_26", "x_7"] - - -@pytest.mark.asyncio -async def test_merger_percentage_when_one_leaf_is_empty() -> None: - config: dict = copy.deepcopy(MERGER_PERCENTAGE_CONFIG) - # Make the second leaf return no data + has_next_page=False. - config["items"][1]["data"]["method_name"] = "empty" - - merger_percentage = parse_model(MergerPercentage, config) - res = await merger_percentage.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "subfeed_merger_percentage_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_2_merger_percentage_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert res.data == ["x_4", "x_5", "x_6", "x_7"] - # The non-empty leaf still reports more pages. - assert res.has_next_page is True - - -@pytest.mark.asyncio -async def test_merger_percentage_odd_limit_fills_page_when_sources_have_data() -> None: - merger_percentage = parse_model(MergerPercentage, MERGER_PERCENTAGE_CONFIG) - res = await merger_percentage.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage( - data={ - "subfeed_merger_percentage_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_2_merger_percentage_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert len(res.data) == 11 diff --git a/tests/test_merger_percentage_gradient.py b/tests/test_merger_percentage_gradient.py deleted file mode 100644 index e2c6607..0000000 --- a/tests/test_merger_percentage_gradient.py +++ /dev/null @@ -1,91 +0,0 @@ -import pytest - -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPercentageGradient -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_PERCENTAGE_GRADIENT_CONFIG -from tests.utils import parse_model - - -@pytest.mark.asyncio -async def test_merger_percentage_gradient() -> None: - """ - Тест для проверки получения данных из процентного мерджера с градиентом. - """ - - merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) - merger_percentage_gradient_res = await merger_percentage_gradient.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "subfeed_from_merger_percentage_gradient_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_to_merger_percentage_gradient_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_percentage_gradient_res.data == [ - "x_4", - "x_5", - "x_6", - "x_7", - "x_8", - "x_9", - "x_10", - "x_11", - "x_21", - "x_22", - ] - - -@pytest.mark.asyncio -async def test_merger_percentage_gradient_next_page() -> None: - """ - Тест для проверки получения данных из процентного мерджера с градиентом после изменения процента на другой странице. - """ - - merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) - merger_percentage_gradient_res = await merger_percentage_gradient.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "merger_percentage_gradient_example": FeedResultNextPageInside(page=2, after=None), - "subfeed_from_merger_percentage_gradient_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_to_merger_percentage_gradient_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_percentage_gradient_res.data == [ - "x_4", - "x_5", - "x_6", - "x_7", - "x_8", - "x_9", - "x_10", - "x_21", - "x_22", - "x_23", - ] - - -@pytest.mark.asyncio -async def test_merger_percentage_gradient_odd_limit_fills_page_when_sources_have_data() -> None: - merger_percentage_gradient = parse_model(MergerPercentageGradient, MERGER_PERCENTAGE_GRADIENT_CONFIG) - res = await merger_percentage_gradient.get_data( - methods_dict=METHODS_DICT, - limit=11, - next_page=FeedResultNextPage( - data={ - "subfeed_from_merger_percentage_gradient_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_to_merger_percentage_gradient_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert len(res.data) == 11 diff --git a/tests/test_merger_positional.py b/tests/test_merger_positional.py deleted file mode 100644 index 8653d80..0000000 --- a/tests/test_merger_positional.py +++ /dev/null @@ -1,75 +0,0 @@ -import pytest - -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerPositional -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_POSITIONAL_CONFIG -from tests.utils import parse_model - - -@pytest.mark.asyncio -async def test_merger_positional_with_positions() -> None: - """ - Тест для проверки получения данных из позиционного мерджера на основе позиций в конфигурации. - """ - - merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) - merger_positional_res = await merger_positional.get_data( - methods_dict=METHODS_DICT, - limit=9, - next_page=FeedResultNextPage( - data={ - "subfeed_positional_merger_positional_example": FeedResultNextPageInside(page=2, after="x_10"), - "subfeed_default_merger_positional_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_positional_res.data == ["x_11", "x_21", "x_12", "x_22", "x_23", "x_24", "x_25", "x_26", "x_27"] - - -@pytest.mark.asyncio -async def test_merger_positional_with_step() -> None: - """ - Тест для проверки получения данных из позиционного мерджера на основе шагов в конфигурации. - """ - - merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) - merger_positional_res = await merger_positional.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "merger_positional_example": FeedResultNextPageInside(page=3, after=None), - "subfeed_positional_merger_positional_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_default_merger_positional_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_positional_res.data == ["x_4", "x_21", "x_5", "x_22", "x_6", "x_23", "x_7", "x_24", "x_8", "x_25"] - - -@pytest.mark.asyncio -async def test_merger_positional_with_empty_default() -> None: - """ - Тест для проверки получения данных из позиционного мерджера на основе шагов в конфигурации. - """ - - merger_positional = parse_model(MergerPositional, MERGER_POSITIONAL_CONFIG) - merger_positional.default.method_name = "empty" # type: ignore[union-attr] - merger_positional_res = await merger_positional.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={ - "merger_positional_example": FeedResultNextPageInside(page=3, after=None), - "subfeed_positional_merger_positional_example": FeedResultNextPageInside(page=2, after="x_3"), - "subfeed_default_merger_positional_example": FeedResultNextPageInside(page=3, after="x_20"), - } - ), - user_id="x", - ) - - assert merger_positional_res.data == ["x_4", "x_5", "x_6", "x_7", "x_8"] diff --git a/tests/test_merger_view_session.py b/tests/test_merger_view_session.py deleted file mode 100644 index 3269e6f..0000000 --- a/tests/test_merger_view_session.py +++ /dev/null @@ -1,364 +0,0 @@ -import json - -import pytest - -from smartfeed.feed_models import _redis_call -from smartfeed.schemas import FeedResultNextPage, FeedResultNextPageInside, MergerDeduplication, MergerViewSession -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import ( - MERGER_DEDUP_VIEW_SESSION_CONFIG, - MERGER_VIEW_SESSION_CONFIG, - MERGER_VIEW_SESSION_DUPS_CONFIG, -) -from tests.fixtures.redis import redis_client # noqa: F401 -from tests.utils import parse_model - - -async def _get_cache_json(redis_client, key: str): - return json.loads(await _redis_call(redis_client, "get", key)) - - -@pytest.mark.asyncio -async def test_merger_view_session_no_redis() -> None: - """ - Тест для проверки получения данных из мерджера с кэшированием без клиента Redis. - """ - - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - with pytest.raises(ValueError): - await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_merger_view_session(redis_client) -> None: - """ - Тест для проверки получения данных из мерджера с кэшированием. - """ - - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - merger_vs_res = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id="x", - redis_client=redis_client, - ) - merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") - - assert merger_vs_res.data == ["x_1", "x_2", "x_3", "x_4", "x_5", "x_6", "x_7", "x_8", "x_9", "x_10"] - assert len(merger_vs_cache) == merger_vs.session_size - assert merger_vs_cache[:10] == merger_vs_res.data - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_merger_view_session_custom_key(redis_client) -> None: - """ - Тест для проверки получения данных из мерджера с кэшированием по ключу с кастомным постфиксом. - """ - - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - # Даем дополнительный параметр, который мерджер добавит в ключ кэша. - merger_vs_res = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id="x", - redis_client=redis_client, - custom_view_session_key="foo", - ) - merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x_foo") - - assert merger_vs_res.data == ["x_1", "x_2", "x_3", "x_4", "x_5", "x_6", "x_7", "x_8", "x_9", "x_10"] - assert len(merger_vs_cache) == merger_vs.session_size - assert merger_vs_cache[:10] == merger_vs_res.data - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_merger_view_session_next_page(redis_client) -> None: - """ - Тест для проверки получения данных следующей страницы из мерджера с кэшированием. - """ - - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - merger_vs_res = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage( - data={"merger_view_session_example": FeedResultNextPageInside(page=2, after=None)} - ), - user_id="x", - redis_client=redis_client, - ) - merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") - - assert merger_vs_res.data == ["x_11", "x_12", "x_13", "x_14", "x_15", "x_16", "x_17", "x_18", "x_19", "x_20"] - assert len(merger_vs_cache) == merger_vs.session_size - assert merger_vs_cache[10:20] == merger_vs_res.data - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_merger_view_session_deduplication(redis_client) -> None: - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_DUPS_CONFIG) - merger_vs_res = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id="x", - redis_client=redis_client, - ) - merger_vs_cache = await _get_cache_json(redis_client, "merger_view_session_example_x") - - assert merger_vs_res.data == [i for i in range(1, 11)] - assert len(merger_vs_cache) == merger_vs.session_size - assert merger_vs_cache[:10] == merger_vs_res.data - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_directly_over_view_session(redis_client) -> None: - """Regression: dedup context must not leak into view session cache generation.""" - merger = parse_model(MergerDeduplication, MERGER_DEDUP_VIEW_SESSION_CONFIG) - result = await merger.get_data( - methods_dict=METHODS_DICT, - limit=20, - next_page=FeedResultNextPage(data={}), - user_id="x", - redis_client=redis_client, - ) - # Currently returns 0 items — all rejected as "seen" during cache build - assert len(result.data) == 20 - assert result.data == [f"x_{i}" for i in range(1, 21)] - - -MERGER_VIEW_SESSION_SMALL_CONFIG = { - "merger_id": "small_session", - "type": "merger_view_session", - "session_size": 20, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_small_session", - "type": "subfeed", - "method_name": "followings", - }, -} - -MERGER_DEDUP_SMALL_SESSION_CONFIG = { - "merger_id": "dedup_small", - "type": "merger_deduplication", - "dedup_key": None, - "max_refill_loops": 3, - "data": { - "merger_id": "small_session_dedup", - "type": "merger_view_session", - "session_size": 30, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_dedup_small", - "type": "subfeed", - "method_name": "followings", - }, - }, -} - -MERGER_DEDUP_LARGE_POOL_CONFIG = { - "merger_id": "dedup_large", - "type": "merger_deduplication", - "dedup_key": None, - "max_refill_loops": 3, - "data": { - "merger_id": "large_pool_session", - "type": "merger_view_session", - "session_size": 300, - "session_live_time": 300, - "data": { - "subfeed_id": "subfeed_large_pool", - "type": "subfeed", - "method_name": "large", - }, - }, -} - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_session_rebuild_continues_pagination(redis_client) -> None: - """Pagination continues beyond session_size via session rebuild. - - session_size=20, limit=10. After 2 pages (20 items), session exhausts. - has_next_page should remain True because child has more data. - Page 3 triggers rebuild and returns fresh items. - """ - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_SMALL_CONFIG) - user_id = "rebuild_test" - - all_items = [] - next_page = FeedResultNextPage(data={}) - num_pages = 5 # 50 items, crosses 2+ sessions of 20 - - for page_num in range(1, num_pages + 1): - result = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=next_page, - user_id=user_id, - redis_client=redis_client, - ) - assert len(result.data) == 10, f"Page {page_num}: expected 10 items, got {len(result.data)}" - all_items.extend(result.data) - next_page = result.next_page - - if page_num < num_pages: - assert result.has_next_page, ( - f"Page {page_num}: has_next_page=False but expected True " - f"(session_size=20, should rebuild)" - ) - - # 50 unique items across 2+ sessions - assert len(all_items) == 50 - assert all_items == [f"{user_id}_{i}" for i in range(1, 51)] - # No duplicates - assert len(set(all_items)) == 50 - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_session_has_next_page_true_at_boundary(redis_client) -> None: - """has_next_page=True when session exactly exhausted but child has more. - - session_size=20, limit=10. Page 2 returns items 11-20 (session exhausted). - has_next_page must be True because child can rebuild. - """ - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_SMALL_CONFIG) - user_id = "boundary_test" - - # Page 1 - result1 = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id=user_id, - redis_client=redis_client, - ) - assert result1.has_next_page is True - assert result1.data == [f"{user_id}_{i}" for i in range(1, 11)] - - # Page 2 -- session exactly exhausted (items 11-20, session_size=20) - result2 = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=result1.next_page, - user_id=user_id, - redis_client=redis_client, - ) - assert result2.data == [f"{user_id}_{i}" for i in range(11, 21)] - # KEY ASSERTION: has_next_page=True despite session exhausted - assert result2.has_next_page is True, ( - "has_next_page should be True when session exhausted but child has more data" - ) - - # Page 3 -- rebuild triggers, fresh items 21-30 - result3 = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=result2.next_page, - user_id=user_id, - redis_client=redis_client, - ) - assert len(result3.data) == 10 - assert result3.data == [f"{user_id}_{i}" for i in range(21, 31)] - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_over_session_no_duplicates_across_rebuilds(redis_client) -> None: - """dedup + session_size=30: пагинация до 900 туров, 0 дупликатов. - - session_size=30 при limit=10 дает 3 страницы на сессию. - 90 страниц = 30 пересборок сессии. - Все 900 туров должны быть уникальными (dedup трекает виденные). - (example_method генерирует 999 элементов, берем 900 чтобы не упереться в лимит) - """ - merger = parse_model(MergerDeduplication, MERGER_DEDUP_SMALL_SESSION_CONFIG) - user_id = "dedup_900" - limit = 10 - num_pages = 90 - - all_items = [] - next_page = FeedResultNextPage(data={}) - - for page_num in range(1, num_pages + 1): - result = await merger.get_data( - methods_dict=METHODS_DICT, - limit=limit, - next_page=next_page, - user_id=user_id, - redis_client=redis_client, - ) - assert len(result.data) == limit, ( - f"Page {page_num}: expected {limit} items, got {len(result.data)}" - ) - all_items.extend(result.data) - next_page = result.next_page - - if page_num < num_pages: - assert result.has_next_page, ( - f"Page {page_num}: has_next_page=False, got only {len(all_items)} items" - ) - - assert len(all_items) == 900 - unique = set(all_items) - assert len(unique) == 900, ( - f"Found {900 - len(unique)} duplicates among 900 items" - ) - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_dedup_over_session_1500_unique_tours(redis_client) -> None: - """dedup + session_size=300, pool=5000: 1500 unique tours, 0 duplicates. - - session_size=300 with limit=150 gives 2 pages per session. - 10 pages = 5 session rebuilds = 1500 tours. - All must be unique (dedup tracks seen items across sessions). - """ - merger = parse_model(MergerDeduplication, MERGER_DEDUP_LARGE_POOL_CONFIG) - user_id = "dedup_1500" - limit = 150 - num_pages = 10 - - all_items = [] - next_page = FeedResultNextPage(data={}) - - for page_num in range(1, num_pages + 1): - result = await merger.get_data( - methods_dict=METHODS_DICT, - limit=limit, - next_page=next_page, - user_id=user_id, - redis_client=redis_client, - ) - assert len(result.data) == limit, ( - f"Page {page_num}: expected {limit} items, got {len(result.data)}" - ) - all_items.extend(result.data) - next_page = result.next_page - - if page_num < num_pages: - assert result.has_next_page, ( - f"Page {page_num}: has_next_page=False after only {len(all_items)} items" - ) - - assert len(all_items) == 1500 - unique = set(all_items) - assert len(unique) == 1500, ( - f"Found {1500 - len(unique)} duplicates among 1500 items" - ) diff --git a/tests/test_mixers.py b/tests/test_mixers.py new file mode 100644 index 0000000..3f76398 --- /dev/null +++ b/tests/test_mixers.py @@ -0,0 +1,71 @@ +import pytest +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +@pytest.fixture +def ctx(): + return ExecutionContext(session_id="s1", methods_dict=METHODS, redis=None) + + +@pytest.mark.asyncio +async def test_executor_runs_subfeed(ctx): + node = SubFeed(subfeed_id="items", method_name="items") + result = await run_executor.run(node, ctx, limit=5, cursor={}) + assert len(result.data) == 5 + assert result.data[0]["_smartfeed_debug_info"]["source"] == "items" + + +from smartfeed.models.mixers import MergerPercentage, MergerPercentageItem, MergerAppend, MergerPositional + + +@pytest.mark.asyncio +async def test_percentage_40_60(ctx): + node = MergerPercentage( + node_id="mix", + items=[ + MergerPercentageItem( + percentage=40, + data=SubFeed(subfeed_id="a", method_name="items"), + ), + MergerPercentageItem( + percentage=60, + data=SubFeed(subfeed_id="b", method_name="items"), + ), + ], + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(result.data) == 10 + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert sources.count("a") == 4 + assert sources.count("b") == 6 + + +@pytest.mark.asyncio +async def test_append_two_subfeeds(ctx): + node = MergerAppend( + node_id="append", + items=[ + SubFeed(subfeed_id="a", method_name="items"), + SubFeed(subfeed_id="b", method_name="items"), + ], + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(result.data) == 10 + + +@pytest.mark.asyncio +async def test_positional_inserts_at_positions(ctx): + node = MergerPositional( + node_id="pos", + positions=[1, 3], + positional=SubFeed(subfeed_id="promo", method_name="items"), + default=SubFeed(subfeed_id="regular", method_name="items"), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert sources[0] == "promo" # position 1 + assert sources[1] == "regular" + assert sources[2] == "promo" # position 3 diff --git a/tests/test_parsing_config.py b/tests/test_parsing_config.py deleted file mode 100644 index bf61007..0000000 --- a/tests/test_parsing_config.py +++ /dev/null @@ -1,58 +0,0 @@ -import pytest - -from smartfeed.manager import FeedManager -from smartfeed.schemas import ( - FeedConfig, - MergerAppend, - MergerDeduplication, - MergerPercentage, - MergerPercentageGradient, - MergerPercentageItem, - MergerPositional, - MergerViewSession, - SubFeed, -) -from tests.fixtures.configs import METHODS_DICT, PARSING_CONFIG_FIXTURE, PARSING_DEDUP_CONFIG_FIXTURE - - -@pytest.mark.asyncio -async def test_parsing_config() -> None: - """ - Тест для проверки парсинга JSON-файла конфигурации. - """ - - feed_manager = FeedManager(config=PARSING_CONFIG_FIXTURE, methods_dict=METHODS_DICT) - - # Feed Config. - assert isinstance(feed_manager.feed_config, FeedConfig) - # Merger Positional. - assert isinstance(feed_manager.feed_config.feed, MergerPositional) - # Merger Append. - assert isinstance(feed_manager.feed_config.feed.positional, MergerAppend) - # SubFeed with SubFeed Params. - assert isinstance(feed_manager.feed_config.feed.positional.items[0], SubFeed) - # Merger Percentage Gradient. - assert isinstance(feed_manager.feed_config.feed.positional.items[1], MergerPercentageGradient) - # Merger View Session. - assert isinstance(feed_manager.feed_config.feed.positional.items[2], MergerViewSession) - # Merger Percentage. - assert isinstance(feed_manager.feed_config.feed.default, MergerPercentage) - # Merger Percentage Item. - assert isinstance(feed_manager.feed_config.feed.default.items[0], MergerPercentageItem) - # SubFeed without SubFeed Params. - assert isinstance(feed_manager.feed_config.feed.positional.items[1].item_from.data, SubFeed) - assert isinstance(feed_manager.feed_config.feed.positional.items[1].item_to.data, SubFeed) - assert isinstance(feed_manager.feed_config.feed.positional.items[2].data, SubFeed) - # SubFeed with Raise Exception False. - assert isinstance(feed_manager.feed_config.feed.default.items[0].data, SubFeed) - assert feed_manager.feed_config.feed.default.items[0].data.raise_error is False - - -@pytest.mark.asyncio -async def test_parsing_config_deduplication_merger() -> None: - feed_manager = FeedManager(config=PARSING_DEDUP_CONFIG_FIXTURE, methods_dict=METHODS_DICT) - - assert isinstance(feed_manager.feed_config, FeedConfig) - assert isinstance(feed_manager.feed_config.feed, MergerDeduplication) - # Deduplication merger is a wrapper around a single child feed. - assert isinstance(feed_manager.feed_config.feed.data, (MergerPercentage, SubFeed)) diff --git a/tests/test_positional_no_refill_high_priority.py b/tests/test_positional_no_refill_high_priority.py deleted file mode 100644 index 8b8fd81..0000000 --- a/tests/test_positional_no_refill_high_priority.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Positional subfeed with highest dedup_priority must NOT be refilled -when it returns fewer items than requested slots. - -Reproduces production bug: TopSort (promo) returns fewer ads than -positional slots → _refill_deficits hardcodes has_next_page=True -in initial state → pointless refill calls even though the subfeed -already signalled has_next_page=False. -""" - -import pytest - -from smartfeed.schemas import FeedResultNextPage, MergerDeduplication -from tests.fixtures import dedup_helpers as dh -from tests.utils import parse_model - - -def _make_counting_method(items): - """Offset-paged method that also counts how many times it was called.""" - call_count = {"value": 0} - - async def _method(user_id, limit, next_page, **kwargs): - call_count["value"] += 1 - from smartfeed.schemas import FeedResultClient - - offset = int(next_page.after or 0) - result_data = items[offset : offset + limit] - next_page.after = offset + len(result_data) - next_page.page += 1 - has_next_page = (offset + len(result_data)) < len(items) - return FeedResultClient(data=result_data, next_page=next_page, has_next_page=has_next_page) - - return _method, call_count - - -@pytest.mark.asyncio -async def test_positional_high_priority_no_refill_when_exhausted(): - """When the positional subfeed (dedup_priority=1, highest) returns fewer - items than requested AND has_next_page=False, it must NOT be refilled. - - Setup: - - positional (promo): only 2 items, dedup_priority=1 - - default: 100 items, dedup_priority=0 - - positions=[1,3,5,7] → needs 4 promo items - - promo returns 2 out of 4 → has_next_page=False - - Expected: promo called exactly ONCE (no refill) - """ - promo_items = dh.make_items("promo", 1001, 1003) # only 2 items - default_items = dh.make_items("default", 1, 101) # 100 items, no overlap - - promo_method, promo_calls = _make_counting_method(promo_items) - default_method, default_calls = _make_counting_method(default_items) - - methods_dict = { - "promo": promo_method, - "default": default_method, - } - - config = dh._dedup_config( - "dedup_wrapper", - dh._positional_config( - "pos_mix", - positions=[1, 3, 5, 7], - positional=dh._subfeed("sf_promo", "promo", dedup_priority=1), - default=dh._subfeed("sf_default", "default", dedup_priority=0), - ), - max_refill_loops=5, - overfetch_factor=1, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=20, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) > 0 - - # The critical assertion: promo must be called only once. - # Before the fix, it was called 1 + max_refill_loops times - # because _refill_deficits hardcoded has_next_page=True. - assert promo_calls["value"] == 1, ( - f"Promo was called {promo_calls['value']} times, expected 1. " - f"Refill should not retry a subfeed that returned has_next_page=False." - ) - - -@pytest.mark.asyncio -async def test_positional_high_priority_no_refill_even_with_overfetch(): - """Same as above but with overfetch_factor > 1 to ensure the fix - works regardless of overfetch settings.""" - promo_items = dh.make_items("promo", 1001, 1003) # only 2 items - default_items = dh.make_items("default", 1, 101) - - promo_method, promo_calls = _make_counting_method(promo_items) - default_method, _ = _make_counting_method(default_items) - - methods_dict = { - "promo": promo_method, - "default": default_method, - } - - config = dh._dedup_config( - "dedup_wrapper", - dh._positional_config( - "pos_mix", - positions=[1, 3, 5, 7], - positional=dh._subfeed("sf_promo", "promo", dedup_priority=1), - default=dh._subfeed("sf_default", "default", dedup_priority=0), - ), - max_refill_loops=5, - overfetch_factor=5, - ) - - merger = parse_model(MergerDeduplication, config) - res = await merger.get_data( - methods_dict=methods_dict, - user_id="u", - limit=20, - next_page=FeedResultNextPage(data={}), - ) - - assert len(res.data) > 0 - assert promo_calls["value"] == 1, ( - f"Promo was called {promo_calls['value']} times with overfetch_factor=5. " - f"Expected 1 — exhausted subfeed must not be refilled." - ) diff --git a/tests/test_redis_live.py b/tests/test_redis_live.py deleted file mode 100644 index 32f59a7..0000000 --- a/tests/test_redis_live.py +++ /dev/null @@ -1,169 +0,0 @@ -import asyncio -import time - -import pytest -import redis - -from smartfeed.schemas import FeedResultNextPage, MergerViewSession -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG -from tests.utils import parse_model - - -class RedisReplicationSimulator: - """ - Симулятор задержки репликации Redis для тестирования проблемы кластера. - """ - - def __init__(self, real_client): - self.real_client = real_client - self.write_delay = 0.1 # Задержка для имитации репликации - self.pending_writes = {} # Ключи которые только что записали - - def exists(self, cache_key): - return self.real_client.exists(cache_key) - - def set(self, name, value, ex=None): - # Записываем в реальный Redis - result = self.real_client.set(name, value, ex=ex) - # Помечаем что этот ключ только что записан (имитация репликации) - self.pending_writes[name] = time.time() - return result - - def get(self, name): - # Если ключ только что записан (в течение write_delay секунд), возвращаем None - if name in self.pending_writes: - write_time = self.pending_writes[name] - if time.time() - write_time < self.write_delay: - return None # Имитация задержки репликации - else: - # Задержка прошла, можно удалить из pending - del self.pending_writes[name] - - # Обычное чтение из Redis - return self.real_client.get(name) - - -@pytest.mark.asyncio -async def test_redis_replication_delay_problem(): - """ - Тест для воспроизведения проблемы репликации Redis с использованием - RedisReplicationSimulator для имитации задержки. - """ - - # Подключаемся к Redis (должен быть запущен локально) - try: - real_client = redis.Redis(host="localhost", port=6379, db=0) - real_client.ping() # Проверяем соединение - except (redis.ConnectionError, redis.ResponseError): - pytest.skip("Redis not available for live testing") - - # Очищаем тестовый ключ - test_key = "test_merger_view_session_test_user" - real_client.delete(test_key) - - # Используем симулятор задержки репликации - redis_client = RedisReplicationSimulator(real_client) - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - - print("\n=== Демонстрация проблемы с задержкой репликации ===") - - try: - # Этот вызов должен воспроизвести проблему с оригинальным кодом - result = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=10, - next_page=FeedResultNextPage(data={}), - user_id="test_user", - redis_client=redis_client, - ) - - print("✅ Исправление работает! Получили результат без ошибки:") - print(f" Данные: {result.data[:5]}... (показаны первые 5)") - print(f" Размер: {len(result.data)}") - print(f" Есть следующая страница: {result.has_next_page}") - - # Проверяем что получили валидные данные - assert len(result.data) == 10 - assert result.data[0] == "test_user_1" - assert result.has_next_page is True - - except TypeError as e: - if "the JSON object must be str, bytes or bytearray, not NoneType" in str(e): - print("❌ Проблема НЕ исправлена! Все еще получаем TypeError") - raise - else: - print(f"❓ Неожиданная ошибка: {e}") - raise - - finally: - # Очистка - real_client.delete(test_key) - real_client.close() - - -@pytest.mark.asyncio -async def test_redis_multiple_requests(): - """ - Тест множественных запросов для проверки стабильности исправления. - """ - - try: - real_client = redis.Redis(host="localhost", port=6379, db=0) - real_client.ping() - except (redis.ConnectionError, redis.ResponseError): - pytest.skip("Redis not available for live testing") - - test_key = "test_merger_multiple_test_user" - real_client.delete(test_key) - - redis_client = RedisReplicationSimulator(real_client) - merger_vs = parse_model(MergerViewSession, MERGER_VIEW_SESSION_CONFIG) - - print("\n=== Тест множественных запросов ===") - - try: - # Первый запрос - создает кэш - result1 = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=5, - next_page=FeedResultNextPage(data={}), - user_id="test_user", - redis_client=redis_client, - ) - - print(f"Первый запрос: получили {len(result1.data)} элементов") - - # Ждем чтобы задержка репликации прошла - await asyncio.sleep(0.2) - - # Второй запрос - должен использовать кэш - from smartfeed.schemas import FeedResultNextPageInside - - result2 = await merger_vs.get_data( - methods_dict=METHODS_DICT, - limit=5, - next_page=FeedResultNextPage( - data={"merger_view_session_example": FeedResultNextPageInside(page=2, after=None)} - ), - user_id="test_user", - redis_client=redis_client, - ) - - print(f"Второй запрос: получили {len(result2.data)} элементов") - print(f"Данные второй страницы: {result2.data}") - - # Проверяем что получили разные данные (пагинация работает) - assert result1.data != result2.data - assert len(result2.data) == 5 - - print("✅ Множественные запросы работают корректно!") - - finally: - real_client.delete(test_key) - real_client.close() - - -if __name__ == "__main__": - # Для запуска напрямую без pytest - asyncio.run(test_redis_replication_delay_problem()) diff --git a/tests/test_redis_lock.py b/tests/test_redis_lock.py new file mode 100644 index 0000000..fdc1e15 --- /dev/null +++ b/tests/test_redis_lock.py @@ -0,0 +1,33 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.execution.redis_lock import RedisLock + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.mark.asyncio +async def test_lock_acquired(redis): + async with RedisLock(redis, "test:lock") as acquired: + assert acquired is True + assert await redis.exists("test:lock") + assert not await redis.exists("test:lock") + + +@pytest.mark.asyncio +async def test_lock_contention(redis): + async with RedisLock(redis, "test:lock") as first: + assert first is True + async with RedisLock(redis, "test:lock") as second: + assert second is False + + +@pytest.mark.asyncio +async def test_lock_released_after_first(redis): + async with RedisLock(redis, "test:lock"): + pass + async with RedisLock(redis, "test:lock") as acquired: + assert acquired is True diff --git a/tests/test_resilience.py b/tests/test_resilience.py new file mode 100644 index 0000000..1707a33 --- /dev/null +++ b/tests/test_resilience.py @@ -0,0 +1,250 @@ +"""Tests for resilience: source failures, positional accuracy, rerank error handling.""" + +import pytest +import fakeredis.aioredis + +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperRerank, WrapperDedup +from smartfeed.models.mixers import MergerPercentage, MergerPercentageItem, MergerPositional +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +async def make_source(prefix, user_id, limit, next_page, **kwargs): + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": f"{prefix}_{start + i}", "val": prefix} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +async def make_failing_source(user_id, limit, next_page, **kwargs): + raise RuntimeError("source unavailable") + + +async def make_promo(user_id, limit, next_page, **kwargs): + data = [{"id": f"promo_{i}", "val": "promo"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def rerank_reverse(items, session_id): + return list(reversed(items)) + + +async def rerank_crash(items, session_id): + raise RuntimeError("ML service down") + + +METHODS = { + "source_a": lambda *a, **kw: make_source("a", *a, **kw), + "source_b": lambda *a, **kw: make_source("b", *a, **kw), + "failing": make_failing_source, + "promo": make_promo, + "rerank_ok": rerank_reverse, + "rerank_crash": rerank_crash, +} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="test", methods_dict=METHODS, redis=redis) + + +# --------------------------------------------------------------------------- +# 1. Source failure: fill page from remaining source +# --------------------------------------------------------------------------- + +class TestSourceFailureFallback: + + @pytest.mark.asyncio + async def test_one_source_fails_page_filled_from_other(self, ctx): + """If one source in percentage mix fails, the other fills the page.""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + data=MergerPercentage( + node_id="mix", + items=[ + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="good", method_name="source_a"), + ), + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="bad", method_name="failing", raise_error=False), + ), + ], + ), + ) + result = await run_executor.run(node, ctx, limit=20, cursor={}) + # Should have items from source_a only (source_b failed gracefully) + assert len(result.data) > 0 + for item in result.data: + assert item["val"] == "a" + + @pytest.mark.asyncio + async def test_all_sources_fail_returns_empty(self, ctx): + """If all sources fail, return empty gracefully.""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + data=MergerPercentage( + node_id="mix", + items=[ + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="bad1", method_name="failing", raise_error=False), + ), + MergerPercentageItem( + percentage=50, + data=SubFeed(subfeed_id="bad2", method_name="failing", raise_error=False), + ), + ], + ), + ) + result = await run_executor.run(node, ctx, limit=20, cursor={}) + assert result.data == [] + + +# --------------------------------------------------------------------------- +# 2. Positional: strict position enforcement +# --------------------------------------------------------------------------- + +class TestPositionalStrictPositions: + + @pytest.mark.asyncio + async def test_promo_at_exact_positions(self, ctx): + """Promo items must be at positions 1, 3, 5, 7 (1-indexed).""" + node = MergerPositional( + node_id="pos", + positions=[1, 3, 5, 7], + positional=SubFeed(subfeed_id="promo", method_name="promo"), + default=SubFeed(subfeed_id="regular", method_name="source_a"), + ) + result = await run_executor.run(node, ctx, limit=20, cursor={}) + assert len(result.data) == 20 + + for pos in [1, 3, 5, 7]: + item = result.data[pos - 1] # 1-indexed -> 0-indexed + source = item.get("_smartfeed_debug_info", {}).get("source", "") + assert source == "promo", ( + f"Position {pos}: expected promo, got {source} (id={item.get('id')})" + ) + + @pytest.mark.asyncio + async def test_non_promo_positions_filled_by_default(self, ctx): + """Non-promo positions must be filled by default source.""" + node = MergerPositional( + node_id="pos", + positions=[1, 3, 5, 7], + positional=SubFeed(subfeed_id="promo", method_name="promo"), + default=SubFeed(subfeed_id="regular", method_name="source_a"), + ) + result = await run_executor.run(node, ctx, limit=20, cursor={}) + + non_promo_positions = [i for i in range(20) if (i + 1) not in [1, 3, 5, 7]] + for pos in non_promo_positions: + item = result.data[pos] + source = item.get("_smartfeed_debug_info", {}).get("source", "") + assert source == "regular", ( + f"Position {pos + 1}: expected regular, got {source}" + ) + + @pytest.mark.asyncio + async def test_positional_with_dedup_preserves_positions(self, ctx): + """Even with outer dedup, promo positions must be preserved.""" + node = Wrapper( + node_id="dedup_outer", + dedup=WrapperDedup(dedup_key="id", overfetch_factor=3), + data=MergerPositional( + node_id="pos", + positions=[1, 3, 5, 7], + positional=SubFeed(subfeed_id="promo", method_name="promo", dedup_priority=10), + default=SubFeed(subfeed_id="regular", method_name="source_a", dedup_priority=1), + ), + ) + result = await run_executor.run(node, ctx, limit=20, cursor={}) + assert len(result.data) == 20 + + for pos in [1, 3, 5, 7]: + item = result.data[pos - 1] + source = item.get("_smartfeed_debug_info", {}).get("source", "") + assert source == "promo", ( + f"Position {pos}: expected promo after dedup, got {source}" + ) + + +# --------------------------------------------------------------------------- +# 3. Rerank failure: raise_error=True vs False +# --------------------------------------------------------------------------- + +class TestRerankFailureHandling: + + @pytest.mark.asyncio + async def test_rerank_crash_raise_error_true(self, ctx): + """With raise_error=True (default), rerank failure crashes the request.""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_crash", raise_error=True), + data=SubFeed(subfeed_id="src", method_name="source_a"), + ) + with pytest.raises(RuntimeError, match="ML service down"): + await run_executor.run(node, ctx, limit=10, cursor={}) + + @pytest.mark.asyncio + async def test_rerank_crash_raise_error_false(self, ctx): + """With raise_error=False, rerank failure keeps original order.""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_crash", raise_error=False), + data=SubFeed(subfeed_id="src", method_name="source_a"), + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(result.data) == 10 + # Items should be in original order (not reversed, since rerank_crash failed) + ids = [item["id"] for item in result.data] + assert ids[0] == "a_0" # original order preserved + + @pytest.mark.asyncio + async def test_rerank_ok_actually_reranks(self, ctx): + """Verify working rerank changes order (control test).""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_ok"), + data=SubFeed(subfeed_id="src", method_name="source_a"), + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(result.data) == 10 + # rerank_ok reverses, so last item of original is first + ids = [item["id"] for item in result.data] + assert ids[0] != "a_0" # order changed + + @pytest.mark.asyncio + async def test_rerank_fail_soft_page2_still_works(self, ctx): + """After fail-soft rerank, page 2 from cache still works.""" + node = Wrapper( + node_id="w", + cache=WrapperCache(session_size=50, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_crash", raise_error=False), + data=SubFeed(subfeed_id="src", method_name="source_a"), + ) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + assert len(r1.data) == 10 + assert len(r2.data) == 10 + # Different pages + ids1 = {item["id"] for item in r1.data} + ids2 = {item["id"] for item in r2.data} + assert not ids1 & ids2 # no overlap diff --git a/tests/test_seen_store_unit.py b/tests/test_seen_store_unit.py deleted file mode 100644 index 2ae77c3..0000000 --- a/tests/test_seen_store_unit.py +++ /dev/null @@ -1,64 +0,0 @@ - -import pytest - -from smartfeed.feed_models import _redis_call -from smartfeed.policies.dedup_utils import decode_seen_from_cursor -from smartfeed.policies.seen_store import CursorSeenStore, RedisSeenStore -from tests.fixtures.redis import redis_client # noqa: F401 - - -@pytest.mark.asyncio -async def test_cursor_seen_store_set_max_and_commit_roundtrip() -> None: - store = CursorSeenStore.from_after(after=None, cursor_compress=True, cursor_max_keys=None) - store.set_max("a", 1) - store.set_max("a", 1) # no-op - store.set_max("a", 0) # no-op (lower) - store.set_max("b", 2) - - after = await store.commit() - decoded = decode_seen_from_cursor(after) - assert decoded == {"a": 1, "b": 2} - - -@pytest.mark.asyncio -async def test_cursor_seen_store_commit_keeps_previous_cursor_state() -> None: - store = CursorSeenStore.from_after( - after={"v": 2, "seen": [["a", 1], ["b", 2]]}, - cursor_compress=False, - cursor_max_keys=None, - ) - store.set_max("c", 3) - - after = await store.commit() - decoded = decode_seen_from_cursor(after) - assert decoded == {"a": 1, "b": 2, "c": 3} - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_redis_seen_store_prefetch_set_max_commit_and_reset(redis_client) -> None: - key = "test_seen_store" - await _redis_call(redis_client, "delete", key) - # Pre-seed zset state - await _redis_call(redis_client, "zadd", key, mapping={"a": 5.0}) - - store = RedisSeenStore.create(redis_client=redis_client, redis_key=key, ttl_seconds=60) - - await store.prefetch(["a", "a", "b"]) # duplicates - assert store.get("a") == 5 - assert store.get("b") is None - - store.set_max("a", 3) # should not reduce existing - store.set_max("b", 2) - - await store.commit() - - # New state should be present in redis - scores = list(await _redis_call(redis_client, "zmscore", key, ["a", "b"])) - assert scores == [5.0, 2.0] - - await store.reset() - scores_after_reset = list(await _redis_call(redis_client, "zmscore", key, ["a", "b"])) - assert scores_after_reset == [None, None] - - await _redis_call(redis_client, "delete", key) diff --git a/tests/test_sub_feed.py b/tests/test_sub_feed.py deleted file mode 100644 index 11645d4..0000000 --- a/tests/test_sub_feed.py +++ /dev/null @@ -1,78 +0,0 @@ -import pytest - -from smartfeed.schemas import FeedResultNextPage, SubFeed -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.subfeeds import ( - SUBFEED_CONFIG, - SUBFEED_CONFIG_NO_RAISE_ERROR, - SUBFEED_CONFIG_RAISE_ERROR, - SUBFEED_WITH_PARAMS_CONFIG, -) - - -@pytest.mark.asyncio -async def test_sub_feed() -> None: - """ - Тест для проверки получения данных из субфида (без параметров). - """ - - sub_feed = SubFeed.model_validate(SUBFEED_CONFIG) - sub_feed_data = await sub_feed.get_data( - methods_dict=METHODS_DICT, - limit=15, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - assert sub_feed_data.data == [f"x_{i}" for i in range(1, 16)] - - -@pytest.mark.asyncio -async def test_sub_feed_with_params() -> None: - """ - Тест для проверки получения данных из субфида (с параметрами). - """ - - sub_feed = SubFeed.model_validate(SUBFEED_WITH_PARAMS_CONFIG) - sub_feed_data = await sub_feed.get_data( - methods_dict=METHODS_DICT, - limit=15, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - assert sub_feed_data.data == [f"x_{i}" for i in range(1, 11)] - - -@pytest.mark.asyncio -async def test_sub_feed_raise_error() -> None: - """ - Тест для проверки получения данных из субфида (без параметров). - """ - - sub_feed = SubFeed.model_validate(SUBFEED_CONFIG_RAISE_ERROR) - - with pytest.raises(Exception): - await sub_feed.get_data( - methods_dict=METHODS_DICT, - limit=15, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - -@pytest.mark.asyncio -async def test_sub_feed_no_raise_error() -> None: - """ - Тест для проверки получения данных из субфида (без параметров). - """ - - sub_feed = SubFeed.model_validate(SUBFEED_CONFIG_NO_RAISE_ERROR) - sub_feed_data = await sub_feed.get_data( - methods_dict=METHODS_DICT, - limit=15, - next_page=FeedResultNextPage(data={}), - user_id="x", - ) - - assert sub_feed_data.data == [] diff --git a/tests/test_subfeed.py b/tests/test_subfeed.py new file mode 100644 index 0000000..1097dd0 --- /dev/null +++ b/tests/test_subfeed.py @@ -0,0 +1,51 @@ +import pytest +from smartfeed.models.subfeed import SubFeed +from tests.conftest import METHODS + + +@pytest.mark.asyncio +async def test_subfeed_returns_data(): + sf = SubFeed(subfeed_id="items", method_name="items") + result = await sf.execute( + methods_dict=METHODS, session_id="s1", limit=5, cursor={} + ) + assert len(result.data) == 5 + assert result.has_next_page is True + + +@pytest.mark.asyncio +async def test_subfeed_stamps_debug_info(): + sf = SubFeed(subfeed_id="my_source", method_name="items") + result = await sf.execute( + methods_dict=METHODS, session_id="s1", limit=3, cursor={} + ) + for i, item in enumerate(result.data): + info = item["_smartfeed_debug_info"] + assert info["source"] == "my_source" + + +@pytest.mark.asyncio +async def test_subfeed_raise_error(): + sf = SubFeed(subfeed_id="err", method_name="error") + with pytest.raises(RuntimeError, match="subfeed error"): + await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor={}) + + +@pytest.mark.asyncio +async def test_subfeed_swallow_error(): + sf = SubFeed(subfeed_id="err", method_name="error", raise_error=False) + result = await sf.execute( + methods_dict=METHODS, session_id="s1", limit=5, cursor={} + ) + assert result.data == [] + assert result.has_next_page is False + + +@pytest.mark.asyncio +async def test_subfeed_cursor_propagation(): + sf = SubFeed(subfeed_id="items", method_name="items") + r1 = await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor={}) + r2 = await sf.execute( + methods_dict=METHODS, session_id="s1", limit=5, cursor=r1.next_page + ) + assert r2.data[0]["id"] == 5 # page 2 starts at id=5 diff --git a/tests/test_view_session_unit.py b/tests/test_view_session_unit.py deleted file mode 100644 index d237185..0000000 --- a/tests/test_view_session_unit.py +++ /dev/null @@ -1,53 +0,0 @@ -from dataclasses import dataclass - -import pytest - -from smartfeed.feed_models import _redis_call -from smartfeed.schemas import FeedResultNextPage, MergerViewSession -from tests.fixtures.configs import METHODS_DICT -from tests.fixtures.mergers import MERGER_VIEW_SESSION_CONFIG -from tests.fixtures.redis import redis_client # noqa: F401 -from tests.utils import parse_model - - -@dataclass -class _ItemWithAttr: - id: str - - -def test_get_dedup_key_supports_dict_and_attr_and_raises_on_missing() -> None: - cfg = dict(MERGER_VIEW_SESSION_CONFIG) - cfg.update({"deduplicate": True, "dedup_key": "id"}) - merger = parse_model(MergerViewSession, cfg) - - assert merger._get_dedup_key_or_attr({"id": "x"}) == "x" - assert merger._get_dedup_key_or_attr(_ItemWithAttr(id="y")) == "y" - - with pytest.raises(AssertionError): - merger._get_dedup_key_or_attr({"nope": 1}) - - -@pytest.mark.parametrize("redis_client", ["sync", "async"], indirect=True) -@pytest.mark.asyncio -async def test_view_session_shuffle_applies_to_result(redis_client, monkeypatch) -> None: - import smartfeed.mergers.view_session as vs_mod - - cfg = dict(MERGER_VIEW_SESSION_CONFIG) - cfg.update({"shuffle": True}) - merger = parse_model(MergerViewSession, cfg) - cache_key = f"{merger.merger_id}_x" - await _redis_call(redis_client, "delete", cache_key) - - # Make shuffle deterministic: reverse in-place - monkeypatch.setattr(vs_mod, "shuffle", lambda data: data.reverse()) - - res = await merger.get_data( - methods_dict=METHODS_DICT, - limit=5, - next_page=FeedResultNextPage(data={}), - user_id="x", - redis_client=redis_client, - ) - - assert res.data == ["x_5", "x_4", "x_3", "x_2", "x_1"] - await _redis_call(redis_client, "delete", cache_key) diff --git a/tests/test_wrapper_cache.py b/tests/test_wrapper_cache.py new file mode 100644 index 0000000..6933bbd --- /dev/null +++ b/tests/test_wrapper_cache.py @@ -0,0 +1,77 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models.wrapper import Wrapper, WrapperCache +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=METHODS, redis=redis) + + +def _make_wrapper(node_id="cached", session_size=50, session_ttl=300): + return Wrapper( + node_id=node_id, + cache=WrapperCache(session_size=session_size, session_ttl=session_ttl), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + + +@pytest.mark.asyncio +async def test_cache_cold_build(ctx, redis): + node = _make_wrapper() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(result.data) == 10 + # Some Redis key should exist + all_keys = list(await redis.keys("sf:*")) + assert len(all_keys) >= 1 + + +@pytest.mark.asyncio +async def test_cache_warm_hit(ctx): + node = _make_wrapper() + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + # Different page + assert r2.data[0]["id"] != r1.data[0]["id"] + assert len(r2.data) == 10 + + +@pytest.mark.asyncio +async def test_generation_id_in_cursor(ctx): + node = _make_wrapper() + result = await run_executor.run(node, ctx, limit=10, cursor={}) + assert "gen" in result.next_page.get("cached", {}) + + +@pytest.mark.asyncio +async def test_stale_gen_resets_to_page1(ctx, redis): + node = _make_wrapper() + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + # Flush to simulate TTL expiry + await redis.flushall() + stale_cursor = {"cached": {"page": 5, "gen": "stale_nonce"}} + r2 = await run_executor.run(node, ctx, limit=10, cursor=stale_cursor) + # Should reset: page 2 in next cursor (just served page 1) + assert r2.next_page["cached"]["page"] == 2 + assert r2.next_page["cached"]["gen"] != "stale_nonce" + + +@pytest.mark.asyncio +async def test_wrapper_without_cache_passthrough(ctx): + """Wrapper without cache passes through to child directly.""" + node = Wrapper( + node_id="passthrough", + data=SubFeed(subfeed_id="items", method_name="items"), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + assert len(result.data) == 5 diff --git a/tests/test_wrapper_dedup.py b/tests/test_wrapper_dedup.py new file mode 100644 index 0000000..22fe0cd --- /dev/null +++ b/tests/test_wrapper_dedup.py @@ -0,0 +1,86 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models.wrapper import Wrapper, WrapperDedup +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.mixers import MergerAppend +from smartfeed.models.base import FeedResult +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +async def make_dupes_a(user_id, limit, next_page, **kw): + data = [{"id": i, "val": f"a_{i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def make_dupes_b(user_id, limit, next_page, **kw): + data = [{"id": i, "val": f"b_{i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +DEDUP_METHODS = {"dupes_a": make_dupes_a, "dupes_b": make_dupes_b} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=DEDUP_METHODS, redis=redis) + + +@pytest.mark.asyncio +async def test_dedup_removes_duplicates(ctx): + node = Wrapper( + node_id="deduped", + dedup=WrapperDedup(dedup_key="id", overfetch_factor=2), + data=MergerAppend( + node_id="mix", + items=[ + SubFeed(subfeed_id="a", method_name="dupes_a"), + SubFeed(subfeed_id="b", method_name="dupes_b"), + ], + ), + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + ids = [item["id"] for item in result.data] + assert len(ids) == len(set(ids)) + + +@pytest.mark.asyncio +async def test_dedup_priority_higher_wins(ctx): + node = Wrapper( + node_id="deduped", + dedup=WrapperDedup(dedup_key="id"), + data=MergerAppend( + node_id="mix", + items=[ + SubFeed(subfeed_id="a", method_name="dupes_a", dedup_priority=5), + SubFeed(subfeed_id="b", method_name="dupes_b", dedup_priority=1), + ], + ), + ) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + for item in result.data: + assert item["val"].startswith("a_") + + +@pytest.mark.asyncio +async def test_dedup_equal_priority_first_seen_wins(ctx): + node = Wrapper( + node_id="deduped", + dedup=WrapperDedup(dedup_key="id"), + data=MergerAppend( + node_id="mix", + items=[ + SubFeed(subfeed_id="a", method_name="dupes_a", dedup_priority=1), + SubFeed(subfeed_id="b", method_name="dupes_b", dedup_priority=1), + ], + ), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + for item in result.data: + assert item["val"].startswith("a_") diff --git a/tests/test_wrapper_full_pipeline.py b/tests/test_wrapper_full_pipeline.py new file mode 100644 index 0000000..492f4b9 --- /dev/null +++ b/tests/test_wrapper_full_pipeline.py @@ -0,0 +1,72 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models import FeedConfig +from smartfeed.models.base import FeedResult +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +async def source_a(user_id, limit, next_page, **kw): + data = [{"id": i, "val": f"a_{i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def source_b(user_id, limit, next_page, **kw): + data = [{"id": i + limit // 2, "val": f"b_{i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def reverse_rerank(items, session_id): + return list(reversed(items)) + + +PIPELINE_METHODS = {"source_a": source_a, "source_b": source_b, "reverse": reverse_rerank} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="test_full", methods_dict=PIPELINE_METHODS, redis=redis) + + +@pytest.mark.asyncio +async def test_full_pipeline(ctx): + config = FeedConfig.parse_obj({ + "version": "2", + "feed": { + "type": "wrapper", "node_id": "full", + "cache": {"session_size": 30, "session_ttl": 300}, + "rerank": {"method_name": "reverse"}, + "dedup": {"dedup_key": "id", "overfetch_factor": 2}, + "data": { + "type": "merger_percentage", "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "source_a", "dedup_priority": 5}}, + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "source_b", "dedup_priority": 1}}, + ], + }, + }, + }) + + r1 = await run_executor.run(config.feed, ctx, limit=10, cursor={}) + + # No duplicate ids + ids = [item["id"] for item in r1.data] + assert len(ids) == len(set(ids)) + + # Has generation_id + assert "gen" in r1.next_page.get("full", {}) + + # Has debug info + for item in r1.data: + info = item["_smartfeed_debug_info"] + assert "source" in info + + # Page 2 from cache + r2 = await run_executor.run(config.feed, ctx, limit=10, cursor=r1.next_page) + assert len(r2.data) > 0 diff --git a/tests/test_wrapper_rerank.py b/tests/test_wrapper_rerank.py new file mode 100644 index 0000000..83608b6 --- /dev/null +++ b/tests/test_wrapper_rerank.py @@ -0,0 +1,83 @@ +import pytest +import fakeredis.aioredis + +from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperRerank +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +async def reverse_rerank(items, session_id): + return list(reversed(items)) + + +async def bad_rerank(items, session_id): + return items[:1] # contract violation + + +RERANK_METHODS = {**METHODS, "reverse": reverse_rerank, "bad": bad_rerank} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=RERANK_METHODS, redis=redis) + + +@pytest.mark.asyncio +async def test_rerank_with_cache(ctx): + node = Wrapper( + node_id="reranked", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="reverse"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + assert len(result.data) == 5 + # Rerank reversed the order: first item should have highest original id + id_first = result.data[0]["id"] + id_last = result.data[-1]["id"] + assert id_first > id_last # reversed = descending ids + + +@pytest.mark.asyncio +async def test_rerank_without_cache(ctx): + node = Wrapper( + node_id="reranked", + rerank=WrapperRerank(method_name="reverse"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + assert len(result.data) == 5 + + +@pytest.mark.asyncio +async def test_rerank_contract_violation(ctx): + node = Wrapper( + node_id="bad", + cache=WrapperCache(session_size=10, session_ttl=300), + rerank=WrapperRerank(method_name="bad"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + with pytest.raises(ValueError, match="must return exactly"): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +@pytest.mark.asyncio +async def test_rerank_stamps_rerank_position(ctx): + node = Wrapper( + node_id="reranked", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="reverse"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + result = await run_executor.run(node, ctx, limit=5, cursor={}) + for item in result.data: + info = item["_smartfeed_debug_info"]["reranked"] + assert "rerank_position" in info + assert "smartfeed_position" in info diff --git a/tests/test_wrapper_shared_cache.py b/tests/test_wrapper_shared_cache.py new file mode 100644 index 0000000..5d707ee --- /dev/null +++ b/tests/test_wrapper_shared_cache.py @@ -0,0 +1,96 @@ +import pytest +import fakeredis.aioredis +from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperRerank +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor +from tests.conftest import METHODS + + +async def rerank_reverse(items, session_id): + return list(reversed(items)) + +async def rerank_identity(items, session_id): + return items + +SHARED_METHODS = {**METHODS, "rerank_a": rerank_reverse, "rerank_b": rerank_identity} + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="s1", methods_dict=SHARED_METHODS, redis=redis) + +@pytest.mark.asyncio +async def test_shared_cache_two_rerankers(ctx): + w_a = Wrapper( + node_id="a", cache_key="shared", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_a"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + w_b = Wrapper( + node_id="b", cache_key="shared", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_b"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + r_a = await run_executor.run(w_a, ctx, limit=5, cursor={}) + r_b = await run_executor.run(w_b, ctx, limit=5, cursor={}) + assert len(r_a.data) == 5 + assert len(r_b.data) == 5 + # Different rerank -> different order + assert r_a.data != r_b.data + + +@pytest.mark.asyncio +async def test_shared_cache_base_written_once(ctx, redis): + """Both wrappers use same cache_key -- base cache key should exist once.""" + w_a = Wrapper( + node_id="a", cache_key="shared_base", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_a"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + w_b = Wrapper( + node_id="b", cache_key="shared_base", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_b"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + await run_executor.run(w_a, ctx, limit=5, cursor={}) + await run_executor.run(w_b, ctx, limit=5, cursor={}) + + # Each wrapper writes its own reranked cache under node_id + # But both share the base key (cache_key="shared_base") + all_keys = [k.decode() if isinstance(k, bytes) else k for k in await redis.keys("sf:*")] + # The shared base key should exist + base_keys = [k for k in all_keys if ":shared_base:" in k and not k.endswith(":a:") and not k.endswith(":b:")] + assert len(base_keys) >= 1 + + +@pytest.mark.asyncio +async def test_shared_cache_warm_read(ctx, redis): + """After cold build, reading warm cache returns same base data, different rerank.""" + w_a = Wrapper( + node_id="a", cache_key="shared_warm", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_a"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + w_b = Wrapper( + node_id="b", cache_key="shared_warm", + cache=WrapperCache(session_size=20, session_ttl=300), + rerank=WrapperRerank(method_name="rerank_b"), + data=SubFeed(subfeed_id="items", method_name="items"), + ) + r_a1 = await run_executor.run(w_a, ctx, limit=5, cursor={}) + # Second call with cursor (warm path for w_a) + r_a2 = await run_executor.run(w_a, ctx, limit=5, cursor=r_a1.next_page) + assert len(r_a2.data) == 5 + + # w_b also reads the shared base -- should work fine + r_b1 = await run_executor.run(w_b, ctx, limit=5, cursor={}) + assert len(r_b1.data) == 5 diff --git a/tests/utils.py b/tests/utils.py deleted file mode 100644 index af29848..0000000 --- a/tests/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -from __future__ import annotations - -from smartfeed.pydantic_compat import parse_model - -__all__ = ["parse_model"] From c20a324511cf6c92a1393c245ecf82c5b83c956d Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Wed, 8 Apr 2026 18:42:17 +0300 Subject: [PATCH 35/41] WIP --- ARCHITECTURE.md | 304 +++++++++++------------------ CHANGELOG.md | 59 +++++- README.md | 369 +++++++++++++++++------------------- smartfeed/models/wrapper.py | 31 ++- 4 files changed, 364 insertions(+), 399 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index a9ba69c..4823ad9 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -1,238 +1,150 @@ -# SmartFeed Architecture (medium-brief) +# SmartFeed Architecture -## 1) What SmartFeed does +## Overview -SmartFeed builds one paginated feed from multiple client-provided sources (“subfeeds”) using a declarative tree config: +SmartFeed builds a paginated feed from multiple async data sources using a tree config. -- **Leaf**: `SubFeed` (calls one client method) -- **Mergers**: compose children (`append`, `distribute`, `positional`, `percentage`, `percentage_gradient`, `view_session`) -- **Wrapper**: `MergerDeduplication` (changes execution semantics around one child) +``` +FeedManager.get_feed(session_id, limit, cursor) + -> executor.run(root_node, ctx, limit, cursor) + -> tree of nodes execute in parallel + -> stamp smartfeed_position on results + -> return FeedResult +``` -Core runtime: +## Node Types -- parse config -> create request `ExecutionContext` -> run tree via shared `Executor` -> return `FeedResult` + `next_page`. +Three types of nodes: +``` +SubFeed (leaf) -- calls async method from methods_dict +Wrapper (pipeline) -- fetch -> dedup -> rerank -> cache -> paginate +Mixer (coordinator) -- runs children in parallel, assembles results +``` -## 2) Public surfaces and core data types +## Execution Flow -### Public entrypoint +``` + FeedManager + | + executor.run() + | + +---------+---------+ + | | + node.execute() node.build_mix_plan() + (SubFeed,Wrapper) (Mixers) + | + asyncio.gather(children) + | + plan.assemble() +``` -- `FeedManager(config, methods_dict, redis_client=None)` - - `get_data(user_id, limit, next_page, **params) -> FeedResult` +### executor.py -`methods_dict` maps config `method_name` strings to host-app callables. +Module-level functions (no class): -### Config schema surface +- `run(node, ctx, limit, cursor)` -- dispatches to `node.execute()` or `node.build_mix_plan()` +- `_execute_mix(plan, ctx, cursor)` -- runs MixPlan children in parallel via `asyncio.gather` -`smartfeed.schemas` keeps stable imports for: +### Wrapper Pipeline -- `FeedConfig`: top-level model (`version`, `feed`) -- `FeedTypes`: discriminated union by `type` +``` +fetch child (session_size or limit) + -> dedup (priority arbitration, seen-set from Redis) + -> rerank (call methods_dict callable) + -> cache (write to Redis with TTL) + -> paginate (slice by page from cursor) +``` -### Cursor / pagination models +Two paths: +- **Cached** (`self.cache` set): warm = read cache + paginate. Cold = full pipeline. +- **Passthrough** (no cache): fetch + dedup (with Redis seen-set) + rerank per page. -- `FeedResultNextPageInside`: one node cursor (`page`, `after`) -- `FeedResultNextPage`: full-tree cursor map (`data: {node_id -> FeedResultNextPageInside}`) +### MixPlan -### Result models +Mixers return `MixPlan(children, assemble)`: -- `FeedResultClient`: required return type of client subfeed methods -- `FeedResult`: normalized return type of any SmartFeed node +```python +MixChild(node_id="mix_0", node=subfeed_a, demand=8) # ask for 8 items +MixChild(node_id="mix_1", node=subfeed_b, demand=12) # ask for 12 items +``` +Executor runs all children via `asyncio.gather`, passes results to `assemble(buffers, child_cursors)`. -## 3) Node interface contract +## Key Models -All nodes inherit `BaseFeedConfigModel` and are executed through: +### Config (Pydantic v2) -- `get_data(methods_dict, user_id, limit, next_page, redis_client=None, ctx=None, **params) -> FeedResult` +``` +FeedConfig(version, feed: FeedNode) +FeedNode = Union[SubFeed, Wrapper, MergerPercentage, MergerPositional, ...] +``` -Important notes: +Discriminated union on `type` field. Forward refs rebuilt at import time. -- If a node implements `build_plan(...)`, executor uses the plan path. -- Base `get_data(...)` delegates back to executor and expects `build_plan(...)` to exist. -- Every node has `dedup_priority: int` (used by dedup arbitration/refill ordering). +### Runtime +``` +ExecutionContext(session_id, methods_dict, redis) +FeedResult(data: list, next_page: dict, has_next_page: bool) +``` -## 4) ExecutionContext +### Output -`ExecutionContext` is per-request state propagated through the tree: +``` +SmartFeedDebugInfo(source, smartfeed_position, rerank_position?, rrf_score?, ...) +FeedItem(id, _smartfeed_debug_info: SmartFeedDebugInfo) +``` -- `methods_dict`, `user_id`, `redis_client` -- `executor` (lazy via `ensure_executor()`) -- optional policy/settings: - - `dedup`: `DeduplicationPolicy` when dedup wrapper is active - - `refill_settings`: `RefillExecutionSettings(overfetch_factor, max_refill_loops)` +## Redis State -Responsibilities: +``` +sf:{session_id}:{node_id}:{config_hash} -- cached session data (JSON list) +sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (gen, child_cursor, child_has_next) +sf:{session_id}:{node_id}:{config_hash}:seen -- dedup seen-set (Redis SET) +sf:{session_id}:{cache_key}:{hash} -- shared base cache (for A/B testing) +sf:{session_id}:{node_id}:{hash}:lock -- distributed lock (SETNX) +``` -- centralize shared plumbing (executor + redis client) -- keep execution policies out of user params +TTL = inactivity timeout. Refreshed on every access via pipeline EXPIRE. +## Dedup -## 5) Executor (runtime engine) +Two dedup paths, unified `_dedup(data, seen=None)` method: -Primary entry: +1. **Cached path** (`_fetch_and_dedup`): overfetch + refill loop. Re-dedup combined batches. +2. **Passthrough path** (`_passthrough`): Redis seen-set for cross-page dedup. SADD new keys after each page. -- `Executor.run(node, ctx, limit, next_page, **params) -> FeedResult` +Priority arbitration: higher `dedup_priority` wins. Equal = first-seen. -Execution strategy: +## Config Hash -1. **Plan-first** - - `build_plan(...)` -> execute returned `Plan` - - otherwise call node `get_data(...)` -2. **Centralized concurrency** - - child runs use executor-managed `asyncio.gather(...)` -3. **Dedup/refill hooks** - - for non-slot nodes with `ctx.dedup`, run `DedupRuntime.run_node_with_dedup_refill(...)` - - for `SlotsPlan`, dedup/refill is handled inside slot execution +`BaseNode.config_hash()` = md5(model_dump_json(sort_keys=True))[:8]. -`SlotsPlan` execution highlights: +Used in Redis keys. Any config change = different hash = fresh cache. -1. collect unique owners + demand per owner -2. fetch owners concurrently (with optional `owner_fetch_limits` overrides) -3. merge only changed cursor keys (`CursorMap.merge_delta`) -4. apply: - - dedup arbitration + refill (`apply_slots_plan_dedup`) when `ctx.dedup` exists - - refill-only deficits (`apply_slots_plan_refill`) when only `ctx.refill_settings` exists -5. consume slot schedule and call `assemble(...)` +## Generation ID -When dedup is active for a slots plan, owners are executed with `dedup=None` in owner context so global arbitration stays centralized. +Cached wrapper stamps `gen` (random nonce) in cursor and `:meta`. Stale gen = rebuild from scratch. +## Shared Cache -## 6) Plans: declarative execution +Two wrappers with same `cache_key` share base data. Each applies its own rerank. RedisLock prevents duplicate fetches. -Plans separate “what to run” from “how to run it”. +## File Structure -- `CallablePlan(fn)` - - node-provided async function with custom flow, still executed by executor - -- `SlotsPlan(ctx, limit, next_page, params, slots, assemble, owner_fetch_limits=None)` - - `slots`: ordered `SlotSpec(owner, max_count)` schedule - - `assemble(output, merged_next_page, owner_results)`: builds final `FeedResult` - - -## 7) Mergers and leaf responsibilities - -### SubFeed (leaf) - -- derives its local cursor from `next_page.data[subfeed_id]` (defaults page=1/after=None) -- calls `methods_dict[method_name]` -- passes only params present in method signature + `subfeed_params` -- async methods are awaited; sync methods run via `asyncio.to_thread(...)` -- `raise_error=False` converts method failure into empty `FeedResultClient` -- optional `shuffle` then normalizes to `FeedResult` - -### Slot-based mergers - -These build `SlotsPlan`: - -- `MergerAppend`: concatenation (optional shuffle) -- `MergerAppendDistribute` (`type="merger_distribute"`): append then redistribute by `distribution_key` -- `MergerPositional`: page-local slot ownership for `positional` vs `default`, keeps its own merger cursor -- `MergerPercentage`: integer allocation by percentages; when total is exactly 100, remainder is distributed to avoid underfill -- `MergerPercentageGradient`: two-owner percentage curve across the page, then advances merger page cursor - -### MergerViewSession (Redis-backed session cache) - -Goal: cache a session-sized list and serve slices. - -Flow: - -1. build cache key: `{merger_id}_{user_id}` + optional suffix from `custom_view_session_key` -2. check Redis `exists`; if no cache or no merger cursor in request -> regenerate session -3. on hit, `get`; if Redis returns `None` unexpectedly, regenerate -4. on generation: execute child once for `session_size`, optional dedup, store JSON with TTL -5. return page slice and increment merger cursor page -6. optional `shuffle` is applied to returned page slice (cache payload is not reshuffled) - -### MergerDeduplication (single-child wrapper) - -Goal: deduplicate while keeping child mix/slot semantics. - -Key behavior: - -- fresh session when merger cursor is absent or `page <= 0` - - reset descendant cursors - - for Redis backend, reset Redis seen-state key -- seen-state backend: - - `cursor`: encoded into merger cursor `after` - - `redis`: ZSET `dedup:{merger_id}:{user_id}` (+ optional custom suffix) -- builds `DeduplicationPolicy` + child `ExecutionContext(dedup=..., refill_settings=...)` -- executes child via shared executor, commits store, writes merger cursor (`page+1`, `after` for cursor backend) - -Refill/overfetch behavior: - -- duplicates trigger bounded refill loops (`max_refill_loops`) -- overfetch (`overfetch_factor`) is applied only for rewindable integer-offset cursors -- when overfetch is used, leaf cursor is rewound to inspected-count to avoid skipping unseen items - - -## 8) Dedup policy + seen stores - -### DeduplicationPolicy - -Owns key extraction + acceptance rules: - -- entity key from `dedup_key` + `missing_key_policy` -- reject duplicates already seen in current response (`seen_request_set`) -- compare candidate priority vs persisted seen priority - -Capabilities: - -- batched prefetch from store -- per-owner arbitration with deterministic tie-break: `(-dedup_priority, owner_rank, item_rank)` -- ordered single-stream acceptance (`accept_batch`) returning accepted items + inspected count - -### Seen stores - -- `CursorSeenStore` - - in-cursor map of `{key -> max_priority}` - - optional compression + max-key trimming at commit - -- `RedisSeenStore` - - cached reads via `redis_zmscore(...)` - - buffered writes via `redis_zadd_and_expire(...)` - - -## 9) Redis/JSON helpers - -- `_redis_call(client, method_name, *args, **kwargs)` - - async redis client: direct await - - sync redis client: `asyncio.to_thread(...)` - -Other helpers: - -- `jsonlib`: thin `orjson` wrapper compatible with package usage (`dumps`/`loads`) -- `dedup_utils`: cursor encode/decode + Redis ZSET helper fallbacks (`zmscore` / pipeline) - - -## 10) End-to-end call flows - -### A) Standard request (no view session, no dedup) - -1. `FeedManager.get_data(...)` builds `ExecutionContext` -2. `Executor.run(root, ctx, limit, next_page)` -3. recursive execution via plans or direct `get_data(...)` -4. returns `FeedResult(data, next_page, has_next_page)` - -### B) Slot-based merger request - -1. merger returns `SlotsPlan` -2. executor fetches owners concurrently -3. optional arbitration/refill runs -4. slots are consumed in schedule order -5. `assemble(...)` builds final result - -### C) Dedup wrapper request - -1. wrapper creates store + policy and child context -2. child executes under dedup/refill control -3. executor performs acceptance/arbitration + bounded refills -4. store commits; wrapper writes merger cursor state - -### D) View-session request - -1. wrapper resolves cache key -2. cache miss/new session -> regenerate and cache -3. cache hit -> load session list from Redis -4. return requested slice + advanced merger page +``` +smartfeed/ + models/ + base.py -- BaseNode, FeedResult, SmartFeedDebugInfo, FeedItem, config_hash + subfeed.py -- SubFeed (leaf) + wrapper.py -- Wrapper (cache + dedup + rerank pipeline) + mixers.py -- Percentage, Positional, Append, Distribute, Gradient + __init__.py -- FeedNode union, FeedConfig, exports + execution/ + executor.py -- run(), _execute_mix() (module-level functions) + context.py -- ExecutionContext + plans.py -- MixPlan, MixChild + redis_lock.py -- RedisLock (SETNX async context manager) + manager.py -- FeedManager (public entry point) +``` diff --git a/CHANGELOG.md b/CHANGELOG.md index 094e455..12fd4a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,62 @@ -# 0.2.0 (2025-11-25) +# Changelog + +## 0.3.0 (2026-04-08) + +SmartFeed v2: complete rewrite. + +### Breaking Changes + +- Unified `Wrapper` node replaces `MergerViewSession`, `MergerDeduplication`, and external rerank hack +- `FeedManager.get_data()` -> `FeedManager.get_feed(session_id, limit, cursor)` +- `user_id` -> `session_id` in ExecutionContext +- Cursors are plain dicts (not `FeedResultNextPage` Pydantic models) +- `FeedResultClient` removed (use `FeedResult` everywhere) +- `merger_id` renamed to `node_id` in configs +- Config `version: "2"` required +- Pydantic v2 required (v1 compat removed) +- `jsonlib.py` removed (orjson used directly) + +### New Features + +- **Wrapper** with optional pipeline stages: cache, dedup, rerank +- **Rerank callable** registered in `methods_dict`, referenced by `method_name` in config +- **Cross-page dedup** via Redis seen-set (SET with TTL) +- **Shared cache** via `cache_key` for A/B testing with different rerankers +- **Generation ID** for stale cursor detection +- **TTL touch** (inactivity timeout) on every cache access +- **Overfetch + refill** for dedup in both cached and passthrough paths +- **`raise_error`** on WrapperRerank: fail-hard or fail-soft mode +- **`_smartfeed_debug_info`** stamped automatically: source, positions, per-wrapper data +- **SmartFeedDebugInfo / FeedItem** Pydantic models for typed output +- **`config_hash`** for automatic cache invalidation on config changes + +### Removed + +- `MergerViewSession` (replaced by `Wrapper(cache=...)`) +- `MergerDeduplication` (replaced by `Wrapper(dedup=...)`) +- `DedupRuntime` (~453 lines) +- `SlotsPlan`, `SlotSpec`, `CallablePlan` (replaced by `MixPlan`) +- `CursorMap` with `merge_delta` / overfetch +- `DeduplicationPolicy`, `CursorSeenStore`, `RedisSeenStore` +- `pydantic_compat.py` +- `jsonlib.py` +- `examples/` +- `mergers/` directory (mixers now in `models/mixers.py`) +- `policies/` directory +- `feed_models.py`, `schemas.py` + +### Stats + +| Metric | v1 | v2 | +|--------|----|----| +| Python files | 24 | 11 | +| Lines of code | 2838 | ~1300 | +| Test files | 26 | 17 | + +## 0.2.0 (2025-11-25) * Bump dependency versions -# 0.1.0 (2024-12-26) +## 0.1.0 (2024-12-26) * Initial build and publish to PyPI diff --git a/README.md b/README.md index a86e8a4..1e13123 100644 --- a/README.md +++ b/README.md @@ -1,229 +1,214 @@ # SmartFeed -Python-package для формирования ленты (Feed) из клиентских данных с заданной конфигурацией. - -## Содержание: - -- [Использование](#использование) - - [Установка](#установка) - - [Формирование конфигурации](#формирование-конфигурации) - - [MergerDeduplication (дедупликация)](#mergerdeduplication-дедупликация) - - [Параметры MergerDeduplication](#параметры-mergerdeduplication) - - [Важные нюансы (сброс, cursor/redis, overfetch)](#важные-нюансы-сброс-cursorredis-overfetch) - - [Требования к клиентскому методу](#требования-к-клиентскому-методу) - - [Запуск](#запуск) - -## Использование - -### Установка - -``` -poetry add git+ssh://git@github.com:epoch8/looky-timeline.git -``` - -### Формирование конфигурации - -Конфигурация каждого фида должна быть словарем следующего вида: -``` -"version": "1", -"view_session": True or False, -"session_size": 800, # по умолчанию 100 -"session_live_time": 500, # по умолчанию 300 -"feed": { - "merger_id": "merger_pos", - "type": "merger_positional", - "positions": [1, 3, 15], - "start": 17, - "end": 200, - "step": 2, - "positional": { - "subfeed_id": "sf_positional", - "type": "subfeed", - "method_name": "ads", - "subfeed_params": { - "limit_to_return": 10, - }, - }, - "default": { - "merger_id": "merger_percent", - "type": "merger_percentage", - "shuffle": False, - "items": [ - { - "percentage": 40, - "data": { - "subfeed_id": "sf_1_default_merger_of_main", - "type": "subfeed", - "method_name": "followings", - }, - }, - { - "percentage": 60, - "data": { - "subfeed_id": "sf_2_default_merger_of_main", - "type": "subfeed", - "method_name": "ads", - "raise_error": True or False (default = True), - }, - }, - ], - }, -}, -``` - -### MergerDeduplication (дедупликация) - -MergerDeduplication — обёртка над одним дочерним узлом (merger или subfeed), которая удаляет дубли по ключу. - -Ключевые свойства реализации: - -- Дедупликация выполняется на уровне листьев (SubFeed), а не пост-обработкой результата мерджера. - Это важно: вложенные мерджеры (positional/percentage/gradient/append/distribute) сохраняют свои правила смешивания. - Если элемент удалён как дубль, MergerDeduplication «дозапросит» следующий элемент из того же источника. -- Состояние «уже видели» может храниться: - - в курсоре (state_backend="cursor") — удобно без Redis, но курсор может расти; - - в Redis (state_backend="redis") — удобно для большого состояния. - -Пример: обернуть существующую конфигурацию фида дедупликацией: +Feed orchestrator. Builds one paginated feed from multiple async data sources using a declarative JSON config. -```json -{ - "version": "1", +## Quick Start + +```python +from smartfeed.manager import FeedManager + +config = { + "version": "2", "feed": { - "merger_id": "dedup_main", - "type": "merger_deduplication", - "dedup_key": "id", - "missing_key_policy": "error", - "state_backend": "cursor", - "cursor_compress": true, - "cursor_max_keys": 2000, - "overfetch_factor": 2, - "max_refill_loops": 20, + "type": "wrapper", + "node_id": "main", + "cache": {"session_size": 300, "session_ttl": 300}, + "dedup": {"dedup_key": "id", "overfetch_factor": 3}, + "rerank": {"method_name": "my_rerank"}, "data": { - "merger_id": "merger_percent", "type": "merger_percentage", + "node_id": "mix", "items": [ - { - "percentage": 60, - "data": { - "subfeed_id": "sf_posts", - "type": "subfeed", - "method_name": "posts", - "dedup_priority": 10 - } - }, - { - "percentage": 40, - "data": { - "subfeed_id": "sf_ads", - "type": "subfeed", - "method_name": "ads", - "dedup_priority": 0 - } - } + {"percentage": 40, "data": {"type": "subfeed", "subfeed_id": "source_a", "method_name": "source_a"}}, + {"percentage": 60, "data": {"type": "subfeed", "subfeed_id": "source_b", "method_name": "source_b"}} ] } } } + +async def source_a(user_id, limit, next_page, **kwargs): + # Return FeedResult(data=[...], next_page={...}, has_next_page=True) + ... + +async def my_rerank(items, session_id): + # Return same items in new order. len(result) == len(items). + return sorted(items, key=lambda x: x.get("score", 0), reverse=True) + +manager = FeedManager(config=config, methods_dict={"source_a": source_a, "source_b": source_b, "my_rerank": my_rerank}, redis_client=redis) +result = await manager.get_feed(session_id="user_123", limit=20, cursor={}) +# result.data -- list of items +# result.next_page -- pass back as cursor for next page +# result.has_next_page -- whether more pages available ``` -В примере выше, если `posts` и `ads` отдают объекты с одинаковым `id`, то «побеждает» источник с большим `dedup_priority`. +## Node Types -### Параметры MergerDeduplication +### SubFeed (leaf) -Обязательные поля: +Calls an async function from `methods_dict`. -- `merger_id: str` — уникальный ID мерджера. -- `type: "merger_deduplication"` -- `data` — ровно один дочерний узел (subfeed или merger). +```json +{"type": "subfeed", "subfeed_id": "tours", "method_name": "get_tours", "dedup_priority": 1} +``` -Поля дедупликации: +Fields: +- `subfeed_id` -- unique ID, used in cursors and `_smartfeed_debug_info.source` +- `method_name` -- key in `methods_dict` +- `subfeed_params` -- static kwargs passed to method (default: `{}`) +- `raise_error` -- if `false`, swallow errors and return empty (default: `true`) +- `shuffle` -- shuffle results (default: `false`) +- `dedup_priority` -- higher = wins dedup conflicts (default: `0`) -- `dedup_key: str | null` — имя ключа/атрибута для поиска дублей. - - если `null`, ключом считается сам объект (подходит, когда объекты уже hashable/строковые). -- `missing_key_policy: "error" | "keep" | "drop"` (default: `"error"`) - - `error`: выбросить ошибку, если у элемента нет `dedup_key`; - - `keep`: сохранить элемент, даже если ключа нет; - - `drop`: выкинуть элемент без ключа. +### Wrapper (unified cache + dedup + rerank) -Состояние seen (межстраничная дедупликация): +Single node with optional pipeline stages: `fetch -> dedup -> rerank -> cache -> paginate`. -- `state_backend: "cursor" | "redis"` (default: `"cursor"`) -- `state_ttl_seconds: int` (default: `3600`) — TTL для Redis состояния (только для backend=`redis`). -- `cursor_compress: bool` (default: `true`) — сжимать seen-состояние в cursor backend. -- `cursor_max_keys: int | null` — ограничить размер seen-состояния в cursor backend (полезно для контроля размера курсора). +```json +{ + "type": "wrapper", + "node_id": "pipeline", + "cache": {"session_size": 300, "session_ttl": 300}, + "dedup": {"dedup_key": "id", "overfetch_factor": 3, "max_refill_loops": 5}, + "rerank": {"method_name": "my_rerank", "raise_error": true}, + "dedup_priority": 3, + "data": { ... } +} +``` -Производительность/поведение: +All stages are optional. Combinations: -- `overfetch_factor: int` (default: `1`) — «перезапрос» внутри листьев, чтобы быстрее добрать `limit` без множества рефиллов. -- `max_refill_loops: int` (default: `20`) — верхняя граница количества дозапросов на один лист. +| cache | dedup | rerank | Behavior | +|-------|-------|--------|----------| +| yes | no | no | Session cache with Redis pagination | +| no | yes | no | Per-page dedup with Redis seen-set | +| yes | yes | yes | Full pipeline: fetch -> dedup -> rerank -> cache | +| yes | no | yes | Rerank + cache (no dedup) | +| no | no | yes | Per-page rerank (no cache, stateless) | -### Важные нюансы (сброс, cursor/redis, overfetch) +#### Cache -- Сброс состояния при `page <= 0` или отсутствии курсора для `merger_id`. - - MergerDeduplication воспринимает это как «fresh session» и очищает курсоры всех дочерних узлов. - - Для backend=`redis` дополнительно удаляет ключ состояния в Redis. +```json +"cache": {"session_size": 300, "session_ttl": 300, "cache_key": null} +``` -- Если `state_backend="redis"`, нужно передать `redis_client` в `FeedManager`. - - Ключ состояния в Redis строится как `dedup:{merger_id}:{user_id}`. - - Можно добавить суффикс через параметр запроса `custom_deduplication_key` (или `custom_view_session_key`), - чтобы разделять состояния для разных режимов выдачи. +- `session_size` -- items to fetch from child and cache +- `session_ttl` -- Redis TTL (seconds). Acts as inactivity timeout: refreshed on every access +- `cache_key` -- explicit key for shared cache between wrappers (A/B testing) -- Приоритет (`dedup_priority`) — это приоритет победы при конфликте дублей, а не порядок вывода. - - Больше `dedup_priority` → элемент «побеждает» и будет считаться seen с этим приоритетом. - - Это поле доступно у всех узлов (merger/subfeed) и используется MergerDeduplication при дедупликации. +#### Dedup -- overfetch работает безопасно только для «перематываемых» курсоров. - - Сейчас overfetch включается только если `next_page.after` у листа — целочисленный offset. - - Если `after` — строка/словарь/любой другой объект, он считается непрозрачным и overfetch не применяется. +```json +"dedup": {"dedup_key": "id", "missing_key_policy": "error", "overfetch_factor": 3, "max_refill_loops": 5, "state_ttl": 300} +``` + +- `dedup_key` -- field name for duplicate detection +- `missing_key_policy` -- `"error"` | `"keep"` | `"drop"` (default: `"error"`) +- `overfetch_factor` -- fetch `limit * factor` to compensate for dedup removals (default: `1`) +- `max_refill_loops` -- max refill iterations (default: `20`) +- `state_ttl` -- TTL for Redis seen-set (default: `300`) + +Cross-page dedup: seen keys stored in Redis SET `sf:{session_id}:{node_id}:{hash}:seen`. -- Главный реальный bottleneck в дедупликации — не обёртки/копии, а рефиллы. - - Если дублей много и upstream-методы дорогие, стоит аккуратно подобрать `overfetch_factor` и `max_refill_loops`. +#### Rerank -### Требования к клиентскому методу +```json +"rerank": {"method_name": "my_rerank", "raise_error": true} +``` -Клиентский метод для получения данных должен обязательно включать в себя следующие параметры: -- **user_id: Any** - ID объекта, на который ориентируемся при получении данных субфида. -- **limit: int** - Количество возвращаемых данных. -- **next_page: FeedResultNextPageInside** - Объект курсора пагинации, формируется на стороне клиента после обработки данных. +- `method_name` -- key in `methods_dict`, must be `async def(items, session_id) -> items` +- `raise_error` -- `true` = crash on rerank failure, `false` = keep original order (default: `true`) -Возвращаемый тип данных: **FeedResultClient**. +Contract: callable returns exactly `len(items)` items. Only reorders. -### Запуск -Для получения ленты с помощью SmartFeed нужно выполнить следующий код: +### Mixer Nodes +**MergerPercentage** -- split by percentage: +```json +{"type": "merger_percentage", "node_id": "mix", "items": [{"percentage": 40, "data": {...}}, {"percentage": 60, "data": {...}}]} ``` -from smartfeed.manager import FeedManager -from smartfeed.schemas import FeedResult, FeedResultNextPage, FeedResultNextPageInside -from client.services import ClientService +**MergerPositional** -- insert at specific positions (1-indexed): +```json +{"type": "merger_positional", "node_id": "pos", "positions": [1, 3, 5, 7], "positional": {...}, "default": {...}} +``` + +**MergerAppend** -- concatenate children: +```json +{"type": "merger_append", "node_id": "append", "items": [{...}, {...}]} +``` + +**MergerAppendDistribute** -- round-robin by key: +```json +{"type": "merger_distribute", "node_id": "diverse", "distribution_key": "operator_id", "items": [{...}]} +``` + +**MergerPercentageGradient** -- shift ratio over pages: +```json +{"type": "merger_percentage_gradient", "node_id": "grad", "item_from": {"percentage": 80, "data": {...}}, "item_to": {"percentage": 20, "data": {...}}, "step": 10, "size_to_step": 30} +``` + +## dedup_priority + +Every node has `dedup_priority: int = 0`. Higher = more important. + +- When two items have the same dedup key, the one with higher priority wins +- Equal priority: first-seen wins (order from mixer children list) +- Wrapper/mixer with `dedup_priority != 0` overrides all children in that subtree +- Purely config-based, not stored in items + +## _smartfeed_debug_info -config = {} # получаем конфигурацию фида -methods_dict = { - "method_1": ClientService().method_1, - "method_2": ClientService().method_2, - # и т.д. +SmartFeed stamps metadata on every item: + +```python +item["_smartfeed_debug_info"] = { + "source": "recommended_tours", # subfeed_id + "smartfeed_position": 5, # final position (set by FeedManager) + "strategy": "model_hot_users", # from subfeed (optional) + # Per-wrapper positions: + "pipeline": {"smartfeed_position": 15, "rerank_position": 3}, + # Rerank fields (set by rerank callable): + "rerank_position": 1, + "rrf_score": 0.032, + "feature_score": 10683198.0, } -# для конфигурации view_session = False, -# Redis передавать небязательно -redis_client = redis.Redis() - -feed_manager = FeedManager( - config=config, - methods_dict=methods_dict, - redis_client=redis_client, -) - -user_id = "sjjdj?" # любой тип данных -limit = 100 -next_page = FeedResultNextPage( - data={ - "subfeed_id": FeedResultNextPageInside(page=1, after=None), - } -) -data: FeedResult = await feed_manager.get_data( - user_id=user_id, - limit=limit, - next_page=next_page, -) -``` \ No newline at end of file +``` + +## Redis Keys + +``` +sf:{session_id}:{node_id}:{config_hash} -- cached data +sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (child cursor, has_next, gen) +sf:{session_id}:{node_id}:{config_hash}:seen -- dedup seen-set (SET) +``` + +`config_hash` = md5 of config subtree JSON. Change any param -> different hash -> fresh cache. + +TTL is an inactivity timeout: refreshed on every request. Cache lives while user scrolls. + +## Generation ID + +Cached wrappers stamp a `gen` nonce in cursor and Redis `:meta`. If user returns with stale `gen` after cache expired -> full rebuild from page 1. + +## Subfeed Method Signature + +```python +async def my_source(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: + return FeedResult(data=[...], next_page={"page": 2, "after": ...}, has_next_page=True) +``` + +## Installation + +``` +pip install epoch8-smartfeed +``` + +Requires: Python 3.9+, Pydantic v2, redis (async), orjson. + +## Testing + +```bash +pytest tests/ -v +``` + +Requires: fakeredis, pytest-asyncio. diff --git a/smartfeed/models/wrapper.py b/smartfeed/models/wrapper.py index 808ea57..eb22179 100644 --- a/smartfeed/models/wrapper.py +++ b/smartfeed/models/wrapper.py @@ -54,6 +54,7 @@ async def execute( ctx: Any = None, **params: Any, ) -> FeedResult: + """Entry point. Routes to cached or passthrough path.""" if self.cache is None or ctx is None or ctx.redis is None: return await self._passthrough(methods_dict, session_id, limit, cursor, ctx) @@ -64,7 +65,7 @@ async def execute( # -- dedup ----------------------------------------------------------------- def _resolve_dedup_priorities(self) -> Dict[str, int]: - """Walk the config subtree to build {subfeed_id: priority} map.""" + """Build {subfeed_id: effective_priority} by walking the config tree.""" result: Dict[str, int] = {} self._collect_priorities(self.data, override_priority=0, out=result) return result @@ -73,6 +74,7 @@ def _resolve_dedup_priorities(self) -> Dict[str, int]: def _collect_priorities( node: BaseNode, override_priority: int, out: Dict[str, int] ) -> None: + """Recursive walk: propagate override_priority down, write SubFeed priorities to out.""" from .subfeed import SubFeed # Determine the effective priority for this subtree. @@ -103,12 +105,7 @@ def _collect_priorities( Wrapper._collect_priorities(child, effective, out) def _dedup(self, data: List, seen: Optional[set] = None) -> List: - """Deduplicate items by dedup_key with priority arbitration. - - Args: - data: items to deduplicate - seen: optional external seen-set (for cross-page dedup). Mutated in-place. - """ + """Remove duplicates by dedup_key. Higher dedup_priority wins, equal = first-seen wins.""" if not self.dedup: return data @@ -165,6 +162,7 @@ def _dedup(self, data: List, seen: Optional[set] = None) -> List: async def _apply_rerank( self, data: list, methods_dict: dict, session_id: str ) -> list: + """Call rerank callable from methods_dict. Validates output length.""" if not self.rerank: return data rerank_fn = methods_dict[self.rerank.method_name] @@ -186,7 +184,7 @@ async def _apply_rerank( # -- debug stamping -------------------------------------------------------- def _stamp_pre_rerank(self, data: list) -> None: - """Stamp position before rerank (order from tree/dedup).""" + """Stamp smartfeed_position on each item (position before rerank).""" for i, item in enumerate(data): if isinstance(item, dict): item.setdefault("_smartfeed_debug_info", {})[self.node_id] = { @@ -194,7 +192,7 @@ def _stamp_pre_rerank(self, data: list) -> None: } def _stamp_post_rerank(self, data: list) -> None: - """Stamp position after rerank.""" + """Stamp rerank_position on each item (position after rerank).""" for i, item in enumerate(data): if isinstance(item, dict): item.setdefault("_smartfeed_debug_info", {}).setdefault(self.node_id, {})["rerank_position"] = i @@ -209,6 +207,7 @@ async def _passthrough( cursor: dict, ctx: Any, ) -> FeedResult: + """No-cache path: fetch -> dedup (with overfetch/refill) -> rerank -> return.""" if self.dedup: # Load persisted seen-set from Redis (cross-page dedup) seen_keys = await self._load_seen_set(ctx, session_id) @@ -258,9 +257,11 @@ async def _passthrough( ) def _seen_set_key(self, session_id: str) -> str: + """Redis key for cross-page dedup seen-set.""" return f"sf:{session_id}:{self.node_id}:{self.config_hash()}:seen" async def _load_seen_set(self, ctx: Any, session_id: str) -> set: + """Load seen keys from Redis SET for cross-page dedup.""" if not ctx or not ctx.redis: return set() key = self._seen_set_key(session_id) @@ -268,6 +269,7 @@ async def _load_seen_set(self, ctx: Any, session_id: str) -> set: return {m.decode() if isinstance(m, bytes) else m for m in members} if members else set() async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None: + """Append new keys to Redis seen-set and refresh TTL.""" if not ctx or not ctx.redis or not new_keys: return key = self._seen_set_key(session_id) @@ -278,10 +280,12 @@ async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None # -- cache helpers --------------------------------------------------------- def _base_key(self, session_id: str) -> str: + """Redis key prefix for this wrapper's cache data.""" key_part = self.cache_key or self.node_id return f"sf:{session_id}:{key_part}:{self.config_hash()}" async def _read_cache(self, ctx: Any, session_id: str) -> Optional[List]: + """Read cached session data from Redis. Returns None on miss.""" base = self._base_key(session_id) raw = await ctx.redis.get(base) if raw is None: @@ -289,6 +293,7 @@ async def _read_cache(self, ctx: Any, session_id: str) -> Optional[List]: return orjson.loads(raw) async def _read_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: + """Read cache metadata (gen, child_cursor, child_has_next) from Redis.""" base = self._base_key(session_id) raw = await ctx.redis.get(f"{base}:meta") if raw is None: @@ -304,6 +309,7 @@ async def _write_cache( child_cursor: Dict, child_has_next: bool = False, ) -> None: + """Write session data and metadata to Redis with TTL.""" base = self._base_key(session_id) ttl = self.cache.session_ttl @@ -317,6 +323,7 @@ async def _write_cache( await pipe.execute() async def _touch_ttl(self, ctx: Any, session_id: str) -> None: + """Refresh TTL on data and meta keys (keeps cache alive while user scrolls).""" if not self.cache: return base = self._base_key(session_id) @@ -332,6 +339,7 @@ def _paginate( self, data: List, limit: int, page: int, gen: str, child_has_next: bool = False, ) -> FeedResult: + """Slice cached data for the requested page and build cursor.""" start = (page - 1) * limit end = start + limit page_data = data[start:end] @@ -360,6 +368,7 @@ async def _execute_with_cache( cursor: dict, ctx: Any, ) -> FeedResult: + """Cached path: warm hit -> paginate, stale/miss -> cold build.""" my_cursor = cursor.get(self.node_id, {}) cursor_gen = my_cursor.get("gen") cursor_page = my_cursor.get("page", 1) @@ -400,6 +409,7 @@ def _base_shared_key(self, session_id: str) -> str: return f"sf:{session_id}:{self.cache_key}:{self.data.config_hash()}" async def _read_shared_base(self, ctx: Any, session_id: str) -> Optional[List]: + """Read shared base data from Redis. Returns None on miss.""" key = self._base_shared_key(session_id) raw = await ctx.redis.get(key) if raw is None: @@ -410,6 +420,7 @@ async def _write_shared_base( self, ctx: Any, session_id: str, data: List, child_cursor: Dict, child_has_next: bool = False, ) -> None: + """Write shared base data and meta to Redis with TTL.""" key = self._base_shared_key(session_id) ttl = self.cache.session_ttl pipe = ctx.redis.pipeline() @@ -422,6 +433,7 @@ async def _write_shared_base( await pipe.execute() async def _read_shared_base_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: + """Read shared base metadata (child_cursor, child_has_next) from Redis.""" key = self._base_shared_key(session_id) raw = await ctx.redis.get(f"{key}:meta") if raw is None: @@ -474,6 +486,7 @@ async def _cold_build( ctx: Any, child_cursor: Dict, ) -> FeedResult: + """Build full session: fetch -> dedup -> rerank -> write cache -> paginate page 1.""" session_size = self.cache.session_size if self.cache_key is not None: From 321b99da97df923a0faa4b4d1ea2f1c1c4059223 Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Wed, 8 Apr 2026 19:09:13 +0300 Subject: [PATCH 36/41] update pyproject --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f6e40c0..077b7cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ packages = [ [tool.poetry.dependencies] python = ">=3.9" orjson = ">=3.9.0" -pydantic = ">=1.10.7" +pydantic = ">=2.0" redis = ">=4.5.5" [tool.poetry.group.dev.dependencies] From 350486c126a94b67db2e907d2e79827ead4a48e8 Mon Sep 17 00:00:00 2001 From: Shakirov Renat Date: Thu, 9 Apr 2026 17:57:14 +0300 Subject: [PATCH 37/41] new tests --- tests/test_continuation.py | 187 ++++++++++++++++++++++++++++++++ tests/test_distribute.py | 211 +++++++++++++++++++++++++++++++++++++ tests/test_gradient.py | 157 +++++++++++++++++++++++++++ 3 files changed, 555 insertions(+) create mode 100644 tests/test_continuation.py create mode 100644 tests/test_distribute.py create mode 100644 tests/test_gradient.py diff --git a/tests/test_continuation.py b/tests/test_continuation.py new file mode 100644 index 0000000..c354834 --- /dev/null +++ b/tests/test_continuation.py @@ -0,0 +1,187 @@ +"""Tests for Wrapper cache exhaustion and continuation cursor rebuild. + +Scenario: + Wrapper(cache, session_size=20) + limit=10. + - Page 1: cold build -> fetches 20 items, caches them, returns items 0-9 + - Page 2: warm hit -> returns items 10-19 + - Page 3: cache exhausted (start=20 >= 20 cached) -> cold rebuild from + child_cursor stored in meta; new generation spawned. + +The child source tracks its own page via next_page so the rebuild continues +from item 20 onwards with no gap and no overlap. +""" + +import pytest +import fakeredis.aioredis + +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.wrapper import Wrapper, WrapperCache +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +# --------------------------------------------------------------------------- +# Infinite stateful source: returns items based on page cursor +# --------------------------------------------------------------------------- + +async def infinite_source(user_id, limit, next_page, **kw): + """Returns `limit` items starting from (page-1)*limit. IDs are globally unique.""" + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": start + i, "val": f"item_{start + i}"} for i in range(limit)] + return FeedResult( + data=data, + next_page={"page": page + 1}, + has_next_page=True, + ) + + +METHODS = {"source": infinite_source} + + +@pytest.fixture +def redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.fixture +def ctx(redis): + return ExecutionContext(session_id="test_session", methods_dict=METHODS, redis=redis) + + +def _make_wrapper(session_size=20): + return Wrapper( + node_id="w", + cache=WrapperCache(session_size=session_size, session_ttl=300), + data=SubFeed(subfeed_id="src", method_name="source"), + ) + + +# --------------------------------------------------------------------------- +# Helper: collect three pages +# --------------------------------------------------------------------------- + +async def _three_pages(ctx, session_size=20, limit=10): + node = _make_wrapper(session_size=session_size) + r1 = await run_executor.run(node, ctx, limit=limit, cursor={}) + r2 = await run_executor.run(node, ctx, limit=limit, cursor=r1.next_page) + r3 = await run_executor.run(node, ctx, limit=limit, cursor=r2.next_page) + return r1, r2, r3 + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_pages_1_and_2_served_from_cache(ctx): + """Pages 1 and 2 are served from the same generation (warm hits).""" + node = _make_wrapper(session_size=20) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + gen1 = r1.next_page["w"]["gen"] + + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + gen2 = r2.next_page["w"]["gen"] + + # Both pages share the same generation id + assert gen1 == gen2 + + +@pytest.mark.asyncio +async def test_page3_triggers_rebuild_with_new_generation(ctx): + """Page 3 exhausts the 20-item cache, triggering a cold rebuild with a new gen.""" + r1, r2, r3 = await _three_pages(ctx) + + gen_before = r2.next_page["w"]["gen"] + gen_after = r3.next_page["w"]["gen"] + + assert gen_before != gen_after, "Page 3 rebuild must produce a new generation id" + + +@pytest.mark.asyncio +async def test_page3_items_continue_from_page2_no_gap(ctx): + """Items on page 3 immediately follow items on page 2 (no gap, no overlap).""" + r1, r2, r3 = await _three_pages(ctx) + + ids_p2 = [item["id"] for item in r2.data] + ids_p3 = [item["id"] for item in r3.data] + + last_p2 = max(ids_p2) + first_p3 = min(ids_p3) + + # No gap: the first item of page 3 immediately follows the last item of page 2 + assert first_p3 == last_p2 + 1, ( + f"Expected page 3 to start at {last_p2 + 1}, got {first_p3}" + ) + + +@pytest.mark.asyncio +async def test_page3_items_no_overlap_with_page1_and_page2(ctx): + """Items on page 3 must not repeat any item from pages 1 or 2.""" + r1, r2, r3 = await _three_pages(ctx) + + seen_ids = {item["id"] for item in r1.data} | {item["id"] for item in r2.data} + p3_ids = {item["id"] for item in r3.data} + + overlap = seen_ids & p3_ids + assert not overlap, f"Page 3 overlaps with earlier pages: {overlap}" + + +@pytest.mark.asyncio +async def test_page3_has_correct_item_count(ctx): + """Page 3 (post-rebuild) still returns exactly `limit` items.""" + _, _, r3 = await _three_pages(ctx) + assert len(r3.data) == 10 + + +@pytest.mark.asyncio +async def test_all_ids_across_three_pages_are_unique(ctx): + """All item IDs across pages 1, 2 and 3 are distinct.""" + r1, r2, r3 = await _three_pages(ctx) + + all_ids = ( + [item["id"] for item in r1.data] + + [item["id"] for item in r2.data] + + [item["id"] for item in r3.data] + ) + assert len(all_ids) == len(set(all_ids)), "Duplicate IDs found across pages" + + +@pytest.mark.asyncio +async def test_all_ids_are_globally_sequential(ctx): + """Items across pages 1-3 form a gapless sequence starting from 0.""" + r1, r2, r3 = await _three_pages(ctx) + + all_ids = sorted( + item["id"] + for page in (r1.data, r2.data, r3.data) + for item in page + ) + expected = list(range(30)) + assert all_ids == expected, f"Expected {expected}, got {all_ids}" + + +@pytest.mark.asyncio +async def test_continuation_works_for_smaller_session_size(ctx): + """Continuation rebuild works correctly with session_size=6 and limit=3.""" + node = _make_wrapper(session_size=6) + + # Pages 1 and 2 from cache + r1 = await run_executor.run(node, ctx, limit=3, cursor={}) + r2 = await run_executor.run(node, ctx, limit=3, cursor=r1.next_page) + # Page 3 triggers rebuild + r3 = await run_executor.run(node, ctx, limit=3, cursor=r2.next_page) + + # New generation + assert r3.next_page["w"]["gen"] != r2.next_page["w"]["gen"] + + # No overlap + seen = {item["id"] for item in r1.data} | {item["id"] for item in r2.data} + new = {item["id"] for item in r3.data} + assert not (seen & new) + + # No gap + last_p2 = max(item["id"] for item in r2.data) + first_p3 = min(item["id"] for item in r3.data) + assert first_p3 == last_p2 + 1 diff --git a/tests/test_distribute.py b/tests/test_distribute.py new file mode 100644 index 0000000..f94bfcf --- /dev/null +++ b/tests/test_distribute.py @@ -0,0 +1,211 @@ +import pytest + +from smartfeed.models.base import FeedResult +from smartfeed.models.mixers import MergerAppendDistribute +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +async def make_operators(user_id, limit, next_page, **kw): + """Source producing items with operator_id cycling through op_0, op_1, op_2.""" + data = [ + {"id": i, "operator_id": f"op_{i % 3}", "val": f"tour_{i}"} + for i in range(limit) + ] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def make_scored(user_id, limit, next_page, **kw): + """Source with a numeric score field for testing sorting_key. + + Items are in ascending score order so that descending sort visibly reorders them. + op_0 gets scores 1, 3, 5, ... and op_1 gets scores 2, 4, 6, ... + """ + data = [ + {"id": i, "operator_id": f"op_{i % 2}", "score": i + 1, "val": f"tour_{i}"} + for i in range(limit) + ] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +async def make_no_key(user_id, limit, next_page, **kw): + """Source producing items WITHOUT operator_id (to test missing key behaviour).""" + data = [{"id": i, "val": f"tour_{i}"} for i in range(limit)] + return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + + +METHODS = { + "operators": make_operators, + "scored": make_scored, + "no_key": make_no_key, +} + + +@pytest.fixture +def ctx(): + return ExecutionContext(session_id="s1", methods_dict=METHODS, redis=None) + + +# --------------------------------------------------------------------------- +# Basic round-robin +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_basic_round_robin_no_consecutive_same_operator(ctx): + """Consecutive items should not share the same operator_id (round-robin effect).""" + node = MergerAppendDistribute( + node_id="dist", + items=[SubFeed(subfeed_id="ops", method_name="operators")], + distribution_key="operator_id", + ) + result = await run_executor.run(node, ctx, limit=9, cursor={}) + + assert len(result.data) == 9 + for i in range(len(result.data) - 1): + assert ( + result.data[i]["operator_id"] != result.data[i + 1]["operator_id"] + ), f"Consecutive duplicate operator at positions {i} and {i+1}" + + +@pytest.mark.asyncio +async def test_round_robin_all_operators_represented(ctx): + """All three operators must appear in a full round-robin of 9 items.""" + node = MergerAppendDistribute( + node_id="dist", + items=[SubFeed(subfeed_id="ops", method_name="operators")], + distribution_key="operator_id", + ) + result = await run_executor.run(node, ctx, limit=9, cursor={}) + + operator_ids = {item["operator_id"] for item in result.data} + assert operator_ids == {"op_0", "op_1", "op_2"} + + +@pytest.mark.asyncio +async def test_round_robin_result_length_matches_limit(ctx): + """Result must contain exactly `limit` items.""" + node = MergerAppendDistribute( + node_id="dist", + items=[SubFeed(subfeed_id="ops", method_name="operators")], + distribution_key="operator_id", + ) + for lim in (3, 6, 9): + result = await run_executor.run(node, ctx, limit=lim, cursor={}) + assert len(result.data) == lim + + +# --------------------------------------------------------------------------- +# sorting_key sorts before distributing +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_sorting_key_desc_highest_score_first(ctx): + """With sorting_key='score' and sorting_desc=True, higher scores appear earlier.""" + node = MergerAppendDistribute( + node_id="dist_sorted", + items=[SubFeed(subfeed_id="scored", method_name="scored")], + distribution_key="operator_id", + sorting_key="score", + sorting_desc=True, + ) + result = await run_executor.run(node, ctx, limit=4, cursor={}) + + assert len(result.data) == 4 + # The highest score items per operator group should appear first. + # Verify that score of first item per operator_id group is higher than second. + from collections import defaultdict + by_op = defaultdict(list) + for item in result.data: + by_op[item["operator_id"]].append(item["score"]) + + for op, scores in by_op.items(): + if len(scores) >= 2: + assert scores[0] >= scores[1], ( + f"Operator {op}: expected descending scores, got {scores}" + ) + + +@pytest.mark.asyncio +async def test_sorting_key_asc_lowest_score_first(ctx): + """With sorting_key='score' and sorting_desc=False, lower scores appear earlier.""" + node = MergerAppendDistribute( + node_id="dist_asc", + items=[SubFeed(subfeed_id="scored", method_name="scored")], + distribution_key="operator_id", + sorting_key="score", + sorting_desc=False, + ) + result = await run_executor.run(node, ctx, limit=4, cursor={}) + + assert len(result.data) == 4 + # All scores in the output must be <= max available score in source + # (weak check: just confirm the sort changed the order vs no sorting) + scores = [item["score"] for item in result.data] + # First item of each operator should have a smaller score than last + assert scores[0] <= scores[-1] or len(set(scores)) == 1 + + +@pytest.mark.asyncio +async def test_no_sorting_key_preserves_source_order(ctx): + """Without sorting_key, source order is preserved before distribution.""" + node_no_sort = MergerAppendDistribute( + node_id="dist_nosort", + items=[SubFeed(subfeed_id="scored", method_name="scored")], + distribution_key="operator_id", + ) + node_sorted = MergerAppendDistribute( + node_id="dist_sorted", + items=[SubFeed(subfeed_id="scored", method_name="scored")], + distribution_key="operator_id", + sorting_key="score", + sorting_desc=True, + ) + + r_nosort = await run_executor.run(node_no_sort, ctx, limit=4, cursor={}) + r_sorted = await run_executor.run(node_sorted, ctx, limit=4, cursor={}) + + ids_nosort = [item["id"] for item in r_nosort.data] + ids_sorted = [item["id"] for item in r_sorted.data] + + # Sorted and unsorted should produce different orderings (source has varying scores) + assert ids_nosort != ids_sorted + + +# --------------------------------------------------------------------------- +# Missing distribution_key raises KeyError +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_missing_distribution_key_raises(ctx): + """Items without the distribution_key should raise KeyError.""" + node = MergerAppendDistribute( + node_id="dist_bad", + items=[SubFeed(subfeed_id="nokey", method_name="no_key")], + distribution_key="operator_id", + ) + with pytest.raises(KeyError): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +# --------------------------------------------------------------------------- +# Multiple subfeeds +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_multiple_subfeeds_merged_before_distribute(ctx): + """Items from multiple subfeeds are pooled before round-robin distribution.""" + node = MergerAppendDistribute( + node_id="dist_multi", + items=[ + SubFeed(subfeed_id="ops_a", method_name="operators"), + SubFeed(subfeed_id="ops_b", method_name="operators"), + ], + distribution_key="operator_id", + ) + result = await run_executor.run(node, ctx, limit=6, cursor={}) + + assert len(result.data) == 6 + # Still no consecutive duplicates + for i in range(len(result.data) - 1): + assert result.data[i]["operator_id"] != result.data[i + 1]["operator_id"] diff --git a/tests/test_gradient.py b/tests/test_gradient.py new file mode 100644 index 0000000..9b49d47 --- /dev/null +++ b/tests/test_gradient.py @@ -0,0 +1,157 @@ +import pytest + +from smartfeed.models.base import FeedResult +from smartfeed.models.mixers import MergerPercentageGradient, MergerPercentageItem +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +async def make_source(prefix, user_id, limit, next_page, **kw): + page = next_page.get("page", 1) + start = (page - 1) * limit + data = [{"id": f"{prefix}_{start + i}", "val": prefix} for i in range(limit)] + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) + + +METHODS = { + "source_a": lambda *a, **kw: make_source("a", *a, **kw), + "source_b": lambda *a, **kw: make_source("b", *a, **kw), +} + + +@pytest.fixture +def ctx(): + return ExecutionContext(session_id="s1", methods_dict=METHODS, redis=None) + + +def _make_gradient(pct_from=80, pct_to=20, step=10, size_to_step=10): + return MergerPercentageGradient( + node_id="grad", + item_from=MergerPercentageItem( + percentage=pct_from, + data=SubFeed(subfeed_id="source_a", method_name="source_a"), + ), + item_to=MergerPercentageItem( + percentage=pct_to, + data=SubFeed(subfeed_id="source_b", method_name="source_b"), + ), + step=step, + size_to_step=size_to_step, + ) + + +@pytest.mark.asyncio +async def test_basic_80_20_page1(ctx): + """On page 1 with 80/20 split and limit=10: 8 from source_a, 2 from source_b.""" + node = _make_gradient(pct_from=80, pct_to=20, step=10, size_to_step=10) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + assert len(result.data) == 10 + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + assert sources.count("source_a") == 8 + assert sources.count("source_b") == 2 + + +@pytest.mark.asyncio +async def test_both_sources_present(ctx): + """Items from both sources appear in the result.""" + node = _make_gradient(pct_from=80, pct_to=20, step=10, size_to_step=10) + result = await run_executor.run(node, ctx, limit=10, cursor={}) + + sources = {item["_smartfeed_debug_info"]["source"] for item in result.data} + assert "source_a" in sources + assert "source_b" in sources + + +@pytest.mark.asyncio +async def test_ratio_shifts_on_later_page(ctx): + """Later pages have a higher proportion of source_b than page 1 (gradient effect).""" + node = _make_gradient(pct_from=80, pct_to=20, step=10, size_to_step=10) + + # Page 1 + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + sources_p1 = [item["_smartfeed_debug_info"]["source"] for item in r1.data] + a_count_p1 = sources_p1.count("source_a") + b_count_p1 = sources_p1.count("source_b") + + # Page 2 (cursor carries page state for the gradient node) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + sources_p2 = [item["_smartfeed_debug_info"]["source"] for item in r2.data] + a_count_p2 = sources_p2.count("source_a") + b_count_p2 = sources_p2.count("source_b") + + # The gradient should shift: fewer source_a and more source_b on page 2 + assert a_count_p2 < a_count_p1 + assert b_count_p2 > b_count_p1 + + +@pytest.mark.asyncio +async def test_step_parameter_controls_shift(ctx): + """A larger step causes a bigger ratio change between pages.""" + node_small_step = _make_gradient(pct_from=80, pct_to=20, step=5, size_to_step=10) + node_large_step = _make_gradient(pct_from=80, pct_to=20, step=20, size_to_step=10) + + # Compute ratio difference between page 1 and page 2 for each node + def _b_count(node, cursor): + import asyncio + r = asyncio.get_event_loop().run_until_complete( + run_executor.run(node, ctx, limit=10, cursor=cursor) + ) + sources = [item["_smartfeed_debug_info"]["source"] for item in r.data] + return sources.count("source_b"), r.next_page + + # Use coroutine approach instead + r1_small = await run_executor.run(node_small_step, ctx, limit=10, cursor={}) + r2_small = await run_executor.run(node_small_step, ctx, limit=10, cursor=r1_small.next_page) + + r1_large = await run_executor.run(node_large_step, ctx, limit=10, cursor={}) + r2_large = await run_executor.run(node_large_step, ctx, limit=10, cursor=r1_large.next_page) + + def b_count(result): + sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] + return sources.count("source_b") + + shift_small = b_count(r2_small) - b_count(r1_small) + shift_large = b_count(r2_large) - b_count(r1_large) + + # Larger step => bigger shift in ratio + assert shift_large >= shift_small + + +@pytest.mark.asyncio +async def test_result_length_matches_limit(ctx): + """Result always contains exactly `limit` items.""" + node = _make_gradient() + for lim in (5, 10, 15): + result = await run_executor.run(node, ctx, limit=lim, cursor={}) + assert len(result.data) == lim + + +@pytest.mark.asyncio +async def test_cursor_advances_page(ctx): + """next_page cursor increments the internal page counter for the gradient node.""" + node = _make_gradient() + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + # The gradient stores its own page in cursor under its node_id + assert r1.next_page.get("grad", {}).get("page") == 2 + + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + assert r2.next_page.get("grad", {}).get("page") == 3 + + +@pytest.mark.asyncio +async def test_gradient_eventually_reaches_100_pct_to(ctx): + """After enough pages the ratio fully shifts to source_b (100%).""" + # step=50 will flip to 100% after 2 pages + node = _make_gradient(pct_from=80, pct_to=20, step=50, size_to_step=10) + + cursor = {} + for _ in range(5): + r = await run_executor.run(node, ctx, limit=10, cursor=cursor) + cursor = r.next_page + + sources = [item["_smartfeed_debug_info"]["source"] for item in r.data] + # At full shift, all items come from source_b + assert sources.count("source_b") == 10 + assert sources.count("source_a") == 0 From d799de8b8607652a040710c1cb4c762c09011080 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Mon, 6 Jul 2026 10:16:13 +0100 Subject: [PATCH 38/41] Debug session, overfetch killed bc YAGNI. --- ARCHITECTURE.md | 25 +- README.md | 388 ++++++++++++++------ pyproject.toml | 4 +- smartfeed/execution/redis_lock.py | 7 +- smartfeed/models/subfeed.py | 6 +- smartfeed/models/wrapper.py | 262 ++++++++----- tests/sources.py | 205 +++++++++++ tests/test_bug_cold_build_stampede.py | 64 ++++ tests/test_bug_overfetch_item_loss.py | 76 ++++ tests/test_bug_redis_lock_decode.py | 64 ++++ tests/test_bug_shared_cache_continuation.py | 64 ++++ tests/test_bug_softfail_cursor_leak.py | 76 ++++ tests/test_bug_variable_page_size.py | 61 +++ tests/test_config_parsing.py | 6 +- tests/test_cursors.py | 4 +- tests/test_dedup_correctness.py | 276 ++++++++++++++ tests/test_dedup_edge_cases.py | 100 +++++ tests/test_dedup_priority.py | 2 +- tests/test_high_priority.py | 14 +- tests/test_invariants.py | 121 ++++++ tests/test_pagination_edge_cases.py | 92 +++++ tests/test_wrapper_cache.py | 6 +- tests/test_wrapper_full_pipeline.py | 2 +- 23 files changed, 1686 insertions(+), 239 deletions(-) create mode 100644 tests/sources.py create mode 100644 tests/test_bug_cold_build_stampede.py create mode 100644 tests/test_bug_overfetch_item_loss.py create mode 100644 tests/test_bug_redis_lock_decode.py create mode 100644 tests/test_bug_shared_cache_continuation.py create mode 100644 tests/test_bug_softfail_cursor_leak.py create mode 100644 tests/test_bug_variable_page_size.py create mode 100644 tests/test_dedup_correctness.py create mode 100644 tests/test_dedup_edge_cases.py create mode 100644 tests/test_invariants.py create mode 100644 tests/test_pagination_edge_cases.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 4823ad9..941623a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -99,23 +99,30 @@ FeedItem(id, _smartfeed_debug_info: SmartFeedDebugInfo) ## Redis State ``` -sf:{session_id}:{node_id}:{config_hash} -- cached session data (JSON list) -sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (gen, child_cursor, child_has_next) -sf:{session_id}:{node_id}:{config_hash}:seen -- dedup seen-set (Redis SET) -sf:{session_id}:{cache_key}:{hash} -- shared base cache (for A/B testing) -sf:{session_id}:{node_id}:{hash}:lock -- distributed lock (SETNX) +sf:{session_id}:{node_id}:{config_hash} -- cached session data (JSON list) +sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (gen, child_cursor, child_has_next) +sf:{session_id}:{node_id}:{config_hash}:coldlock -- cold-build lock (SETNX, ttl 30s) +sf:{session_id}:{node_id}:{config_hash}:seen -- session-scoped dedup seen-set (Redis SET) +sf:{session_id}:{cache_key}:{child_hash}:{segment} -- shared base segment (per continuation window) +sf:{session_id}:{cache_key}:{child_hash}:{segment}:lock -- shared segment cold-build lock (SETNX) ``` TTL = inactivity timeout. Refreshed on every access via pipeline EXPIRE. +Cursor for a cached wrapper is `{node_id: {offset, gen}}` (absolute offset). ## Dedup -Two dedup paths, unified `_dedup(data, seen=None)` method: +Two dedup paths, unified `_dedup(data, seen=None)` method. Both fetch exactly the +outstanding deficit and refill until the target is filled or the child is exhausted +(no over-fetch, so no item loss; no early give-up, so no short pages), and both carry +a session-scoped Redis seen-set so dedup holds across the whole scroll: -1. **Cached path** (`_fetch_and_dedup`): overfetch + refill loop. Re-dedup combined batches. -2. **Passthrough path** (`_passthrough`): Redis seen-set for cross-page dedup. SADD new keys after each page. +1. **Cached path** (`_fetch_and_dedup` / `_cold_build`): seen-set carried across cold + rebuilds, reset on a fresh scroll (empty cursor). +2. **Passthrough path** (`_passthrough`): seen-set persisted across pages; SADD after each page. -Priority arbitration: higher `dedup_priority` wins. Equal = first-seen. +Priority arbitration: higher `dedup_priority` wins even when the higher-priority copy is +not first-seen; equal priority = first-seen. ## Config Hash diff --git a/README.md b/README.md index 1e13123..1ab3e10 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,37 @@ # SmartFeed -Feed orchestrator. Builds one paginated feed from multiple async data sources using a declarative JSON config. +Feed orchestrator. Builds one paginated feed from many independent async data sources, described by a declarative JSON/dict config tree. You register your own async fetch functions; SmartFeed handles parallel fetch, cross-source dedup, optional rerank, per-session caching, and cursor-based pagination. + +This README is written for engineers evaluating how it works and how much integration effort it takes. For internal architecture and control flow, see `ARCHITECTURE.md`. + +- [Install](#install) +- [Quick start](#quick-start) +- [How it works](#how-it-works) +- [Public API](#public-api) +- [Callable contracts you implement](#callable-contracts-you-implement) +- [Config reference](#config-reference) +- [Wrapper pipeline semantics](#wrapper-pipeline-semantics) +- [Dedup and dedup_priority](#dedup-and-dedup_priority) +- [Redis keys and state](#redis-keys-and-state) +- [Error handling and resilience](#error-handling-and-resilience) +- [`_smartfeed_debug_info`](#_smartfeed_debug_info) +- [Integration checklist and footguns](#integration-checklist-and-footguns) +- [Testing](#testing) +- [Integration effort: honest read](#integration-effort-honest-read) + +## Install -## Quick Start +``` +pip install epoch8-smartfeed +``` + +Requires: Python 3.9+, **pydantic v2**, orjson. An async Redis client (`redis.asyncio.Redis`) is required only if any Wrapper uses `cache` or cross-page `dedup`; without it those features silently fall back to uncached behavior. + +## Quick start ```python from smartfeed.manager import FeedManager +from smartfeed.models.base import FeedResult config = { "version": "2", @@ -13,202 +39,320 @@ config = { "type": "wrapper", "node_id": "main", "cache": {"session_size": 300, "session_ttl": 300}, - "dedup": {"dedup_key": "id", "overfetch_factor": 3}, + "dedup": {"dedup_key": "id"}, "rerank": {"method_name": "my_rerank"}, + # NOTE: cache_key is a field on the wrapper, NOT inside "cache". + # "cache_key": "shared_pool", "data": { "type": "merger_percentage", "node_id": "mix", "items": [ {"percentage": 40, "data": {"type": "subfeed", "subfeed_id": "source_a", "method_name": "source_a"}}, - {"percentage": 60, "data": {"type": "subfeed", "subfeed_id": "source_b", "method_name": "source_b"}} - ] - } - } + {"percentage": 60, "data": {"type": "subfeed", "subfeed_id": "source_b", "method_name": "source_b"}}, + ], + }, + }, } +# A subfeed fetch function. **kwargs is REQUIRED (see Callable contracts). async def source_a(user_id, limit, next_page, **kwargs): - # Return FeedResult(data=[...], next_page={...}, has_next_page=True) - ... + # ... fetch your data ... + return FeedResult(data=[{"id": 1}, {"id": 2}], next_page={"page": 2}, has_next_page=True) + +async def source_b(user_id, limit, next_page, **kwargs): + return FeedResult(data=[{"id": 3}], next_page={}, has_next_page=False) +# A rerank function. MUST return exactly len(items) items (reorder only). async def my_rerank(items, session_id): - # Return same items in new order. len(result) == len(items). return sorted(items, key=lambda x: x.get("score", 0), reverse=True) -manager = FeedManager(config=config, methods_dict={"source_a": source_a, "source_b": source_b, "my_rerank": my_rerank}, redis_client=redis) -result = await manager.get_feed(session_id="user_123", limit=20, cursor={}) -# result.data -- list of items -# result.next_page -- pass back as cursor for next page -# result.has_next_page -- whether more pages available +manager = FeedManager( + config=config, + methods_dict={"source_a": source_a, "source_b": source_b, "my_rerank": my_rerank}, + redis_client=redis, # async redis.asyncio.Redis; needed for cache/dedup persistence +) + +result = await manager.get_feed(session_id="user_123", limit=20, cursor=None) +# result.data -> list of item dicts, each stamped with _smartfeed_debug_info +# result.next_page -> opaque cursor; pass back verbatim as `cursor` next call +# result.has_next_page -> bool ``` -## Node Types +Config is validated when `FeedManager` is constructed. A malformed config raises `pydantic.ValidationError` at construction, not on the first `get_feed`. -### SubFeed (leaf) +## How it works -Calls an async function from `methods_dict`. +A feed is a tree of nodes described in config. There are three node kinds (seven concrete `type` values): -```json -{"type": "subfeed", "subfeed_id": "tours", "method_name": "get_tours", "dedup_priority": 1} -``` +- **SubFeed** (`subfeed`) — a leaf. Calls one of your async functions from `methods_dict` and adapts its result into SmartFeed's shape, tagging each item with its source. +- **Wrapper** (`wrapper`) — a pipeline over one child. Optional stages applied in order: dedup, rerank, cache, paginate. Each stage is independent. +- **Mixer** — a combiner over several children, run concurrently and merged by a rule. Five concrete types: `merger_percentage`, `merger_append`, `merger_positional`, `merger_percentage_gradient`, `merger_distribute`. "Mixer" is a documentation umbrella, not a class. -Fields: -- `subfeed_id` -- unique ID, used in cursors and `_smartfeed_debug_info.source` -- `method_name` -- key in `methods_dict` -- `subfeed_params` -- static kwargs passed to method (default: `{}`) -- `raise_error` -- if `false`, swallow errors and return empty (default: `true`) -- `shuffle` -- shuffle results (default: `false`) -- `dedup_priority` -- higher = wins dedup conflicts (default: `0`) +The executor walks the tree. For each node it either calls `execute()` (SubFeed, Wrapper) or builds a mix plan and runs the children concurrently via `asyncio.gather`, then assembles them. Nodes nest arbitrarily. Each node paginates independently: the returned `next_page` cursor is a dict keyed per node, and you round-trip it verbatim. -### Wrapper (unified cache + dedup + rerank) +"Concurrent" here means `asyncio` concurrency on one event loop (sibling sources are awaited together, so page latency is close to the slowest source, not the sum). It is not multi-core. -Single node with optional pipeline stages: `fetch -> dedup -> rerank -> cache -> paginate`. +## Public API -```json -{ - "type": "wrapper", - "node_id": "pipeline", - "cache": {"session_size": 300, "session_ttl": 300}, - "dedup": {"dedup_key": "id", "overfetch_factor": 3, "max_refill_loops": 5}, - "rerank": {"method_name": "my_rerank", "raise_error": true}, - "dedup_priority": 3, - "data": { ... } -} +The entire public surface is one class with two methods. + +```python +FeedManager(config: dict, methods_dict: dict, redis_client: Optional[Any] = None) +``` +Parses and validates `config` into a `FeedConfig` immediately (pydantic v2). `methods_dict` maps names to your async fetch/rerank callables. `redis_client` must be an async `redis.asyncio.Redis`; if `None`, every Wrapper with `cache`/`dedup` degrades to uncached passthrough. Holds no per-request state. + +```python +async get_feed(session_id: str, limit: int, cursor: Optional[dict] = None) -> FeedResult ``` +`cursor=None` is treated as the first page (`{}`). Builds a fresh execution context per call and delegates to the executor, then stamps each dict item's final 0-based page position. Returns `FeedResult(data, next_page, has_next_page)`. `next_page` is opaque; store it and pass it back as `cursor`. -All stages are optional. Combinations: +There is **no request-parameters channel**: only `session_id`, `limit`, and the cursor are threaded to sources. Per-request context must be baked into `subfeed_params` when you build the config, or captured in closures when you build `methods_dict` for that request. -| cache | dedup | rerank | Behavior | -|-------|-------|--------|----------| -| yes | no | no | Session cache with Redis pagination | -| no | yes | no | Per-page dedup with Redis seen-set | -| yes | yes | yes | Full pipeline: fetch -> dedup -> rerank -> cache | -| yes | no | yes | Rerank + cache (no dedup) | -| no | no | yes | Per-page rerank (no cache, stateless) | +`FeedResult` (`smartfeed.models.base`): `data: list`, `next_page: dict`, `has_next_page: bool`. `data` is a bare list; the pipeline assumes dict items and silently skips non-dict items in stamping/dedup. -#### Cache +## Callable contracts you implement -```json -"cache": {"session_size": 300, "session_ttl": 300, "cache_key": null} +### SubFeed fetch function + +```python +async def fetch(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: ... ``` -- `session_size` -- items to fetch from child and cache -- `session_ttl` -- Redis TTL (seconds). Acts as inactivity timeout: refreshed on every access -- `cache_key` -- explicit key for shared cache between wrappers (A/B testing) +- **`**kwargs` is mandatory.** The executor injects `ctx=` into every call. A function without `**kwargs` raises `TypeError` on every invocation. (And if that source has `raise_error=False`, the `TypeError` is swallowed into an empty result, so a signature bug looks exactly like "this source had no items today".) +- `user_id` is the feed `session_id` (not necessarily an account id). +- `limit` is this leaf's computed demand for the page. Under a mixer it is a fraction of the page; under a dedup wrapper it is inflated by `overfetch_factor` or a refill deficit. +- `next_page` is this subfeed's own cursor slice (`cursor.get(subfeed_id, {})`), opaque. Echo back whatever you need on the next page. +- Return a `FeedResult` (or any object exposing `.data` / `.next_page` / `.has_next_page`; there is no `isinstance` check). +- On success each dict item is stamped with `_smartfeed_debug_info.source = subfeed_id`. Non-dict items are passed through untouched. -#### Dedup +### Rerank function -```json -"dedup": {"dedup_key": "id", "missing_key_policy": "error", "overfetch_factor": 3, "max_refill_loops": 5, "state_ttl": 300} +```python +async def rerank(items: list, session_id: str) -> list: ... ``` -- `dedup_key` -- field name for duplicate detection -- `missing_key_policy` -- `"error"` | `"keep"` | `"drop"` (default: `"error"`) -- `overfetch_factor` -- fetch `limit * factor` to compensate for dedup removals (default: `1`) -- `max_refill_loops` -- max refill iterations (default: `20`) -- `state_ttl` -- TTL for Redis seen-set (default: `300`) +- Must return a list of **exactly `len(items)`** items. Reordering or replacing is allowed; adding or dropping is not — a length mismatch always raises `ValueError`, even with `raise_error=False`. +- `items` are already fetched and deduped, in pre-rerank order. +- Exceptions thrown by your function are governed by `rerank.raise_error` (see [Error handling](#error-handling-and-resilience)). -Cross-page dedup: seen keys stored in Redis SET `sf:{session_id}:{node_id}:{hash}:seen`. +## Config reference -#### Rerank +All defaults below are byte-exact from the code. `FeedConfig` is the top-level shape: -```json -"rerank": {"method_name": "my_rerank", "raise_error": true} -``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `version` | `str` | required | Descriptive tag. Carried but never read by the code. | +| `feed` | `FeedNode` | required | Root node. Discriminated union on `type` over the 7 node types. | + +### SubFeed — `type: "subfeed"` + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `subfeed_id` | `str` | required | Logical name. Stamped as `_smartfeed_debug_info.source`, is the cursor namespace key, and the identity dedup priority resolves against. | +| `method_name` | `str` | required | Key into `methods_dict`. Looked up **before** the try block, so a missing/typo'd name raises `KeyError` even with `raise_error=False`. | +| `subfeed_params` | `dict` | `{}` | Static kwargs merged into the fetch call. A key colliding with `user_id`/`limit`/`next_page`/`ctx` raises `TypeError` at call time (not validated at parse time). | +| `raise_error` | `bool` | `True` | `True`: exceptions propagate. `False`: swallow, return an empty page. On the swallow path `next_page` is the full incoming cursor unchanged (not re-nested under `subfeed_id`). | +| `shuffle` | `bool` | `False` | Shuffle this source's page in place before stamping. Per-page only. | +| `dedup_priority` | `int` | `0` | See [dedup_priority](#dedup-and-dedup_priority). | + +### Wrapper — `type: "wrapper"` + +Wraps exactly one child with optional stages. `cache`, `dedup`, `rerank` are all independent and optional (all 8 on/off combinations parse). + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | Namespaces this wrapper's cursor, its cache keys (when `cache_key` unset), and its debug sub-dict. | +| `data` | node | required | The single child node. | +| `cache` | `WrapperCache?` | `None` | Enables per-session caching. Requires a live `redis_client` too, else passthrough. | +| `dedup` | `WrapperDedup?` | `None` | Enables duplicate removal with overfetch/refill. | +| `rerank` | `WrapperRerank?` | `None` | Enables a rerank pass (after fetch+dedup). | +| `cache_key` | `str?` | `None` | **Sibling of `cache`, not inside it.** Enables a shared base cache across wrappers with the same key (see [Wrapper pipeline](#wrapper-pipeline-semantics)). No-op unless `cache` is also set. | +| `dedup_priority` | `int` | `0` | See below. | + +`WrapperCache`: + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `session_size` | `int` | `300` | Items built and cached per session (cold build). Pages are sliced from this. | +| `session_ttl` | `int` | `300` | Redis TTL (s) on the data + `:meta` keys. Refreshed on every warm read — an inactivity timeout, not a fixed expiry. | -- `method_name` -- key in `methods_dict`, must be `async def(items, session_id) -> items` -- `raise_error` -- `true` = crash on rerank failure, `false` = keep original order (default: `true`) +`WrapperDedup`: -Contract: callable returns exactly `len(items)` items. Only reorders. +| Field | Type | Default | Meaning | +|---|---|---|---| +| `dedup_key` | `str` | required | Item field for identity. Compared as `str(item[key])`, so `1` and `"1"` collide. | +| `missing_key_policy` | `"error" \| "keep" \| "drop"` | `"error"` | When an item lacks `dedup_key`: `error` raises, `keep` passes it through un-deduped, `drop` discards it. | +| `overfetch_factor` | `int` | `4` | **Deprecated, no effect.** The dedup fetch now pulls exactly the outstanding deficit and refills until the target is filled or the child is exhausted. Kept for config back-compat. | +| `max_refill_loops` | `int` | `2` | **Deprecated, no effect.** Refill now continues until the page is full or the source runs out (bounded by an internal safety cap), so pages are not cut short. | +| `state_ttl` | `int` | `300` | TTL (s) on the Redis seen-set (session-scoped dedup state, both cache and passthrough paths). | -### Mixer Nodes +`WrapperRerank`: + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `method_name` | `str` | required | Key into `methods_dict`. `KeyError` if missing. | +| `raise_error` | `bool` | `True` | Governs only exceptions thrown by your rerank function. Does **not** govern the length-mismatch check, which always raises. | + +### Mixers + +**MergerPercentage** — `type: "merger_percentage"`. Splits the page by fixed percentages, then concatenates a solid block per branch in config order (it does not interleave). -**MergerPercentage** -- split by percentage: ```json -{"type": "merger_percentage", "node_id": "mix", "items": [{"percentage": 40, "data": {...}}, {"percentage": 60, "data": {...}}]} +{"type": "merger_percentage", "node_id": "mix", + "items": [{"percentage": 40, "data": {…}}, {"percentage": 60, "data": {…}}]} ``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | Prefix for internal child ids. | +| `items` | `list[{percentage:int, data:node}]` | required | Per-branch demand = `limit * percentage // 100`. Rounding remainder is topped up (largest-remainder) **only when the percentages sum to exactly 100**. `percentage` has no range validator; sum is not enforced. | +| `dedup_priority` | `int` | `0` | | + +There is no `[:limit]` trim: percentages summing over 100 overflow the page, under 100 underfill. + +**MergerAppend** — `type: "merger_append"`. Equal-share split, concatenated in order, then trimmed to `limit`. -**MergerPositional** -- insert at specific positions (1-indexed): ```json -{"type": "merger_positional", "node_id": "pos", "positions": [1, 3, 5, 7], "positional": {...}, "default": {...}} +{"type": "merger_append", "node_id": "append", "items": [{…}, {…}]} ``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | | +| `items` | `list[node]` | required | Each child gets `limit // n`; the remainder goes one-each to the first children. Empty `items` yields an empty page (no error). | +| `dedup_priority` | `int` | `0` | | + +**MergerPositional** — `type: "merger_positional"`. Places `positional` items at fixed slots, `default` items everywhere else. -**MergerAppend** -- concatenate children: ```json -{"type": "merger_append", "node_id": "append", "items": [{...}, {...}]} +{"type": "merger_positional", "node_id": "pos", + "positions": [1, 3, 5, 7], "positional": {…}, "default": {…}} ``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | | +| `positions` | `list[int]` | `[]` | 1-indexed slots, relative to **each page's** `1..limit` window (they repeat every page; they are not absolute feed positions). | +| `positional` | node | required | Fills the pinned slots. Demand = count of in-range positions. | +| `default` | node | required | Fills the rest. Demand = `limit - pos_count`, fixed regardless of actual returns. | +| `dedup_priority` | `int` | `0` | | + +If a source runs dry the other backfills; if the positional source is exhausted the page can underfill (default demand is fixed). + +**MergerPercentageGradient** — `type: "merger_percentage_gradient"`. A two-child percentage merger whose ratio shifts by `step` points every `size_to_step` cumulative items, tracked across pages via its own cursor counter. -**MergerAppendDistribute** -- round-robin by key: ```json -{"type": "merger_distribute", "node_id": "diverse", "distribution_key": "operator_id", "items": [{...}]} +{"type": "merger_percentage_gradient", "node_id": "grad", + "item_from": {"percentage": 80, "data": {…}}, "item_to": {"percentage": 20, "data": {…}}, + "step": 10, "size_to_step": 30} ``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | Also the merged-cursor key holding `{"page": n}`. | +| `item_from` | `{percentage:int, data:node}` | required | Starting branch. Its share falls toward 0. | +| `item_to` | `{percentage:int, data:node}` | required | Ending branch. Its share rises toward 100, then clamps. | +| `step` | `int` | required | Percentage points shifted per segment. | +| `size_to_step` | `int` | required | Items per shift. Used as a `range()` step; `0` raises at runtime, not at parse time. | +| `dedup_priority` | `int` | `0` | | + +**MergerAppendDistribute** — `type: "merger_distribute"` (note the literal). Round-robin by a key so no two adjacent items share it. -**MergerPercentageGradient** -- shift ratio over pages: ```json -{"type": "merger_percentage_gradient", "node_id": "grad", "item_from": {"percentage": 80, "data": {...}}, "item_to": {"percentage": 20, "data": {...}}, "step": 10, "size_to_step": 30} +{"type": "merger_distribute", "node_id": "diverse", + "distribution_key": "operator_id", "sorting_key": "score", "sorting_desc": true, "items": [{…}]} ``` +| Field | Type | Default | Meaning | +|---|---|---|---| +| `node_id` | `str` | required | | +| `items` | `list[node]` | required | **Every child is asked for the full `limit`** (up to `N * limit` fetched), then pooled and trimmed to `limit`. Higher backend load than items shown. | +| `distribution_key` | `str` | required | Group key. `entry[key]` with no missing-key policy — `KeyError` if an item lacks it. | +| `sorting_key` | `str?` | `None` | If set, the whole pool is sorted by this key before grouping. `KeyError` if missing. | +| `sorting_desc` | `bool` | `False` | Sort direction. Ignored if `sorting_key` is `None`. | +| `dedup_priority` | `int` | `0` | | -## dedup_priority +## Wrapper pipeline semantics -Every node has `dedup_priority: int = 0`. Higher = more important. +A Wrapper has two paths, chosen at runtime: -- When two items have the same dedup key, the one with higher priority wins -- Equal priority: first-seen wins (order from mixer children list) -- Wrapper/mixer with `dedup_priority != 0` overrides all children in that subtree -- Purely config-based, not stored in items +- **Cached** (`cache` set **and** `redis_client` present). First request for a session is a cold build: fetch `session_size` items, dedup, rerank, store the batch in Redis, serve page 1. Later pages slice out of the stored batch with no refetch. When the batch is exhausted the pipeline continues from where the child left off (it does not restart the source). Concurrent first requests for a session are serialized by a cold-build lock, so the child is fetched once. Rerank runs once per cold build over the whole batch, not per page. +- **Passthrough** (no `cache`, or no Redis). The pipeline reruns every page. -## _smartfeed_debug_info +**Dedup contract (session-scoped).** Within one continuous forward scroll, a user sees each item at most once — no matter how far they scroll, and across the internal rebuild boundaries the cursor crosses. Dedup is scoped to the session, not across sessions: re-sending an empty cursor (or page 1) starts a **fresh** scroll and may repeat items already shown in a previous scroll. Both paths honor this via a Redis seen-set (`:seen`): the passthrough path persists shown ids across pages; the cached path carries the seen-set across cold rebuilds and resets it when a fresh scroll begins (empty cursor). Neither path over-fetches, so no unique item is silently skipped, and refill fills each page to `limit` unless the source is genuinely exhausted. -SmartFeed stamps metadata on every item: +**Shared base cache (`cache_key`).** Two wrappers with different `node_id` but the same `cache_key` share one fetched+deduped base dataset (keyed by `cache_key` + the child's config hash + the continuation position), each applying its own rerank on top. Each page-window is its own shared segment, so a shared wrapper paginates past `session_size` normally. A distributed lock ensures only one caller does each cold fetch. Caveat: whichever wrapper wins the cold-build race bakes in **its** `session_size` and dedup settings for all sharers; a sibling with different settings reads that same pool. The shared base dedups within each segment; for strict cross-scroll uniqueness prefer a non-shared cached wrapper. Use `cache_key` for A/B testing two rerankers over one item pool; do not rely on per-sharer dedup/size differences. -```python -item["_smartfeed_debug_info"] = { - "source": "recommended_tours", # subfeed_id - "smartfeed_position": 5, # final position (set by FeedManager) - "strategy": "model_hot_users", # from subfeed (optional) - # Per-wrapper positions: - "pipeline": {"smartfeed_position": 15, "rerank_position": 3}, - # Rerank fields (set by rerank callable): - "rerank_position": 1, - "rrf_score": 0.032, - "feature_score": 10683198.0, -} -``` +## Dedup and dedup_priority -## Redis Keys +Every node has `dedup_priority: int = 0`. When a Wrapper's dedup sees the same key from two sources, the item whose source subtree has the higher effective priority wins; equal priority means first-seen wins (order follows the mixer's child list). Priority propagates down: a node with a non-zero `dedup_priority` overrides its whole subtree unless a descendant sets its own non-zero value. It is config-only and never written onto items. + +## Redis keys and state + +All SmartFeed state lives in Redis under `sf:{session_id}:...`. `config_hash` is the first 8 hex chars of an md5 over the node's sorted JSON dump, so any config change moves the key and old data ages out. ``` -sf:{session_id}:{node_id}:{config_hash} -- cached data -sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (child cursor, has_next, gen) -sf:{session_id}:{node_id}:{config_hash}:seen -- dedup seen-set (SET) +sf:{session_id}:{cache_key or node_id}:{config_hash} cached session batch (JSON) +sf:{session_id}:{cache_key or node_id}:{config_hash}:meta gen nonce, child cursor, child has_next +sf:{session_id}:{cache_key or node_id}:{config_hash}:coldlock cold-build lock (ttl 30s) +sf:{session_id}:{node_id}:{config_hash}:seen dedup seen-set (session-scoped; cache + passthrough) +sf:{session_id}:{cache_key}:{child_config_hash}:{segment} shared base segment (when cache_key set) +sf:{session_id}:{cache_key}:{child_config_hash}:{segment}:meta shared base segment meta +sf:{session_id}:{cache_key}:{child_config_hash}:{segment}:lock shared segment cold-build lock (ttl 30s) ``` -`config_hash` = md5 of config subtree JSON. Change any param -> different hash -> fresh cache. +(`{segment}` is `"0"` for page 1 and a short hash of the continuation cursor for later windows.) -TTL is an inactivity timeout: refreshed on every request. Cache lives while user scrolls. +TTL (`session_ttl`) is refreshed on every warm read, so an active user's cache does not expire mid-scroll. The `:seen` set has its own `state_ttl`. Each cold build stamps a random `gen` nonce into the cursor and `:meta`; if a returning cursor's `gen` no longer matches (cache expired or rebuilt), SmartFeed rebuilds cleanly from page 1 rather than resuming into an inconsistent state. -## Generation ID +Note the `next_page` cursor shape is config-dependent: a cached wrapper hides the child cursor behind `{node_id: {offset, gen}}` (an absolute offset, so a client may change `limit` between pages without skipping or repeating items), an uncached wrapper exposes the raw child cursor, and the gradient injects `{node_id: {page}}`. Adding or removing a `cache` block changes the cursor your clients round-trip. -Cached wrappers stamp a `gen` nonce in cursor and Redis `:meta`. If user returns with stale `gen` after cache expired -> full rebuild from page 1. +## Error handling and resilience -## Subfeed Method Signature +Error handling is **opt-in per node**. By default (`raise_error=True` everywhere) one failing source or rerank fails the entire `get_feed` call, including sibling branches that already succeeded — `asyncio.gather` runs without `return_exceptions=True`, and `FeedManager` has no try/except of its own. For fail-soft behavior, set `raise_error=False` on the specific SubFeed/rerank nodes. -```python -async def my_source(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResult: - return FeedResult(data=[...], next_page={"page": 2, "after": ...}, has_next_page=True) -``` +Four sharp edges an integrator will hit: -## Installation +1. **Missing `method_name`** is looked up before the try block, so it raises `KeyError` regardless of `raise_error`. +2. **Rerank length mismatch** always raises `ValueError`, regardless of `rerank.raise_error`. +3. **A subfeed function without `**kwargs`** raises `TypeError` on every call (the injected `ctx` has nowhere to go). With `raise_error=False` this is swallowed into an empty page — a signature bug that reads as "no items". +4. **`subfeed_params` key collisions** (with `user_id`/`limit`/`next_page`/`ctx`) raise `TypeError` at call time, with the same silent-swallow caveat under `raise_error=False`. -``` -pip install epoch8-smartfeed -``` +Silent behavior downgrades (no exception): `cache` without a `redis_client`; `cache_key` without `cache`; divergent shared-`cache_key` settings. These ship without error and only show up as degraded cache/dedup metrics. + +## `_smartfeed_debug_info` + +Every dict item carries a `_smartfeed_debug_info` bundle. The core writes these (all positions are **0-based** at runtime — the model comments saying "1-based" are stale): + +- `source` — the producing `subfeed_id`. +- `smartfeed_position` — final 0-based page position (set by `FeedManager`). Wrappers also stamp a pre-rerank position under their `node_id`. +- `rerank_position` — 0-based post-rerank index under the wrapper's `node_id`, only when rerank is configured. -Requires: Python 3.9+, Pydantic v2, redis (async), orjson. +Other documented fields (`strategy`, `rrf_score`, `feature_score`, `feature_position`, `total_reranked`, `raw_params`) are **conventions only**; the core never writes them. Your rerank callable may attach them (the bundle allows extra keys). `FeedItem` and `SmartFeedDebugInfo` are typing aids and are never instantiated at runtime — the pipeline operates on plain dicts. + +## Integration checklist and footguns + +What you must provide: + +- One async fetch function per `method_name`, each accepting `**kwargs` and returning a `FeedResult`. +- Any rerank functions, each returning exactly `len(items)` items. +- The `methods_dict` (a plain dict; no registration API). +- The config tree. pydantic checks types and shapes only; **semantic correctness is yours** — matching `dedup_key`s across siblings, percentages summing to 100, `cache_key` placed on the Wrapper, distribution/sorting keys present on items. +- An async `redis.asyncio.Redis` if you want caching or cross-page dedup persistence. +- Explicit `raise_error=False` on each node you want to fail soft (no global switch). +- Keep cached items orjson-serializable (`str/int/float/bool/None/list/dict`). A `Decimal` that works in passthrough throws once caching serializes the batch. + +Footguns, condensed: + +- `**kwargs` is load-bearing on every subfeed function (absorbs injected `ctx`). +- `cache_key` goes on the Wrapper, not inside `cache`, and is a no-op without `cache`. +- `cache` without Redis, and shared `cache_key` with divergent settings, degrade silently. +- Percentage mixers do not interleave and do not enforce sum==100; positional slots are per-page; distribute over-fetches and has no missing-key policy. +- Cursor shape depends on whether a `cache` block is present. ## Testing -```bash +- `fakeredis.aioredis.FakeRedis()` is the standard fixture; the whole integration is testable without a real Redis. +- `pytest` + `pytest-asyncio` (`@pytest.mark.asyncio`). Both are pinned dev deps. +- `tests/conftest.py` has a minimal `METHODS` dict (items / empty / error) as a starting point. +- There is no built-in validator for your subfeed signatures. Add a smoke test that runs each configured method through `get_feed` once before deploy — specifically to catch the `**kwargs`/`ctx` and `subfeed_params`-collision classes, which surface only at call time. + +``` pytest tests/ -v ``` -Requires: fakeredis, pytest-asyncio. diff --git a/pyproject.toml b/pyproject.toml index 077b7cb..ee6a9e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,8 +21,8 @@ black = "^23.3.0" fakeredis = "^2.34.1" isort = "^5.12.0" mypy = "^1.3.0" -pytest = "^7.3.1" -pytest-asyncio = "^0.21.0" +pytest = ">=8.2" +pytest-asyncio = ">=1.0" types-redis = "^4.5.5.2" [tool.black] diff --git a/smartfeed/execution/redis_lock.py b/smartfeed/execution/redis_lock.py index 86b6daf..dec85aa 100644 --- a/smartfeed/execution/redis_lock.py +++ b/smartfeed/execution/redis_lock.py @@ -32,5 +32,8 @@ async def __aexit__(self, *exc): if not self._owned: return val = await self._redis.get(self._key) - if val and val.decode() == self._token: - await self._redis.delete(self._key) + if val is not None: + # Redis clients configured with decode_responses=True return str, not bytes. + token = val.decode() if isinstance(val, bytes) else val + if token == self._token: + await self._redis.delete(self._key) diff --git a/smartfeed/models/subfeed.py b/smartfeed/models/subfeed.py index 7d8fa72..8b7a45f 100644 --- a/smartfeed/models/subfeed.py +++ b/smartfeed/models/subfeed.py @@ -36,7 +36,11 @@ async def execute( except Exception: if self.raise_error: raise - return FeedResult(data=[], next_page=cursor, has_next_page=False) + # Scope the failure cursor to THIS subfeed (its own, unadvanced position) + # so it retries next page without clobbering sibling cursors on merge. + return FeedResult( + data=[], next_page={self.subfeed_id: subfeed_cursor}, has_next_page=False + ) if self.shuffle: shuffle(result.data) diff --git a/smartfeed/models/wrapper.py b/smartfeed/models/wrapper.py index eb22179..febbe04 100644 --- a/smartfeed/models/wrapper.py +++ b/smartfeed/models/wrapper.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import hashlib +import json import secrets from typing import Any, Dict, List, Literal, Optional @@ -10,6 +12,11 @@ from smartfeed.execution import executor as _executor from .base import BaseNode, FeedResult, coerce_feed_node +# Safety bound for the dedup refill loop: stop scanning after this many fetched items +# per `target` when the child keeps yielding duplicates (guards against a misbehaving +# source that returns has_next=True forever). Well above any realistic duplicate rate. +_REFILL_SCAN_FACTOR = 50 + class WrapperCache(BaseModel): session_size: int = 300 @@ -24,6 +31,10 @@ class WrapperRerank(BaseModel): class WrapperDedup(BaseModel): dedup_key: str missing_key_policy: Literal["error", "keep", "drop"] = "error" + # Deprecated / no effect: the dedup fetch now pulls exactly the outstanding + # deficit and refills until the target is filled or the child is exhausted, so + # it neither over-fetches (no item loss) nor gives up early (no short pages). + # Kept for config back-compat. overfetch_factor: int = 4 max_refill_loops: int = 2 state_ttl: int = 300 # TTL for Redis seen-set (seconds) @@ -135,8 +146,10 @@ def _dedup(self, data: List, seen: Optional[set] = None) -> List: key_val = str(item[key_field]) - # Cross-page: already seen on previous pages - if key_val in seen: + # Cross-page / cross-rebuild: already shown earlier in this scroll. + # `key_val in batch` means it is a *within-batch* duplicate, which must + # fall through to priority arbitration below rather than be skipped here. + if key_val in seen and key_val not in batch: continue # In-batch priority arbitration @@ -207,33 +220,33 @@ async def _passthrough( cursor: dict, ctx: Any, ) -> FeedResult: - """No-cache path: fetch -> dedup (with overfetch/refill) -> rerank -> return.""" + """No-cache path: fetch -> dedup (refill to a full page) -> rerank -> return.""" if self.dedup: - # Load persisted seen-set from Redis (cross-page dedup) + # Cross-page seen-set from Redis (persists shown ids across pages). seen_keys = await self._load_seen_set(ctx, session_id) - overfetch = self.dedup.overfetch_factor - max_loops = self.dedup.max_refill_loops - fetch_size = limit * overfetch current_cursor = dict(cursor) - - result = await _executor.run(self.data, ctx, fetch_size, current_cursor) - data = self._dedup(result.data, seen_keys) - current_cursor = result.next_page - has_next = result.has_next_page - - loop = 0 - while len(data) < limit and has_next and loop < max_loops: - deficit = limit - len(data) - result = await _executor.run(self.data, ctx, deficit * overfetch, current_cursor) - data.extend(self._dedup(result.data, seen_keys)) + data: List = [] + has_next = True + scanned = 0 + scan_cap = max(limit, 1) * _REFILL_SCAN_FACTOR + + # Passthrough has no buffer, so it must never over-fetch: pull exactly the + # outstanding deficit each round and refill until the page is full or the + # child is exhausted. Nothing is discarded, so nothing is lost. Keep scanning + # through duplicate runs (don't stop on a single all-duplicate window), bounded + # by scan_cap so a duplicate-only source can't spin forever. + while len(data) < limit and has_next and scanned < scan_cap: + need = limit - len(data) + result = await _executor.run(self.data, ctx, need, current_cursor) current_cursor = result.next_page has_next = result.has_next_page - loop += 1 - - data = data[:limit] + if not result.data: + break # child yielded nothing -> stop instead of spinning + scanned += len(result.data) + data.extend(self._dedup(result.data, seen_keys)) - # Persist seen-set to Redis for next page + # Persist seen-set to Redis for the next page. new_keys = set() key_field = self.dedup.dedup_key for item in data: @@ -277,6 +290,12 @@ async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None ttl = self.dedup.state_ttl await ctx.redis.expire(key, ttl) + async def _reset_seen_set(self, ctx: Any, session_id: str) -> None: + """Clear the seen-set so a fresh scroll (empty cursor) starts from page 1.""" + if not ctx or not ctx.redis: + return + await ctx.redis.delete(self._seen_set_key(session_id)) + # -- cache helpers --------------------------------------------------------- def _base_key(self, session_id: str) -> str: @@ -336,18 +355,21 @@ async def _touch_ttl(self, ctx: Any, session_id: str) -> None: # -- pagination ------------------------------------------------------------ def _paginate( - self, data: List, limit: int, page: int, gen: str, + self, data: List, limit: int, offset: int, gen: str, child_has_next: bool = False, ) -> FeedResult: - """Slice cached data for the requested page and build cursor.""" - start = (page - 1) * limit - end = start + limit - page_data = data[start:end] + """Slice cached data at an absolute offset and build the next cursor. + + The cursor carries an absolute offset (not a page number), so pagination stays + correct even if the client varies `limit` between requests. + """ + end = offset + limit + page_data = data[offset:end] has_next = end < len(data) or child_has_next next_cursor = { self.node_id: { - "page": page + 1, + "offset": end, "gen": gen, } } @@ -371,7 +393,7 @@ async def _execute_with_cache( """Cached path: warm hit -> paginate, stale/miss -> cold build.""" my_cursor = cursor.get(self.node_id, {}) cursor_gen = my_cursor.get("gen") - cursor_page = my_cursor.get("page", 1) + cursor_offset = my_cursor.get("offset", 0) # Warm path: try to read existing cache if cursor_gen: @@ -379,17 +401,16 @@ async def _execute_with_cache( if meta and meta.get("gen") == cursor_gen: cached_data = await self._read_cache(ctx, session_id) if cached_data is not None: - # Check if we have enough data for this page - start = (cursor_page - 1) * limit - if start < len(cached_data): + # Serve from cache while this offset is still within the batch. + if cursor_offset < len(cached_data): await self._touch_ttl(ctx, session_id) child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) return self._paginate( - cached_data, limit, cursor_page, cursor_gen, + cached_data, limit, cursor_offset, cursor_gen, child_has_next=child_has_next, ) # Cache exhausted: rebuild with continuation cursor - return await self._cold_build( + return await self._cold_build_locked( methods_dict, session_id, limit, @@ -398,43 +419,92 @@ async def _execute_with_cache( ) # Cold path: no gen or stale gen -> fresh build - return await self._cold_build( + return await self._cold_build_locked( methods_dict, session_id, limit, ctx, child_cursor={} ) - # -- shared base cache helpers ---------------------------------------------- + async def _cold_build_locked( + self, + methods_dict: dict, + session_id: str, + limit: int, + ctx: Any, + child_cursor: Dict, + ) -> FeedResult: + """Cold build guarded by a lock so concurrent first requests fetch the child + ONCE: one caller builds and writes the cache, the rest wait and serve from it. + A fresh/continuation build always serves the new batch from offset 0.""" + from smartfeed.execution.redis_lock import RedisLock + + lock_key = f"{self._base_key(session_id)}:coldlock" + async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: + if acquired: + return await self._cold_build( + methods_dict, session_id, limit, ctx, child_cursor + ) + # Another coroutine is building -- wait for its cache, then serve page 1. + for _ in range(50): + await asyncio.sleep(0.1) + meta = await self._read_meta(ctx, session_id) + if meta: + cached = await self._read_cache(ctx, session_id) + if cached is not None: + child_has_next = meta.get( + "child_has_next", bool(meta.get("child_cursor")) + ) + return self._paginate( + cached, limit, 0, meta.get("gen", ""), + child_has_next=child_has_next, + ) + # Fallback: builder never wrote -> build ourselves. + return await self._cold_build( + methods_dict, session_id, limit, ctx, child_cursor + ) - def _base_shared_key(self, session_id: str) -> str: - """Key for the shared base cache (deduped, NOT reranked), keyed by cache_key.""" - return f"sf:{session_id}:{self.cache_key}:{self.data.config_hash()}" + # -- shared base cache helpers ---------------------------------------------- - async def _read_shared_base(self, ctx: Any, session_id: str) -> Optional[List]: - """Read shared base data from Redis. Returns None on miss.""" - key = self._base_shared_key(session_id) + @staticmethod + def _cursor_segment(child_cursor: Dict) -> str: + """Stable short tag for a continuation position, so each page-window of the + shared base is its own segment (page 1 = "0").""" + if not child_cursor: + return "0" + raw = json.dumps(child_cursor, sort_keys=True, default=str) + return hashlib.md5(raw.encode()).hexdigest()[:8] + + def _base_shared_key(self, session_id: str, child_cursor: Dict) -> str: + """Key for a shared-base segment (deduped, NOT reranked), keyed by cache_key + and the continuation position so continuations don't collide with page 1.""" + seg = self._cursor_segment(child_cursor) + return f"sf:{session_id}:{self.cache_key}:{self.data.config_hash()}:{seg}" + + async def _read_shared_base(self, ctx: Any, session_id: str, child_cursor: Dict) -> Optional[List]: + """Read a shared-base segment from Redis. Returns None on miss.""" + key = self._base_shared_key(session_id, child_cursor) raw = await ctx.redis.get(key) if raw is None: return None return orjson.loads(raw) async def _write_shared_base( - self, ctx: Any, session_id: str, data: List, child_cursor: Dict, - child_has_next: bool = False, + self, ctx: Any, session_id: str, child_cursor: Dict, data: List, + child_cursor_out: Dict, child_has_next: bool = False, ) -> None: - """Write shared base data and meta to Redis with TTL.""" - key = self._base_shared_key(session_id) + """Write a shared-base segment and its meta to Redis with TTL.""" + key = self._base_shared_key(session_id, child_cursor) ttl = self.cache.session_ttl pipe = ctx.redis.pipeline() pipe.set(key, orjson.dumps(data), ex=ttl) pipe.set( f"{key}:meta", - orjson.dumps({"child_cursor": child_cursor, "child_has_next": child_has_next}), + orjson.dumps({"child_cursor": child_cursor_out, "child_has_next": child_has_next}), ex=ttl, ) await pipe.execute() - async def _read_shared_base_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: - """Read shared base metadata (child_cursor, child_has_next) from Redis.""" - key = self._base_shared_key(session_id) + async def _read_shared_base_meta(self, ctx: Any, session_id: str, child_cursor: Dict) -> Optional[Dict]: + """Read a shared-base segment's metadata (child_cursor, child_has_next).""" + key = self._base_shared_key(session_id, child_cursor) raw = await ctx.redis.get(f"{key}:meta") if raw is None: return None @@ -445,38 +515,41 @@ async def _fetch_and_dedup( ctx: Any, target: int, child_cursor: Dict, + seen: Optional[set] = None, ) -> tuple: - """Fetch target items from child with overfetch + refill loop for dedup. + """Collect up to `target` deduped items from the child. + + Fetches exactly the outstanding deficit each round and keeps EVERY survivor + (never truncates, never advances the child cursor past an unconsumed unique), + so no item is lost. Refills until `target` unique items are collected or the + child is exhausted, so pages fill under dedup attrition instead of coming up + short. `seen`, if given, holds keys already shown earlier in this scroll (so + they are not repeated across a rebuild boundary); `_dedup` mutates it in place + with the keys emitted here. Returns (data, child_cursor, child_has_next). """ - overfetch = self.dedup.overfetch_factor if self.dedup else 1 - max_loops = self.dedup.max_refill_loops if self.dedup else 0 - fetch_size = target * overfetch cursor = dict(child_cursor) - - # First fetch - result = await _executor.run(self.data, ctx, fetch_size, cursor) - data = self._dedup(result.data) if self.dedup else result.data - cursor = result.next_page - has_next = result.has_next_page - - # Refill loop: keep fetching if dedup ate too many items. - # Accumulate seen keys so refill batches are deduped against all prior items. - loop = 0 - while len(data) < target and has_next and loop < max_loops: - deficit = target - len(data) - refill_size = deficit * overfetch - result = await _executor.run(self.data, ctx, refill_size, cursor) - # Dedup refill against ALL accumulated data - combined = data + result.data - combined = self._dedup(combined) if self.dedup else combined - data = combined + data: List = [] + has_next = True + scanned = 0 + scan_cap = max(target, 1) * _REFILL_SCAN_FACTOR + + # Keep scanning through duplicate runs until `target` unique items are + # collected or the child is exhausted; bounded by scan_cap so a + # duplicate-only source can't spin forever. + while len(data) < target and has_next and scanned < scan_cap: + need = target - len(data) + result = await _executor.run(self.data, ctx, need, cursor) cursor = result.next_page has_next = result.has_next_page - loop += 1 + if not result.data: + break # child yielded nothing -> stop instead of spinning + scanned += len(result.data) + new = self._dedup(result.data, seen) if self.dedup else result.data + data.extend(new) - return data[:target], cursor, has_next + return data, cursor, has_next async def _cold_build( self, @@ -491,12 +564,12 @@ async def _cold_build( if self.cache_key is not None: # Shared cache path: base data (fetch + dedup) is shared across wrappers - # with the same cache_key; rerank is applied per-wrapper after. + # with the same cache_key; rerank is applied per-wrapper after. Each + # continuation position is a distinct shared segment (see _base_shared_key). base_data = await self._build_shared_base( methods_dict, session_id, ctx, child_cursor ) - # Read child_cursor from shared base meta for continuation - shared_meta = await self._read_shared_base_meta(ctx, session_id) + shared_meta = await self._read_shared_base_meta(ctx, session_id, child_cursor) child_cursor_out = shared_meta.get("child_cursor", {}) if shared_meta else {} child_has_next = shared_meta.get("child_has_next", False) if shared_meta else False # Per-wrapper: stamp pre-rerank, apply rerank, stamp post-rerank @@ -506,10 +579,21 @@ async def _cold_build( if self.rerank: self._stamp_post_rerank(data) else: - # Standard path: fetch -> dedup (with overfetch/refill) -> rerank -> stamp -> cache + # Standard path. Session-scoped seen-set: reset on a fresh scroll (empty + # cursor) so re-requesting page 1 returns page 1, and carry it across + # rebuilds so a shown id never repeats later in the same scroll. + seen: Optional[set] = None + if self.dedup: + if not child_cursor: + await self._reset_seen_set(ctx, session_id) + seen = set() + else: + seen = await self._load_seen_set(ctx, session_id) data, child_cursor_out, child_has_next = await self._fetch_and_dedup( - ctx, session_size, child_cursor + ctx, session_size, child_cursor, seen=seen ) + if self.dedup: + await self._save_seen_set(ctx, session_id, seen) self._stamp_pre_rerank(data) data = await self._apply_rerank(data, methods_dict, session_id) if self.rerank: @@ -517,7 +601,7 @@ async def _cold_build( gen = secrets.token_hex(4) await self._write_cache(ctx, session_id, data, gen, child_cursor_out, child_has_next) - return self._paginate(data, limit, 1, gen, child_has_next=child_has_next) + return self._paginate(data, limit, 0, gen, child_has_next=child_has_next) async def _build_shared_base( self, @@ -531,10 +615,10 @@ async def _build_shared_base( from smartfeed.execution.redis_lock import RedisLock session_size = self.cache.session_size - lock_key = f"{self._base_shared_key(session_id)}:lock" + lock_key = f"{self._base_shared_key(session_id, child_cursor)}:lock" - # Fast path: base already cached - existing = await self._read_shared_base(ctx, session_id) + # Fast path: this segment already cached + existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing @@ -542,24 +626,26 @@ async def _build_shared_base( async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: if acquired: # Double-check after acquiring lock - existing = await self._read_shared_base(ctx, session_id) + existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing base_data, child_cursor_out, child_has_next = await self._fetch_and_dedup( ctx, session_size, child_cursor ) - await self._write_shared_base(ctx, session_id, base_data, child_cursor_out, child_has_next) + await self._write_shared_base( + ctx, session_id, child_cursor, base_data, child_cursor_out, child_has_next + ) return base_data else: - # Another coroutine holds the lock -- poll until base cache appears + # Another coroutine holds the lock -- poll until the segment appears for _ in range(50): await asyncio.sleep(0.1) - existing = await self._read_shared_base(ctx, session_id) + existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing # Fallback: fetch ourselves if lock holder never wrote - existing = await self._read_shared_base(ctx, session_id) + existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing data, _, _ = await self._fetch_and_dedup(ctx, session_size, child_cursor) diff --git a/tests/sources.py b/tests/sources.py new file mode 100644 index 0000000..c25561d --- /dev/null +++ b/tests/sources.py @@ -0,0 +1,205 @@ +"""Shared test harness for dedup / pagination correctness. + +The existing test sources are mostly non-stateful (they ignore ``next_page`` and +re-emit the same page forever), which is exactly why the dedup/pagination bugs +were not caught. This module provides: + + * ``ScriptedSource`` -- a *stateful*, offset-paginated source over a FINITE + pool. The pool may contain repeated ids (cross-page duplicates). It records + every id it emits and how many times it was called, so tests can assert + coverage ("no unique item was lost") and fetch counts ("cold build fetched + once"). Finiteness is what lets us prove "no loss": drain to exhaustion and + compare the union of pages against the known unique set. + + * pool factories (unique / duplicate-injected / overlapping). + + * a ``drain`` helper that pages to exhaustion with an infinite-loop guard, plus + invariant assertions (coverage / no-duplicates / page-fullness). + +Cursor convention for ScriptedSource: ``{"pos": }``. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, Optional + +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperDedup, WrapperRerank +from smartfeed.execution.context import ExecutionContext +from smartfeed.execution import executor as run_executor + + +# --------------------------------------------------------------------------- +# Source +# --------------------------------------------------------------------------- + +class ScriptedSource: + """Stateful, offset-paginated source over a finite pool of item dicts.""" + + def __init__(self, pool: List[Dict], latency: float = 0.0, id_key: str = "id"): + self.pool = [dict(it) for it in pool] + self.latency = latency + self.id_key = id_key + self.calls = 0 + self.emitted: List[Any] = [] + + @property + def unique_ids(self) -> set: + return {it[self.id_key] for it in self.pool if self.id_key in it} + + async def __call__(self, user_id, limit, next_page, **kw) -> FeedResult: + self.calls += 1 + pos = (next_page or {}).get("pos", 0) + if self.latency: + await asyncio.sleep(self.latency) + chunk = self.pool[pos: pos + limit] + # Fresh copies so downstream in-place stamping never mutates the pool. + data = [dict(it) for it in chunk] + for it in data: + self.emitted.append(it.get(self.id_key)) + new_pos = pos + len(chunk) + return FeedResult( + data=data, + next_page={"pos": new_pos}, + has_next_page=new_pos < len(self.pool), + ) + + +# --------------------------------------------------------------------------- +# Pool factories (all deterministic -- no RNG, so failures are reproducible) +# --------------------------------------------------------------------------- + +def unique_pool(n: int, start: int = 0, val_prefix: str = "u") -> List[Dict]: + """Ids start..start+n-1, no duplicates. Isolates *item loss* from dedup.""" + return [{"id": start + k, "val": f"{val_prefix}{start + k}"} for k in range(n)] + + +def dup_pool(n_unique: int, every: int = 4) -> List[Dict]: + """n_unique distinct ids, with an earlier id re-inserted every ``every`` slots. + + The re-inserted copies are *cross-page* duplicates once paginated. The unique + universe is still exactly range(n_unique), so ``assert_full_coverage`` derives + expectations straight from the pool. + """ + pool: List[Dict] = [] + for k in range(n_unique): + pool.append({"id": k, "val": f"v{k}"}) + if k >= 2 and k % every == 0: + dup_id = k - 2 + pool.append({"id": dup_id, "val": f"v{dup_id}"}) # re-emit an earlier id + return pool + + +def adjacent_dup_pool(n_unique: int) -> List[Dict]: + """Each id emitted twice, back-to-back -- both copies land in the same batch, + so within-batch dedup must catch them (no cross-rebuild ambiguity).""" + pool: List[Dict] = [] + for k in range(n_unique): + pool.append({"id": k, "val": f"v{k}"}) + pool.append({"id": k, "val": f"v{k}"}) + return pool + + +# --------------------------------------------------------------------------- +# Context / node builders +# --------------------------------------------------------------------------- + +def make_ctx(methods: Dict, redis=None, session_id: str = "s1") -> ExecutionContext: + return ExecutionContext(session_id=session_id, methods_dict=methods, redis=redis) + + +def subfeed(subfeed_id: str, method_name: str, **kw) -> SubFeed: + return SubFeed(subfeed_id=subfeed_id, method_name=method_name, **kw) + + +def wrapper( + child, + *, + node_id: str = "w", + session_size: Optional[int] = None, + session_ttl: int = 300, + dedup_key: Optional[str] = None, + overfetch_factor: int = 4, + max_refill_loops: int = 2, + missing_key_policy: str = "error", + rerank_method: Optional[str] = None, + rerank_raise: bool = True, + cache_key: Optional[str] = None, +) -> Wrapper: + """Build a Wrapper with the requested pipeline stages.""" + cache = None + if session_size is not None: + cache = WrapperCache(session_size=session_size, session_ttl=session_ttl) + dedup = None + if dedup_key is not None: + dedup = WrapperDedup( + dedup_key=dedup_key, + overfetch_factor=overfetch_factor, + max_refill_loops=max_refill_loops, + missing_key_policy=missing_key_policy, + ) + rerank = None + if rerank_method is not None: + rerank = WrapperRerank(method_name=rerank_method, raise_error=rerank_raise) + return Wrapper( + node_id=node_id, + cache=cache, + dedup=dedup, + rerank=rerank, + cache_key=cache_key, + data=child, + ) + + +# --------------------------------------------------------------------------- +# Drain + invariant assertions +# --------------------------------------------------------------------------- + +async def drain(node, ctx, limit: int, max_pages: int = 500, cursor=None) -> List[List]: + """Page until ``has_next_page`` is False. Returns the list of page-data lists. + + Raises AssertionError if pagination does not terminate within ``max_pages`` -- + this is the guard that catches infinite-loop bugs (e.g. shared-cache + continuation re-serving page 1 forever). + """ + cursor = cursor or {} + pages: List[List] = [] + for _ in range(max_pages): + r = await run_executor.run(node, ctx, limit=limit, cursor=cursor) + pages.append(r.data) + if not r.has_next_page: + return pages + cursor = r.next_page + raise AssertionError( + f"pagination did not terminate within {max_pages} pages " + f"(infinite loop / non-advancing cursor)" + ) + + +def flat_ids(pages: List[List], key: str = "id") -> List: + return [it[key] for pg in pages for it in pg if isinstance(it, dict) and key in it] + + +def assert_no_duplicates(pages: List[List], key: str = "id") -> None: + ids = flat_ids(pages, key) + seen, dupes = set(), set() + for x in ids: + (dupes if x in seen else seen).add(x) + assert not dupes, f"{len(dupes)} id(s) served more than once: {sorted(dupes)[:20]}" + + +def assert_full_coverage(pages: List[List], expected: set, key: str = "id") -> None: + got = set(flat_ids(pages, key)) + missing = set(expected) - got + extra = got - set(expected) + assert not missing, f"{len(missing)} unique id(s) NEVER served (lost): {sorted(missing)[:20]}" + assert not extra, f"served id(s) outside the source universe: {sorted(extra)[:20]}" + + +def assert_pages_full(pages: List[List], limit: int) -> None: + """Every page except the last must be exactly ``limit`` (dedup must refill, + not short-change, until the source is genuinely exhausted).""" + for i, pg in enumerate(pages[:-1]): + assert len(pg) == limit, f"page {i} returned {len(pg)} items, expected {limit}" diff --git a/tests/test_bug_cold_build_stampede.py b/tests/test_bug_cold_build_stampede.py new file mode 100644 index 0000000..1a354e7 --- /dev/null +++ b/tests/test_bug_cold_build_stampede.py @@ -0,0 +1,64 @@ +"""BUG #5 -- cold-build thundering herd on the standard (non-shared) cached path. + +Only the cache_key path takes a RedisLock. For a normal cached wrapper, N +concurrent first requests for one session each fetch the full session_size from +the child and each write a different gen; last writer wins, so the losers' gens +no longer match meta and they rebuild again on their next page. Upstream +amplification + cache thrash. + +The source has latency so all N callers enter the cold build before any of them +writes the cache -- a deterministic race. Fails today (N fetches); passes once +the standard path is guarded like the shared path. +""" + +import asyncio + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +N = 5 + + +@pytest.mark.asyncio +async def test_concurrent_cold_build_fetches_child_once(): + src = S.ScriptedSource(S.unique_pool(200), latency=0.05) + node = S.wrapper(S.subfeed("src", "src"), session_size=50) + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + results = await asyncio.gather(*[ + run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) + ]) + + assert all(len(r.data) == 10 for r in results) + # One winner should fetch the child; the rest should read the cache it wrote. + assert src.calls == 1, f"expected 1 cold fetch, got {src.calls} (thundering herd)" + + +@pytest.mark.asyncio +async def test_concurrent_cold_build_single_generation(): + src = S.ScriptedSource(S.unique_pool(200), latency=0.05) + node = S.wrapper(S.subfeed("src", "src"), session_size=50) + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + results = await asyncio.gather(*[ + run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) + ]) + gens = {r.next_page["w"]["gen"] for r in results} + assert len(gens) == 1, f"concurrent callers got {len(gens)} different generations: {gens}" + + +@pytest.mark.asyncio +async def test_shared_cache_path_already_locks_control(): + """Control: the cache_key path DOES take a lock, so it fetches once.""" + src = S.ScriptedSource(S.unique_pool(200), latency=0.05) + node = S.wrapper(S.subfeed("src", "src"), node_id="w", + session_size=50, cache_key="shared") + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + await asyncio.gather(*[ + run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) + ]) + assert src.calls == 1, f"shared path should fetch once, got {src.calls}" diff --git a/tests/test_bug_overfetch_item_loss.py b/tests/test_bug_overfetch_item_loss.py new file mode 100644 index 0000000..11529fc --- /dev/null +++ b/tests/test_bug_overfetch_item_loss.py @@ -0,0 +1,76 @@ +"""BUG #1 -- dedup silently drops the overfetched remainder. + +Both dedup paths fetch ``target * overfetch_factor`` items, keep ``target``, and +return the child cursor advanced past *all* fetched items. The surviving-but- +truncated items are consumed from the source and never served on any page. + + * _fetch_and_dedup (cached): wrapper.py:456 (fetch_size) + :479 (data[:target]) + * _passthrough (no cache): wrapper.py:217 (fetch_size) + :234 (data[:limit]) + +These tests drain a finite source and assert full coverage. They fail today +(~75% of items lost at overfetch=4) and must pass once the remainder is +preserved instead of discarded. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +@pytest.mark.asyncio +async def test_cached_dedup_covers_all_unique_items(): + src = S.ScriptedSource(S.unique_pool(200)) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, + dedup_key="id", overfetch_factor=4) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=200) + S.assert_full_coverage(pages, set(range(200))) + + +@pytest.mark.asyncio +async def test_passthrough_dedup_covers_all_unique_items(): + src = S.ScriptedSource(S.unique_pool(200)) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=4) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=200) + S.assert_full_coverage(pages, set(range(200))) + + +@pytest.mark.asyncio +async def test_cached_dedup_covers_all_with_real_duplicates(): + """With a source that genuinely contains duplicates, dedup should remove the + dupes but still surface every *unique* id. Coverage = the unique universe.""" + pool = S.dup_pool(150, every=4) # unique universe is range(150), plus cross-page dupes + src = S.ScriptedSource(pool) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, + dedup_key="id", overfetch_factor=4) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_full_coverage(pages, set(range(150))) + + +@pytest.mark.asyncio +async def test_overfetch_1_is_lossless_control(): + """Control: with overfetch_factor=1 there is no discard, so no loss. Proves the + harness is sound and localizes the bug to overfetch>1.""" + src = S.ScriptedSource(S.unique_pool(200)) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, + dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_full_coverage(pages, set(range(200))) + + +@pytest.mark.asyncio +async def test_no_dedup_is_lossless_control(): + """Control: a plain cached wrapper (no dedup) never overfetches, never loses.""" + src = S.ScriptedSource(S.unique_pool(200)) + node = S.wrapper(S.subfeed("src", "src"), session_size=20) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=200) + S.assert_full_coverage(pages, set(range(200))) diff --git a/tests/test_bug_redis_lock_decode.py b/tests/test_bug_redis_lock_decode.py new file mode 100644 index 0000000..ffec6bf --- /dev/null +++ b/tests/test_bug_redis_lock_decode.py @@ -0,0 +1,64 @@ +"""BUG #4 -- RedisLock assumes bytes; crashes under decode_responses=True. + +redis_lock.py:35 does ``val.decode()`` in __aexit__. With a Redis client +configured ``decode_responses=True`` (a very common setup), the reply is already +str and release raises AttributeError, breaking the shared-cache cold build. The +rest of the codebase (seen-set load, wrapper.py:269) already handles both +bytes/str, so this is an internal inconsistency. + +Also included: a green characterization test showing the token check itself is +correct (a lock whose value was overwritten by another owner is NOT deleted on +release). The residual issue there is the non-atomic get-then-delete (TOCTOU), +which cannot be reproduced deterministically without concurrency injection and is +documented in the review rather than asserted here. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution.redis_lock import RedisLock +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_lock_release_with_decode_responses(): + redis = fakeredis.aioredis.FakeRedis(decode_responses=True) + async with RedisLock(redis, "k", ttl=5) as acquired: + assert acquired is True + # Release must have deleted the key without raising. + assert not await redis.exists("k") + + +@pytest.mark.asyncio +async def test_shared_cold_build_with_decode_responses(): + redis = fakeredis.aioredis.FakeRedis(decode_responses=True) + src = S.ScriptedSource(S.unique_pool(100)) + node = S.wrapper(S.subfeed("src", "src"), node_id="w", + session_size=20, cache_key="shared") + ctx = S.make_ctx({"src": src}, redis=redis) + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r.data) == 10 + + +@pytest.mark.asyncio +async def test_lock_bytes_mode_is_fine_control(): + redis = fakeredis.aioredis.FakeRedis() # bytes mode (default) + async with RedisLock(redis, "k", ttl=5) as acquired: + assert acquired is True + assert not await redis.exists("k") + + +@pytest.mark.asyncio +async def test_lock_does_not_delete_foreign_token(): + """Characterization: if another owner overwrote the key, release leaves it. + (The token compare-and-delete is correct in the non-racy case.)""" + redis = fakeredis.aioredis.FakeRedis() + lock = RedisLock(redis, "k", ttl=30) + acquired = await lock.__aenter__() + assert acquired is True + # Simulate expiry + re-acquire by a different owner. + await redis.set("k", b"someone-elses-token") + await lock.__aexit__(None, None, None) + # Foreign token must survive -- we must not delete a lock we no longer own. + assert await redis.get("k") == b"someone-elses-token" diff --git a/tests/test_bug_shared_cache_continuation.py b/tests/test_bug_shared_cache_continuation.py new file mode 100644 index 0000000..b863d2c --- /dev/null +++ b/tests/test_bug_shared_cache_continuation.py @@ -0,0 +1,64 @@ +"""BUG #2 -- shared cache (cache_key) cannot paginate past session_size. + +When a cache_key wrapper exhausts its per-wrapper cache, the warm path calls +_cold_build with the continuation child_cursor, but _build_shared_base +(wrapper.py:536-539) fast-paths on the already-cached base and returns it, +ignoring the continuation cursor. Every "continuation" re-serves the first +session_size items under a fresh gen -> the feed loops on page 1 forever. + +These tests fail today (non-termination / repeats) and must pass once the shared +base supports continuation. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +def _shared_wrapper(): + return S.wrapper(S.subfeed("src", "src"), node_id="w", + session_size=20, cache_key="shared") + + +@pytest.mark.asyncio +async def test_shared_cache_paginates_to_exhaustion(): + src = S.ScriptedSource(S.unique_pool(100)) + node = _shared_wrapper() + ctx = S.make_ctx({"src": src}, redis=_redis()) + # Guard is low: a looping feed re-serves the same 20 items and never ends. + pages = await S.drain(node, ctx, limit=10, max_pages=40) + S.assert_no_duplicates(pages) + S.assert_full_coverage(pages, set(range(100))) + + +@pytest.mark.asyncio +async def test_shared_cache_continuation_advances_past_session(): + """Page 3 (first page after the 20-item session is exhausted) must contain new + ids, not repeat page 1.""" + src = S.ScriptedSource(S.unique_pool(100)) + node = _shared_wrapper() + ctx = S.make_ctx({"src": src}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + r3 = await run_executor.run(node, ctx, limit=10, cursor=r2.next_page) + ids_seen = {it["id"] for it in r1.data} | {it["id"] for it in r2.data} + ids_p3 = {it["id"] for it in r3.data} + assert not (ids_seen & ids_p3), f"page 3 repeats earlier ids: {sorted(ids_seen & ids_p3)}" + + +@pytest.mark.asyncio +async def test_shared_cache_within_session_is_fine_control(): + """Control: within the first session_size items, shared cache paginates correctly.""" + src = S.ScriptedSource(S.unique_pool(100)) + node = _shared_wrapper() + ctx = S.make_ctx({"src": src}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + ids = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] + assert ids == list(range(20)) diff --git a/tests/test_bug_softfail_cursor_leak.py b/tests/test_bug_softfail_cursor_leak.py new file mode 100644 index 0000000..32a8f2b --- /dev/null +++ b/tests/test_bug_softfail_cursor_leak.py @@ -0,0 +1,76 @@ +"""BUG #6 -- a soft-failed subfeed returns the whole inbound cursor. + +On raise_error=False, SubFeed.execute returns ``next_page=cursor`` -- the entire +feed cursor -- instead of ``{subfeed_id: ...}`` (subfeed.py:39). _merge_cursor +(mixers.py:12-16) does ``merged.update(c)`` per child in child order, so a failed +child that sorts AFTER a healthy sibling dumps all inbound keys and overwrites +the sibling's fresh cursor with the stale inbound value. + +Fails today; passes once the failure path returns a cursor scoped to its own id. +""" + +import pytest + +from tests import sources as S +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.mixers import MergerAppend +from smartfeed.execution import executor as run_executor + + +async def good_source(user_id, limit, next_page, **kw): + pos = next_page.get("pos", 0) + data = [{"id": pos + i} for i in range(limit)] + return FeedResult(data=data, next_page={"pos": pos + limit}, has_next_page=True) + + +async def failing_source(user_id, limit, next_page, **kw): + raise RuntimeError("upstream down") + + +METHODS = {"good": good_source, "bad": failing_source} + + +def _node(): + # child order [good, bad] so the failed child's cursor merges LAST. + return MergerAppend(node_id="top", items=[ + SubFeed(subfeed_id="good_src", method_name="good"), + SubFeed(subfeed_id="bad_src", method_name="bad", raise_error=False), + ]) + + +@pytest.mark.asyncio +async def test_softfail_does_not_clobber_sibling_cursor(): + ctx = S.make_ctx(METHODS, redis=None) + # Inbound cursor carries a STALE position for good_src. + inbound = {"good_src": {"pos": 5}} + r = await run_executor.run(_node(), ctx, limit=10, cursor=inbound) + # good_src is asked for demand 5 (MergerAppend split of limit 10) from pos 5, + # so its fresh cursor is pos 10. The failed sibling must not overwrite it. + assert r.next_page.get("good_src", {}).get("pos") == 10, ( + f"good_src cursor was clobbered by the failed sibling: {r.next_page}" + ) + + +@pytest.mark.asyncio +async def test_softfail_cursor_is_scoped_to_own_id(): + ctx = S.make_ctx(METHODS, redis=None) + inbound = {"good_src": {"pos": 5}, "unrelated": {"x": 1}} + r = await run_executor.run(_node(), ctx, limit=10, cursor=inbound) + # The failed subfeed must not resurrect unrelated inbound cursor keys. + assert "unrelated" not in r.next_page, ( + f"failed subfeed leaked unrelated inbound cursor keys: {r.next_page}" + ) + + +@pytest.mark.asyncio +async def test_softfail_clobber_reserves_items_downstream(): + """The user-facing consequence: because the failed sibling rewinds good_src's + cursor, a subsequent page re-serves items already shown.""" + ctx = S.make_ctx(METHODS, redis=None) + r1 = await run_executor.run(_node(), ctx, limit=10, cursor={}) + r2 = await run_executor.run(_node(), ctx, limit=10, cursor=r1.next_page) + r3 = await run_executor.run(_node(), ctx, limit=10, cursor=r2.next_page) + ids = [it["id"] for pg in (r1.data, r2.data, r3.data) for it in pg] + dupes = sorted({x for x in ids if ids.count(x) > 1}) + assert not dupes, f"items re-served across pages due to cursor clobber: {dupes}" diff --git a/tests/test_bug_variable_page_size.py b/tests/test_bug_variable_page_size.py new file mode 100644 index 0000000..40a1c74 --- /dev/null +++ b/tests/test_bug_variable_page_size.py @@ -0,0 +1,61 @@ +"""BUG #3 -- changing `limit` between pages on a cached wrapper skips/repeats. + +The cursor stores a page *number*; the offset is recomputed as (page-1)*limit at +read time (_paginate, wrapper.py:343). If the client changes page size between +requests, the page-number -> offset mapping is inconsistent and items are +skipped (growing limit) or repeated (shrinking limit). + +Fails today; passes once the cursor tracks an absolute offset (or page size is +pinned into the cursor). +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +def _cached(): + # session_size large enough that both pages come from one cached batch. + return S.wrapper(S.subfeed("src", "src"), session_size=200) + + +@pytest.mark.asyncio +async def test_growing_limit_does_not_skip(): + src = S.ScriptedSource(S.unique_pool(200)) + node = _cached() + ctx = S.make_ctx({"src": src}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) # expect ids 0..9 + r2 = await run_executor.run(node, ctx, limit=50, cursor=r1.next_page) # expect 10..59 + served = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] + assert served == list(range(60)), f"expected contiguous 0..59, got gaps: {served[:15]}..." + + +@pytest.mark.asyncio +async def test_shrinking_limit_does_not_repeat(): + src = S.ScriptedSource(S.unique_pool(200)) + node = _cached() + ctx = S.make_ctx({"src": src}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=50, cursor={}) # ids 0..49 + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) # expect 50..59 + served = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] + assert len(served) == len(set(served)), ( + f"repeated ids on shrink: {sorted({x for x in served if served.count(x) > 1})}" + ) + + +@pytest.mark.asyncio +async def test_constant_limit_is_fine_control(): + """Control: with a constant limit the page cursor works correctly.""" + src = S.ScriptedSource(S.unique_pool(200)) + node = _cached() + ctx = S.make_ctx({"src": src}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + served = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] + assert served == list(range(20)) diff --git a/tests/test_config_parsing.py b/tests/test_config_parsing.py index 3ebe991..f9db222 100644 --- a/tests/test_config_parsing.py +++ b/tests/test_config_parsing.py @@ -23,7 +23,7 @@ def test_config_hash_changes_with_params(): def test_parse_simple_subfeed(): - cfg = FeedConfig.parse_obj({ + cfg = FeedConfig.model_validate({ "version": "2", "feed": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, }) @@ -31,7 +31,7 @@ def test_parse_simple_subfeed(): def test_parse_wrapper_with_cache(): - cfg = FeedConfig.parse_obj({ + cfg = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "w", @@ -43,7 +43,7 @@ def test_parse_wrapper_with_cache(): def test_parse_nested_config(): - cfg = FeedConfig.parse_obj({ + cfg = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "outer", diff --git a/tests/test_cursors.py b/tests/test_cursors.py index 0e3bb42..06b23ec 100644 --- a/tests/test_cursors.py +++ b/tests/test_cursors.py @@ -19,7 +19,7 @@ def ctx(redis): @pytest.mark.asyncio async def test_cached_wrapper_hides_child_cursors(ctx): - config = FeedConfig.parse_obj({ + config = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "cached", @@ -39,7 +39,7 @@ async def test_cached_wrapper_hides_child_cursors(ctx): @pytest.mark.asyncio async def test_uncached_wrapper_exposes_child_cursors(ctx): - config = FeedConfig.parse_obj({ + config = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "passthrough", diff --git a/tests/test_dedup_correctness.py b/tests/test_dedup_correctness.py new file mode 100644 index 0000000..02664a4 --- /dev/null +++ b/tests/test_dedup_correctness.py @@ -0,0 +1,276 @@ +"""Dedup correctness with sources that genuinely contain duplicates. + +Zero-duplicate sources can prove *item loss* but not *dedup correctness*. These +tests inject real duplicates (overlapping sources, cross-page repeats, within- +batch repeats) and assert dedup removes exactly the dupes while surfacing every +unique id. + +These exercise dedup across overlapping sources, cross-page/cross-rebuild repeats, +priority arbitration, and every merger type, at both overfetch=1 and overfetch=4. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.mixers import ( + MergerAppend, + MergerAppendDistribute, + MergerPercentage, + MergerPercentageItem, + MergerPositional, +) +from smartfeed.execution import executor as run_executor + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +# --------------------------------------------------------------------------- +# Overlapping sources: every unique id exactly once +# --------------------------------------------------------------------------- + +def _overlap_ctx(): + a = S.ScriptedSource(S.unique_pool(100, start=0)) # ids 0..99 + b = S.ScriptedSource(S.unique_pool(100, start=50)) # ids 50..149 (overlap 50..99) + return a, b, S.make_ctx({"a": a, "b": b}, redis=_redis()) + + +@pytest.mark.asyncio +async def test_overlapping_sources_cached_each_id_once(): + """session_size > universe -> one batch -> within-batch dedup handles overlap.""" + a, b, ctx = _overlap_ctx() + merger = MergerAppend(node_id="m", items=[ + SubFeed(subfeed_id="a", method_name="a"), + SubFeed(subfeed_id="b", method_name="b"), + ]) + node = S.wrapper(merger, session_size=500, dedup_key="id", overfetch_factor=4) + pages = await S.drain(node, ctx, limit=10, max_pages=100) + S.assert_no_duplicates(pages) + S.assert_full_coverage(pages, set(range(150))) + + +@pytest.mark.asyncio +async def test_overlapping_sources_passthrough_each_id_once(): + """Passthrough cross-page dedup via the Redis seen-set (overfetch=1 -> no loss).""" + a, b, ctx = _overlap_ctx() + merger = MergerAppend(node_id="m", items=[ + SubFeed(subfeed_id="a", method_name="a"), + SubFeed(subfeed_id="b", method_name="b"), + ]) + node = S.wrapper(merger, dedup_key="id", overfetch_factor=1) + pages = await S.drain(node, ctx, limit=10, max_pages=200) + S.assert_no_duplicates(pages) + S.assert_full_coverage(pages, set(range(150))) + + +# --------------------------------------------------------------------------- +# Within-batch high-density duplicates collapse +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_adjacent_duplicates_collapse_within_batch(): + src = S.ScriptedSource(S.adjacent_dup_pool(80)) # each id emitted twice, back-to-back + node = S.wrapper(S.subfeed("src", "src"), session_size=500, dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=100) + S.assert_no_duplicates(pages) + S.assert_full_coverage(pages, set(range(80))) + + +# --------------------------------------------------------------------------- +# Cross-page duplicate suppression (passthrough seen-set) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_cross_page_duplicate_suppressed_passthrough(): + # id 3 appears at pos 3 and again at pos 25 (a later page). + pool = S.unique_pool(30) + pool.insert(25, {"id": 3, "val": "v3-again"}) + src = S.ScriptedSource(pool) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=5, max_pages=50) + ids = S.flat_ids(pages) + assert ids.count(3) == 1, f"cross-page duplicate not suppressed: id 3 served {ids.count(3)}x" + + +# --------------------------------------------------------------------------- +# Cross-REBUILD dedup: within one scroll, a shown id must not reappear after a +# rebuild boundary (session-scoped seen-set carried across cold rebuilds). +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_cached_dedup_suppresses_across_rebuild_boundary(): + # id 3 at pos 3 (batch 1) and pos 20 (batch 2, after session_size=20 boundary). + pool = S.unique_pool(20) + pool.append({"id": 3, "val": "v3-again"}) # pos 20 + src = S.ScriptedSource(pool) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=50) + assert S.flat_ids(pages).count(3) == 1 + + +# --------------------------------------------------------------------------- +# Priority arbitration +# --------------------------------------------------------------------------- + +def _prio_wrapper(): + merger = MergerAppend(node_id="m", items=[ + SubFeed(subfeed_id="lo", method_name="x", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="x", dedup_priority=5), + ]) + return S.wrapper(merger, dedup_key="id") + + +def _stamped(id_, source): + return {"id": id_, "_smartfeed_debug_info": {"source": source}} + + +def _winner_source(item): + return (item.get("_smartfeed_debug_info") or {}).get("source") + + +# Regression guard for BUG#7: in-batch priority arbitration must run even when the +# higher-priority copy is NOT the first-seen one. The seen-check skips only true +# cross-page dupes (`key in seen and key not in batch`), so a within-batch duplicate +# still reaches priority arbitration. +@pytest.mark.asyncio +async def test_priority_higher_source_wins_when_listed_second(): + w = _prio_wrapper() + # low-prio at index 0, high-prio dup at index 1: higher priority must win. + out = w._dedup([_stamped(7, "lo"), _stamped(7, "hi")]) + assert len(out) == 1 + assert _winner_source(out[0]) == "hi", f"lower-priority first-seen copy won: {out}" + + +@pytest.mark.asyncio +async def test_priority_higher_source_wins_integration_merger_order(): + """Integration: MergerAppend lists the low-priority source first, so its copy is + first-seen; the high-priority source (listed second) must still win the key.""" + lo = S.ScriptedSource([{"id": 0, "val": "lo"}]) + hi = S.ScriptedSource([{"id": 0, "val": "hi"}]) + merger = MergerAppend(node_id="m", items=[ + SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), + ]) + node = S.wrapper(merger, session_size=100, dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"lo": lo, "hi": hi}, redis=_redis()) + r = await run_executor.run(node, ctx, limit=10, cursor={}) + kept = [it for it in r.data if it["id"] == 0] + assert len(kept) == 1 + assert kept[0]["val"] == "hi", f"higher-priority source lost to first-seen: {kept}" + + +@pytest.mark.asyncio +async def test_priority_preserved_when_high_arrives_in_refill(): + """Simulates the refill recombination: _dedup(data + refill) where the high- + priority copy only shows up in the refill batch. It must still win, and the + stale copy must be removed (assert the list directly so a fix that keeps the + dupe cannot pass).""" + w = _prio_wrapper() + data = [_stamped(1, "lo"), _stamped(7, "lo")] # first fetch + refill = [_stamped(7, "hi"), _stamped(2, "hi")] # refill batch + out = w._dedup(data + refill) + assert [it["id"] for it in out] == [1, 7, 2], f"expected one copy each in order: {out}" + assert _winner_source(out[1]) == "hi", f"high-priority refill copy lost: {out}" + + +@pytest.mark.asyncio +async def test_priority_defaults_zero_without_source_stamp(): + w = _prio_wrapper() + out = w._dedup([{"id": 8}, _stamped(8, "hi")]) # unstamped (prio 0) then hi (prio 5) + assert len(out) == 1 + assert _winner_source(out[0]) == "hi", f"unstamped prio-0 copy beat higher priority: {out}" + + +@pytest.mark.asyncio +async def test_priority_cross_page_shown_item_wins_documented(): + """Characterization: cross-page, an already-shown (lower-priority) item cannot be + retracted when a higher-priority duplicate appears on a later page.""" + # page 1 shows id 5 from the low-priority source; page 2 has id 5 from high-prio. + lo = S.ScriptedSource([{"id": 5, "val": "lo"}, {"id": 6, "val": "lo"}]) + hi = S.ScriptedSource([{"id": 5, "val": "hi"}, {"id": 7, "val": "hi"}]) + merger = MergerAppend(node_id="m", items=[ + SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), + ]) + node = S.wrapper(merger, dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"lo": lo, "hi": hi}, redis=_redis()) + r1 = await run_executor.run(node, ctx, limit=1, cursor={}) + # Whatever is shown first for id 5 is final; it is never served twice. + r2 = await run_executor.run(node, ctx, limit=3, cursor=r1.next_page) + ids = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] + assert ids.count(5) == 1 + + +# --------------------------------------------------------------------------- +# Dedup coverage holds through every merger type (regression guard for BUG#1), +# for both overfetch=1 and overfetch=4. +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("of", [1, 4]) +@pytest.mark.asyncio +async def test_dedup_over_distribute_covers_all(of): + pool = [{"id": k, "grp": k % 5} for k in range(120)] + src = S.ScriptedSource(pool) + merger = MergerAppendDistribute(node_id="m", items=[SubFeed(subfeed_id="a", method_name="a")], + distribution_key="grp") + node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) + ctx = S.make_ctx({"a": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_full_coverage(pages, set(range(120))) + S.assert_no_duplicates(pages) + + +@pytest.mark.parametrize("of", [1, 4]) +@pytest.mark.asyncio +async def test_dedup_over_percentage_covers_all(of): + src = S.ScriptedSource(S.unique_pool(120)) + merger = MergerPercentage(node_id="m", items=[ + MergerPercentageItem(percentage=100, data=SubFeed(subfeed_id="a", method_name="a")), + ]) + node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) + ctx = S.make_ctx({"a": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_full_coverage(pages, set(range(120))) + S.assert_no_duplicates(pages) + + +@pytest.mark.parametrize("of", [1, 4]) +@pytest.mark.asyncio +async def test_dedup_over_positional_covers_all(of): + a = S.ScriptedSource(S.unique_pool(120, start=0)) # ids 0..119 + b = S.ScriptedSource(S.unique_pool(60, start=1000)) # ids 1000..1059 + merger = MergerPositional(node_id="m", positions=[1, 3, 5], + positional=SubFeed(subfeed_id="a", method_name="a"), + default=SubFeed(subfeed_id="b", method_name="b")) + node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) + ctx = S.make_ctx({"a": a, "b": b}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_full_coverage(pages, set(range(120)) | set(range(1000, 1060))) + S.assert_no_duplicates(pages) + + +# --------------------------------------------------------------------------- +# Regression guard for BUG#8: dense duplicates must not produce short non-final +# pages -- the refill loop scans through duplicate runs until the page is full. +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_dense_duplicates_still_fill_pages(): + # Each unique id 1..79 is followed by three copies of a hot id (0): only ~1 in 4 + # fetched items is new, so a page needs several child fetches to fill. + pool = [] + for k in range(1, 80): + pool.append({"id": k}) + pool += [{"id": 0}] * 3 + src = S.ScriptedSource(pool) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", + overfetch_factor=1, max_refill_loops=2) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=10, max_pages=300) + S.assert_pages_full(pages, limit=10) diff --git a/tests/test_dedup_edge_cases.py b/tests/test_dedup_edge_cases.py new file mode 100644 index 0000000..ad02ca8 --- /dev/null +++ b/tests/test_dedup_edge_cases.py @@ -0,0 +1,100 @@ +"""Dedup edge cases: missing-key policies, non-dict items, key stringification. + +Contract/characterization tests. Where behavior is arguably wrong (int-vs-str key +collision) the test documents CURRENT behavior with a pointer to the review, so a +future change is a conscious decision rather than a silent regression. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.models.base import FeedResult + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +# --------------------------------------------------------------------------- +# missing_key_policy +# --------------------------------------------------------------------------- + +async def _missing_key_source(user_id, limit, next_page, **kw): + # Second item lacks the dedup key "id". + return FeedResult( + data=[{"id": 1, "val": "a"}, {"val": "no-id"}, {"id": 2, "val": "b"}], + next_page={"pos": 0}, has_next_page=False, + ) + + +@pytest.mark.asyncio +async def test_missing_key_policy_error_raises(): + src = _missing_key_source + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", + overfetch_factor=1, missing_key_policy="error") + ctx = S.make_ctx({"src": src}, redis=_redis()) + from smartfeed.execution import executor as run_executor + with pytest.raises(KeyError, match="Dedup key"): + await run_executor.run(node, ctx, limit=10, cursor={}) + + +@pytest.mark.asyncio +async def test_missing_key_policy_keep_passes_through(): + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", + overfetch_factor=1, missing_key_policy="keep") + ctx = S.make_ctx({"src": _missing_key_source}, redis=_redis()) + from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r.data) == 3 # the no-id item is kept un-deduped + + +@pytest.mark.asyncio +async def test_missing_key_policy_drop_removes(): + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", + overfetch_factor=1, missing_key_policy="drop") + ctx = S.make_ctx({"src": _missing_key_source}, redis=_redis()) + from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) + ids = [it.get("id") for it in r.data] + assert ids == [1, 2] # the no-id item is dropped + + +# --------------------------------------------------------------------------- +# non-dict items +# --------------------------------------------------------------------------- + +async def _non_dict_source(user_id, limit, next_page, **kw): + return FeedResult(data=[{"id": 1}, "a-string", {"id": 2}, 42], + next_page={"pos": 0}, has_next_page=False) + + +@pytest.mark.asyncio +async def test_non_dict_items_pass_through_dedup(): + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": _non_dict_source}, redis=_redis()) + from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert "a-string" in r.data and 42 in r.data # non-dicts survive + + +# --------------------------------------------------------------------------- +# key stringification: int 1 vs str "1" +# --------------------------------------------------------------------------- + +async def _mixed_type_key_source(user_id, limit, next_page, **kw): + return FeedResult(data=[{"id": 1, "val": "int"}, {"id": "1", "val": "str"}], + next_page={"pos": 0}, has_next_page=False) + + +@pytest.mark.asyncio +async def test_int_and_str_key_currently_collide_characterization(): + """CURRENT BEHAVIOR: dedup does str(item[key]), so id=1 (int) and id="1" (str) + are treated as the SAME key and one is dropped. See review -- decide whether + mixed-type ids should be distinct. This test pins today's behavior; flip it if + the product decision is that they must be distinct.""" + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": _mixed_type_key_source}, redis=_redis()) + from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r.data) == 1, "int 1 and str '1' currently collide via str() coercion" diff --git a/tests/test_dedup_priority.py b/tests/test_dedup_priority.py index f1be5a1..5cb73f4 100644 --- a/tests/test_dedup_priority.py +++ b/tests/test_dedup_priority.py @@ -33,7 +33,7 @@ def ctx(redis): @pytest.mark.asyncio async def test_wrapper_dedup_priority_override(ctx): """Inner wrapper with dedup_priority=3 overrides child subfeed priorities.""" - config = FeedConfig.parse_obj({ + config = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "outer", diff --git a/tests/test_high_priority.py b/tests/test_high_priority.py index 69e1781..694fc5b 100644 --- a/tests/test_high_priority.py +++ b/tests/test_high_priority.py @@ -370,14 +370,18 @@ class TestMissingDedupKeyPolicies: @staticmethod async def _make_items_with_missing_key(user_id, limit, next_page, **kw): - """Half of the items have 'id', half do not.""" + """Half of the items have 'id', half do not. Stateful (advances by page) so a + dedup refill fetches genuinely new items instead of regenerating seen ids.""" + page = next_page.get("page", 1) + start = (page - 1) * limit data = [] for i in range(limit): - if i % 2 == 0: - data.append({"id": i, "val": f"has_id_{i}"}) + gid = start + i + if gid % 2 == 0: + data.append({"id": gid, "val": f"has_id_{gid}"}) else: - data.append({"val": f"no_id_{i}"}) # no 'id' key - return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) + data.append({"val": f"no_id_{gid}"}) # no 'id' key + return FeedResult(data=data, next_page={"page": page + 1}, has_next_page=True) def _make_methods(self): return {"missing_key": self._make_items_with_missing_key} diff --git a/tests/test_invariants.py b/tests/test_invariants.py new file mode 100644 index 0000000..3223fbd --- /dev/null +++ b/tests/test_invariants.py @@ -0,0 +1,121 @@ +"""Property/invariant backbone for dedup + pagination. + +Each test drains a FINITE source to exhaustion and asserts one universal +property (coverage / no-duplicates / page-fullness / termination) across the +config matrix, so a single property covers a whole class of configurations. +Uses unique (duplicate-free) pools so coverage/no-duplicate assertions are +unambiguous; duplicate-heavy dedup semantics live in test_dedup_correctness.py. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S + + +POOL_N = 200 + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +def _build(cfg): + """cfg keys: cache(bool), dedup(bool), overfetch(int), rerank(bool).""" + child = S.subfeed("src", "src") + node = S.wrapper( + child, + session_size=50 if cfg["cache"] else None, + dedup_key="id" if cfg["dedup"] else None, + overfetch_factor=cfg.get("overfetch", 4), + rerank_method="identity" if cfg.get("rerank") else None, + ) + return node + + +async def identity_rerank(items, session_id): + return items + + +def _ctx(src): + return S.make_ctx({"src": src, "identity": identity_rerank}, redis=_redis()) + + +def _p(id_, marks=(), **cfg): + return pytest.param(cfg, id=id_, marks=marks) + + +# Coverage must hold for every overfetch value and both paths: no item is ever lost. +COVERAGE_CONFIGS = [ + _p("cache-nodedup", cache=True, dedup=False, overfetch=4), + _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), + _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), + _p("passthrough-dedup-of1", cache=False, dedup=True, overfetch=1), + _p("cache-dedup-of2", cache=True, dedup=True, overfetch=2), + _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), + _p("passthrough-dedup-of2", cache=False, dedup=True, overfetch=2), + _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), + _p("cache-dedup-of4-rerank", cache=True, dedup=True, overfetch=4, rerank=True), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cfg", COVERAGE_CONFIGS) +async def test_coverage_no_item_lost(cfg): + """Draining a finite unique source serves every id exactly once -- nothing lost.""" + src = S.ScriptedSource(S.unique_pool(POOL_N)) + node = _build(cfg) + pages = await S.drain(node, _ctx(src), limit=10, max_pages=200) + S.assert_full_coverage(pages, set(range(POOL_N))) + + +# Same matrix WITHOUT the LOSS marks: item loss shows up as coverage gaps, NOT as +# duplicates or non-termination, so these properties hold for every config. +NODUP_CONFIGS = [ + _p("cache-nodedup", cache=True, dedup=False, overfetch=4), + _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), + _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), + _p("passthrough-dedup-of1", cache=False, dedup=True, overfetch=1), + _p("cache-dedup-of2", cache=True, dedup=True, overfetch=2), + _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), + _p("passthrough-dedup-of2", cache=False, dedup=True, overfetch=2), + _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), + _p("cache-dedup-of4-rerank", cache=True, dedup=True, overfetch=4, rerank=True), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cfg", NODUP_CONFIGS) +async def test_no_duplicates_across_pages(cfg): + """No id is ever served twice (unique source; loss must not manifest as dupes).""" + src = S.ScriptedSource(S.unique_pool(POOL_N)) + node = _build(cfg) + pages = await S.drain(node, _ctx(src), limit=10, max_pages=200) + S.assert_no_duplicates(pages) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cfg", [ + _p("cache-nodedup", cache=True, dedup=False, overfetch=4), + _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), + _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), + _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), + _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), +]) +async def test_pages_are_full_until_exhaustion(cfg): + """Every non-final page has exactly `limit` items (loss does not shorten pages, + so this passes even where coverage fails -- keeping the two concerns separate).""" + src = S.ScriptedSource(S.unique_pool(POOL_N)) + node = _build(cfg) + pages = await S.drain(node, _ctx(src), limit=10, max_pages=200) + S.assert_pages_full(pages, limit=10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cfg", NODUP_CONFIGS) +async def test_pagination_terminates(cfg): + """A finite source must produce a terminating pagination (has_next -> False).""" + src = S.ScriptedSource(S.unique_pool(POOL_N)) + node = _build(cfg) + # drain raises if it does not terminate; reaching this line means it did. + await S.drain(node, _ctx(src), limit=10, max_pages=200) diff --git a/tests/test_pagination_edge_cases.py b/tests/test_pagination_edge_cases.py new file mode 100644 index 0000000..afe5282 --- /dev/null +++ b/tests/test_pagination_edge_cases.py @@ -0,0 +1,92 @@ +"""Pagination edge cases: has_next_page exactness at exhaustion, empty sources, +limit larger than the pool, and shuffle interacting with cross-page dedup. + +These are green controls -- they pin correct behavior that a regression could +break (a phantom trailing empty page, a lost final item, shuffle defeating the +seen-set). +""" + +import math + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.models.subfeed import SubFeed +from smartfeed.execution import executor as run_executor + + +def _redis(): + return fakeredis.aioredis.FakeRedis() + + +# --------------------------------------------------------------------------- +# has_next_page exactness: no phantom trailing empty page, no lost final item +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +@pytest.mark.parametrize("n,limit", [(25, 10), (30, 10), (20, 10), (40, 10), (5, 10), (1, 10)]) +async def test_has_next_exact_passthrough(n, limit): + src = S.ScriptedSource(S.unique_pool(n)) + node = S.wrapper(S.subfeed("src", "src")) # passthrough, no cache + ctx = S.make_ctx({"src": src}, redis=None) + pages = await S.drain(node, ctx, limit=limit, max_pages=50) + assert sum(len(p) for p in pages) == n, "not every item was served" + assert all(len(p) > 0 for p in pages), "phantom trailing empty page" + assert len(pages) == max(1, math.ceil(n / limit)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n,limit", [(25, 10), (30, 10), (20, 10), (45, 10)]) +async def test_has_next_exact_cached(n, limit): + src = S.ScriptedSource(S.unique_pool(n)) + node = S.wrapper(S.subfeed("src", "src"), session_size=20) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=limit, max_pages=50) + assert sum(len(p) for p in pages) == n + assert all(len(p) > 0 for p in pages), "phantom trailing empty page" + + +# --------------------------------------------------------------------------- +# Empty source +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_empty_source_returns_empty_no_next(): + src = S.ScriptedSource([]) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id", overfetch_factor=4) + ctx = S.make_ctx({"src": src}, redis=_redis()) + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert r.data == [] + assert r.has_next_page is False + + +# --------------------------------------------------------------------------- +# limit larger than the whole pool +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_limit_larger_than_pool(): + src = S.ScriptedSource(S.unique_pool(5)) + node = S.wrapper(S.subfeed("src", "src")) + ctx = S.make_ctx({"src": src}, redis=None) + pages = await S.drain(node, ctx, limit=10, max_pages=10) + assert [it["id"] for pg in pages for it in pg] == list(range(5)) + + +# --------------------------------------------------------------------------- +# shuffle=True must not defeat cross-page dedup (seen-set) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_shuffle_does_not_defeat_cross_page_dedup(): + pool = S.unique_pool(30) + pool.insert(25, {"id": 3, "val": "again"}) # id 3 re-emitted on a later page + src = S.ScriptedSource(pool) + node = S.wrapper(SubFeed(subfeed_id="src", method_name="src", shuffle=True), + dedup_key="id", overfetch_factor=1) + ctx = S.make_ctx({"src": src}, redis=_redis()) + pages = await S.drain(node, ctx, limit=5, max_pages=50) + ids = S.flat_ids(pages) + assert ids.count(3) == 1, "shuffle let a cross-page duplicate through the seen-set" + S.assert_full_coverage(pages, set(range(30))) diff --git a/tests/test_wrapper_cache.py b/tests/test_wrapper_cache.py index 6933bbd..d05fd41 100644 --- a/tests/test_wrapper_cache.py +++ b/tests/test_wrapper_cache.py @@ -59,10 +59,10 @@ async def test_stale_gen_resets_to_page1(ctx, redis): r1 = await run_executor.run(node, ctx, limit=10, cursor={}) # Flush to simulate TTL expiry await redis.flushall() - stale_cursor = {"cached": {"page": 5, "gen": "stale_nonce"}} + stale_cursor = {"cached": {"offset": 40, "gen": "stale_nonce"}} r2 = await run_executor.run(node, ctx, limit=10, cursor=stale_cursor) - # Should reset: page 2 in next cursor (just served page 1) - assert r2.next_page["cached"]["page"] == 2 + # Should reset: served page 1 (offset 0), next cursor points at offset 10 + assert r2.next_page["cached"]["offset"] == 10 assert r2.next_page["cached"]["gen"] != "stale_nonce" diff --git a/tests/test_wrapper_full_pipeline.py b/tests/test_wrapper_full_pipeline.py index 492f4b9..bbea611 100644 --- a/tests/test_wrapper_full_pipeline.py +++ b/tests/test_wrapper_full_pipeline.py @@ -36,7 +36,7 @@ def ctx(redis): @pytest.mark.asyncio async def test_full_pipeline(ctx): - config = FeedConfig.parse_obj({ + config = FeedConfig.model_validate({ "version": "2", "feed": { "type": "wrapper", "node_id": "full", From fb553af42d3180c50e97b96669c23a04ab700b8a Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Mon, 6 Jul 2026 13:10:51 +0100 Subject: [PATCH 39/41] Lint fixes. --- smartfeed/execution/redis_lock.py | 6 +-- smartfeed/models/__init__.py | 3 +- smartfeed/models/base.py | 30 ++++++++------ smartfeed/models/mixers.py | 44 ++++++++++---------- smartfeed/models/subfeed.py | 4 +- smartfeed/models/wrapper.py | 69 ++++++++++++++++--------------- 6 files changed, 79 insertions(+), 77 deletions(-) diff --git a/smartfeed/execution/redis_lock.py b/smartfeed/execution/redis_lock.py index dec85aa..655fd75 100644 --- a/smartfeed/execution/redis_lock.py +++ b/smartfeed/execution/redis_lock.py @@ -23,12 +23,10 @@ def __init__(self, redis: AsyncRedis, key: str, ttl: int = 10): self._owned = False async def __aenter__(self) -> bool: - self._owned = bool( - await self._redis.set(self._key, self._token, nx=True, ex=self._ttl) - ) + self._owned = bool(await self._redis.set(self._key, self._token, nx=True, ex=self._ttl)) return self._owned - async def __aexit__(self, *exc): + async def __aexit__(self, *exc: object) -> None: if not self._owned: return val = await self._redis.get(self._key) diff --git a/smartfeed/models/__init__.py b/smartfeed/models/__init__.py index e50a531..1cdb995 100644 --- a/smartfeed/models/__init__.py +++ b/smartfeed/models/__init__.py @@ -42,7 +42,7 @@ class FeedConfig(BaseModel): version: str - feed: FeedNode # type: ignore[valid-type] + feed: FeedNode # Rebuild forward refs so that nested FeedNode fields resolve correctly @@ -66,7 +66,6 @@ class FeedConfig(BaseModel): "FeedResult", "SmartFeedDebugInfo", "FeedItem", - "SubFeed", "MergerPercentage", "MergerPercentageItem", diff --git a/smartfeed/models/base.py b/smartfeed/models/base.py index 9df9e0f..96a6fb6 100644 --- a/smartfeed/models/base.py +++ b/smartfeed/models/base.py @@ -2,29 +2,32 @@ import hashlib import json -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict, Field +if TYPE_CHECKING: + from smartfeed.execution.plans import MixPlan + class SmartFeedDebugInfo(BaseModel): """Metadata stamped by SmartFeed on every item.""" # Always present (stamped by SubFeed) - source: str # subfeed_id (regular_tours, recommended_tours, promo_tours) + source: str # subfeed_id (regular_tours, recommended_tours, promo_tours) # Stamped by FeedManager (top-level position in final feed) - smartfeed_position: Optional[int] = None # 0-based position in page + smartfeed_position: Optional[int] = None # 0-based position in page # Stamped by subfeed methods (optional, source-specific) - strategy: Optional[str] = None # ML model strategy (model_hot_users, model_cold_users) + strategy: Optional[str] = None # ML model strategy (model_hot_users, model_cold_users) # Stamped by rerank callable (optional, present only when rerank is configured) - rerank_position: Optional[int] = None # position after rerank (1-based) - rrf_score: Optional[float] = None # RRF score - feature_score: Optional[float] = None # feature score from ES coefficients - feature_position: Optional[int] = None # rank by feature score (1-based) - total_reranked: Optional[int] = None # total items in rerank batch + rerank_position: Optional[int] = None # position after rerank (1-based) + rrf_score: Optional[float] = None # RRF score + feature_score: Optional[float] = None # feature score from ES coefficients + feature_position: Optional[int] = None # rank by feature score (1-based) + total_reranked: Optional[int] = None # total items in rerank batch raw_params: Optional[Dict[str, Any]] = None # raw tour params from Redis model_config = ConfigDict(extra="allow") @@ -34,9 +37,7 @@ class FeedItem(BaseModel): """One item in the feed output. Wraps tour data + SmartFeed metadata.""" id: Any - smartfeed_debug_info: Optional[SmartFeedDebugInfo] = Field( - default=None, alias="_smartfeed_debug_info" - ) + smartfeed_debug_info: Optional[SmartFeedDebugInfo] = Field(default=None, alias="_smartfeed_debug_info") model_config = ConfigDict(extra="allow", populate_by_name=True) @@ -49,6 +50,10 @@ def config_hash(self) -> str: raw = json.dumps(self.model_dump(), sort_keys=True, default=str) return hashlib.md5(raw.encode()).hexdigest()[:8] + def build_mix_plan(self, *, ctx: Any, limit: int, cursor: dict) -> "MixPlan": + """Mixer nodes override this; leaf nodes (SubFeed, Wrapper) use execute() instead.""" + raise NotImplementedError(f"{type(self).__name__} is not a mixer node") + class FeedResult(BaseModel): data: List @@ -67,4 +72,5 @@ def coerce_feed_node(value: Any) -> Any: # Lazy import to break circular dependency from smartfeed.models import FeedNode # noqa: PLC0415 from pydantic import TypeAdapter + return TypeAdapter(FeedNode).validate_python(value) diff --git a/smartfeed/models/mixers.py b/smartfeed/models/mixers.py index 3fb4241..10b43d9 100644 --- a/smartfeed/models/mixers.py +++ b/smartfeed/models/mixers.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections import defaultdict, deque -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional, Tuple from pydantic import BaseModel, model_validator @@ -20,6 +20,7 @@ def _merge_cursor(child_cursors: Dict[str, dict]) -> dict: # MergerPercentageItem # --------------------------------------------------------------------------- + class MergerPercentageItem(BaseModel): percentage: int data: Any # BaseNode subclass @@ -36,6 +37,7 @@ def _coerce_data(cls, values: Any) -> Any: # MergerPercentage # --------------------------------------------------------------------------- + class MergerPercentage(BaseNode): type: Literal["merger_percentage"] = "merger_percentage" node_id: str @@ -81,7 +83,7 @@ def build_mix_plan( def assemble( buffers: Dict[str, list], child_cursors: Dict[str, dict], - ): + ) -> Tuple[List[Any], dict]: # Simple concatenation: demand already ensures correct proportions merged_data: List[Any] = [] for child in children: @@ -95,6 +97,7 @@ def assemble( # MergerAppend # --------------------------------------------------------------------------- + class MergerAppend(BaseNode): type: Literal["merger_append"] = "merger_append" node_id: str @@ -107,10 +110,7 @@ def _coerce_items(cls, values: Any) -> Any: if isinstance(values, dict): raw_items = values.get("items") if isinstance(raw_items, list): - values["items"] = [ - coerce_feed_node(item) if isinstance(item, dict) else item - for item in raw_items - ] + values["items"] = [coerce_feed_node(item) if isinstance(item, dict) else item for item in raw_items] return values def build_mix_plan( @@ -123,8 +123,10 @@ def build_mix_plan( # Each child gets equal demand share; leftover goes to first children n = len(self.items) if n == 0: - def assemble_empty(buffers, child_cursors): + + def assemble_empty(buffers: Dict[str, list], child_cursors: Dict[str, dict]) -> Tuple[List[Any], dict]: return [], _merge_cursor(child_cursors) + return MixPlan(children=[], assemble=assemble_empty) base_demand = limit // n @@ -142,7 +144,7 @@ def assemble_empty(buffers, child_cursors): def assemble( buffers: Dict[str, list], child_cursors: Dict[str, dict], - ): + ) -> Tuple[List[Any], dict]: merged_data: List[Any] = [] for child in children: merged_data.extend(buffers.get(child.node_id, [])) @@ -156,12 +158,13 @@ def assemble( # MergerPositional # --------------------------------------------------------------------------- + class MergerPositional(BaseNode): type: Literal["merger_positional"] = "merger_positional" node_id: str positions: List[int] = [] positional: Any # BaseNode subclass - default: Any # BaseNode subclass + default: Any # BaseNode subclass dedup_priority: int = 0 @model_validator(mode="before") @@ -201,7 +204,7 @@ def build_mix_plan( def assemble( buffers: Dict[str, list], child_cursors: Dict[str, dict], - ): + ) -> Tuple[List[Any], dict]: pos_items = deque(buffers.get(positional_child.node_id, [])) def_items = deque(buffers.get(default_child.node_id, [])) @@ -224,6 +227,7 @@ def assemble( # MergerPercentageGradient # --------------------------------------------------------------------------- + class MergerPercentageGradient(BaseNode): """Percentage-based merger that shifts the ratio over pages.""" @@ -253,11 +257,7 @@ def _calculate_demands(self, page: int, limit: int) -> tuple: percentage_to = 100 if i > start_position: - iter_limit = ( - (limit * page - start_position) - if i > limit * page - else (i - start_position) - ) + iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) start_position = i from_take = iter_limit * percentage_from // 100 to_take = iter_limit - from_take @@ -295,7 +295,7 @@ def build_mix_plan( def assemble( buffers: Dict[str, list], child_cursors: Dict[str, dict], - ): + ) -> Tuple[List[Any], dict]: from_data = list(buffers.get(from_child.node_id, [])) to_data = list(buffers.get(to_child.node_id, [])) result: List[Any] = [] @@ -303,8 +303,8 @@ def assemble( for seg in segments: ft = int(seg["from_take"]) tt = int(seg["to_take"]) - result.extend(from_data[fi: fi + ft]) - result.extend(to_data[ti: ti + tt]) + result.extend(from_data[fi : fi + ft]) + result.extend(to_data[ti : ti + tt]) fi += ft ti += tt merged_cur = _merge_cursor(child_cursors) @@ -318,6 +318,7 @@ def assemble( # MergerAppendDistribute # --------------------------------------------------------------------------- + class MergerAppendDistribute(BaseNode): """Append merger that round-robins items by a distribution key.""" @@ -335,10 +336,7 @@ def _coerce_items(cls, values: Any) -> Any: if isinstance(values, dict): raw_items = values.get("items") if isinstance(raw_items, list): - values["items"] = [ - coerce_feed_node(item) if isinstance(item, dict) else item - for item in raw_items - ] + values["items"] = [coerce_feed_node(item) if isinstance(item, dict) else item for item in raw_items] return values def _uniform_distribute(self, data: list) -> list: @@ -382,7 +380,7 @@ def build_mix_plan( def assemble( buffers: Dict[str, list], child_cursors: Dict[str, dict], - ): + ) -> Tuple[List[Any], dict]: all_items: List[Any] = [] for child in children: all_items.extend(buffers.get(child.node_id, [])) diff --git a/smartfeed/models/subfeed.py b/smartfeed/models/subfeed.py index 8b7a45f..61b7724 100644 --- a/smartfeed/models/subfeed.py +++ b/smartfeed/models/subfeed.py @@ -38,9 +38,7 @@ async def execute( raise # Scope the failure cursor to THIS subfeed (its own, unadvanced position) # so it retries next page without clobbering sibling cursors on merge. - return FeedResult( - data=[], next_page={self.subfeed_id: subfeed_cursor}, has_next_page=False - ) + return FeedResult(data=[], next_page={self.subfeed_id: subfeed_cursor}, has_next_page=False) if self.shuffle: shuffle(result.data) diff --git a/smartfeed/models/wrapper.py b/smartfeed/models/wrapper.py index febbe04..8cf1928 100644 --- a/smartfeed/models/wrapper.py +++ b/smartfeed/models/wrapper.py @@ -69,9 +69,7 @@ async def execute( if self.cache is None or ctx is None or ctx.redis is None: return await self._passthrough(methods_dict, session_id, limit, cursor, ctx) - return await self._execute_with_cache( - methods_dict, session_id, limit, cursor, ctx - ) + return await self._execute_with_cache(methods_dict, session_id, limit, cursor, ctx) # -- dedup ----------------------------------------------------------------- @@ -82,9 +80,7 @@ def _resolve_dedup_priorities(self) -> Dict[str, int]: return result @staticmethod - def _collect_priorities( - node: BaseNode, override_priority: int, out: Dict[str, int] - ) -> None: + def _collect_priorities(node: BaseNode, override_priority: int, out: Dict[str, int]) -> None: """Recursive walk: propagate override_priority down, write SubFeed priorities to out.""" from .subfeed import SubFeed @@ -97,8 +93,7 @@ def _collect_priorities( return # Walk children: look in known child-bearing attributes - for attr in ("items", "data", "positional", "default", - "item_from", "item_to"): + for attr in ("items", "data", "positional", "default", "item_from", "item_to"): child = getattr(node, attr, None) if child is None: continue @@ -172,9 +167,7 @@ def _dedup(self, data: List, seen: Optional[set] = None) -> List: # -- rerank ---------------------------------------------------------------- - async def _apply_rerank( - self, data: list, methods_dict: dict, session_id: str - ) -> list: + async def _apply_rerank(self, data: list, methods_dict: dict, session_id: str) -> list: """Call rerank callable from methods_dict. Validates output length.""" if not self.rerank: return data @@ -189,8 +182,7 @@ async def _apply_rerank( return data if len(result) != original_len: raise ValueError( - f"Rerank '{self.rerank.method_name}' must return exactly " - f"{original_len} items, got {len(result)}" + f"Rerank '{self.rerank.method_name}' must return exactly " f"{original_len} items, got {len(result)}" ) return result @@ -287,6 +279,7 @@ async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None return key = self._seen_set_key(session_id) await ctx.redis.sadd(key, *new_keys) + assert self.dedup is not None, "seen-set only written when dedup configured" ttl = self.dedup.state_ttl await ctx.redis.expire(key, ttl) @@ -329,6 +322,7 @@ async def _write_cache( child_has_next: bool = False, ) -> None: """Write session data and metadata to Redis with TTL.""" + assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" base = self._base_key(session_id) ttl = self.cache.session_ttl @@ -355,7 +349,11 @@ async def _touch_ttl(self, ctx: Any, session_id: str) -> None: # -- pagination ------------------------------------------------------------ def _paginate( - self, data: List, limit: int, offset: int, gen: str, + self, + data: List, + limit: int, + offset: int, + gen: str, child_has_next: bool = False, ) -> FeedResult: """Slice cached data at an absolute offset and build the next cursor. @@ -406,7 +404,10 @@ async def _execute_with_cache( await self._touch_ttl(ctx, session_id) child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) return self._paginate( - cached_data, limit, cursor_offset, cursor_gen, + cached_data, + limit, + cursor_offset, + cursor_gen, child_has_next=child_has_next, ) # Cache exhausted: rebuild with continuation cursor @@ -419,9 +420,7 @@ async def _execute_with_cache( ) # Cold path: no gen or stale gen -> fresh build - return await self._cold_build_locked( - methods_dict, session_id, limit, ctx, child_cursor={} - ) + return await self._cold_build_locked(methods_dict, session_id, limit, ctx, child_cursor={}) async def _cold_build_locked( self, @@ -439,9 +438,7 @@ async def _cold_build_locked( lock_key = f"{self._base_key(session_id)}:coldlock" async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: if acquired: - return await self._cold_build( - methods_dict, session_id, limit, ctx, child_cursor - ) + return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) # Another coroutine is building -- wait for its cache, then serve page 1. for _ in range(50): await asyncio.sleep(0.1) @@ -449,17 +446,16 @@ async def _cold_build_locked( if meta: cached = await self._read_cache(ctx, session_id) if cached is not None: - child_has_next = meta.get( - "child_has_next", bool(meta.get("child_cursor")) - ) + child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) return self._paginate( - cached, limit, 0, meta.get("gen", ""), + cached, + limit, + 0, + meta.get("gen", ""), child_has_next=child_has_next, ) # Fallback: builder never wrote -> build ourselves. - return await self._cold_build( - methods_dict, session_id, limit, ctx, child_cursor - ) + return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) # -- shared base cache helpers ---------------------------------------------- @@ -487,10 +483,16 @@ async def _read_shared_base(self, ctx: Any, session_id: str, child_cursor: Dict) return orjson.loads(raw) async def _write_shared_base( - self, ctx: Any, session_id: str, child_cursor: Dict, data: List, - child_cursor_out: Dict, child_has_next: bool = False, + self, + ctx: Any, + session_id: str, + child_cursor: Dict, + data: List, + child_cursor_out: Dict, + child_has_next: bool = False, ) -> None: """Write a shared-base segment and its meta to Redis with TTL.""" + assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" key = self._base_shared_key(session_id, child_cursor) ttl = self.cache.session_ttl pipe = ctx.redis.pipeline() @@ -560,15 +562,14 @@ async def _cold_build( child_cursor: Dict, ) -> FeedResult: """Build full session: fetch -> dedup -> rerank -> write cache -> paginate page 1.""" + assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" session_size = self.cache.session_size if self.cache_key is not None: # Shared cache path: base data (fetch + dedup) is shared across wrappers # with the same cache_key; rerank is applied per-wrapper after. Each # continuation position is a distinct shared segment (see _base_shared_key). - base_data = await self._build_shared_base( - methods_dict, session_id, ctx, child_cursor - ) + base_data = await self._build_shared_base(methods_dict, session_id, ctx, child_cursor) shared_meta = await self._read_shared_base_meta(ctx, session_id, child_cursor) child_cursor_out = shared_meta.get("child_cursor", {}) if shared_meta else {} child_has_next = shared_meta.get("child_has_next", False) if shared_meta else False @@ -593,6 +594,7 @@ async def _cold_build( ctx, session_size, child_cursor, seen=seen ) if self.dedup: + assert seen is not None await self._save_seen_set(ctx, session_id, seen) self._stamp_pre_rerank(data) data = await self._apply_rerank(data, methods_dict, session_id) @@ -614,6 +616,7 @@ async def _build_shared_base( only one caller fetches from child on a cold build. Others wait and read.""" from smartfeed.execution.redis_lock import RedisLock + assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" session_size = self.cache.session_size lock_key = f"{self._base_shared_key(session_id, child_cursor)}:lock" From cb63d5eb47b176972d5aec5dd03267e71712d6b0 Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 9 Jul 2026 15:35:43 +0100 Subject: [PATCH 40/41] Cleanup + bugfixes, tests added, linter switched to pyright. --- .github/workflows/build-and-publish.yaml | 39 +- .github/workflows/lint.yaml | 17 +- .github/workflows/tests.yaml | 11 +- .gitignore | 4 +- ARCHITECTURE.md | 30 +- LICENSE | 2 +- Makefile | 26 +- README.md | 47 +- docker-compose.yml | 9 - smartfeed/__init__.py | 39 +- smartfeed/execution/context.py | 4 +- smartfeed/execution/executor.py | 35 +- smartfeed/execution/plans.py | 10 +- smartfeed/execution/redis_lock.py | 28 +- smartfeed/manager.py | 47 +- smartfeed/models/__init__.py | 58 ++- smartfeed/models/base.py | 74 ++- smartfeed/models/mixers.py | 175 ++++--- smartfeed/models/subfeed.py | 17 +- smartfeed/models/wrapper.py | 505 ++++++++++++-------- smartfeed/py.typed | 0 tests/sources.py | 17 +- tests/test_async_concurrency.py | 17 +- tests/test_bug_cold_build_stampede.py | 17 +- tests/test_bug_lock_loser_gives_up_early.py | 52 ++ tests/test_bug_lock_release_not_atomic.py | 62 +++ tests/test_bug_nested_wrapper_priority.py | 65 +++ tests/test_bug_overfetch_item_loss.py | 13 +- tests/test_bug_passthrough_fresh_scroll.py | 34 ++ tests/test_bug_redis_lock_decode.py | 7 +- tests/test_bug_seen_ttl_not_touched.py | 50 ++ tests/test_bug_shared_cache_continuation.py | 5 +- tests/test_bug_shared_cache_dedup.py | 39 ++ tests/test_bug_softfail_cursor_leak.py | 23 +- tests/test_bug_stale_waiter.py | 46 ++ tests/test_bug_torn_warm_read.py | 86 ++++ tests/test_bug_variable_page_size.py | 12 +- tests/test_config_parsing.py | 67 +-- tests/test_config_validation.py | 138 ++++++ tests/test_continuation.py | 19 +- tests/test_cursors.py | 42 +- tests/test_dedup_correctness.py | 136 +++--- tests/test_dedup_edge_cases.py | 34 +- tests/test_dedup_priority.py | 39 +- tests/test_distribute.py | 21 +- tests/test_gradient.py | 7 +- tests/test_gradient_demands.py | 93 ++++ tests/test_high_priority.py | 36 +- tests/test_input_validation.py | 156 ++++++ tests/test_invariants.py | 49 +- tests/test_manager.py | 68 +++ tests/test_medium_priority.py | 36 +- tests/test_mixers.py | 8 +- tests/test_pagination_edge_cases.py | 9 +- tests/test_resilience.py | 21 +- tests/test_shared_lock_contention.py | 43 ++ tests/test_sharp_edges.py | 82 ++++ tests/test_silent_degradation.py | 30 ++ tests/test_subfeed.py | 18 +- tests/test_wrapper_cache.py | 2 +- tests/test_wrapper_dedup.py | 2 +- tests/test_wrapper_full_pipeline.py | 52 +- tests/test_wrapper_rerank.py | 3 + tests/test_wrapper_shared_cache.py | 31 +- 64 files changed, 2169 insertions(+), 795 deletions(-) delete mode 100644 docker-compose.yml create mode 100644 smartfeed/py.typed create mode 100644 tests/test_bug_lock_loser_gives_up_early.py create mode 100644 tests/test_bug_lock_release_not_atomic.py create mode 100644 tests/test_bug_nested_wrapper_priority.py create mode 100644 tests/test_bug_passthrough_fresh_scroll.py create mode 100644 tests/test_bug_seen_ttl_not_touched.py create mode 100644 tests/test_bug_shared_cache_dedup.py create mode 100644 tests/test_bug_stale_waiter.py create mode 100644 tests/test_bug_torn_warm_read.py create mode 100644 tests/test_config_validation.py create mode 100644 tests/test_gradient_demands.py create mode 100644 tests/test_input_validation.py create mode 100644 tests/test_manager.py create mode 100644 tests/test_shared_lock_contention.py create mode 100644 tests/test_sharp_edges.py create mode 100644 tests/test_silent_degradation.py diff --git a/.github/workflows/build-and-publish.yaml b/.github/workflows/build-and-publish.yaml index a93aa53..3976002 100644 --- a/.github/workflows/build-and-publish.yaml +++ b/.github/workflows/build-and-publish.yaml @@ -3,10 +3,10 @@ name: Build and Publish to PyPI on: push: tags: - - 'v*' + - "v[0-9]+.[0-9]+.[0-9]+" jobs: - build: + test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -19,9 +19,35 @@ jobs: - name: Install Poetry run: pipx install poetry - - name: Set version from tag - if: startsWith(github.ref, 'refs/tags/v') - run: poetry version "${GITHUB_REF_NAME#v}" + - name: Check tag matches pyproject version + run: | + PYPROJECT_VERSION="$(poetry version -s)" + TAG_VERSION="${GITHUB_REF_NAME#v}" + if [ "$PYPROJECT_VERSION" != "$TAG_VERSION" ]; then + echo "Tag v$TAG_VERSION does not match pyproject version $PYPROJECT_VERSION" >&2 + exit 1 + fi + + - name: Install dependencies + run: poetry install --all-extras + + - name: Run tests + run: poetry run pytest + + build: + runs-on: ubuntu-latest + needs: + - test + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install Poetry + run: pipx install poetry - name: Build distributions run: poetry build @@ -35,7 +61,7 @@ jobs: publish-to-test-pypi: environment: name: testpypi - url: https://pypi.org/project/epoch8-smartfeed/ + url: https://test.pypi.org/project/epoch8-smartfeed/ permissions: id-token: write runs-on: ubuntu-latest @@ -54,6 +80,7 @@ jobs: publish-to-pypi: environment: name: pypi + url: https://pypi.org/project/epoch8-smartfeed/ permissions: id-token: write runs-on: ubuntu-latest diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index c5996e3..0e0b1a7 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -5,7 +5,10 @@ on: paths: - ".github/workflows/lint.yaml" - "smartfeed/**" + - "tests/**" - "pyproject.toml" + - "poetry.lock" + pull_request: jobs: lint: @@ -13,8 +16,6 @@ jobs: strategy: matrix: python-version: - # - "3.8" - # - "3.9" - "3.10" steps: @@ -37,10 +38,16 @@ jobs: run: | poetry install --all-extras - - name: Lint with mypy + - name: Lint with ruff run: | - poetry run mypy smartfeed + pip install ruff==0.15.20 + ruff check smartfeed tests + + - name: Type-check with pyright + run: | + pip install pyright==1.1.409 + pyright --pythonpath "$(which python)" smartfeed tests - name: Lint with black run: | - poetry run black --check smartfeed + poetry run black --check smartfeed tests diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 14b0d53..f9b8b51 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -5,7 +5,10 @@ on: paths: - ".github/workflows/tests.yaml" - "smartfeed/**" + - "tests/**" - "pyproject.toml" + - "poetry.lock" + pull_request: jobs: test: @@ -13,18 +16,12 @@ jobs: strategy: matrix: python-version: - # - "3.8" - "3.9" - "3.10" + - "3.11" - "3.12" - "3.13" - services: - redis: - image: redis - ports: - - 6379:6379 - steps: - uses: actions/checkout@v3 diff --git a/.gitignore b/.gitignore index 95f8006..4d10ba0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,7 @@ __pycache__/ local_settings.py db.sqlite3 media - +frontend/ ### Linux ### *~ @@ -210,4 +210,4 @@ ENV/ /site # Poetry Lock -*.lock \ No newline at end of file +*.lock diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 941623a..6a4b1c0 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -43,7 +43,7 @@ Mixer (coordinator) -- runs children in parallel, assembles results Module-level functions (no class): -- `run(node, ctx, limit, cursor)` -- dispatches to `node.execute()` or `node.build_mix_plan()` +- `run(node, ctx, limit, cursor)` -- dispatches on `isinstance(node, MixerNode)`: mixers build a MixPlan, leaf/pipeline nodes execute() - `_execute_mix(plan, ctx, cursor)` -- runs MixPlan children in parallel via `asyncio.gather` ### Wrapper Pipeline @@ -91,22 +91,25 @@ FeedResult(data: list, next_page: dict, has_next_page: bool) ### Output -``` -SmartFeedDebugInfo(source, smartfeed_position, rerank_position?, rrf_score?, ...) -FeedItem(id, _smartfeed_debug_info: SmartFeedDebugInfo) -``` +Items are plain dicts; each carries a `_smartfeed_debug_info` dict bundle +(source, smartfeed_position, rerank_position when reranked -- all 0-based). ## Redis State ``` -sf:{session_id}:{node_id}:{config_hash} -- cached session data (JSON list) -sf:{session_id}:{node_id}:{config_hash}:meta -- metadata (gen, child_cursor, child_has_next) -sf:{session_id}:{node_id}:{config_hash}:coldlock -- cold-build lock (SETNX, ttl 30s) -sf:{session_id}:{node_id}:{config_hash}:seen -- session-scoped dedup seen-set (Redis SET) -sf:{session_id}:{cache_key}:{child_hash}:{segment} -- shared base segment (per continuation window) -sf:{session_id}:{cache_key}:{child_hash}:{segment}:lock -- shared segment cold-build lock (SETNX) +sf:{session_id}:{cache_key or node_id}:{config_hash} -- cached session batch (Redis LIST of orjson items) +sf:{session_id}:{cache_key or node_id}:{config_hash}:meta -- metadata (gen, child_cursor, child_has_next) +sf:{session_id}:{cache_key or node_id}:{config_hash}:coldlock -- cold-build lock (SETNX, ttl 10s) +sf:{session_id}:{node_id}:{config_hash}:seen -- session-scoped dedup seen-set (Redis SET) +sf:{session_id}:{cache_key}:{child_hash}:{segment} -- shared base segment (blob, per continuation window) +sf:{session_id}:{cache_key}:{child_hash}:{segment}:meta -- shared segment meta (child cursor / has_next) +sf:{session_id}:{cache_key}:{child_hash}:{segment}:lock -- shared segment cold-build lock (SETNX, ttl 10s) ``` +Warm page reads are one MULTI/EXEC pipeline: GET meta + LLEN + LRANGE of the page +window + EXPIRE on data/meta/seen -- an atomic snapshot that transfers only the +requested window, never the whole batch. + TTL = inactivity timeout. Refreshed on every access via pipeline EXPIRE. Cursor for a cached wrapper is `{node_id: {offset, gen}}` (absolute offset). @@ -126,7 +129,8 @@ not first-seen; equal priority = first-seen. ## Config Hash -`BaseNode.config_hash()` = md5(model_dump_json(sort_keys=True))[:8]. +`BaseNode.config_hash()` = md5(json.dumps(model_dump(), sort_keys=True, default=str))[:8], +memoized per node instance (config is immutable after validation). Used in Redis keys. Any config change = different hash = fresh cache. @@ -143,7 +147,7 @@ Two wrappers with same `cache_key` share base data. Each applies its own rerank. ``` smartfeed/ models/ - base.py -- BaseNode, FeedResult, SmartFeedDebugInfo, FeedItem, config_hash + base.py -- BaseNode, MixerNode, FeedResult, config_hash subfeed.py -- SubFeed (leaf) wrapper.py -- Wrapper (cache + dedup + rerank pipeline) mixers.py -- Percentage, Positional, Append, Distribute, Gradient diff --git a/LICENSE b/LICENSE index 90cc1c9..f95ad24 100644 --- a/LICENSE +++ b/LICENSE @@ -186,7 +186,7 @@ same "printed page" as the copyright notice for easier identification within third-party archives. - Copyright 2023 Epoch8 + Copyright 2023-2026 Epoch8 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/Makefile b/Makefile index 682148e..aba1647 100644 --- a/Makefile +++ b/Makefile @@ -1,27 +1,11 @@ lint: - mypy smartfeed - black --check smartfeed + ruff check smartfeed tests + pyright smartfeed tests + black --check smartfeed tests format: - black --verbose smartfeed tests + black smartfeed tests isort smartfeed tests test: - pytest -s -vv -k "not test_merger_view_session" - -test_cache: - pytest -s -vv -k "test_merger_view_session" - -.PHONY: test_async_chart charting - -# Runs only the async loop block + Chrome trace test. -# Writes trace.json next to this Makefile (project root). -test_async_chart: - rm -f ./trace.json - SMARTFEED_CHROME_TRACE=./trace.json pytest -q tests/test_async_loop_blocks_trace.py - @echo "\nWrote trace: $(CURDIR)/trace.json" - @echo "Open Chrome -> chrome://tracing -> Load -> select trace.json" - -# Convenience target: generate the trace + try to open chrome://tracing. -charting: test_async_chart - -@open -a "Google Chrome" "chrome://tracing" 2>/dev/null || true + pytest diff --git a/README.md b/README.md index 1ab3e10..b599c2e 100644 --- a/README.md +++ b/README.md @@ -27,14 +27,17 @@ pip install epoch8-smartfeed Requires: Python 3.9+, **pydantic v2**, orjson. An async Redis client (`redis.asyncio.Redis`) is required only if any Wrapper uses `cache` or cross-page `dedup`; without it those features silently fall back to uncached behavior. +Note the distribution name: an unrelated 2015 project owns the PyPI name `smartfeed` and installs the **same** top-level `smartfeed` package — install `epoch8-smartfeed`, and do not co-install the other one. + ## Quick start ```python +from redis.asyncio import Redis + from smartfeed.manager import FeedManager from smartfeed.models.base import FeedResult config = { - "version": "2", "feed": { "type": "wrapper", "node_id": "main", @@ -66,10 +69,12 @@ async def source_b(user_id, limit, next_page, **kwargs): async def my_rerank(items, session_id): return sorted(items, key=lambda x: x.get("score", 0), reverse=True) +redis = Redis() # or None: cache/dedup then degrade to per-page passthrough + manager = FeedManager( config=config, methods_dict={"source_a": source_a, "source_b": source_b, "my_rerank": my_rerank}, - redis_client=redis, # async redis.asyncio.Redis; needed for cache/dedup persistence + redis_client=redis, # needed for cache/dedup persistence ) result = await manager.get_feed(session_id="user_123", limit=20, cursor=None) @@ -78,7 +83,7 @@ result = await manager.get_feed(session_id="user_123", limit=20, cursor=None) # result.has_next_page -> bool ``` -Config is validated when `FeedManager` is constructed. A malformed config raises `pydantic.ValidationError` at construction, not on the first `get_feed`. +Config is validated when `FeedManager` is constructed. A malformed config raises `pydantic.ValidationError` at construction, not on the first `get_feed`. `subfeed_id` and `node_id` share one namespace (cursor keys, Redis keys) and must be unique across the whole tree — duplicates raise at construction. ## How it works @@ -97,7 +102,7 @@ The executor walks the tree. For each node it either calls `execute()` (SubFeed, The entire public surface is one class with two methods. ```python -FeedManager(config: dict, methods_dict: dict, redis_client: Optional[Any] = None) +FeedManager(config: dict, methods_dict: dict, redis_client: Optional[redis.asyncio.Redis] = None) ``` Parses and validates `config` into a `FeedConfig` immediately (pydantic v2). `methods_dict` maps names to your async fetch/rerank callables. `redis_client` must be an async `redis.asyncio.Redis`; if `None`, every Wrapper with `cache`/`dedup` degrades to uncached passthrough. Holds no per-request state. @@ -106,9 +111,11 @@ async get_feed(session_id: str, limit: int, cursor: Optional[dict] = None) -> Fe ``` `cursor=None` is treated as the first page (`{}`). Builds a fresh execution context per call and delegates to the executor, then stamps each dict item's final 0-based page position. Returns `FeedResult(data, next_page, has_next_page)`. `next_page` is opaque; store it and pass it back as `cursor`. +Inputs are validated: a non-positive or non-int `limit` and an empty `session_id` raise `ValueError`; a non-dict `cursor` raises `TypeError`. Cursor **contents** are treated as untrusted (they round-trip through clients): a malformed per-node entry (non-dict slice, negative or non-int `offset`, non-str `gen`, non-positive-int `page`) raises `ValueError` naming the node, regardless of any `raise_error` setting. + There is **no request-parameters channel**: only `session_id`, `limit`, and the cursor are threaded to sources. Per-request context must be baked into `subfeed_params` when you build the config, or captured in closures when you build `methods_dict` for that request. -`FeedResult` (`smartfeed.models.base`): `data: list`, `next_page: dict`, `has_next_page: bool`. `data` is a bare list; the pipeline assumes dict items and silently skips non-dict items in stamping/dedup. +`FeedResult` (`smartfeed.models.base`): `data: List[Any]`, `next_page: dict`, `has_next_page: bool`. Items are usually dicts; non-dict items are passed through untouched by stamping/dedup (hence the deliberately loose element type). ## Callable contracts you implement @@ -120,7 +127,7 @@ async def fetch(user_id: str, limit: int, next_page: dict, **kwargs) -> FeedResu - **`**kwargs` is mandatory.** The executor injects `ctx=` into every call. A function without `**kwargs` raises `TypeError` on every invocation. (And if that source has `raise_error=False`, the `TypeError` is swallowed into an empty result, so a signature bug looks exactly like "this source had no items today".) - `user_id` is the feed `session_id` (not necessarily an account id). -- `limit` is this leaf's computed demand for the page. Under a mixer it is a fraction of the page; under a dedup wrapper it is inflated by `overfetch_factor` or a refill deficit. +- `limit` is this leaf's computed demand for the page. Under a mixer it is a fraction of the page; under a dedup wrapper it is the outstanding refill deficit. - `next_page` is this subfeed's own cursor slice (`cursor.get(subfeed_id, {})`), opaque. Echo back whatever you need on the next page. - Return a `FeedResult` (or any object exposing `.data` / `.next_page` / `.has_next_page`; there is no `isinstance` check). - On success each dict item is stamped with `_smartfeed_debug_info.source = subfeed_id`. Non-dict items are passed through untouched. @@ -141,9 +148,10 @@ All defaults below are byte-exact from the code. `FeedConfig` is the top-level s | Field | Type | Default | Meaning | |---|---|---|---| -| `version` | `str` | required | Descriptive tag. Carried but never read by the code. | | `feed` | `FeedNode` | required | Root node. Discriminated union on `type` over the 7 node types. | +Unknown top-level keys (e.g. a legacy `version` tag) are ignored. + ### SubFeed — `type: "subfeed"` | Field | Type | Default | Meaning | @@ -151,7 +159,7 @@ All defaults below are byte-exact from the code. `FeedConfig` is the top-level s | `subfeed_id` | `str` | required | Logical name. Stamped as `_smartfeed_debug_info.source`, is the cursor namespace key, and the identity dedup priority resolves against. | | `method_name` | `str` | required | Key into `methods_dict`. Looked up **before** the try block, so a missing/typo'd name raises `KeyError` even with `raise_error=False`. | | `subfeed_params` | `dict` | `{}` | Static kwargs merged into the fetch call. A key colliding with `user_id`/`limit`/`next_page`/`ctx` raises `TypeError` at call time (not validated at parse time). | -| `raise_error` | `bool` | `True` | `True`: exceptions propagate. `False`: swallow, return an empty page. On the swallow path `next_page` is the full incoming cursor unchanged (not re-nested under `subfeed_id`). | +| `raise_error` | `bool` | `True` | `True`: exceptions propagate. `False`: swallow, return an empty page. On the swallow path `next_page` is `{subfeed_id: }` with `has_next_page=False`, so the source retries the same position next page without clobbering sibling cursors on merge. | | `shuffle` | `bool` | `False` | Shuffle this source's page in place before stamping. Per-page only. | | `dedup_priority` | `int` | `0` | See [dedup_priority](#dedup-and-dedup_priority). | @@ -182,9 +190,9 @@ Wraps exactly one child with optional stages. `cache`, `dedup`, `rerank` are all |---|---|---|---| | `dedup_key` | `str` | required | Item field for identity. Compared as `str(item[key])`, so `1` and `"1"` collide. | | `missing_key_policy` | `"error" \| "keep" \| "drop"` | `"error"` | When an item lacks `dedup_key`: `error` raises, `keep` passes it through un-deduped, `drop` discards it. | -| `overfetch_factor` | `int` | `4` | **Deprecated, no effect.** The dedup fetch now pulls exactly the outstanding deficit and refills until the target is filled or the child is exhausted. Kept for config back-compat. | -| `max_refill_loops` | `int` | `2` | **Deprecated, no effect.** Refill now continues until the page is full or the source runs out (bounded by an internal safety cap), so pages are not cut short. | -| `state_ttl` | `int` | `300` | TTL (s) on the Redis seen-set (session-scoped dedup state, both cache and passthrough paths). | +| `state_ttl` | `int` | `300` | TTL (s) on the Redis seen-set (session-scoped dedup state, both cache and passthrough paths). Refreshed on every warm read, like the cache TTL. | + +(The pre-release `overfetch_factor` / `max_refill_loops` knobs are gone: the dedup fetch pulls exactly the outstanding deficit and refills until the page is full or the child is exhausted. Unknown keys in the config are ignored, so configs still carrying them parse.) `WrapperRerank`: @@ -248,8 +256,8 @@ If a source runs dry the other backfills; if the positional source is exhausted | `node_id` | `str` | required | Also the merged-cursor key holding `{"page": n}`. | | `item_from` | `{percentage:int, data:node}` | required | Starting branch. Its share falls toward 0. | | `item_to` | `{percentage:int, data:node}` | required | Ending branch. Its share rises toward 100, then clamps. | -| `step` | `int` | required | Percentage points shifted per segment. | -| `size_to_step` | `int` | required | Items per shift. Used as a `range()` step; `0` raises at runtime, not at parse time. | +| `step` | `int` | required | Percentage points shifted per segment. Must be > 0 (validated at parse). | +| `size_to_step` | `int` | required | Items per shift. Must be > 0 (validated at parse). | | `dedup_priority` | `int` | `0` | | **MergerAppendDistribute** — `type: "merger_distribute"` (note the literal). Round-robin by a key so no two adjacent items share it. @@ -287,13 +295,13 @@ Every node has `dedup_priority: int = 0`. When a Wrapper's dedup sees the same k All SmartFeed state lives in Redis under `sf:{session_id}:...`. `config_hash` is the first 8 hex chars of an md5 over the node's sorted JSON dump, so any config change moves the key and old data ages out. ``` -sf:{session_id}:{cache_key or node_id}:{config_hash} cached session batch (JSON) +sf:{session_id}:{cache_key or node_id}:{config_hash} cached session batch (Redis LIST of orjson items) sf:{session_id}:{cache_key or node_id}:{config_hash}:meta gen nonce, child cursor, child has_next -sf:{session_id}:{cache_key or node_id}:{config_hash}:coldlock cold-build lock (ttl 30s) +sf:{session_id}:{cache_key or node_id}:{config_hash}:coldlock cold-build lock (ttl 10s) sf:{session_id}:{node_id}:{config_hash}:seen dedup seen-set (session-scoped; cache + passthrough) sf:{session_id}:{cache_key}:{child_config_hash}:{segment} shared base segment (when cache_key set) sf:{session_id}:{cache_key}:{child_config_hash}:{segment}:meta shared base segment meta -sf:{session_id}:{cache_key}:{child_config_hash}:{segment}:lock shared segment cold-build lock (ttl 30s) +sf:{session_id}:{cache_key}:{child_config_hash}:{segment}:lock shared segment cold-build lock (ttl 10s) ``` (`{segment}` is `"0"` for page 1 and a short hash of the continuation cursor for later windows.) @@ -306,24 +314,25 @@ Note the `next_page` cursor shape is config-dependent: a cached wrapper hides th Error handling is **opt-in per node**. By default (`raise_error=True` everywhere) one failing source or rerank fails the entire `get_feed` call, including sibling branches that already succeeded — `asyncio.gather` runs without `return_exceptions=True`, and `FeedManager` has no try/except of its own. For fail-soft behavior, set `raise_error=False` on the specific SubFeed/rerank nodes. -Four sharp edges an integrator will hit: +Five sharp edges an integrator will hit: 1. **Missing `method_name`** is looked up before the try block, so it raises `KeyError` regardless of `raise_error`. 2. **Rerank length mismatch** always raises `ValueError`, regardless of `rerank.raise_error`. 3. **A subfeed function without `**kwargs`** raises `TypeError` on every call (the injected `ctx` has nowhere to go). With `raise_error=False` this is swallowed into an empty page — a signature bug that reads as "no items". 4. **`subfeed_params` key collisions** (with `user_id`/`limit`/`next_page`/`ctx`) raise `TypeError` at call time, with the same silent-swallow caveat under `raise_error=False`. +5. **Malformed cursor contents** raise `ValueError` regardless of `raise_error` — cursors are validated as untrusted client input (see [Public API](#public-api)). Catch it at your API boundary and treat it as a bad request, not a server error. Silent behavior downgrades (no exception): `cache` without a `redis_client`; `cache_key` without `cache`; divergent shared-`cache_key` settings. These ship without error and only show up as degraded cache/dedup metrics. ## `_smartfeed_debug_info` -Every dict item carries a `_smartfeed_debug_info` bundle. The core writes these (all positions are **0-based** at runtime — the model comments saying "1-based" are stale): +Every dict item carries a `_smartfeed_debug_info` bundle — a plain dict; there is no model class behind it. The core writes these (all positions **0-based**): - `source` — the producing `subfeed_id`. - `smartfeed_position` — final 0-based page position (set by `FeedManager`). Wrappers also stamp a pre-rerank position under their `node_id`. - `rerank_position` — 0-based post-rerank index under the wrapper's `node_id`, only when rerank is configured. -Other documented fields (`strategy`, `rrf_score`, `feature_score`, `feature_position`, `total_reranked`, `raw_params`) are **conventions only**; the core never writes them. Your rerank callable may attach them (the bundle allows extra keys). `FeedItem` and `SmartFeedDebugInfo` are typing aids and are never instantiated at runtime — the pipeline operates on plain dicts. +Any other fields (`strategy`, `rrf_score`, `feature_score`, `feature_position`, ...) are **conventions only**; the core never writes them — your fetch/rerank callables may attach whatever extra keys they like. The pipeline operates on plain dicts throughout. ## Integration checklist and footguns diff --git a/docker-compose.yml b/docker-compose.yml deleted file mode 100644 index 8657196..0000000 --- a/docker-compose.yml +++ /dev/null @@ -1,9 +0,0 @@ -version: '3.8' - -services: - sfredis: - image: redis - container_name: sfredis - command: [ 'redis-server', '--port', '6379' ] - ports: - - '6379:6379' diff --git a/smartfeed/__init__.py b/smartfeed/__init__.py index cb612a7..c6a7df9 100644 --- a/smartfeed/__init__.py +++ b/smartfeed/__init__.py @@ -1,2 +1,39 @@ -from .models import * # noqa from .manager import FeedManager +from .models import ( + BaseNode, + FeedConfig, + FeedNode, + FeedResult, + MergerAppend, + MergerAppendDistribute, + MergerPercentage, + MergerPercentageGradient, + MergerPercentageItem, + MergerPositional, + MixerNode, + SubFeed, + Wrapper, + WrapperCache, + WrapperDedup, + WrapperRerank, +) + +__all__ = [ + "BaseNode", + "FeedConfig", + "FeedManager", + "FeedNode", + "FeedResult", + "MergerAppend", + "MergerAppendDistribute", + "MergerPercentage", + "MergerPercentageGradient", + "MergerPercentageItem", + "MergerPositional", + "MixerNode", + "SubFeed", + "Wrapper", + "WrapperCache", + "WrapperDedup", + "WrapperRerank", +] diff --git a/smartfeed/execution/context.py b/smartfeed/execution/context.py index 0f2bd27..a87d6e4 100644 --- a/smartfeed/execution/context.py +++ b/smartfeed/execution/context.py @@ -9,5 +9,5 @@ @dataclass class ExecutionContext: session_id: str - methods_dict: Dict[str, Callable] - redis: Optional[AsyncRedis] = None + methods_dict: Dict[str, Callable[..., Any]] + redis: Optional[AsyncRedis[Any]] = None diff --git a/smartfeed/execution/executor.py b/smartfeed/execution/executor.py index dace07e..8eb555e 100644 --- a/smartfeed/execution/executor.py +++ b/smartfeed/execution/executor.py @@ -1,29 +1,30 @@ from __future__ import annotations import asyncio -from typing import Any +from typing import Any, Dict, List + +from smartfeed.models.base import BaseNode, FeedResult, MixerNode -from smartfeed.models.base import BaseNode, FeedResult from .context import ExecutionContext from .plans import MixPlan -async def run(node: BaseNode, ctx: ExecutionContext, limit: int, cursor: dict) -> FeedResult: - """Execute a feed node. Dispatches to node.execute() or build_mix_plan().""" - if hasattr(node, "execute"): - return await node.execute( - methods_dict=ctx.methods_dict, - session_id=ctx.session_id, - limit=limit, - cursor=cursor, - ctx=ctx, - ) +async def run(node: BaseNode, ctx: ExecutionContext, limit: int, cursor: Dict[str, Any]) -> FeedResult: + """Execute a feed node. Mixers fan out via a MixPlan; leaf/pipeline nodes execute directly.""" + if isinstance(node, MixerNode): + plan = node.build_mix_plan(ctx=ctx, limit=limit, cursor=cursor) + return await _execute_mix(plan, ctx, cursor) - plan: MixPlan = node.build_mix_plan(ctx=ctx, limit=limit, cursor=cursor) - return await _execute_mix(plan, ctx, cursor) + return await node.execute( + methods_dict=ctx.methods_dict, + session_id=ctx.session_id, + limit=limit, + cursor=cursor, + ctx=ctx, + ) -async def _execute_mix(plan: MixPlan, ctx: ExecutionContext, cursor: dict) -> FeedResult: +async def _execute_mix(plan: MixPlan, ctx: ExecutionContext, cursor: Dict[str, Any]) -> FeedResult: """Execute a MixPlan: run children in parallel, assemble results.""" if not plan.children: merged, merged_cursor = plan.assemble({}, {}) @@ -32,8 +33,8 @@ async def _execute_mix(plan: MixPlan, ctx: ExecutionContext, cursor: dict) -> Fe tasks = [run(c.node, ctx, c.demand, cursor) for c in plan.children] results = await asyncio.gather(*tasks) - buffers = {} - child_cursors = {} + buffers: Dict[str, List[Any]] = {} + child_cursors: Dict[str, Dict[str, Any]] = {} any_has_next = False for child, result in zip(plan.children, results): buffers[child.node_id] = result.data diff --git a/smartfeed/execution/plans.py b/smartfeed/execution/plans.py index 3bd0cfc..fab4cd3 100644 --- a/smartfeed/execution/plans.py +++ b/smartfeed/execution/plans.py @@ -1,17 +1,21 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Callable, List +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple + +if TYPE_CHECKING: + from smartfeed.models.base import BaseNode @dataclass class MixChild: node_id: str - node: Any # BaseNode + node: BaseNode demand: int @dataclass class MixPlan: children: List[MixChild] - assemble: Callable # (buffers: dict[str,list], cursors: dict[str,dict]) -> (list, dict) + # (buffers: {child_id: items}, cursors: {child_id: cursor}) -> (merged items, merged cursor) + assemble: Callable[[Dict[str, List[Any]], Dict[str, Dict[str, Any]]], Tuple[List[Any], Dict[str, Any]]] diff --git a/smartfeed/execution/redis_lock.py b/smartfeed/execution/redis_lock.py index 655fd75..6bfbb54 100644 --- a/smartfeed/execution/redis_lock.py +++ b/smartfeed/execution/redis_lock.py @@ -1,8 +1,10 @@ from __future__ import annotations import secrets +from typing import Any from redis.asyncio import Redis as AsyncRedis +from redis.exceptions import WatchError class RedisLock: @@ -15,7 +17,7 @@ class RedisLock: ... # someone else holds it """ - def __init__(self, redis: AsyncRedis, key: str, ttl: int = 10): + def __init__(self, redis: "AsyncRedis[Any]", key: str, ttl: int = 10): self._redis = redis self._key = key self._ttl = ttl @@ -29,9 +31,21 @@ async def __aenter__(self) -> bool: async def __aexit__(self, *exc: object) -> None: if not self._owned: return - val = await self._redis.get(self._key) - if val is not None: - # Redis clients configured with decode_responses=True return str, not bytes. - token = val.decode() if isinstance(val, bytes) else val - if token == self._token: - await self._redis.delete(self._key) + # Atomic compare-and-delete: WATCH the key so that if it changes between our + # read and the DELETE (our TTL expired and someone else acquired), the EXEC + # aborts instead of deleting the new holder's lock. + async with self._redis.pipeline(transaction=True) as pipe: + try: + await pipe.watch(self._key) + val = await self._redis.get(self._key) + if val is None: + return # lock expired; nothing to release + # Redis clients configured with decode_responses=True return str, not bytes. + token = val.decode() if isinstance(val, bytes) else val + if token != self._token: + return # lock expired and re-acquired by someone else; not ours + pipe.multi() + pipe.delete(self._key) + await pipe.execute() + except WatchError: + pass # key changed under us -> it is no longer ours to delete diff --git a/smartfeed/manager.py b/smartfeed/manager.py index d4ca18c..a385776 100644 --- a/smartfeed/manager.py +++ b/smartfeed/manager.py @@ -1,9 +1,11 @@ from __future__ import annotations -from typing import Any, Dict, Optional +from typing import Any, Callable, Dict, Optional + +from redis.asyncio import Redis as AsyncRedis -from .execution.context import ExecutionContext from .execution import executor as _executor +from .execution.context import ExecutionContext from .models import FeedConfig from .models.base import FeedResult @@ -11,16 +13,11 @@ class FeedManager: def __init__( self, - config: Dict, - methods_dict: Dict, - redis_client: Optional[Any] = None, + config: Dict[str, Any], + methods_dict: Dict[str, Callable[..., Any]], + redis_client: Optional[AsyncRedis[Any]] = None, ) -> None: - if hasattr(FeedConfig, "model_validate"): - # Pydantic v2 - self.config = FeedConfig.model_validate(config) - else: - # Pydantic v1 - self.config = FeedConfig.parse_obj(config) + self.config = FeedConfig.model_validate(config) self.methods_dict = methods_dict self.redis_client = redis_client @@ -28,9 +25,16 @@ async def get_feed( self, session_id: str, limit: int, - cursor: Optional[Dict] = None, + cursor: Optional[Dict[str, Any]] = None, ) -> FeedResult: - cursor = cursor or {} + # Boundary guards: these inputs come from the integrator's request handler + # (and the cursor round-trips through untrusted clients); fail loudly here + # instead of slicing/keying garbage deeper in the pipeline. + if not isinstance(session_id, str) or not session_id: + raise ValueError(f"session_id must be a non-empty string, got {session_id!r}") + if not isinstance(limit, int) or isinstance(limit, bool) or limit <= 0: + raise ValueError(f"limit must be a positive int, got {limit!r}") + cursor = self._validated_cursor(cursor) ctx = ExecutionContext( session_id=session_id, methods_dict=self.methods_dict, @@ -41,3 +45,20 @@ async def get_feed( if isinstance(item, dict): item.setdefault("_smartfeed_debug_info", {})["smartfeed_position"] = i return result + + @staticmethod + def _validated_cursor(cursor: object) -> Dict[str, Any]: + """Runtime gate for the declared Optional[dict] contract. + + Untyped callers can pass anything, so the check runs on an `object` view: + checking the already-narrowed parameter would make the raise provably dead + code to type analysis (pyright reportUnreachable). + """ + if cursor is None: + return {} + if not isinstance(cursor, dict): + raise TypeError( + f"cursor must be the dict from the previous FeedResult.next_page (or None), " + f"got {type(cursor).__name__}" + ) + return cursor diff --git a/smartfeed/models/__init__.py b/smartfeed/models/__init__.py index 1cdb995..84c519f 100644 --- a/smartfeed/models/__init__.py +++ b/smartfeed/models/__init__.py @@ -1,20 +1,20 @@ from __future__ import annotations -from typing import Annotated, Union +from typing import Annotated, Dict, Union -from pydantic import Field +from pydantic import BaseModel, Field, model_validator -from .base import BaseNode, FeedResult, SmartFeedDebugInfo, FeedItem -from .subfeed import SubFeed +from .base import BaseNode, FeedResult, MixerNode from .mixers import ( + MergerAppend, + MergerAppendDistribute, MergerPercentage, + MergerPercentageGradient, MergerPercentageItem, - MergerAppend, MergerPositional, - MergerPercentageGradient, - MergerAppendDistribute, ) -from .wrapper import Wrapper, WrapperCache, WrapperRerank, WrapperDedup +from .subfeed import SubFeed +from .wrapper import Wrapper, WrapperCache, WrapperDedup, WrapperRerank # --------------------------------------------------------------------------- # FeedNode: discriminated union of all node types @@ -37,13 +37,48 @@ # FeedConfig: top-level config model # --------------------------------------------------------------------------- -from pydantic import BaseModel + +def _walk_ids(node: BaseNode, seen: Dict[str, str]) -> None: + """Collect every subfeed_id/node_id in the tree, raising on a duplicate. + + Note the BaseNode check comes BEFORE the generic has-a-.data check: a Wrapper is + both a node AND has .data, and must be visited itself, not skipped over. + """ + node_id = getattr(node, "subfeed_id", None) or getattr(node, "node_id", None) + if node_id is not None: + kind = type(node).__name__ + if node_id in seen: + raise ValueError( + f"duplicate id '{node_id}' ({seen[node_id]} and {kind}): subfeed_id and node_id share " + f"one namespace and must be unique across the config tree -- duplicates silently " + f"collide in cursors and Redis keys" + ) + seen[node_id] = kind + for attr in ("items", "data", "positional", "default", "item_from", "item_to"): + child = getattr(node, attr, None) + if child is None or child is node: + continue + children = child if isinstance(child, list) else [child] + for item in children: + if isinstance(item, BaseNode): + _walk_ids(item, seen) + else: + # MergerPercentageItem and friends wrap their node in .data + inner = getattr(item, "data", None) + if isinstance(inner, BaseNode): + _walk_ids(inner, seen) class FeedConfig(BaseModel): - version: str + # Top-level keys other than `feed` (e.g. a legacy `version` tag) are ignored. feed: FeedNode + @model_validator(mode="after") + def _validate_unique_ids(self) -> "FeedConfig": + seen: Dict[str, str] = {} + _walk_ids(self.feed, seen) + return self + # Rebuild forward refs so that nested FeedNode fields resolve correctly _types_ns = {"FeedNode": FeedNode} @@ -63,9 +98,8 @@ class FeedConfig(BaseModel): __all__ = [ "BaseNode", + "MixerNode", "FeedResult", - "SmartFeedDebugInfo", - "FeedItem", "SubFeed", "MergerPercentage", "MergerPercentageItem", diff --git a/smartfeed/models/base.py b/smartfeed/models/base.py index 96a6fb6..69b3cf5 100644 --- a/smartfeed/models/base.py +++ b/smartfeed/models/base.py @@ -4,59 +4,52 @@ import json from typing import TYPE_CHECKING, Any, Dict, List, Optional -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, PrivateAttr if TYPE_CHECKING: + from smartfeed.execution.context import ExecutionContext from smartfeed.execution.plans import MixPlan -class SmartFeedDebugInfo(BaseModel): - """Metadata stamped by SmartFeed on every item.""" - - # Always present (stamped by SubFeed) - source: str # subfeed_id (regular_tours, recommended_tours, promo_tours) - - # Stamped by FeedManager (top-level position in final feed) - smartfeed_position: Optional[int] = None # 0-based position in page - - # Stamped by subfeed methods (optional, source-specific) - strategy: Optional[str] = None # ML model strategy (model_hot_users, model_cold_users) - - # Stamped by rerank callable (optional, present only when rerank is configured) - rerank_position: Optional[int] = None # position after rerank (1-based) - rrf_score: Optional[float] = None # RRF score - feature_score: Optional[float] = None # feature score from ES coefficients - feature_position: Optional[int] = None # rank by feature score (1-based) - total_reranked: Optional[int] = None # total items in rerank batch - raw_params: Optional[Dict[str, Any]] = None # raw tour params from Redis - - model_config = ConfigDict(extra="allow") - - -class FeedItem(BaseModel): - """One item in the feed output. Wraps tour data + SmartFeed metadata.""" +class BaseNode(BaseModel): + dedup_priority: int = 0 - id: Any - smartfeed_debug_info: Optional[SmartFeedDebugInfo] = Field(default=None, alias="_smartfeed_debug_info") + # Config is immutable once FeedManager validates it, so the hash is computed once + # per node instance instead of on every Redis key construction. + _config_hash: Optional[str] = PrivateAttr(default=None) - model_config = ConfigDict(extra="allow", populate_by_name=True) + def config_hash(self) -> str: + if self._config_hash is None: + # json.dumps with sort_keys for deterministic hashing + raw = json.dumps(self.model_dump(), sort_keys=True, default=str) + self._config_hash = hashlib.md5(raw.encode()).hexdigest()[:8] + return self._config_hash + async def execute( + self, + methods_dict: Dict[str, Any], + session_id: str, + limit: int, + cursor: Dict[str, Any], + ctx: Optional[ExecutionContext] = None, + **params: Any, + ) -> FeedResult: + """Leaf/pipeline nodes (SubFeed, Wrapper) override this; mixers build a MixPlan instead.""" + raise NotImplementedError(f"{type(self).__name__} does not execute directly") -class BaseNode(BaseModel): - dedup_priority: int = 0 - def config_hash(self) -> str: - # Use json.dumps with sort_keys for deterministic hashing - raw = json.dumps(self.model_dump(), sort_keys=True, default=str) - return hashlib.md5(raw.encode()).hexdigest()[:8] +class MixerNode(BaseNode): + """Base for combiner nodes. The executor dispatches on this type: mixers return a + MixPlan of children to run concurrently instead of executing directly.""" - def build_mix_plan(self, *, ctx: Any, limit: int, cursor: dict) -> "MixPlan": - """Mixer nodes override this; leaf nodes (SubFeed, Wrapper) use execute() instead.""" - raise NotImplementedError(f"{type(self).__name__} is not a mixer node") + def build_mix_plan(self, *, ctx: ExecutionContext, limit: int, cursor: Dict[str, Any]) -> MixPlan: + raise NotImplementedError class FeedResult(BaseModel): - data: List + # Items are usually dicts; non-dict items are passed through untouched by + # stamping/dedup, so the element type is deliberately Any. + data: List[Any] next_page: Dict[str, Any] has_next_page: bool @@ -70,7 +63,8 @@ def coerce_feed_node(value: Any) -> Any: if not isinstance(value, dict): return value # Lazy import to break circular dependency - from smartfeed.models import FeedNode # noqa: PLC0415 from pydantic import TypeAdapter + from smartfeed.models import FeedNode + return TypeAdapter(FeedNode).validate_python(value) diff --git a/smartfeed/models/mixers.py b/smartfeed/models/mixers.py index 10b43d9..1c25ac6 100644 --- a/smartfeed/models/mixers.py +++ b/smartfeed/models/mixers.py @@ -1,16 +1,19 @@ from __future__ import annotations from collections import defaultdict, deque -from typing import Any, Dict, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, Field, model_validator -from .base import BaseNode, coerce_feed_node from ..execution.plans import MixChild, MixPlan +from .base import MixerNode, coerce_feed_node +if TYPE_CHECKING: + from ..execution.context import ExecutionContext -def _merge_cursor(child_cursors: Dict[str, dict]) -> dict: - merged: dict = {} + +def _merge_cursor(child_cursors: Dict[str, Dict[str, Any]]) -> Dict[str, Any]: + merged: Dict[str, Any] = {} for c in child_cursors.values(): merged.update(c) return merged @@ -38,7 +41,7 @@ def _coerce_data(cls, values: Any) -> Any: # --------------------------------------------------------------------------- -class MergerPercentage(BaseNode): +class MergerPercentage(MixerNode): type: Literal["merger_percentage"] = "merger_percentage" node_id: str items: List[MergerPercentageItem] @@ -47,15 +50,15 @@ class MergerPercentage(BaseNode): def build_mix_plan( self, *, - ctx: Any, + ctx: ExecutionContext, limit: int, - cursor: dict, + cursor: Dict[str, Any], ) -> MixPlan: total_pct = sum(item.percentage for item in self.items) # Compute per-child demand with remainder distribution demands: List[int] = [] - remainders: List[tuple] = [] + remainders: List[Tuple[int, int]] = [] for idx, item in enumerate(self.items): raw = limit * item.percentage child_limit = raw // 100 @@ -81,9 +84,9 @@ def build_mix_plan( ] def assemble( - buffers: Dict[str, list], - child_cursors: Dict[str, dict], - ) -> Tuple[List[Any], dict]: + buffers: Dict[str, List[Any]], + child_cursors: Dict[str, Dict[str, Any]], + ) -> Tuple[List[Any], Dict[str, Any]]: # Simple concatenation: demand already ensures correct proportions merged_data: List[Any] = [] for child in children: @@ -98,7 +101,7 @@ def assemble( # --------------------------------------------------------------------------- -class MergerAppend(BaseNode): +class MergerAppend(MixerNode): type: Literal["merger_append"] = "merger_append" node_id: str items: List[Any] # list of BaseNode subclasses @@ -116,15 +119,17 @@ def _coerce_items(cls, values: Any) -> Any: def build_mix_plan( self, *, - ctx: Any, + ctx: ExecutionContext, limit: int, - cursor: dict, + cursor: Dict[str, Any], ) -> MixPlan: # Each child gets equal demand share; leftover goes to first children n = len(self.items) if n == 0: - def assemble_empty(buffers: Dict[str, list], child_cursors: Dict[str, dict]) -> Tuple[List[Any], dict]: + def assemble_empty( + buffers: Dict[str, List[Any]], child_cursors: Dict[str, Dict[str, Any]] + ) -> Tuple[List[Any], Dict[str, Any]]: return [], _merge_cursor(child_cursors) return MixPlan(children=[], assemble=assemble_empty) @@ -142,9 +147,9 @@ def assemble_empty(buffers: Dict[str, list], child_cursors: Dict[str, dict]) -> ] def assemble( - buffers: Dict[str, list], - child_cursors: Dict[str, dict], - ) -> Tuple[List[Any], dict]: + buffers: Dict[str, List[Any]], + child_cursors: Dict[str, Dict[str, Any]], + ) -> Tuple[List[Any], Dict[str, Any]]: merged_data: List[Any] = [] for child in children: merged_data.extend(buffers.get(child.node_id, [])) @@ -159,7 +164,7 @@ def assemble( # --------------------------------------------------------------------------- -class MergerPositional(BaseNode): +class MergerPositional(MixerNode): type: Literal["merger_positional"] = "merger_positional" node_id: str positions: List[int] = [] @@ -179,9 +184,9 @@ def _coerce_children(cls, values: Any) -> Any: def build_mix_plan( self, *, - ctx: Any, + ctx: ExecutionContext, limit: int, - cursor: dict, + cursor: Dict[str, Any], ) -> MixPlan: # positions are 1-indexed; determine how many positional slots fall within [1..limit] pos_slots = [p for p in self.positions if 1 <= p <= limit] @@ -202,9 +207,9 @@ def build_mix_plan( children = [positional_child, default_child] def assemble( - buffers: Dict[str, list], - child_cursors: Dict[str, dict], - ) -> Tuple[List[Any], dict]: + buffers: Dict[str, List[Any]], + child_cursors: Dict[str, Dict[str, Any]], + ) -> Tuple[List[Any], Dict[str, Any]]: pos_items = deque(buffers.get(positional_child.node_id, [])) def_items = deque(buffers.get(default_child.node_id, [])) @@ -228,56 +233,90 @@ def assemble( # --------------------------------------------------------------------------- -class MergerPercentageGradient(BaseNode): +class MergerPercentageGradient(MixerNode): """Percentage-based merger that shifts the ratio over pages.""" type: Literal["merger_percentage_gradient"] = "merger_percentage_gradient" node_id: str item_from: MergerPercentageItem item_to: MergerPercentageItem - step: int - size_to_step: int + step: int = Field(gt=0) # a zero/negative shift cannot progress (use merger_percentage for a fixed split) + size_to_step: int = Field(gt=0) # segment length; 0 would make the gradient walk undefined dedup_priority: int = 0 - def _calculate_demands(self, page: int, limit: int) -> tuple: - percentage_from = self.item_from.percentage - percentage_to = self.item_to.percentage - start_position = limit * (page - 1) - first_iter = True + def _percentages_at(self, k: int) -> Tuple[int, int]: + """Percentage split at 1-indexed gradient segment k. + + Closed form of "shift by `step` once per segment while `to` < 100, clamping + to (0, 100) on overshoot" -- so demand computation never has to replay the + shift history from segment 1. + """ + pct_from = self.item_from.percentage + pct_to = self.item_to.percentage + if k <= 1 or pct_to >= 100: + return pct_from, pct_to + # The walk freezes at the first shift where `to` reaches 100 (guard stops + # shifting) or an overshoot clamps the pair to (0, 100). + freeze = min( + -((pct_to - 100) // self.step), # ceil((100 - pct_to) / step): `to` reaches 100 + pct_from // self.step + 1, # first shift where `from` would pass 0 + ) + shifts = min(k - 1, freeze) + raw_from = pct_from - self.step * shifts + raw_to = pct_to + self.step * shifts + if raw_to > 100 or raw_from < 0: + return 0, 100 + return raw_from, raw_to + + def _calculate_demands(self, page: int, limit: int) -> Tuple[int, int, List[Dict[str, int]]]: + """Demands and interleave segments for THIS page's window only. + + Iterates just the gradient segments overlapping [limit*(page-1), limit*page) + -- at most limit/size_to_step + 2 of them -- instead of replaying the whole + scroll history, so per-request work is bounded by the page size no matter + how deep (or how tampered) `page` is. + """ + page_start = limit * (page - 1) + page_end = limit * page limit_from = 0 limit_to = 0 - segments: List[Dict] = [] - - for i in range(self.size_to_step, limit * page + self.size_to_step, self.size_to_step): - if not first_iter and percentage_to < 100: - percentage_from -= self.step - percentage_to += self.step - if percentage_to > 100 or percentage_from < 0: - percentage_from = 0 - percentage_to = 100 - - if i > start_position: - iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) - start_position = i - from_take = iter_limit * percentage_from // 100 - to_take = iter_limit - from_take - limit_from += from_take - limit_to += to_take - segments.append({"limit": iter_limit, "from_take": from_take, "to_take": to_take}) - - if first_iter: - first_iter = False + segments: List[Dict[str, int]] = [] + + first_seg = page_start // self.size_to_step + 1 + last_seg = -(-page_end // self.size_to_step) # ceil + for k in range(first_seg, last_seg + 1): + lo = max(page_start, (k - 1) * self.size_to_step) + hi = min(k * self.size_to_step, page_end) + iter_limit = hi - lo + if iter_limit <= 0: + continue + pct_from, _ = self._percentages_at(k) + from_take = iter_limit * pct_from // 100 + to_take = iter_limit - from_take + limit_from += from_take + limit_to += to_take + segments.append({"limit": iter_limit, "from_take": from_take, "to_take": to_take}) return limit_from, limit_to, segments def build_mix_plan( self, *, - ctx: Any, + ctx: ExecutionContext, limit: int, - cursor: dict, + cursor: Dict[str, Any], ) -> MixPlan: - page = cursor.get(self.node_id, {}).get("page", 1) + node_cursor = cursor.get(self.node_id, {}) + if not isinstance(node_cursor, dict): + raise ValueError( + f"MergerPercentageGradient '{self.node_id}': cursor entry must be a dict, " + f"got {type(node_cursor).__name__}" + ) + page = node_cursor.get("page", 1) + if not isinstance(page, int) or isinstance(page, bool) or page < 1: + raise ValueError( + f"MergerPercentageGradient '{self.node_id}': cursor 'page' must be a positive int, got {page!r}" + ) limit_from, limit_to, segments = self._calculate_demands(page, limit) from_child = MixChild( @@ -293,9 +332,9 @@ def build_mix_plan( children = [from_child, to_child] def assemble( - buffers: Dict[str, list], - child_cursors: Dict[str, dict], - ) -> Tuple[List[Any], dict]: + buffers: Dict[str, List[Any]], + child_cursors: Dict[str, Dict[str, Any]], + ) -> Tuple[List[Any], Dict[str, Any]]: from_data = list(buffers.get(from_child.node_id, [])) to_data = list(buffers.get(to_child.node_id, [])) result: List[Any] = [] @@ -319,7 +358,7 @@ def assemble( # --------------------------------------------------------------------------- -class MergerAppendDistribute(BaseNode): +class MergerAppendDistribute(MixerNode): """Append merger that round-robins items by a distribution key.""" type: Literal["merger_distribute"] = "merger_distribute" @@ -339,11 +378,11 @@ def _coerce_items(cls, values: Any) -> Any: values["items"] = [coerce_feed_node(item) if isinstance(item, dict) else item for item in raw_items] return values - def _uniform_distribute(self, data: list) -> list: + def _uniform_distribute(self, data: List[Any]) -> List[Any]: if self.sorting_key: data = sorted(data, key=lambda x: x[self.sorting_key], reverse=self.sorting_desc) - grouped_entries: Dict[Any, deque] = defaultdict(deque) + grouped_entries: Dict[Any, deque[Any]] = defaultdict(deque) for entry in data: grouped_entries[entry[self.distribution_key]].append(entry) @@ -364,9 +403,9 @@ def _uniform_distribute(self, data: list) -> list: def build_mix_plan( self, *, - ctx: Any, + ctx: ExecutionContext, limit: int, - cursor: dict, + cursor: Dict[str, Any], ) -> MixPlan: children = [ MixChild( @@ -378,9 +417,9 @@ def build_mix_plan( ] def assemble( - buffers: Dict[str, list], - child_cursors: Dict[str, dict], - ) -> Tuple[List[Any], dict]: + buffers: Dict[str, List[Any]], + child_cursors: Dict[str, Dict[str, Any]], + ) -> Tuple[List[Any], Dict[str, Any]]: all_items: List[Any] = [] for child in children: all_items.extend(buffers.get(child.node_id, [])) diff --git a/smartfeed/models/subfeed.py b/smartfeed/models/subfeed.py index 61b7724..5f5ceb2 100644 --- a/smartfeed/models/subfeed.py +++ b/smartfeed/models/subfeed.py @@ -1,10 +1,13 @@ from __future__ import annotations from random import shuffle -from typing import Any, Dict, List, Literal +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional from .base import BaseNode, FeedResult +if TYPE_CHECKING: + from smartfeed.execution.context import ExecutionContext + class SubFeed(BaseNode): type: Literal["subfeed"] = "subfeed" @@ -16,20 +19,28 @@ class SubFeed(BaseNode): async def execute( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - cursor: dict, + cursor: Dict[str, Any], + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: method = methods_dict[self.method_name] subfeed_cursor = cursor.get(self.subfeed_id, {}) + # Like the missing-method KeyError, this raises regardless of raise_error: + # a malformed cursor is client tampering, not a source failure to fail soft on. + if not isinstance(subfeed_cursor, dict): + raise ValueError( + f"SubFeed '{self.subfeed_id}': cursor entry must be a dict, got {type(subfeed_cursor).__name__}" + ) try: result: FeedResult = await method( user_id=session_id, limit=limit, next_page=subfeed_cursor, + ctx=ctx, **params, **self.subfeed_params, ) diff --git a/smartfeed/models/wrapper.py b/smartfeed/models/wrapper.py index 8cf1928..e61e5ad 100644 --- a/smartfeed/models/wrapper.py +++ b/smartfeed/models/wrapper.py @@ -4,19 +4,30 @@ import hashlib import json import secrets -from typing import Any, Dict, List, Literal, Optional - -from pydantic import BaseModel, model_validator +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Tuple import orjson +from pydantic import BaseModel, PrivateAttr, model_validator + from smartfeed.execution import executor as _executor + from .base import BaseNode, FeedResult, coerce_feed_node +if TYPE_CHECKING: + from smartfeed.execution.context import ExecutionContext + # Safety bound for the dedup refill loop: stop scanning after this many fetched items # per `target` when the child keeps yielding duplicates (guards against a misbehaving # source that returns has_next=True forever). Well above any realistic duplicate rate. _REFILL_SCAN_FACTOR = 50 +# Cold-build lock: how long a builder may hold the lock, and how a loser waits it out. +# The loser's poll budget covers the FULL lock TTL -- giving up earlier and rebuilding +# without the lock would double-fetch the child and write a competing generation. +_COLD_LOCK_TTL = 10 +_LOCK_POLL_INTERVAL = 0.1 +_LOCK_POLL_ROUNDS = int(_COLD_LOCK_TTL / _LOCK_POLL_INTERVAL) + class WrapperCache(BaseModel): session_size: int = 300 @@ -31,12 +42,6 @@ class WrapperRerank(BaseModel): class WrapperDedup(BaseModel): dedup_key: str missing_key_policy: Literal["error", "keep", "drop"] = "error" - # Deprecated / no effect: the dedup fetch now pulls exactly the outstanding - # deficit and refills until the target is filled or the child is exhausted, so - # it neither over-fetches (no item loss) nor gives up early (no short pages). - # Kept for config back-compat. - overfetch_factor: int = 4 - max_refill_loops: int = 2 state_ttl: int = 300 # TTL for Redis seen-set (seconds) @@ -49,6 +54,9 @@ class Wrapper(BaseNode): data: Any # BaseNode subclass; typed as Any to support discriminated union deserialization cache_key: Optional[str] = None + # Config is immutable after validation, so the priority map is computed once. + _priorities: Optional[Dict[str, int]] = PrivateAttr(default=None) + @model_validator(mode="before") @classmethod def _coerce_data(cls, values: Any) -> Any: @@ -58,11 +66,11 @@ def _coerce_data(cls, values: Any) -> Any: async def execute( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - cursor: dict, - ctx: Any = None, + cursor: Dict[str, Any], + ctx: Optional[ExecutionContext] = None, **params: Any, ) -> FeedResult: """Entry point. Routes to cached or passthrough path.""" @@ -74,10 +82,15 @@ async def execute( # -- dedup ----------------------------------------------------------------- def _resolve_dedup_priorities(self) -> Dict[str, int]: - """Build {subfeed_id: effective_priority} by walking the config tree.""" - result: Dict[str, int] = {} - self._collect_priorities(self.data, override_priority=0, out=result) - return result + """Build {subfeed_id: effective_priority} by walking the config tree. + + The config never changes after validation, so the walk runs once per + wrapper instance instead of once per dedup round.""" + if self._priorities is None: + result: Dict[str, int] = {} + self._collect_priorities(self.data, override_priority=0, out=result) + self._priorities = result + return self._priorities @staticmethod def _collect_priorities(node: BaseNode, override_priority: int, out: Dict[str, int]) -> None: @@ -92,25 +105,25 @@ def _collect_priorities(node: BaseNode, override_priority: int, out: Dict[str, i out[node.subfeed_id] = effective return - # Walk children: look in known child-bearing attributes + # Walk children: look in known child-bearing attributes. + # The BaseNode check comes BEFORE the .data unwrap: a Wrapper is the one + # node that ALSO has .data, and it must be visited itself so its own + # dedup_priority participates in propagation. for attr in ("items", "data", "positional", "default", "item_from", "item_to"): child = getattr(node, attr, None) if child is None: continue - if isinstance(child, list): - for item in child: - child_node = getattr(item, "data", item) - if isinstance(child_node, BaseNode): - Wrapper._collect_priorities(child_node, effective, out) - elif isinstance(child, BaseModel): - # MergerPercentageItem has .data - child_node = getattr(child, "data", child) - if isinstance(child_node, BaseNode): - Wrapper._collect_priorities(child_node, effective, out) - elif isinstance(child, BaseNode): - Wrapper._collect_priorities(child, effective, out) - - def _dedup(self, data: List, seen: Optional[set] = None) -> List: + children = child if isinstance(child, list) else [child] + for item in children: + if isinstance(item, BaseNode): + Wrapper._collect_priorities(item, effective, out) + elif isinstance(item, BaseModel): + # MergerPercentageItem and friends wrap their node in .data + inner = getattr(item, "data", None) + if isinstance(inner, BaseNode): + Wrapper._collect_priorities(inner, effective, out) + + def _dedup(self, data: List[Any], seen: Optional[Set[str]] = None) -> List[Any]: """Remove duplicates by dedup_key. Higher dedup_priority wins, equal = first-seen wins.""" if not self.dedup: return data @@ -122,8 +135,8 @@ def _dedup(self, data: List, seen: Optional[set] = None) -> List: seen = set() # {key_val: (priority, index_in_result)} for in-batch priority arbitration - batch: Dict[Any, tuple] = {} - result: List = [] + batch: Dict[str, Tuple[int, int]] = {} + result: List[Any] = [] for item in data: if not isinstance(item, dict): @@ -167,28 +180,28 @@ def _dedup(self, data: List, seen: Optional[set] = None) -> List: # -- rerank ---------------------------------------------------------------- - async def _apply_rerank(self, data: list, methods_dict: dict, session_id: str) -> list: + async def _apply_rerank(self, data: List[Any], methods_dict: Dict[str, Any], session_id: str) -> List[Any]: """Call rerank callable from methods_dict. Validates output length.""" if not self.rerank: return data rerank_fn = methods_dict[self.rerank.method_name] original_len = len(data) try: - result = await rerank_fn(data, session_id) + reranked: List[Any] = await rerank_fn(data, session_id) except Exception: if self.rerank.raise_error: raise # Fail-soft: keep original order return data - if len(result) != original_len: + if len(reranked) != original_len: raise ValueError( - f"Rerank '{self.rerank.method_name}' must return exactly " f"{original_len} items, got {len(result)}" + f"Rerank '{self.rerank.method_name}' must return exactly " f"{original_len} items, got {len(reranked)}" ) - return result + return reranked # -- debug stamping -------------------------------------------------------- - def _stamp_pre_rerank(self, data: list) -> None: + def _stamp_pre_rerank(self, data: List[Any]) -> None: """Stamp smartfeed_position on each item (position before rerank).""" for i, item in enumerate(data): if isinstance(item, dict): @@ -196,7 +209,7 @@ def _stamp_pre_rerank(self, data: list) -> None: "smartfeed_position": i, } - def _stamp_post_rerank(self, data: list) -> None: + def _stamp_post_rerank(self, data: List[Any]) -> None: """Stamp rerank_position on each item (position after rerank).""" for i, item in enumerate(data): if isinstance(item, dict): @@ -206,19 +219,26 @@ def _stamp_post_rerank(self, data: list) -> None: async def _passthrough( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - cursor: dict, - ctx: Any, + cursor: Dict[str, Any], + ctx: Optional[ExecutionContext], ) -> FeedResult: """No-cache path: fetch -> dedup (refill to a full page) -> rerank -> return.""" + assert ctx is not None, "wrapper execution requires ctx (the executor always provides it)" if self.dedup: - # Cross-page seen-set from Redis (persists shown ids across pages). - seen_keys = await self._load_seen_set(ctx, session_id) + # Session-scoped seen-set, mirroring the cached path: an empty cursor is a + # fresh scroll, so reset the set (re-requesting page 1 returns page 1); + # otherwise load it so shown ids persist across pages. + if not cursor: + await self._reset_seen_set(ctx, session_id) + seen_keys: Set[str] = set() + else: + seen_keys = await self._load_seen_set(ctx, session_id) current_cursor = dict(cursor) - data: List = [] + data: List[Any] = [] has_next = True scanned = 0 scan_cap = max(limit, 1) * _REFILL_SCAN_FACTOR @@ -239,7 +259,7 @@ async def _passthrough( data.extend(self._dedup(result.data, seen_keys)) # Persist seen-set to Redis for the next page. - new_keys = set() + new_keys: Set[str] = set() key_field = self.dedup.dedup_key for item in data: if isinstance(item, dict) and key_field in item: @@ -265,7 +285,7 @@ def _seen_set_key(self, session_id: str) -> str: """Redis key for cross-page dedup seen-set.""" return f"sf:{session_id}:{self.node_id}:{self.config_hash()}:seen" - async def _load_seen_set(self, ctx: Any, session_id: str) -> set: + async def _load_seen_set(self, ctx: Optional[ExecutionContext], session_id: str) -> Set[str]: """Load seen keys from Redis SET for cross-page dedup.""" if not ctx or not ctx.redis: return set() @@ -273,97 +293,143 @@ async def _load_seen_set(self, ctx: Any, session_id: str) -> set: members = await ctx.redis.smembers(key) return {m.decode() if isinstance(m, bytes) else m for m in members} if members else set() - async def _save_seen_set(self, ctx: Any, session_id: str, new_keys: set) -> None: - """Append new keys to Redis seen-set and refresh TTL.""" + async def _save_seen_set(self, ctx: Optional[ExecutionContext], session_id: str, new_keys: Set[str]) -> None: + """Append new keys to Redis seen-set and refresh TTL (one round-trip).""" if not ctx or not ctx.redis or not new_keys: return key = self._seen_set_key(session_id) - await ctx.redis.sadd(key, *new_keys) assert self.dedup is not None, "seen-set only written when dedup configured" - ttl = self.dedup.state_ttl - await ctx.redis.expire(key, ttl) + pipe = ctx.redis.pipeline() + pipe.sadd(key, *new_keys) + pipe.expire(key, self.dedup.state_ttl) + await pipe.execute() - async def _reset_seen_set(self, ctx: Any, session_id: str) -> None: + async def _reset_seen_set(self, ctx: Optional[ExecutionContext], session_id: str) -> None: """Clear the seen-set so a fresh scroll (empty cursor) starts from page 1.""" if not ctx or not ctx.redis: return await ctx.redis.delete(self._seen_set_key(session_id)) # -- cache helpers --------------------------------------------------------- + # + # The session batch is stored as a Redis LIST of individually-encoded items, so a + # warm page read transfers and parses only the requested window (LRANGE), never + # the whole batch. Shared-base segments stay single blobs: they are always read + # whole (each sharing wrapper reranks the full segment). def _base_key(self, session_id: str) -> str: """Redis key prefix for this wrapper's cache data.""" key_part = self.cache_key or self.node_id return f"sf:{session_id}:{key_part}:{self.config_hash()}" - async def _read_cache(self, ctx: Any, session_id: str) -> Optional[List]: - """Read cached session data from Redis. Returns None on miss.""" - base = self._base_key(session_id) - raw = await ctx.redis.get(base) - if raw is None: - return None - return orjson.loads(raw) - - async def _read_meta(self, ctx: Any, session_id: str) -> Optional[Dict]: + async def _read_meta(self, ctx: ExecutionContext, session_id: str) -> Optional[Dict[str, Any]]: """Read cache metadata (gen, child_cursor, child_has_next) from Redis.""" + assert ctx.redis is not None, "cached path requires redis" base = self._base_key(session_id) raw = await ctx.redis.get(f"{base}:meta") if raw is None: return None - return orjson.loads(raw) + meta: Dict[str, Any] = orjson.loads(raw) + return meta async def _write_cache( self, - ctx: Any, + ctx: ExecutionContext, session_id: str, - data: List, + data: List[Any], gen: str, - child_cursor: Dict, + child_cursor: Dict[str, Any], child_has_next: bool = False, ) -> None: - """Write session data and metadata to Redis with TTL.""" + """Replace session data (as a LIST of encoded items) and metadata atomically.""" assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" + assert ctx.redis is not None, "cached path requires redis" base = self._base_key(session_id) ttl = self.cache.session_ttl + try: + encoded = [orjson.dumps(item) for item in data] + meta = orjson.dumps({"gen": gen, "child_cursor": child_cursor, "child_has_next": child_has_next}) + except TypeError as e: + raise TypeError( + f"Wrapper '{self.node_id}': cache is enabled, so items and child cursors must be " + f"orjson-serializable (str/int/float/bool/None/list/dict): {e}" + ) from e pipe = ctx.redis.pipeline() - pipe.set(base, orjson.dumps(data), ex=ttl) - pipe.set( - f"{base}:meta", - orjson.dumps({"gen": gen, "child_cursor": child_cursor, "child_has_next": child_has_next}), - ex=ttl, - ) + pipe.delete(base) + if encoded: + pipe.rpush(base, *encoded) + pipe.expire(base, ttl) + pipe.set(f"{base}:meta", meta, ex=ttl) await pipe.execute() - async def _touch_ttl(self, ctx: Any, session_id: str) -> None: - """Refresh TTL on data and meta keys (keeps cache alive while user scrolls).""" - if not self.cache: - return + async def _read_snapshot_and_touch( + self, ctx: ExecutionContext, session_id: str, offset: int, limit: int + ) -> Tuple[Optional[Dict[str, Any]], List[Any], int]: + """Read meta + one page window as ONE atomic snapshot and refresh TTLs. + + A single MULTI/EXEC round-trip: reading meta and data as separate commands + would let a concurrent rebuild land in between, pairing old-gen meta with the + replacement batch (torn read) -- `_write_cache` swaps both keys in one + transaction, so reads go through one transaction too. Only the requested + window is transferred and parsed (LRANGE), not the whole session batch. The + TTL refresh rides along: data/meta get session_ttl and the dedup seen-set gets + state_ttl, so dedup state cannot expire mid-scroll while the cache it guards + is being kept alive. + + Returns (meta, page_items, total_batch_len); meta is None on miss. + """ + assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" + assert ctx.redis is not None, "cached path requires redis" base = self._base_key(session_id) ttl = self.cache.session_ttl pipe = ctx.redis.pipeline() + pipe.get(f"{base}:meta") + pipe.llen(base) + pipe.lrange(base, offset, offset + limit - 1) pipe.expire(base, ttl) pipe.expire(f"{base}:meta", ttl) - await pipe.execute() + if self.dedup: + pipe.expire(self._seen_set_key(session_id), self.dedup.state_ttl) + results = await pipe.execute() + raw_meta, total, raw_page = results[0], results[1], results[2] + meta: Optional[Dict[str, Any]] = orjson.loads(raw_meta) if raw_meta is not None else None + page = [orjson.loads(x) for x in raw_page] + return meta, page, int(total) + + async def _serve_new_generation( + self, ctx: ExecutionContext, session_id: str, limit: int, stale_gen: Optional[str] + ) -> Optional[FeedResult]: + """Serve page 1 of the current batch, but only if it is a NEW generation. + + Used by cold-build waiters: the outgoing batch's keys are never deleted before + a rebuild, so a batch still carrying `stale_gen` is the one being replaced -- + serving it would re-show items and hand back an instantly-stale cursor. + """ + meta, page, total = await self._read_snapshot_and_touch(ctx, session_id, 0, limit) + if meta is None or meta.get("gen") == stale_gen: + return None + child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) + return self._page_result(page, 0, limit, total, meta.get("gen", ""), child_has_next) # -- pagination ------------------------------------------------------------ - def _paginate( + def _page_result( self, - data: List, - limit: int, + page_data: List[Any], offset: int, + limit: int, + total: int, gen: str, child_has_next: bool = False, ) -> FeedResult: - """Slice cached data at an absolute offset and build the next cursor. + """Build a page response from an already-sliced window of the session batch. The cursor carries an absolute offset (not a page number), so pagination stays correct even if the client varies `limit` between requests. """ end = offset + limit - page_data = data[offset:end] - has_next = end < len(data) or child_has_next + has_next = end < total or child_has_next next_cursor = { self.node_id: { @@ -382,85 +448,122 @@ def _paginate( async def _execute_with_cache( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - cursor: dict, - ctx: Any, + cursor: Dict[str, Any], + ctx: ExecutionContext, ) -> FeedResult: """Cached path: warm hit -> paginate, stale/miss -> cold build.""" + # Cursors round-trip through untrusted clients: reject malformed contents + # loudly instead of crashing mid-slice or serving a wrong page. my_cursor = cursor.get(self.node_id, {}) + if not isinstance(my_cursor, dict): + raise ValueError( + f"Wrapper '{self.node_id}': cursor entry must be a dict, " + f"got {type(my_cursor).__name__} (pass back next_page verbatim)" + ) cursor_gen = my_cursor.get("gen") + if cursor_gen is not None and not isinstance(cursor_gen, str): + raise ValueError(f"Wrapper '{self.node_id}': cursor 'gen' must be a string, got {cursor_gen!r}") cursor_offset = my_cursor.get("offset", 0) + if not isinstance(cursor_offset, int) or isinstance(cursor_offset, bool) or cursor_offset < 0: + raise ValueError( + f"Wrapper '{self.node_id}': cursor 'offset' must be a non-negative int, got {cursor_offset!r}" + ) # Warm path: try to read existing cache if cursor_gen: - meta = await self._read_meta(ctx, session_id) + meta, page, total = await self._read_snapshot_and_touch(ctx, session_id, cursor_offset, limit) if meta and meta.get("gen") == cursor_gen: - cached_data = await self._read_cache(ctx, session_id) - if cached_data is not None: - # Serve from cache while this offset is still within the batch. - if cursor_offset < len(cached_data): - await self._touch_ttl(ctx, session_id) - child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) - return self._paginate( - cached_data, - limit, - cursor_offset, - cursor_gen, - child_has_next=child_has_next, - ) - # Cache exhausted: rebuild with continuation cursor - return await self._cold_build_locked( - methods_dict, - session_id, - limit, - ctx, - child_cursor=meta.get("child_cursor", {}), - ) + # Serve from cache while this offset is still within the batch. + if cursor_offset < total: + child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) + return self._page_result(page, cursor_offset, limit, total, cursor_gen, child_has_next) + # Cache exhausted: rebuild with continuation cursor + return await self._cold_build_locked( + methods_dict, + session_id, + limit, + ctx, + child_cursor=meta.get("child_cursor", {}), + stale_gen=cursor_gen, + ) + # Stale gen or evicted cache -> fresh build replacing whatever batch exists now. + return await self._cold_build_locked( + methods_dict, + session_id, + limit, + ctx, + child_cursor={}, + stale_gen=meta.get("gen") if meta else None, + ) - # Cold path: no gen or stale gen -> fresh build - return await self._cold_build_locked(methods_dict, session_id, limit, ctx, child_cursor={}) + # Fresh scroll (no gen): rebuild, replacing the current batch if one exists. + # Its gen is the baseline a lock waiter must NOT accept as the new build. + meta = await self._read_meta(ctx, session_id) + return await self._cold_build_locked( + methods_dict, + session_id, + limit, + ctx, + child_cursor={}, + stale_gen=meta.get("gen") if meta else None, + ) async def _cold_build_locked( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - ctx: Any, - child_cursor: Dict, + ctx: ExecutionContext, + child_cursor: Dict[str, Any], + stale_gen: Optional[str] = None, ) -> FeedResult: - """Cold build guarded by a lock so concurrent first requests fetch the child - ONCE: one caller builds and writes the cache, the rest wait and serve from it. - A fresh/continuation build always serves the new batch from offset 0.""" + """Cold build guarded by a lock so concurrent rebuilds fetch the child ONCE: + one caller builds and writes the cache, the rest wait for it and serve the + new batch from offset 0. + + `stale_gen` is the generation being replaced (None if none existed). The old + keys are not deleted before a rebuild, so a waiter must ignore metas still + carrying `stale_gen`: accepting the first meta it sees would re-serve the + outgoing batch and hand back a cursor that is stale the moment the builder + finishes. Waiters poll for the full lock TTL (the builder can hold it no + longer); if no new generation appears by then the builder is dead, so the + second round retries the lock, and only after that do we build unlocked as + a last resort. + """ from smartfeed.execution.redis_lock import RedisLock + assert ctx.redis is not None, "cached path requires redis" lock_key = f"{self._base_key(session_id)}:coldlock" - async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: - if acquired: - return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) - # Another coroutine is building -- wait for its cache, then serve page 1. - for _ in range(50): - await asyncio.sleep(0.1) + for _ in range(2): + async with RedisLock(ctx.redis, lock_key, ttl=_COLD_LOCK_TTL) as acquired: + if acquired: + # Another caller may have finished this same rebuild while we + # queued for the lock -- serve its batch instead of re-fetching. + meta = await self._read_meta(ctx, session_id) + if meta and meta.get("gen") != stale_gen: + served = await self._serve_new_generation(ctx, session_id, limit, stale_gen) + if served is not None: + return served + return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) + # Another coroutine is building -- wait for it to publish a NEW generation. + for _ in range(_LOCK_POLL_ROUNDS): + await asyncio.sleep(_LOCK_POLL_INTERVAL) meta = await self._read_meta(ctx, session_id) - if meta: - cached = await self._read_cache(ctx, session_id) - if cached is not None: - child_has_next = meta.get("child_has_next", bool(meta.get("child_cursor"))) - return self._paginate( - cached, - limit, - 0, - meta.get("gen", ""), - child_has_next=child_has_next, - ) - # Fallback: builder never wrote -> build ourselves. - return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) + if meta and meta.get("gen") != stale_gen: + served = await self._serve_new_generation(ctx, session_id, limit, stale_gen) + if served is not None: + return served + # Two rounds without acquiring or seeing a new generation: build unlocked so + # the request still completes. + return await self._cold_build(methods_dict, session_id, limit, ctx, child_cursor) # -- shared base cache helpers ---------------------------------------------- @staticmethod - def _cursor_segment(child_cursor: Dict) -> str: + def _cursor_segment(child_cursor: Dict[str, Any]) -> str: """Stable short tag for a continuation position, so each page-window of the shared base is its own segment (page 1 = "0").""" if not child_cursor: @@ -468,57 +571,70 @@ def _cursor_segment(child_cursor: Dict) -> str: raw = json.dumps(child_cursor, sort_keys=True, default=str) return hashlib.md5(raw.encode()).hexdigest()[:8] - def _base_shared_key(self, session_id: str, child_cursor: Dict) -> str: + def _base_shared_key(self, session_id: str, child_cursor: Dict[str, Any]) -> str: """Key for a shared-base segment (deduped, NOT reranked), keyed by cache_key and the continuation position so continuations don't collide with page 1.""" seg = self._cursor_segment(child_cursor) return f"sf:{session_id}:{self.cache_key}:{self.data.config_hash()}:{seg}" - async def _read_shared_base(self, ctx: Any, session_id: str, child_cursor: Dict) -> Optional[List]: + async def _read_shared_base( + self, ctx: ExecutionContext, session_id: str, child_cursor: Dict[str, Any] + ) -> Optional[List[Any]]: """Read a shared-base segment from Redis. Returns None on miss.""" + assert ctx.redis is not None, "cached path requires redis" key = self._base_shared_key(session_id, child_cursor) raw = await ctx.redis.get(key) if raw is None: return None - return orjson.loads(raw) + segment: List[Any] = orjson.loads(raw) + return segment async def _write_shared_base( self, - ctx: Any, + ctx: ExecutionContext, session_id: str, - child_cursor: Dict, - data: List, - child_cursor_out: Dict, + child_cursor: Dict[str, Any], + data: List[Any], + child_cursor_out: Dict[str, Any], child_has_next: bool = False, ) -> None: """Write a shared-base segment and its meta to Redis with TTL.""" assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" + assert ctx.redis is not None, "cached path requires redis" key = self._base_shared_key(session_id, child_cursor) ttl = self.cache.session_ttl + try: + payload = orjson.dumps(data) + meta = orjson.dumps({"child_cursor": child_cursor_out, "child_has_next": child_has_next}) + except TypeError as e: + raise TypeError( + f"Wrapper '{self.node_id}': cache is enabled, so items and child cursors must be " + f"orjson-serializable (str/int/float/bool/None/list/dict): {e}" + ) from e pipe = ctx.redis.pipeline() - pipe.set(key, orjson.dumps(data), ex=ttl) - pipe.set( - f"{key}:meta", - orjson.dumps({"child_cursor": child_cursor_out, "child_has_next": child_has_next}), - ex=ttl, - ) + pipe.set(key, payload, ex=ttl) + pipe.set(f"{key}:meta", meta, ex=ttl) await pipe.execute() - async def _read_shared_base_meta(self, ctx: Any, session_id: str, child_cursor: Dict) -> Optional[Dict]: + async def _read_shared_base_meta( + self, ctx: ExecutionContext, session_id: str, child_cursor: Dict[str, Any] + ) -> Optional[Dict[str, Any]]: """Read a shared-base segment's metadata (child_cursor, child_has_next).""" + assert ctx.redis is not None, "cached path requires redis" key = self._base_shared_key(session_id, child_cursor) raw = await ctx.redis.get(f"{key}:meta") if raw is None: return None - return orjson.loads(raw) + meta: Dict[str, Any] = orjson.loads(raw) + return meta async def _fetch_and_dedup( self, - ctx: Any, + ctx: ExecutionContext, target: int, - child_cursor: Dict, - seen: Optional[set] = None, - ) -> tuple: + child_cursor: Dict[str, Any], + seen: Optional[Set[str]] = None, + ) -> Tuple[List[Any], Dict[str, Any], bool]: """Collect up to `target` deduped items from the child. Fetches exactly the outstanding deficit each round and keeps EVERY survivor @@ -532,7 +648,11 @@ async def _fetch_and_dedup( Returns (data, child_cursor, child_has_next). """ cursor = dict(child_cursor) - data: List = [] + data: List[Any] = [] + # One seen-set across ALL refill rounds: a fresh set per round would re-admit + # round-1 ids, letting duplicates survive within one build. + if seen is None: + seen = set() has_next = True scanned = 0 scan_cap = max(target, 1) * _REFILL_SCAN_FACTOR @@ -555,11 +675,11 @@ async def _fetch_and_dedup( async def _cold_build( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, limit: int, - ctx: Any, - child_cursor: Dict, + ctx: ExecutionContext, + child_cursor: Dict[str, Any], ) -> FeedResult: """Build full session: fetch -> dedup -> rerank -> write cache -> paginate page 1.""" assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" @@ -583,19 +703,22 @@ async def _cold_build( # Standard path. Session-scoped seen-set: reset on a fresh scroll (empty # cursor) so re-requesting page 1 returns page 1, and carry it across # rebuilds so a shown id never repeats later in the same scroll. - seen: Optional[set] = None + seen: Optional[Set[str]] = None + preloaded: Set[str] = set() if self.dedup: if not child_cursor: await self._reset_seen_set(ctx, session_id) seen = set() else: seen = await self._load_seen_set(ctx, session_id) + preloaded = set(seen) data, child_cursor_out, child_has_next = await self._fetch_and_dedup( ctx, session_size, child_cursor, seen=seen ) if self.dedup: assert seen is not None - await self._save_seen_set(ctx, session_id, seen) + # SADD only the delta: everything in `preloaded` is already in Redis. + await self._save_seen_set(ctx, session_id, seen - preloaded) self._stamp_pre_rerank(data) data = await self._apply_rerank(data, methods_dict, session_id) if self.rerank: @@ -603,20 +726,21 @@ async def _cold_build( gen = secrets.token_hex(4) await self._write_cache(ctx, session_id, data, gen, child_cursor_out, child_has_next) - return self._paginate(data, limit, 0, gen, child_has_next=child_has_next) + return self._page_result(data[:limit], 0, limit, len(data), gen, child_has_next) async def _build_shared_base( self, - methods_dict: dict, + methods_dict: Dict[str, Any], session_id: str, - ctx: Any, - child_cursor: Dict, - ) -> List: + ctx: ExecutionContext, + child_cursor: Dict[str, Any], + ) -> List[Any]: """Fetch and dedup the shared base data, using a distributed lock to ensure only one caller fetches from child on a cold build. Others wait and read.""" from smartfeed.execution.redis_lock import RedisLock assert self.cache is not None, "cached path; execute() routes cache=None to passthrough" + assert ctx.redis is not None, "cached path requires redis" session_size = self.cache.session_size lock_key = f"{self._base_shared_key(session_id, child_cursor)}:lock" @@ -625,31 +749,32 @@ async def _build_shared_base( if existing is not None: return existing - # Try to acquire the lock - async with RedisLock(ctx.redis, lock_key, ttl=30) as acquired: - if acquired: - # Double-check after acquiring lock - existing = await self._read_shared_base(ctx, session_id, child_cursor) - if existing is not None: - return existing - - base_data, child_cursor_out, child_has_next = await self._fetch_and_dedup( - ctx, session_size, child_cursor - ) - await self._write_shared_base( - ctx, session_id, child_cursor, base_data, child_cursor_out, child_has_next - ) - return base_data - else: - # Another coroutine holds the lock -- poll until the segment appears - for _ in range(50): - await asyncio.sleep(0.1) + # Segments are only ever created (never rebuilt in place), so "segment exists" + # is the completion signal. Waiters poll for the full lock TTL; if the segment + # never appears the builder is dead, so the second round retries the lock, and + # only after that do we fetch unlocked as a last resort. + for _ in range(2): + async with RedisLock(ctx.redis, lock_key, ttl=_COLD_LOCK_TTL) as acquired: + if acquired: + # Double-check after acquiring lock existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing - # Fallback: fetch ourselves if lock holder never wrote + + base_data, child_cursor_out, child_has_next = await self._fetch_and_dedup( + ctx, session_size, child_cursor + ) + await self._write_shared_base( + ctx, session_id, child_cursor, base_data, child_cursor_out, child_has_next + ) + return base_data + # Another coroutine holds the lock -- poll until the segment appears. + for _ in range(_LOCK_POLL_ROUNDS): + await asyncio.sleep(_LOCK_POLL_INTERVAL) existing = await self._read_shared_base(ctx, session_id, child_cursor) if existing is not None: return existing - data, _, _ = await self._fetch_and_dedup(ctx, session_size, child_cursor) - return data + # Two rounds without acquiring or seeing the segment: fetch unlocked so the + # request still completes. + data, _, _ = await self._fetch_and_dedup(ctx, session_size, child_cursor) + return data diff --git a/smartfeed/py.typed b/smartfeed/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/tests/sources.py b/tests/sources.py index c25561d..28f954f 100644 --- a/tests/sources.py +++ b/tests/sources.py @@ -22,7 +22,7 @@ from __future__ import annotations import asyncio -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional from smartfeed.models.base import FeedResult from smartfeed.models.subfeed import SubFeed @@ -35,6 +35,7 @@ # Source # --------------------------------------------------------------------------- + class ScriptedSource: """Stateful, offset-paginated source over a finite pool of item dicts.""" @@ -54,7 +55,7 @@ async def __call__(self, user_id, limit, next_page, **kw) -> FeedResult: pos = (next_page or {}).get("pos", 0) if self.latency: await asyncio.sleep(self.latency) - chunk = self.pool[pos: pos + limit] + chunk = self.pool[pos : pos + limit] # Fresh copies so downstream in-place stamping never mutates the pool. data = [dict(it) for it in chunk] for it in data: @@ -71,6 +72,7 @@ async def __call__(self, user_id, limit, next_page, **kw) -> FeedResult: # Pool factories (all deterministic -- no RNG, so failures are reproducible) # --------------------------------------------------------------------------- + def unique_pool(n: int, start: int = 0, val_prefix: str = "u") -> List[Dict]: """Ids start..start+n-1, no duplicates. Isolates *item loss* from dedup.""" return [{"id": start + k, "val": f"{val_prefix}{start + k}"} for k in range(n)] @@ -106,6 +108,7 @@ def adjacent_dup_pool(n_unique: int) -> List[Dict]: # Context / node builders # --------------------------------------------------------------------------- + def make_ctx(methods: Dict, redis=None, session_id: str = "s1") -> ExecutionContext: return ExecutionContext(session_id=session_id, methods_dict=methods, redis=redis) @@ -121,9 +124,7 @@ def wrapper( session_size: Optional[int] = None, session_ttl: int = 300, dedup_key: Optional[str] = None, - overfetch_factor: int = 4, - max_refill_loops: int = 2, - missing_key_policy: str = "error", + missing_key_policy: Literal["error", "keep", "drop"] = "error", rerank_method: Optional[str] = None, rerank_raise: bool = True, cache_key: Optional[str] = None, @@ -136,8 +137,6 @@ def wrapper( if dedup_key is not None: dedup = WrapperDedup( dedup_key=dedup_key, - overfetch_factor=overfetch_factor, - max_refill_loops=max_refill_loops, missing_key_policy=missing_key_policy, ) rerank = None @@ -157,6 +156,7 @@ def wrapper( # Drain + invariant assertions # --------------------------------------------------------------------------- + async def drain(node, ctx, limit: int, max_pages: int = 500, cursor=None) -> List[List]: """Page until ``has_next_page`` is False. Returns the list of page-data lists. @@ -173,8 +173,7 @@ async def drain(node, ctx, limit: int, max_pages: int = 500, cursor=None) -> Lis return pages cursor = r.next_page raise AssertionError( - f"pagination did not terminate within {max_pages} pages " - f"(infinite loop / non-advancing cursor)" + f"pagination did not terminate within {max_pages} pages " f"(infinite loop / non-advancing cursor)" ) diff --git a/tests/test_async_concurrency.py b/tests/test_async_concurrency.py index 65e2cf4..46f10c6 100644 --- a/tests/test_async_concurrency.py +++ b/tests/test_async_concurrency.py @@ -4,13 +4,12 @@ import asyncio import time from dataclasses import dataclass, field -from typing import Any, List, Optional +from typing import Optional import pytest from smartfeed.models.base import FeedResult from smartfeed.models.subfeed import SubFeed -from smartfeed.models.wrapper import Wrapper, WrapperCache, WrapperDedup from smartfeed.models.mixers import MergerAppend, MergerPercentage, MergerPercentageItem, MergerPositional from smartfeed.execution.context import ExecutionContext from smartfeed.execution import executor as run_executor @@ -20,9 +19,11 @@ # Helpers: concurrency tracker + loop block monitor # --------------------------------------------------------------------------- + @dataclass class ConcurrencyTracker: """Tracks how many leaf calls are in-flight at the same time.""" + current: int = 0 peak: int = 0 _lock: asyncio.Lock = field(default_factory=asyncio.Lock) @@ -72,6 +73,7 @@ async def _run(self): # Slow subfeed factory # --------------------------------------------------------------------------- + def make_slow_subfeed_method( source: str, latency: float, @@ -107,6 +109,7 @@ async def method(user_id: str, limit: int, next_page: dict, **kw) -> FeedResult: # Tests # --------------------------------------------------------------------------- + class TestParallelSubfeeds: """Verify that sibling subfeeds in a mixer run concurrently.""" @@ -121,7 +124,6 @@ async def test_two_subfeeds_run_in_parallel(self): "slow_b": make_slow_subfeed_method("b", latency, tracker, calls), } ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) - executor = None node = MergerAppend( node_id="par", @@ -156,7 +158,6 @@ async def test_three_percentage_subfeeds_parallel(self): "slow_c": make_slow_subfeed_method("c", latency, tracker, calls), } ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) - executor = None node = MergerPercentage( node_id="pct", @@ -187,7 +188,6 @@ async def test_positional_branches_parallel(self): "slow_regular": make_slow_subfeed_method("regular", latency, tracker, calls), } ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) - executor = None node = MergerPositional( node_id="pos", @@ -222,7 +222,6 @@ async def test_loop_not_blocked_under_load(self): subfeeds.append(SubFeed(subfeed_id=name, method_name=name)) ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) - executor = None node = MergerAppend(node_id="big", items=subfeeds) @@ -243,12 +242,8 @@ async def test_nested_mixers_parallel(self): calls = {} latency = 0.05 - methods = { - f"slow_{x}": make_slow_subfeed_method(x, latency, tracker, calls) - for x in ("a", "b", "c", "d") - } + methods = {f"slow_{x}": make_slow_subfeed_method(x, latency, tracker, calls) for x in ("a", "b", "c", "d")} ctx = ExecutionContext(session_id="s1", methods_dict=methods, redis=None) - executor = None node = MergerAppend( node_id="outer", diff --git a/tests/test_bug_cold_build_stampede.py b/tests/test_bug_cold_build_stampede.py index 1a354e7..8e675d2 100644 --- a/tests/test_bug_cold_build_stampede.py +++ b/tests/test_bug_cold_build_stampede.py @@ -1,4 +1,4 @@ -"""BUG #5 -- cold-build thundering herd on the standard (non-shared) cached path. +"""Regression: cold-build thundering herd on the standard (non-shared) cached path. Only the cache_key path takes a RedisLock. For a normal cached wrapper, N concurrent first requests for one session each fetch the full session_size from @@ -29,9 +29,7 @@ async def test_concurrent_cold_build_fetches_child_once(): node = S.wrapper(S.subfeed("src", "src"), session_size=50) ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) - results = await asyncio.gather(*[ - run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) - ]) + results = await asyncio.gather(*[run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N)]) assert all(len(r.data) == 10 for r in results) # One winner should fetch the child; the rest should read the cache it wrote. @@ -44,9 +42,7 @@ async def test_concurrent_cold_build_single_generation(): node = S.wrapper(S.subfeed("src", "src"), session_size=50) ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) - results = await asyncio.gather(*[ - run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) - ]) + results = await asyncio.gather(*[run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N)]) gens = {r.next_page["w"]["gen"] for r in results} assert len(gens) == 1, f"concurrent callers got {len(gens)} different generations: {gens}" @@ -55,10 +51,7 @@ async def test_concurrent_cold_build_single_generation(): async def test_shared_cache_path_already_locks_control(): """Control: the cache_key path DOES take a lock, so it fetches once.""" src = S.ScriptedSource(S.unique_pool(200), latency=0.05) - node = S.wrapper(S.subfeed("src", "src"), node_id="w", - session_size=50, cache_key="shared") + node = S.wrapper(S.subfeed("src", "src"), node_id="w", session_size=50, cache_key="shared") ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) - await asyncio.gather(*[ - run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N) - ]) + await asyncio.gather(*[run_executor.run(node, ctx, limit=10, cursor={}) for _ in range(N)]) assert src.calls == 1, f"shared path should fetch once, got {src.calls}" diff --git a/tests/test_bug_lock_loser_gives_up_early.py b/tests/test_bug_lock_loser_gives_up_early.py new file mode 100644 index 0000000..44e7b20 --- /dev/null +++ b/tests/test_bug_lock_loser_gives_up_early.py @@ -0,0 +1,52 @@ +"""Regression: the cold-build lock loser gives up after 5s of +polling while the winner's lock lives 30s. + +`_cold_build_locked` polls 50 x 0.1s and then falls through to an UNLOCKED +`_cold_build`. Any cold build slower than 5s therefore gets a second, unguarded +child fetch racing the real builder: double upstream load and two competing +gens, one of which is orphaned (that client's next page triggers yet another +full rebuild). + +Correct behavior: the loser waits out the winner (poll budget >= lock TTL, or +re-attempt the lock) so the child is fetched exactly once. + +Timing is compressed uniformly (all asyncio.sleep x0.1) so the >5s-nominal +build finishes in <1s of real time; the poll-budget/build-time ratio that +triggers the bug is preserved. +""" + +import asyncio + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_slow_cold_build_still_fetches_child_once(monkeypatch): + real_sleep = asyncio.sleep + + async def compressed_sleep(delay: float) -> None: + await real_sleep(delay * 0.1) + + monkeypatch.setattr(asyncio, "sleep", compressed_sleep) + + # Nominal 8s child latency: longer than the old code's 5s poll budget, + # comfortably shorter than the 10s lock TTL the fixed code waits out. + src = S.ScriptedSource(S.unique_pool(100), latency=8.0) + node = S.wrapper(S.subfeed("src", "src"), session_size=50) + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + a, b = await asyncio.gather( + run_executor.run(node, ctx, limit=10, cursor={}), + run_executor.run(node, ctx, limit=10, cursor={}), + ) + + assert src.calls == 1, ( + f"expected 1 cold fetch (loser waits for the lock holder), got {src.calls}: " + f"the loser gave up polling before the winner finished and rebuilt without the lock" + ) + gens = {a.next_page["w"]["gen"], b.next_page["w"]["gen"]} + assert len(gens) == 1, f"concurrent callers got competing generations: {gens}" diff --git a/tests/test_bug_lock_release_not_atomic.py b/tests/test_bug_lock_release_not_atomic.py new file mode 100644 index 0000000..da105aa --- /dev/null +++ b/tests/test_bug_lock_release_not_atomic.py @@ -0,0 +1,62 @@ +"""Regression: RedisLock release is GET-then-DELETE, not atomic. + +If the lock expires after the release-path GET returned the holder's own token +but before the DELETE executes, and another coroutine acquires the lock in that +window, the DELETE removes the NEW holder's lock: mutual exclusion is silently +lost for a third contender. + +Correct behavior: release must be compare-and-delete in one atomic step (Lua), +so a holder can only ever delete its own token. This test forces the expiry + +re-acquire interleaving deterministically via a hook between GET and DELETE. + +(The related no-TTL-extension concern -- a build outliving the 30s +lock TTL -- is a design decision on the fix and is not encoded here.) +""" + +from typing import Awaitable, Callable, Optional + +import pytest +import fakeredis.aioredis + +from smartfeed.execution.redis_lock import RedisLock + +GetHook = Callable[[str, Optional[bytes]], Awaitable[None]] + + +class HookedRedis: + """Proxy that awaits a hook after each GET returns (to interleave expiry + + re-acquisition between the lock's release GET and DELETE).""" + + def __init__(self, inner): + self._inner = inner + self.on_get: Optional[GetHook] = None + + def __getattr__(self, name): + return getattr(self._inner, name) + + async def get(self, key): + val = await self._inner.get(key) + if self.on_get is not None: + await self.on_get(key, val) + return val + + +@pytest.mark.asyncio +async def test_release_never_deletes_a_successors_lock(): + redis = fakeredis.aioredis.FakeRedis() + hooked = HookedRedis(redis) + + async def expire_and_reacquire(key, val): + hooked.on_get = None # fire once, on A's release GET + await redis.delete("lk") # A's lock TTL expires... + await redis.set("lk", b"token-B") # ...and B acquires in the window + + # HookedRedis duck-types the two Redis calls the lock makes; it is deliberately not a Redis[Any]. + lock_a = RedisLock(hooked, "lk", ttl=30) # pyright: ignore[reportArgumentType] + async with lock_a as acquired: + assert acquired + hooked.on_get = expire_and_reacquire # arm just before release + + assert ( + await redis.get("lk") == b"token-B" + ), "A's release deleted B's lock: GET-then-DELETE must be an atomic compare-and-delete" diff --git a/tests/test_bug_nested_wrapper_priority.py b/tests/test_bug_nested_wrapper_priority.py new file mode 100644 index 0000000..51357ce --- /dev/null +++ b/tests/test_bug_nested_wrapper_priority.py @@ -0,0 +1,65 @@ +"""Regression: _collect_priorities skips a nested Wrapper's own +dedup_priority. + +The walk's generic BaseModel branch does `getattr(child, "data", child)` to unwrap +MergerPercentageItem-style containers -- but a Wrapper is the one BaseNode that +ALSO has `.data`, so a Wrapper sitting directly in a mixer's `items` list (or as +`positional`/`default`) gets unwrapped into its child and its own dedup_priority +override never propagates. + +Documented contract (README "Dedup and dedup_priority"): "a node with a non-zero +dedup_priority overrides its whole subtree" -- for ANY node, wrappers included. + +Both sources emit the same id; the copy from the priority-5 subtree must win the +dedup arbitration even though the priority-0 copy is first-seen. +""" + +import pytest + +from smartfeed.models.base import FeedResult +from smartfeed.models.mixers import MergerAppend, MergerPositional +from smartfeed.models.subfeed import SubFeed +from smartfeed.models.wrapper import Wrapper +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +def _source(who: str): + async def fetch(user_id, limit, next_page, **kw): + return FeedResult(data=[{"id": 1, "who": who}], next_page={}, has_next_page=False) + + return fetch + + +def _ctx(): + return S.make_ctx({"a": _source("a"), "b": _source("b")}) + + +@pytest.mark.asyncio +async def test_wrapper_in_items_list_priority_wins(): + boosted = Wrapper(node_id="boost", dedup_priority=5, data=SubFeed(subfeed_id="b", method_name="b")) + merger = MergerAppend(node_id="mix", items=[SubFeed(subfeed_id="a", method_name="a"), boosted]) + node = S.wrapper(merger, dedup_key="id") + + r = await run_executor.run(node, _ctx(), limit=2, cursor={}) + assert len(r.data) == 1 + assert r.data[0]["who"] == "b", ( + f"the priority-5 wrapper subtree must win dedup, got copy from '{r.data[0]['who']}' " + f"(nested wrapper's dedup_priority was skipped in the priority walk)" + ) + + +@pytest.mark.asyncio +async def test_wrapper_as_positional_child_priority_wins(): + boosted = Wrapper(node_id="boost", dedup_priority=5, data=SubFeed(subfeed_id="b", method_name="b")) + merger = MergerPositional( + node_id="pos", + positions=[2], + positional=boosted, + default=SubFeed(subfeed_id="a", method_name="a"), + ) + node = S.wrapper(merger, dedup_key="id") + + r = await run_executor.run(node, _ctx(), limit=2, cursor={}) + assert len(r.data) == 1 + assert r.data[0]["who"] == "b", f"the priority-5 wrapper subtree must win dedup, got copy from '{r.data[0]['who']}'" diff --git a/tests/test_bug_overfetch_item_loss.py b/tests/test_bug_overfetch_item_loss.py index 11529fc..2934511 100644 --- a/tests/test_bug_overfetch_item_loss.py +++ b/tests/test_bug_overfetch_item_loss.py @@ -1,4 +1,4 @@ -"""BUG #1 -- dedup silently drops the overfetched remainder. +"""Regression: dedup silently drops the overfetched remainder. Both dedup paths fetch ``target * overfetch_factor`` items, keep ``target``, and return the child cursor advanced past *all* fetched items. The surviving-but- @@ -25,8 +25,7 @@ def _redis(): @pytest.mark.asyncio async def test_cached_dedup_covers_all_unique_items(): src = S.ScriptedSource(S.unique_pool(200)) - node = S.wrapper(S.subfeed("src", "src"), session_size=20, - dedup_key="id", overfetch_factor=4) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=200) S.assert_full_coverage(pages, set(range(200))) @@ -35,7 +34,7 @@ async def test_cached_dedup_covers_all_unique_items(): @pytest.mark.asyncio async def test_passthrough_dedup_covers_all_unique_items(): src = S.ScriptedSource(S.unique_pool(200)) - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=4) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=200) S.assert_full_coverage(pages, set(range(200))) @@ -47,8 +46,7 @@ async def test_cached_dedup_covers_all_with_real_duplicates(): dupes but still surface every *unique* id. Coverage = the unique universe.""" pool = S.dup_pool(150, every=4) # unique universe is range(150), plus cross-page dupes src = S.ScriptedSource(pool) - node = S.wrapper(S.subfeed("src", "src"), session_size=20, - dedup_key="id", overfetch_factor=4) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_full_coverage(pages, set(range(150))) @@ -59,8 +57,7 @@ async def test_overfetch_1_is_lossless_control(): """Control: with overfetch_factor=1 there is no discard, so no loss. Proves the harness is sound and localizes the bug to overfetch>1.""" src = S.ScriptedSource(S.unique_pool(200)) - node = S.wrapper(S.subfeed("src", "src"), session_size=20, - dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_full_coverage(pages, set(range(200))) diff --git a/tests/test_bug_passthrough_fresh_scroll.py b/tests/test_bug_passthrough_fresh_scroll.py new file mode 100644 index 0000000..05f5d74 --- /dev/null +++ b/tests/test_bug_passthrough_fresh_scroll.py @@ -0,0 +1,34 @@ +"""Regression: passthrough dedup never resets the seen-set on a +fresh scroll. + +Documented contract (README "Dedup contract", CHANGELOG 1.0.0): "re-sending an +empty cursor (or page 1) starts a FRESH scroll and may repeat items already +shown in a previous scroll". The cached path honors this (`_reset_seen_set` in +`_cold_build`); the passthrough path only ever loads and appends to the +seen-set, so a page refresh (empty cursor) silently filters out everything the +previous scroll showed instead of serving page 1 again. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_empty_cursor_starts_fresh_scroll_on_passthrough(): + src = S.ScriptedSource(S.unique_pool(50)) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") # no cache -> passthrough + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + ids_scroll_1 = [it["id"] for it in r1.data] + + # New scroll: empty cursor again (page refresh / F5) + r2 = await run_executor.run(node, ctx, limit=10, cursor={}) + ids_scroll_2 = [it["id"] for it in r2.data] + + assert ids_scroll_2 == ids_scroll_1, ( + f"fresh scroll must restart from page 1, got {ids_scroll_2} " f"(previous scroll's seen-set was never reset)" + ) diff --git a/tests/test_bug_redis_lock_decode.py b/tests/test_bug_redis_lock_decode.py index ffec6bf..541189b 100644 --- a/tests/test_bug_redis_lock_decode.py +++ b/tests/test_bug_redis_lock_decode.py @@ -1,4 +1,4 @@ -"""BUG #4 -- RedisLock assumes bytes; crashes under decode_responses=True. +"""Regression: RedisLock assumes bytes; crashes under decode_responses=True. redis_lock.py:35 does ``val.decode()`` in __aexit__. With a Redis client configured ``decode_responses=True`` (a very common setup), the reply is already @@ -10,7 +10,7 @@ correct (a lock whose value was overwritten by another owner is NOT deleted on release). The residual issue there is the non-atomic get-then-delete (TOCTOU), which cannot be reproduced deterministically without concurrency injection and is -documented in the review rather than asserted here. +documented separately rather than asserted here. """ import pytest @@ -34,8 +34,7 @@ async def test_lock_release_with_decode_responses(): async def test_shared_cold_build_with_decode_responses(): redis = fakeredis.aioredis.FakeRedis(decode_responses=True) src = S.ScriptedSource(S.unique_pool(100)) - node = S.wrapper(S.subfeed("src", "src"), node_id="w", - session_size=20, cache_key="shared") + node = S.wrapper(S.subfeed("src", "src"), node_id="w", session_size=20, cache_key="shared") ctx = S.make_ctx({"src": src}, redis=redis) r = await run_executor.run(node, ctx, limit=10, cursor={}) assert len(r.data) == 10 diff --git a/tests/test_bug_seen_ttl_not_touched.py b/tests/test_bug_seen_ttl_not_touched.py new file mode 100644 index 0000000..ac57d2d --- /dev/null +++ b/tests/test_bug_seen_ttl_not_touched.py @@ -0,0 +1,50 @@ +"""Regression: `_touch_ttl` refreshes cache data+meta but never +the dedup seen-set. + +Every warm page read refreshes the TTL on the data and :meta keys (inactivity +timeout), but the :seen key keeps its original state_ttl countdown. A slow +scroll therefore keeps the cache alive indefinitely while the seen-set silently +expires; the next continuation rebuild then loads an empty seen-set and +re-serves already-shown ids, breaking the documented "each item at most once +per scroll" contract. + +Correct behavior: a warm read that refreshes data/meta must refresh :seen too +(it is the same inactivity signal, and dedup state must not outlive-lag the +cache it guards). +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_warm_read_refreshes_seen_set_ttl_alongside_cache(): + redis = fakeredis.aioredis.FakeRedis() + src = S.ScriptedSource(S.unique_pool(30)) + node = S.wrapper(S.subfeed("src", "src"), session_size=10, dedup_key="id") + ctx = S.make_ctx({"src": src}, redis=redis) + + r1 = await run_executor.run(node, ctx, limit=5, cursor={}) # cold build + + # Age every key close to expiry, then serve a warm page. + keys = sorted(k.decode() for k in await redis.keys("*")) + seen_keys = [k for k in keys if k.endswith(":seen")] + assert seen_keys, "dedup wrapper must have written a :seen key" + for k in keys: + await redis.expire(k, 7) + + await run_executor.run(node, ctx, limit=5, cursor=r1.next_page) # warm hit + + # Control: data/meta are refreshed today (documented inactivity-timeout behavior). + base_key = [k for k in keys if not k.endswith((":seen", ":meta"))][0] + assert await redis.ttl(base_key) > 7, "warm read should refresh the cache TTL" + + # The bug: :seen stays on its old countdown while the cache it guards lives on. + seen_ttl = await redis.ttl(seen_keys[0]) + assert seen_ttl > 7, ( + f":seen TTL was not refreshed on the warm read (still {seen_ttl}s) -- " + f"the seen-set can expire mid-scroll while the cache stays alive, causing duplicates" + ) diff --git a/tests/test_bug_shared_cache_continuation.py b/tests/test_bug_shared_cache_continuation.py index b863d2c..ce05de0 100644 --- a/tests/test_bug_shared_cache_continuation.py +++ b/tests/test_bug_shared_cache_continuation.py @@ -1,4 +1,4 @@ -"""BUG #2 -- shared cache (cache_key) cannot paginate past session_size. +"""Regression: shared cache (cache_key) cannot paginate past session_size. When a cache_key wrapper exhausts its per-wrapper cache, the warm path calls _cold_build with the continuation child_cursor, but _build_shared_base @@ -22,8 +22,7 @@ def _redis(): def _shared_wrapper(): - return S.wrapper(S.subfeed("src", "src"), node_id="w", - session_size=20, cache_key="shared") + return S.wrapper(S.subfeed("src", "src"), node_id="w", session_size=20, cache_key="shared") @pytest.mark.asyncio diff --git a/tests/test_bug_shared_cache_dedup.py b/tests/test_bug_shared_cache_dedup.py new file mode 100644 index 0000000..c0cec12 --- /dev/null +++ b/tests/test_bug_shared_cache_dedup.py @@ -0,0 +1,39 @@ +"""Regression: shared cache_key path drops cross-round dedup. + +`_cold_build` (shared branch) calls `_fetch_and_dedup` with seen=None, and +`_dedup` then allocates a FRESH seen set on every refill round (wrapper.py:121). +Round 2's refill re-admits ids already collected in round 1, so duplicates +survive WITHIN one shared segment and are served to the user -- contradicting +README's "the shared base dedups within each segment". + +The identical config without cache_key dedups correctly (control test). +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S + + +POOL = S.dup_pool(20, every=4) # earlier ids re-emitted -> duplicates across refill rounds + + +@pytest.mark.asyncio +async def test_shared_cache_scroll_has_no_duplicates(): + src = S.ScriptedSource(POOL) + node = S.wrapper(S.subfeed("src", "src"), session_size=10, dedup_key="id", cache_key="shared") + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + pages = await S.drain(node, ctx, limit=5) + S.assert_no_duplicates(pages) + + +@pytest.mark.asyncio +async def test_standard_cache_scroll_has_no_duplicates_control(): + """Control: same config minus cache_key is clean today.""" + src = S.ScriptedSource(POOL) + node = S.wrapper(S.subfeed("src", "src"), session_size=10, dedup_key="id") + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + pages = await S.drain(node, ctx, limit=5) + S.assert_no_duplicates(pages) diff --git a/tests/test_bug_softfail_cursor_leak.py b/tests/test_bug_softfail_cursor_leak.py index 32a8f2b..c80590d 100644 --- a/tests/test_bug_softfail_cursor_leak.py +++ b/tests/test_bug_softfail_cursor_leak.py @@ -1,4 +1,4 @@ -"""BUG #6 -- a soft-failed subfeed returns the whole inbound cursor. +"""Regression: a soft-failed subfeed returns the whole inbound cursor. On raise_error=False, SubFeed.execute returns ``next_page=cursor`` -- the entire feed cursor -- instead of ``{subfeed_id: ...}`` (subfeed.py:39). _merge_cursor @@ -33,10 +33,13 @@ async def failing_source(user_id, limit, next_page, **kw): def _node(): # child order [good, bad] so the failed child's cursor merges LAST. - return MergerAppend(node_id="top", items=[ - SubFeed(subfeed_id="good_src", method_name="good"), - SubFeed(subfeed_id="bad_src", method_name="bad", raise_error=False), - ]) + return MergerAppend( + node_id="top", + items=[ + SubFeed(subfeed_id="good_src", method_name="good"), + SubFeed(subfeed_id="bad_src", method_name="bad", raise_error=False), + ], + ) @pytest.mark.asyncio @@ -47,9 +50,9 @@ async def test_softfail_does_not_clobber_sibling_cursor(): r = await run_executor.run(_node(), ctx, limit=10, cursor=inbound) # good_src is asked for demand 5 (MergerAppend split of limit 10) from pos 5, # so its fresh cursor is pos 10. The failed sibling must not overwrite it. - assert r.next_page.get("good_src", {}).get("pos") == 10, ( - f"good_src cursor was clobbered by the failed sibling: {r.next_page}" - ) + assert ( + r.next_page.get("good_src", {}).get("pos") == 10 + ), f"good_src cursor was clobbered by the failed sibling: {r.next_page}" @pytest.mark.asyncio @@ -58,9 +61,7 @@ async def test_softfail_cursor_is_scoped_to_own_id(): inbound = {"good_src": {"pos": 5}, "unrelated": {"x": 1}} r = await run_executor.run(_node(), ctx, limit=10, cursor=inbound) # The failed subfeed must not resurrect unrelated inbound cursor keys. - assert "unrelated" not in r.next_page, ( - f"failed subfeed leaked unrelated inbound cursor keys: {r.next_page}" - ) + assert "unrelated" not in r.next_page, f"failed subfeed leaked unrelated inbound cursor keys: {r.next_page}" @pytest.mark.asyncio diff --git a/tests/test_bug_stale_waiter.py b/tests/test_bug_stale_waiter.py new file mode 100644 index 0000000..4a102f2 --- /dev/null +++ b/tests/test_bug_stale_waiter.py @@ -0,0 +1,46 @@ +"""Regression: cold-rebuild lock waiter serves the STALE cache batch. + +When a rebuild races (double-fired page request at a cache-exhaustion boundary), +the lock loser polls `_read_meta` and accepts the FIRST meta it sees. The old +data/meta keys are never deleted before a rebuild, so the loser instantly serves +the OLD batch from offset 0 under the OLD gen: the user is re-shown page-1 items +mid-scroll, and the loser's next cursor is stale (the builder writes a new gen), +forcing a full feed restart on the following page. + +Correct behavior: no concurrent response may repeat ids already shown earlier in +this scroll. (Both responses serving the same NEW page is fine -- same cursor, +idempotent retry.) + +The child has latency 0.3s, so the winner is still building when the loser's +0.1s poll finds the stale meta -- the race is deterministic. +""" + +import asyncio + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_concurrent_continuation_rebuild_does_not_reserve_shown_items(): + src = S.ScriptedSource(S.unique_pool(100), latency=0.3) + node = S.wrapper(S.subfeed("src", "src"), session_size=20) + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + shown = {it["id"] for it in r1.data} | {it["id"] for it in r2.data} + + # cursor offset == len(batch) -> continuation rebuild; double-fire it + a, b = await asyncio.gather( + run_executor.run(node, ctx, limit=10, cursor=r2.next_page), + run_executor.run(node, ctx, limit=10, cursor=r2.next_page), + ) + ids_a = [it["id"] for it in a.data] + ids_b = [it["id"] for it in b.data] + + assert not (set(ids_a) & shown), f"response A re-served already-shown ids: {sorted(set(ids_a) & shown)}" + assert not (set(ids_b) & shown), f"response B re-served already-shown ids: {sorted(set(ids_b) & shown)}" diff --git a/tests/test_bug_torn_warm_read.py b/tests/test_bug_torn_warm_read.py new file mode 100644 index 0000000..4a2fa3f --- /dev/null +++ b/tests/test_bug_torn_warm_read.py @@ -0,0 +1,86 @@ +"""Regression: torn warm read -- meta and data are two sequential +GETs, so a rebuild landing between them pairs OLD-gen meta with the NEW batch. + +`_execute_with_cache` reads `:meta` (gen matches the cursor -> proceed), then +reads the data key. `_write_cache` replaces both keys atomically in one +pipeline, so a rebuild that lands between the two GETs makes the warm reader +serve the REPLACEMENT batch at the OLD scroll's offset under the OLD gen: +items from a batch the client's gen never belonged to, while page 2 of the +original batch is silently skipped. + +Correct behavior: the served page must be consistent with ONE snapshot -- +either the old batch at the old offset, or a clean rebuild (new gen from +offset 0). This test forces the interleaving deterministically by triggering a +fresh-scroll rebuild from a hook between the two GETs. +""" + +from typing import Awaitable, Callable, Optional + +import pytest +import fakeredis.aioredis + +from smartfeed.models.base import FeedResult +from tests import sources as S +from smartfeed.execution import executor as run_executor + +GetHook = Callable[[str, Optional[bytes]], Awaitable[None]] + + +class HookedRedis: + """Proxy that awaits a hook after each GET returns (to interleave a racing + writer between the wrapper's meta GET and data GET).""" + + def __init__(self, inner): + self._inner = inner + self.on_get: Optional[GetHook] = None + + def __getattr__(self, name): + return getattr(self._inner, name) + + async def get(self, key): + val = await self._inner.get(key) + if self.on_get is not None: + await self.on_get(key, val) + return val + + +class ShiftingSource: + """Returns a disjoint id range on every call, so each rebuild produces a + recognizably different batch.""" + + def __init__(self): + self.calls = 0 + + async def __call__(self, user_id, limit, next_page, **kw) -> FeedResult: + base = self.calls * 100 + self.calls += 1 + data = [{"id": base + i, "val": f"v{base + i}"} for i in range(limit)] + return FeedResult(data=data, next_page={}, has_next_page=False) + + +@pytest.mark.asyncio +async def test_warm_read_is_snapshot_consistent_under_racing_rebuild(): + hooked = HookedRedis(fakeredis.aioredis.FakeRedis()) + node = S.wrapper(S.subfeed("src", "src"), session_size=10) + ctx = S.make_ctx({"src": ShiftingSource()}, redis=hooked) + + r1 = await run_executor.run(node, ctx, limit=5, cursor={}) # batch [0..9], gen g1 + gen_1 = r1.next_page["w"]["gen"] + assert [it["id"] for it in r1.data] == [0, 1, 2, 3, 4] + + async def rebuild_between_reads(key, val): + if not key.endswith(":meta") or val is None: + return + hooked.on_get = None # fire once + # A fresh scroll (empty cursor) rebuilds the batch -> [100..109], new gen. + await run_executor.run(node, ctx, limit=5, cursor={}) + + hooked.on_get = rebuild_between_reads + r2 = await run_executor.run(node, ctx, limit=5, cursor=r1.next_page) + + ids_2 = [it["id"] for it in r2.data] + gen_2 = r2.next_page["w"]["gen"] + assert ids_2 == [5, 6, 7, 8, 9] or gen_2 != gen_1, ( + f"torn read: served {ids_2} from the replacement batch under the old gen -- " + f"meta and data must be read as one snapshot (single MGET/pipeline)" + ) diff --git a/tests/test_bug_variable_page_size.py b/tests/test_bug_variable_page_size.py index 40a1c74..a1652f2 100644 --- a/tests/test_bug_variable_page_size.py +++ b/tests/test_bug_variable_page_size.py @@ -1,4 +1,4 @@ -"""BUG #3 -- changing `limit` between pages on a cached wrapper skips/repeats. +"""Regression: changing `limit` between pages on a cached wrapper skips/repeats. The cursor stores a page *number*; the offset is recomputed as (page-1)*limit at read time (_paginate, wrapper.py:343). If the client changes page size between @@ -30,7 +30,7 @@ async def test_growing_limit_does_not_skip(): src = S.ScriptedSource(S.unique_pool(200)) node = _cached() ctx = S.make_ctx({"src": src}, redis=_redis()) - r1 = await run_executor.run(node, ctx, limit=10, cursor={}) # expect ids 0..9 + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) # expect ids 0..9 r2 = await run_executor.run(node, ctx, limit=50, cursor=r1.next_page) # expect 10..59 served = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] assert served == list(range(60)), f"expected contiguous 0..59, got gaps: {served[:15]}..." @@ -41,12 +41,12 @@ async def test_shrinking_limit_does_not_repeat(): src = S.ScriptedSource(S.unique_pool(200)) node = _cached() ctx = S.make_ctx({"src": src}, redis=_redis()) - r1 = await run_executor.run(node, ctx, limit=50, cursor={}) # ids 0..49 + r1 = await run_executor.run(node, ctx, limit=50, cursor={}) # ids 0..49 r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) # expect 50..59 served = [it["id"] for it in r1.data] + [it["id"] for it in r2.data] - assert len(served) == len(set(served)), ( - f"repeated ids on shrink: {sorted({x for x in served if served.count(x) > 1})}" - ) + assert len(served) == len( + set(served) + ), f"repeated ids on shrink: {sorted({x for x in served if served.count(x) > 1})}" @pytest.mark.asyncio diff --git a/tests/test_config_parsing.py b/tests/test_config_parsing.py index f9db222..8cdf5de 100644 --- a/tests/test_config_parsing.py +++ b/tests/test_config_parsing.py @@ -1,6 +1,5 @@ -import pytest from smartfeed.models.base import BaseNode -from smartfeed.models import FeedConfig +from smartfeed.models import FeedConfig, SubFeed, Wrapper class DummyNode(BaseNode): @@ -23,38 +22,52 @@ def test_config_hash_changes_with_params(): def test_parse_simple_subfeed(): - cfg = FeedConfig.model_validate({ - "version": "2", - "feed": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, - }) + cfg = FeedConfig.model_validate( + { + "version": "2", + "feed": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, + } + ) + assert isinstance(cfg.feed, SubFeed) assert cfg.feed.subfeed_id == "x" def test_parse_wrapper_with_cache(): - cfg = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "w", - "cache": {"session_size": 100, "session_ttl": 300}, - "data": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, - }, - }) + cfg = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "w", + "cache": {"session_size": 100, "session_ttl": 300}, + "data": {"type": "subfeed", "subfeed_id": "x", "method_name": "items"}, + }, + } + ) + assert isinstance(cfg.feed, Wrapper) + assert cfg.feed.cache is not None assert cfg.feed.cache.session_size == 100 def test_parse_nested_config(): - cfg = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "outer", - "dedup": {"dedup_key": "id"}, - "data": { - "type": "merger_percentage", "node_id": "mix", - "items": [ - {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "items"}}, - {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "items"}}, - ], + cfg = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "outer", + "dedup": {"dedup_key": "id"}, + "data": { + "type": "merger_percentage", + "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "items"}}, + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "items"}}, + ], + }, }, - }, - }) + } + ) + assert isinstance(cfg.feed, Wrapper) + assert cfg.feed.dedup is not None assert cfg.feed.dedup.dedup_key == "id" diff --git a/tests/test_config_validation.py b/tests/test_config_validation.py new file mode 100644 index 0000000..3f7e8dc --- /dev/null +++ b/tests/test_config_validation.py @@ -0,0 +1,138 @@ +"""Config-construction validation: duplicate +subfeed_id/node_id anywhere in the tree, and non-positive gradient step / +size_to_step, are rejected at parse time instead of corrupting state or +crashing at runtime. +""" + +import pytest +from pydantic import ValidationError + +from smartfeed.models import ( + FeedConfig, + MergerAppendDistribute, + MergerPercentage, + MergerPercentageGradient, + MergerPositional, +) + + +def _cfg(feed: dict) -> dict: + return {"version": "2", "feed": feed} + + +def _sub(sid: str) -> dict: + return {"type": "subfeed", "subfeed_id": sid, "method_name": "m"} + + +def test_duplicate_subfeed_ids_rejected(): + feed = {"type": "merger_append", "node_id": "mix", "items": [_sub("a"), _sub("a")]} + with pytest.raises(ValidationError, match="duplicate id 'a'"): + FeedConfig.model_validate(_cfg(feed)) + + +def test_subfeed_id_colliding_with_node_id_rejected(): + """subfeed_id and node_id share one cursor namespace -- a collision across + kinds is just as corrupting as within one kind.""" + feed = {"type": "wrapper", "node_id": "x", "data": _sub("x")} + with pytest.raises(ValidationError, match="duplicate id 'x'"): + FeedConfig.model_validate(_cfg(feed)) + + +def test_duplicate_across_nested_branches_rejected(): + feed = { + "type": "merger_percentage", + "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "wrapper", "node_id": "w1", "data": _sub("a")}}, + {"percentage": 50, "data": {"type": "wrapper", "node_id": "w2", "data": _sub("a")}}, + ], + } + with pytest.raises(ValidationError, match="duplicate id 'a'"): + FeedConfig.model_validate(_cfg(feed)) + + +def test_duplicate_wrapper_node_ids_rejected(): + feed = { + "type": "merger_append", + "node_id": "mix", + "items": [ + {"type": "wrapper", "node_id": "w", "data": _sub("a")}, + {"type": "wrapper", "node_id": "w", "data": _sub("b")}, + ], + } + with pytest.raises(ValidationError, match="duplicate id 'w'"): + FeedConfig.model_validate(_cfg(feed)) + + +def test_unique_ids_parse_ok(): + feed = { + "type": "merger_percentage", + "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "wrapper", "node_id": "w1", "data": _sub("a")}}, + {"percentage": 50, "data": {"type": "wrapper", "node_id": "w2", "data": _sub("b")}}, + ], + } + cfg = FeedConfig.model_validate(_cfg(feed)) + assert isinstance(cfg.feed, MergerPercentage) + assert cfg.feed.node_id == "mix" + + +def _gradient_feed(step: int, size_to_step: int) -> dict: + return { + "type": "merger_percentage_gradient", + "node_id": "g", + "item_from": {"percentage": 80, "data": _sub("a")}, + "item_to": {"percentage": 20, "data": _sub("b")}, + "step": step, + "size_to_step": size_to_step, + } + + +def test_gradient_zero_size_to_step_rejected_at_parse(): + with pytest.raises(ValidationError, match="size_to_step"): + FeedConfig.model_validate(_cfg(_gradient_feed(step=10, size_to_step=0))) + + +def test_gradient_non_positive_step_rejected_at_parse(): + for bad_step in (0, -5): + with pytest.raises(ValidationError, match="step"): + FeedConfig.model_validate(_cfg(_gradient_feed(step=bad_step, size_to_step=10))) + + +def test_gradient_positive_params_parse_ok(): + cfg = FeedConfig.model_validate(_cfg(_gradient_feed(step=10, size_to_step=30))) + assert isinstance(cfg.feed, MergerPercentageGradient) + assert cfg.feed.step == 10 + + +def test_positional_round_trips_via_public_config(): + feed = { + "type": "merger_positional", + "node_id": "pos", + "positions": [1, 3], + "positional": _sub("p"), + "default": _sub("d"), + } + cfg = FeedConfig.model_validate(_cfg(feed)) + assert isinstance(cfg.feed, MergerPositional) + assert cfg.feed.positions == [1, 3] + + +def test_distribute_round_trips_via_public_config(): + feed = { + "type": "merger_distribute", + "node_id": "dist", + "distribution_key": "op", + "sorting_key": "score", + "sorting_desc": True, + "items": [_sub("x"), _sub("y")], + } + cfg = FeedConfig.model_validate(_cfg(feed)) + assert isinstance(cfg.feed, MergerAppendDistribute) + assert cfg.feed.distribution_key == "op" + + +def test_missing_required_field_raises(): + with pytest.raises(ValidationError): + FeedConfig.model_validate(_cfg({"type": "subfeed", "subfeed_id": "a"})) # method_name missing diff --git a/tests/test_continuation.py b/tests/test_continuation.py index c354834..ba626fe 100644 --- a/tests/test_continuation.py +++ b/tests/test_continuation.py @@ -25,6 +25,7 @@ # Infinite stateful source: returns items based on page cursor # --------------------------------------------------------------------------- + async def infinite_source(user_id, limit, next_page, **kw): """Returns `limit` items starting from (page-1)*limit. IDs are globally unique.""" page = next_page.get("page", 1) @@ -62,6 +63,7 @@ def _make_wrapper(session_size=20): # Helper: collect three pages # --------------------------------------------------------------------------- + async def _three_pages(ctx, session_size=20, limit=10): node = _make_wrapper(session_size=session_size) r1 = await run_executor.run(node, ctx, limit=limit, cursor={}) @@ -74,6 +76,7 @@ async def _three_pages(ctx, session_size=20, limit=10): # Tests # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_pages_1_and_2_served_from_cache(ctx): """Pages 1 and 2 are served from the same generation (warm hits).""" @@ -111,9 +114,7 @@ async def test_page3_items_continue_from_page2_no_gap(ctx): first_p3 = min(ids_p3) # No gap: the first item of page 3 immediately follows the last item of page 2 - assert first_p3 == last_p2 + 1, ( - f"Expected page 3 to start at {last_p2 + 1}, got {first_p3}" - ) + assert first_p3 == last_p2 + 1, f"Expected page 3 to start at {last_p2 + 1}, got {first_p3}" @pytest.mark.asyncio @@ -140,11 +141,7 @@ async def test_all_ids_across_three_pages_are_unique(ctx): """All item IDs across pages 1, 2 and 3 are distinct.""" r1, r2, r3 = await _three_pages(ctx) - all_ids = ( - [item["id"] for item in r1.data] - + [item["id"] for item in r2.data] - + [item["id"] for item in r3.data] - ) + all_ids = [item["id"] for item in r1.data] + [item["id"] for item in r2.data] + [item["id"] for item in r3.data] assert len(all_ids) == len(set(all_ids)), "Duplicate IDs found across pages" @@ -153,11 +150,7 @@ async def test_all_ids_are_globally_sequential(ctx): """Items across pages 1-3 form a gapless sequence starting from 0.""" r1, r2, r3 = await _three_pages(ctx) - all_ids = sorted( - item["id"] - for page in (r1.data, r2.data, r3.data) - for item in page - ) + all_ids = sorted(item["id"] for page in (r1.data, r2.data, r3.data) for item in page) expected = list(range(30)) assert all_ids == expected, f"Expected {expected}, got {all_ids}" diff --git a/tests/test_cursors.py b/tests/test_cursors.py index 06b23ec..031ec76 100644 --- a/tests/test_cursors.py +++ b/tests/test_cursors.py @@ -19,16 +19,18 @@ def ctx(redis): @pytest.mark.asyncio async def test_cached_wrapper_hides_child_cursors(ctx): - config = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "cached", - "cache": {"session_size": 50, "session_ttl": 300}, - "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, - }, - }) - - executor = None + config = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "cached", + "cache": {"session_size": 50, "session_ttl": 300}, + "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, + }, + } + ) + result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) # Only wrapper cursor, NOT subfeed cursor @@ -39,15 +41,17 @@ async def test_cached_wrapper_hides_child_cursors(ctx): @pytest.mark.asyncio async def test_uncached_wrapper_exposes_child_cursors(ctx): - config = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "passthrough", - "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, - }, - }) - - executor = None + config = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "passthrough", + "data": {"type": "subfeed", "subfeed_id": "items", "method_name": "items"}, + }, + } + ) + result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) # Child cursor is visible (no cache to hide it) diff --git a/tests/test_dedup_correctness.py b/tests/test_dedup_correctness.py index 02664a4..a72a141 100644 --- a/tests/test_dedup_correctness.py +++ b/tests/test_dedup_correctness.py @@ -32,9 +32,10 @@ def _redis(): # Overlapping sources: every unique id exactly once # --------------------------------------------------------------------------- + def _overlap_ctx(): - a = S.ScriptedSource(S.unique_pool(100, start=0)) # ids 0..99 - b = S.ScriptedSource(S.unique_pool(100, start=50)) # ids 50..149 (overlap 50..99) + a = S.ScriptedSource(S.unique_pool(100, start=0)) # ids 0..99 + b = S.ScriptedSource(S.unique_pool(100, start=50)) # ids 50..149 (overlap 50..99) return a, b, S.make_ctx({"a": a, "b": b}, redis=_redis()) @@ -42,11 +43,14 @@ def _overlap_ctx(): async def test_overlapping_sources_cached_each_id_once(): """session_size > universe -> one batch -> within-batch dedup handles overlap.""" a, b, ctx = _overlap_ctx() - merger = MergerAppend(node_id="m", items=[ - SubFeed(subfeed_id="a", method_name="a"), - SubFeed(subfeed_id="b", method_name="b"), - ]) - node = S.wrapper(merger, session_size=500, dedup_key="id", overfetch_factor=4) + merger = MergerAppend( + node_id="m", + items=[ + SubFeed(subfeed_id="a", method_name="a"), + SubFeed(subfeed_id="b", method_name="b"), + ], + ) + node = S.wrapper(merger, session_size=500, dedup_key="id") pages = await S.drain(node, ctx, limit=10, max_pages=100) S.assert_no_duplicates(pages) S.assert_full_coverage(pages, set(range(150))) @@ -56,11 +60,14 @@ async def test_overlapping_sources_cached_each_id_once(): async def test_overlapping_sources_passthrough_each_id_once(): """Passthrough cross-page dedup via the Redis seen-set (overfetch=1 -> no loss).""" a, b, ctx = _overlap_ctx() - merger = MergerAppend(node_id="m", items=[ - SubFeed(subfeed_id="a", method_name="a"), - SubFeed(subfeed_id="b", method_name="b"), - ]) - node = S.wrapper(merger, dedup_key="id", overfetch_factor=1) + merger = MergerAppend( + node_id="m", + items=[ + SubFeed(subfeed_id="a", method_name="a"), + SubFeed(subfeed_id="b", method_name="b"), + ], + ) + node = S.wrapper(merger, dedup_key="id") pages = await S.drain(node, ctx, limit=10, max_pages=200) S.assert_no_duplicates(pages) S.assert_full_coverage(pages, set(range(150))) @@ -70,10 +77,11 @@ async def test_overlapping_sources_passthrough_each_id_once(): # Within-batch high-density duplicates collapse # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_adjacent_duplicates_collapse_within_batch(): src = S.ScriptedSource(S.adjacent_dup_pool(80)) # each id emitted twice, back-to-back - node = S.wrapper(S.subfeed("src", "src"), session_size=500, dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), session_size=500, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=100) S.assert_no_duplicates(pages) @@ -84,13 +92,14 @@ async def test_adjacent_duplicates_collapse_within_batch(): # Cross-page duplicate suppression (passthrough seen-set) # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_cross_page_duplicate_suppressed_passthrough(): # id 3 appears at pos 3 and again at pos 25 (a later page). pool = S.unique_pool(30) pool.insert(25, {"id": 3, "val": "v3-again"}) src = S.ScriptedSource(pool) - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=5, max_pages=50) ids = S.flat_ids(pages) @@ -102,13 +111,14 @@ async def test_cross_page_duplicate_suppressed_passthrough(): # rebuild boundary (session-scoped seen-set carried across cold rebuilds). # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_cached_dedup_suppresses_across_rebuild_boundary(): # id 3 at pos 3 (batch 1) and pos 20 (batch 2, after session_size=20 boundary). pool = S.unique_pool(20) pool.append({"id": 3, "val": "v3-again"}) # pos 20 src = S.ScriptedSource(pool) - node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=50) assert S.flat_ids(pages).count(3) == 1 @@ -118,11 +128,15 @@ async def test_cached_dedup_suppresses_across_rebuild_boundary(): # Priority arbitration # --------------------------------------------------------------------------- + def _prio_wrapper(): - merger = MergerAppend(node_id="m", items=[ - SubFeed(subfeed_id="lo", method_name="x", dedup_priority=1), - SubFeed(subfeed_id="hi", method_name="x", dedup_priority=5), - ]) + merger = MergerAppend( + node_id="m", + items=[ + SubFeed(subfeed_id="lo", method_name="x", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="x", dedup_priority=5), + ], + ) return S.wrapper(merger, dedup_key="id") @@ -134,7 +148,7 @@ def _winner_source(item): return (item.get("_smartfeed_debug_info") or {}).get("source") -# Regression guard for BUG#7: in-batch priority arbitration must run even when the +# Regression guard: in-batch priority arbitration must run even when the # higher-priority copy is NOT the first-seen one. The seen-check skips only true # cross-page dupes (`key in seen and key not in batch`), so a within-batch duplicate # still reaches priority arbitration. @@ -153,11 +167,14 @@ async def test_priority_higher_source_wins_integration_merger_order(): first-seen; the high-priority source (listed second) must still win the key.""" lo = S.ScriptedSource([{"id": 0, "val": "lo"}]) hi = S.ScriptedSource([{"id": 0, "val": "hi"}]) - merger = MergerAppend(node_id="m", items=[ - SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), - SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), - ]) - node = S.wrapper(merger, session_size=100, dedup_key="id", overfetch_factor=1) + merger = MergerAppend( + node_id="m", + items=[ + SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), + ], + ) + node = S.wrapper(merger, session_size=100, dedup_key="id") ctx = S.make_ctx({"lo": lo, "hi": hi}, redis=_redis()) r = await run_executor.run(node, ctx, limit=10, cursor={}) kept = [it for it in r.data if it["id"] == 0] @@ -172,8 +189,8 @@ async def test_priority_preserved_when_high_arrives_in_refill(): stale copy must be removed (assert the list directly so a fix that keeps the dupe cannot pass).""" w = _prio_wrapper() - data = [_stamped(1, "lo"), _stamped(7, "lo")] # first fetch - refill = [_stamped(7, "hi"), _stamped(2, "hi")] # refill batch + data = [_stamped(1, "lo"), _stamped(7, "lo")] # first fetch + refill = [_stamped(7, "hi"), _stamped(2, "hi")] # refill batch out = w._dedup(data + refill) assert [it["id"] for it in out] == [1, 7, 2], f"expected one copy each in order: {out}" assert _winner_source(out[1]) == "hi", f"high-priority refill copy lost: {out}" @@ -194,11 +211,14 @@ async def test_priority_cross_page_shown_item_wins_documented(): # page 1 shows id 5 from the low-priority source; page 2 has id 5 from high-prio. lo = S.ScriptedSource([{"id": 5, "val": "lo"}, {"id": 6, "val": "lo"}]) hi = S.ScriptedSource([{"id": 5, "val": "hi"}, {"id": 7, "val": "hi"}]) - merger = MergerAppend(node_id="m", items=[ - SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), - SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), - ]) - node = S.wrapper(merger, dedup_key="id", overfetch_factor=1) + merger = MergerAppend( + node_id="m", + items=[ + SubFeed(subfeed_id="lo", method_name="lo", dedup_priority=1), + SubFeed(subfeed_id="hi", method_name="hi", dedup_priority=5), + ], + ) + node = S.wrapper(merger, dedup_key="id") ctx = S.make_ctx({"lo": lo, "hi": hi}, redis=_redis()) r1 = await run_executor.run(node, ctx, limit=1, cursor={}) # Whatever is shown first for id 5 is final; it is never served twice. @@ -208,47 +228,51 @@ async def test_priority_cross_page_shown_item_wins_documented(): # --------------------------------------------------------------------------- -# Dedup coverage holds through every merger type (regression guard for BUG#1), -# for both overfetch=1 and overfetch=4. +# Dedup coverage holds through every merger type (regression guard). # --------------------------------------------------------------------------- -@pytest.mark.parametrize("of", [1, 4]) + @pytest.mark.asyncio -async def test_dedup_over_distribute_covers_all(of): +async def test_dedup_over_distribute_covers_all(): pool = [{"id": k, "grp": k % 5} for k in range(120)] src = S.ScriptedSource(pool) - merger = MergerAppendDistribute(node_id="m", items=[SubFeed(subfeed_id="a", method_name="a")], - distribution_key="grp") - node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) + merger = MergerAppendDistribute( + node_id="m", items=[SubFeed(subfeed_id="a", method_name="a")], distribution_key="grp" + ) + node = S.wrapper(merger, session_size=20, dedup_key="id") ctx = S.make_ctx({"a": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_full_coverage(pages, set(range(120))) S.assert_no_duplicates(pages) -@pytest.mark.parametrize("of", [1, 4]) @pytest.mark.asyncio -async def test_dedup_over_percentage_covers_all(of): +async def test_dedup_over_percentage_covers_all(): src = S.ScriptedSource(S.unique_pool(120)) - merger = MergerPercentage(node_id="m", items=[ - MergerPercentageItem(percentage=100, data=SubFeed(subfeed_id="a", method_name="a")), - ]) - node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) + merger = MergerPercentage( + node_id="m", + items=[ + MergerPercentageItem(percentage=100, data=SubFeed(subfeed_id="a", method_name="a")), + ], + ) + node = S.wrapper(merger, session_size=20, dedup_key="id") ctx = S.make_ctx({"a": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_full_coverage(pages, set(range(120))) S.assert_no_duplicates(pages) -@pytest.mark.parametrize("of", [1, 4]) @pytest.mark.asyncio -async def test_dedup_over_positional_covers_all(of): - a = S.ScriptedSource(S.unique_pool(120, start=0)) # ids 0..119 - b = S.ScriptedSource(S.unique_pool(60, start=1000)) # ids 1000..1059 - merger = MergerPositional(node_id="m", positions=[1, 3, 5], - positional=SubFeed(subfeed_id="a", method_name="a"), - default=SubFeed(subfeed_id="b", method_name="b")) - node = S.wrapper(merger, session_size=20, dedup_key="id", overfetch_factor=of) +async def test_dedup_over_positional_covers_all(): + a = S.ScriptedSource(S.unique_pool(120, start=0)) # ids 0..119 + b = S.ScriptedSource(S.unique_pool(60, start=1000)) # ids 1000..1059 + merger = MergerPositional( + node_id="m", + positions=[1, 3, 5], + positional=SubFeed(subfeed_id="a", method_name="a"), + default=SubFeed(subfeed_id="b", method_name="b"), + ) + node = S.wrapper(merger, session_size=20, dedup_key="id") ctx = S.make_ctx({"a": a, "b": b}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_full_coverage(pages, set(range(120)) | set(range(1000, 1060))) @@ -256,10 +280,11 @@ async def test_dedup_over_positional_covers_all(of): # --------------------------------------------------------------------------- -# Regression guard for BUG#8: dense duplicates must not produce short non-final +# Regression guard: dense duplicates must not produce short non-final # pages -- the refill loop scans through duplicate runs until the page is full. # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_dense_duplicates_still_fill_pages(): # Each unique id 1..79 is followed by three copies of a hot id (0): only ~1 in 4 @@ -269,8 +294,7 @@ async def test_dense_duplicates_still_fill_pages(): pool.append({"id": k}) pool += [{"id": 0}] * 3 src = S.ScriptedSource(pool) - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", - overfetch_factor=1, max_refill_loops=2) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=10, max_pages=300) S.assert_pages_full(pages, limit=10) diff --git a/tests/test_dedup_edge_cases.py b/tests/test_dedup_edge_cases.py index ad02ca8..e8ab4f3 100644 --- a/tests/test_dedup_edge_cases.py +++ b/tests/test_dedup_edge_cases.py @@ -1,7 +1,7 @@ """Dedup edge cases: missing-key policies, non-dict items, key stringification. Contract/characterization tests. Where behavior is arguably wrong (int-vs-str key -collision) the test documents CURRENT behavior with a pointer to the review, so a +collision) the test documents CURRENT behavior, so a future change is a conscious decision rather than a silent regression. """ @@ -20,41 +20,43 @@ def _redis(): # missing_key_policy # --------------------------------------------------------------------------- + async def _missing_key_source(user_id, limit, next_page, **kw): # Second item lacks the dedup key "id". return FeedResult( data=[{"id": 1, "val": "a"}, {"val": "no-id"}, {"id": 2, "val": "b"}], - next_page={"pos": 0}, has_next_page=False, + next_page={"pos": 0}, + has_next_page=False, ) @pytest.mark.asyncio async def test_missing_key_policy_error_raises(): src = _missing_key_source - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", - overfetch_factor=1, missing_key_policy="error") + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", missing_key_policy="error") ctx = S.make_ctx({"src": src}, redis=_redis()) from smartfeed.execution import executor as run_executor + with pytest.raises(KeyError, match="Dedup key"): await run_executor.run(node, ctx, limit=10, cursor={}) @pytest.mark.asyncio async def test_missing_key_policy_keep_passes_through(): - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", - overfetch_factor=1, missing_key_policy="keep") + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", missing_key_policy="keep") ctx = S.make_ctx({"src": _missing_key_source}, redis=_redis()) from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) assert len(r.data) == 3 # the no-id item is kept un-deduped @pytest.mark.asyncio async def test_missing_key_policy_drop_removes(): - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", - overfetch_factor=1, missing_key_policy="drop") + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", missing_key_policy="drop") ctx = S.make_ctx({"src": _missing_key_source}, redis=_redis()) from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) ids = [it.get("id") for it in r.data] assert ids == [1, 2] # the no-id item is dropped @@ -64,16 +66,17 @@ async def test_missing_key_policy_drop_removes(): # non-dict items # --------------------------------------------------------------------------- + async def _non_dict_source(user_id, limit, next_page, **kw): - return FeedResult(data=[{"id": 1}, "a-string", {"id": 2}, 42], - next_page={"pos": 0}, has_next_page=False) + return FeedResult(data=[{"id": 1}, "a-string", {"id": 2}, 42], next_page={"pos": 0}, has_next_page=False) @pytest.mark.asyncio async def test_non_dict_items_pass_through_dedup(): - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") ctx = S.make_ctx({"src": _non_dict_source}, redis=_redis()) from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) assert "a-string" in r.data and 42 in r.data # non-dicts survive @@ -82,9 +85,11 @@ async def test_non_dict_items_pass_through_dedup(): # key stringification: int 1 vs str "1" # --------------------------------------------------------------------------- + async def _mixed_type_key_source(user_id, limit, next_page, **kw): - return FeedResult(data=[{"id": 1, "val": "int"}, {"id": "1", "val": "str"}], - next_page={"pos": 0}, has_next_page=False) + return FeedResult( + data=[{"id": 1, "val": "int"}, {"id": "1", "val": "str"}], next_page={"pos": 0}, has_next_page=False + ) @pytest.mark.asyncio @@ -93,8 +98,9 @@ async def test_int_and_str_key_currently_collide_characterization(): are treated as the SAME key and one is dropped. See review -- decide whether mixed-type ids should be distinct. This test pins today's behavior; flip it if the product decision is that they must be distinct.""" - node = S.wrapper(S.subfeed("src", "src"), dedup_key="id", overfetch_factor=1) + node = S.wrapper(S.subfeed("src", "src"), dedup_key="id") ctx = S.make_ctx({"src": _mixed_type_key_source}, redis=_redis()) from smartfeed.execution import executor as run_executor + r = await run_executor.run(node, ctx, limit=10, cursor={}) assert len(r.data) == 1, "int 1 and str '1' currently collide via str() coercion" diff --git a/tests/test_dedup_priority.py b/tests/test_dedup_priority.py index 5cb73f4..3486496 100644 --- a/tests/test_dedup_priority.py +++ b/tests/test_dedup_priority.py @@ -33,25 +33,30 @@ def ctx(redis): @pytest.mark.asyncio async def test_wrapper_dedup_priority_override(ctx): """Inner wrapper with dedup_priority=3 overrides child subfeed priorities.""" - config = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "outer", - "dedup": {"dedup_key": "id"}, - "data": { - "type": "merger_append", "node_id": "top", - "items": [ - {"type": "subfeed", "subfeed_id": "promo", "method_name": "promo", "dedup_priority": 10}, - { - "type": "wrapper", "node_id": "inner", "dedup_priority": 3, - "data": {"type": "subfeed", "subfeed_id": "regular", "method_name": "regular"}, - }, - ], + config = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "outer", + "dedup": {"dedup_key": "id"}, + "data": { + "type": "merger_append", + "node_id": "top", + "items": [ + {"type": "subfeed", "subfeed_id": "promo", "method_name": "promo", "dedup_priority": 10}, + { + "type": "wrapper", + "node_id": "inner", + "dedup_priority": 3, + "data": {"type": "subfeed", "subfeed_id": "regular", "method_name": "regular"}, + }, + ], + }, }, - }, - }) + } + ) - executor = None result = await run_executor.run(config.feed, ctx, limit=10, cursor={}) # promo(10) > inner wrapper override(3) -> all promo diff --git a/tests/test_distribute.py b/tests/test_distribute.py index f94bfcf..7e40c15 100644 --- a/tests/test_distribute.py +++ b/tests/test_distribute.py @@ -9,10 +9,7 @@ async def make_operators(user_id, limit, next_page, **kw): """Source producing items with operator_id cycling through op_0, op_1, op_2.""" - data = [ - {"id": i, "operator_id": f"op_{i % 3}", "val": f"tour_{i}"} - for i in range(limit) - ] + data = [{"id": i, "operator_id": f"op_{i % 3}", "val": f"item_{i}"} for i in range(limit)] return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) @@ -22,16 +19,13 @@ async def make_scored(user_id, limit, next_page, **kw): Items are in ascending score order so that descending sort visibly reorders them. op_0 gets scores 1, 3, 5, ... and op_1 gets scores 2, 4, 6, ... """ - data = [ - {"id": i, "operator_id": f"op_{i % 2}", "score": i + 1, "val": f"tour_{i}"} - for i in range(limit) - ] + data = [{"id": i, "operator_id": f"op_{i % 2}", "score": i + 1, "val": f"item_{i}"} for i in range(limit)] return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) async def make_no_key(user_id, limit, next_page, **kw): """Source producing items WITHOUT operator_id (to test missing key behaviour).""" - data = [{"id": i, "val": f"tour_{i}"} for i in range(limit)] + data = [{"id": i, "val": f"item_{i}"} for i in range(limit)] return FeedResult(data=data, next_page={"page": 2}, has_next_page=True) @@ -51,6 +45,7 @@ def ctx(): # Basic round-robin # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_basic_round_robin_no_consecutive_same_operator(ctx): """Consecutive items should not share the same operator_id (round-robin effect).""" @@ -99,6 +94,7 @@ async def test_round_robin_result_length_matches_limit(ctx): # sorting_key sorts before distributing # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_sorting_key_desc_highest_score_first(ctx): """With sorting_key='score' and sorting_desc=True, higher scores appear earlier.""" @@ -115,15 +111,14 @@ async def test_sorting_key_desc_highest_score_first(ctx): # The highest score items per operator group should appear first. # Verify that score of first item per operator_id group is higher than second. from collections import defaultdict + by_op = defaultdict(list) for item in result.data: by_op[item["operator_id"]].append(item["score"]) for op, scores in by_op.items(): if len(scores) >= 2: - assert scores[0] >= scores[1], ( - f"Operator {op}: expected descending scores, got {scores}" - ) + assert scores[0] >= scores[1], f"Operator {op}: expected descending scores, got {scores}" @pytest.mark.asyncio @@ -176,6 +171,7 @@ async def test_no_sorting_key_preserves_source_order(ctx): # Missing distribution_key raises KeyError # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_missing_distribution_key_raises(ctx): """Items without the distribution_key should raise KeyError.""" @@ -192,6 +188,7 @@ async def test_missing_distribution_key_raises(ctx): # Multiple subfeeds # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_multiple_subfeeds_merged_before_distribute(ctx): """Items from multiple subfeeds are pooled before round-robin distribution.""" diff --git a/tests/test_gradient.py b/tests/test_gradient.py index 9b49d47..6957759 100644 --- a/tests/test_gradient.py +++ b/tests/test_gradient.py @@ -95,9 +95,8 @@ async def test_step_parameter_controls_shift(ctx): # Compute ratio difference between page 1 and page 2 for each node def _b_count(node, cursor): import asyncio - r = asyncio.get_event_loop().run_until_complete( - run_executor.run(node, ctx, limit=10, cursor=cursor) - ) + + r = asyncio.get_event_loop().run_until_complete(run_executor.run(node, ctx, limit=10, cursor=cursor)) sources = [item["_smartfeed_debug_info"]["source"] for item in r.data] return sources.count("source_b"), r.next_page @@ -147,9 +146,11 @@ async def test_gradient_eventually_reaches_100_pct_to(ctx): node = _make_gradient(pct_from=80, pct_to=20, step=50, size_to_step=10) cursor = {} + r = None for _ in range(5): r = await run_executor.run(node, ctx, limit=10, cursor=cursor) cursor = r.next_page + assert r is not None sources = [item["_smartfeed_debug_info"]["source"] for item in r.data] # At full shift, all items come from source_b diff --git a/tests/test_gradient_demands.py b/tests/test_gradient_demands.py new file mode 100644 index 0000000..d275920 --- /dev/null +++ b/tests/test_gradient_demands.py @@ -0,0 +1,93 @@ +"""The gradient demand computation was rewritten from an O(limit * page) replay +of the whole shift history into an O(limit / size_to_step) walk over the current +page window with a closed-form percentage (a tampered or +legitimately deep `page` must not grow per-request work). + +This file pins (a) exact output equivalence with the original algorithm, kept +verbatim below as the reference, and (b) the bounded-work property. +""" + +from typing import Dict, List + +from smartfeed.models.mixers import MergerPercentageGradient, MergerPercentageItem +from smartfeed.models.subfeed import SubFeed + + +def _reference_demands(pct_from: int, pct_to: int, step: int, size_to_step: int, page: int, limit: int) -> tuple: + """The pre-rewrite algorithm, verbatim -- the behavioral reference.""" + percentage_from = pct_from + percentage_to = pct_to + start_position = limit * (page - 1) + first_iter = True + limit_from = 0 + limit_to = 0 + segments: List[Dict] = [] + + for i in range(size_to_step, limit * page + size_to_step, size_to_step): + if not first_iter and percentage_to < 100: + percentage_from -= step + percentage_to += step + if percentage_to > 100 or percentage_from < 0: + percentage_from = 0 + percentage_to = 100 + + if i > start_position: + iter_limit = (limit * page - start_position) if i > limit * page else (i - start_position) + start_position = i + from_take = iter_limit * percentage_from // 100 + to_take = iter_limit - from_take + limit_from += from_take + limit_to += to_take + segments.append({"limit": iter_limit, "from_take": from_take, "to_take": to_take}) + + if first_iter: + first_iter = False + + return limit_from, limit_to, segments + + +def _node(pct_from: int, pct_to: int, step: int, size_to_step: int) -> MergerPercentageGradient: + return MergerPercentageGradient( + node_id="g", + item_from=MergerPercentageItem(percentage=pct_from, data=SubFeed(subfeed_id="a", method_name="a")), + item_to=MergerPercentageItem(percentage=pct_to, data=SubFeed(subfeed_id="b", method_name="b")), + step=step, + size_to_step=size_to_step, + ) + + +def test_window_rewrite_matches_reference_exactly(): + cases = 0 + for pct_from, pct_to in [(80, 20), (100, 0), (0, 100), (50, 70), (60, 20), (90, 30), (50, 50), (70, 0)]: + for step in (1, 3, 10, 25, 40, 100): + for size_to_step in (1, 2, 3, 7, 10, 30): + node = _node(pct_from, pct_to, step, size_to_step) + for limit in (1, 3, 5, 10, 13): + for page in (1, 2, 3, 5, 8): + got = node._calculate_demands(page, limit) + want = _reference_demands(pct_from, pct_to, step, size_to_step, page, limit) + assert got == want, ( + f"mismatch at from={pct_from} to={pct_to} step={step} " + f"size_to_step={size_to_step} page={page} limit={limit}: {got} != {want}" + ) + cases += 1 + assert cases == 8 * 6 * 6 * 5 * 5 + + +def test_deep_page_work_is_bounded_by_page_size(): + node = _node(80, 20, 10, 10) + limit = 20 + limit_from, limit_to, segments = node._calculate_demands(page=10**6, limit=limit) + # Work is bounded by the page window, not the scroll depth. + assert len(segments) <= limit // 10 + 2 + # By page 10**6 the gradient long since clamped to (0, 100). + assert (limit_from, limit_to) == (0, limit) + + +def test_page_window_is_fully_allocated(): + """from_take + to_take always sum to the page window, whatever the params.""" + node = _node(80, 20, 7, 9) + for page in (1, 2, 4, 100): + limit_from, limit_to, segments = node._calculate_demands(page, 10) + assert limit_from + limit_to == 10 + assert sum(s["limit"] for s in segments) == 10 diff --git a/tests/test_high_priority.py b/tests/test_high_priority.py index 694fc5b..046dc2a 100644 --- a/tests/test_high_priority.py +++ b/tests/test_high_priority.py @@ -135,7 +135,8 @@ async def make_fixed_a(user_id, limit, next_page, **kw): """Always returns ids 0..limit-1 regardless of page.""" data = [{"id": i, "val": "a"} for i in range(limit)] return FeedResult( - data=data, next_page={"page": next_page.get("page", 1) + 1}, + data=data, + next_page={"page": next_page.get("page", 1) + 1}, has_next_page=True, ) @@ -143,7 +144,8 @@ async def make_fixed_b(user_id, limit, next_page, **kw): """Always returns ids 0..limit-1 (same as a) -- full overlap.""" data = [{"id": i, "val": "b"} for i in range(limit)] return FeedResult( - data=data, next_page={"page": next_page.get("page", 1) + 1}, + data=data, + next_page={"page": next_page.get("page", 1) + 1}, has_next_page=True, ) @@ -217,9 +219,9 @@ async def test_dedup_with_cache_across_pages(self, redis): cursor = result.next_page # With cache, dedup holds across ALL pages served from the same session - assert len(all_ids) == len(set(all_ids)), ( - "Duplicates found across cached pages -- dedup should cover the entire session" - ) + assert len(all_ids) == len( + set(all_ids) + ), "Duplicates found across cached pages -- dedup should cover the entire session" # --------------------------------------------------------------------------- @@ -265,15 +267,13 @@ async def test_page3_continues_after_rebuild(self, redis): # No overlap between pages all_ids = ids_p1 + ids_p2 + ids_p3 - assert len(all_ids) == len(set(all_ids)), ( - "Overlap detected between pages after session rebuild" - ) + assert len(all_ids) == len(set(all_ids)), "Overlap detected between pages after session rebuild" # No gap: the maximum id from page 2 + 1 should be the minimum id from page 3 # (since make_sequential produces strictly sequential ids) - assert min(ids_p3) == max(ids_p2) + 1, ( - f"Gap detected: page 2 ended at {max(ids_p2)}, page 3 starts at {min(ids_p3)}" - ) + assert ( + min(ids_p3) == max(ids_p2) + 1 + ), f"Gap detected: page 2 ended at {max(ids_p2)}, page 3 starts at {min(ids_p3)}" # --------------------------------------------------------------------------- @@ -322,9 +322,7 @@ async def test_has_next_true_at_boundary(self, redis): # Page 3 -- rebuild triggered r3 = await run_executor.run(node, ctx, limit=10, cursor=r2.next_page) assert len(r3.data) == 10, "Page 3 should return data after rebuild" - assert r3.data[0]["id"] != r1.data[0]["id"], ( - "Page 3 should contain new items, not a repeat of page 1" - ) + assert r3.data[0]["id"] != r1.data[0]["id"], "Page 3 should contain new items, not a repeat of page 1" @pytest.mark.asyncio async def test_page2_has_next_page_reflects_boundary(self, redis): @@ -495,9 +493,9 @@ async def test_exhausted_promo_fills_with_default(self): # Only 2 promo items total promo_positions = [i for i, src in enumerate(sources) if src == "promo"] - assert len(promo_positions) == 2, ( - f"Expected exactly 2 promo items, got {len(promo_positions)} at indices {promo_positions}" - ) + assert ( + len(promo_positions) == 2 + ), f"Expected exactly 2 promo items, got {len(promo_positions)} at indices {promo_positions}" # Positions 5 and 7 fall back to default (promo exhausted) # These are at index 4 and 6 in the result (0-indexed) since only 2 @@ -549,6 +547,4 @@ async def counting_promo(user_id, limit, next_page, **kw): await run_executor.run(node, ctx, limit=20, cursor={}) # The promo subfeed should be called exactly once - assert call_count == 1, ( - f"Promo subfeed was called {call_count} times; expected exactly 1" - ) + assert call_count == 1, f"Promo subfeed was called {call_count} times; expected exactly 1" diff --git a/tests/test_input_validation.py b/tests/test_input_validation.py new file mode 100644 index 0000000..ac56b93 --- /dev/null +++ b/tests/test_input_validation.py @@ -0,0 +1,156 @@ +"""Input validation at the public boundary and on untrusted cursor contents. + +Policy: integrator-level type errors (limit, session_id, cursor type) AND +tampered per-node cursor contents raise loudly (ValueError/TypeError) instead +of crashing mid-slice, serving wrong pages, or silently degrading. +""" + +import decimal + +import pytest +import fakeredis.aioredis + +from smartfeed.manager import FeedManager +from smartfeed.models.base import FeedResult +from smartfeed.models.mixers import MergerPercentageGradient, MergerPercentageItem +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +CONFIG = {"version": "2", "feed": {"type": "subfeed", "subfeed_id": "src", "method_name": "src"}} + + +async def _one_item(user_id, limit, next_page, **kw): + return FeedResult(data=[{"id": 1}], next_page={}, has_next_page=False) + + +def _manager(): + return FeedManager(config=CONFIG, methods_dict={"src": _one_item}, redis_client=None) + + +# -- get_feed boundary (#14) ------------------------------------------------- + + +@pytest.mark.asyncio +async def test_limit_zero_and_negative_raise(): + mgr = _manager() + for bad in (0, -3): + with pytest.raises(ValueError, match="limit"): + await mgr.get_feed(session_id="u", limit=bad) + + +@pytest.mark.asyncio +async def test_limit_non_int_raises(): + with pytest.raises(ValueError, match="limit"): + # Deliberately ill-typed: proves get_feed rejects a non-int limit at runtime. + await _manager().get_feed(session_id="u", limit="10") # pyright: ignore[reportArgumentType] + + +@pytest.mark.asyncio +async def test_cursor_wrong_type_raises(): + with pytest.raises(TypeError, match="cursor"): + # Deliberately ill-typed: proves get_feed rejects a non-dict cursor at runtime. + await _manager().get_feed(session_id="u", limit=10, cursor="abc") # pyright: ignore[reportArgumentType] + + +@pytest.mark.asyncio +async def test_empty_session_id_raises(): + with pytest.raises(ValueError, match="session_id"): + await _manager().get_feed(session_id="", limit=10) + + +# -- cached-wrapper cursor tampering (#13) ------------------------------------ + + +def _cached_wrapper_ctx(): + src = S.ScriptedSource(S.unique_pool(30)) + node = S.wrapper(S.subfeed("src", "src"), session_size=10) + ctx = S.make_ctx({"src": src}, redis=fakeredis.aioredis.FakeRedis()) + return node, ctx + + +@pytest.mark.asyncio +async def test_tampered_offset_rejected(): + node, ctx = _cached_wrapper_ctx() + r1 = await run_executor.run(node, ctx, limit=5, cursor={}) + gen = r1.next_page["w"]["gen"] + for bad_offset in (-4, "5", 2.5, None): + with pytest.raises(ValueError, match="Wrapper 'w'.*offset"): + await run_executor.run(node, ctx, limit=5, cursor={"w": {"offset": bad_offset, "gen": gen}}) + + +@pytest.mark.asyncio +async def test_tampered_gen_type_rejected(): + node, ctx = _cached_wrapper_ctx() + await run_executor.run(node, ctx, limit=5, cursor={}) + with pytest.raises(ValueError, match="Wrapper 'w'.*gen"): + await run_executor.run(node, ctx, limit=5, cursor={"w": {"offset": 5, "gen": 123}}) + + +@pytest.mark.asyncio +async def test_node_cursor_entry_not_dict_rejected(): + node, ctx = _cached_wrapper_ctx() + with pytest.raises(ValueError, match="Wrapper 'w'.*must be a dict"): + await run_executor.run(node, ctx, limit=5, cursor={"w": "junk"}) + + +@pytest.mark.asyncio +async def test_huge_offset_is_valid_and_triggers_continuation(): + """A huge-but-valid int offset is not tampering: it means 'past the batch', + which is the documented cache-exhausted continuation path.""" + node, ctx = _cached_wrapper_ctx() + r1 = await run_executor.run(node, ctx, limit=5, cursor={}) + gen = r1.next_page["w"]["gen"] + r = await run_executor.run(node, ctx, limit=5, cursor={"w": {"offset": 10**9, "gen": gen}}) + assert len(r.data) == 5 # continuation rebuild served a full fresh page + + +@pytest.mark.asyncio +async def test_subfeed_cursor_slice_not_dict_rejected(): + node = S.subfeed("src", "src") + ctx = S.make_ctx({"src": S.ScriptedSource(S.unique_pool(5))}) + with pytest.raises(ValueError, match="SubFeed 'src'"): + await run_executor.run(node, ctx, limit=5, cursor={"src": "junk"}) + + +# -- gradient cursor tampering (#13/#16) --------------------------------------- + + +@pytest.mark.asyncio +async def test_gradient_tampered_page_rejected(): + node = MergerPercentageGradient( + node_id="grad", + item_from=MergerPercentageItem(percentage=80, data=S.subfeed("a", "a")), + item_to=MergerPercentageItem(percentage=20, data=S.subfeed("b", "b")), + step=10, + size_to_step=10, + ) + ctx = S.make_ctx({"a": S.ScriptedSource(S.unique_pool(50)), "b": S.ScriptedSource(S.unique_pool(50, start=100))}) + for bad_cursor in ({"grad": {"page": 0}}, {"grad": {"page": -1}}, {"grad": {"page": "2"}}, {"grad": "junk"}): + with pytest.raises(ValueError, match="grad"): + await run_executor.run(node, ctx, limit=10, cursor=bad_cursor) + + +# -- cache serialization asymmetry (#17) --------------------------------------- + + +async def _decimal_source(user_id, limit, next_page, **kw): + return FeedResult(data=[{"id": 1, "price": decimal.Decimal("9.99")}], next_page={}, has_next_page=False) + + +@pytest.mark.asyncio +async def test_non_serializable_item_with_cache_raises_named_error(): + node = S.wrapper(S.subfeed("src", "src"), session_size=10) + ctx = S.make_ctx({"src": _decimal_source}, redis=fakeredis.aioredis.FakeRedis()) + with pytest.raises(TypeError, match="Wrapper 'w'.*orjson-serializable"): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +@pytest.mark.asyncio +async def test_non_serializable_item_passthrough_works_control(): + """Control: the same item is fine without cache -- the asymmetry is documented, + and the cached path's error above must name the node so it is diagnosable.""" + node = S.wrapper(S.subfeed("src", "src")) + ctx = S.make_ctx({"src": _decimal_source}, redis=None) + r = await run_executor.run(node, ctx, limit=5, cursor={}) + assert len(r.data) == 1 diff --git a/tests/test_invariants.py b/tests/test_invariants.py index 3223fbd..0fbcc73 100644 --- a/tests/test_invariants.py +++ b/tests/test_invariants.py @@ -21,13 +21,12 @@ def _redis(): def _build(cfg): - """cfg keys: cache(bool), dedup(bool), overfetch(int), rerank(bool).""" + """cfg keys: cache(bool), dedup(bool), rerank(bool).""" child = S.subfeed("src", "src") node = S.wrapper( child, session_size=50 if cfg["cache"] else None, dedup_key="id" if cfg["dedup"] else None, - overfetch_factor=cfg.get("overfetch", 4), rerank_method="identity" if cfg.get("rerank") else None, ) return node @@ -45,17 +44,13 @@ def _p(id_, marks=(), **cfg): return pytest.param(cfg, id=id_, marks=marks) -# Coverage must hold for every overfetch value and both paths: no item is ever lost. +# Coverage must hold on both paths: no item is ever lost. COVERAGE_CONFIGS = [ - _p("cache-nodedup", cache=True, dedup=False, overfetch=4), - _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), - _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), - _p("passthrough-dedup-of1", cache=False, dedup=True, overfetch=1), - _p("cache-dedup-of2", cache=True, dedup=True, overfetch=2), - _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), - _p("passthrough-dedup-of2", cache=False, dedup=True, overfetch=2), - _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), - _p("cache-dedup-of4-rerank", cache=True, dedup=True, overfetch=4, rerank=True), + _p("cache-nodedup", cache=True, dedup=False), + _p("passthrough-nodedup", cache=False, dedup=False), + _p("cache-dedup", cache=True, dedup=True), + _p("passthrough-dedup", cache=False, dedup=True), + _p("cache-dedup-rerank", cache=True, dedup=True, rerank=True), ] @@ -72,15 +67,11 @@ async def test_coverage_no_item_lost(cfg): # Same matrix WITHOUT the LOSS marks: item loss shows up as coverage gaps, NOT as # duplicates or non-termination, so these properties hold for every config. NODUP_CONFIGS = [ - _p("cache-nodedup", cache=True, dedup=False, overfetch=4), - _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), - _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), - _p("passthrough-dedup-of1", cache=False, dedup=True, overfetch=1), - _p("cache-dedup-of2", cache=True, dedup=True, overfetch=2), - _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), - _p("passthrough-dedup-of2", cache=False, dedup=True, overfetch=2), - _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), - _p("cache-dedup-of4-rerank", cache=True, dedup=True, overfetch=4, rerank=True), + _p("cache-nodedup", cache=True, dedup=False), + _p("passthrough-nodedup", cache=False, dedup=False), + _p("cache-dedup", cache=True, dedup=True), + _p("passthrough-dedup", cache=False, dedup=True), + _p("cache-dedup-rerank", cache=True, dedup=True, rerank=True), ] @@ -95,13 +86,15 @@ async def test_no_duplicates_across_pages(cfg): @pytest.mark.asyncio -@pytest.mark.parametrize("cfg", [ - _p("cache-nodedup", cache=True, dedup=False, overfetch=4), - _p("passthrough-nodedup", cache=False, dedup=False, overfetch=4), - _p("cache-dedup-of1", cache=True, dedup=True, overfetch=1), - _p("cache-dedup-of4", cache=True, dedup=True, overfetch=4), - _p("passthrough-dedup-of4", cache=False, dedup=True, overfetch=4), -]) +@pytest.mark.parametrize( + "cfg", + [ + _p("cache-nodedup", cache=True, dedup=False), + _p("passthrough-nodedup", cache=False, dedup=False), + _p("cache-dedup", cache=True, dedup=True), + _p("passthrough-dedup", cache=False, dedup=True), + ], +) async def test_pages_are_full_until_exhaustion(cfg): """Every non-final page has exactly `limit` items (loss does not shorten pages, so this passes even where coverage fails -- keeping the two concerns separate).""" diff --git a/tests/test_manager.py b/tests/test_manager.py new file mode 100644 index 0000000..1e1ada3 --- /dev/null +++ b/tests/test_manager.py @@ -0,0 +1,68 @@ +"""FeedManager-level tests: the sole public entrypoint was +previously exercised only indirectly through executor.run(). +""" + +import pytest +from pydantic import ValidationError + +from smartfeed.manager import FeedManager +from smartfeed.models.base import FeedResult + + +async def make_items(user_id, limit, next_page, **kw): + page = next_page.get("page", 1) + start = (page - 1) * limit + return FeedResult( + data=[{"id": start + i} for i in range(limit)], + next_page={"page": page + 1}, + has_next_page=True, + ) + + +CONFIG = { + "feed": { + "type": "merger_percentage", + "node_id": "mix", + "items": [ + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "items"}}, + {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "items"}}, + ], + } +} + + +def test_malformed_config_raises_at_construction(): + with pytest.raises(ValidationError): + FeedManager(config={"feed": {"type": "wrapper"}}, methods_dict={}) # node_id/data missing + + +def test_unknown_node_type_raises_at_construction(): + with pytest.raises(ValidationError): + FeedManager(config={"feed": {"type": "no_such_node"}}, methods_dict={}) + + +@pytest.mark.asyncio +async def test_positions_stamped_0_based_contiguous(): + mgr = FeedManager(config=CONFIG, methods_dict={"items": make_items}) + r = await mgr.get_feed(session_id="u", limit=10) + positions = [it["_smartfeed_debug_info"]["smartfeed_position"] for it in r.data] + assert positions == list(range(len(r.data))) + assert len(r.data) == 10 + + +@pytest.mark.asyncio +async def test_cursor_round_trip_through_manager(): + mgr = FeedManager(config=CONFIG, methods_dict={"items": make_items}) + r1 = await mgr.get_feed(session_id="u", limit=10) + r2 = await mgr.get_feed(session_id="u", limit=10, cursor=r1.next_page) + ids1 = [it["id"] for it in r1.data] + ids2 = [it["id"] for it in r2.data] + assert ids2 and ids2 != ids1 # page 2 advanced both sources + + +@pytest.mark.asyncio +async def test_legacy_version_key_is_ignored(): + """Configs carrying the removed top-level `version` tag still parse.""" + mgr = FeedManager(config={"version": "2", **CONFIG}, methods_dict={"items": make_items}) + r = await mgr.get_feed(session_id="u", limit=4) + assert len(r.data) == 4 diff --git a/tests/test_medium_priority.py b/tests/test_medium_priority.py index 8f53609..edaeec3 100644 --- a/tests/test_medium_priority.py +++ b/tests/test_medium_priority.py @@ -167,9 +167,7 @@ async def test_result_contains_non_empty_child_items(self): assert len(result.data) > 0, "Expected items from the non-empty child" sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] - assert all(src == "a" for src in sources), ( - f"Expected all items from 'a', got sources: {sources}" - ) + assert all(src == "a" for src in sources), f"Expected all items from 'a', got sources: {sources}" @pytest.mark.asyncio async def test_has_next_page_reflects_non_empty_child(self): @@ -217,9 +215,7 @@ async def test_result_contains_only_non_empty_source(self): assert len(result.data) > 0, "Expected items from the non-empty source" sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] - assert all(src == "items" for src in sources), ( - f"Expected all items from 'items', got sources: {sources}" - ) + assert all(src == "items" for src in sources), f"Expected all items from 'items', got sources: {sources}" @pytest.mark.asyncio async def test_no_items_from_empty_source(self): @@ -271,9 +267,9 @@ async def test_odd_limit_returns_exactly_11(self): result = await run_executor.run(node, ctx, limit=11, cursor={}) - assert len(result.data) == 11, ( - f"Expected exactly 11 items for odd limit with 40/60 split, got {len(result.data)}" - ) + assert ( + len(result.data) == 11 + ), f"Expected exactly 11 items for odd limit with 40/60 split, got {len(result.data)}" @pytest.mark.asyncio async def test_odd_limit_split_sums_to_limit(self): @@ -337,13 +333,9 @@ async def test_positional_items_at_step_positions(self): position = i + 1 # 1-indexed source = item["_smartfeed_debug_info"]["source"] if position in pos_set: - assert source == "positional", ( - f"Position {position} (index {i}) should be 'positional', got '{source}'" - ) + assert source == "positional", f"Position {position} (index {i}) should be 'positional', got '{source}'" else: - assert source == "default", ( - f"Position {position} (index {i}) should be 'default', got '{source}'" - ) + assert source == "default", f"Position {position} (index {i}) should be 'default', got '{source}'" @pytest.mark.asyncio async def test_step_positions_count(self): @@ -360,8 +352,7 @@ async def test_step_positions_count(self): sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] assert sources.count("positional") == len(self._STEP_POSITIONS), ( - f"Expected {len(self._STEP_POSITIONS)} positional items, " - f"got {sources.count('positional')}" + f"Expected {len(self._STEP_POSITIONS)} positional items, " f"got {sources.count('positional')}" ) assert sources.count("default") == 20 - len(self._STEP_POSITIONS) @@ -389,9 +380,9 @@ async def test_only_positional_items_when_default_empty(self): assert len(result.data) > 0, "Expected positional items in the result" sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] - assert all(src == "positional" for src in sources), ( - f"Expected all items from 'positional', got sources: {sources}" - ) + assert all( + src == "positional" for src in sources + ), f"Expected all items from 'positional', got sources: {sources}" @pytest.mark.asyncio async def test_positional_items_fill_configured_positions(self): @@ -410,9 +401,6 @@ async def test_positional_items_fill_configured_positions(self): # With no default items, the assemble loop skips non-positional slots # and only places items at the configured positional slots. # Positional items appear but non-positional slots are vacant. - pos_set = set(positions) for i, item in enumerate(result.data): source = item["_smartfeed_debug_info"]["source"] - assert source == "positional", ( - f"Index {i} should be 'positional' (only source), got '{source}'" - ) + assert source == "positional", f"Index {i} should be 'positional' (only source), got '{source}'" diff --git a/tests/test_mixers.py b/tests/test_mixers.py index 3f76398..488752a 100644 --- a/tests/test_mixers.py +++ b/tests/test_mixers.py @@ -2,6 +2,7 @@ from smartfeed.models.subfeed import SubFeed from smartfeed.execution.context import ExecutionContext from smartfeed.execution import executor as run_executor +from smartfeed.models.mixers import MergerPercentage, MergerPercentageItem, MergerAppend, MergerPositional from tests.conftest import METHODS @@ -18,9 +19,6 @@ async def test_executor_runs_subfeed(ctx): assert result.data[0]["_smartfeed_debug_info"]["source"] == "items" -from smartfeed.models.mixers import MergerPercentage, MergerPercentageItem, MergerAppend, MergerPositional - - @pytest.mark.asyncio async def test_percentage_40_60(ctx): node = MergerPercentage( @@ -66,6 +64,6 @@ async def test_positional_inserts_at_positions(ctx): ) result = await run_executor.run(node, ctx, limit=5, cursor={}) sources = [item["_smartfeed_debug_info"]["source"] for item in result.data] - assert sources[0] == "promo" # position 1 + assert sources[0] == "promo" # position 1 assert sources[1] == "regular" - assert sources[2] == "promo" # position 3 + assert sources[2] == "promo" # position 3 diff --git a/tests/test_pagination_edge_cases.py b/tests/test_pagination_edge_cases.py index afe5282..3a8b21a 100644 --- a/tests/test_pagination_edge_cases.py +++ b/tests/test_pagination_edge_cases.py @@ -24,6 +24,7 @@ def _redis(): # has_next_page exactness: no phantom trailing empty page, no lost final item # --------------------------------------------------------------------------- + @pytest.mark.asyncio @pytest.mark.parametrize("n,limit", [(25, 10), (30, 10), (20, 10), (40, 10), (5, 10), (1, 10)]) async def test_has_next_exact_passthrough(n, limit): @@ -51,10 +52,11 @@ async def test_has_next_exact_cached(n, limit): # Empty source # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_empty_source_returns_empty_no_next(): src = S.ScriptedSource([]) - node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id", overfetch_factor=4) + node = S.wrapper(S.subfeed("src", "src"), session_size=20, dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) r = await run_executor.run(node, ctx, limit=10, cursor={}) assert r.data == [] @@ -65,6 +67,7 @@ async def test_empty_source_returns_empty_no_next(): # limit larger than the whole pool # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_limit_larger_than_pool(): src = S.ScriptedSource(S.unique_pool(5)) @@ -78,13 +81,13 @@ async def test_limit_larger_than_pool(): # shuffle=True must not defeat cross-page dedup (seen-set) # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_shuffle_does_not_defeat_cross_page_dedup(): pool = S.unique_pool(30) pool.insert(25, {"id": 3, "val": "again"}) # id 3 re-emitted on a later page src = S.ScriptedSource(pool) - node = S.wrapper(SubFeed(subfeed_id="src", method_name="src", shuffle=True), - dedup_key="id", overfetch_factor=1) + node = S.wrapper(SubFeed(subfeed_id="src", method_name="src", shuffle=True), dedup_key="id") ctx = S.make_ctx({"src": src}, redis=_redis()) pages = await S.drain(node, ctx, limit=5, max_pages=50) ids = S.flat_ids(pages) diff --git a/tests/test_resilience.py b/tests/test_resilience.py index 1707a33..ee7bc4c 100644 --- a/tests/test_resilience.py +++ b/tests/test_resilience.py @@ -15,6 +15,7 @@ # Helpers # --------------------------------------------------------------------------- + async def make_source(prefix, user_id, limit, next_page, **kwargs): page = next_page.get("page", 1) start = (page - 1) * limit @@ -63,8 +64,8 @@ def ctx(redis): # 1. Source failure: fill page from remaining source # --------------------------------------------------------------------------- -class TestSourceFailureFallback: +class TestSourceFailureFallback: @pytest.mark.asyncio async def test_one_source_fails_page_filled_from_other(self, ctx): """If one source in percentage mix fails, the other fills the page.""" @@ -119,8 +120,8 @@ async def test_all_sources_fail_returns_empty(self, ctx): # 2. Positional: strict position enforcement # --------------------------------------------------------------------------- -class TestPositionalStrictPositions: +class TestPositionalStrictPositions: @pytest.mark.asyncio async def test_promo_at_exact_positions(self, ctx): """Promo items must be at positions 1, 3, 5, 7 (1-indexed).""" @@ -136,9 +137,7 @@ async def test_promo_at_exact_positions(self, ctx): for pos in [1, 3, 5, 7]: item = result.data[pos - 1] # 1-indexed -> 0-indexed source = item.get("_smartfeed_debug_info", {}).get("source", "") - assert source == "promo", ( - f"Position {pos}: expected promo, got {source} (id={item.get('id')})" - ) + assert source == "promo", f"Position {pos}: expected promo, got {source} (id={item.get('id')})" @pytest.mark.asyncio async def test_non_promo_positions_filled_by_default(self, ctx): @@ -155,16 +154,14 @@ async def test_non_promo_positions_filled_by_default(self, ctx): for pos in non_promo_positions: item = result.data[pos] source = item.get("_smartfeed_debug_info", {}).get("source", "") - assert source == "regular", ( - f"Position {pos + 1}: expected regular, got {source}" - ) + assert source == "regular", f"Position {pos + 1}: expected regular, got {source}" @pytest.mark.asyncio async def test_positional_with_dedup_preserves_positions(self, ctx): """Even with outer dedup, promo positions must be preserved.""" node = Wrapper( node_id="dedup_outer", - dedup=WrapperDedup(dedup_key="id", overfetch_factor=3), + dedup=WrapperDedup(dedup_key="id"), data=MergerPositional( node_id="pos", positions=[1, 3, 5, 7], @@ -178,17 +175,15 @@ async def test_positional_with_dedup_preserves_positions(self, ctx): for pos in [1, 3, 5, 7]: item = result.data[pos - 1] source = item.get("_smartfeed_debug_info", {}).get("source", "") - assert source == "promo", ( - f"Position {pos}: expected promo after dedup, got {source}" - ) + assert source == "promo", f"Position {pos}: expected promo after dedup, got {source}" # --------------------------------------------------------------------------- # 3. Rerank failure: raise_error=True vs False # --------------------------------------------------------------------------- -class TestRerankFailureHandling: +class TestRerankFailureHandling: @pytest.mark.asyncio async def test_rerank_crash_raise_error_true(self, ctx): """With raise_error=True (default), rerank failure crashes the request.""" diff --git a/tests/test_shared_lock_contention.py b/tests/test_shared_lock_contention.py new file mode 100644 index 0000000..e1f7b00 --- /dev/null +++ b/tests/test_shared_lock_contention.py @@ -0,0 +1,43 @@ +"""Two wrappers sharing a cache_key race one cold build: the shared-segment lock +must serialize them so the child is fetched exactly once (the _build_shared_base +waiter path had no concurrent coverage). + +The child has latency, so both wrappers are guaranteed to be inside the cold +build window together; their own coldlocks differ (different wrapper config +hashes), so contention lands on the shared-segment lock itself. +""" + +import asyncio + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +async def identity(items, session_id): + return items + + +async def reverse(items, session_id): + return list(reversed(items)) + + +@pytest.mark.asyncio +async def test_concurrent_shared_cold_build_fetches_child_once(): + src = S.ScriptedSource(S.unique_pool(100), latency=0.3) + ctx = S.make_ctx({"src": src, "identity": identity, "reverse": reverse}, redis=fakeredis.aioredis.FakeRedis()) + w_a = S.wrapper(S.subfeed("src", "src"), node_id="a", session_size=20, cache_key="pool", rerank_method="identity") + w_b = S.wrapper(S.subfeed("src", "src"), node_id="b", session_size=20, cache_key="pool", rerank_method="reverse") + + a, b = await asyncio.gather( + run_executor.run(w_a, ctx, limit=10, cursor={}), + run_executor.run(w_b, ctx, limit=10, cursor={}), + ) + + assert src.calls == 1, f"shared base must be fetched once, got {src.calls}" + assert len(a.data) == 10 and len(b.data) == 10 + # Same shared base, per-wrapper rerank: b serves a's window reversed-ish (both + # windows come from the same 20-item segment). + assert {it["id"] for it in a.data} | {it["id"] for it in b.data} <= set(range(20)) diff --git a/tests/test_sharp_edges.py b/tests/test_sharp_edges.py new file mode 100644 index 0000000..93388e6 --- /dev/null +++ b/tests/test_sharp_edges.py @@ -0,0 +1,82 @@ +"""README's documented "sharp edges" pinned as tests. + +Edge 2 (rerank length mismatch always raises) and edges 1/3/4 previously had no +coverage in the exact configurations the README warns about. Edge 5 (cursor +validation) lives in test_input_validation.py. +""" + +import pytest + +from smartfeed.models.base import FeedResult +from smartfeed.models.subfeed import SubFeed +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +# -- edge 1: missing method_name ---------------------------------------------- + + +@pytest.mark.asyncio +async def test_missing_method_raises_keyerror_even_fail_soft(): + """method_name is resolved BEFORE the try block: KeyError regardless of raise_error.""" + node = SubFeed(subfeed_id="s", method_name="no_such_method", raise_error=False) + ctx = S.make_ctx({}) + with pytest.raises(KeyError, match="no_such_method"): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +# -- edge 3: fetch function without **kwargs ----------------------------------- + + +async def no_kwargs_fetch(user_id, limit, next_page): + return FeedResult(data=[{"id": 1}], next_page={}, has_next_page=False) + + +@pytest.mark.asyncio +async def test_no_kwargs_source_raises_typeerror(): + node = S.subfeed("s", "src") + ctx = S.make_ctx({"src": no_kwargs_fetch}) + with pytest.raises(TypeError): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +@pytest.mark.asyncio +async def test_no_kwargs_source_swallowed_to_empty_page_when_fail_soft(): + """The documented footgun: under raise_error=False the signature bug reads + exactly like 'this source had no items today'.""" + node = SubFeed(subfeed_id="s", method_name="src", raise_error=False) + ctx = S.make_ctx({"src": no_kwargs_fetch}) + r = await run_executor.run(node, ctx, limit=5, cursor={}) + assert r.data == [] + assert r.has_next_page is False + + +# -- edge 4: subfeed_params colliding with injected kwargs ---------------------- + + +@pytest.mark.asyncio +async def test_subfeed_params_collision_raises_typeerror(): + async def fine_fetch(user_id, limit, next_page, **kw): + return FeedResult(data=[], next_page={}, has_next_page=False) + + node = SubFeed(subfeed_id="s", method_name="src", subfeed_params={"limit": 3}) + ctx = S.make_ctx({"src": fine_fetch}) + with pytest.raises(TypeError): + await run_executor.run(node, ctx, limit=5, cursor={}) + + +# -- edge 2: rerank length mismatch always raises ------------------------------- + + +@pytest.mark.asyncio +async def test_rerank_length_mismatch_raises_even_fail_soft(): + """rerank.raise_error=False covers only exceptions FROM the callable; the + length contract is enforced unconditionally.""" + + async def drops_one(items, session_id): + return items[:-1] + + node = S.wrapper(S.subfeed("s", "src"), rerank_method="drops_one", rerank_raise=False) + ctx = S.make_ctx({"src": S.ScriptedSource(S.unique_pool(10)), "drops_one": drops_one}) + with pytest.raises(ValueError, match="must return exactly"): + await run_executor.run(node, ctx, limit=5, cursor={}) diff --git a/tests/test_silent_degradation.py b/tests/test_silent_degradation.py new file mode 100644 index 0000000..c19a0d2 --- /dev/null +++ b/tests/test_silent_degradation.py @@ -0,0 +1,30 @@ +"""Documented silent downgrades: these configurations must +degrade gracefully, not crash -- and must not write state they cannot honor. +""" + +import pytest +import fakeredis.aioredis + +from tests import sources as S +from smartfeed.execution import executor as run_executor + + +@pytest.mark.asyncio +async def test_cache_without_redis_degrades_to_passthrough(): + node = S.wrapper(S.subfeed("s", "src"), session_size=20) + ctx = S.make_ctx({"src": S.ScriptedSource(S.unique_pool(50))}, redis=None) + r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + assert [it["id"] for it in r1.data] == list(range(10)) + # Passthrough exposes the raw child cursor and keeps paginating without Redis. + r2 = await run_executor.run(node, ctx, limit=10, cursor=r1.next_page) + assert [it["id"] for it in r2.data] == list(range(10, 20)) + + +@pytest.mark.asyncio +async def test_cache_key_without_cache_is_pure_passthrough(): + redis = fakeredis.aioredis.FakeRedis() + node = S.wrapper(S.subfeed("s", "src"), cache_key="shared") # no cache block + ctx = S.make_ctx({"src": S.ScriptedSource(S.unique_pool(50))}, redis=redis) + r = await run_executor.run(node, ctx, limit=10, cursor={}) + assert len(r.data) == 10 + assert await redis.keys("*") == [] # no cache/shared/lock keys written diff --git a/tests/test_subfeed.py b/tests/test_subfeed.py index 1097dd0..2e3e3eb 100644 --- a/tests/test_subfeed.py +++ b/tests/test_subfeed.py @@ -6,9 +6,7 @@ @pytest.mark.asyncio async def test_subfeed_returns_data(): sf = SubFeed(subfeed_id="items", method_name="items") - result = await sf.execute( - methods_dict=METHODS, session_id="s1", limit=5, cursor={} - ) + result = await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor={}) assert len(result.data) == 5 assert result.has_next_page is True @@ -16,10 +14,8 @@ async def test_subfeed_returns_data(): @pytest.mark.asyncio async def test_subfeed_stamps_debug_info(): sf = SubFeed(subfeed_id="my_source", method_name="items") - result = await sf.execute( - methods_dict=METHODS, session_id="s1", limit=3, cursor={} - ) - for i, item in enumerate(result.data): + result = await sf.execute(methods_dict=METHODS, session_id="s1", limit=3, cursor={}) + for item in result.data: info = item["_smartfeed_debug_info"] assert info["source"] == "my_source" @@ -34,9 +30,7 @@ async def test_subfeed_raise_error(): @pytest.mark.asyncio async def test_subfeed_swallow_error(): sf = SubFeed(subfeed_id="err", method_name="error", raise_error=False) - result = await sf.execute( - methods_dict=METHODS, session_id="s1", limit=5, cursor={} - ) + result = await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor={}) assert result.data == [] assert result.has_next_page is False @@ -45,7 +39,5 @@ async def test_subfeed_swallow_error(): async def test_subfeed_cursor_propagation(): sf = SubFeed(subfeed_id="items", method_name="items") r1 = await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor={}) - r2 = await sf.execute( - methods_dict=METHODS, session_id="s1", limit=5, cursor=r1.next_page - ) + r2 = await sf.execute(methods_dict=METHODS, session_id="s1", limit=5, cursor=r1.next_page) assert r2.data[0]["id"] == 5 # page 2 starts at id=5 diff --git a/tests/test_wrapper_cache.py b/tests/test_wrapper_cache.py index d05fd41..104883a 100644 --- a/tests/test_wrapper_cache.py +++ b/tests/test_wrapper_cache.py @@ -56,7 +56,7 @@ async def test_generation_id_in_cursor(ctx): @pytest.mark.asyncio async def test_stale_gen_resets_to_page1(ctx, redis): node = _make_wrapper() - r1 = await run_executor.run(node, ctx, limit=10, cursor={}) + await run_executor.run(node, ctx, limit=10, cursor={}) # Flush to simulate TTL expiry await redis.flushall() stale_cursor = {"cached": {"offset": 40, "gen": "stale_nonce"}} diff --git a/tests/test_wrapper_dedup.py b/tests/test_wrapper_dedup.py index 22fe0cd..cedff78 100644 --- a/tests/test_wrapper_dedup.py +++ b/tests/test_wrapper_dedup.py @@ -36,7 +36,7 @@ def ctx(redis): async def test_dedup_removes_duplicates(ctx): node = Wrapper( node_id="deduped", - dedup=WrapperDedup(dedup_key="id", overfetch_factor=2), + dedup=WrapperDedup(dedup_key="id"), data=MergerAppend( node_id="mix", items=[ diff --git a/tests/test_wrapper_full_pipeline.py b/tests/test_wrapper_full_pipeline.py index bbea611..9e98ebf 100644 --- a/tests/test_wrapper_full_pipeline.py +++ b/tests/test_wrapper_full_pipeline.py @@ -36,22 +36,44 @@ def ctx(redis): @pytest.mark.asyncio async def test_full_pipeline(ctx): - config = FeedConfig.model_validate({ - "version": "2", - "feed": { - "type": "wrapper", "node_id": "full", - "cache": {"session_size": 30, "session_ttl": 300}, - "rerank": {"method_name": "reverse"}, - "dedup": {"dedup_key": "id", "overfetch_factor": 2}, - "data": { - "type": "merger_percentage", "node_id": "mix", - "items": [ - {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "a", "method_name": "source_a", "dedup_priority": 5}}, - {"percentage": 50, "data": {"type": "subfeed", "subfeed_id": "b", "method_name": "source_b", "dedup_priority": 1}}, - ], + config = FeedConfig.model_validate( + { + "version": "2", + "feed": { + "type": "wrapper", + "node_id": "full", + "cache": {"session_size": 30, "session_ttl": 300}, + "rerank": {"method_name": "reverse"}, + "dedup": { + "dedup_key": "id", + }, + "data": { + "type": "merger_percentage", + "node_id": "mix", + "items": [ + { + "percentage": 50, + "data": { + "type": "subfeed", + "subfeed_id": "a", + "method_name": "source_a", + "dedup_priority": 5, + }, + }, + { + "percentage": 50, + "data": { + "type": "subfeed", + "subfeed_id": "b", + "method_name": "source_b", + "dedup_priority": 1, + }, + }, + ], + }, }, - }, - }) + } + ) r1 = await run_executor.run(config.feed, ctx, limit=10, cursor={}) diff --git a/tests/test_wrapper_rerank.py b/tests/test_wrapper_rerank.py index 83608b6..727a7f2 100644 --- a/tests/test_wrapper_rerank.py +++ b/tests/test_wrapper_rerank.py @@ -54,6 +54,9 @@ async def test_rerank_without_cache(ctx): ) result = await run_executor.run(node, ctx, limit=5, cursor={}) assert len(result.data) == 5 + # The rerank must actually run: "reverse" flips the order, so a silently + # skipped rerank (ascending ids) fails here. + assert result.data[0]["id"] > result.data[-1]["id"] @pytest.mark.asyncio diff --git a/tests/test_wrapper_shared_cache.py b/tests/test_wrapper_shared_cache.py index 5d707ee..3181fac 100644 --- a/tests/test_wrapper_shared_cache.py +++ b/tests/test_wrapper_shared_cache.py @@ -10,29 +10,36 @@ async def rerank_reverse(items, session_id): return list(reversed(items)) + async def rerank_identity(items, session_id): return items + SHARED_METHODS = {**METHODS, "rerank_a": rerank_reverse, "rerank_b": rerank_identity} + @pytest.fixture def redis(): return fakeredis.aioredis.FakeRedis() + @pytest.fixture def ctx(redis): return ExecutionContext(session_id="s1", methods_dict=SHARED_METHODS, redis=redis) + @pytest.mark.asyncio async def test_shared_cache_two_rerankers(ctx): w_a = Wrapper( - node_id="a", cache_key="shared", + node_id="a", + cache_key="shared", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_a"), data=SubFeed(subfeed_id="items", method_name="items"), ) w_b = Wrapper( - node_id="b", cache_key="shared", + node_id="b", + cache_key="shared", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_b"), data=SubFeed(subfeed_id="items", method_name="items"), @@ -49,13 +56,15 @@ async def test_shared_cache_two_rerankers(ctx): async def test_shared_cache_base_written_once(ctx, redis): """Both wrappers use same cache_key -- base cache key should exist once.""" w_a = Wrapper( - node_id="a", cache_key="shared_base", + node_id="a", + cache_key="shared_base", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_a"), data=SubFeed(subfeed_id="items", method_name="items"), ) w_b = Wrapper( - node_id="b", cache_key="shared_base", + node_id="b", + cache_key="shared_base", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_b"), data=SubFeed(subfeed_id="items", method_name="items"), @@ -67,21 +76,27 @@ async def test_shared_cache_base_written_once(ctx, redis): # But both share the base key (cache_key="shared_base") all_keys = [k.decode() if isinstance(k, bytes) else k for k in await redis.keys("sf:*")] # The shared base key should exist - base_keys = [k for k in all_keys if ":shared_base:" in k and not k.endswith(":a:") and not k.endswith(":b:")] - assert len(base_keys) >= 1 + # Shared segments are sf:{sid}:{cache_key}:{child_hash}:{segment} (4 colons); + # the per-wrapper caches sf:{sid}:{cache_key}:{wrapper_hash} have only 3. + segment_keys = [ + k for k in all_keys if ":shared_base:" in k and k.count(":") == 4 and not k.endswith((":meta", ":lock")) + ] + assert len(segment_keys) == 1, f"exactly one shared base segment expected, got {segment_keys}" @pytest.mark.asyncio async def test_shared_cache_warm_read(ctx, redis): """After cold build, reading warm cache returns same base data, different rerank.""" w_a = Wrapper( - node_id="a", cache_key="shared_warm", + node_id="a", + cache_key="shared_warm", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_a"), data=SubFeed(subfeed_id="items", method_name="items"), ) w_b = Wrapper( - node_id="b", cache_key="shared_warm", + node_id="b", + cache_key="shared_warm", cache=WrapperCache(session_size=20, session_ttl=300), rerank=WrapperRerank(method_name="rerank_b"), data=SubFeed(subfeed_id="items", method_name="items"), From 637822a5904f2fee56c78dc1e0fe3d9e6bd8038d Mon Sep 17 00:00:00 2001 From: Pavel Kochetov Date: Thu, 9 Jul 2026 16:18:44 +0100 Subject: [PATCH 41/41] Changelog commited separately. --- CHANGELOG.md | 69 +++++++++++++++-------------- TODO.md | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 35 deletions(-) create mode 100644 TODO.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 12fd4a0..c41b80e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,57 +1,56 @@ # Changelog -## 0.3.0 (2026-04-08) +## 1.0.0 (2026-07-07) -SmartFeed v2: complete rewrite. - -### Breaking Changes - -- Unified `Wrapper` node replaces `MergerViewSession`, `MergerDeduplication`, and external rerank hack -- `FeedManager.get_data()` -> `FeedManager.get_feed(session_id, limit, cursor)` -- `user_id` -> `session_id` in ExecutionContext -- Cursors are plain dicts (not `FeedResultNextPage` Pydantic models) -- `FeedResultClient` removed (use `FeedResult` everywhere) -- `merger_id` renamed to `node_id` in configs -- Config `version: "2"` required -- Pydantic v2 required (v1 compat removed) -- `jsonlib.py` removed (orjson used directly) +SmartFeed v2: complete rewrite. Changes below are relative to the last published release, **0.2.0**. (The pre-rewrite internal working tree contained additional intermediate modules — never part of a published release — that are not listed here.) See [MIGRATION.md](MIGRATION.md) for step-by-step upgrade instructions. ### New Features - **Wrapper** with optional pipeline stages: cache, dedup, rerank -- **Rerank callable** registered in `methods_dict`, referenced by `method_name` in config +- **Session-scoped dedup contract**: within one continuous scroll a user sees each item at most once, no matter how far they scroll; an empty cursor starts a fresh scroll - **Cross-page dedup** via Redis seen-set (SET with TTL) +- **Deficit-based refill**: dedup returns full pages without discarding unique items (both cached and passthrough paths) +- **Rerank callable** registered in `methods_dict`, referenced by `method_name` in config - **Shared cache** via `cache_key` for A/B testing with different rerankers - **Generation ID** for stale cursor detection - **TTL touch** (inactivity timeout) on every cache access -- **Overfetch + refill** for dedup in both cached and passthrough paths - **`raise_error`** on WrapperRerank: fail-hard or fail-soft mode -- **`_smartfeed_debug_info`** stamped automatically: source, positions, per-wrapper data -- **SmartFeedDebugInfo / FeedItem** Pydantic models for typed output +- **`_smartfeed_debug_info`** stamped automatically: source, positions, per-wrapper data (plain dicts, all positions 0-based) - **`config_hash`** for automatic cache invalidation on config changes +- **Config validation at construction**: duplicate `subfeed_id`/`node_id` anywhere in the tree are rejected (they share one cursor/Redis namespace); gradient `step`/`size_to_step` must be > 0 +- **Input validation**: `get_feed` rejects non-positive `limit`, empty `session_id`, non-dict `cursor`; malformed per-node cursor contents (tampered `offset`/`gen`/`page`) raise `ValueError` naming the node +- **Session batch stored as a Redis LIST**: warm page reads transfer and parse only the requested window (single atomic pipeline: meta + LLEN + LRANGE + TTL refresh) +- **PEP 561 `py.typed` marker**: annotations are visible to downstream type checkers -### Removed +### Breaking Changes -- `MergerViewSession` (replaced by `Wrapper(cache=...)`) -- `MergerDeduplication` (replaced by `Wrapper(dedup=...)`) -- `DedupRuntime` (~453 lines) -- `SlotsPlan`, `SlotSpec`, `CallablePlan` (replaced by `MixPlan`) -- `CursorMap` with `merge_delta` / overfetch -- `DeduplicationPolicy`, `CursorSeenStore`, `RedisSeenStore` -- `pydantic_compat.py` -- `jsonlib.py` +- Models moved from `smartfeed.schemas` to `smartfeed.models` (`FeedManager` keeps its name) +- `FeedManager.get_data(user_id, limit, next_page, **params)` -> `FeedManager.get_feed(session_id, limit, cursor)` +- `user_id` -> `session_id` in ExecutionContext +- Cursors are plain dicts, not `FeedResultNextPage` Pydantic models +- Config: `merger_view_session` -> unified `wrapper` node with optional `cache` / `dedup` / `rerank` stages (`session_live_time` -> `cache.session_ttl`) +- Config: `merger_id` -> `node_id` on merger and wrapper nodes (subfeed `subfeed_id` is unchanged) +- Config: merger-level `shuffle` removed; only `SubFeed.shuffle` remains +- Config: `MergerPositional.start` / `end` / `step` removed; `positions` are now page-relative (were absolute across pages) +- Config: top-level `version` is no longer required or read (unknown top-level keys are ignored, so configs carrying it still parse) +- `FeedResultClient`, `FeedResultNextPage` removed (use `FeedResult` and plain-dict cursors) +- Pydantic v2 required (v1 support removed) + +### Removed (from the 0.2.0 public surface) + +- `smartfeed.schemas` module (models now live in `smartfeed.models`) +- `MergerViewSession` node (replaced by `Wrapper(cache=..., dedup=...)`) +- merger-level `shuffle`; `MergerPositional.start` / `end` / `step` +- `FeedResultClient`, `FeedResultNextPage` - `examples/` -- `mergers/` directory (mixers now in `models/mixers.py`) -- `policies/` directory -- `feed_models.py`, `schemas.py` -### Stats +### Stats (v2 vs the pre-rewrite internal working tree) -| Metric | v1 | v2 | -|--------|----|----| +| Metric | pre-rewrite | v2 | +|--------|-------------|----| | Python files | 24 | 11 | -| Lines of code | 2838 | ~1300 | -| Test files | 26 | 17 | +| Lines of code | 2838 | ~1650 | +| Test files | 26 | 43 | ## 0.2.0 (2025-11-25) diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..84c2ce1 --- /dev/null +++ b/TODO.md @@ -0,0 +1,121 @@ +# TODO — next version + +Planning notes for the next batch of changes. Current error handling works but is +far from ideal; this file scopes what "ideal" looks like and the open questions. + +> Compat reminder (from CLAUDE.md): runtime is **Python 3.9+**. Anything below that +> relies on 3.11-only APIs (e.g. `Exception.add_note`) needs a fallback. + +--- + +## Theme: error handling overhaul + +### Current state (for reference) + +- **Fail-hard is the default** (`SubFeed.raise_error = True`). On failure the exception + is re-raised at `smartfeed/models/subfeed.py:47`, then propagates through the mixer's + `asyncio.gather(...)` at `smartfeed/execution/executor.py:34` (no `return_exceptions`), + so it bubbles out of `get_feed` raw. +- **Fail-soft is opt-in** (`raise_error=False`). The subfeed swallows the exception and + returns `FeedResult(data=[], next_page={subfeed_id: }, has_next_page=False)` + (`subfeed.py:47`). It is **completely silent** — no log, no debug stamp (empty data = + nothing to stamp), so the caller cannot tell "source failed" from "source was empty". +- **Backfill of a dead source's slots** only happens via the wrapper refill loop + (`wrapper.py:251` passthrough-with-dedup, `wrapper.py:663` cached cold build). A bare + mixer or a non-dedup passthrough does **not** redistribute — the dead child's demand + share just yields nothing and the page comes back short. +- The package currently has **no logging and no `print`** anywhere — there is no logger + to hook into yet. + +--- + +### 1. Fail-hard: propagate a clear, understandable error + +**Problem.** When a hard failure crashes the whole feed, the exception that surfaces is +the raw error from the user's method. It carries no smartfeed-level context: *which* +subfeed (`subfeed_id`), *where* in the config tree, on *which* page/cursor. Debugging a +crash means guessing which source blew up. + +**Desired.** The propagated error names the failing node and preserves the original +traceback — e.g. "SubFeed 'trending_posts' raised while fetching page N" chained to the +original exception, so `raise ... from e` keeps the root cause visible. + +**Ideas / anchors.** +- Catch at `subfeed.py:47` in the `raise_error=True` branch; attach context, then re-raise. +- Preserve the original: either a dedicated `SubFeedError`/`SmartFeedError` that chains via + `from e`, or annotate in place. **`Exception.add_note` is 3.11+** — needs a 3.9/3.10 + fallback (wrap in a custom exception, or prepend context to a re-raised copy). +- Consider the `gather` semantics: with the default, only the *first* exception is raised + and sibling failures are discarded. Decide whether to keep first-wins or switch to + `return_exceptions=True` and aggregate (e.g. surface an ExceptionGroup / a combined + message) — `executor.py:34`. + +**Open questions.** +- Introduce a real exception hierarchy (`SmartFeedError` base) or keep raising the + original type with added context? +- Should the error object carry structured fields (node_id, cursor slice, page) for + programmatic handling, not just a message? + +--- + +### 2. Fail-soft: observability + self-heal without a dedup node + +Two sub-goals. + +**2a. Make soft failures observable (not silent).** + +Today a soft failure vanishes. We want the failure recorded somewhere the integrator can +see. Options to weigh: +- Emit via the stdlib `logging` module (`logging.getLogger("smartfeed")`) — library-correct, + but there's no logger anywhere yet, so this is a from-scratch decision on logger name / + levels / what we log. +- Surface it in the result: a feed-level `errors` channel (per-source failure list) on + `FeedResult`, or a debug entry. Note: item-level `_smartfeed_debug_info` won't work for a + failed source because `data=[]` — there's no item to stamp. Needs a feed-level home. +- Optional user callback / hook (`on_source_error=...`) so the integrator decides. +- (Explicitly *not* `print` — a published library must not print.) + +**Open questions.** +- Where does a soft error live: logs only, `FeedResult` field, callback, or several? +- Does adding an `errors` field to `FeedResult` break the documented "byte-exact" result + shape? If so, gate it / version it. + +**2b. Self-heal (backfill dead-source positions) even with no dedup node.** *(ideal / exploratory — no chosen design yet)* + +Today, surviving siblings only backfill a dead source's slots when a dedup wrapper (or a +cached cold build) sits above, because only the wrapper runs a deficit-refill loop. Goal: +the page fills from healthy sources even when there's no dedup wrapper in the tree. + +**Tensions / constraints to respect.** +- **Do not reintroduce over-fetch.** Over-fetching demand as a buffer was deliberately + removed (commit `d799de8` "overfetch killed bc YAGNI"; pinned by + `tests/test_bug_overfetch_item_loss.py`). Any backfill design must not silently discard + fetched items or re-open that bug. +- The current mixer model fans out **once** via `gather` then `assemble` — there is no + re-fetch round at the mixer level today. + +**Candidate directions (pick / prototype later).** +- Generalize the wrapper's deficit-refill loop **down into mixers / the executor**, so any + mixer can re-request the outstanding deficit from its surviving children (mirrors + `wrapper.py:663`, without requiring dedup). Decide where it lives: `executor.run` for + mixers, or a new `MixerNode` capability. +- Decide proportion behavior when a source dies: do the survivors' percentages + **renormalize** to cover the gap, or just fill in configured order? +- Bound it like the wrapper does (`_REFILL_SCAN_FACTOR`) so a permanently-dead or + duplicate-only source can't spin the loop forever. + +**Open questions.** +- Should self-heal be always-on, or opt-in per mixer (a `backfill=True` flag)? +- Interaction with soft-fail's unadvanced self-scoped cursor (`subfeed.py:52`): a source + that fails every round must not stall the loop — needs a "give up on this child" signal. + +--- + +## Cross-cutting (don't forget when implementing) + +- Add `tests/test_bug_*.py` pins for each fixed behavior (repo convention — don't fold into + existing files). +- Update `README.md` "Error philosophy" + the `ARCHITECTURE.md` internals if behavior/keys + move; add a `CHANGELOG.md` entry. +- Keep public annotations precise (`py.typed`); run `make lint` (ruff + pyright + black). +- Honor 3.9+ compat for any exception-notes / logging approach.