From 4a3b541452cfc55175b941bb51796472703809b2 Mon Sep 17 00:00:00 2001
From: keakon <keakon@gmail.com>
Date: Thu, 28 Mar 2024 17:09:23 +0800
Subject: [PATCH] optimize checking traceback level

---
 delayed/__init__.py |  2 +-
 delayed/worker.py   | 19 ++++++++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/delayed/__init__.py b/delayed/__init__.py
index 7d7519b..321d147 100644
--- 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
index 90c9aee..e0fb92c 100644
--- 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: