diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b0b81b..9353145 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,9 +4,11 @@ on: push: branches: [main] pull_request: - branches: [main] workflow_call: # Allow release.yml to call this workflow +env: + UV_PYTHON: "3.11" + jobs: lint: name: Lint @@ -23,7 +25,7 @@ jobs: run: uv python install 3.11 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --extra dev --extra all - name: Run ruff check run: uv run ruff check . @@ -46,13 +48,15 @@ jobs: run: uv python install 3.11 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --extra dev --extra all - name: Run mypy run: uv run mypy src/arc_client --ignore-missing-imports test: name: Test (Python ${{ matrix.python-version }}) + env: + UV_PYTHON: ${{ matrix.python-version }} runs-on: ubuntu-latest strategy: fail-fast: false @@ -71,10 +75,12 @@ jobs: run: uv python install ${{ matrix.python-version }} - name: Install dependencies - run: uv sync --extra dev + run: uv sync --extra dev --extra all - name: Run tests - run: uv run pytest tests/unit -v --tb=short + run: | + uv run python --version + uv run pytest tests/unit -v --tb=short test-platforms: name: Test (${{ matrix.os }}) @@ -96,10 +102,12 @@ jobs: run: uv python install 3.11 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --extra dev --extra all - name: Run tests - run: uv run pytest tests/unit -v --tb=short + run: | + uv run python --version + uv run pytest tests/unit -v --tb=short coverage: name: Coverage @@ -116,7 +124,7 @@ jobs: run: uv python install 3.11 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --extra dev --extra all - name: Run tests with coverage run: uv run pytest tests/unit --cov=src/arc_client --cov-report=xml diff --git a/CHANGELOG.md b/CHANGELOG.md index 447d024..ab5948b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ All notable changes to `arc-tsdb-client` are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [0.2.0] - 2026-07-03 +## [0.2.0] - Unreleased A large feature-parity release bringing the SDK up to date with the Arc server, plus bug fixes and the first real HTTP-layer test coverage. @@ -28,9 +28,10 @@ plus bug fixes and the first real HTTP-layer test coverage. `mqtt`. - **zstd write compression** (now the default; gzip and none still supported). Compression is configurable via a `Compression` enum, a string, or a bool. -- **Automatic retries** with exponential backoff + jitter for connection errors, - 429 (honoring `Retry-After`), and 503. Configurable via `max_retries` / - `retry_backoff`. +- **Automatic retries** with exponential backoff + jitter for read-only + connection/transport failures, 429 (honoring `Retry-After`) and 503. Mutations + retry only connection establishment; data DELETE never retries. Configure + via `max_retries` / `retry_backoff`. - New exceptions: `ArcNotSupportedError` (501) and `ArcLicenseError` (403 for a missing enterprise-license feature; subclass of `ArcAuthenticationError`). - `ArcPermissionError` (403: the token lacks the permission the endpoint @@ -65,6 +66,9 @@ plus bug fixes and the first real HTTP-layer test coverage. ### Fixed +- Buffered close waits for active writes, retains unacknowledged batches and + surfaces failures without replay. Typed errors retain structured diagnostics. + - Polars datetime → columnar conversion used a fragile private attribute that raised at runtime; it now casts via the public `polars` API and normalizes the series' `time_unit` to microseconds. diff --git a/src/arc_client/config.py b/src/arc_client/config.py index c3f7c83..0e4caae 100644 --- a/src/arc_client/config.py +++ b/src/arc_client/config.py @@ -27,8 +27,8 @@ class ClientConfig(BaseModel): max_retries: int = Field( default=3, ge=0, - description="Max automatic retries for transient failures " - "(connection errors, 429, 503). 0 disables retries.", + description="Max retries for read-only transient failures or mutation connection " + "establishment failures. Data DELETE never retries. 0 disables retries.", ) retry_backoff: float = Field( default=0.5, diff --git a/src/arc_client/exceptions.py b/src/arc_client/exceptions.py index a0244ff..fad56fd 100644 --- a/src/arc_client/exceptions.py +++ b/src/arc_client/exceptions.py @@ -2,13 +2,29 @@ from __future__ import annotations -from typing import Optional +from typing import Any, Optional class ArcError(Exception): - """Base exception for arc-client.""" - - pass + """Base exception; HTTP failures retain structured diagnostics when available.""" + + _response_body: Optional[dict[str, Any]] = None + + @property + def response_body(self) -> Optional[dict[str, Any]]: + """Retain diagnostics through existing higher-level exception wrappers.""" + error: Optional[BaseException] = self + seen: set[int] = set() + while isinstance(error, ArcError) and id(error) not in seen: + seen.add(id(error)) + if error._response_body is not None: + return error._response_body + error = error.__cause__ + return None + + @response_body.setter + def response_body(self, value: Optional[dict[str, Any]]) -> None: + self._response_body = value class ArcConnectionError(ArcError): diff --git a/src/arc_client/http/async_http.py b/src/arc_client/http/async_http.py index cb9e32d..bb6598c 100644 --- a/src/arc_client/http/async_http.py +++ b/src/arc_client/http/async_http.py @@ -52,24 +52,37 @@ async def _request( ) -> httpx.Response: """Send a request with retry/backoff for transient failures. - Retries connection/timeout errors and 429/503 responses up to - ``config.max_retries`` times, honoring ``Retry-After`` when present. + Reads retry transient failures up to ``config.max_retries`` times. + Mutations retry only connection establishment failures; data DELETE + never retries. Unknown write outcomes are returned to the caller. """ request_kwargs = self._prepare_request_kwargs(headers=headers, **kwargs) client = self._get_client() last_exc: Optional[Exception] = None + read_only = method.upper() in {"GET", "HEAD", "OPTIONS"} + # DELETE may partially commit. Other mutations may be retried only + # when no application request could have reached the server. + data_delete = method.upper() == "POST" and path.rstrip("/") == "/api/v1/delete" for attempt in range(self.config.max_retries + 1): try: response = await client.request(method, path, **request_kwargs) except RETRYABLE_TRANSPORT_ERRORS as e: last_exc = e - if attempt >= self.config.max_retries: + if ( + data_delete + or not (read_only or isinstance(e, (httpx.ConnectError, httpx.ConnectTimeout))) + or attempt >= self.config.max_retries + ): break await asyncio.sleep(compute_backoff(attempt, self.config.retry_backoff)) continue - if is_retryable_status(response.status_code) and attempt < self.config.max_retries: + if ( + read_only + and is_retryable_status(response.status_code) + and attempt < self.config.max_retries + ): retry_after = parse_retry_after(response) await asyncio.sleep( compute_backoff(attempt, self.config.retry_backoff, retry_after) diff --git a/src/arc_client/http/base.py b/src/arc_client/http/base.py index ffa588b..1cb8a3f 100644 --- a/src/arc_client/http/base.py +++ b/src/arc_client/http/base.py @@ -13,6 +13,7 @@ from arc_client.exceptions import ( ArcAuthenticationError, ArcConnectionError, + ArcError, ArcLicenseError, ArcNotFoundError, ArcNotSupportedError, @@ -100,6 +101,19 @@ def _is_permission_data_unavailable(message: object) -> bool: def handle_response_error(response: httpx.Response) -> None: + """Raise the existing typed error, retaining the server's diagnostic object.""" + try: + _raise_response_error(response) + except ArcError as error: + try: + body = response.json() + except ValueError: + body = None + error.response_body = body if isinstance(body, dict) else None + raise + + +def _raise_response_error(response: httpx.Response) -> None: """Handle HTTP error responses and raise appropriate exceptions. Raises: @@ -185,7 +199,7 @@ def handle_connection_error(error: Exception, url: str) -> None: raise ArcConnectionError(f"Failed to connect to {url}: {error}") from error -# Transport errors that are safe to retry (connection refused, DNS, timeouts). +# Transient transport errors eligible for retry on read-only requests. RETRYABLE_TRANSPORT_ERRORS = ( httpx.ConnectError, httpx.ConnectTimeout, diff --git a/src/arc_client/http/sync_http.py b/src/arc_client/http/sync_http.py index b4beb7c..2e46e52 100644 --- a/src/arc_client/http/sync_http.py +++ b/src/arc_client/http/sync_http.py @@ -52,24 +52,37 @@ def _request( ) -> httpx.Response: """Send a request with retry/backoff for transient failures. - Retries connection/timeout errors and 429/503 responses up to - ``config.max_retries`` times, honoring ``Retry-After`` when present. + Reads retry transient failures up to ``config.max_retries`` times. + Mutations retry only connection establishment failures; data DELETE + never retries. Unknown write outcomes are returned to the caller. """ request_kwargs = self._prepare_request_kwargs(headers=headers, **kwargs) client = self._get_client() last_exc: Optional[Exception] = None + read_only = method.upper() in {"GET", "HEAD", "OPTIONS"} + # DELETE may partially commit. Other mutations may be retried only + # when no application request could have reached the server. + data_delete = method.upper() == "POST" and path.rstrip("/") == "/api/v1/delete" for attempt in range(self.config.max_retries + 1): try: response = client.request(method, path, **request_kwargs) except RETRYABLE_TRANSPORT_ERRORS as e: last_exc = e - if attempt >= self.config.max_retries: + if ( + data_delete + or not (read_only or isinstance(e, (httpx.ConnectError, httpx.ConnectTimeout))) + or attempt >= self.config.max_retries + ): break time.sleep(compute_backoff(attempt, self.config.retry_backoff)) continue - if is_retryable_status(response.status_code) and attempt < self.config.max_retries: + if ( + read_only + and is_retryable_status(response.status_code) + and attempt < self.config.max_retries + ): retry_after = parse_retry_after(response) time.sleep(compute_backoff(attempt, self.config.retry_backoff, retry_after)) continue diff --git a/src/arc_client/ingestion/async_buffered.py b/src/arc_client/ingestion/async_buffered.py index c984115..6f5c582 100644 --- a/src/arc_client/ingestion/async_buffered.py +++ b/src/arc_client/ingestion/async_buffered.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio -import contextlib import time from collections import defaultdict from typing import TYPE_CHECKING, Any @@ -38,9 +37,12 @@ def __init__( self._last_flush_time = time.monotonic() self._lock = asyncio.Lock() self._closed = False + self._closing = False + self._flush_error: BaseException | None = None # Background task that flushes on flush_interval even when idle. self._flush_task: asyncio.Task[None] | None = None + self._stop_event = asyncio.Event() async def write(self, record: dict[str, Any]) -> None: """Write a single record to the buffer.""" @@ -65,6 +67,9 @@ async def write(self, record: dict[str, Any]) -> None: columns[key] = [value] async with self._lock: + self._raise_if_failed() + if self._closed or self._closing: + raise RuntimeError("Buffered writer is closed") self._buffers[measurement].append(columns) self._record_counts[measurement] += 1 @@ -86,6 +91,9 @@ async def write_columnar( num_records = len(next(iter(columns.values()))) async with self._lock: + self._raise_if_failed() + if self._closed or self._closing: + raise RuntimeError("Buffered writer is closed") self._buffers[measurement].append(columns) self._record_counts[measurement] += num_records @@ -96,25 +104,40 @@ async def write_columnar( await self._flush_all_unlocked() async def flush(self) -> None: - """Manually flush all buffered data.""" + """Flush all buffered data, waiting for any in-flight write. + + A failed write leaves its data pending and makes the buffer terminal: + subsequent flush/close calls raise that failure without retrying data + whose server outcome may be unknown. Producers wait while I/O runs. + """ async with self._lock: await self._flush_all_unlocked() async def _flush_loop(self) -> None: """Periodically flush measurements whose data has aged past the interval.""" try: - while True: - await asyncio.sleep(self._flush_interval) + while not self._stop_event.is_set(): + try: + await asyncio.wait_for(self._stop_event.wait(), self._flush_interval) + return + except asyncio.TimeoutError: + pass async with self._lock: - if self._closed: + if self._closed or self._stop_event.is_set(): return if ( self._buffers and time.monotonic() - self._last_flush_time >= self._flush_interval ): await self._flush_all_unlocked() - except asyncio.CancelledError: - return + except Exception as error: + # Foreground flush/close will surface this error. Do not retry an + # uncertain outcome, or leave a task exception unobserved. + self._flush_error = error + + def _raise_if_failed(self) -> None: + if self._flush_error is not None: + raise self._flush_error async def _flush_measurement_unlocked(self, measurement: str) -> None: """Flush a single measurement. Must hold lock.""" @@ -122,20 +145,21 @@ async def _flush_measurement_unlocked(self, measurement: str) -> None: return merged = self._merge_columnar(self._buffers[measurement]) - self._buffers[measurement] = [] - self._record_counts[measurement] = 0 - - # Release lock during I/O - self._lock.release() + # Retain data and the lock until acknowledgement. Concurrent timer, + # manual, and close flushes must never send the same batch twice. try: await self._client.write_columnar(measurement, merged) - finally: - await self._lock.acquire() + except BaseException as error: + self._flush_error = error + raise + self._buffers[measurement] = [] + self._record_counts[measurement] = 0 self._last_flush_time = time.monotonic() async def _flush_all_unlocked(self) -> None: """Flush all measurements. Must hold lock.""" + self._raise_if_failed() measurements = list(self._buffers.keys()) for measurement in measurements: if self._buffers[measurement]: @@ -176,12 +200,13 @@ async def close(self) -> None: if self._closed: return - # Cancel the timer task before the final flush so it can't race. + self._closing = True + # Wake an idle timer, but let any in-flight write finish. Shielding + # also prevents cancellation of close's caller from cancelling I/O. + self._stop_event.set() task = self._flush_task if task is not None: - task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await task + await asyncio.shield(task) self._flush_task = None async with self._lock: diff --git a/src/arc_client/ingestion/buffered.py b/src/arc_client/ingestion/buffered.py index d666fbd..9f23843 100644 --- a/src/arc_client/ingestion/buffered.py +++ b/src/arc_client/ingestion/buffered.py @@ -58,6 +58,8 @@ def __init__( self._last_flush_time = time.monotonic() self._lock = threading.Lock() self._closed = False + self._closing = False + self._flush_error: BaseException | None = None # Background timer that flushes on flush_interval even when no writes # arrive. Started on __enter__ and torn down on close(). @@ -99,6 +101,9 @@ def write(self, record: dict[str, Any]) -> None: columns[key] = [value] with self._lock: + self._raise_if_failed() + if self._closed or self._closing: + raise RuntimeError("Buffered writer is closed") self._buffers[measurement].append(columns) self._record_counts[measurement] += 1 @@ -127,6 +132,9 @@ def write_columnar( num_records = len(next(iter(columns.values()))) with self._lock: + self._raise_if_failed() + if self._closed or self._closing: + raise RuntimeError("Buffered writer is closed") self._buffers[measurement].append(columns) self._record_counts[measurement] += num_records @@ -139,7 +147,12 @@ def write_columnar( self._flush_all() def flush(self) -> None: - """Manually flush all buffered data.""" + """Flush all buffered data, waiting for any in-flight write. + + A failed write leaves its data pending and makes the buffer terminal: + subsequent flush/close calls raise that failure without retrying data + whose server outcome may be unknown. Producers wait while I/O runs. + """ with self._lock: self._flush_all() @@ -155,15 +168,25 @@ def _start_flush_thread(self) -> None: def _flush_loop(self) -> None: """Periodically flush measurements whose data has aged past the interval.""" - while not self._stop_event.wait(self._flush_interval): + try: + while not self._stop_event.wait(self._flush_interval): + with self._lock: + if self._closed or self._stop_event.is_set(): + return + if ( + self._buffers + and time.monotonic() - self._last_flush_time >= self._flush_interval + ): + self._flush_all() + except Exception as error: + # Keep the error observable on foreground operations; never retry + # an uncertain write automatically from the background thread. with self._lock: - if self._closed: - return - if ( - self._buffers - and time.monotonic() - self._last_flush_time >= self._flush_interval - ): - self._flush_all() + self._flush_error = error + + def _raise_if_failed(self) -> None: + if self._flush_error is not None: + raise self._flush_error def _flush_measurement(self, measurement: str) -> None: """Flush a single measurement's buffer. Must hold lock.""" @@ -173,21 +196,21 @@ def _flush_measurement(self, measurement: str) -> None: # Merge all columnar batches into one merged = self._merge_columnar(self._buffers[measurement]) - # Clear buffer - self._buffers[measurement] = [] - self._record_counts[measurement] = 0 - - # Write to Arc (release lock during I/O) - self._lock.release() + # Keep ownership and data until the write is acknowledged. Holding + # the lock across I/O serializes timer, manual, and close flushes. try: self._client.write_columnar(measurement, merged) - finally: - self._lock.acquire() + except BaseException as error: + self._flush_error = error + raise + self._buffers[measurement] = [] + self._record_counts[measurement] = 0 self._last_flush_time = time.monotonic() def _flush_all(self) -> None: """Flush all measurements. Must hold lock.""" + self._raise_if_failed() measurements = list(self._buffers.keys()) for measurement in measurements: if self._buffers[measurement]: @@ -229,11 +252,12 @@ def close(self) -> None: if self._closed: return - # Stop the timer thread first so it can't race the final flush. + self._closing = True + # Stop new timer work and wait for any active write to finish. self._stop_event.set() thread = self._flush_thread if thread is not None: - thread.join(timeout=self._flush_interval + 1.0) + thread.join() self._flush_thread = None with self._lock: diff --git a/tests/unit/http/test_error_diagnostics.py b/tests/unit/http/test_error_diagnostics.py new file mode 100644 index 0000000..0d5a7f8 --- /dev/null +++ b/tests/unit/http/test_error_diagnostics.py @@ -0,0 +1,38 @@ +"""Structured HTTP diagnostics survive existing public exception wrappers.""" + +import inspect +from contextlib import AsyncExitStack + +import pytest + +from arc_client import ArcClient, AsyncArcClient +from arc_client.exceptions import ArcError + + +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("operation", ["query", "write", "cq-execute"]) +async def test_wrapper_retains_response_body(httpx_mock, async_mode, operation): + body = {"error": "Permission denied: read required", "detail": "server diagnostic"} + httpx_mock.add_response(status_code=403, json=body) + async with AsyncExitStack() as stack: + if async_mode: + client = await stack.enter_async_context(AsyncArcClient(max_retries=0)) + else: + client = stack.enter_context(ArcClient(max_retries=0)) + with pytest.raises(ArcError) as failure: + if operation == "query": + value = client.query.query("SELECT 1") + elif operation == "write": + value = client.write.write_line_protocol("cpu value=1i") + else: + value = client.continuous_queries.execute(1) + if inspect.isawaitable(value): + await value + assert failure.value.response_body == body + + +def test_error_body_handles_cause_cycle(): + first, second = ArcError("first"), ArcError("second") + first.__cause__ = second + second.__cause__ = first + assert first.response_body is None diff --git a/tests/unit/http/test_retry_safety.py b/tests/unit/http/test_retry_safety.py new file mode 100644 index 0000000..27bde19 --- /dev/null +++ b/tests/unit/http/test_retry_safety.py @@ -0,0 +1,67 @@ +"""Mutation outcomes must not be hidden by automatically replaying writes.""" + +import inspect +from contextlib import AsyncExitStack + +import httpx +import pytest + +from arc_client import ArcClient, AsyncArcClient +from arc_client.exceptions import ArcError + + +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "error", [httpx.ReadTimeout, httpx.WriteTimeout, httpx.RemoteProtocolError] +) +async def test_uncertain_write_is_sent_once(httpx_mock, async_mode, error): + calls = [] + + def failed_write(request): + calls.append(request.content) + raise error("unknown write outcome") + + httpx_mock.add_callback(failed_write, is_reusable=True) + async with AsyncExitStack() as stack: + if async_mode: + client = await stack.enter_async_context(AsyncArcClient(max_retries=3, retry_backoff=0)) + else: + client = stack.enter_context(ArcClient(max_retries=3, retry_backoff=0)) + with pytest.raises(ArcError): + result = client.write.write_line_protocol("cpu,host=a value=1i") + if inspect.isawaitable(result): + await result + assert len(calls) == 1 + + +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("status", [429, 503]) +async def test_write_http_failure_is_not_replayed(httpx_mock, async_mode, status): + httpx_mock.add_response( + status_code=status, json={"error": "write outcome unknown"}, is_reusable=True + ) + async with AsyncExitStack() as stack: + if async_mode: + client = await stack.enter_async_context(AsyncArcClient(max_retries=3, retry_backoff=0)) + else: + client = stack.enter_context(ArcClient(max_retries=3, retry_backoff=0)) + with pytest.raises(ArcError): + result = client.write.write_line_protocol("cpu,host=a value=1i") + if inspect.isawaitable(result): + await result + assert len(httpx_mock.get_requests()) == 1 + + +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +async def test_presend_connection_failure_can_retry_write(httpx_mock, async_mode): + httpx_mock.add_exception(httpx.ConnectError("connection not established")) + httpx_mock.add_response(json={"success": True}) + async with AsyncExitStack() as stack: + if async_mode: + client = await stack.enter_async_context(AsyncArcClient(max_retries=1, retry_backoff=0)) + else: + client = stack.enter_context(ArcClient(max_retries=1, retry_backoff=0)) + result = client.write.write_line_protocol("cpu,host=a value=1i") + if inspect.isawaitable(result): + await result + assert len(httpx_mock.get_requests()) == 2 diff --git a/tests/unit/test_buffered_lifecycle.py b/tests/unit/test_buffered_lifecycle.py new file mode 100644 index 0000000..b83e31a --- /dev/null +++ b/tests/unit/test_buffered_lifecycle.py @@ -0,0 +1,286 @@ +"""Deterministic ownership, shutdown, and failure checks for buffered writes.""" + +import asyncio +import threading + +import pytest + +from arc_client.ingestion.async_buffered import AsyncBufferedWriter +from arc_client.ingestion.buffered import BufferedWriter + + +class SyncWriter: + def __init__(self, error=None): + self.entered = threading.Event() + self.release = threading.Event() + self.calls = [] + self.committed = [] + self.error = error + + def write_columnar(self, measurement, columns): + self.calls.append((measurement, columns)) + self.entered.set() + assert self.release.wait(3), "test did not release write" + if self.error is not None: + raise self.error + self.committed.append((measurement, columns)) + + +class AsyncWriter: + def __init__(self, error=None): + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.calls = [] + self.committed = [] + self.error = error + self.cancelled = False + + async def write_columnar(self, measurement, columns): + self.calls.append((measurement, columns)) + self.entered.set() + try: + await asyncio.wait_for(self.release.wait(), 3) + except asyncio.CancelledError: + self.cancelled = True + raise + if self.error is not None: + raise self.error + self.committed.append((measurement, columns)) + + +async def next_turn(): + """Let already scheduled tasks reach their next suspension without sleeping.""" + reached = asyncio.Event() + asyncio.get_running_loop().call_soon(reached.set) + await reached.wait() + + +async def start_async_timer(buffer): + buffer._last_flush_time = 0 + buffer._flush_interval = 0.001 + await buffer.__aenter__() + + +def start_sync_timer(buffer): + buffer._last_flush_time = 0 + buffer._flush_interval = 0.001 + buffer.__enter__() + + +async def test_async_close_waits_for_active_timer_write_without_cancelling(): + writer = AsyncWriter() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1], "value": [2]}) + await start_async_timer(buffer) + await asyncio.wait_for(writer.entered.wait(), 3) + closing = asyncio.create_task(buffer.close()) + await next_turn() + writer.release.set() + await asyncio.wait_for(closing, 3) + assert writer.cancelled is False + assert len(writer.committed) == len(writer.calls) == 1 + assert buffer.pending_count == 0 + + +async def test_async_manual_flush_and_close_serialize_with_timer(): + writer = AsyncWriter() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + await start_async_timer(buffer) + await asyncio.wait_for(writer.entered.wait(), 3) + flushing = asyncio.create_task(buffer.flush()) + closing = asyncio.create_task(buffer.close()) + await next_turn() + writer.release.set() + await asyncio.wait_for(asyncio.gather(flushing, closing), 3) + assert writer.cancelled is False + assert len(writer.committed) == len(writer.calls) == 1 + + +async def test_async_concurrent_closes_do_not_duplicate_inflight_write(): + writer = AsyncWriter() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + await start_async_timer(buffer) + await asyncio.wait_for(writer.entered.wait(), 3) + first = asyncio.create_task(buffer.close()) + second = asyncio.create_task(buffer.close()) + await next_turn() + writer.release.set() + await asyncio.wait_for(asyncio.gather(first, second), 3) + assert len(writer.committed) == len(writer.calls) == 1 + + +async def test_async_cancelling_close_propagates_without_cancelling_timer_io(): + writer = AsyncWriter() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + await start_async_timer(buffer) + await asyncio.wait_for(writer.entered.wait(), 3) + closing = asyncio.create_task(buffer.close()) + await next_turn() + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + assert writer.cancelled is False + assert buffer.pending_count == 1 + writer.release.set() + await asyncio.wait_for(buffer.close(), 3) + assert len(writer.committed) == len(writer.calls) == 1 + + +async def test_async_cancelling_manual_flush_retains_data_and_propagates(): + writer = AsyncWriter() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + flushing = asyncio.create_task(buffer.flush()) + await asyncio.wait_for(writer.entered.wait(), 3) + flushing.cancel() + with pytest.raises(asyncio.CancelledError): + await flushing + assert buffer.pending_count == 1 + with pytest.raises(asyncio.CancelledError): + await buffer.close() + assert len(writer.calls) == 1 + + +def test_sync_close_and_manual_flush_wait_for_active_timer(monkeypatch): + writer = SyncWriter() + buffer = BufferedWriter(writer, flush_interval=60) + buffer.write_columnar("cpu", {"time": [1]}) + start_sync_timer(buffer) + assert writer.entered.wait(3) + joining = threading.Event() + timer_join = buffer._flush_thread.join + + def observed_join(*args, **kwargs): + joining.set() + timer_join(*args, **kwargs) + + monkeypatch.setattr(buffer._flush_thread, "join", observed_join) + closing = threading.Thread(target=buffer.close) + closing.start() + assert joining.wait(3) + flushing = threading.Thread(target=buffer.flush) + flushing.start() + writer.release.set() + closing.join(3) + flushing.join(3) + assert not closing.is_alive() and not flushing.is_alive() + assert len(writer.committed) == len(writer.calls) == 1 + assert buffer.pending_count == 0 + + +async def test_async_idle_close_wakes_timer_immediately(): + buffer = AsyncBufferedWriter(AsyncWriter(), flush_interval=3600) + await buffer.__aenter__() + task = buffer._flush_task + await next_turn() + await asyncio.wait_for(buffer.close(), 1) + assert task.done() + + +@pytest.mark.parametrize("background", [False, True]) +async def test_async_failure_retains_all_unsent_data_and_close_surfaces_error(background): + error = RuntimeError("write outcome unknown") + writer = AsyncWriter(error) + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + await buffer.write_columnar("mem", {"time": [2]}) + writer.release.set() + if background: + await start_async_timer(buffer) + await asyncio.wait_for(writer.entered.wait(), 3) + # Observe task completion, allowing the implementation to record the error. + await asyncio.gather(buffer._flush_task, return_exceptions=True) + else: + with pytest.raises(RuntimeError, match="outcome unknown"): + await buffer.flush() + assert buffer.pending_count == 2 + for operation in (buffer.close, buffer.flush, buffer.close): + with pytest.raises(RuntimeError) as caught: + await operation() + assert caught.value is error + assert len(writer.calls) == 1 + + +@pytest.mark.parametrize("background", [False, True]) +def test_sync_failure_retains_all_unsent_data_and_close_surfaces_error(background): + error = RuntimeError("write outcome unknown") + writer = SyncWriter(error) + buffer = BufferedWriter(writer, flush_interval=60) + buffer.write_columnar("cpu", {"time": [1]}) + buffer.write_columnar("mem", {"time": [2]}) + writer.release.set() + if background: + start_sync_timer(buffer) + assert writer.entered.wait(3) + buffer._flush_thread.join(3) + assert not buffer._flush_thread.is_alive() + else: + with pytest.raises(RuntimeError, match="outcome unknown"): + buffer.flush() + assert buffer.pending_count == 2 + for operation in (buffer.close, buffer.flush, buffer.close): + with pytest.raises(RuntimeError) as caught: + operation() + assert caught.value is error + assert len(writer.calls) == 1 + + +async def test_async_successful_batches_clear_but_failed_batch_remains(): + class FailSecond: + calls = [] + + async def write_columnar(self, measurement, columns): + self.calls.append(measurement) + if measurement == "mem": + raise RuntimeError("second batch failed") + + writer = FailSecond() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1]}) + await buffer.write_columnar("mem", {"time": [2]}) + with pytest.raises(RuntimeError): + await buffer.flush() + assert writer.calls == ["cpu", "mem"] + assert buffer.pending_count == 1 + + +def test_sync_successful_batches_clear_but_failed_batch_remains(): + class FailSecond: + calls = [] + + def write_columnar(self, measurement, columns): + self.calls.append(measurement) + if measurement == "mem": + raise RuntimeError("second batch failed") + + writer = FailSecond() + buffer = BufferedWriter(writer, flush_interval=60) + buffer.write_columnar("cpu", {"time": [1]}) + buffer.write_columnar("mem", {"time": [2]}) + with pytest.raises(RuntimeError): + buffer.flush() + assert writer.calls == ["cpu", "mem"] + assert buffer.pending_count == 1 + + +async def test_async_sparse_padding_survives_lifecycle_changes(): + writer = AsyncWriter() + writer.release.set() + buffer = AsyncBufferedWriter(writer, flush_interval=60) + await buffer.write_columnar("cpu", {"time": [1], "a": [10]}) + await buffer.write_columnar("cpu", {"time": [2], "b": [20]}) + await buffer.close() + assert writer.committed == [("cpu", {"time": [1, 2], "a": [10, None], "b": [None, 20]})] + + +def test_sync_sparse_padding_survives_lifecycle_changes(): + writer = SyncWriter() + writer.release.set() + buffer = BufferedWriter(writer, flush_interval=60) + buffer.write_columnar("cpu", {"time": [1], "a": [10]}) + buffer.write_columnar("cpu", {"time": [2], "b": [20]}) + buffer.close() + assert writer.committed == [("cpu", {"time": [1, 2], "a": [10, None], "b": [None, 20]})]