From b01a799a9425274de0e8ed7f4b2025b735ac525f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20Kr=C4=8Dek?= <roman.krcek@tietoevry.com>
Date: Tue, 13 May 2025 18:14:36 +0200
Subject: [PATCH] Fix timing of async functions and make histogram finer at
 lower ranges

---
 requirements.txt                     | 1 +
 telegram_downloader_bot/telemetry.py | 7 ++++---
 telegram_downloader_bot/utils.py     | 5 +++--
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index d9f067a..658e514 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,5 +6,6 @@ tiktok_downloader==0.3.5
 uvloop==0.19.0
 tgcrypto==1.2.5
 prometheus-client==0.21.1
+prometheus-async==25.1.0
 pydantic-settings==2.5.2
 pydantic==2.9.2
\ No newline at end of file
diff --git a/telegram_downloader_bot/telemetry.py b/telegram_downloader_bot/telemetry.py
index ae9f75c..8262ee3 100644
--- a/telegram_downloader_bot/telemetry.py
+++ b/telegram_downloader_bot/telemetry.py
@@ -1,10 +1,10 @@
 from prometheus_client import Histogram, start_http_server
 
 DOWNLOAD_DURATION = Histogram(
-    'tt_download_time',
-    'Time taken to download a single tiktok video',
+    'download_time_seconds',
+    'Time taken to download a single media item',
     ['service'],
-    buckets=[0.1, 0.5, 1, 2, 5, 10, 30, 60, 120]
+    buckets=[0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100]
 )
 
 FILE_SIZE_BYTES = Histogram(
@@ -13,6 +13,7 @@ FILE_SIZE_BYTES = Histogram(
     ['service'],
     buckets=[
         1e6,    # 1 MB
+        2e6,    # 2 MB
         5e6,    # 5 MB
         10e6,   # 10 MB
         25e6,   # 25 MB
diff --git a/telegram_downloader_bot/utils.py b/telegram_downloader_bot/utils.py
index 0420060..5576c65 100644
--- a/telegram_downloader_bot/utils.py
+++ b/telegram_downloader_bot/utils.py
@@ -4,6 +4,7 @@ import re
 
 from datetime import datetime
 from hashlib import sha256
+from prometheus_async.aio import time as async_time
 from pyrogram import Client
 from pyrogram.types import Message
 from tiktok_downloader import snaptik
@@ -59,7 +60,7 @@ def get_user_folder(message: Message) -> os.path:
     return user_folder
 
 
-@DOWNLOAD_DURATION.labels(service='telegram').time()
+@async_time(DOWNLOAD_DURATION.labels(service='telegram'))
 async def handle_media_message_contents(client: Client,
                                         message: Message):
     """Detect what kind of media is being sent over from the user.
@@ -132,7 +133,7 @@ async def check_if_tt_downloaded(tt_hash: str) -> bool:
     return tt_hash in all_tt_hashes
 
 
-@DOWNLOAD_DURATION.labels(service='tiktok').time()
+@async_time(DOWNLOAD_DURATION.labels(service='tiktok'))
 async def download_tt_video(url: str) -> str:
     """Downloads tiktok video from a given URL.
     Makes sure the video integrity is correct."""