Improved search speed by reusing providers pools

11 files changed, 236 insertions, 35 deletions

diff --git a/libs/subliminal/providers/__init__.py b/libs/subliminal/providers/__init__.py
index 882236ca3..1de7b53c6 100644
--- a/libs/subliminal/providers/__init__.py
+++ b/libs/subliminal/providers/__init__.py
@@ -161,5 +161,8 @@ class Provider(object):
         """
         raise NotImplementedError
 
+    def ping(self):
+        return True
+
     def __repr__(self):
         return '<%s [%r]>' % (self.__class__.__name__, self.video_types)
diff --git a/libs/subliminal_patch/core.py b/libs/subliminal_patch/core.py
index f4605b1eb..9ed698c09 100644
--- a/libs/subliminal_patch/core.py
+++ b/libs/subliminal_patch/core.py
@@ -5,6 +5,7 @@ import json
 import re
 import os
 import logging
+import datetime
 import socket
 import traceback
 import time
@@ -55,6 +56,8 @@ REMOVE_CRAP_FROM_FILENAME = re.compile(r"(?i)(?:([\s_-]+(?:obfuscated|scrambled|
 
 SUBTITLE_EXTENSIONS = ('.srt', '.sub', '.smi', '.txt', '.ssa', '.ass', '.mpl', '.vtt')
 
+_POOL_LIFETIME = datetime.timedelta(hours=12)
+
 
 def remove_crap_from_fn(fn):
     # in case of the second regex part, the legit release group name will be in group(2), if it's followed by [string]
@@ -69,7 +72,7 @@ class SZProviderPool(ProviderPool):
     def __init__(self, providers=None, provider_configs=None, blacklist=None, ban_list=None, throttle_callback=None,
                  pre_download_hook=None, post_download_hook=None, language_hook=None):
         #: Name of providers to use
-        self.providers = providers
+        self.providers = set(providers or [])
 
         #: Provider configuration
         self.provider_configs = provider_configs or {}
@@ -91,9 +94,69 @@ class SZProviderPool(ProviderPool):
         self.post_download_hook = post_download_hook
         self.language_hook = language_hook
 
+        self._born = time.time()
+
         if not self.throttle_callback:
             self.throttle_callback = lambda x, y: x
 
+    def update(self, providers, provider_configs, blacklist, ban_list):
+        # Check if the pool was initialized enough hours ago
+        self._check_lifetime()
+
+        # Check if any new provider has been added
+        updated = set(providers) != self.providers or ban_list != self.ban_list
+        removed_providers = list(sorted(self.providers - set(providers)))
+        new_providers = list(sorted(set(providers) - self.providers))
+
+        # Terminate and delete removed providers from instance
+        for removed in removed_providers:
+            try:
+                del self[removed]
+                # If the user has updated the providers but hasn't made any
+                # subtitle searches yet, the removed provider won't be in the
+                # self dictionary
+            except KeyError:
+                pass
+
+        if updated:
+            logger.debug("Removed providers: %s", removed_providers)
+            logger.debug("New providers: %s", new_providers)
+
+            self.discarded_providers.difference_update(new_providers)
+            self.providers.difference_update(removed_providers)
+            self.providers.update(list(providers))
+
+        self.blacklist = blacklist
+
+        # Restart providers with new configs
+        for key, val in provider_configs.items():
+            # key: provider's name; val: config dict
+            old_val = self.provider_configs.get(key)
+
+            if old_val == val:
+                continue
+
+            logger.debug("Restarting provider: %s", key)
+            try:
+                provider = provider_registry[key](**val)
+                provider.initialize()
+            except Exception as error:
+                self.throttle_callback(key, error)
+            else:
+                self.initialized_providers[key] = provider
+                updated = True
+
+        self.provider_configs = provider_configs
+
+        return updated
+
+    def _check_lifetime(self):
+        # This method is used to avoid possible memory leaks
+        if abs(self._born - time.time()) > _POOL_LIFETIME.seconds:
+            logger.info("%s elapsed. Terminating providers", _POOL_LIFETIME)
+            self._born = time.time()
+            self.terminate()
+
     def __enter__(self):
         return self
 
@@ -170,17 +233,7 @@ class SZProviderPool(ProviderPool):
         logger.info('Listing subtitles with provider %r and languages %r', provider, provider_languages)
         results = []
         try:
-            try:
-                results = self[provider].list_subtitles(video, provider_languages)
-            except ResponseNotReady:
-                logger.error('Provider %r response error, reinitializing', provider)
-                try:
-                    self[provider].terminate()
-                    self[provider].initialize()
-                    results = self[provider].list_subtitles(video, provider_languages)
-                except:
-                    logger.error('Provider %r reinitialization error: %s', provider, traceback.format_exc())
-
+            results = self[provider].list_subtitles(video, provider_languages)
             seen = []
             out = []
             for s in results:
@@ -198,16 +251,13 @@ class SZProviderPool(ProviderPool):
                         continue
                 if s.id in seen:
                     continue
+
                 s.plex_media_fps = float(video.fps) if video.fps else None
                 out.append(s)
                 seen.append(s.id)
 
             return out
 
-        except (requests.Timeout, socket.timeout) as e:
-            logger.error('Provider %r timed out', provider)
-            self.throttle_callback(provider, e)
-
         except Exception as e:
             logger.exception('Unexpected error in provider %r: %s', provider, traceback.format_exc())
             self.throttle_callback(provider, e)
@@ -289,16 +339,6 @@ class SZProviderPool(ProviderPool):
                 logger.error('Provider %r connection error', subtitle.provider_name)
                 self.throttle_callback(subtitle.provider_name, e)
 
-            except ResponseNotReady as e:
-                logger.error('Provider %r response error, reinitializing', subtitle.provider_name)
-                try:
-                    self[subtitle.provider_name].terminate()
-                    self[subtitle.provider_name].initialize()
-                except:
-                    logger.error('Provider %r reinitialization error: %s', subtitle.provider_name,
-                                 traceback.format_exc())
-                    self.throttle_callback(subtitle.provider_name, e)
-
             except rarfile.BadRarFile:
                 logger.error('Malformed RAR file from provider %r, skipping subtitle.', subtitle.provider_name)
                 logger.debug("RAR Traceback: %s", traceback.format_exc())
diff --git a/libs/subliminal_patch/core_persistent.py b/libs/subliminal_patch/core_persistent.py
new file mode 100644
index 000000000..ee55b12a8
--- /dev/null
+++ b/libs/subliminal_patch/core_persistent.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+from __future__ import absolute_import
+
+from collections import defaultdict
+import logging
+import time
+
+from subliminal.core import check_video
+
+logger = logging.getLogger(__name__)
+
+# list_all_subtitles, list_supported_languages, list_supported_video_types, download_subtitles, download_best_subtitles
+def list_all_subtitles(videos, languages, pool_instance):
+    listed_subtitles = defaultdict(list)
+
+    # return immediatly if no video passed the checks
+    if not videos:
+        return listed_subtitles
+
+    for video in videos:
+        logger.info("Listing subtitles for %r", video)
+        subtitles = pool_instance.list_subtitles(
+            video, languages - video.subtitle_languages
+        )
+        listed_subtitles[video].extend(subtitles)
+        logger.info("Found %d subtitle(s)", len(subtitles))
+
+    return listed_subtitles
+
+
+def list_supported_languages(pool_instance):
+    return pool_instance.list_supported_languages()
+
+
+def list_supported_video_types(pool_instance):
+    return pool_instance.list_supported_video_types()
+
+
+def download_subtitles(subtitles, pool_instance):
+    for subtitle in subtitles:
+        logger.info("Downloading subtitle %r with score %s", subtitle, subtitle.score)
+        pool_instance.download_subtitle(subtitle)
+
+
+def download_best_subtitles(
+    videos,
+    languages,
+    pool_instance,
+    min_score=0,
+    hearing_impaired=False,
+    only_one=False,
+    compute_score=None,
+    throttle_time=0,
+    score_obj=None,
+):
+    downloaded_subtitles = defaultdict(list)
+
+    # check videos
+    checked_videos = []
+    for video in videos:
+        if not check_video(video, languages=languages, undefined=only_one):
+            logger.info("Skipping video %r", video)
+            continue
+        checked_videos.append(video)
+
+    # return immediately if no video passed the checks
+    if not checked_videos:
+        return downloaded_subtitles
+
+    got_multiple = len(checked_videos) > 1
+
+    # download best subtitles
+    for video in checked_videos:
+        logger.info("Downloading best subtitles for %r", video)
+        subtitles = pool_instance.download_best_subtitles(
+            pool_instance.list_subtitles(video, languages - video.subtitle_languages),
+            video,
+            languages,
+            min_score=min_score,
+            hearing_impaired=hearing_impaired,
+            only_one=only_one,
+            compute_score=compute_score,
+            score_obj=score_obj,
+        )
+        logger.info("Downloaded %d subtitle(s)", len(subtitles))
+        downloaded_subtitles[video].extend(subtitles)
+
+        if got_multiple and throttle_time:
+            logger.debug("Waiting %ss before continuing ...", throttle_time)
+            time.sleep(throttle_time)
+
+    return downloaded_subtitles
diff --git a/libs/subliminal_patch/providers/__init__.py b/libs/subliminal_patch/providers/__init__.py
index ced4694f3..b8a850b58 100644
--- a/libs/subliminal_patch/providers/__init__.py
+++ b/libs/subliminal_patch/providers/__init__.py
@@ -1,8 +1,11 @@
 # coding=utf-8
 
 from __future__ import absolute_import
+
+import functools
 import importlib
 import os
+import logging
 import subliminal
 from subliminal.providers import Provider as _Provider
 from subliminal.subtitle import Subtitle as _Subtitle
@@ -14,11 +17,50 @@ from subzero.lib.io import get_viable_encoding
 import six
 
 
+logger = logging.getLogger(__name__)
+
+
 class Provider(_Provider):
     hash_verifiable = False
     hearing_impaired_verifiable = False
     skip_wrong_fps = True
 
+    def ping(self):
+        """Check if the provider is alive."""
+        return True
+
+
+def reinitialize_on_error(exceptions: tuple, attempts=1):
+    """Method decorator for Provider class. It will reinitialize the instance
+    and re-run the method in case of exceptions.
+
+    :param exceptions: tuple of expected exceptions
+    :param attempts: number of attempts to call the method
+    """
+
+    def real_decorator(method):
+        @functools.wraps(method)
+        def wrapper(self, *args, **kwargs):
+            inc = 1
+            while True:
+                try:
+                    return method(self, *args, **kwargs)
+                except exceptions as error:
+                    if inc > attempts:
+                        raise
+
+                    logger.exception(error)
+                    logger.debug("Reinitializing %s instance (%s attempt)", self, inc)
+
+                    self.terminate()
+                    self.initialize()
+
+                    inc += 1
+
+        return wrapper
+
+    return real_decorator
+
 
 # register providers
 # fixme: this is bad
diff --git a/libs/subliminal_patch/providers/addic7ed.py b/libs/subliminal_patch/providers/addic7ed.py
index ebdd7ae2b..e7ff5e2c4 100644
--- a/libs/subliminal_patch/providers/addic7ed.py
+++ b/libs/subliminal_patch/providers/addic7ed.py
@@ -11,12 +11,14 @@ from urllib.parse import quote_plus
 import babelfish
 from dogpile.cache.api import NO_VALUE
 from requests import Session
+from requests.exceptions import RequestException
 from subliminal.cache import region
 from subliminal.video import Episode, Movie
 from subliminal.exceptions import DownloadLimitExceeded, AuthenticationError, ConfigurationError
 from subliminal.providers.addic7ed import Addic7edProvider as _Addic7edProvider, \
     Addic7edSubtitle as _Addic7edSubtitle, ParserBeautifulSoup
 from subliminal.subtitle import fix_line_ending
+from subliminal_patch.providers import reinitialize_on_error
 from subliminal_patch.utils import sanitize
 from subliminal_patch.exceptions import TooManyRequests
 from subliminal_patch.pitcher import pitchers, load_verification, store_verification
@@ -550,6 +552,7 @@ class Addic7edProvider(_Addic7edProvider):
 
         return subtitles
 
+    @reinitialize_on_error((RequestException,), attempts=1)
     def list_subtitles(self, video, languages):
         if isinstance(video, Episode):
             # lookup show_id
@@ -586,6 +589,7 @@ class Addic7edProvider(_Addic7edProvider):
 
         return []
 
+    @reinitialize_on_error((RequestException,), attempts=1)
     def download_subtitle(self, subtitle):
         last_dls = region.get("addic7ed_dls")
         now = datetime.datetime.now()
diff --git a/libs/subliminal_patch/providers/legendasdivx.py b/libs/subliminal_patch/providers/legendasdivx.py
index 2f7ac60e7..e4a5ab292 100644
--- a/libs/subliminal_patch/providers/legendasdivx.py
+++ b/libs/subliminal_patch/providers/legendasdivx.py
@@ -10,6 +10,7 @@ from requests.exceptions import HTTPError
 import rarfile
 
 from guessit import guessit
+from requests.exceptions import RequestException
 from subliminal.cache import region
 from subliminal.exceptions import ConfigurationError, AuthenticationError, ServiceUnavailable, DownloadLimitExceeded
 from subliminal.providers import ParserBeautifulSoup
@@ -18,7 +19,7 @@ from subliminal.utils import sanitize, sanitize_release_group
 from subliminal.video import Episode, Movie
 from subliminal_patch.exceptions import TooManyRequests, IPAddressBlocked
 from subliminal_patch.http import RetryingCFSession
-from subliminal_patch.providers import Provider
+from subliminal_patch.providers import Provider, reinitialize_on_error
 from subliminal_patch.score import get_scores, framerate_equal
 from subliminal_patch.subtitle import Subtitle, guess_matches
 from subzero.language import Language
@@ -260,6 +261,7 @@ class LegendasdivxProvider(Provider):
             )
         return subtitles
 
+    @reinitialize_on_error((RequestException,), attempts=1)
     def query(self, video, languages):
 
         _searchurl = self.searchurl
@@ -362,7 +364,8 @@ class LegendasdivxProvider(Provider):
 
     def list_subtitles(self, video, languages):
         return self.query(video, languages)
-
+    
+    @reinitialize_on_error((RequestException,), attempts=1)
     def download_subtitle(self, subtitle):
 
         try:
diff --git a/libs/subliminal_patch/providers/legendastv.py b/libs/subliminal_patch/providers/legendastv.py
index 638f332fb..6ee07ee11 100644
--- a/libs/subliminal_patch/providers/legendastv.py
+++ b/libs/subliminal_patch/providers/legendastv.py
@@ -8,8 +8,10 @@ from subliminal.exceptions import ConfigurationError
 from subliminal.providers.legendastv import LegendasTVSubtitle as _LegendasTVSubtitle, \
     LegendasTVProvider as _LegendasTVProvider, Episode, Movie, guessit, sanitize, region, type_map, \
     raise_for_status, json, SHOW_EXPIRATION_TIME, title_re, season_re, datetime, pytz, NO_VALUE, releases_key, \
-    SUBTITLE_EXTENSIONS, language_converters
+    SUBTITLE_EXTENSIONS, language_converters, ServiceUnavailable
 
+from requests.exceptions import RequestException
+from subliminal_patch.providers import reinitialize_on_error
 from subliminal_patch.subtitle import guess_matches
 from subzero.language import Language
 
@@ -184,6 +186,7 @@ class LegendasTVProvider(_LegendasTVProvider):
 
         return titles_found
 
+    @reinitialize_on_error((RequestException, ServiceUnavailable), attempts=1)
     def query(self, language, titles, season=None, episode=None, year=None, imdb_id=None):
         # search for titles
         titles_found = self.search_titles(titles, season, year, imdb_id)
diff --git a/libs/subliminal_patch/providers/opensubtitles.py b/libs/subliminal_patch/providers/opensubtitles.py
index cfa144670..2918fd6ce 100644
--- a/libs/subliminal_patch/providers/opensubtitles.py
+++ b/libs/subliminal_patch/providers/opensubtitles.py
@@ -18,6 +18,7 @@ from subliminal.providers.opensubtitles import OpenSubtitlesProvider as _OpenSub
     DownloadLimitReached, InvalidImdbid, UnknownUserAgent, DisabledUserAgent, OpenSubtitlesError
 from .mixins import ProviderRetryMixin
 from subliminal.subtitle import fix_line_ending
+from subliminal_patch.providers import reinitialize_on_error
 from subliminal_patch.http import SubZeroRequestsTransport
 from subliminal_patch.utils import sanitize, fix_inconsistent_naming
 from subliminal.cache import region
@@ -236,7 +237,7 @@ class OpenSubtitlesProvider(ProviderRetryMixin, _OpenSubtitlesProvider):
     def terminate(self):
         self.server = None
         self.token = None
-
+    
     def list_subtitles(self, video, languages):
         """
         :param video:
