Apply threading lock to HLS DRM preparation

Without this, if two threads started at the same time there was a very good chance they would run the code and license twice, which is unnecessary.
This commit is contained in:
rlaphoenix 2023-03-04 11:41:10 +00:00
parent 0bceb772c2
commit f8166f098c
1 changed files with 32 additions and 29 deletions

View File

@ -11,7 +11,7 @@ from functools import partial
from hashlib import md5 from hashlib import md5
from pathlib import Path from pathlib import Path
from queue import Queue from queue import Queue
from threading import Event from threading import Event, Lock
from typing import Any, Callable, Optional, Union from typing import Any, Callable, Optional, Union
import m3u8 import m3u8
@ -214,6 +214,8 @@ class HLS:
log.error("Track's HLS playlist has no segments, expecting an invariant M3U8 playlist.") log.error("Track's HLS playlist has no segments, expecting an invariant M3U8 playlist.")
sys.exit(1) sys.exit(1)
drm_lock = Lock()
def download_segment(filename: str, segment: m3u8.Segment, init_data: Queue, segment_key: Queue) -> int: def download_segment(filename: str, segment: m3u8.Segment, init_data: Queue, segment_key: Queue) -> int:
if stop_event.is_set(): if stop_event.is_set():
# the track already started downloading, but another failed or was stopped # the track already started downloading, but another failed or was stopped
@ -221,34 +223,35 @@ class HLS:
segment_save_path = (save_dir / filename).with_suffix(".mp4") segment_save_path = (save_dir / filename).with_suffix(".mp4")
newest_segment_key = segment_key.get() with drm_lock:
try: newest_segment_key = segment_key.get()
if segment.key and newest_segment_key[1] != segment.key: try:
try: if segment.key and newest_segment_key[1] != segment.key:
drm = HLS.get_drm( try:
# TODO: We append master.keys because m3u8 class only puts the last EXT-X-KEY drm = HLS.get_drm(
# to the segment.key property, not supporting multi-drm scenarios. # TODO: We append master.keys because m3u8 class only puts the last EXT-X-KEY
# By re-adding every single EXT-X-KEY found, we can at least try to get # to the segment.key property, not supporting multi-drm scenarios.
# a suitable key. However, it may not match the right segment/timeframe! # By re-adding every single EXT-X-KEY found, we can at least try to get
# It will try to use the first key provided where possible. # a suitable key. However, it may not match the right segment/timeframe!
keys=[segment.key] + master.keys, # It will try to use the first key provided where possible.
proxy=proxy keys=[segment.key] + master.keys,
) proxy=proxy
except NotImplementedError as e: )
log.error(str(e)) except NotImplementedError as e:
sys.exit(1) log.error(str(e))
else: sys.exit(1)
if drm: else:
drm = drm[0] # just use the first supported DRM system for now if drm:
log.debug("Got segment key, %s", drm) drm = drm[0] # just use the first supported DRM system for now
if isinstance(drm, Widevine): log.debug("Got segment key, %s", drm)
# license and grab content keys if isinstance(drm, Widevine):
if not license_widevine: # license and grab content keys
raise ValueError("license_widevine func must be supplied to use Widevine DRM") if not license_widevine:
license_widevine(drm) raise ValueError("license_widevine func must be supplied to use Widevine DRM")
newest_segment_key = (drm, segment.key) license_widevine(drm)
finally: newest_segment_key = (drm, segment.key)
segment_key.put(newest_segment_key) finally:
segment_key.put(newest_segment_key)
if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment): if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment):
return 0 return 0