kenzuya/unshackle
1
0
Fork 2
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-11-11 19:45:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(downloader): keep retrying while a cut connection makes progress

Browse Source 
Reset the attempt budget when a resumed cut delivers bytes, close the stream in a finally, skip the ranged path for a finished file, and add a throttle-cut profile to fault_server.
This commit is contained in:
imSp4rky
2026-08-24 14:32:03 -06:00
parent d9d6115054
commit e5e1ce14c6
2 changed files with 135 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/fault_server.py
+62 -8
View File
@@ -15,6 +15,7 @@ Fault profiles:
rate-limit reply 429 + Retry-After during recurring throttle windows
(window_open seconds serving / window_closed seconds throttled)
tail-slow heavily throttle the last K segments' bodies
throttle-cut throttle every body and close it (FIN) after N seconds or N bytes
flaky-first fail attempt 1 per path (503), succeed on retry
loopback no faults; payloads served from memory (overhead ceiling)
@@ -24,6 +25,7 @@ download paths exercise correctly, and an unsatisfiable resume gets a 416.
Standalone:
uv run python scripts/fault_server.py --profile rate-limit --segments 32
uv run python scripts/fault_server.py --profile stall --stall-secs 4 --port 8080
uv run python scripts/fault_server.py --profile throttle-cut --segments 1 --seg-size 71303168
Importable: build a ``FaultServer``, serve it on a thread, and hash output against
``segment_payload(path, size)``.
@@ -74,6 +76,8 @@ class FaultProfile:
tail_rate_kib: int = 256 # throttle rate for tail-slow segments (KiB/s)
flaky_first: bool = False # fail attempt 1 per path, succeed after
flaky_status: int = 503 # status returned by a flaky-first failure
cut_secs: float = 0.0 # close the body cleanly (FIN) after this long (0 = never)
cut_bytes: int = 0 # close the body cleanly (FIN) after this many bytes (0 = never)
PROFILES: dict[str, FaultProfile] = {
@@ -88,6 +92,10 @@ PROFILES: dict[str, FaultProfile] = {
),
"tail-slow": FaultProfile("tail-slow", tail_slow=4, tail_rate_kib=256),
"flaky-first": FaultProfile("flaky-first", flaky_first=True, flaky_status=503),
# a CDN that throttles one flow and drops it on a timer: every connection delivers bytes
# and then ends short, so each retry makes progress but no single attempt finishes. Not
# seed-gated like stall and reset, because the real host cuts every attempt, not a share.
"throttle-cut": FaultProfile("throttle-cut", body_rate_kib=480, cut_secs=6.5, cut_bytes=5 * MIB // 2),
}
@@ -235,7 +243,14 @@ class FaultHandler(BaseHTTPRequestHandler):
rate = prof.tail_rate_kib
stall = server.decides("stall", self.path, attempt, prof.stall_pct)
reset = server.decides("reset", self.path, attempt, prof.reset_pct)
self._serve(server.payload(self.path), stall=stall, reset=reset, rate_kib=rate)
self._serve(
server.payload(self.path),
stall=stall,
reset=reset,
rate_kib=rate,
cut_secs=prof.cut_secs,
cut_bytes=prof.cut_bytes,
)
def _send_status(self, code: int) -> None:
server: FaultServer = self.server # type: ignore[assignment]
@@ -252,7 +267,15 @@ class FaultHandler(BaseHTTPRequestHandler):
self.send_header("Content-Length", "0")
self.end_headers()
def _serve(self, full: bytes, stall: bool = False, reset: bool = False, rate_kib: int = 0) -> None:
def _serve(
self,
full: bytes,
stall: bool = False,
reset: bool = False,
rate_kib: int = 0,
cut_secs: float = 0.0,
cut_bytes: int = 0,
) -> None:
server: FaultServer = self.server # type: ignore[assignment]
rng = self._parse_range(len(full))
if rng is not None and rng[0] >= len(full): # unsatisfiable resume: a real CDN answers 416
@@ -294,20 +317,47 @@ class FaultHandler(BaseHTTPRequestHandler):
self.wfile.write(body[half:])
self.wfile.flush()
else:
self._write_body(body, rate_kib)
self._write_body(body, rate_kib, cut_secs, cut_bytes)
except (BrokenPipeError, ConnectionError, OSError):
self.close_connection = True
def _write_body(self, body: bytes, rate_kib: int) -> None:
if rate_kib <= 0:
def _cut(self) -> None:
"""End the body early with a clean half-close.
FIN rather than RST: the client reads fewer bytes than Content-Length promised and
raises IncompleteRead, which is the resumable short read a throttling CDN produces.
An RST (see the reset fault) is a different error class on the client side.
"""
try:
self.wfile.flush()
self.connection.shutdown(socket.SHUT_WR)
except OSError:
pass
self.close_connection = True
def _write_body(self, body: bytes, rate_kib: int, cut_secs: float = 0.0, cut_bytes: int = 0) -> None:
if rate_kib <= 0 and cut_secs <= 0 and cut_bytes <= 0:
self.wfile.write(body)
return
chunk = 64 * 1024 # throttle: fixed-size chunks paced to rate_kib KiB/s
per_chunk = chunk / (rate_kib * 1024)
per_chunk = chunk / (rate_kib * 1024) if rate_kib > 0 else 0.0
deadline = time.monotonic() + cut_secs if cut_secs > 0 else None
sent = 0
for off in range(0, len(body), chunk):
self.wfile.write(body[off : off + chunk])
size = min(chunk, len(body) - off)
if cut_bytes > 0:
size = min(size, cut_bytes - sent)
self.wfile.write(body[off : off + size])
self.wfile.flush()
time.sleep(per_chunk)
sent += size
if cut_bytes > 0 and sent >= cut_bytes:
self._cut()
return
if deadline is not None and time.monotonic() >= deadline:
self._cut()
return
if per_chunk:
time.sleep(per_chunk)
def resolve_profile(name: str, overrides: Optional[dict[str, Any]] = None) -> FaultProfile:
@@ -338,6 +388,8 @@ def _build_arg_parser() -> argparse.ArgumentParser:
p.add_argument("--body-rate-kib", type=int, help="Override: throttle every body to this KiB/s (0 = off).")
p.add_argument("--tail-slow", type=int, help="Override: throttle the last K segments.")
p.add_argument("--tail-rate-kib", type=int, help="Override: tail-slow throttle rate (KiB/s).")
p.add_argument("--cut-secs", type=float, help="Override: close each body (FIN) after this long (0 = never).")
p.add_argument("--cut-bytes", type=int, help="Override: close each body (FIN) after this many bytes (0 = never).")
return p
@@ -353,6 +405,8 @@ def main() -> None:
"body_rate_kib": args.body_rate_kib,
"tail_slow": args.tail_slow,
"tail_rate_kib": args.tail_rate_kib,
"cut_secs": args.cut_secs,
"cut_bytes": args.cut_bytes,
}
profile = resolve_profile(args.profile, overrides)
server = FaultServer(profile, args.seg_size, args.segments, args.seed, port=args.port)
unshackle/core/downloaders/requests.py
+73 -48
View File
@@ -29,6 +29,10 @@ from unshackle.core.utilities import get_debug_logger, get_extension
MAX_ATTEMPTS = 5
RETRY_WAIT = 2
# a cut that still delivered bytes is resumable and refreshes the attempt budget; bound it
# so a server that sends a little and closes, over and over, still terminates
MAX_RESUMES = 500
MIN_RESUME_PROGRESS = 64 * 1024
PROGRESS_WINDOW = 2
# read timeout bounds the gap between chunks so a quiet connection errors instead of hanging
@@ -710,6 +714,7 @@ def download(
attempts = 1
written = 0
resumes = 0
while True:
if claimed is not None and claimed():
return
@@ -719,6 +724,7 @@ def download(
return
if not part_mode:
written = 0
secured = written if part_mode else resume_offset
last_speed_refresh = _time()
try:
@@ -791,10 +797,16 @@ def download(
content_range = parse_content_range(
stream.headers.get("Content-Range") or stream.headers.get("content-range")
)
if content_range is None or content_range[0] != part_offset + written or content_range[1] > part_end:
if (
content_range is None
or content_range[0] != part_offset + written
or content_range[1] > part_end
or (content_range[2] is not None and content_range[2] <= part_end)
):
# a wrong start would land bytes at the wrong offset in the shared file,
# and an end past part_end would spill into a sibling part's region; fail
# the attempt before any byte is written
# an end past part_end would spill into a sibling part's region, and a
# total at or below part_end means the resource changed under the part
# plan; fail the attempt before any byte is written
raise IOError(f"part {part_offset}-{part_end} got Content-Range {content_range!r}")
if use_rnet:
content_encoded = is_content_encoded(
@@ -855,6 +867,7 @@ def download(
chunk_size = adaptive_chunk_size(content_length)
if limiter:
chunk_size = min(chunk_size, max(8192, int(limiter.rate / 4)))
secured = written if part_mode else resume_offset
total_size = (resume_offset + content_length) if resumed and content_length > 0 else content_length
if not segmented and not part_mode:
@@ -901,47 +914,46 @@ def download(
emit_progress = (not segmented) or part_mode
if register is not None:
register(stream, True)
for chunk in chunks:
if DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()):
break
if claimed is not None and claimed():
# close the handle or Windows can't delete the stray .!dev at merge (WinError 32)
if register is not None:
register(stream, False)
try:
stream.close()
except Exception:
pass
return
if limiter is not None:
limiter.consume(len(chunk))
_write(chunk)
download_size = len(chunk)
written += download_size
if emit_progress:
_bytes_since_yield += download_size
_data_accumulated += download_size
now = _time()
time_since = now - last_speed_refresh
if time_since > PROGRESS_WINDOW:
yield dict(advance=_bytes_since_yield)
_bytes_since_yield = 0
if not part_mode:
download_speed = math.ceil(_data_accumulated / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
_data_accumulated = 0
if emit_progress and _bytes_since_yield > 0:
yield dict(advance=_bytes_since_yield)
if register is not None:
register(stream, False)
# finally, not a trailing close: when the consumer abandons this generator it
# raises GeneratorExit at a yield, and the leaked stream then holds a
# pool_block=True connection forever, plus the .!dev handle on Windows. The
# same finally closes a superseded hedge loser, or Windows can't delete its
# stray .!dev at merge (WinError 32)
try:
stream.close()
except Exception:
pass
for chunk in chunks:
if DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()):
break
if claimed is not None and claimed():
return
if limiter is not None:
limiter.consume(len(chunk))
_write(chunk)
download_size = len(chunk)
written += download_size
if emit_progress:
_bytes_since_yield += download_size
_data_accumulated += download_size
now = _time()
time_since = now - last_speed_refresh
if time_since > PROGRESS_WINDOW:
yield dict(advance=_bytes_since_yield)
_bytes_since_yield = 0
if not part_mode:
download_speed = math.ceil(_data_accumulated / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
_data_accumulated = 0
if emit_progress and _bytes_since_yield > 0:
yield dict(advance=_bytes_since_yield)
finally:
if register is not None:
register(stream, False)
try:
stream.close()
except Exception:
pass
aborted = abort is not None and abort.is_set()
if (DOWNLOAD_CANCELLED.is_set() or aborted) and (not content_length or written < content_length):
@@ -981,14 +993,22 @@ def download(
if not cancelled:
raise
return
if on_retry is not None:
if not part_mode:
resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0
# a CDN that throttles a flow and closes it every few seconds would otherwise
# exhaust MAX_ATTEMPTS partway into a large part; a cut with no progress (or a
# slice item, which rewrites from scratch) still spends an attempt, and only
# those reach on_retry so the adaptive controller does not slow a healthy flow
progress = written if part_mode else (0 if item_range else resume_offset)
if progress - secured >= MIN_RESUME_PROGRESS and resumes < MAX_RESUMES:
resumes += 1
attempts = 0
elif on_retry is not None:
try:
on_retry(exc)
except Exception:
pass
if not part_mode:
resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0
delay = retry_sleep(exc, attempts)
delay = retry_sleep(exc, max(1, attempts))
if abort is not None:
# interruptible nap: backoff can reach MAX_BACKOFF, and teardown always
# sets the batch abort, so a parked worker must wake and exit at once
@@ -1462,7 +1482,12 @@ def requests(
url_item = urls[0]
try:
ranged_used = False
if max_workers > 1 and speed_limiter is None and not has_range_header(url_item):
if (
max_workers > 1
and speed_limiter is None
and not has_range_header(url_item)
and not url_item["save_path"].exists()
):
total_size, supports_ranges = probe_ranged(url_item["url"], session)
if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE:
try: