Merge pull request #10 from osuAkatsuki/user-one-country-on-rankings

Ensure user can only be placed on a single country ranking

main.py

@@ -54,6 +54,19 @@ async def disconnect() -> None:
 STR_TO_INT_MODE = {"std": 0, "taiko": 1, "ctb": 2, "mania": 3}
 
 
+async def get_all_country_rankings_keys(rank_key: str) -> list[str]:
+    country_rankings_keys = []
+
+    cursor = None
+    while cursor != 0:
+        cursor, keys = await redis.scan(cursor=cursor or 0, match=f"{rank_key}:*")
+
+        for key in keys:
+            country_rankings_keys.append(key.decode())
+
+    return country_rankings_keys
+
+
 async def recalc_ranks() -> None:
     print("Recalculating all user ranks")
 
@@ -83,31 +96,43 @@ async def recalc_ranks() -> None:
             )
             users = cast(list[dict[str, Any]], users)
 
+            rank_key = f"ripple:{redis_board}:{mode}"
+
+            country_ranking_keys = await get_all_country_rankings_keys(rank_key)
+
             async with redis.pipeline() as pipe:
                 for user in users:
-                    inactive_days = (start_time - user["latest_pp_awarded"]) / 60 / 60 / 24
+                    inactive_days = (
+                        (start_time - user["latest_pp_awarded"]) / 60 / 60 / 24
+                    )
+                    country = user["country"].lower()
+
+                    user_country_rank_key = f"{rank_key}:{country}"
+
+                    # delete all country rankings for a user besides their current country
+                    # regardless of if they are inactive or not
+                    for key in country_ranking_keys:
+                        # ensure that we never delete the key of their own country
+                        # or it may cause a temporary ranking shift
+                        if key == user_country_rank_key:
+                            continue
+
+                        await pipe.zrem(key, user["id"])
 
                     if inactive_days < 60 and user["privileges"] & 1 and user["pp"] > 0:
                         await pipe.zadd(
-                            f"ripple:{redis_board}:{mode}",
+                            rank_key,
                             {user["id"]: user["pp"]},
                         )
 
-                        country = user["country"].lower()
                         if country != "xx":
                             await pipe.zadd(
-                                f"ripple:{redis_board}:{mode}:{country}",
+                                user_country_rank_key,
                                 {user["id"]: user["pp"]},
                             )
                     else:
-                        await pipe.zrem(f"ripple:{redis_board}:{mode}", user["id"])
-
-                        country = user["country"].lower()
-                        if country != "xx":
-                            await pipe.zrem(
-                                f"ripple:{redis_board}:{mode}:{country}",
-                                user["id"],
-                            )
+                        await pipe.zrem(rank_key, user["id"])
+                        await pipe.zrem(user_country_rank_key, user["id"])
 
                 await pipe.execute()
 
@@ -317,7 +342,7 @@ async def main() -> None:
     print(f"Finished running cron in {time.time() - start_time:.2f} seconds")
 
 
-uvloop.install()
 if __name__ == "__main__":
+    uvloop.install()
     loop = asyncio.get_event_loop()
     loop.run_until_complete(main())