From 01a6b336a728c192ff718d2c0081e2babbf440ca Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Fri, 25 Sep 2026 07:21:19 +0000
Subject: [PATCH] Fix memory leak in `ydb.aio.retry_operation`

When a coroutine yields an exception out of `retry_operation` using `YdbRetryOperationFinalResult.exc`, a reference to the exception is retained in `result.exc` after it gets raised. When dealing with exceptions that don't participate fully in cyclic garbage collection (such as Pydantic's `ValidationError` underlying a custom error), this reference cycle (`YdbRetryOperationFinalResult` -> `Exception` -> traceback -> generator frame -> locals -> `YdbRetryOperationFinalResult`) causes coroutine local objects (payloads) to stay alive in memory.

This change avoids retaining the reference to the completed exception object within the generator frame once it is re-raised, fixing the memory leak. A memory leak regression test was also added reflecting a complex chained exception scenario.

Co-authored-by: makhalin <2647652+makhalin@users.noreply.github.com>
---
 tests/test_retries.py | 51 +++++++++++++++++++++++++++++++++++++++++++
 ydb/retries.py        |  4 +++-
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_retries.py

diff --git a/tests/test_retries.py b/tests/test_retries.py
new file mode 100644
index 000000000..c49eda54e
--- /dev/null
+++ b/tests/test_retries.py
@@ -0,0 +1,51 @@
+import asyncio
+import gc
+import weakref
+
+import pytest
+import ydb
+
+
+class Payload:
+    def __init__(self):
+        self.data = b"x" * 1024 * 1024
+
+
+class BadRequestError(Exception):
+    pass
+
+
+class CException(Exception):
+    # Simulate an exception that might not participate well in GC
+    pass
+
+
+@pytest.mark.asyncio
+async def test_retry_operation_memory_leak():
+    weak_refs = []
+
+    async def my_coro():
+        payload = Payload()
+        weak_refs.append(weakref.ref(payload))
+
+        # We simulate a reference cycle without relying on Pydantic's Rust-based errors.
+        # By chaining exceptions, we build a traceback that points back to this frame.
+        try:
+            raise CException("inner")
+        except CException as e:
+            raise BadRequestError("bad request") from e
+
+    try:
+        await ydb.aio.retry_operation(my_coro)
+    except BadRequestError:
+        pass
+
+    # Run the event loop briefly and trigger garbage collection
+    for _ in range(3):
+        await asyncio.sleep(0)
+        gc.collect()
+
+    # The payload should be garbage collected because the reference cycle
+    # involving the generator yielded result, the traceback, and the local frame
+    # was explicitly broken by setting result.exc = None in the retry implementation.
+    assert not any(w() is not None for w in weak_refs)
diff --git a/ydb/retries.py b/ydb/retries.py
index c151e3d21..fb9cf8a64 100644
--- a/ydb/retries.py
+++ b/ydb/retries.py
@@ -119,7 +119,9 @@ def retry_operation_impl(
             yield result
 
             if result.exc is not None:
-                raise result.exc
+                exc = result.exc
+                result.exc = None
+                raise exc
 
         except issues.Error as e:
             status = e