refactor(hosts): legacy storage

Author: M03ED

app/core/hosts.py

@@ -1,12 +1,15 @@
+from asyncio import Lock
+from copy import deepcopy
+
 from sqlalchemy.ext.asyncio import AsyncSession
+from aiocache import cached
 
 from app import on_startup
 from app.core.manager import core_manager
 from app.db import GetDB
-from app.db.crud.host import get_host_by_id, get_hosts, get_or_create_inbound
+from app.db.crud.host import get_host_by_id, get_hosts, upsert_inbounds
 from app.db.models import ProxyHost, ProxyHostSecurity
-from app.models.host import MuxSettings, TransportSettings
-from app.utils.store import DictStorage
+from app.models.host import MuxSettings, TransportSettings, BaseHost
 
 
 def _prepare_host_data(host: ProxyHost) -> dict:
@@ -40,18 +43,23 @@ def _prepare_host_data(host: ProxyHost) -> dict:
     }
 
 
-@DictStorage
-async def hosts(storage: dict, db: AsyncSession):
-    inbounds_list = await core_manager.get_inbounds()
-    for tag in inbounds_list:
-        await get_or_create_inbound(db, tag)
+class HostManager:
+    def __init__(self):
+        self._hosts = {}
+        self._lock = Lock()
+
+    async def setup(self, db: AsyncSession):
+        db_hosts = await get_hosts(db)
+        await self.add_hosts(db, db_hosts)
 
-    storage.clear()
-    db_hosts = await get_hosts(db)
+    async def _reset_cache(self):
+        await self.get_hosts.cache.clear()
 
-    for host in db_hosts:
+    @staticmethod
+    async def _prepare_host_entry(db: AsyncSession, host: BaseHost, inbounds_list: list[str]) -> tuple[int, dict] | None:
         if host.is_disabled or (host.inbound_tag not in inbounds_list):
