Fix memory leak issue (#49)

beaverhabits/api.py

@@ -1,4 +1,3 @@
-import logging
 from datetime import datetime
 
 from fastapi import APIRouter, Depends, FastAPI, HTTPException

beaverhabits/configs.py

@@ -19,6 +19,7 @@ class StorageType(Enum):
 
 class Settings(BaseSettings):
     ENV: str = "dev"
+    DEBUG: bool = False
     SENTRY_DSN: str = ""
 
     # UI config

beaverhabits/main.py

@@ -4,6 +4,7 @@
 
 import sentry_sdk
 from fastapi import FastAPI
+from nicegui import ui
 
 from beaverhabits.api import init_api_routes
 from beaverhabits.app import crud
@@ -33,7 +34,7 @@ async def lifespan(_: FastAPI):
     def startup():
         loop = asyncio.get_running_loop()
         loop.set_debug(True)
-        loop.slow_callback_duration = 0.05
+        loop.slow_callback_duration = 0.01
 
 
 @app.get("/health")
@@ -58,6 +59,46 @@ async def user_count():
     sentry_sdk.init(settings.SENTRY_DSN)
 
 
+if settings.DEBUG:
+    import gc
+    import os
+    from collections import Counter
+
+    import psutil
+    from psutil._common import bytes2human
+
+    class MemoryMonitor:
+
+        def __init__(self) -> None:
+            self.last_mem: int = 0
+            self.obj_count: dict[str, int] = {}
+
+        def print_stats(self) -> None:
+            gc.collect()
+            memory = psutil.Process(os.getpid()).memory_info().rss
+            growth = memory - self.last_mem
+            print(bytes2human(memory), end=" ")
+            print(
+                bytes2human(growth, r"%(value)+.1f%(symbol)s") if growth else "",
+                end=" ",
+            )
+            counter: Counter = Counter()
+            for obj in gc.get_objects():
+                counter[type(obj).__name__] += 1
+            for cls, count in counter.items():
+                prev_count = self.obj_count.get(cls, 0)
+                if count != prev_count:
+                    print(f"{cls}={count} ({count - prev_count:+})", end=" ")
+                    self.obj_count[cls] = count
+            self.obj_count = {
+                cls: count for cls, count in self.obj_count.items() if count > 0
+            }
+            print()
+            self.last_mem = memory
+
+    monitor = MemoryMonitor()
+    ui.timer(5.0, monitor.print_stats)
+
 if __name__ == "__main__":
     print(
         'Please start the app with the "uvicorn"'

beaverhabits/storage/dict.py

@@ -2,8 +2,6 @@
 import logging
 from dataclasses import dataclass, field
 
-import pytz
-
 from beaverhabits.storage.storage import CheckedRecord, Habit, HabitList, HabitStatus
 from beaverhabits.utils import generate_short_hash
 

beaverhabits/storage/user_file.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 from typing import Optional
 
+from fastapi_users_db_sqlalchemy import UUID_ID
 from nicegui.storage import PersistentDict
 
 from beaverhabits.app.db import User
@@ -13,17 +14,17 @@
 
 class UserDiskStorage(UserStorage[DictHabitList]):
     def __init__(self):
-        self.user: dict[User, PersistentDict] = {}
+        self.user: dict[UUID_ID, PersistentDict] = {}
 
     def _get_persistent_dict(self, user: User) -> PersistentDict:
-        if user in self.user:
-            return self.user[user]
+        if user.id in self.user:
+            return self.user[user.id]
 
         path = Path(f"{USER_DATA_FOLDER}/{str(user.email)}.json")
         d = PersistentDict(path, encoding="utf-8")
 
         # Cache the persistent dict
-        self.user[user] = d
+        self.user[user.id] = d
 
         return d
 

pyproject.toml

@@ -31,6 +31,7 @@ dev = [
     "beautifulsoup4<5.0.0,>=4.12.3",
     "pytest<9.0.0,>=8.1.1",
     "pip>=24.3.1",
+    "psutil>=6.1.1",
 ]
 
 [project.scripts]