From fe845284c4f01345c6dfaf0631860e8fe0534df3 Mon Sep 17 00:00:00 2001
From: pukkandan <pukkandan.ytdlp@gmail.com>
Date: Sat, 20 Mar 2021 08:49:02 +0530
Subject: [PATCH] [aria2c] Support retry/abort unavailable fragments (#173)

Co-authored by: Damiano Amatruda <damiano.amatruda@outlook.com>
---
 yt_dlp/downloader/external.py | 76 +++++++++++++++++++++--------------
 1 file changed, 46 insertions(+), 30 deletions(-)

diff --git a/yt_dlp/downloader/external.py b/yt_dlp/downloader/external.py
index 176105c66..63e430ac9 100644
--- a/yt_dlp/downloader/external.py
+++ b/yt_dlp/downloader/external.py
@@ -24,7 +24,6 @@ from ..utils import (
     cli_bool_option,
     cli_configuration_args,
     encodeFilename,
-    error_to_compat_str,
     encodeArgument,
     handle_youtubedl_headers,
     check_executable,
@@ -117,19 +116,42 @@ class ExternalFD(FileDownloader):
 
         self._debug_cmd(cmd)
 
-        p = subprocess.Popen(
-            cmd, stderr=subprocess.PIPE)
-        _, stderr = process_communicate_or_kill(p)
-        if p.returncode != 0:
-            self.to_stderr(stderr.decode('utf-8', 'replace'))
-
         if 'fragments' in info_dict:
-            file_list = []
+            fragment_retries = self.params.get('fragment_retries', 0)
+            skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True)
+
+            count = 0
+            while count <= fragment_retries:
+                p = subprocess.Popen(
+                    cmd, stderr=subprocess.PIPE)
+                _, stderr = process_communicate_or_kill(p)
+                if p.returncode == 0:
+                    break
+                # TODO: Decide whether to retry based on error code
+                # https://aria2.github.io/manual/en/html/aria2c.html#exit-status
+                self.to_stderr(stderr.decode('utf-8', 'replace'))
+                count += 1
+                if count <= fragment_retries:
+                    self.to_screen(
+                        '[%s] Got error. Retrying fragments (attempt %d of %s)...'
+                        % (self.get_basename(), count, self.format_retries(fragment_retries)))
+            if count > fragment_retries:
+                if not skip_unavailable_fragments:
+                    self.report_error('Giving up after %s fragment retries' % fragment_retries)
+                    return -1
+
             dest, _ = sanitize_open(tmpfilename, 'wb')
-            for i, fragment in enumerate(info_dict['fragments']):
-                file = '%s-Frag%d' % (tmpfilename, i)
+            for frag_index, fragment in enumerate(info_dict['fragments']):
+                fragment_filename = '%s-Frag%d' % (tmpfilename, frag_index)
+                try:
+                    src, _ = sanitize_open(fragment_filename, 'rb')
+                except IOError:
+                    if skip_unavailable_fragments and frag_index > 1:
+                        self.to_screen('[%s] Skipping fragment %d ...' % (self.get_basename(), frag_index))
+                        continue
+                    self.report_error('Unable to open fragment %d' % frag_index)
+                    return -1
                 decrypt_info = fragment.get('decrypt_info')
-                src, _ = sanitize_open(file, 'rb')
                 if decrypt_info:
                     if decrypt_info['METHOD'] == 'AES-128':
                         iv = decrypt_info.get('IV')
@@ -146,20 +168,16 @@ class ExternalFD(FileDownloader):
                     fragment_data = src.read()
                     dest.write(fragment_data)
                 src.close()
-                file_list.append(file)
+                if not self.params.get('keep_fragments', False):
+                    os.remove(encodeFilename(fragment_filename))
             dest.close()
-            if not self.params.get('keep_fragments', False):
-                for file_path in file_list:
-                    try:
-                        os.remove(file_path)
-                    except OSError as ose:
-                        self.report_error("Unable to delete file %s; %s" % (file_path, error_to_compat_str(ose)))
-                try:
-                    file_path = '%s.frag.urls' % tmpfilename
-                    os.remove(file_path)
-                except OSError as ose:
-                    self.report_error("Unable to delete file %s; %s" % (file_path, error_to_compat_str(ose)))
-
+            os.remove(encodeFilename('%s.frag.urls' % tmpfilename))
+        else:
+            p = subprocess.Popen(
+                cmd, stderr=subprocess.PIPE)
+            _, stderr = process_communicate_or_kill(p)
+            if p.returncode != 0:
+                self.to_stderr(stderr.decode('utf-8', 'replace'))
         return p.returncode
 
     def _prepare_url(self, info_dict, url):
@@ -276,17 +294,15 @@ class Aria2cFD(ExternalFD):
         cmd += ['--auto-file-renaming=false']
 
         if 'fragments' in info_dict:
-            cmd += verbose_level_args
-            cmd += ['--uri-selector', 'inorder', '--download-result=hide']
+            cmd += ['--file-allocation=none', '--uri-selector=inorder']
             url_list_file = '%s.frag.urls' % tmpfilename
             url_list = []
-            for i, fragment in enumerate(info_dict['fragments']):
-                tmpsegmentname = '%s-Frag%d' % (os.path.basename(tmpfilename), i)
-                url_list.append('%s\n\tout=%s' % (fragment['url'], tmpsegmentname))
+            for frag_index, fragment in enumerate(info_dict['fragments']):
+                fragment_filename = '%s-Frag%d' % (os.path.basename(tmpfilename), frag_index)
+                url_list.append('%s\n\tout=%s' % (fragment['url'], fragment_filename))
             stream, _ = sanitize_open(url_list_file, 'wb')
             stream.write('\n'.join(url_list).encode('utf-8'))
             stream.close()
-
             cmd += ['-i', url_list_file]
         else:
             cmd += ['--', info_dict['url']]