Implement pack command timeout (#26)

Author: vberlier

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+- Added `pack` command timeout and updated `beet` and `lectern`
+
 ## [0.7.0] - 2021-04-23
 
 ### Added

commanderbot_ext/ext/pack/pack_cog.py

@@ -1,11 +1,7 @@
 import asyncio
 import io
-from contextlib import asynccontextmanager
 from logging import Logger, getLogger
-from typing import AsyncIterator
 
-from beet import FormattedPipelineException
-from beet.toolchain.utils import format_exc
 from discord import File, Message
 from discord.ext.commands import Bot, Cog, Context, command
 
@@ -17,6 +13,7 @@ def __init__(self, bot: Bot, **options):
         self.bot: Bot = bot
         self.log: Logger = getLogger(self.qualified_name)
         self.project_config = options
+        self.build_timeout = options.pop("timeout", 5)
 
     @command(name="pack")
     async def cmd_pack(self, ctx: Context):
@@ -28,40 +25,28 @@ async def cmd_pack(self, ctx: Context):
         author = message.author.display_name
         message_content = message.content.split("\n", 1)[-1]
 
-        self.log.info("%s: Running build for %s.", message.id, author)
+        self.log.info("%s - Running build for %s.", message.id, author)
 
         loop = asyncio.get_running_loop()
 
-        async with self.error_handler(ctx):
-            attachments = await loop.run_in_executor(
-                None, generate_packs, author, self.project_config, message_content
-            )
-
-            if attachments:
-                files = [
-                    File(io.BytesIO(data), filename=filename)
-                    for filename, data in attachments.items()
-                ]
-                await ctx.send(files=files)
-            else:
-                await message.add_reaction("🤔")
-
-        self.log.info("%s: Done.", message.id)
-
-    @asynccontextmanager
-    async def error_handler(self, ctx: Context) -> AsyncIterator[None]:
-        try:
-            yield
-        except FormattedPipelineException as exc:
-            message = exc.message
-            exception = exc.__cause__ if exc.format_cause else None
-        except Exception as exc:
-            message = "An unhandled exception occurred. This could be a bug."
-            exception = exc
+        build_output, attachments = await loop.run_in_executor(
+            None,
+            generate_packs,
+            self.project_config,
+            self.build_timeout,
+            author,
+            message_content,
+        )
+
+        content = f"```{joined}```" if (joined := "\n\n".join(build_output)) else ""
+        files = [
+            File(io.BytesIO(data), filename=filename)
+            for filename, data in attachments.items()
+        ]
+
+        if content or files:
+            await ctx.send(content, files=files)
         else:
-            return
-
-        if exception:
-            message += f"\n\n{format_exc(exception)}"
+            await message.add_reaction("🤔")
 
-        await ctx.send(f"```{message}```")
+        self.log.info("%s - Done.", message.id)

commanderbot_ext/ext/pack/pack_generate.py

@@ -1,50 +1,98 @@
 import io
 import os
-from typing import Dict
+from multiprocessing import Process, Queue
+from queue import Empty
+from typing import Dict, List, Tuple
 from zipfile import ZipFile
 
-from beet import Context, ProjectConfig, config_error_handler, run_beet
+from beet import (
+    Context,
+    FormattedPipelineException,
+    ProjectConfig,
+    config_error_handler,
+    run_beet,
+)
 from beet.core.utils import JsonDict
+from beet.toolchain.utils import format_exc
 from lectern import Document
 
+BuildResult = Tuple[List[str], Dict[str, bytes]]
+
 
 def generate_packs(
+    project_config: JsonDict,
+    build_timeout: float,
+    project_name: str,
+    message_content: str,
+) -> BuildResult:
+    q: "Queue[BuildResult]" = Queue()
+
+    p = Process(target=worker, args=(q, project_name, project_config, message_content))
+    p.start()
+    p.join(timeout=build_timeout)
+
+    try:
+        return q.get_nowait()
+    except Empty:
+        p.kill()
+        return [
+            f"Timeout exceeded. The message took more than {build_timeout} seconds to process."
+        ], {}
+
+
+def worker(
+    q: "Queue[BuildResult]",
     project_name: str,
     project_config: JsonDict,
     message_content: str,
-) -> Dict[str, bytes]:
+):
     project_directory = os.getcwd()
-    packs: Dict[str, bytes] = {}
 
-    message_config = {
+    base_config = {
         "name": project_name,
         "pipeline": [__name__],
         "meta": {
             "source": message_content,
+            "build_output": [],
+            "build_attachments": {},
         },
     }
 
-    with config_error_handler():
-        config = (
-            ProjectConfig(**project_config)
-            .resolve(project_directory)
-            .with_defaults(ProjectConfig(**message_config).resolve(project_directory))
-        )
+    try:
+        with config_error_handler():
+            config = (
+                ProjectConfig(**project_config)
+                .resolve(project_directory)
+                .with_defaults(ProjectConfig(**base_config).resolve(project_directory))
+            )
+
+        with run_beet(config) as ctx:
+            build_output = ctx.meta["build_output"]
+            attachments = ctx.meta["build_attachments"]
+
+            for pack in ctx.packs:
+                if pack:
+                    fp = io.BytesIO()
 
-    with run_beet(config) as ctx:
-        for pack in ctx.packs:
-            if not pack:
-                continue
+                    with ZipFile(fp, mode="w") as output:
+                        pack.dump(output)
+                        output.writestr("source.md", message_content)
 
-            fp = io.BytesIO()
+                    attachments[f"{pack.name}.zip"] = fp.getvalue()
 
-            with ZipFile(fp, mode="w") as output:
-                pack.dump(output)
-                output.writestr("source.md", message_content)
+            q.put((build_output, attachments))
+            return
+    except FormattedPipelineException as exc:
+        build_output = [exc.message]
+        exception = exc.__cause__ if exc.format_cause else None
+    except Exception as exc:
+        build_output = ["An unhandled exception occurred. This could be a bug."]
+        exception = exc
 
-            packs[f"{pack.name}.zip"] = fp.getvalue()
+    if exception:
+        build_output.append(format_exc(exception))
 
-    return packs
+    q.put((build_output, {}))
 
 
 def beet_default(ctx: Context):