optimize checking traceback level

keakon · Mar 28, 2024 · 4a3b541 · 4a3b541
1 parent ae26632
commit 4a3b541
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/delayed/__init__.py b/delayed/__init__.py
@@ -1,3 +1,3 @@
 # -*- coding: utf-8 -*-
 
-__version__ = '1.0.1b1'
+__version__ = '1.0.1b2'
diff --git a/delayed/worker.py b/delayed/worker.py
@@ -6,7 +6,7 @@
 import sys
 import threading
 import time
-import traceback
+import sys
 from typing import Union
 
 from .constants import DEFAULT_SLEEP_TIME, MAX_SLEEP_TIME, Status
@@ -58,12 +58,25 @@ def run(self):  # pragma: no cover
                         except Exception:
                             logger.exception('Failed to execute task %s.', task._func_path)
 
+                            need_retry = False
                             _, _, exc_traceback = sys.exc_info()
-                            if len(traceback.format_tb(exc_traceback)) > 2:
+                            if exc_traceback:
+                                tb_next = exc_traceback.tb_next
+                                if tb_next:
+                                    tb_next2 = tb_next.tb_next
+                                    if tb_next2:
+                                        # invalid call, should not be retried
+                                        need_retry = True
+                                        # delete tracebacks to avoid memory leak
+                                        del tb_next2
+                                    del tb_next
+                                del exc_traceback
+
+                            if need_retry:
                                 self._requeue_task(task)
                             else:
-                                # invalid call, should not be retried
                                 self._release_task()
+                            self._release_task()
                         else:
                             self._release_task()
         finally: