diff --git a/python/examples/flight/memory_budget.md b/python/examples/flight/memory_budget.md new file mode 100644 index 000000000000..983a111e534a --- /dev/null +++ b/python/examples/flight/memory_budget.md @@ -0,0 +1,127 @@ + + +# Cooperative payload admission with Flight DoExchange + +`memory_budget.py` illustrates the application-level workaround discussed in +[GH-37900](https://github.com/apache/arrow/issues/37900#issuecomment-1737277716). +A sender announces the next batch's size and waits for the receiver's grant. +The receiver shares a byte budget across exchanges, consumes granted batches, +and releases each reservation before acknowledging consumption. + +This example does **not** implement admission before gRPC allocates memory. +Its budget counts logical int64 value bytes in granted application requests. +It does not cap process RSS, transport read-ahead, schemas, control frames, IPC +metadata, decompression, allocator overhead, or sender allocations. A sender +that ignores the protocol can transmit a large frame before the receiver +validates it. This example assumes cooperating peers. + +## Run + +Use a Python environment with PyArrow including Flight and pytest installed. +From the repository root: + +```sh +python -m pytest -q python/examples/flight/test_memory_budget.py +python python/examples/flight/memory_budget_benchmark.py +``` + +The benchmark starts both endpoints on localhost and prints JSON. It uses four +clients, a 65,536-byte budget, 32,768-byte batches, and a simulated one-millisecond +consumer delay. Each client attempts 100 batches; one additional oversized batch +must be refused. Counts depend on thread scheduling. These tests live with the +example and must be invoked explicitly. + +## Use the example + +Run this from `python/examples/flight`, or put that directory on `PYTHONPATH`: + +```python +import pyarrow as pa +import pyarrow.flight as flight +from memory_budget import BudgetedFlightServer, ByteBudget, SCHEMA, send_batches + +budget = ByteBudget(64 * 1024) + +def consume(batch): + # Finish processing here. Do not retain the batch or its buffers. + print(batch.num_rows) + +with BudgetedFlightServer(("localhost", 0), budget, consume) as server: + with flight.FlightClient(("localhost", server.port)) as client: + batch = pa.record_batch([pa.array(range(1024), type=pa.int64())], + schema=SCHEMA) + print(send_batches(client, [batch], timeout=10)) +print(budget.snapshot()) +``` + +## Protocol and ownership + +The private example command is `example:cooperative-payload-budget:v1`. Control +messages contain a one-byte opcode and an eight-byte unsigned size in network +byte order. Exactly nine metadata bytes are accepted. The fixed schema has one +int64 column named `value`; null values and empty batches are rejected. The +announced size is `rows * 8`. No compression is enabled by the example client. + +1. Send `ANNOUNCE(size)` as application metadata, without batch data. +2. The receiver validates the size and tries to reserve it under a shared lock. + It replies `REFUSE(size)` immediately when the budget is full. The sender + skips that batch; retry policy belongs to the caller. +3. After `GRANT(size)`, send one batch. The receiver validates its schema and + logical size, then calls the synchronous consumer while holding the + reservation. Schema transmission for the first batch is outside accounting. +4. Release the reservation and reply `CONSUMED(size)`. The sender waits for this + acknowledgement before announcing another batch. + +Only one batch can be outstanding per compliant stream. Concurrent consumers +may run on different Flight handler threads, so the supplied callback must be +thread-safe. Budget exhaustion does not queue requests or busy-wait. The default +per-batch limit is 16 MiB of logical values; it is a protocol check performed +after Flight receives each frame, not a transport allocation limit. + +Reservations are released on normal consumption, callback failure, invalid +payload, early EOF, cancellation, and deadline errors. The client requires a +finite positive deadline for the complete RPC. An arbitrary peer without a +deadline can hold a grant indefinitely; the server does not enforce a lease. +A consumer that blocks must arrange its own cancellation. A consumer that +retains buffers must extend ownership/accounting itself before using this +pattern in an application. + +## Measurement + +`reserved_bytes` and `peak_reserved_bytes` describe the logical admission +counter, not allocated memory. The benchmark also reports the combined +client/server process's lifetime peak RSS and Arrow's default memory pool +metrics. RSS is reported on macOS/Linux and is not current usage; Arrow's pool excludes +gRPC allocations. +Both include effects unrelated to the reservation counter. Elapsed time includes +server/client setup, simulated consumption, and shutdown. This workload checks +contention and cleanup; it is not a transport throughput comparison. + +For example, vary available capacity while keeping the workload fixed: + +```sh +python python/examples/flight/memory_budget_benchmark.py --budget-bytes 65536 +python python/examples/flight/memory_budget_benchmark.py --budget-bytes 131072 +``` + +There are no Flight wire-protocol, native allocator, or C++ API changes. Applying +this pattern to downloads would place the same budget and grant decision at the +client receiver. It requires an application protocol agreed by both peers and +does not interoperate automatically with generic Flight clients. diff --git a/python/examples/flight/memory_budget.py b/python/examples/flight/memory_budget.py new file mode 100644 index 000000000000..264bdb90b062 --- /dev/null +++ b/python/examples/flight/memory_budget.py @@ -0,0 +1,250 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Cooperative payload admission over Flight DoExchange. + +This example accounts for logical int64 value bytes admitted by the +application. +It does NOT limit gRPC allocations, process RSS, schema/control messages, IPC +metadata, decompression, or allocator overhead. Both peers must follow the +protocol. Flight delivers data only after the transport has allocated it, so +checking a received batch cannot protect against an uncooperative sender. + +A sender announces one nonempty batch, waits for a grant, sends that batch, and +waits for a consumed acknowledgement before announcing the next batch. A shared +receiver budget refuses requests immediately when full; it does not queue them. +The consumer must finish synchronously and must not retain batches or buffers. +Client deadlines bound a stalled cooperative exchange; an arbitrary peer can +hold a reservation indefinitely, and a blocking consumer must implement its own +cancellation. This is an application example, not a standard Flight protocol. +""" + +import contextlib +from dataclasses import dataclass +import math +import struct +import threading + +import pyarrow as pa +import pyarrow.flight as flight + + +SCHEMA = pa.schema([("value", pa.int64())]) +COMMAND = b"example:cooperative-payload-budget:v1" +ANNOUNCE, GRANT, REFUSE, CONSUMED = range(1, 5) +CONTROL = struct.Struct("!BQ") +MAX_UINT64 = (1 << 64) - 1 + + +def _positive_size(size): + # bool is an int subclass, but is not a meaningful byte count. + if type(size) is not int or not 0 < size <= MAX_UINT64: + raise ValueError("byte count must be a positive uint64 integer") + return size + + +@dataclass(frozen=True) +class BudgetStats: + capacity_bytes: int + reserved_bytes: int + peak_reserved_bytes: int + admitted: int + refused: int + + +class ByteBudget: + """Thread-safe, immediate admission of application payload reservations.""" + + def __init__(self, capacity_bytes): + self._capacity_bytes = _positive_size(capacity_bytes) + self._lock = threading.Lock() + self._reserved = self._peak = self._admitted = self._refused = 0 + + @property + def capacity_bytes(self): + return self._capacity_bytes + + def try_reserve(self, size): + """Return an owned reservation, or None without waiting when full.""" + _positive_size(size) + with self._lock: + if size > self.capacity_bytes - self._reserved: + self._refused += 1 + return None + reservation = _Reservation(self, size) + self._reserved += size + self._peak = max(self._peak, self._reserved) + self._admitted += 1 + return reservation + + def snapshot(self): + with self._lock: + return BudgetStats(self.capacity_bytes, self._reserved, self._peak, + self._admitted, self._refused) + + +class _Reservation: + def __init__(self, budget, size): + self._budget = budget + self._size = size + self._released = False + + def release(self): + """Release once; repeated release is harmless and returns False.""" + with self._budget._lock: + if self._released: + return False + self._budget._reserved -= self._size + self._released = True + return True + + def __enter__(self): + with self._budget._lock: + if self._released: + raise RuntimeError("reservation already released") + return self + + def __exit__(self, *exc): + self.release() + + +def control_message(kind, size): + _positive_size(size) + if kind not in (ANNOUNCE, GRANT, REFUSE, CONSUMED): + raise ValueError("unknown control message") + return CONTROL.pack(kind, size) + + +def read_control(chunk, allowed): + """Validate before parsing; Flight has already allocated the frame.""" + metadata = chunk.app_metadata + if (chunk.data is not None or metadata is None or + len(metadata) != CONTROL.size): + raise ValueError("expected a nine-byte metadata-only control message") + kind, size = CONTROL.unpack(metadata) + if kind not in allowed: + raise ValueError("unexpected control message") + _positive_size(size) + return kind, size + + +def payload_size(batch): + """Logical value bytes for the example's non-null int64 schema only.""" + if (not batch.schema.equals(SCHEMA, check_metadata=True) or + batch.column(0).null_count): + raise ValueError("expected one non-null int64 column named value") + return _positive_size(batch.num_rows * 8) + + +class BudgetedFlightServer(flight.FlightServerBase): + """Consume admitted batches and share one budget across all exchanges.""" + + def __init__(self, location, budget, consume, *, + max_batch_bytes=16 * 1024 * 1024): + self.budget = budget + self.consume = consume + self.max_batch_bytes = _positive_size(max_batch_bytes) + super().__init__(location) + + def do_exchange(self, context, descriptor, reader, writer): + if (descriptor.descriptor_type != flight.DescriptorType.CMD or + descriptor.command != COMMAND): + raise ValueError("unknown example protocol") + while True: + try: + chunk = reader.read_chunk() + except StopIteration: + return + _, size = read_control(chunk, {ANNOUNCE}) + if size > self.max_batch_bytes or size % 8: + raise ValueError("announced payload exceeds the batch limit " + "or is not a multiple of eight") + reservation = self.budget.try_reserve(size) + if reservation is None: + writer.write_metadata(control_message(REFUSE, size)) + continue + # Covers grant write failure, premature EOF, validation failure, + # callback exceptions, and transport cancellation/deadline errors. + with reservation: + writer.write_metadata(control_message(GRANT, size)) + try: + chunk = reader.read_chunk() + except StopIteration as exc: + raise ValueError("stream ended before granted batch") \ + from exc + try: + if (chunk.data is None or + chunk.app_metadata is not None or + payload_size(chunk.data) != size): + raise ValueError("batch does not match its grant") + self.consume(chunk.data) + finally: + # Do not keep the last batch alive while waiting for the + # next announcement. Consumers must also drop references. + chunk = None + writer.write_metadata(control_message(CONSUMED, size)) + + +def send_batches(client, batches, *, timeout=10.0): + """Send cooperatively; return admitted/refused counts without retrying. + + The timeout covers the whole RPC. Callers decide when to retry refused + batches; this example never silently retries or buffers an unbounded queue. + Batches are already allocated by the sender before they are announced. + """ + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("timeout must be finite and positive") + writer, reader = client.do_exchange( + flight.FlightDescriptor.for_command(COMMAND), + options=flight.FlightCallOptions(timeout=timeout)) + admitted = refused = 0 + started = False + try: + for batch in batches: + size = payload_size(batch) + writer.write_metadata(control_message(ANNOUNCE, size)) + kind, reply_size = read_control( + reader.read_chunk(), {GRANT, REFUSE}) + if reply_size != size: + raise ValueError("reply does not match announcement") + if kind == REFUSE: + refused += 1 + continue + if not started: + writer.begin(SCHEMA) + started = True + writer.write_batch(batch) + _, reply_size = read_control(reader.read_chunk(), {CONSUMED}) + if reply_size != size: + raise ValueError("acknowledgement does not match batch") + admitted += 1 + writer.done_writing() + try: + reader.read_chunk() + except StopIteration: + pass + else: + raise ValueError("unexpected message after end of exchange") + writer.close() + except BaseException: + # In particular, cancel if producing/validating a batch failed while + # the server was waiting. Preserve the original exception on close. + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + raise + return {"admitted": admitted, "refused": refused} diff --git a/python/examples/flight/memory_budget_benchmark.py b/python/examples/flight/memory_budget_benchmark.py new file mode 100644 index 000000000000..ca7836969caa --- /dev/null +++ b/python/examples/flight/memory_budget_benchmark.py @@ -0,0 +1,133 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Localhost workload for cooperative Flight payload admission. + +The simulated consumer delay keeps reservations occupied to expose contention. +Elapsed time is not a general Flight throughput benchmark. The process contains +both endpoints; RSS includes native transport/Arrow allocations and is not +bounded by the application's logical payload budget. +""" + +import argparse +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict +import json +import platform +import sys +import time + +import pyarrow as pa +import pyarrow.flight as flight + +from memory_budget import ( + BudgetedFlightServer, ByteBudget, SCHEMA, send_batches, +) + + +def peak_rss_bytes(): + # getrusage reports bytes on macOS and KiB on Linux. Only assign units + # on those platforms. + if sys.platform != "darwin" and not sys.platform.startswith("linux"): + return None + import resource + raw = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform == "darwin": + return raw + return raw * 1024 + + +def run_benchmark(*, workers=4, batches=100, rows=4096, + budget_bytes=65536, consumer_delay=0.001): + if min(workers, batches, rows, budget_bytes) <= 0: + raise ValueError("workload sizes must be positive") + if max(rows * 8, budget_bytes + 8) > 16 * 1024 * 1024: + raise ValueError("benchmark payloads must fit the 16 MiB batch limit") + if not 0 <= consumer_delay <= 1: + raise ValueError("consumer delay must be between zero and one second") + budget = ByteBudget(budget_bytes) + batch = pa.record_batch([pa.array(range(rows), type=pa.int64())], + schema=SCHEMA) + # Deliberately over budget to exercise refusal even with one worker. + oversized = pa.record_batch( + [pa.array(range(budget_bytes // 8 + 1), type=pa.int64())], + schema=SCHEMA) + pool = pa.default_memory_pool() + before = {"rss_high_water_bytes": peak_rss_bytes(), + "arrow_pool_bytes_allocated": pool.bytes_allocated()} + + def consume(batch): + time.sleep(consumer_delay) + + start = time.perf_counter() + with BudgetedFlightServer(("localhost", 0), budget, consume) as server: + def send(_): + with flight.FlightClient(("localhost", server.port)) as client: + return send_batches( + client, (batch for _ in range(batches)), + timeout=max(30, batches * consumer_delay * 4)) + + with ThreadPoolExecutor(max_workers=workers) as executor: + outcomes = list(executor.map(send, range(workers))) + with flight.FlightClient(("localhost", server.port)) as client: + outcomes.append(send_batches(client, [oversized])) + elapsed = time.perf_counter() - start + stats = budget.snapshot() + assert stats.reserved_bytes == 0 + assert stats.peak_reserved_bytes <= budget_bytes + assert stats.admitted == sum(item["admitted"] for item in outcomes) + assert stats.refused == sum(item["refused"] for item in outcomes) + return { + "python": platform.python_version(), "pyarrow": pa.__version__, + "platform": platform.platform(), + "workload": {"workers": workers, "batches_per_worker": batches, + "rows_per_batch": rows, + "payload_bytes_per_batch": rows * 8, + "consumer_delay_seconds": consumer_delay, + "extra_oversized_attempts": 1}, + "elapsed_seconds": elapsed, + "budget": asdict(stats), + "memory_before": before, + "memory_after": { + "rss_high_water_bytes": peak_rss_bytes(), + "arrow_pool_bytes_allocated": pool.bytes_allocated(), + "arrow_pool_lifetime_peak_bytes": pool.max_memory()}, + "metric_definitions": { + "reserved_bytes": "Application logical int64 value bytes granted " + "but not yet consumed or released on failure.", + "rss_high_water_bytes": "Lifetime peak RSS of the combined " + "client/server process; not current RSS.", + "arrow_pool_bytes_allocated": "Current bytes in Arrow's default " + "pool; excludes gRPC allocations.", + "arrow_pool_lifetime_peak_bytes": "Lifetime high-water allocation " + "in Arrow's default pool.", + "elapsed_seconds": "Localhost workload including setup, " + "simulated consumer delay, and shutdown."}, + "scope": "Cooperative application admission only; no cap on gRPC " + "allocation, read-ahead, IPC metadata, schemas, " + "decompression, allocator overhead, sender memory, " + "or process RSS."} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--batches", type=int, default=100) + parser.add_argument("--rows", type=int, default=4096) + parser.add_argument("--budget-bytes", type=int, default=65536) + parser.add_argument("--consumer-delay", type=float, default=0.001) + print(json.dumps(run_benchmark(**vars(parser.parse_args())), indent=2)) diff --git a/python/examples/flight/test_memory_budget.py b/python/examples/flight/test_memory_budget.py new file mode 100644 index 000000000000..aea00d1776e3 --- /dev/null +++ b/python/examples/flight/test_memory_budget.py @@ -0,0 +1,259 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Unit and localhost RPC tests for the cooperative admission example. + +Run explicitly with: pytest python/examples/flight/test_memory_budget.py +""" + +from concurrent.futures import ThreadPoolExecutor +import contextlib +import threading +import time + +import pyarrow as pa +import pyarrow.flight as flight +import pytest + +from memory_budget import ( + ANNOUNCE, COMMAND, CONSUMED, CONTROL, GRANT, MAX_UINT64, SCHEMA, + BudgetedFlightServer, ByteBudget, control_message, payload_size, + read_control, send_batches, +) + + +def make_batch(rows=8): + return pa.record_batch([pa.array(range(rows), type=pa.int64())], + schema=SCHEMA) + + +@contextlib.contextmanager +def exchange(client, timeout=5): + writer, reader = client.do_exchange( + flight.FlightDescriptor.for_command(COMMAND), + options=flight.FlightCallOptions(timeout=timeout)) + try: + yield writer, reader + finally: + reader.cancel() + with contextlib.suppress(pa.ArrowException): + writer.close() + + +def wait_for_release(budget): + deadline = time.monotonic() + 5 + while budget.snapshot().reserved_bytes and time.monotonic() < deadline: + time.sleep(0.005) + assert budget.snapshot().reserved_bytes == 0 + + +@pytest.mark.parametrize("size", [0, -1, True, 1.5, "8", MAX_UINT64 + 1]) +def test_invalid_sizes(size): + with pytest.raises(ValueError): + ByteBudget(size) + budget = ByteBudget(64) + with pytest.raises(ValueError): + budget.try_reserve(size) + assert budget.snapshot().reserved_bytes == 0 + + +def test_reservation_lifetime(): + budget = ByteBudget(64) + reservation = budget.try_reserve(64) + assert budget.try_reserve(1) is None + assert reservation.release() + assert not reservation.release() + with pytest.raises(RuntimeError, match="already released"): + with reservation: + pass + with pytest.raises(RuntimeError, match="consumer failed"): + with budget.try_reserve(64): + raise RuntimeError("consumer failed") + assert budget.snapshot().reserved_bytes == 0 + assert budget.snapshot().peak_reserved_bytes == 64 + with pytest.raises(AttributeError): + budget.capacity_bytes = 128 + + +def test_concurrent_reservations_do_not_overcommit(): + budget = ByteBudget(64) + barrier = threading.Barrier(16) + + def reserve(): + reservation = budget.try_reserve(16) + barrier.wait(timeout=5) + if reservation: + reservation.release() + return reservation is not None + + with ThreadPoolExecutor(max_workers=16) as pool: + admitted = sum(pool.map(lambda _: reserve(), range(16))) + stats = budget.snapshot() + assert admitted == stats.admitted == 4 + assert stats.refused == 12 + assert stats.peak_reserved_bytes == 64 + assert stats.reserved_bytes == 0 + + +@pytest.mark.parametrize("raw", [b"", b"x" * 10, CONTROL.pack(99, 8), + CONTROL.pack(ANNOUNCE, 0)]) +def test_control_validation(raw): + chunk = flight.FlightStreamChunk(None, pa.py_buffer(raw)) + with pytest.raises(ValueError): + read_control(chunk, {ANNOUNCE}) + + +def test_payload_restrictions(): + assert payload_size(make_batch()) == 64 + for batch in [make_batch(0), pa.record_batch([[None]], schema=SCHEMA), + pa.record_batch([[1]], names=["other"]), + pa.record_batch([[1.0]], names=["value"])]: + with pytest.raises(ValueError): + payload_size(batch) + + +def test_rpc_accept_refuse_and_reuse(): + budget = ByteBudget(64) + seen = [] + + def consume(batch): + assert budget.snapshot().reserved_bytes == batch.num_rows * 8 + seen.append(batch.num_rows) + + with BudgetedFlightServer(("localhost", 0), budget, consume) as server, \ + flight.FlightClient(("localhost", server.port)) as client: + assert send_batches(client, [make_batch(8), make_batch(16), + make_batch(8)]) == { + "admitted": 2, "refused": 1} + assert send_batches(client, [make_batch(8)]) == { + "admitted": 1, "refused": 0} + assert seen == [8, 8, 8] + stats = budget.snapshot() + assert stats.admitted == 3 + assert stats.refused == 1 + assert stats.peak_reserved_bytes == 64 + assert stats.reserved_bytes == 0 + + +def test_rpc_concurrent_refusal_while_consumer_owns_batch(): + budget = ByteBudget(64) + entered = threading.Event() + finish = threading.Event() + + def consume(batch): + entered.set() + assert finish.wait(timeout=5) + + with BudgetedFlightServer(("localhost", 0), budget, consume) as server, \ + flight.FlightClient(("localhost", server.port)) as first, \ + flight.FlightClient(("localhost", server.port)) as second, \ + ThreadPoolExecutor(max_workers=1) as pool: + pending = pool.submit(send_batches, first, [make_batch()]) + try: + assert entered.wait(timeout=5) + assert send_batches(second, [make_batch()]) == { + "admitted": 0, "refused": 1} + assert budget.snapshot().reserved_bytes == 64 + finally: + finish.set() + assert pending.result(timeout=5)["admitted"] == 1 + assert send_batches(second, [make_batch()])["admitted"] == 1 + assert budget.snapshot().reserved_bytes == 0 + assert budget.snapshot().peak_reserved_bytes == 64 + + +def test_rpc_consumer_exception_releases_reservation(): + budget = ByteBudget(64) + fail = True + + def consume(batch): + nonlocal fail + if fail: + fail = False + raise RuntimeError("consumer failed") + + with BudgetedFlightServer(("localhost", 0), budget, consume) as server, \ + flight.FlightClient(("localhost", server.port)) as client: + with pytest.raises(pa.ArrowException, match="consumer failed"): + send_batches(client, [make_batch()]) + wait_for_release(budget) + assert send_batches(client, [make_batch()])["admitted"] == 1 + + +@pytest.mark.parametrize("failure", ["cancel", "deadline", "eof", "size"]) +def test_rpc_incomplete_or_mismatched_batch_releases(failure): + budget = ByteBudget(64) + with BudgetedFlightServer(("localhost", 0), budget, + lambda batch: None) as server, \ + flight.FlightClient(("localhost", server.port)) as client: + timeout = 0.2 if failure == "deadline" else 5 + with exchange(client, timeout=timeout) as (writer, reader): + writer.write_metadata(control_message(ANNOUNCE, 64)) + assert read_control(reader.read_chunk(), {GRANT}) == (GRANT, 64) + assert budget.snapshot().reserved_bytes == 64 + if failure == "cancel": + reader.cancel() + elif failure == "deadline": + with pytest.raises(pa.ArrowException): + reader.read_chunk() + else: + if failure == "eof": + writer.done_writing() + else: + writer.begin(SCHEMA) + writer.write_batch(make_batch(4)) + with pytest.raises(pa.ArrowException, + match="before granted|match"): + reader.read_chunk() + wait_for_release(budget) + assert send_batches(client, [make_batch()])["admitted"] == 1 + + +@pytest.mark.parametrize("raw", [b"x" * 10, CONTROL.pack(ANNOUNCE, 0), + CONTROL.pack(ANNOUNCE, 7), + CONTROL.pack(ANNOUNCE, 128), + CONTROL.pack(CONSUMED, 64)]) +def test_rpc_rejects_invalid_announcements(raw): + budget = ByteBudget(64) + with BudgetedFlightServer(("localhost", 0), budget, lambda batch: None, + max_batch_bytes=64) as server, \ + flight.FlightClient(("localhost", server.port)) as client: + with exchange(client) as (writer, reader): + writer.write_metadata(raw) + with pytest.raises(pa.ArrowException): + reader.read_chunk() + assert budget.snapshot().admitted == 0 + assert budget.snapshot().reserved_bytes == 0 + + +def test_rpc_rejects_data_before_announcement(): + budget = ByteBudget(64) + with BudgetedFlightServer(("localhost", 0), budget, + lambda batch: None) as server, \ + flight.FlightClient(("localhost", server.port)) as client: + with exchange(client) as (writer, reader): + writer.begin(SCHEMA) + writer.write_batch(make_batch()) + with pytest.raises(pa.ArrowException, match="control message"): + reader.read_chunk() + assert budget.snapshot().admitted == 0 + + +@pytest.mark.parametrize("timeout", [0, -1, float("inf"), float("nan")]) +def test_invalid_client_deadline(timeout): + with pytest.raises(ValueError, match="timeout"): + send_batches(None, [], timeout=timeout)