Merge pull request #4071 from Spinnich/fix/random-pick-constant-time

perf(roms): pick a random rom without paging to a random offset

Author: gantoine

backend/endpoints/roms/__init__.py

@@ -898,6 +898,69 @@ def get_rom_identifiers(
     return [r.id for r in db_roms]
 
 
+@protected_route(router.get, "/random", [Scope.ROMS_READ])
+def get_random_rom(
+    request: Request,
+    platform_ids: Annotated[
+        list[int] | None,
+        Query(
+            description=(
+                "Platform internal ids. Multiple values are allowed by repeating the"
+                " parameter, and the pick will match any of the values."
+            ),
+        ),
+    ] = None,
+    collection_id: Annotated[
+        int | None,
+        Query(description="Collection internal id.", ge=1),
+    ] = None,
+    virtual_collection_id: Annotated[
+        str | None,
+        Query(description="Virtual collection internal id."),
+    ] = None,
+    smart_collection_id: Annotated[
+        int | None,
+        Query(description="Smart collection internal id.", ge=1),
+    ] = None,
+) -> SimpleRomSchema | None:
+    """Retrieve one rom picked at random, or null when the scope holds none.
+
+    Sampled on the primary key instead of paged to, so the pick doesn't get
+    slower as the library grows.
+    """
+    perms = get_permissions(request)
+
+    base_query, _ = db_rom_handler.get_roms_query(user_id=request.user.id)
+    query = db_rom_handler.filter_roms(
+        query=base_query,
+        user_id=request.user.id,
+        hidden_platform_ids=perms.hidden_platform_ids,  # type: ignore
+        hidden_rom_ids=perms.hidden_rom_ids,  # type: ignore
+        platform_ids=platform_ids,
+        collection_id=collection_id,
+        virtual_collection_id=virtual_collection_id,
+        smart_collection_id=smart_collection_id,
+        include_related=False,
+    )
+
+    rom_id = db_rom_handler.get_random_rom_id(query=query)
+    if rom_id is None:
+        return None
+
+    rom = db_rom_handler.get_rom_simple(rom_id)
+    if not rom:
+        return None
+
+    # The fetch is by raw id, so it re-checks the row it actually loaded rather
+    # than trusting the filter that chose the id: a rom that moved to a hidden
+    # platform in between was picked under its old one. Reads no database, and
+    # null keeps a hidden rom indistinguishable from an empty scope.
+    if not perms.can_see_rom(rom.id, rom.platform_id):
+        return None
+
+    return SimpleRomSchema.from_orm_with_request(rom, request)
+
+
 @protected_route(
     router.get,
     "/download",

backend/handler/database/roms_handler.py

@@ -2,6 +2,7 @@
 import hashlib
 import json
 import re
+import secrets
 from collections.abc import Iterable, Sequence
 from datetime import datetime
 from typing import Any, NamedTuple
@@ -140,6 +141,11 @@
 # stored lowercase, which keeps the lookup an indexed equality.
 HEX_DIGEST_REGEX = re.compile(r"[0-9a-fA-F]+")
 
+# Primary keys offered per round trip when picking a random rom. Sixteen
+# lands a hit ~99% of the time on a library occupying a quarter of its id
+# range, which is what deletions leave behind on a long-lived instance.
+RANDOM_ID_SAMPLE_SIZE = 16
+
 # CRC32 (8), MD5 and RetroAchievements (32), SHA-1 (40).
 ROM_HASH_COLUMNS_BY_DIGEST_LENGTH: dict[int, tuple[QueryableAttribute, ...]] = {
     8: (Rom.crc_hash,),
@@ -1642,6 +1648,51 @@ def get_rom_count(
             or 0
         )
 
+    @begin_session
+    def get_random_rom_id(
+        self,
+        query: Query,
+        *,
+        session: Session = None,  # type: ignore
+    ) -> int | None:
+        """Pick one rom id at random, uniformly, from a filtered query.
+
+        Two mechanisms, both uniform. First a batch of random primary keys is
+        offered to the query: every id in the table is equally likely to be
+        offered, so any hit is an unbiased pick, and the whole batch costs one
+        index lookup per candidate no matter how large the library is.
+
+        A batch misses when the query matches too little of the id space (a
+        scoped gallery, an id range left full of gaps by deletions). It then
+        falls back to counting the set and taking the row at a random position
+        in it. That reads no rows, only index entries, since the statement
+        selects nothing but the id.
+        """
+        id_query = query.order_by(None).with_only_columns(Rom.id)  # type: ignore
+
+        # Bounds come from the table rather than the filtered set: they only
+        # need to cover it, and MIN/MAX over an untouched primary key are two
+        # index seeks, where the same pair over a joined and filtered set is a
+        # scan of it.
+        lowest, highest = session.execute(
+            select(func.min(Rom.id), func.max(Rom.id))
+        ).one()
+        if lowest is None or highest is None:
+            return None
+
+        span = highest - lowest + 1
+        candidates = {
+            lowest + secrets.randbelow(span) for _ in range(RANDOM_ID_SAMPLE_SIZE)
+        }
+        hits = session.scalars(id_query.where(Rom.id.in_(candidates))).all()
+        if hits:
+            return secrets.choice(hits)
+
+        total = self.get_rom_count(query=query, session=session)
+        if total == 0:
+            return None
+        return session.scalar(id_query.limit(1).offset(secrets.randbelow(total)))
+
     @begin_session
     def get_roms_by_fs_name(
         self,