yacloud/shadowmire
1
0
Fork 0
mirror of https://github.com/taoky/shadowmire.git synced 2025-07-08 09:12:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add timeout for requests.Session and XMLRPC

Browse Source
This commit is contained in:
taoky 2024-09-27 18:56:29 +08:00
parent 4664c23ad5
commit 1ca3e23bcf
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
shadowmire.py
View File

@ -22,6 +22,8 @@ from concurrent.futures import Future, ThreadPoolExecutor, as_completed
import signal import signal
import tomllib import tomllib
from copy import deepcopy from copy import deepcopy
import functools
from http.client import HTTPConnection
import requests import requests
import click import click
@ -285,9 +287,20 @@ class CustomXMLRPCTransport(xmlrpc.client.Transport):
user_agent = USER_AGENT user_agent = USER_AGENT
def make_connection(self, host: tuple[str, dict[str, str]] | str) -> HTTPConnection:
conn = super().make_connection(host)
if conn.timeout is None:
# 2 min timeout
conn.timeout = 120
return conn
def create_requests_session() -> requests.Session: def create_requests_session() -> requests.Session:
s = requests.Session() s = requests.Session()
# hardcode 1min timeout for connect & read for now
# https://requests.readthedocs.io/en/latest/user/advanced/#timeouts
# A hack to overwrite get() method
s.get_orig, s.get = s.get, functools.partial(s.get, timeout=(60, 60)) # type: ignore
retries = Retry(total=3, backoff_factor=0.1) retries = Retry(total=3, backoff_factor=0.1)
s.mount("http://", HTTPAdapter(max_retries=retries)) s.mount("http://", HTTPAdapter(max_retries=retries))
s.mount("https://", HTTPAdapter(max_retries=retries)) s.mount("https://", HTTPAdapter(max_retries=retries))