Add caching

Implement idea from
https://raw.githubusercontent.com/nakidai/cptc/d69da87a74ea1fc18b73ec7660dad7d590e28ddd/TODO

4 files changed, 126 insertions, 9 deletions

diff --git a/cache/index.json b/cache/index.json
new file mode 100644
index 0000000..9e26dfe
--- /dev/null
+++ b/cache/index.json
@@ -0,0 +1 @@
+{}
\ No newline at end of file
diff --git a/src/petthecord/bot.py b/src/petthecord/bot.py
index 0725943..9f8e975 100644
--- a/src/petthecord/bot.py
+++ b/src/petthecord/bot.py
@@ -30,19 +30,42 @@ class PatTheCordCog(commands.Cog):
 
 
 class Bot(commands.Bot):
-    def __init__(self, host: str = "127.0.0.1", port: int = 8080) -> None:
+    def __init__(
+        self,
+
+        host: str = "127.0.0.1",
+        port: int = 8080,
+        caching: bool = True,
+        cache_path: str = "/var/cache/petthecord",
+        cache_lifetime: int = 86400,
+        cache_gc_delay: int = 14400,
+    ) -> None:
         super().__init__(
             command_prefix="!",
             intents=Intents.default()
         )
         self._host = host
         self._port = port
+        self._caching = caching
+        self._cache_path = cache_path
+        self._cache_lifetime = cache_lifetime
+        self._cache_gc_delay = cache_gc_delay
 
     async def on_ready(self) -> None:
         await self.add_cog(PatTheCordCog(self))
         await self.tree.sync()
 
-        runner = AppRunner(Server(self))
+        server = Server(
+            self,
+            self._caching,
+            self._cache_path,
+            self._cache_lifetime,
+            self._cache_gc_delay
+        )
+        runner = AppRunner(server)
         await runner.setup()
         site = TCPSite(runner, self._host, self._port)
         await site.start()
+
+        if self._caching:
+            await server.clean_cache()
diff --git a/src/petthecord/main.py b/src/petthecord/main.py
index b6f3450..2e2047a 100644
--- a/src/petthecord/main.py
+++ b/src/petthecord/main.py
@@ -23,9 +23,42 @@ def main() -> None:
         metavar="HOST",
         help="Bind IP"
     )
+    parser.add_argument(
+        "-d", "--cache-dir",
+        default="/var/cache/petthecord",
+        metavar="PATH",
+        help="Directory for cache storing"
+    )
+    parser.add_argument(
+        "-n", "--no-cache",
+        action="store_true",
+        default=False,
+        help="Turn off the cache"
+    )
+    parser.add_argument(
+        "-l", "--cache-lifetime",
+        default=86400,
+        type=int,
+        metavar="TIME",
+        help="Lifetime of cached avatar in seconds"
+    )
+    parser.add_argument(
+        "-s", "--cache-gc-delay",
+        default=14400,
+        type=int,
+        metavar="TIME",
+        help="Delay between cache's garbage collector runs in seconds"
+    )
     args = parser.parse_args()
 
-    bot = Bot(args.host, args.port)
+    bot = Bot(
+        args.host,
+        args.port,
+        not args.no_cache,
+        args.cache_dir,
+        args.cache_lifetime,
+        args.cache_gc_delay
+    )
     if (token := getenv("PETTHECORD_TOKEN")) is not None:
         bot.run(token)
     elif (token_path := getenv("PETTHECORD_TOKEN_FILE")) is not None:
diff --git a/src/petthecord/server.py b/src/petthecord/server.py
index 35cbca7..7ffe812 100644
--- a/src/petthecord/server.py
+++ b/src/petthecord/server.py
@@ -1,4 +1,10 @@
+from asyncio import sleep
+from json import load, dumps
 from io import BytesIO
+from os import listdir, remove
+from os.path import getmtime
+from pathlib import Path
+from time import time
 from typing import NoReturn
 
 from aiohttp.web import Application, StreamResponse, get, HTTPFound, Request, Response
@@ -8,10 +14,31 @@ from petpetgif import petpet
 
 
 class Server(Application):
-    def __init__(self, client: commands.Bot) -> None:
+    def __init__(
+        self,
+
+        client: commands.Bot,
+        caching: bool = True,
+        cache_path: str = "/var/cache/petthecord",
+        cache_lifetime: int = 86400,
+        cache_gc_delay: int = 14400,
+    ) -> None:
         self.client = client
+        self.caching = caching
+        self.cache_path = Path(cache_path)
+        self.cache_lifetime = cache_lifetime
+        self.cache_gc_delay = cache_gc_delay
         super().__init__()
 
+        if self.caching:
+            index_path = self.cache_path / "index.json"
+            if not index_path.exists():
+                with open(index_path, "w") as f:
+                    f.write("{}")
+
+            with open(index_path, "r") as f:
+                self.cache = load(f)
+
         self.add_routes(
             [
                 get("/{uid}", self.petpet),
@@ -23,6 +50,7 @@ class Server(Application):
         raise HTTPFound("https://github.com/nakidai/petthecord")
 
     async def petpet(self, request: Request) -> StreamResponse:
+        print(request)
         try:
             uid = int(request.match_info["uid"][:request.match_info["uid"].find('.')])
         except ValueError:
@@ -38,9 +66,41 @@ class Server(Application):
         if user.avatar is None:
             return Response(status=404)
 
-        image = await user.avatar.read()
-        dest = BytesIO()
-        petpet.make(BytesIO(image), dest)
-        dest.seek(0)
+        avatar_path = str(self.cache_path / f"{user.id}_{user.avatar.key}.gif")
+        if self.caching:
+            if (path := self.cache.get(user.id)) != avatar_path:
+                if path:
+                    remove(path)
+                self.cache[user.id] = avatar_path
+                with open(self.cache_path / "index.json", "w") as f:
+                    f.write(dumps(self.cache))
+
+            if not Path(avatar_path).exists():
+                with open(avatar_path, "wb") as f:
+                    image = await user.avatar.read()
+                    petpet.make(BytesIO(image), f)
+
+            Path(avatar_path).touch()
+
+            with open(avatar_path, "rb") as f:
+                return Response(body=f.read(), content_type="image/png")
+        else:
+            with BytesIO() as f:
+                image = await user.avatar.read()
+                petpet.make(BytesIO(image), f)
+
+                f.seek(0)
+
+                return Response(body=f.read(), content_type="image/png")
+
+    async def clean_cache(self) -> None:
+        for filename in listdir(self.cache_path):
+            path = (self.cache_path / filename)
+            if path.is_file() and filename != "index.json":
+                if (time() - getmtime(path) > self.cache_lifetime):
+                    del self.cache[filename.split('_')[0]]
+                    remove(path)
+        with open(self.cache_path / "index.json", "w") as f:
+            f.write(dumps(self.cache))
 
-        return Response(body=dest.getvalue(), content_type="image/png")
+        await sleep(self.cache_gc_delay)