From 43807746ff6d7ebd8d1a219c0f6438a2607b7290 Mon Sep 17 00:00:00 2001
From: Dean Lee <deanlee3@gmail.com>
Date: Tue, 3 Dec 2024 13:30:54 +0800
Subject: [PATCH] athenad: fix memory leak by closing Response objects (#34101)

* fix memory leak by closing Response

* use with
---
 system/athena/athenad.py            | 12 ++++++------
 system/athena/tests/test_athenad.py |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/system/athena/athenad.py b/system/athena/athenad.py
index 909e99fa71546e..78285f018e3f56 100755
--- a/system/athena/athenad.py
+++ b/system/athena/athenad.py
@@ -258,13 +258,13 @@ def upload_handler(end_event: threading.Event) -> None:
           sz = -1
 
         cloudlog.event("athena.upload_handler.upload_start", fn=fn, sz=sz, network_type=network_type, metered=metered, retry_count=item.retry_count)
-        response = _do_upload(item, partial(cb, sm, item, tid, end_event))
 
-        if response.status_code not in (200, 201, 401, 403, 412):
-          cloudlog.event("athena.upload_handler.retry", status_code=response.status_code, fn=fn, sz=sz, network_type=network_type, metered=metered)
-          retry_upload(tid, end_event)
-        else:
-          cloudlog.event("athena.upload_handler.success", fn=fn, sz=sz, network_type=network_type, metered=metered)
+        with _do_upload(item, partial(cb, sm, item, tid, end_event)) as response:
+          if response.status_code not in (200, 201, 401, 403, 412):
+            cloudlog.event("athena.upload_handler.retry", status_code=response.status_code, fn=fn, sz=sz, network_type=network_type, metered=metered)
+            retry_upload(tid, end_event)
+          else:
+            cloudlog.event("athena.upload_handler.success", fn=fn, sz=sz, network_type=network_type, metered=metered)
 
         UploadQueueCache.cache(upload_queue)
       except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, requests.exceptions.SSLError):
diff --git a/system/athena/tests/test_athenad.py b/system/athena/tests/test_athenad.py
index a4dcbef67a7bd5..a6bfc689306533 100644
--- a/system/athena/tests/test_athenad.py
+++ b/system/athena/tests/test_athenad.py
@@ -240,7 +240,7 @@ def test_upload_handler(self, host):
   @with_upload_handler
   def test_upload_handler_retry(self, mocker, host, status, retry):
     mock_put = mocker.patch('requests.put')
-    mock_put.return_value.status_code = status
+    mock_put.return_value.__enter__.return_value.status_code = status
     fn = self._create_file('qlog.zst')
     item = athenad.UploadItem(path=fn, url=f"{host}/qlog.zst", headers={}, created_at=int(time.time()*1000), id='', allow_cellular=True)