diff --git a/_test_unstructured_client/unit/test_streamed_split_cache.py b/_test_unstructured_client/unit/test_streamed_split_cache.py new file mode 100644 index 00000000..fb38fe32 --- /dev/null +++ b/_test_unstructured_client/unit/test_streamed_split_cache.py @@ -0,0 +1,572 @@ +"""Cached split response ownership, retry and cancellation regressions.""" +import asyncio +import io +import gzip +import gc +import threading +from pathlib import Path + +import httpx +import httpcore +import pytest + +from unstructured_client._hooks.custom import request_utils as ru +from unstructured_client._hooks.custom import split_pdf_hook as hook_module +from unstructured_client._hooks.custom.split_pdf_hook import SplitPdfHook + + +class Body(httpx.AsyncByteStream): + def __init__(self, chunks=(), error=None, gate=None): + self.chunks = chunks + self.error = error + self.gate = gate + self.closed = False + + async def __aiter__(self): + for chunk in self.chunks: + yield chunk + if self.gate is not None: + await self.gate.wait() + if self.error is not None: + raise self.error + + async def aclose(self): + self.closed = True + + +def policy(retry=False): + return ru.RetryConfig( + "backoff" if retry else "none", + ru.BackoffStrategy(1, 1, 1, 5, min_attempts=2, absolute_max_elapsed_time_ms=1000), + True, + ) + + +async def call(client, directory, *, retry=False, limiter=None, source=None): + return await ru.call_api_async( + client, httpx.Request("POST", "https://example.test/partition", content=b"input"), + source if source is not None else io.BytesIO(b"input"), + limiter if limiter is not None else asyncio.Semaphore(1), + retry_config=policy(retry), cache_response_dir=str(directory) if directory else None, + ) + + +@pytest.mark.asyncio +async def test_cached_response_preserves_bytes_and_path_contract(tmp_path): + payload = b'{"text":"caf\xc3\xa9"}\n' + b'x' * 150000 + body = Body([payload[:17], payload[17:]]) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body, headers={"Content-Type": "application/x-ndjson"}) + )) as client: + response = await call(client, tmp_path) + assert Path(response.text).read_bytes() == payload + assert body.closed + assert len(list(tmp_path.iterdir())) == 1 + assert not list(tmp_path.glob("*.partial")) + # A server-provided content-type does not mint the SDK's trusted final-file marker. + assert ru.ELEMENTS_FILE_EXTENSION_KEY not in response.extensions + + +@pytest.mark.asyncio +async def test_noncached_response_stays_buffered(tmp_path): + body = Body([b"one", b"two"]) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + response = await call(client, None) + assert response.content == b"onetwo" + assert body.closed + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_body_retry_uses_fresh_partial_and_closes_each_response(tmp_path): + bodies = [] + observed = [] + def handler(request): + observed.append(list(tmp_path.iterdir())) + body = Body([b"partial"], httpx.ReadError("broken")) if not bodies else Body([b"complete"]) + bodies.append(body) + return httpx.Response(200, stream=body) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + response = await call(client, tmp_path, retry=True) + assert len(bodies) == 2 + assert observed == [[], []] + assert all(body.closed for body in bodies) + assert Path(response.text).read_bytes() == b"complete" + assert len(list(tmp_path.iterdir())) == 1 + + +@pytest.mark.asyncio +async def test_persistent_body_failure_leaves_no_cache(tmp_path): + bodies = [] + def handler(request): + body = Body([b"partial"], httpx.ReadTimeout("stalled")) + bodies.append(body) + return httpx.Response(200, stream=body) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + with pytest.raises(httpx.ReadTimeout): + await call(client, tmp_path, retry=True) + assert len(bodies) >= 2 + assert all(body.closed for body in bodies) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_cancellation_settles_repeatedly_cancelled_write_before_cleanup(tmp_path, monkeypatch): + current = asyncio.current_task() + all_tasks = asyncio.all_tasks + monkeypatch.setattr(hook_module.asyncio, "all_tasks", lambda: all_tasks() - {current}) + entered, release = threading.Event(), threading.Event() + original = ru._write_all + def slow_write(file, content): + entered.set() + assert release.wait(3) + original(file, content) + monkeypatch.setattr(ru, "_write_all", slow_write) + body = Body([b"x" * 65536]) + limiter = asyncio.Semaphore(1) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + task = asyncio.create_task(call(client, tmp_path, limiter=limiter)) + try: + assert await asyncio.to_thread(entered.wait, 2) + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + assert not task.done() + assert limiter.locked() + assert list(tmp_path.glob("*.partial")) + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert body.closed + assert not limiter.locked() + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_disk_failure_is_not_network_retried(tmp_path, monkeypatch): + bodies = [] + def handler(request): + body = Body([b"content"]) + bodies.append(body) + return httpx.Response(200, stream=body) + def failed_write(file, content): + raise OSError("disk full") + monkeypatch.setattr(ru, "_write_all", failed_write) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + with pytest.raises(OSError, match="disk full"): + await call(client, tmp_path, retry=True) + assert len(bodies) == 1 + assert bodies[0].closed + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_file_close_failure_cannot_publish(tmp_path, monkeypatch): + original = ru.os.fdopen + class FailedClose: + def __init__(self, file): self.file = file + @property + def closed(self): return self.file.closed + def write(self, data): return self.file.write(data) + def close(self): + self.file.close() + raise OSError("close failed") + monkeypatch.setattr(ru.os, "fdopen", lambda *a, **kw: FailedClose(original(*a, **kw))) + body = Body([b"content"]) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + with pytest.raises(OSError, match="close failed"): + await call(client, tmp_path) + assert body.closed + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_response_close_failure_cannot_publish(tmp_path): + class FailedClose(Body): + async def aclose(self): + self.closed = True + raise OSError("response close failed") + body = FailedClose([b"content"]) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + with pytest.raises(OSError, match="response close failed"): + await call(client, tmp_path) + assert body.closed + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry", [False, True]) +async def test_503_body_readable_and_retry_unchanged(tmp_path, retry): + bodies = [] + def handler(request): + status = 503 if not bodies else 200 + body = Body([b"try later" if status == 503 else b"done"]) + bodies.append(body) + return httpx.Response(status, stream=body) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + response = await call(client, tmp_path, retry=retry) + if retry: + assert len(bodies) == 2 + assert Path(response.text).read_bytes() == b"done" + else: + assert response.status_code == 503 + assert response.content == b"try later" + assert not list(tmp_path.iterdir()) + assert all(body.closed for body in bodies) + + +@pytest.mark.asyncio +async def test_admission_is_held_until_body_and_cache_write_finish(tmp_path, monkeypatch): + gate = asyncio.Event() + body = Body([b"a" * 65536], gate=gate) + entered, release = threading.Event(), threading.Event() + original = ru._write_all + def slow_write(file, content): + entered.set() + assert release.wait(3) + original(file, content) + monkeypatch.setattr(ru, "_write_all", slow_write) + limiter = asyncio.Semaphore(1) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + task = asyncio.create_task(call(client, tmp_path, limiter=limiter)) + try: + assert await asyncio.to_thread(entered.wait, 2) + assert limiter.locked() + release.set() + await asyncio.sleep(0) + assert limiter.locked() + gate.set() + result = await task + finally: + release.set() + gate.set() + assert not limiter.locked() + assert Path(result.text).read_bytes() == b"a" * 65536 + + +@pytest.mark.asyncio +async def test_hook_cached_operations_keep_separate_paths_and_close_input(tmp_path): + hook = SplitPdfHook() + dirs = [tmp_path / "one", tmp_path / "two"] + for directory in dirs: directory.mkdir() + sources = [open(tmp_path / f"input-{i}", "w+b") for i in range(2)] + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=Body([request.content])) + )) as client: + results = await asyncio.gather(*(hook.call_api_partial( + pdf_chunk_request=httpx.Request("POST", "https://example.test", content=str(i).encode()), + pdf_chunk_file=sources[i], limiter=asyncio.Semaphore(1), _operation_id=str(i), + chunk_index=1, page_number=1, async_client=client, retry_config=policy(), + cache_tmp_data_feature=True, temp_dir_path=str(dirs[i]), + ) for i in range(2))) + assert [Path(r.text).read_bytes() for r in results] == [b"0", b"1"] + assert all(source.closed for source in sources) + assert all(Path(result.text).parent == directory for result, directory in zip(results, dirs)) + + +@pytest.mark.asyncio +async def test_missing_cache_directory_closes_owned_input_without_request(tmp_path): + hook = SplitPdfHook() + source = open(tmp_path / "input", "w+b") + def unexpected_request(request): + raise AssertionError("invalid cache directory must fail before send") + async with httpx.AsyncClient(transport=httpx.MockTransport(unexpected_request)) as client: + with pytest.raises(RuntimeError, match="Temp directory"): + await hook.call_api_partial( + pdf_chunk_request=httpx.Request("POST", "https://example.test"), + pdf_chunk_file=source, limiter=asyncio.Semaphore(1), _operation_id="missing", + chunk_index=1, page_number=1, async_client=client, retry_config=policy(), + cache_tmp_data_feature=True, temp_dir_path=None, + ) + assert source.closed + + +@pytest.mark.asyncio +async def test_cancellation_during_httpx_automatic_close_settles_before_cleanup(tmp_path, monkeypatch): + current = asyncio.current_task() + all_tasks = asyncio.all_tasks + monkeypatch.setattr(hook_module.asyncio, "all_tasks", lambda: all_tasks() - {current}) + entered = asyncio.Event() + release = asyncio.Event() + class SlowClose(Body): + async def aclose(self): + entered.set() + await release.wait() + self.closed = True + body = SlowClose([b"content"]) + limiter = asyncio.Semaphore(1) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + task = asyncio.create_task(call(client, tmp_path, limiter=limiter)) + await asyncio.wait_for(entered.wait(), 2) + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + assert not task.done() + assert limiter.locked() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert body.closed + assert not list(tmp_path.iterdir()) + assert not limiter.locked() + + +@pytest.mark.asyncio +async def test_compressed_cached_response_decodes_entity_only(tmp_path): + payload = b'{"text":"hello"}\n' + body = Body([gzip.compress(payload)]) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body, headers={"Content-Encoding": "gzip"}) + )) as client: + response = await call(client, tmp_path) + assert Path(response.text).read_bytes() == payload + assert "content-encoding" not in response.headers + + +@pytest.mark.asyncio +async def test_cleanup_close_error_preserves_body_retry(tmp_path, monkeypatch): + original = ru.os.fdopen + opened = [] + class FailedClose: + def __init__(self, file): + self.file = file + @property + def closed(self): + return self.file.closed + def write(self, content): + return self.file.write(content) + def close(self): + self.file.close() + raise OSError("secondary close failure") + def fdopen(*args, **kwargs): + file = original(*args, **kwargs) + opened.append(file) + return FailedClose(file) if len(opened) == 1 else file + monkeypatch.setattr(ru.os, "fdopen", fdopen) + bodies = [] + def handler(request): + body = Body([b"partial"], httpx.ReadError("broken")) if not bodies else Body([b"complete"]) + bodies.append(body) + return httpx.Response(200, stream=body) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + response = await call(client, tmp_path, retry=True) + assert len(bodies) == 2 + assert all(body.closed for body in bodies) + assert Path(response.text).read_bytes() == b"complete" + assert not list(tmp_path.glob("*.partial")) + + +@pytest.mark.asyncio +async def test_failed_publication_cleans_reserved_final_and_keeps_existing_artifact(tmp_path, monkeypatch): + existing = tmp_path / "existing.json" + existing.write_bytes(b"previous result") + body = Body([b"new result"]) + def failed_replace(source, destination): + assert Path(destination).exists() + assert Path(destination) != existing + raise OSError("publication failure") + monkeypatch.setattr(ru.os, "replace", failed_replace) + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=body) + )) as client: + with pytest.raises(OSError, match="publication failure"): + await call(client, tmp_path, retry=True) + assert body.closed + assert list(tmp_path.iterdir()) == [existing] + assert existing.read_bytes() == b"previous result" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.parametrize("transport_kind", ["equivalent", "httpcore"]) +async def test_iterator_finalizer_close_stays_owned_through_failed_write(tmp_path, monkeypatch, cancel, transport_kind): + current = asyncio.current_task() + all_tasks = asyncio.all_tasks + monkeypatch.setattr(hook_module.asyncio, "all_tasks", lambda: all_tasks() - {current}) + file_close_entered, file_close_release = threading.Event(), threading.Event() + transport_close_entered, transport_close_release = asyncio.Event(), asyncio.Event() + original_fdopen = ru.os.fdopen + class SlowFileClose: + def __init__(self, file): + self.file = file + @property + def closed(self): + return self.file.closed + def write(self, content): + raise OSError("primary disk failure") + def close(self): + file_close_entered.set() + assert file_close_release.wait(3) + self.file.close() + class FinalizerClose(Body): + closing = False + completed = False + closer = None + async def __aiter__(self): + try: + yield b"x" * 65536 + yield b"unread" + finally: + await self.aclose() + async def aclose(self): + if self.closing: + return + self.closing = True + self.closer = asyncio.current_task() + transport_close_entered.set() + await transport_close_release.wait() + self.completed = True + self.closed = True + class NetworkStream(httpcore.AsyncNetworkStream): + def __init__(self, gated): + payload = b"x" * 131072 if gated else b"[]" + self.data = (b"HTTP/1.1 200 OK\r\nContent-Length: " + str(len(payload)).encode() + + b"\r\nConnection: close\r\n\r\n" + payload) + self.gated = gated + self.completed = False + self.closer = None + async def read(self, max_bytes, timeout=None): + result, self.data = self.data[:max_bytes], self.data[max_bytes:] + return result + async def write(self, buffer, timeout=None): + pass + async def start_tls(self, ssl_context, server_hostname=None, timeout=None): + return self + async def aclose(self): + self.closer = asyncio.current_task() + if self.gated: + transport_close_entered.set() + await transport_close_release.wait() + self.completed = True + class Backend(httpcore.AsyncNetworkBackend): + def __init__(self): + self.connections = [] + async def connect_tcp(self, host, port, **kwargs): + stream = NetworkStream(gated=not self.connections) + self.connections.append(stream) + return stream + backend = Backend() + first = FinalizerClose() + if transport_kind == "httpcore": + transport = httpx.AsyncHTTPTransport() + transport._pool = httpcore.AsyncConnectionPool(network_backend=backend) + else: + bodies = [first, Body([b"[]"])] + transport = httpx.MockTransport(lambda request: httpx.Response(200, stream=bodies.pop(0))) + monkeypatch.setattr(ru.os, "fdopen", lambda *a, **kw: SlowFileClose(original_fdopen(*a, **kw))) + limiter = asyncio.Semaphore(1) + async with httpx.AsyncClient(transport=transport) as client: + task = asyncio.create_task(call(client, tmp_path, limiter=limiter)) + try: + assert await asyncio.to_thread(file_close_entered.wait, 2) + if transport_kind == "httpcore": + first = backend.connections[0] + # Give abandoned async generators a chance to initiate transport close. + gc.collect() + for _ in range(20): + await asyncio.sleep(0) + file_close_release.set() + await asyncio.wait_for(transport_close_entered.wait(), 2) + for _ in range(20): + await asyncio.sleep(0) + assert not task.done(), "attempt released admission before actual transport close" + assert limiter.locked() + assert ru._is_settling_io_task(first.closer) + if cancel: + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + hook_module._cancel_running_tasks() + await asyncio.sleep(0) + assert not task.done() + assert not first.closer.cancelled() + transport_close_release.set() + with pytest.raises(asyncio.CancelledError if cancel else OSError, match=None if cancel else "primary disk failure"): + await task + assert first.completed + assert not limiter.locked() + assert not list(tmp_path.iterdir()) + monkeypatch.setattr(ru.os, "fdopen", original_fdopen) + response = await call(client, tmp_path, limiter=limiter) + assert Path(response.text).read_bytes() == b"[]" + if transport_kind == "httpcore": + assert len(backend.connections) == 2 + assert all(connection.completed for connection in backend.connections) + assert not transport._pool.connections + finally: + file_close_release.set() + transport_close_release.set() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream_close_fails", [False, True]) +async def test_iterator_close_failure_cannot_publish_or_mask_stream_close(tmp_path, stream_close_fails): + class Iterator: + yielded = False + close_attempted = False + def __aiter__(self): + return self + async def __anext__(self): + if self.yielded: + raise StopAsyncIteration + self.yielded = True + return b"complete" + async def aclose(self): + self.close_attempted = True + raise ValueError("secondary iterator close failure") + class Stream(httpx.AsyncByteStream): + iterator = Iterator() + close_attempted = False + def __aiter__(self): + return self.iterator + async def aclose(self): + self.close_attempted = True + if stream_close_fails: + raise OSError("primary stream close failure") + stream = Stream() + response = httpx.Response(200, stream=ru._SettledResponseStream(stream)) + with pytest.raises(OSError if stream_close_fails else ValueError, + match="primary stream close failure" if stream_close_fails else "secondary iterator close failure"): + await ru._cache_streamed_response(response, str(tmp_path)) + assert stream.close_attempted + assert stream.iterator.close_attempted + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_self_iterating_stream_is_closed_once(tmp_path): + class Stream(httpx.AsyncByteStream): + yielded = False + closes = 0 + def __aiter__(self): + return self + async def __anext__(self): + if self.yielded: + raise StopAsyncIteration + self.yielded = True + return b"complete" + async def aclose(self): + self.closes += 1 + assert self.closes == 1 + stream = Stream() + wrapped = ru._SettledResponseStream(stream) + response = httpx.Response(200, stream=wrapped) + cached = await ru._cache_streamed_response(response, str(tmp_path)) + await wrapped.aclose() + assert stream.closes == 1 + assert Path(cached.text).read_bytes() == b"complete" diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index 96382d4a..d37939dc 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -6,7 +6,9 @@ import logging import os import tempfile -from typing import Tuple, Any, BinaryIO, Optional, TextIO +import sys +import weakref +from typing import Tuple, Any, AsyncIterator, Awaitable, BinaryIO, Optional, TextIO, TypeVar import httpx from httpx._multipart import DataField, FileField @@ -186,6 +188,155 @@ def create_pdf_chunk_request( ) +_IOResult = TypeVar("_IOResult") +_settling_tasks: weakref.WeakSet[asyncio.Future[Any]] = weakref.WeakSet() + + +def _is_settling_io_task(task: asyncio.Future[Any]) -> bool: + """The sync worker may cancel request tasks, but must let owned I/O settle.""" + return task in _settling_tasks + + +async def _settle_io(awaitable: Awaitable[_IOResult]) -> _IOResult: + """Finish an owned I/O operation before cancellation can release its resources.""" + task = asyncio.ensure_future(awaitable) + _settling_tasks.add(task) + try: + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + except Exception: + break + if not task.cancelled(): + task.exception() + raise + finally: + _settling_tasks.discard(task) + + +class _SettledResponseStream(httpx.AsyncByteStream): + """Protect HTTPX's automatic EOF close as well as explicit cleanup.""" + + def __init__(self, stream: httpx.AsyncByteStream): + self._stream = stream + self._iterator: Optional[AsyncIterator[bytes]] = None + self._close_task: Optional[asyncio.Task[None]] = None + + async def __aiter__(self): + # Retain the transport iterator when a cache write interrupts iteration. + # Otherwise async-generator finalization can start an unowned close that + # sets closed flags before its actual transport cleanup has settled. + if self._iterator is None: + self._iterator = self._stream.__aiter__() + async for chunk in self._iterator: + yield chunk + + async def _close(self) -> None: + try: + await self._stream.aclose() + finally: + preserving_error = sys.exc_info()[0] is not None + iterator = self._iterator + self._iterator = None + if iterator is not None and iterator is not self._stream and hasattr(iterator, "aclose"): + try: + await iterator.aclose() + except Exception: + if not preserving_error: + raise + logger.warning("Failed to finalize a failed response iterator", exc_info=True) + + async def aclose(self) -> None: + # Every close entry point joins the same operation, including HTTPX's + # automatic EOF close, which sets response.is_closed before awaiting us. + if self._close_task is None: + self._close_task = asyncio.create_task(self._close()) + await _settle_io(self._close_task) + + +def _write_all(file: BinaryIO, content: bytes) -> None: + remaining = memoryview(content) + while remaining: + written = file.write(remaining) + if not written: + raise OSError("Cache write made no progress") + remaining = remaining[written:] + + +async def _cache_streamed_response(response: httpx.Response, directory: str) -> httpx.Response: + """Publish a path only after the HTTP entity and owned I/O close cleanly.""" + path: Optional[str] = None + final_path: Optional[str] = None + file: Optional[BinaryIO] = None + published = False + try: + fd, path = tempfile.mkstemp(prefix="chunk-", suffix=".partial", dir=directory) + try: + file = os.fdopen(fd, "wb", buffering=0) + except BaseException: + os.close(fd) + raise + async for content in response.aiter_bytes(chunk_size=64 * 1024): + await _settle_io(asyncio.get_running_loop().run_in_executor(None, _write_all, file, content)) + await _settle_io(asyncio.get_running_loop().run_in_executor(None, file.close)) + file = None + await _settle_io(response.aclose()) + # Reserve the final name exclusively: a reused partial-file stem must + # never overwrite a previously returned cache artifact. No final name + # exists until the entity, file and response have all closed cleanly. + final_fd, final_path = tempfile.mkstemp(prefix="chunk-", suffix=".json", dir=directory) + os.close(final_fd) + cached = httpx.Response( + status_code=response.status_code, + # Entity bytes are decoded before caching. The synthetic payload is + # a filename, so wire compression/framing headers no longer apply. + headers={key: value for key, value in response.headers.items() + if key.lower() not in {"content-encoding", "content-length", "transfer-encoding"}}, + content=final_path.encode(), + extensions=response.extensions, + ) + os.replace(path, final_path) + published = True + return cached + finally: + preserving_error = sys.exc_info()[0] is not None + try: + try: + if file is not None and not file.closed: + await _settle_io(asyncio.get_running_loop().run_in_executor(None, file.close)) + except Exception: + if not preserving_error: + raise + logger.warning("Failed to close a failed chunk cache", exc_info=True) + finally: + try: + try: + if not response.is_closed: + await _settle_io(response.aclose()) + except Exception: + if not preserving_error: + raise + logger.warning("Failed to close a failed chunk response", exc_info=True) + finally: + if not published: + for failed_path in (path, final_path): + if failed_path is None: + continue + try: + os.unlink(failed_path) + except FileNotFoundError: + pass + except OSError: + if not preserving_error: + raise + logger.warning("Failed to remove a failed chunk cache", exc_info=True) + + async def call_api_async( client: httpx.AsyncClient, pdf_chunk_request: httpx.Request, @@ -195,12 +346,27 @@ async def call_api_async( operation_id: Optional[str] = None, chunk_index: Optional[int] = None, page_number: Optional[int] = None, + cache_response_dir: Optional[str] = None, ) -> httpx.Response: retryable_codes = ["5xx"] effective_retry_config = create_split_retry_config(retry_config) async def do_request(): - return await client.send(pdf_chunk_request) + if cache_response_dir is None: + return await client.send(pdf_chunk_request) + response = await client.send(pdf_chunk_request, stream=True) + assert isinstance(response.stream, httpx.AsyncByteStream) + response.stream = _SettledResponseStream(response.stream) + if response.status_code == 200: + # The retry envelope includes body reads and the admission permit + # remains held through cache publication and response close. + return await _cache_streamed_response(response, cache_response_dir) + try: + # Error bodies remain available to retry classification and callers. + await response.aread() + return response + finally: + await _settle_io(response.aclose()) async with limiter: try: diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index cafa29dc..d0ffceb3 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -17,7 +17,6 @@ from pathlib import Path from typing import Any, Coroutine, Optional, Tuple, Union, cast, Generator, BinaryIO -import aiofiles import httpx from httpx import AsyncClient from pypdf import PdfReader, PdfWriter @@ -446,7 +445,8 @@ def _create_transport_error_response(error: BaseException) -> httpx.Response: def _cancel_running_tasks() -> None: for task in asyncio.all_tasks(): - if not task.done(): + # Paired private hooks share ownership of settling response I/O. + if not task.done() and not request_utils._is_settling_io_task(task): # pylint: disable=protected-access task.cancel() @@ -1033,6 +1033,10 @@ async def call_api_partial( page_number, "cached" if cache_tmp_data_feature else "memory", ) + if cache_tmp_data_feature and temp_dir_path is None: + if not isinstance(pdf_chunk_file, io.BytesIO) and not pdf_chunk_file.closed: + pdf_chunk_file.close() + raise RuntimeError("Temp directory path not found for cached split PDF operation") response = await request_utils.call_api_async( client=async_client, limiter=limiter, @@ -1042,20 +1046,14 @@ async def call_api_partial( operation_id=_operation_id, chunk_index=chunk_index, page_number=page_number, + cache_response_dir=temp_dir_path if cache_tmp_data_feature else None, ) if response.status_code == 200: if cache_tmp_data_feature: - if temp_dir_path is None: - raise RuntimeError("Temp directory path not found for cached split PDF operation") - # If we get 200, dump the contents to a file and return the path - temp_file_name = f"{temp_dir_path}/{uuid.uuid4()}.json" - async with aiofiles.open(temp_file_name, mode='wb') as temp_file: - # Avoid reading the entire response into memory - async for bytes_chunk in response.aiter_bytes(): - await temp_file.write(bytes_chunk) - # we save the path in content attribute to be used in after_success - response._content = temp_file_name.encode() # pylint: disable=protected-access + # call_api_async consumed and closed the response under its + # semaphore/retry scope and returned the operation-owned path. + temp_file_name = response.text logger.debug( "split_pdf event=chunk_cached operation_id=%s chunk_index=%d page_number=%d cache_file=%s", _operation_id, @@ -1063,20 +1061,7 @@ async def call_api_partial( page_number, Path(temp_file_name).name, ) - response = httpx.Response( - status_code=response.status_code, - headers=response.headers, - content=temp_file_name.encode(), - extensions=response.extensions, - ) - else: - response = httpx.Response( - status_code=response.status_code, - headers=response.headers, - content=response.content, - extensions=response.extensions, - ) - else: + if not cache_tmp_data_feature: response = httpx.Response( status_code=response.status_code, headers=response.headers,