add prometheus metrics to download youtube video on an interval

modules/args.py

@@ -31,4 +31,10 @@ def get_args():
         "--cache-state-file",
         help="JSON file to persist cache state on server shutdown and recover on startup. if specified, the server will not empty the cache on shutdown"
     )
+    parser.add_argument(
+        "--download-monitoring-interval",
+        type=int,
+        default=0,
+        help="run download monitoring loop every interval, interval=0 to disable monitoring"
+    )
     return parser.parse_args()

modules/download_interval.py

@@ -0,0 +1,44 @@
+import os
+import shutil
+import time
+import threading
+from modules.metrics import Metrics, MetricsHandler
+from modules.args import get_args
+from pytubefix import YouTube
+from pytubefix.exceptions import RegexMatchError, VideoUnavailable
+
+args = get_args()
+
+def download_loop(interval):
+    video_url = "https://www.youtube.com/watch?v=X96RjH8WC5o"
+    while True:
+        try:
+            start_time = time.time()
+            # Download YouTube video to a temp directory
+            video =  YouTube(video_url)
+            stream = video.streams.get_highest_resolution()
+            os.makedirs("temp", exist_ok=True)
+            download_path = stream.download(output_path="temp")
+            if not os.path.exists(download_path):
+                raise Exception()
+            shutil.rmtree("temp", ignore_errors=True)
+            end_time = time.time()
+
+            # Calculate and update download speed and success metrics
+            download_time = end_time - start_time
+            video_size = stream.filesize
+            if download_time > 0:
+                speed = video_size / download_time 
+            else:
+                speed = "video did not download"
+            MetricsHandler.download_speed.set(speed)
+            MetricsHandler.download_success.labels(reason="video_downloaded").set(1)
+
+        except VideoUnavailable:
+            MetricsHandler.download_success.labels(reason="video_not_found").set(0)
+        except RegexMatchError:
+            MetricsHandler.download_success.labels(reason="regex_mismatch").set(0)
+        except Exception:
+            MetricsHandler.download_success.labels(reason="download_failed").set(0)
+
+        time.sleep(interval)

modules/metrics.py

@@ -42,6 +42,19 @@ class Metrics(enum.Enum):
         prometheus_client.Counter,
     )
 
+    DOWNLOAD_SPEED = (
+        "download_speed",
+        "Speed of YouTube video download in bytes/second",
+        prometheus_client.Gauge
+    )
+
+    DOWNLOAD_SUCCESS = (
+        "download_success",
+        "Success of YouTube video download (1=success, 0=fail)",
+        prometheus_client.Gauge,
+        ["reason"]
+    )
+
     CACHE_SIZE = (
         "cache_size",
         "Total entries in cache",

server.py

@@ -29,7 +29,7 @@
 from modules.args import get_args
 from modules.cache import Cache
 from modules.metrics import MetricsHandler
-
+from modules.download_interval import download_loop
 
 logging.Formatter.converter = time.gmtime
 logging.basicConfig(
@@ -343,6 +343,14 @@ def download_video(config: VideoConfig):
                 write_log_to_client(f"Failed to download {config.url}")
         return video_path
 
+def start_download_monitor():
+    # Run download monitor loop thread in background if an interval is defined
+    if args.download_monitoring_interval > 0:
+        threading.Thread(
+            target = download_loop,
+            args = (args.download_monitoring_interval,),
+            daemon = True
+        ).start()
 
 # Worker Thread for video playing.
 def play_video_worker():
@@ -851,6 +859,7 @@ def signal_handler():
     MetricsHandler.init()
     MetricsHandler.cache_size.set(0)
     MetricsHandler.cache_size_bytes.set(0)
+    start_download_monitor()
 
     # Start up interlude by default
     if args.interlude: