wip

tools/lib/url_file.py

@@ -1,11 +1,10 @@
-import logging
 import os
-import socket
 import time
+import threading
 from hashlib import sha256
-from urllib3 import PoolManager, Retry
-from urllib3.response import BaseHTTPResponse
+from urllib3 import PoolManager
 from urllib3.util import Timeout
+from tenacity import retry, wait_random_exponential, stop_after_attempt
 
 from openpilot.common.file_helpers import atomic_write_in_dir
 from openpilot.system.hardware.hw import Paths
@@ -13,9 +12,8 @@ from openpilot.system.hardware.hw import Paths
 K = 1000
 CHUNK_SIZE = 1000 * K
 
-logging.getLogger("urllib3").setLevel(logging.WARNING)
 
-def hash_256(link: str) -> str:
+def hash_256(link):
   hsh = str(sha256((link.split("?")[0]).encode('utf-8')).hexdigest())
   return hsh
 
@@ -25,25 +23,13 @@ class URLFileException(Exception):
 
 
 class URLFile:
-  _pool_manager: PoolManager|None = None
+  _tlocal = threading.local()
 
-  @staticmethod
-  def reset() -> None:
-    URLFile._pool_manager = None
-
-  @staticmethod
-  def pool_manager() -> PoolManager:
-    if URLFile._pool_manager is None:
-      socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),]
-      retries = Retry(total=5, backoff_factor=0.5, status_forcelist=[409, 429, 503, 504])
-      URLFile._pool_manager = PoolManager(num_pools=10, maxsize=100, socket_options=socket_options, retries=retries)
-    return URLFile._pool_manager
-
-  def __init__(self, url: str, timeout: int=10, debug: bool=False, cache: bool|None=None):
+  def __init__(self, url, debug=False, cache=None):
     self._url = url
-    self._timeout = Timeout(connect=timeout, read=timeout)
     self._pos = 0
-    self._length: int|None = None
+    self._length = None
+    self._local_file = None
     self._debug = debug
     #  True by default, false if FILEREADER_CACHE is defined, but can be overwritten by the cache input
     self._force_download = not int(os.environ.get("FILEREADER_CACHE", "0"))
@@ -53,23 +39,30 @@ class URLFile:
     if not self._force_download:
       os.makedirs(Paths.download_cache_root(), exist_ok=True)
 
+    try:
+      self._http_client = URLFile._tlocal.http_client
+    except AttributeError:
+      self._http_client = URLFile._tlocal.http_client = PoolManager()
+
   def __enter__(self):
     return self
 
-  def __exit__(self, exc_type, exc_value, traceback) -> None:
-    pass
+  def __exit__(self, exc_type, exc_value, traceback):
+    if self._local_file is not None:
+      os.remove(self._local_file.name)
+      self._local_file.close()
+      self._local_file = None
 
-  def _request(self, method: str, url: str, headers: dict[str, str]|None=None) -> BaseHTTPResponse:
-    return URLFile.pool_manager().request(method, url, timeout=self._timeout, headers=headers)
-
-  def get_length_online(self) -> int:
-    response = self._request('HEAD', self._url)
+  @retry(wait=wait_random_exponential(multiplier=1, max=5), stop=stop_after_attempt(3), reraise=True)
+  def get_length_online(self):
+    timeout = Timeout(connect=50.0, read=500.0)
+    response = self._http_client.request('HEAD', self._url, timeout=timeout, preload_content=False)
     if not (200 <= response.status <= 299):
       return -1
     length = response.headers.get('content-length', 0)
     return int(length)
 
-  def get_length(self) -> int:
+  def get_length(self):
     if self._length is not None:
       return self._length
 
@@ -86,7 +79,7 @@ class URLFile:
         file_length.write(str(self._length))
     return self._length
 
-  def read(self, ll: int|None=None) -> bytes:
+  def read(self, ll=None):
     if self._force_download:
       return self.read_aux(ll=ll)
 
@@ -118,9 +111,10 @@ class URLFile:
         self._pos = file_end
         return response
 
-  def read_aux(self, ll: int|None=None) -> bytes:
+  @retry(wait=wait_random_exponential(multiplier=1, max=5), stop=stop_after_attempt(3), reraise=True)
+  def read_aux(self, ll=None):
     download_range = False
-    headers = {}
+    headers = {'Connection': 'keep-alive'}
     if self._pos != 0 or ll is not None:
       if ll is None:
         end = self.get_length() - 1
@@ -134,7 +128,8 @@ class URLFile:
     if self._debug:
       t1 = time.time()
 
-    response = self._request('GET', self._url, headers=headers)
+    timeout = Timeout(connect=50.0, read=500.0)
+    response = self._http_client.request('GET', self._url, timeout=timeout, preload_content=False, headers=headers)
     ret = response.data
 
     if self._debug:
@@ -153,12 +148,9 @@ class URLFile:
     self._pos += len(ret)
     return ret
 
-  def seek(self, pos:int) -> None:
+  def seek(self, pos):
     self._pos = pos
 
   @property
-  def name(self) -> str:
+  def name(self):
     return self._url
-
-
-os.register_at_fork(after_in_child=URLFile.reset)