diff --git a/brozzler/ydl.py b/brozzler/ydl.py index 93ea72e..af4cfe4 100644 --- a/brozzler/ydl.py +++ b/brozzler/ydl.py @@ -28,6 +28,9 @@ import os import json import doublethink import datetime +import threading + +global_ydl_lock = threading.Lock() _orig_webpage_read_content = youtube_dl.extractor.generic.GenericIE._webpage_read_content def _webpage_read_content(self, *args, **kwargs): @@ -162,10 +165,12 @@ def _build_youtube_dl(worker, destdir, site): with open(ctx['filename'], 'rb') as f: # include content-length header to avoid chunked # transfer, which warcprox currently rejects + extra_headers = dict(site.extra_headers()) + extra_headers['content-length'] = size request, response = worker._warcprox_write_record( warcprox_address=worker._proxy_for(site), url=url, warc_type='resource', content_type=mimetype, payload=f, - extra_headers={'content-length': size}) + extra_headers=extra_headers) # consulted by _remember_videos() self.stitch_ups.append({ 'url': url, @@ -182,8 +187,14 @@ def _build_youtube_dl(worker, destdir, site): if worker._using_warcprox(site): self._push_stitched_up_vid_to_warcprox(site, info_dict, ctx) - youtube_dl.downloader.fragment.FragmentFD._finish_frag_download = _finish_frag_download - return super().process_info(info_dict) + # lock this section to prevent race condition between threads that + # want to monkey patch _finish_frag_download() at the same time + with global_ydl_lock: + try: + youtube_dl.downloader.fragment.FragmentFD._finish_frag_download = _finish_frag_download + return super().process_info(info_dict) + finally: + youtube_dl.downloader.fragment.FragmentFD._finish_frag_download = _orig__finish_frag_download def maybe_heartbeat_site_last_claimed(*args, **kwargs): # in case youtube-dl takes a long time, heartbeat site.last_claimed diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 878cd35..f5007a1 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -801,7 +801,10 @@ def test_ydl_stitching(httpd): rr = doublethink.Rethinker('localhost', db='brozzler') frontier = brozzler.RethinkDbFrontier(rr) site = brozzler.Site(rr, { - 'seed': 'http://localhost:%s/site10/' % httpd.server_port}) + 'seed': 'http://localhost:%s/site10/' % httpd.server_port, + 'warcprox_meta': { + 'warc-prefix': 'test_ydl_stitching', + 'captures-table-extra-fields': {'test_id':test_id}}}) brozzler.new_site(frontier, site) # the site should be brozzled fairly quickly @@ -816,11 +819,21 @@ def test_ydl_stitching(httpd): assert len(pages) == 1 page = pages[0] assert len(page.videos) == 6 + stitched_url = 'youtube-dl:00001:http://localhost:%s/site10/' % httpd.server_port assert { 'blame': 'youtube-dl', 'content-length': 267900, 'content-type': 'video/mp4', 'response_code': 204, - 'url': 'youtube-dl:00001:http://localhost:%s/site10/' % httpd.server_port, + 'url': stitched_url, } in page.videos + time.sleep(2) # in case warcprox hasn't finished processing urls + # take a look at the captures table + captures = list(rr.table('captures').filter({'test_id':test_id}).run()) + l = [c for c in captures if c['url'] == stitched_url] + assert len(l) == 1 + c = l[0] + assert c['filename'].startswith('test_ydl_stitching') + assert c['content_type'] == 'video/mp4' + assert c['http_method'] == 'WARCPROX_WRITE_RECORD'