fix: respect sleep-requests arg in specific code paths.

This commit is contained in:
arabcoders 2026-05-07 21:35:40 +03:00
parent ac3c666d32
commit 9b3cc8b310
9 changed files with 149 additions and 36 deletions

View file

@ -187,6 +187,7 @@ class GenericTaskHandler(BaseHandler):
url=url,
no_archive=True,
no_log=True,
budget_sleep=True,
)
if not info:

View file

@ -185,6 +185,7 @@ class RssGenericHandler(BaseHandler):
url=url,
no_archive=True,
no_log=True,
budget_sleep=True,
)
if not info:

View file

@ -107,6 +107,7 @@ class HandleTask(TaskSchema):
no_archive=True,
follow_redirect=False,
sanitize_info=True,
budget_sleep=True,
)
if not ie_info or not isinstance(ie_info, dict):
@ -133,7 +134,7 @@ class HandleTask(TaskSchema):
archive_file: Path = Path(archive_file)
(ie_info, _) = await fetch_info(params, self.url, no_archive=True, follow_redirect=True)
(ie_info, _) = await fetch_info(params, self.url, no_archive=True, follow_redirect=True, budget_sleep=True)
if not ie_info or not isinstance(ie_info, dict):
return (False, "Failed to extract information from URL.")

View file

@ -49,6 +49,17 @@ class ExtractorConfig:
self.wait_threshold = wait_threshold
def _sleep_timeout(config: dict[str, Any], timeout: float, budget_sleep: bool) -> float:
if not budget_sleep:
return timeout
sleep_requests = config.get("sleep_interval_requests")
if not isinstance(sleep_requests, int | float) or sleep_requests <= 0:
return timeout
return timeout + min(float(sleep_requests) * 20, 300.0)
class ExtractorPool(metaclass=Singleton):
"""
Manages process pool and semaphore for video information extraction.
@ -312,6 +323,7 @@ async def fetch_info(
sanitize_info: bool = False,
capture_logs: int | None = None,
extractor_config: ExtractorConfig | None = None,
budget_sleep: bool = False,
**kwargs,
) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
"""
@ -329,6 +341,7 @@ async def fetch_info(
sanitize_info: Sanitize the extracted information
capture_logs: If provided (e.g., logging.WARNING), capture logs
extractor_config: Configuration for the extractor
budget_sleep: Whether to add extra timeout budget for request-sleep-heavy extraction
**kwargs: Additional arguments
Returns:
@ -352,6 +365,7 @@ async def fetch_info(
loop = asyncio.get_running_loop()
safe_config = _sanitize_config(config)
timeout = _sleep_timeout(safe_config, extractor_config.timeout, budget_sleep)
try:
try:
@ -372,7 +386,7 @@ async def fetch_info(
**kwargs,
),
),
timeout=extractor_config.timeout,
timeout=timeout,
)
except TimeoutError:
@ -396,7 +410,7 @@ async def fetch_info(
**kwargs,
),
),
timeout=extractor_config.timeout,
timeout=timeout,
)
finally:
semaphore.release()

View file

@ -208,6 +208,7 @@ async def add(
no_archive=False,
follow_redirect=True,
capture_logs=logging.WARNING,
budget_sleep=True,
)
if not entry:

View file

@ -792,6 +792,7 @@ async def item_nfo_generate(request: Request, queue: DownloadQueue) -> Response:
url=item.info.url,
no_archive=True,
follow_redirect=True,
budget_sleep=True,
)
if not info_dict:

View file

