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 tomllib
from copy import deepcopy
import functools
from http.client import HTTPConnection
import requests
import click
@ -285,9 +287,20 @@ class CustomXMLRPCTransport(xmlrpc.client.Transport):
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:
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)
s.mount("http://", HTTPAdapter(max_retries=retries))
s.mount("https://", HTTPAdapter(max_retries=retries))