Support fractional seconds for files_timeout

This will be helpful in testing, since some machines are able to
process quite a lot in one second, which was the previous lowest
timeout.

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)