@ -18,7 +18,7 @@ def reset_config() -> Generator[None, None, None]:
Config._reset_singleton()
class _JsonRequest:
class _Req:
def __init__(self, payload: Any) -> None:
self._payload = payload
self.body_exists = payload is not None
@ -27,65 +27,65 @@ class _JsonRequest:
return self._payload
class _InspectRequest(_JsonRequest):
class _InspectReq(_Req):
query: dict[str, str] = {}
match_info: dict[str, str] = {}
class _QueryRequest:
class _QueryReq:
def __init__(self, query: dict[str, str]) -> None:
self.query = query
def _patch_to_thread(monkeypatch: pytest.MonkeyPatch, module: Any, config: Config, url: str) -> dict[str, bool]:
called = {"to_thread": False, "validate": False}
def _patch_thread(monkeypatch: pytest.MonkeyPatch, module: Any, config: Config, url: str) -> dict[str, bool]:
seen = {"to_thread": False, "validate": False}
def fake_validate_url(next_url: str, allow_internal: bool = False) -> bool:
called["validate"] = True
seen["validate"] = True
assert next_url == url
assert allow_internal is config.allow_internal_urls
raise ValueError("Invalid hostname.")
async def fake_to_thread(func, *args, **kwargs):
called["to_thread"] = True
seen["to_thread"] = True
return func(*args, **kwargs)
monkeypatch.setattr(module, "validate_url", fake_validate_url)
monkeypatch.setattr(module.asyncio, "to_thread", fake_to_thread)
return called
return seen
@pytest.mark.asyncio
async def test_inspect_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_inspect_thread(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
request = _InspectRequest({"url": "https://bad.example/task"})
called = _patch_to_thread(monkeypatch, tasks_router, config, "https://bad.example/task")
request = _InspectReq({"url": "https://bad.example/task"})
seen = _patch_thread(monkeypatch, tasks_router, config, "https://bad.example/task")
response = await tasks_router.task_handler_inspect(request, handler=None, encoder=None, config=config)
assert response.status == 400
assert json.loads(response.body.decode("utf-8")) == {"error": "Invalid hostname."}
assert called == {"to_thread": True, "validate": True}
assert seen == {"to_thread": True, "validate": True}
@pytest.mark.asyncio
async def test_conditions_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_conditions_thread(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
request = _JsonRequest({"url": "https://bad.example/cond", "condition": "title ~= 'x'"})
called = _patch_to_thread(monkeypatch, conditions_router, config, "https://bad.example/cond")
request = _Req({"url": "https://bad.example/cond", "condition": "title ~= 'x'"})
seen = _patch_thread(monkeypatch, conditions_router, config, "https://bad.example/cond")
response = await conditions_router.conditions_test(request, encoder=None, cache=None, config=config)
assert response.status == 400
assert json.loads(response.body.decode("utf-8")) == {"error": "Invalid hostname."}
assert called == {"to_thread": True, "validate": True}
assert seen == {"to_thread": True, "validate": True}
@pytest.mark.asyncio
async def test_info_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_info_thread(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
request = _QueryRequest({"url": "https://bad.example/info"})
called = _patch_to_thread(monkeypatch, ytdlp_router, config, "https://bad.example/info")
request = _QueryReq({"url": "https://bad.example/info"})
seen = _patch_thread(monkeypatch, ytdlp_router, config, "https://bad.example/info")
response = await ytdlp_router.get_info(request, cache=None, config=config)
@ -95,14 +95,14 @@ async def test_info_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None
"message": "Invalid hostname.",
"error": "Invalid hostname.",
}
assert called == {"to_thread": True, "validate": True}
assert seen == {"to_thread": True, "validate": True}
@pytest.mark.asyncio
async def test_archive_ids_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_archive_ids_thread(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
request = _JsonRequest(["https://bad.example/archive"])
called = _patch_to_thread(monkeypatch, ytdlp_router, config, "https://bad.example/archive")
request = _Req(["https://bad.example/archive"])
seen = _patch_thread(monkeypatch, ytdlp_router, config, "https://bad.example/archive")
response = await ytdlp_router.get_archive_ids(request, config)
@ -117,4 +117,4 @@ async def test_archive_ids_validate_off_thread(monkeypatch: pytest.MonkeyPatch)
"error": "Invalid hostname.",
}
]
assert called == {"to_thread": True, "validate": True}
assert seen == {"to_thread": True, "validate": True}

View file

@ -0,0 +1,94 @@
from __future__ import annotations
import asyncio
from unittest.mock import Mock
import pytest
from app.features.ytdlp import extractor
class _Loop:
def __init__(self) -> None:
self.calls: list[object | None] = []
def run_in_executor(self, executor, func): # noqa: ANN001
self.calls.append(executor)
return func
class _Pool:
def __init__(self) -> None:
self.semaphore = asyncio.Semaphore(1)
self.executor = object()
def get_semaphore(self, _config: extractor.ExtractorConfig) -> asyncio.Semaphore:
return self.semaphore
def get_pool(self, _config: extractor.ExtractorConfig) -> object:
return self.executor
def test_sleep_budget() -> None:
assert extractor._sleep_timeout({}, 70, False) == 70
assert extractor._sleep_timeout({"sleep_interval_requests": 0}, 70, True) == 70
assert extractor._sleep_timeout({"sleep_interval_requests": 3}, 70, True) == 130
assert extractor._sleep_timeout({"sleep_interval_requests": 30}, 70, True) == 370
@pytest.mark.asyncio
async def test_timeout_no_retry(monkeypatch: pytest.MonkeyPatch) -> None:
pool = _Pool()
loop = _Loop()
seen: list[float] = []
async def fake_wait_for(*, fut, timeout):
seen.append(timeout)
raise TimeoutError
monkeypatch.setattr(extractor.ExtractorPool, "get_instance", classmethod(lambda cls: pool))
monkeypatch.setattr(extractor.asyncio, "get_running_loop", lambda: loop)
monkeypatch.setattr(extractor.asyncio, "wait_for", fake_wait_for)
with pytest.raises(TimeoutError):
await extractor.fetch_info(
config={"sleep_interval_requests": 3},
url="https://example.com",
extractor_config=extractor.ExtractorConfig(concurrency=1, timeout=70),
budget_sleep=True,
)
assert loop.calls == [pool.executor]
assert seen == [130]
assert not pool.semaphore.locked()
@pytest.mark.asyncio
async def test_pool_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
pool = _Pool()
loop = _Loop()
expected = ({"id": "ok"}, [])
seen: list[float] = []
async def fake_wait_for(*, fut, timeout):
seen.append(timeout)
if len(loop.calls) == 1:
raise RuntimeError("pool failed")
return fut()
monkeypatch.setattr(extractor.ExtractorPool, "get_instance", classmethod(lambda cls: pool))
monkeypatch.setattr(extractor.asyncio, "get_running_loop", lambda: loop)
monkeypatch.setattr(extractor.asyncio, "wait_for", fake_wait_for)
monkeypatch.setattr(extractor, "extract_info_sync", Mock(return_value=expected))
result = await extractor.fetch_info(
config={"sleep_interval_requests": 3},
url="https://example.com",
extractor_config=extractor.ExtractorConfig(concurrency=1, timeout=70),
budget_sleep=True,
)
assert result == expected
assert loop.calls == [pool.executor, None]
assert seen == [130, 130]
assert not pool.semaphore.locked()

View file

@ -18,7 +18,7 @@ def reset_config() -> Generator[None, None, None]:
Config._reset_singleton()
class _Response:
class _Resp:
def __init__(self, *, status_code: int = 200, content: bytes = b"img", content_type: str = "image/jpeg") -> None:
self.status_code = status_code
self.content = content
@ -26,25 +26,25 @@ class _Response:
@pytest.mark.asyncio
async def test_thumbnail_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_thumb_thread(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
req = make_mocked_request("GET", "/api/thumbnail?url=https://example.com/a.jpg")
req._rel_url = req._rel_url.with_query({"url": "https://example.com/a.jpg"})
called = {"to_thread": False, "validate": False}
seen = {"to_thread": False, "validate": False}
def fake_validate_url(url: str, allow_internal: bool = False) -> bool:
called["validate"] = True
seen["validate"] = True
assert url == "https://example.com/a.jpg"
assert allow_internal is config.allow_internal_urls
return True
async def fake_to_thread(func, *args, **kwargs):
called["to_thread"] = True
seen["to_thread"] = True
return func(*args, **kwargs)
client = AsyncMock()
client.request.return_value = _Response()
client.request.return_value = _Resp()
monkeypatch.setattr(images, "validate_url", fake_validate_url)
monkeypatch.setattr(images.asyncio, "to_thread", fake_to_thread)
@ -70,13 +70,13 @@ async def test_thumbnail_validate_off_thread(monkeypatch: pytest.MonkeyPatch) ->
response = await images.get_thumbnail(req, config)
assert response.status == web.HTTPOk.status_code
assert called["to_thread"] is True
assert called["validate"] is True
assert seen["to_thread"] is True
assert seen["validate"] is True
client.request.assert_awaited_once()
@pytest.mark.asyncio
async def test_thumbnail_validate_rejects(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_thumb_reject(monkeypatch: pytest.MonkeyPatch) -> None:
config = Config.get_instance()
req = make_mocked_request("GET", "/api/thumbnail?url=https://bad.example/a.jpg")
req._rel_url = req._rel_url.with_query({"url": "https://bad.example/a.jpg"})