Add time limits to celery tasks on start

Syslifters · Mar 5, 2024 · 59cf1d2 · 59cf1d2
1 parent d7ecc38
commit 59cf1d2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/api/src/reportcreator_api/conf/settings.py b/api/src/reportcreator_api/conf/settings.py
@@ -452,8 +452,7 @@
 
 # Time limits are only enforced if a broker is configured and an external worker is used (but not in eager mode).
 # Self-hosted SysReptor instances use the eager mode by default, resulting in no PDF rendering time limits being applied.
-CELERY_TASK_SOFT_TIME_LIMIT = config('PDF_RENDERING_TIME_LIMIT', cast=int, default=60)
-CELERY_TASK_TIME_LIMIT = CELERY_TASK_SOFT_TIME_LIMIT + 5
+PDF_RENDERING_TIME_LIMIT = config('PDF_RENDERING_TIME_LIMIT', cast=int, default=60)
 
 
 # History

diff --git a/api/src/reportcreator_api/tasks/rendering/entry.py b/api/src/reportcreator_api/tasks/rendering/entry.py
@@ -148,13 +148,15 @@ async def format_project_template_data(project: PentestProject, project_type: Op
     )
 
 
-async def get_celery_result_async(task, timeout=timedelta(seconds=settings.CELERY_TASK_TIME_LIMIT)):
+async def get_celery_result_async(task, timeout=timedelta(seconds=settings.PDF_RENDERING_TIME_LIMIT + 5)):
     start_time = timezone.now()
     while not task.ready():
         if timezone.now() > start_time + timeout:
             logging.error('PDF rendering task timeout')
             raise TimeoutError('PDF rendering timeout')
         await asyncio.sleep(0.2)
+    if isinstance(task.result, Exception):
+        raise task.result
     return task.result
 
 

diff --git a/api/src/reportcreator_api/tasks/rendering/tasks.py b/api/src/reportcreator_api/tasks/rendering/tasks.py
@@ -6,7 +6,12 @@
 from reportcreator_api.tasks.rendering import render
 
 
-@shared_task(name='reportcreator.render_pdf', expires=settings.CELERY_TASK_TIME_LIMIT)
+@shared_task(
+    name='reportcreator.render_pdf', 
+    soft_time_limit=settings.PDF_RENDERING_TIME_LIMIT, 
+    time_limit=settings.PDF_RENDERING_TIME_LIMIT + 5, 
+    expires=settings.PDF_RENDERING_TIME_LIMIT + 5,
+)
 @log_timing
 def render_pdf_task(*args, **kwargs) -> dict:
     pdf, msgs = render.render_pdf(*args, **kwargs)