add storage safeguard. add locks to prevent race conditions.

main.py

@@ -1,11 +1,18 @@
-from asyncio import create_task, shield, to_thread, wait_for
+from asyncio import Condition, Lock, create_task, shield, to_thread, wait_for
 from sys import argv
 
+import requests
 from discord import Game, Intents, Interaction, app_commands
 from discord.ext import commands
 
 from PICable import PICable
 
+MAX_STORAGE_KB = 10 * 1024 * 1024  # 10GB
+current_storage_kb = 0
+storage_condition = Condition()
+current_repositories = set()
+current_repositories_lock = Lock()
+
 if len(argv) != 3:
     print("Usage:\n\t" + argv[0] + " <discord_token> <github_token>")
     exit(1)
@@ -16,8 +23,6 @@ intents = Intents.default()
 intents.message_content = True
 client = commands.Bot(command_prefix="/", intents=intents)
 
-current_repositories = set()
-
 
 @client.event
 async def on_ready():
@@ -36,21 +41,49 @@ async def ping(interaction: Interaction):
     await interaction.response.send_message(f"Pong! Latency: {latency:.2f}ms")
 
 
+async def handle_picable(owner: str, repository: str, repository_size_kb: int, github_token: str):
+    global current_storage_kb
+    async with storage_condition:
+        await storage_condition.wait_for(lambda: current_storage_kb + repository_size_kb <= MAX_STORAGE_KB)
+        current_storage_kb += repository_size_kb
+
+    print(f"Storage is {current_storage_kb / MAX_STORAGE_KB * 100:.2f}% full.")
+    return await to_thread(PICable, owner, repository, github_token)
+
+
 @client.tree.command(name="picable", description="Check if a Github repository is eligible for PIC.")
 @app_commands.describe(owner="The owner of the repository", repository="The name of the repository")
 async def picable(interaction: Interaction, owner: str, repository: str):
     await interaction.response.defer(thinking=True)
 
     repository_full_name = f"{owner}/{repository}"
-    if repository_full_name in current_repositories:
+    async with current_repositories_lock:
+        if repository_full_name in current_repositories:
+            await interaction.followup.send(
+                f"The repository {repository_full_name} is already being analyzed. Please wait for the current analysis to complete. :warning:"
+            )
+            return
+
+        current_repositories.add(repository_full_name)
+
+    url = f"https://api.github.com/repos/{owner}/{repository}"
+    headers = {"Authorization": f"token {github_token}"}
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    repo_info = response.json()
+    repository_size_kb = repo_info["size"]
+    print(f"Repository size: {repository_size_kb / 1024} MB")
+
+    if repository_size_kb > MAX_STORAGE_KB:
         await interaction.followup.send(
-            f"The repository {repository_full_name} is already being analyzed. Please wait for the current analysis to complete. :warning:"
+            f"The repository {repository_full_name} is too large to analyze. Please try a smaller repository. :warning:"
         )
+        current_repositories.remove(repository_full_name)
         return
 
-    current_repositories.add(repository_full_name)
+    global current_storage_kb
     try:
-        result_future = create_task(to_thread(PICable, owner, repository, github_token))
+        result_future = create_task(handle_picable(owner, repository, repository_size_kb, github_token))
         try:
             result = await wait_for(shield(result_future), timeout=5)
             await interaction.followup.send(result)
@@ -59,12 +92,16 @@ async def picable(interaction: Interaction, owner: str, repository: str):
                 f"The analysis for {repository_full_name} is still taking place. The results will be posted here once the analysis is complete. :clock4:"
             )
             result = await result_future
-            await interaction.channel.send(result)
+            await interaction.channel.send(f"{interaction.user.mention}\n{result}")
     except Exception as e:
         print(f"Error processing PICable check: {e}")
         await interaction.channel.send("An error occurred while processing your request. Please try again later.")
     finally:
-        current_repositories.remove(repository_full_name)
+        async with current_repositories_lock:
+            current_repositories.remove(repository_full_name)
+        async with storage_condition:
+            current_storage_kb -= repository_size_kb
+            storage_condition.notify_all()
 
 
 client.run(discord_token)