Merge pull request #132 from ghostty-org/async-entity-fetching

perf: improve entity fetching flow

app/components/entity_mentions/cache.py

@@ -29,30 +29,31 @@ def __init__(self, ttr: int) -> None:
         self._ttr = dt.timedelta(seconds=ttr)
         self._cache: dict[CacheKey, tuple[dt.datetime, EntityKind, Entity]] = {}
 
-    def _fetch_entity(self, key: CacheKey) -> None:
+    async def _fetch_entity(self, key: CacheKey) -> None:
         repo_name, entity_id = key
-        repo_path = (config.GITHUB_ORG, config.GITHUB_REPOS[repo_name])
+        signature = (config.GITHUB_ORG, config.GITHUB_REPOS[repo_name], entity_id)
         try:
-            entity = gh.rest.issues.get(*repo_path, entity_id).parsed_data
+            entity = (await gh.rest.issues.async_get(*signature)).parsed_data
             kind = "Issue"
             if entity.pull_request:
-                entity = gh.rest.pulls.get(*repo_path, entity_id).parsed_data
+                entity = (await gh.rest.pulls.async_get(*signature)).parsed_data
                 kind = "Pull Request"
         except RequestFailed:
-            entity = get_discussion(*repo_path, entity_id)
+            entity = await get_discussion(*signature)
             kind = "Discussion"
         self._cache[key] = (dt.datetime.now(), kind, cast(Entity, entity))
 
-    def _refresh(self, key: CacheKey) -> None:
+    async def _refresh(self, key: CacheKey) -> None:
         if key not in self._cache:
-            self._fetch_entity(key)
+            await self._fetch_entity(key)
             return
         timestamp, *_ = self._cache[key]
         if dt.datetime.now() - timestamp >= self._ttr:
-            self._fetch_entity(key)
+            await self._fetch_entity(key)
 
-    def __getitem__(self, key: CacheKey) -> tuple[EntityKind, Entity]:
-        self._refresh(key)
+    async def get(self, repo: RepoName, number: int) -> tuple[EntityKind, Entity]:
+        key = repo, number
+        await self._refresh(key)
         _, kind, entity = self._cache[key]
         return kind, entity
 

app/components/entity_mentions/discussions.py

@@ -19,8 +19,8 @@
 """
 
 
-def get_discussion(org: str, name: str, number: int) -> SimpleNamespace:
-    resp = gh.graphql.request(
+async def get_discussion(org: str, name: str, number: int) -> SimpleNamespace:
+    resp = await gh.graphql.arequest(
         DISCUSSION_QUERY, variables={"number": number, "org": org, "repo": name}
     )
     data = resp["repository"]["discussion"]

app/components/entity_mentions/fmt.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
+import asyncio
 import re
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import discord
 from githubkit.versions.latest.models import Issue, PullRequest
@@ -8,6 +11,9 @@
 
 from .cache import Entity, EntityKind, RepoName, entity_cache
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 GITHUB_URL = "https://github.com"
 ENTITY_REGEX = re.compile(r"(?:\b(web|bot|main))?#(\d{1,6})(?!\.\d)\b")
 ENTITY_TEMPLATE = "**{kind} [#{entity.number}](<{entity.html_url}>):** {entity.title}"
@@ -67,25 +73,29 @@ def _format_mention(entity: Entity, kind: EntityKind) -> str:
     return f"{emoji or ":question:"} {headline}\n{subtext}"
 
 
-def entity_message(message: discord.Message) -> tuple[str, int]:
-    matches = dict.fromkeys(m.groups() for m in ENTITY_REGEX.finditer(message.content))
+async def entity_message(message: discord.Message) -> tuple[str, int]:
+    raw_matches = dict.fromkeys(
+        m.groups() for m in ENTITY_REGEX.finditer(message.content)
+    )
     omitted = 0
-    if len(matches) > 10:
+    if len(raw_matches) > 10:
         # Too many mentions, preventing a DoS
-        omitted = len(matches) - 10
-        matches = list(matches)[:10]
-
-    entities: list[str] = []
-    for repo_name, number_ in matches:
-        number = int(number_)
-        try:
-            kind, entity = entity_cache[cast(RepoName, repo_name or "main"), number]
-        except KeyError:
-            continue
-        if entity.number < 10 and repo_name is None:
-            # Ignore single-digit mentions (likely a false positive)
-            continue
-        entities.append(_format_mention(entity, kind))
+        omitted = len(raw_matches) - 10
+        raw_matches = list(raw_matches)[:10]
+
+    matches: Iterator[tuple[RepoName, int]] = (
+        (cast(RepoName, repo_name or "main"), int(number))
+        for repo_name, number in raw_matches
+        # Ignore single-digit mentions (likely a false positive)
+        if repo_name is not None or int(number) >= 10
+    )
+
+    entities = [
+        _format_mention(entity, kind)
+        for kind, entity in await asyncio.gather(
+            *(entity_cache.get(*m) for m in matches)
+        )
+    ]
 
     if len("\n".join(entities)) > 2000:
         while len("\n".join(entities)) > 1975:  # Accounting for omission note

app/components/entity_mentions/integration.py

@@ -72,7 +72,7 @@ async def reply_with_entities(message: discord.Message) -> None:
         )
         return
 
-    msg_content, entity_count = entity_message(message)
+    msg_content, entity_count = await entity_message(message)
     if not entity_count:
         return
 
@@ -95,8 +95,8 @@ async def on_message_delete(message: discord.Message) -> None:
 async def on_message_edit(before: discord.Message, after: discord.Message) -> None:
     if before.content == after.content:
         return
-    old_entites = entity_message(before)
-    new_entities = entity_message(after)
+    old_entites = await entity_message(before)
+    new_entities = await entity_message(after)
     if old_entites == new_entities:
         # Message changed but mentions are the same
         return