perf(feed): a shuffle walks a stored random key instead of sorting the whole scope

Every video carries a random number drawn once (videos.shuffle_key, migration 0118, indexed with
id). A seed picks a point on that ring: the feed's Shuffle sort pages along it and wraps, and the
Home surprise reads a 500-video window of it and shuffles that window by the seed, so no video is
favoured for sitting after a wide gap. The md5(id || seed) order read and sorted the whole feed
scope on every call: ~350 ms on prod per surprise, 381-767 ms per Shuffle page on localdev, now
14 ms and 4 ms.

backend/tests/test_shuffle_walk.py

@@ -0,0 +1,151 @@
+"""A seeded shuffle walks a stored random key instead of sorting the whole scope. DB lane.
+
+Every video carries a random number drawn once (`Video.shuffle_key`). A seed picks a point on that
+ring; the feed's Shuffle sort pages along the ring from there and wraps, and the Home surprise reads
+a window of it and shuffles that window by the seed (`discover.queries.yt_surprise`). Ordering the
+scope by md5(id || seed) cost ~350 ms per call on prod and grew with every Shuffle page.
+"""
+
+from contextlib import contextmanager
+from datetime import datetime, timedelta, timezone
+
+from sqlalchemy import event
+
+from app import feedquery
+from app.catalog.models import Channel, Subscription, Video
+from app.discover import queries
+from app.discover.queries import _FEED_BASE
+
+NOW = datetime.now(timezone.utc)
+KEYS = {"k05": 0.05, "k20": 0.20, "k40": 0.40, "k60": 0.60, "k80": 0.80, "k95": 0.95}
+
+
+def _seed(db, make_user, keys=KEYS):
+    user = make_user()
+    db.add(Channel(id="UCring", title="Ring"))
+    db.flush()
+    db.add(Subscription(user_id=user.id, channel_id="UCring"))
+    db.add_all(
+        [
+            Video(
+                id=vid,
+                channel_id="UCring",
+                title=vid,
+                duration_seconds=20 * 60,
+                published_at=NOW - timedelta(days=i + 1),
+                shuffle_key=key,
+            )
+            for i, (vid, key) in enumerate(keys.items())
+        ]
+    )
+    db.commit()
+    return user
+
+
+def _ring(seed, keys=KEYS):
+    """The ring order from the seed's point: the keys at or after it ascending, then the rest."""
+    r = feedquery.shuffle_start(seed)
+    return sorted(keys, key=lambda v: (keys[v] < r, keys[v]))
+
+
+def _seed_landing_between(lo, hi):
+    return next(s for s in range(1000) if lo < feedquery.shuffle_start(s) < hi)
+
+
+def test_the_walk_starts_at_the_seeds_point_and_wraps(db, make_user):
+    user = _seed(db, make_user)
+    query, _ = feedquery.filtered_query(db, user, **_FEED_BASE)
+    seed = _seed_landing_between(0.45, 0.55)  # starts at k60, wraps to k05 after k95
+    got = [r.id for r in feedquery.shuffle_walk(db, query, seed, 6)]
+    assert got == ["k60", "k80", "k95", "k05", "k20", "k40"]
+    assert [r.id for r in feedquery.shuffle_walk(db, query, seed, 4)] == got[:4]
+
+
+def test_the_feed_shuffle_pages_along_the_ring(db, make_user, as_user):
+    user = _seed(db, make_user)
+    client = as_user(user)
+    for seed in (_seed_landing_between(0.1, 0.15), _seed_landing_between(0.85, 0.9)):
+        walked, cursor = [], None
+        for _ in range(5):
+            url = f"/api/feed?sort=shuffle&seed={seed}&limit=2" + (
+                f"&cursor={cursor}" if cursor else ""
+            )
+            body = client.get(url).json()
+            walked += [it["id"] for it in body["items"]]
+            cursor = body["next_cursor"]
+            if not body["has_more"]:
+                break
+        assert walked == _ring(seed), seed
+
+
+def test_a_new_seed_opens_the_shuffle_elsewhere(db, make_user, as_user):
+    user = _seed(db, make_user)
+    client = as_user(user)
+    first = {
+        seed: client.get(f"/api/feed?sort=shuffle&seed={seed}&limit=2").json()["items"][0]["id"]
+        for seed in (_seed_landing_between(0.1, 0.15), _seed_landing_between(0.65, 0.7))
+    }
+    assert sorted(first.values()) == ["k20", "k80"]
+
+
+@contextmanager
+def _statements(db):
+    seen: list[str] = []
+
+    def grab(conn, cursor, statement, *args):
+        seen.append(statement)
+
+    engine = db.get_bind()
+    event.listen(engine, "before_cursor_execute", grab)
+    try:
+        yield seen
+    finally:
+        event.remove(engine, "before_cursor_execute", grab)
+
+
+def _draws_of_videos(seen):
+    return [s for s in seen if "FROM videos" in s and "LIMIT" in s]
+
+
+def test_neither_draw_orders_the_scope_by_a_hash(db, make_user):
+    # Each draw reads the videos as an ascending range of the key (what ix_videos_shuffle_key serves),
+    # never ordered by a hash or by the ring's CASE (which no index can serve).
+    user = _seed(db, make_user)
+    with _statements(db) as surprise:
+        queries.yt_surprise(db, user, "any", 7, 3)
+    with _statements(db) as page:
+        feedquery.feed_page(
+            db, user, dict(_FEED_BASE), sort="shuffle", seed=7, limit=2, cursor=None, lang="en"
+        )
+    for seen in (surprise, page):
+        draws = _draws_of_videos(seen)
+        assert draws
+        for s in draws:
+            assert "md5(" not in s.lower()
+            assert "ORDER BY videos.shuffle_key, videos.id" in s.replace("\n", " "), s
+
+
+def test_the_surprise_reads_a_window_of_the_ring(db, make_user, monkeypatch):
+    user = _seed(db, make_user)
+    monkeypatch.setattr(queries, "SURPRISE_WINDOW", 3)
+    for seed in range(20):
+        drawn = {r["id"] for r in queries.yt_surprise(db, user, "any", seed, 3)}
+        assert drawn == set(_ring(seed)[:3]), seed
+
+
+def test_a_wide_gap_on_the_ring_does_not_make_a_favourite(db, make_user):
+    # One video after an empty half of the ring, twenty packed into the other half. Taking the first
+    # video after the seed's point would deal the lonely one about half the time; the window is
+    # shuffled by the seed, so it leads about as often as any other.
+    keys = {"lonely": 0.5} | {f"p{i:02}": 0.51 + i * 0.02 for i in range(20)}
+    user = _seed(db, make_user, keys)
+    leads = [queries.yt_surprise(db, user, "any", s, 1)[0]["id"] for s in range(90)]
+    assert leads.count("lonely") < 15, leads.count("lonely")
+
+
+def test_a_new_video_gets_its_own_key(db, make_user):
+    _seed(db, make_user, {})
+    db.add_all([Video(id=f"n{i}", channel_id="UCring", title="n") for i in range(3)])
+    db.commit()
+    keys = [db.get(Video, f"n{i}").shuffle_key for i in range(3)]
+    assert len(set(keys)) == 3 and all(0 <= k < 1 for k in keys)