@@ -272,6 +273,7 @@ class OpenSubtitlesProvider(ProviderRetryMixin, _OpenSubtitlesProvider):
                           use_tag_search=self.use_tag_search, only_foreign=self.only_foreign,
                           also_foreign=self.also_foreign)
 
+    @reinitialize_on_error((NoSession, Unauthorized, OpenSubtitlesError, ServiceUnavailable), attempts=1)
     def query(self, video, languages, hash=None, size=None, imdb_id=None, query=None, season=None, episode=None,
               tag=None, use_tag_search=False, only_foreign=False, also_foreign=False):
         # fill the search criteria
@@ -377,6 +379,7 @@ class OpenSubtitlesProvider(ProviderRetryMixin, _OpenSubtitlesProvider):
 
         return subtitles
 
+    @reinitialize_on_error((NoSession, Unauthorized, OpenSubtitlesError, ServiceUnavailable), attempts=1)
     def download_subtitle(self, subtitle):
         logger.info('Downloading subtitle %r', subtitle)
         response = self.use_token_or_login(
diff --git a/libs/subliminal_patch/providers/opensubtitlescom.py b/libs/subliminal_patch/providers/opensubtitlescom.py
index 94e041ab7..e2a86664c 100644
--- a/libs/subliminal_patch/providers/opensubtitlescom.py
+++ b/libs/subliminal_patch/providers/opensubtitlescom.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import logging
 import os
+import time
 import datetime
 
 from requests import Session, ConnectionError, Timeout, ReadTimeout
@@ -147,15 +148,18 @@ class OpenSubtitlesComProvider(ProviderRetryMixin, Provider):
         self.password = password
         self.video = None
         self.use_hash = use_hash
+        self._started = None
 
     def initialize(self):
-        self.token = region.get("oscom_token", expiration_time=TOKEN_EXPIRATION_TIME)
-        if self.token is NO_VALUE:
-            self.login()
+        self._started = time.time()
+        self.login()
 
     def terminate(self):
         self.session.close()
 
+    def ping(self):
+        return self._started and (time.time() - self._started) < TOKEN_EXPIRATION_TIME
+
     def login(self):
         try:
             r = self.session.post(self.server_url + 'login',
diff --git a/libs/subliminal_patch/providers/subscene.py b/libs/subliminal_patch/providers/subscene.py
index 42f0221b0..66329a779 100644
--- a/libs/subliminal_patch/providers/subscene.py
+++ b/libs/subliminal_patch/providers/subscene.py
@@ -20,13 +20,14 @@ import rarfile
 from babelfish import language_converters
 from guessit import guessit
 from dogpile.cache.api import NO_VALUE
+from requests.exceptions import RequestException
 from subliminal import Episode, ProviderError
 from subliminal.video import Episode, Movie
 from subliminal.exceptions import ConfigurationError, ServiceUnavailable
 from subliminal.utils import sanitize_release_group
 from subliminal.cache import region
 from subliminal_patch.http import RetryingCFSession
-from subliminal_patch.providers import Provider
+from subliminal_patch.providers import Provider, reinitialize_on_error
 from subliminal_patch.providers.mixins import ProviderSubtitleArchiveMixin
 from subliminal_patch.subtitle import Subtitle, guess_matches
 from subliminal_patch.converters.subscene import language_ids, supported_languages
@@ -315,7 +316,9 @@ class SubsceneProvider(Provider, ProviderSubtitleArchiveMixin):
             return search(*args, **kwargs)
         except requests.HTTPError:
             region.delete("subscene_cookies2")
+            raise
 
+    @reinitialize_on_error((RequestException,), attempts=1)
     def query(self, video):
         subtitles = []
         if isinstance(video, Episode):
diff --git a/libs/subliminal_patch/providers/xsubs.py b/libs/subliminal_patch/providers/xsubs.py
index c7f166390..c23328582 100644
--- a/libs/subliminal_patch/providers/xsubs.py
+++ b/libs/subliminal_patch/providers/xsubs.py
@@ -6,6 +6,7 @@ import re
 from subzero.language import Language
 from guessit import guessit
 from requests import Session
+from requests.exceptions import RequestException
 
 from subliminal.providers import ParserBeautifulSoup, Provider
 from subliminal import __short_version__
@@ -16,6 +17,7 @@ from subliminal.subtitle import Subtitle, fix_line_ending
 from subliminal.utils import sanitize, sanitize_release_group
 from subliminal.video import Episode
 from subliminal_patch.subtitle import guess_matches
+from subliminal_patch.providers import reinitialize_on_error
 
 logger = logging.getLogger(__name__)
 article_re = re.compile(r'^([A-Za-z]{1,3}) (.*)$')
@@ -189,7 +191,8 @@ class XSubsProvider(Provider):
                 break
 
         return int(show_id) if show_id else None
-
+    
+    @reinitialize_on_error((RequestException,), attempts=1)
     def query(self, show_id, series, season, year=None, country=None):
         # get the season list of the show
         logger.info('Getting the season list of show id %d', show_id)
@@ -291,6 +294,7 @@ class XSubsProvider(Provider):
 
         return []
 
+    @reinitialize_on_error((RequestException,), attempts=1)
     def download_subtitle(self, subtitle):
         if isinstance(subtitle, XSubsSubtitle):
             # download the subtitle