diff --git a/main.py b/main.py index ffb94de..1a25763 100644 --- a/main.py +++ b/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] + " ") 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)