From 38d94484bb2dfb70cf0035da5ac15e641f76a01a Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 10 Nov 2020 00:21:36 +0200
Subject: [PATCH] Use mautrix utility function for file upload retry

---
 mautrix_telegram/util/file_transfer.py | 31 +++++---------------------
 requirements.txt                       |  2 +-
 2 files changed, 6 insertions(+), 27 deletions(-)

diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py
index 7dddf607..da7d3523 100644
--- a/mautrix_telegram/util/file_transfer.py
+++ b/mautrix_telegram/util/file_transfer.py
@@ -30,8 +30,7 @@ from telethon.errors import (AuthBytesInvalidError, AuthKeyInvalidError, Locatio
                              SecurityError, FileIdInvalidError)
 
 from mautrix.appservice import IntentAPI
-from mautrix.types import ContentURI
-from mautrix.errors import MatrixRequestError, MatrixConnectionError
+from mautrix.util.network_retry import call_with_net_retry
 
 from ..tgclient import MautrixTelegramClient
 from ..db import TelegramFile as DBTelegramFile
@@ -108,28 +107,6 @@ def _location_to_id(location: TypeLocation) -> str:
         return f"{location.volume_id}-{location.local_id}"
 
 
-bad_gateway_sleep = 5
-
-
-async def _try_upload_media(intent: IntentAPI, file: bytes, mime: str, retries: int = 2
-                            ) -> ContentURI:
-    while True:
-        try:
-            return await intent.upload_media(file, mime)
-        except MatrixRequestError as e:
-            if not retries or e.http_status not in (502, 504):
-                raise
-            log.warning("Got gateway error trying to upload media, retrying in "
-                        f"{bad_gateway_sleep} seconds")
-        except MatrixConnectionError as e:
-            if not retries:
-                raise
-            log.warning(f"Got connection error trying to upload media: {e}, retrying in "
-                        f"{bad_gateway_sleep} seconds")
-        await asyncio.sleep(bad_gateway_sleep)
-        retries -= 1
-
-
 async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
                                        thumbnail_loc: TypeLocation, mime_type: str, encrypt: bool,
                                        video: Optional[bytes], custom_data: Optional[bytes] = None,
@@ -168,7 +145,8 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
     if encrypt:
         file, decryption_info = encrypt_attachment(file)
         upload_mime_type = "application/octet-stream"
-    content_uri = await _try_upload_media(intent, file, upload_mime_type)
+    content_uri = await call_with_net_retry(intent.upload_media, file, upload_mime_type,
+                                            _action="upload media")
     if decryption_info:
         decryption_info.url = content_uri
 
@@ -269,7 +247,8 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
         if encrypt and encrypt_attachment:
             file, decryption_info = encrypt_attachment(file)
             upload_mime_type = "application/octet-stream"
-        content_uri = await _try_upload_media(intent, file, upload_mime_type)
+        content_uri = await call_with_net_retry(intent.upload_media, file, upload_mime_type,
+                                                _action="upload media")
         if decryption_info:
             decryption_info.url = content_uri
 
diff --git a/requirements.txt b/requirements.txt
index 3d2861a4..3b0f7c3c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,6 +5,6 @@ python-magic>=0.4,<0.5
 commonmark>=0.8,<0.10
 aiohttp>=3,<4
 yarl>=1,<2
-mautrix>=0.8.1,<0.9
+mautrix>=0.8.2,<0.9
 telethon>=1.17,<1.18
 telethon-session-sqlalchemy>=0.2.14,<0.3