diff options
Diffstat (limited to '')
| -rw-r--r-- | snekbox/nsjail.py | 5 | ||||
| -rw-r--r-- | snekbox/utils/timed.py | 8 |
2 files changed, 8 insertions, 5 deletions
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 1de7b1e..9bf20bf 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -3,6 +3,7 @@ import re import subprocess import sys from collections.abc import Generator +from contextlib import nullcontext from pathlib import Path from tempfile import NamedTemporaryFile from typing import Iterable, TypeVar @@ -56,7 +57,7 @@ class NsJail: memfs_home: str = "home", memfs_output: str = "home", files_limit: int | None = 100, - files_timeout: int | None = 5, + files_timeout: float | None = 5, files_pattern: str = "**/[!_]*", ): """ @@ -267,7 +268,7 @@ class NsJail: # Parse attachments with time limit try: - with time_limit(self.files_timeout): + with time_limit(self.files_timeout) if self.files_timeout else nullcontext(): attachments = fs.files_list( limit=self.files_limit, pattern=self.files_pattern, diff --git a/snekbox/utils/timed.py b/snekbox/utils/timed.py index 11d126c..59a4e7f 100644 --- a/snekbox/utils/timed.py +++ b/snekbox/utils/timed.py @@ -11,7 +11,7 @@ __all__ = ("time_limit",) @contextmanager -def time_limit(timeout: int | None = None) -> Generator[None, None, None]: +def time_limit(timeout: float) -> Generator[None, None, None]: """ Decorator to call a function with a time limit. @@ -25,10 +25,12 @@ def time_limit(timeout: int | None = None) -> Generator[None, None, None]: def signal_handler(_signum, _frame): raise TimeoutError(f"time_limit call timed out after {timeout} seconds.") + # ITIMER_PROF would be more appropriate, but SIGPROF doesn't seem to interrupt sleeps. signal.signal(signal.SIGALRM, signal_handler) - signal.alarm(timeout) + signal.setitimer(signal.ITIMER_REAL, timeout) try: yield finally: - signal.alarm(0) + # Clear the timer if the function finishes early. + signal.setitimer(signal.ITIMER_REAL, 0) |