feat(spectator): support save replays

.gitignore

@@ -207,3 +207,6 @@ marimo/_lsp/
 __marimo__/
 bancho.py-master/*
 .vscode/settings.json
+
+# runtime file
+replays/

app/config.py

@@ -34,7 +34,7 @@ class Settings:
 
     # SignalR 设置
     SIGNALR_NEGOTIATE_TIMEOUT: int = int(os.getenv("SIGNALR_NEGOTIATE_TIMEOUT", "30"))
-    SIGNALR_PING_INTERVAL: int = int(os.getenv("SIGNALR_PING_INTERVAL", "120"))
+    SIGNALR_PING_INTERVAL: int = int(os.getenv("SIGNALR_PING_INTERVAL", "15"))
 
     # Fetcher 设置
     FETCHER_CLIENT_ID: str = os.getenv("FETCHER_CLIENT_ID", "")

app/models/score.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import Literal, TypedDict
 
 from .mods import API_MODS, APIMod, init_mods
@@ -83,6 +83,43 @@ class HitResult(str, Enum):
         )
 
 
+class HitResultInt(IntEnum):
+    PERFECT = 0
+    GREAT = 1
+    GOOD = 2
+    OK = 3
+    MEH = 4
+    MISS = 5
+
+    LARGE_TICK_HIT = 6
+    SMALL_TICK_HIT = 7
+    SLIDER_TAIL_HIT = 8
+
+    LARGE_BONUS = 9
+    SMALL_BONUS = 10
+
+    LARGE_TICK_MISS = 11
+    SMALL_TICK_MISS = 12
+
+    IGNORE_HIT = 13
+    IGNORE_MISS = 14
+
+    NONE = 15
+    COMBO_BREAK = 16
+
+    LEGACY_COMBO_INCREASE = 99
+
+    def is_hit(self) -> bool:
+        return self not in (
+            HitResultInt.NONE,
+            HitResultInt.IGNORE_MISS,
+            HitResultInt.COMBO_BREAK,
+            HitResultInt.LARGE_TICK_MISS,
+            HitResultInt.SMALL_TICK_MISS,
+            HitResultInt.MISS,
+        )
+
+
 class LeaderboardType(Enum):
     GLOBAL = "global"
     FRIENDS = "friends"
@@ -91,6 +128,7 @@ class LeaderboardType(Enum):
 
 
 ScoreStatistics = dict[HitResult, int]
+ScoreStatisticsInt = dict[HitResultInt, int]
 
 
 class SoloScoreSubmissionInfo(BaseModel):
@@ -128,8 +166,8 @@ class SoloScoreSubmissionInfo(BaseModel):
 class LegacyReplaySoloScoreInfo(TypedDict):
     online_id: int
     mods: list[APIMod]
-    statistics: ScoreStatistics
+    statistics: ScoreStatisticsInt
-    maximum_statistics: ScoreStatistics
+    maximum_statistics: ScoreStatisticsInt
     client_version: str
     rank: Rank
     user_id: int

app/models/spectator_hub.py

@@ -4,18 +4,23 @@ import datetime
 from enum import IntEnum
 from typing import Any
 
+from app.models.beatmap import BeatmapRankStatus
+
 from .score import (
-    HitResult,
+    GameMode,
+    ScoreStatisticsInt,
 )
 from .signalr import MessagePackArrayModel
 
 import msgpack
-from pydantic import Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
 
 
 class APIMod(MessagePackArrayModel):
     acronym: str
-    settings: dict[str, Any] = Field(default_factory=dict)
+    settings: dict[str, Any] | list = Field(
+        default_factory=dict
+    )  # FIXME: with settings
 
 
 class SpectatedUserState(IntEnum):
@@ -32,7 +37,7 @@ class SpectatorState(MessagePackArrayModel):
     ruleset_id: int | None = None  # 0,1,2,3
     mods: list[APIMod] = Field(default_factory=list)
     state: SpectatedUserState
-    maximum_statistics: dict[HitResult, int] = Field(default_factory=dict)
+    maximum_statistics: ScoreStatisticsInt = Field(default_factory=dict)
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, SpectatorState):
@@ -54,11 +59,13 @@ class ScoreProcessorStatistics(MessagePackArrayModel):
 
 
 class FrameHeader(MessagePackArrayModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     total_score: int
     acc: float
     combo: int
     max_combo: int
-    statistics: dict[HitResult, int] = Field(default_factory=dict)
+    statistics: ScoreStatisticsInt = Field(default_factory=dict)
     score_processor_statistics: ScoreProcessorStatistics
     received_time: datetime.datetime
     mods: list[APIMod] = Field(default_factory=list)
@@ -78,6 +85,10 @@ class FrameHeader(MessagePackArrayModel):
             return datetime.datetime.fromisoformat(v)
         raise ValueError(f"Cannot convert {type(v)} to datetime")
 
+    @field_serializer("received_time")
+    def serialize_received_time(self, v: datetime.datetime) -> msgpack.ext.Timestamp:
+        return msgpack.ext.Timestamp.from_datetime(v)
+
 
 class ReplayButtonState(IntEnum):
     NONE = 0
@@ -89,7 +100,7 @@ class ReplayButtonState(IntEnum):
 
 
 class LegacyReplayFrame(MessagePackArrayModel):
-    time: int  # from ReplayFrame,the parent of LegacyReplayFrame
+    time: float  # from ReplayFrame,the parent of LegacyReplayFrame
     x: float | None = None
     y: float | None = None
     button_state: ReplayButtonState
@@ -98,3 +109,37 @@ class LegacyReplayFrame(MessagePackArrayModel):
 class FrameDataBundle(MessagePackArrayModel):
     header: FrameHeader
     frames: list[LegacyReplayFrame]
+
+
+# Use for server
+class APIUser(BaseModel):
+    id: int
+    name: str
+
+
+class ScoreInfo(BaseModel):
+    mods: list[APIMod]
+    user: APIUser
+    ruleset: int
+    maximum_statistics: ScoreStatisticsInt
+    id: int | None = None
+    total_score: int | None = None
+    acc: float | None = None
+    max_combo: int | None = None
+    combo: int | None = None
+    statistics: ScoreStatisticsInt = Field(default_factory=dict)
+
+
+class StoreScore(BaseModel):
+    score_info: ScoreInfo
+    replay_frames: list[LegacyReplayFrame] = Field(default_factory=list)
+
+
+class StoreClientState(BaseModel):
+    state: SpectatorState | None
+    beatmap_status: BeatmapRankStatus
+    checksum: str
+    gamemode: GameMode
+    score_token: int
+    watched_user: set[int]
+    score: StoreScore

app/path.py

@@ -5,3 +5,4 @@ from pathlib import Path
 STATIC_DIR = Path(__file__).parent.parent / "static"
 
 REPLAY_DIR = Path(__file__).parent.parent / "replays"
+REPLAY_DIR.mkdir(exist_ok=True)

app/signalr/hub/hub.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import time
+import traceback
 from typing import Any
 
 from app.config import settings
@@ -92,6 +93,12 @@ class Hub:
     async def send_packet(self, client: Client, type: PacketType, packet: list[Any]):
         await client.send_packet(type, packet)
 
+    async def broadcast_call(self, method: str, *args: Any) -> None:
+        tasks = []
+        for client in self.clients.values():
+            tasks.append(self.call_noblock(client, method, *args))
+        await asyncio.gather(*tasks)
+
     async def _listen_client(self, client: Client) -> None:
         jump = False
         while not jump:
@@ -104,13 +111,12 @@ class Hub:
                 self.tasks.add(task)
                 task.add_done_callback(self.tasks.discard)
             except WebSocketDisconnect as e:
-                if e.code == 1005:
-                    continue
                 print(
                     f"Client {client.connection_id} disconnected: {e.code}, {e.reason}"
                 )
                 jump = True
             except Exception as e:
+                traceback.print_exc()
                 print(f"Error in client {client.connection_id}: {e}")
                 jump = True
         await self.remove_client(client.connection_id)
@@ -139,6 +145,7 @@ class Hub:
                     result = e.message
 
                 except Exception as e:
+                    traceback.print_exc()
                     code = ResultKind.ERROR
                     result = str(e)
 

app/signalr/hub/spectator.py

@@ -1,15 +1,278 @@
 from __future__ import annotations
 
-from app.models.spectator_hub import FrameDataBundle, SpectatorState
+import json
+import lzma
+import struct
+import time
+
+from app.database import Beatmap
+from app.database.score import Score
+from app.database.score_token import ScoreToken
+from app.database.user import User
+from app.dependencies.database import engine
+from app.models.beatmap import BeatmapRankStatus
+from app.models.mods import mods_to_int
+from app.models.score import MODE_TO_INT, LegacyReplaySoloScoreInfo, ScoreStatisticsInt
+from app.models.spectator_hub import (
+    APIUser,
+    FrameDataBundle,
+    LegacyReplayFrame,
+    ScoreInfo,
+    SpectatedUserState,
+    SpectatorState,
+    StoreClientState,
+    StoreScore,
+)
+from app.path import REPLAY_DIR
+from app.utils import unix_timestamp_to_windows
 
 from .hub import Client, Hub
 
+from sqlalchemy.orm import joinedload
+from sqlmodel import select
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+READ_SCORE_TIMEOUT = 30
+REPLAY_LATEST_VER = 30000016
+
+
+def encode_uleb128(num: int) -> bytes | bytearray:
+    if num == 0:
+        return b"\x00"
+
+    ret = bytearray()
+
+    while num != 0:
+        ret.append(num & 0x7F)
+        num >>= 7
+        if num != 0:
+            ret[-1] |= 0x80
+
+    return ret
+
+
+def encode_string(s: str) -> bytes:
+    """Write `s` into bytes (ULEB128 & string)."""
+    if s:
+        encoded = s.encode()
+        ret = b"\x0b" + encode_uleb128(len(encoded)) + encoded
+    else:
+        ret = b"\x00"
+
+    return ret
+
+
+def save_replay(
+    ruleset_id: int,
+    md5: str,
+    username: str,
+    score: Score,
+    statistics: ScoreStatisticsInt,
+    maximum_statistics: ScoreStatisticsInt,
+    frames: list[LegacyReplayFrame],
+) -> None:
+    data = bytearray()
+    data.extend(struct.pack("<bi", ruleset_id, REPLAY_LATEST_VER))
+    data.extend(encode_string(md5))
+    data.extend(encode_string(username))
+    data.extend(encode_string(f"lazer-{username}-{score.started_at.isoformat()}"))
+    data.extend(
+        struct.pack(
+            "<hhhhhhihbi",
+            score.n300,
+            score.n100,
+            score.n50,
+            score.ngeki,
+            score.nkatu,
+            score.nmiss,
+            score.total_score,
+            score.max_combo,
+            score.is_perfect_combo,
+            mods_to_int(score.mods),
+        )
+    )
+    data.extend(encode_string(""))  # hp graph
+    data.extend(
+        struct.pack(
+            "<q",
+            unix_timestamp_to_windows(round(score.started_at.timestamp())),
+        )
+    )
+
+    # write frames
+    # FIXME: cannot play in stable
+    frame_strs = []
+    last_time = 0
+    for frame in frames:
+        frame_strs.append(
+            f"{frame.time - last_time}|{frame.x or 0.0}"
+            f"|{frame.y or 0.0}|{frame.button_state.value}"
+        )
+        last_time = frame.time
+    frame_strs.append("-12345|0|0|0")
+
+    compressed = lzma.compress(
+        ",".join(frame_strs).encode("ascii"), format=lzma.FORMAT_ALONE
+    )
+    data.extend(struct.pack("<i", len(compressed)))
+    data.extend(compressed)
+    data.extend(struct.pack("<q", score.id))
+    assert score.id
+    score_info = LegacyReplaySoloScoreInfo(
+        online_id=score.id,
+        mods=score.mods,
+        statistics=statistics,
+        maximum_statistics=maximum_statistics,
+        client_version="",
+        rank=score.rank,
+        user_id=score.user_id,
+        total_score_without_mods=score.total_score_without_mods,
+    )
+    compressed = lzma.compress(
+        json.dumps(score_info).encode(), format=lzma.FORMAT_ALONE
+    )
+    data.extend(struct.pack("<i", len(compressed)))
+    data.extend(compressed)
+
+    replay_path = REPLAY_DIR / f"lazer-{score.type}-{username}-{score.id}.osr"
+    replay_path.write_bytes(data)
+
 
 class SpectatorHub(Hub):
+    def __init__(self) -> None:
+        super().__init__()
+        self.state: dict[int, StoreClientState] = {}
+
     async def BeginPlaySession(
         self, client: Client, score_token: int, state: SpectatorState
-    ) -> None: ...
+    ) -> None:
+        user_id = int(client.connection_id)
+        previous_state = self.state.get(user_id)
+        if previous_state is not None:
+            return
+        if state.beatmap_id is None or state.ruleset_id is None:
+            return
+        async with AsyncSession(engine) as session:
+            async with session.begin():
+                beatmap = (
+                    await session.exec(
+                        select(Beatmap).where(Beatmap.id == state.beatmap_id)
+                    )
+                ).first()
+                if not beatmap:
+                    return
+                user = (
+                    await session.exec(select(User).where(User.id == user_id))
+                ).first()
+                if not user:
+                    return
+                name = user.name
+                store = StoreClientState(
+                    state=state,
+                    beatmap_status=beatmap.beatmap_status,
+                    checksum=beatmap.checksum,
+                    gamemode=beatmap.mode,
+                    score_token=score_token,
+                    watched_user=set(),
+                    score=StoreScore(
+                        score_info=ScoreInfo(
+                            mods=state.mods,
+                            user=APIUser(id=user_id, name=name),
+                            ruleset=state.ruleset_id,
+                            maximum_statistics=state.maximum_statistics,
+                        )
+                    ),
+                )
+                self.state[user_id] = store
+        await self.broadcast_call("UserBeganPlaying", user_id, state.model_dump())
 
-    async def SendFrameData(
+    async def SendFrameData(self, client: Client, frame_data: FrameDataBundle) -> None:
-        self, client: Client, frame_data: FrameDataBundle
+        user_id = int(client.connection_id)
-    ) -> None: ...
+        state = self.state.get(user_id)
+        if not state:
+            return
+        score = state.score
+        if not score:
+            return
+        score.score_info.acc = frame_data.header.acc
+        score.score_info.combo = frame_data.header.combo
+        score.score_info.max_combo = frame_data.header.max_combo
+        score.score_info.statistics = frame_data.header.statistics
+        score.score_info.total_score = frame_data.header.total_score
+        score.score_info.mods = frame_data.header.mods
+        score.replay_frames.extend(frame_data.frames)
+        await self.broadcast_call(
+            "UserSentFrames",
+            user_id,
+            frame_data.model_dump(),
+        )
+
+    async def EndPlaySession(self, client: Client, state: SpectatorState) -> None:
+        print(f"EndPlaySession -> {client.connection_id} {state.model_dump()!r}")
+        user_id = int(client.connection_id)
+        store = self.state.get(user_id)
+        if not store:
+            return
+        score = store.score
+        if not score or not store.score_token:
+            return
+
+        async def _save_replay():
+            async with AsyncSession(engine) as session:
+                async with session:
+                    start_time = time.time()
+                    score_record = None
+                    while time.time() - start_time < READ_SCORE_TIMEOUT:
+                        sub_query = select(ScoreToken.score_id).where(
+                            ScoreToken.id == store.score_token,
+                        )
+                        result = await session.exec(
+                            select(Score)
+                            .options(joinedload(Score.beatmap))  # pyright: ignore[reportArgumentType]
+                            .where(
+                                Score.id == sub_query,
+                                Score.user_id == user_id,
+                            )
+                        )
+                        score_record = result.first()
+                        if score_record:
+                            break
+                    if not score_record:
+                        return
+                    if not score_record.passed:
+                        return
+                    score_record.has_replay = True
+                    await session.commit()
+                    await session.refresh(score_record)
+                    save_replay(
+                        ruleset_id=MODE_TO_INT[store.gamemode],
+                        md5=store.checksum,
+                        username=store.score.score_info.user.name,
+                        score=score_record,
+                        statistics=score.score_info.statistics,
+                        maximum_statistics=score.score_info.maximum_statistics,
+                        frames=score.replay_frames,
+                    )
+
+        if (
+            (
+                BeatmapRankStatus.PENDING
+                < store.beatmap_status
+                <= BeatmapRankStatus.LOVED
+            )
+            and any(
+                k.is_hit() and v > 0 for k, v in score.score_info.statistics.items()
+            )
+            and state.state != SpectatedUserState.Failed
+        ):
+            # save replay
+            await _save_replay()
+
+        del self.state[user_id]
+        if state.state == SpectatedUserState.Playing:
+            state.state = SpectatedUserState.Quit
+        await self.broadcast_call(
+            "UserEndedPlaying",
+            user_id,
+            state.model_dump(),
+        )

app/signalr/packet.py

@@ -27,7 +27,11 @@ class ResultKind(IntEnum):
 def parse_packet(data: bytes) -> tuple[PacketType, list[Any]]:
     length, offset = decode_varint(data)
     message_data = data[offset : offset + length]
-    unpacked = msgpack.unpackb(message_data, raw=False)
+    # FIXME: custom deserializer for APIMod
+    # https://github.com/ppy/osu/blob/master/osu.Game/Online/API/ModSettingsDictionaryFormatter.cs
+    unpacked = msgpack.unpackb(
+        message_data, raw=False, strict_map_key=False, use_list=True
+    )
     return PacketType(unpacked[0]), unpacked[1:]
 
 

app/signalr/utils.py

@@ -2,24 +2,20 @@ from __future__ import annotations
 
 from collections.abc import Callable
 import inspect
+import sys
 from typing import Any, ForwardRef, cast
 
+# https://github.com/pydantic/pydantic/blob/main/pydantic/v1/typing.py#L61-L75
+if sys.version_info < (3, 12, 4):
 
-# https://github.com/pydantic/pydantic/blob/main/pydantic/v1/typing.py#L56-L66
+    def evaluate_forwardref(type_: ForwardRef, globalns: Any, localns: Any) -> Any:
-def evaluate_forwardref(
+        return cast(Any, type_)._evaluate(globalns, localns, recursive_guard=set())
-    type_: ForwardRef,
+else:
-    globalns: Any,
+
-    localns: Any,
+    def evaluate_forwardref(type_: ForwardRef, globalns: Any, localns: Any) -> Any:
-) -> Any:
+        return cast(Any, type_)._evaluate(
-    # Even though it is the right signature for python 3.9,
+            globalns, localns, type_params=(), recursive_guard=set()
-    # mypy complains with
+        )
-    # `error: Too many arguments for "_evaluate" of
-    # "ForwardRef"` hence the cast...
-    return cast(Any, type_)._evaluate(
-        globalns,
-        localns,
-        set(),
-    )
 
 
 def get_annotation(param: inspect.Parameter, globalns: dict[str, Any]) -> Any:

app/utils.py

@@ -30,7 +30,7 @@ import rosu_pp_py as rosu
 
 def unix_timestamp_to_windows(timestamp: int) -> int:
     """Convert a Unix timestamp to a Windows timestamp."""
-    return timestamp * 10_000 + 11_644_473_600_000_000
+    return (timestamp + 62135596800) * 10_000_000
 
 
 async def convert_db_user_to_api_user(db_user: DBUser, ruleset: str = "osu") -> User: