From 960cdf30dac732f18f2be23bc07386c649c43f2d Mon Sep 17 00:00:00 2001
From: Untone <anton.rewin@gmail.com>
Date: Wed, 21 Feb 2024 18:26:18 +0300
Subject: [PATCH] batch-render-follows

---
 main.py             |  4 ++--
 services/follows.py | 30 ++++++++++++++++++++----------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/main.py b/main.py
index 6e0f4f0a..b730a9f2 100644
--- a/main.py
+++ b/main.py
@@ -7,7 +7,7 @@ from ariadne.asgi import GraphQL
 from starlette.applications import Starlette
 from starlette.routing import Route
 
-from services.follows import FollowsCached
+from services.follows import start_cached_follows
 from services.rediscache import redis
 from services.schema import resolvers
 from services.search import search_service
@@ -37,7 +37,7 @@ app = Starlette(
     on_startup=[
         redis.connect,
         ViewedStorage.init,
-        FollowsCached.worker,
+        start_cached_follows,
         search_service.info,
         # start_sentry,
         start,
diff --git a/services/follows.py b/services/follows.py
index 6e463974..76503714 100644
--- a/services/follows.py
+++ b/services/follows.py
@@ -119,17 +119,24 @@ class FollowsCached:
 
     @staticmethod
     async def update_cache():
+        BATCH_SIZE = 30  # Adjust batch size as needed
         with local_session() as session:
-            for author in session.query(Author).all():
-                if isinstance(author, Author):
-                    redis_key = f"user:{author.user}:author"
-                    author_dict = author.dict()
-                    if isinstance(author_dict, dict):
-                        await redis.execute("set", redis_key, json.dumps(author_dict))
-                    follows = await get_author_follows(None, None, user=author.user)
-                    if isinstance(follows, dict):
-                        redis_key = f"user:{author.user}:follows"
-                        await redis.execute("set", redis_key, json.dumps(follows))
+            authors = session.query(Author).all()
+            total_authors = len(authors)
+            for i in range(0, total_authors, BATCH_SIZE):
+                batch_authors = authors[i:i+BATCH_SIZE]
+                await asyncio.gather(*[FollowsCached.update_author_cache(author) for author in batch_authors])
+
+    @staticmethod
+    async def update_author_cache(author):
+        redis_key = f"user:{author.user}:author"
+        author_dict = author.dict()
+        if isinstance(author_dict, dict):
+            await redis.execute("set", redis_key, json.dumps(author_dict))
+        follows = await get_author_follows(None, None, user=author.user)
+        if isinstance(follows, dict):
+            redis_key = f"user:{author.user}:follows"
+            await redis.execute("set", redis_key, json.dumps(follows))
 
     @staticmethod
     async def worker():
@@ -146,3 +153,6 @@ class FollowsCached:
                 break
             except Exception as exc:
                 logger.error(exc)
+
+async def start_cached_follows():
+    await FollowsCached.worker()