Merge pull request #119 from nlevitt/ydl-stitch-fix

WIP youtube-dl stitching fixes

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

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'