diff --git a/modules/args.py b/modules/args.py index 43883a8..69039ba 100644 --- a/modules/args.py +++ b/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() diff --git a/modules/download_interval.py b/modules/download_interval.py new file mode 100644 index 0000000..53f026c --- /dev/null +++ b/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) \ No newline at end of file diff --git a/modules/metrics.py b/modules/metrics.py index 20fbd8c..9274598 100644 --- a/modules/metrics.py +++ b/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", diff --git a/server.py b/server.py index 775f8bd..e38c509 100644 --- a/server.py +++ b/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: