Handling Async Batch Uploads for Large Datasets: Resumable Chunked Transfer with Checkpointing #

from __future__ import annotations

import asyncio
import hashlib
import random
from dataclasses import dataclass
from pathlib import Path

import httpx
from pydantic import BaseModel, Field


class UploadCheckpoint(BaseModel):
    """Crash-safe progress record; the only state that survives a restart."""
    upload_id: str
    file_sha256: str
    chunk_size: int
    total_chunks: int
    committed: list[int] = Field(default_factory=list)  # server-acked sequences

    def save(self, path: Path) -> None:
        path.write_text(self.model_dump_json(indent=2))

    @classmethod
    def load(cls, path: Path) -> "UploadCheckpoint | None":
        if not path.exists():
            return None
        return cls.model_validate_json(path.read_text())


@dataclass
class CircuitBreaker:
    """Opens after `threshold` consecutive server/timeout failures."""
    threshold: int = 5
    cooldown: float = 30.0
    _failures: int = 0
    _opened_at: float | None = None

    def allow(self, now: float) -> bool:
        if self._opened_at is None:
            return True
        if now - self._opened_at >= self.cooldown:  # half-open: allow one probe
            self._opened_at = None
            self._failures = 0
            return True
        return False

    def record_success(self) -> None:
        self._failures = 0
        self._opened_at = None

    def record_failure(self, now: float) -> None:
        self._failures += 1
        if self._failures >= self.threshold:
            self._opened_at = now


def _backoff(attempt: int, base: float = 0.5, cap: float = 30.0) -> float:
    """Exponential backoff with full jitter."""
    return random.uniform(0.0, min(cap, base * (2 ** attempt)))


def _read_chunk(path: Path, seq: int, chunk_size: int) -> bytes:
    """Blocking read of a single chunk by offset — call via asyncio.to_thread."""
    with path.open("rb") as fh:
        fh.seek(seq * chunk_size)
        return fh.read(chunk_size)


async def _put_chunk(
    client: httpx.AsyncClient,
    upload_id: str,
    seq: int,
    data: bytes,
    breaker: CircuitBreaker,
    max_retries: int = 5,
) -> None:
    """Upload one chunk with classification-driven retry; raises only on permanent failure."""
    digest = hashlib.sha256(data).hexdigest()
    headers = {
        "X-Chunk-SHA256": digest,
        # Stable across retries -> the server deduplicates instead of appending.
        "X-Idempotency-Key": f"{upload_id}:{seq}",
    }
    loop = asyncio.get_running_loop()
    for attempt in range(max_retries + 1):
        if not breaker.allow(loop.time()):
            await asyncio.sleep(breaker.cooldown)
        try:
            resp = await client.put(
                f"/uploads/{upload_id}/chunks/{seq}", content=data, headers=headers
            )
        except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout):
            breaker.record_failure(loop.time())
            await asyncio.sleep(_backoff(attempt))
            continue

        if resp.status_code in (200, 201):
            breaker.record_success()
            return
        if resp.status_code == 429:  # honour server-supplied Retry-After
            await asyncio.sleep(float(resp.headers.get("Retry-After", _backoff(attempt))))
            continue
        if 500 <= resp.status_code < 600:
            breaker.record_failure(loop.time())
            await asyncio.sleep(_backoff(attempt))
            continue
        resp.raise_for_status()  # 4xx other than 429 is permanent -> abort
    raise RuntimeError(f"chunk {seq} exhausted {max_retries} retries")


async def upload_large_dataset(
    file_path: Path,
    upload_id: str,
    base_url: str,
    chunk_size: int = 8 * 1024 * 1024,   # 8 MiB
    max_concurrency: int = 4,
    checkpoint_path: Path | None = None,
) -> UploadCheckpoint:
    """Resumable, bounded-memory upload of one large file."""
    checkpoint_path = checkpoint_path or file_path.with_suffix(file_path.suffix + ".ckpt")
    size = file_path.stat().st_size
    total = -(-size // chunk_size)  # ceiling division

    ckpt = UploadCheckpoint.load(checkpoint_path) or UploadCheckpoint(
        upload_id=upload_id,
        file_sha256=await asyncio.to_thread(_whole_file_digest, file_path),
        chunk_size=chunk_size,
        total_chunks=total,
    )
    done: set[int] = set(ckpt.committed)
    breaker = CircuitBreaker()
    semaphore = asyncio.Semaphore(max_concurrency)
    lock = asyncio.Lock()  # guards checkpoint writes across concurrent workers

    async with httpx.AsyncClient(base_url=base_url, timeout=httpx.Timeout(60.0)) as client:
        async def worker(seq: int) -> None:
            async with semaphore:
                data = await asyncio.to_thread(_read_chunk, file_path, seq, chunk_size)
                await _put_chunk(client, upload_id, seq, data, breaker)
                async with lock:  # persist progress ONLY after the ack
                    ckpt.committed.append(seq)
                    ckpt.save(checkpoint_path)

        await asyncio.gather(*(worker(s) for s in range(total) if s not in done))

        resp = await client.post(
            f"/uploads/{upload_id}/complete", json={"sha256": ckpt.file_sha256}
        )
        resp.raise_for_status()

    checkpoint_path.unlink(missing_ok=True)  # remove only after a clean commit
    return ckpt


def _whole_file_digest(path: Path, block: int = 1024 * 1024) -> str:
    sha = hashlib.sha256()
    with path.open("rb") as fh:
        while chunk := fh.read(block):
            sha.update(chunk)
    return sha.hexdigest()

import httpx
import pytest
from pathlib import Path


@pytest.mark.asyncio
async def test_retries_429_then_resumes(tmp_path: Path) -> None:
    src = tmp_path / "dataset.bin"
    src.write_bytes(b"x" * (20 * 1024 * 1024))  # 20 MiB -> 3 chunks at 8 MiB
    attempts: dict[str, int] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path.endswith("/complete"):
            return httpx.Response(201, json={"committed": True})
        key = request.headers["X-Idempotency-Key"]
        attempts[key] = attempts.get(key, 0) + 1
        if attempts[key] == 1:  # first touch of every chunk is throttled
            return httpx.Response(429, headers={"Retry-After": "0"})
        return httpx.Response(200)

    async with httpx.AsyncClient(
        base_url="https://repo.test", transport=httpx.MockTransport(handler)
    ) as client:
        # inject the mocked client by monkeypatching the module's constructor in real code;
        # here we exercise _put_chunk directly for a focused assertion.
        from mymodule import CircuitBreaker, _put_chunk
        await _put_chunk(client, "up-1", 0, b"x" * 8, CircuitBreaker())

    assert attempts["up-1:0"] == 2  # one 429, one success on the SAME sequence