make URLFile safe after fork() (#31309)

* make URLFile safe after fork()

* cache the pool manager in each instance

* type hints
pull/31305/head^2
Greg Hogan 1 year ago committed by GitHub
parent 06bdc69a02
commit a8aa04e6bd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 25
      tools/lib/url_file.py

@ -6,6 +6,7 @@ from hashlib import sha256
from urllib3 import PoolManager from urllib3 import PoolManager
from urllib3.util import Timeout from urllib3.util import Timeout
from tenacity import retry, wait_random_exponential, stop_after_attempt from tenacity import retry, wait_random_exponential, stop_after_attempt
from typing import Optional
from openpilot.common.file_helpers import atomic_write_in_dir from openpilot.common.file_helpers import atomic_write_in_dir
from openpilot.system.hardware.hw import Paths from openpilot.system.hardware.hw import Paths
@ -25,9 +26,12 @@ class URLFileException(Exception):
class URLFile: class URLFile:
_tlocal = threading.local() _pid: Optional[int] = None
_pool_manager: Optional[PoolManager] = None
_pool_manager_lock = threading.Lock()
def __init__(self, url, debug=False, cache=None): def __init__(self, url, debug=False, cache=None):
self._pool_manager = None
self._url = url self._url = url
self._pos = 0 self._pos = 0
self._length = None self._length = None
@ -41,11 +45,6 @@ class URLFile:
if not self._force_download: if not self._force_download:
os.makedirs(Paths.download_cache_root(), exist_ok=True) 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): def __enter__(self):
return self return self
@ -55,10 +54,20 @@ class URLFile:
self._local_file.close() self._local_file.close()
self._local_file = None self._local_file = None
def _http_client(self) -> PoolManager:
if self._pool_manager is None:
pid = os.getpid()
with URLFile._pool_manager_lock:
if URLFile._pid != pid or URLFile._pool_manager is None: # unsafe to share after fork
URLFile._pid = pid
URLFile._pool_manager = PoolManager(num_pools=10, maxsize=10)
self._pool_manager = URLFile._pool_manager
return self._pool_manager
@retry(wait=wait_random_exponential(multiplier=1, max=5), stop=stop_after_attempt(3), reraise=True) @retry(wait=wait_random_exponential(multiplier=1, max=5), stop=stop_after_attempt(3), reraise=True)
def get_length_online(self): def get_length_online(self):
timeout = Timeout(connect=50.0, read=500.0) timeout = Timeout(connect=50.0, read=500.0)
response = self._http_client.request('HEAD', self._url, timeout=timeout, preload_content=False) response = self._http_client().request('HEAD', self._url, timeout=timeout, preload_content=False)
if not (200 <= response.status <= 299): if not (200 <= response.status <= 299):
return -1 return -1
length = response.headers.get('content-length', 0) length = response.headers.get('content-length', 0)
@ -131,7 +140,7 @@ class URLFile:
t1 = time.time() t1 = time.time()
timeout = Timeout(connect=50.0, read=500.0) timeout = Timeout(connect=50.0, read=500.0)
response = self._http_client.request('GET', self._url, timeout=timeout, preload_content=False, headers=headers) response = self._http_client().request('GET', self._url, timeout=timeout, preload_content=False, headers=headers)
ret = response.data ret = response.data
if self._debug: if self._debug:

Loading…
Cancel
Save