From 84995347fccc12f588efe59540418b9070a3c177 Mon Sep 17 00:00:00 2001
From: Gorfiend <Gorfiend@gmail.com>
Date: Wed, 14 Nov 2018 18:37:56 -0500
Subject: [PATCH] [niconico] Add heartbeat logic (closes #14582)

---
 youtube_dl/downloader/common.py  | 38 ++++++++++++++++++++++++++++++--
 youtube_dl/extractor/niconico.py |  9 ++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/youtube_dl/downloader/common.py b/youtube_dl/downloader/common.py
index 5979833c0..2419153e6 100644
--- a/youtube_dl/downloader/common.py
+++ b/youtube_dl/downloader/common.py
@@ -5,8 +5,12 @@ import re
 import sys
 import time
 import random
+import threading
 
-from ..compat import compat_os_name
+from ..compat import (
+    compat_os_name,
+    compat_urllib_request,
+)
 from ..utils import (
     decodeArgument,
     encodeFilename,
@@ -361,7 +365,37 @@ class FileDownloader(object):
                     else '%.2f' % sleep_interval))
             time.sleep(sleep_interval)
 
-        return self.real_download(filename, info_dict)
+        timer = [None]
+        heartbeat_lock = None
+        download_complete = False
+        if 'heartbeat_url'in info_dict:
+            heartbeat_lock = threading.Lock()
+
+            heartbeat_url = info_dict['heartbeat_url']
+            heartbeat_data = info_dict['heartbeat_data']
+            heartbeat_interval = info_dict.get('heartbeat_interval', 30)
+            self.to_screen('[download] Heartbeat with %s second interval...' % heartbeat_interval)
+
+            def heartbeat():
+                try:
+                    compat_urllib_request.urlopen(url=heartbeat_url, data=heartbeat_data)
+                except Exception:
+                    self.to_screen("[download] Heartbeat failed")
+
+                with heartbeat_lock:
+                    if not download_complete:
+                        timer[0] = threading.Timer(heartbeat_interval, heartbeat)
+                        timer[0].start()
+
+            heartbeat()
+
+        try:
+            return self.real_download(filename, info_dict)
+        finally:
+            if heartbeat_lock:
+                with heartbeat_lock:
+                    timer[0].cancel()
+                    download_complete = True
 
     def real_download(self, filename, info_dict):
         """Real download process. Redefine in subclasses."""
diff --git a/youtube_dl/extractor/niconico.py b/youtube_dl/extractor/niconico.py
index 76b412ff1..0d06f0e80 100644
--- a/youtube_dl/extractor/niconico.py
+++ b/youtube_dl/extractor/niconico.py
@@ -254,6 +254,12 @@ class NiconicoIE(InfoExtractor):
                 }
             }).encode())
 
+        # get heartbeat info
+        heartbeat_url = session_api_endpoint['url'] + '/' + session_response['data']['session']['id'] + '?_format=json&_method=PUT'
+        heartbeat_data = json.dumps(session_response['data']).encode()
+        # interval, convert milliseconds to seconds, then halve to make a buffer.
+        heartbeat_interval = session_api_data['heartbeat_lifetime'] / 2000
+
         resolution = video_quality.get('resolution', {})
 
         return {
@@ -264,6 +270,9 @@ class NiconicoIE(InfoExtractor):
             'vbr': float_or_none(video_quality.get('bitrate'), 1000),
             'height': resolution.get('height'),
             'width': resolution.get('width'),
+            'heartbeat_url': heartbeat_url,
+            'heartbeat_data': heartbeat_data,
+            'heartbeat_interval': heartbeat_interval,
         }
 
     def _real_extract(self, url):