diff --git a/server/repl.py b/server/repl.py index 7d0c4ca..9a681b6 100644 --- a/server/repl.py +++ b/server/repl.py @@ -32,6 +32,12 @@ log_lock = asyncio.Lock() +class ReplOOMError(ReplError): + """Raised when the REPL exceeds its memory limit.""" + + pass + + async def log_snippet(uuid: UUID, snippet_id: str, code: str) -> None: if settings.environment == Environment.prod: header = f"[{uuid.hex[:8]}] Running snippet {snippet_id}:" @@ -75,7 +81,9 @@ def __init__( self.header_cmd_response: ReplResponse | None = None self.proc: Process | None = None - self.error_file = tempfile.TemporaryFile("w+") + self.error_file = tempfile.TemporaryFile( + "w+", encoding="utf-8", errors="ignore" + ) self.max_memory_bytes = max_repl_mem * 1024 * 1024 self.max_repl_uses = max_repl_uses @@ -146,6 +154,9 @@ def _preexec() -> None: os.setsid() + self.error_file.seek(0) + self.error_file.truncate() + self.proc = await asyncio.create_subprocess_exec( "lake", "env", @@ -154,7 +165,7 @@ def _preexec() -> None: env=os.environ, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, + stderr=self.error_file, preexec_fn=_preexec, ) @@ -231,10 +242,13 @@ async def send_timeout( ) raise e except LeanError as e: - logger.exception("Lean REPL error: %s", e) + logger.exception("Lean REPL error: {}", e) raise e + except ReplOOMError as e: + logger.error("[{}] REPL OOM: {}", self.uuid.hex[:8], e) + return ReplResponse(id=snippet.id, error=str(e), time=0.0) except ReplError as e: - logger.exception("REPL error: %s", e) + logger.exception("REPL error: {}", e) raise e return ReplResponse( @@ -287,18 +301,26 @@ async def send( logger.error("Broken pipe while writing to REPL stdin") raise LeanError("Lean process broken pipe") except Exception as e: - logger.error("Failed to write to REPL stdin: %s", e) + logger.error("Failed to write to REPL stdin: {}", e) raise LeanError("Failed to write to REPL stdin") logger.debug("Reading response from REPL stdout") raw = await self._read_response() elapsed = loop.time() - start - logger.debug("Raw response from REPL: %r", raw) + if self.proc.returncode not in (None, 0): + self.error_file.seek(0) + err = ( + self.error_file.read().strip() + or "Lean REPL died, probably out of memory" + ) + raise ReplOOMError(err) + + logger.debug("Raw response from REPL: {}", raw) try: resp: CommandResponse | Error = json.loads(raw) except json.JSONDecodeError: - logger.error("JSON decode error: %r", raw) + logger.error("JSON decode error: {}", raw) raise ReplError("JSON decode error") self.error_file.seek(0) @@ -306,7 +328,7 @@ async def send( self.error_file.seek(0) self.error_file.truncate(0) if err: - logger.error("Stderr: %s", err) + logger.error("Stderr: {}", err) raise LeanError(err) elapsed_time = round(elapsed, 6) @@ -336,7 +358,7 @@ async def _read_response(self) -> bytes: break lines.append(chunk) except Exception as e: - logger.error("Failed to read from REPL stdout: %s", e) + logger.error("Failed to read from REPL stdout: {}", e) raise LeanError("Failed to read from REPL stdout") return b"".join(lines)