A seed past a float overflowed the shuffle's start point, and a cursor holding Infinity or a 400-digit key overflowed its int/float read: both answered 500. The routes bound the seed (422), the start point folds it to 32 bits, and the cursor read counts an overflow as a bad cursor (400). A plan test holds that the walk reads ix_videos_shuffle_key with no sort over a realistic scope.
224 lines
8.4 KiB
Python
224 lines
8.4 KiB
Python
"""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.
|
|
"""
|
|
|
|
import base64
|
|
import random
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from sqlalchemy import event, insert, text
|
|
|
|
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_walk_reads_the_ring_index_and_sorts_nothing(db, make_user):
|
|
# The text of the query is not the plan: each arc must be served by ix_videos_shuffle_key in key
|
|
# order (an index on another shape, or none, leaves a Sort of the whole scope). A scope of a few
|
|
# thousand videos, analysed, so the planner weighs it as it does a real one.
|
|
user = _seed(db, make_user)
|
|
rng = random.Random(7)
|
|
db.execute(
|
|
insert(Video),
|
|
[
|
|
{
|
|
"id": f"bulk{i:05}",
|
|
"channel_id": "UCring",
|
|
"title": "b",
|
|
"duration_seconds": 600,
|
|
"published_at": NOW - timedelta(minutes=i),
|
|
"shuffle_key": rng.random(),
|
|
}
|
|
for i in range(4000)
|
|
],
|
|
)
|
|
db.commit()
|
|
db.execute(text("ANALYZE videos"))
|
|
query, _ = feedquery.filtered_query(db, user, **_FEED_BASE)
|
|
captured: list[tuple[str, object]] = []
|
|
|
|
def grab(conn, cursor, statement, params, *args):
|
|
if "FROM videos" in statement and "LIMIT" in statement:
|
|
captured.append((statement, params))
|
|
|
|
engine = db.get_bind()
|
|
event.listen(engine, "before_cursor_execute", grab)
|
|
try:
|
|
seed = _seed_landing_between(0.45, 0.55)
|
|
feedquery.shuffle_walk(db, query, seed, 6) # the first arc, from the seed's point
|
|
feedquery.shuffle_walk(db, query, seed, 6, (1, 0.3, "bulk00001")) # the wrap, resumed
|
|
finally:
|
|
event.remove(engine, "before_cursor_execute", grab)
|
|
assert len(captured) == 2
|
|
conn = db.connection()
|
|
for statement, params in captured:
|
|
plan = "\n".join(r[0] for r in conn.exec_driver_sql("EXPLAIN " + statement, params))
|
|
assert "ix_videos_shuffle_key" in plan, plan
|
|
assert "Sort Key: videos.shuffle_key" not in plan, plan
|
|
|
|
|
|
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_any_seed_lands_on_the_ring():
|
|
# An int too large for a float overflowed the multiplication: a 500 for a hand-typed seed.
|
|
for seed in (10**400, -(10**400), 2**53 + 1, -7):
|
|
assert 0 <= feedquery.shuffle_start(seed) < 1, seed
|
|
|
|
|
|
def test_a_seed_out_of_range_is_refused_not_a_500(db, make_user, as_user):
|
|
client = as_user(_seed(db, make_user))
|
|
huge = "9" * 400
|
|
for url in (
|
|
f"/api/feed?sort=shuffle&seed={huge}",
|
|
f"/api/discover?seed={huge}",
|
|
f"/api/discover/pick?seed={huge}",
|
|
):
|
|
assert client.get(url).status_code == 422, url
|
|
|
|
|
|
def test_a_shuffle_cursor_past_a_float_is_a_bad_cursor(db, make_user, as_user):
|
|
client = as_user(_seed(db, make_user))
|
|
for k in ("[Infinity, 0.5]", f"[0, {'9' * 400}]"):
|
|
raw = '{"k": %s, "id": "k20"}' % k
|
|
cursor = base64.urlsafe_b64encode(raw.encode()).decode()
|
|
assert client.get(f"/api/feed?sort=shuffle&seed=1&cursor={cursor}").status_code == 400, k
|
|
|
|
|
|
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)
|