-            continue
+            return None
+
         downstream = None
         if (
             host.transport_settings
@@ -67,10 +75,47 @@ async def hosts(storage: dict, db: AsyncSession):
         else:
             host_data["downloadSettings"] = None
 
-        storage[host.id] = host_data
+        return (host.id, host_data)
+
+    async def add_host(self, db: AsyncSession, host: BaseHost):
+        await self.add_hosts(db, [host])
+
+    async def add_hosts(self, db: AsyncSession, hosts: list[BaseHost]):
+        serialized_hosts = [BaseHost.model_validate(host) for host in hosts]
+        inbounds_list = await core_manager.get_inbounds()
+        await upsert_inbounds(db, inbounds_list)
+        await db.commit()
+
+        prepared_hosts = []
+        for host in serialized_hosts:
+            result = await self._prepare_host_entry(db, host, inbounds_list)
+            if result:
+                prepared_hosts.append(result)
+
+        # Acquire lock only for updating the dict and cache
+        async with self._lock:
+            for host_id, host_data in prepared_hosts:
+                self._hosts[host_id] = host_data
+            await self._reset_cache()
+
+    async def remove_host(self, id: int):
+        async with self._lock:
+            self._hosts.pop(id, None)
+            await self._reset_cache()
+
+    async def get_host(self, id: int) -> dict | None:
+        async with self._lock:
+            return deepcopy(self._hosts.get(id))
+
+    @cached()
+    async def get_hosts(self) -> dict[int, dict]:
+        async with self._lock:
+            return deepcopy(self._hosts)
+
 
+host_manager: HostManager = HostManager()
 
 @on_startup
 async def initialize_hosts():
     async with GetDB() as db:
-        await hosts.update(db)
+        await host_manager.setup(db)

app/db/crud/group.py

@@ -4,14 +4,16 @@
 from app.db.models import ProxyInbound, Group
 from app.models.group import GroupCreate, GroupModify
 
-from .host import get_or_create_inbound
+from .host import upsert_inbounds
 
 
 async def get_inbounds_by_tags(db: AsyncSession, tags: list[str]) -> list[ProxyInbound]:
     """
-    Retrieves inbounds by their tags.
+    Retrieves or creates inbounds by their tags using efficient bulk upsert.
     """
-    return [(await get_or_create_inbound(db, tag)) for tag in tags]
+    inbounds_map = await upsert_inbounds(db, tags)
+    # Return in the same order as input tags
+    return [inbounds_map[tag] for tag in tags]
 
 
 async def load_group_attrs(group: Group):
@@ -124,13 +126,14 @@ async def modify_group(db: AsyncSession, db_group: Group, modified_group: GroupM
         Group: The updated Group object.
     """
 
+    if modified_group.inbound_tags:
+        inbounds = await get_inbounds_by_tags(db, modified_group.inbound_tags)
+        db_group.inbounds = inbounds
     if db_group.name != modified_group.name:
         db_group.name = modified_group.name
     if modified_group.is_disabled is not None:
         db_group.is_disabled = modified_group.is_disabled
-    if modified_group.inbound_tags:
-        inbounds = await get_inbounds_by_tags(db, modified_group.inbound_tags)
-        db_group.inbounds = inbounds
+
     await db.commit()
     await db.refresh(db_group)
     await load_group_attrs(db_group)

app/db/crud/host.py

@@ -2,35 +2,79 @@
 from enum import Enum
 from typing import List, Optional
 
-from sqlalchemy import select
+from sqlalchemy import bindparam, select
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.mysql import insert as mysql_insert
+from sqlalchemy import insert
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.db.models import ProxyHost, ProxyInbound
 from app.models.host import CreateHost
 
 
+async def upsert_inbounds(db: AsyncSession, inbound_tags: list[str]) -> dict[str, ProxyInbound]:
+    """
+    Efficiently upserts multiple proxy inbounds and returns them.
+    Uses INSERT ... ON CONFLICT DO NOTHING pattern to avoid unnecessary SELECT queries.
+
+    Args:
+        db (AsyncSession): Database session.
+        inbound_tags (List[str]): List of inbound tags to upsert.
+
+    Returns:
+        dict[str, ProxyInbound]: Mapping of tag to ProxyInbound object.
+
+    Note:
+        This function does not commit the transaction. The caller is responsible for committing.
+    """
+    if not inbound_tags:
+        return {}
+
+    # Remove duplicates while preserving order
+    unique_tags = list(dict.fromkeys(inbound_tags))
+
+    dialect = db.bind.dialect.name
+
+    # Build upsert statement based on dialect
+    if dialect == "postgresql":
+        stmt = pg_insert(ProxyInbound).values(tag=bindparam("tag"))
+        stmt = stmt.on_conflict_do_nothing(index_elements=["tag"])
+    elif dialect == "mysql":
+        stmt = mysql_insert(ProxyInbound).values(tag=bindparam("tag"))
+        stmt = stmt.on_duplicate_key_update(tag=ProxyInbound.tag)
+    else:  # SQLite
+        stmt = insert(ProxyInbound).values(tag=bindparam("tag")).prefix_with("OR IGNORE")
+
+    # Execute upsert for all tags
+    params = [{"tag": tag} for tag in unique_tags]
+    await db.execute(stmt, params)
+    await db.flush()  # Flush to make inserted rows visible in this transaction
+
+    # Now select all the inbounds we just upserted
+    select_stmt = select(ProxyInbound).where(ProxyInbound.tag.in_(unique_tags))
+    result = await db.execute(select_stmt)
+    inbounds = result.scalars().all()
+
+    # Return as a mapping
+    return {inbound.tag: inbound for inbound in inbounds}
+
+
 async def get_or_create_inbound(db: AsyncSession, inbound_tag: str) -> ProxyInbound:
     """
     Retrieves or creates a proxy inbound based on the given tag.
 
+    Note: This function is deprecated. Use upsert_inbounds() for better performance,
+    especially when dealing with multiple inbounds.
+
     Args:
         db (AsyncSession): Database session.
         inbound_tag (str): The tag of the inbound.
 
     Returns:
         ProxyInbound: The retrieved or newly created proxy inbound.
     """
-    stmt = select(ProxyInbound).where(ProxyInbound.tag == inbound_tag)
-    result = await db.execute(stmt)
-    inbound = result.scalar_one_or_none()
-
-    if not inbound:
-        inbound = ProxyInbound(tag=inbound_tag)
-        db.add(inbound)
-        await db.commit()
-        await db.refresh(inbound)
-
-    return inbound
+    result = await upsert_inbounds(db, [inbound_tag])
+    return result[inbound_tag]
 
 
 async def get_inbounds_not_in_tags(db: AsyncSession, excluded_tags: List[str]) -> List[ProxyInbound]:

app/node/__init__.py

@@ -28,7 +28,7 @@ async def update_node(self, node: Node) -> PasarGuardNode:
                 except Exception:
                     pass
                 finally:
-                    del self._nodes[node.id]
+                    self._nodes.pop(node.id, None)
 
             new_node = create_node(
                 connection=type_map[node.connection_type],
@@ -56,7 +56,7 @@ async def remove_node(self, id: int) -> None:
                 except Exception:
                     pass
                 finally:
-                    del self._nodes[id]
+                    self._nodes.pop(id, None)
 
     async def get_node(self, id: int) -> PasarGuardNode | None:
         async with self._lock.reader_lock:

app/operation/core.py

@@ -7,7 +7,7 @@
 from app.core.manager import core_manager
 from app.operation import BaseOperation
 from app import notification
-from app.core.hosts import hosts as hosts_storage
+from app.core.hosts import host_manager
 from app.utils.logger import get_logger
 
 
@@ -28,7 +28,7 @@ async def create_core(self, db: AsyncSession, new_core: CoreCreate, admin: Admin
         core = CoreResponse.model_validate(db_core)
         asyncio.create_task(notification.create_core(core, admin.username))
 
-        await hosts_storage.update(db)
+        await host_manager.setup(db)
 
         return core
 
@@ -55,7 +55,7 @@ async def modify_core(
         core = CoreResponse.model_validate(db_core)
         asyncio.create_task(notification.modify_core(core, admin.username))
 
-        await hosts_storage.update(db)
+        await host_manager.setup(db)
 
         return core
 
@@ -72,4 +72,4 @@ async def delete_core(self, db: AsyncSession, core_id: int, admin: AdminDetails)
 
         logger.info(f'core config "{db_core.name}" deleted by admin "{admin.username}"')
 
-        await hosts_storage.update(db)
+        await host_manager.setup(db)

app/operation/host.py

@@ -6,7 +6,7 @@
 from app.models.admin import AdminDetails
 from app.operation import BaseOperation
 from app.db.crud.host import create_host, get_host_by_id, remove_host, get_hosts, modify_host
-from app.core.hosts import hosts as hosts_storage
+from app.core.hosts import host_manager
 from app.utils.logger import get_logger
 
 from app import notification
@@ -49,7 +49,7 @@ async def create_host(self, db: AsyncSession, new_host: CreateHost, admin: Admin
         host = BaseHost.model_validate(db_host)
         asyncio.create_task(notification.create_host(host, admin.username))
 
-        await hosts_storage.update(db)
+        await host_manager.add_host(db, db_host)
 
         return host
 
@@ -70,7 +70,7 @@ async def modify_host(
         host = BaseHost.model_validate(db_host)
         asyncio.create_task(notification.modify_host(host, admin.username))
 
-        await hosts_storage.update(db)
+        await host_manager.add_host(db, db_host)
 
         return host
 
@@ -83,7 +83,7 @@ async def remove_host(self, db: AsyncSession, host_id: int, admin: AdminDetails)
 
         asyncio.create_task(notification.remove_host(host, admin.username))
 
-        await hosts_storage.update(db)
+        await host_manager.remove_host(host.id)
 
     async def modify_hosts(
         self, db: AsyncSession, modified_hosts: list[CreateHost], admin: AdminDetails
@@ -100,7 +100,7 @@ async def modify_hosts(
             else:
                 await modify_host(db, old_host, host)
 
-        await hosts_storage.update(db)
+        await host_manager.add_hosts(db, modified_hosts)
 
         logger.info(f'Host\'s has been modified by admin "{admin.username}"')
 

app/subscription/share.py

@@ -6,7 +6,7 @@
 
 from jdatetime import date as jd
 
-from app.core.hosts import hosts as hosts_storage
+from app.core.hosts import host_manager
 from app.core.manager import core_manager
 from app.db.models import UserStatus
 from app.models.user import UsersResponseWithInbounds
@@ -271,7 +271,7 @@ async def process_inbounds_and_tags(
     reverse=False,
 ) -> list | str:
     proxy_settings = user.proxy_settings.dict()
-    for host in await filter_hosts(hosts_storage.values(), user.status):
+    for host in await filter_hosts((await host_manager.get_hosts()).values(), user.status):
         host_data = await process_host(host, format_variables, user.inbounds, proxy_settings, conf)
         if not host_data:
             continue