|
| 1 | +"""Measurement protocol behind the constant-memory streaming coverage. |
| 2 | +
|
| 3 | +`stream()` claims memory stays bounded no matter how many records the export carries. |
| 4 | +Proving that needs two comparisons on the same live data, which is what |
| 5 | +`profile_streaming_memory` collects in one pass: |
| 6 | +
|
| 7 | +- the allocation peak of a small export against the peak of an export ten times longer, |
| 8 | + so a peak that scales with the record count fails; |
| 9 | +- the peak of the streamed read against the peak of buffering the same export, so the |
| 10 | + measurement is shown to be sensitive enough to see buffering at this volume. |
| 11 | +
|
| 12 | +Every measured window covers a whole read, from the request being opened to the body being |
| 13 | +exhausted. That matters: an implementation that buffers does so before it yields its first |
| 14 | +record, so a window that started mid-stream would measure only the walk over an |
| 15 | +already-materialised list and see nothing wrong. |
| 16 | +
|
| 17 | +Peaks come from `tracemalloc`, which counts Python allocations rather than process resident |
| 18 | +size, so the numbers do not move with allocator or kernel behaviour. Every peak is reported |
| 19 | +net of the memory already live when its window opens, and each window starts with a |
| 20 | +collection, so an object released by an earlier read is never charged to a later one. |
| 21 | +""" |
| 22 | + |
| 23 | +import gc |
| 24 | +import tracemalloc |
| 25 | +from collections.abc import Awaitable, Callable, Iterator |
| 26 | +from contextlib import contextmanager |
| 27 | +from dataclasses import dataclass |
| 28 | +from typing import Any |
| 29 | + |
| 30 | +# Records exported before the first measured window, so connection setup, lazy imports and |
| 31 | +# the first read buffers are paid for outside every comparison. This is a separate read, |
| 32 | +# not a prologue of a measured one, so nothing can hide in it. |
| 33 | +WARMUP_RECORDS = 500 |
| 34 | +# Records exported by the short read that the volume read is compared against. |
| 35 | +SAMPLE_RECORDS = 2000 |
| 36 | +# Records exported by the volume read: ten times the short read, and enough that buffering |
| 37 | +# the same export costs tens of megabytes. |
| 38 | +VOLUME_RECORDS = 20000 |
| 39 | +# How much the volume read's peak may exceed the short read's. Ten times the records must |
| 40 | +# not cost materially more memory; a buffering read lands an order of magnitude above this. |
| 41 | +MAX_PEAK_GROWTH = 3.0 |
| 42 | +# How much more buffering the export must peak than streaming it. Guards against a vacuous |
| 43 | +# pass: on a dataset too small to matter, the two reads would cost about the same. |
| 44 | +MIN_BUFFERING_CONTRAST = 20.0 |
| 45 | + |
| 46 | +# One traced frame per allocation: the call site is never reported, so tracing costs less. |
| 47 | +TRACE_FRAMES = 1 |
| 48 | + |
| 49 | + |
| 50 | +@dataclass(frozen=True) |
| 51 | +class PeakReading: |
| 52 | + records: int |
| 53 | + peak_bytes: int |
| 54 | + |
| 55 | + |
| 56 | +@dataclass(frozen=True) |
| 57 | +class StreamingMemoryProfile: |
| 58 | + sample: PeakReading |
| 59 | + volume: PeakReading |
| 60 | + buffered: PeakReading |
| 61 | + |
| 62 | + |
| 63 | +@contextmanager |
| 64 | +def tracing_allocations() -> Iterator[None]: |
| 65 | + gc.collect() |
| 66 | + tracemalloc.start(TRACE_FRAMES) |
| 67 | + try: |
| 68 | + yield |
| 69 | + finally: |
| 70 | + tracemalloc.stop() |
| 71 | + |
| 72 | + |
| 73 | +def measure_peak(read: Callable[[], int]) -> PeakReading: |
| 74 | + gc.collect() |
| 75 | + tracemalloc.reset_peak() |
| 76 | + baseline_bytes = tracemalloc.get_traced_memory()[0] |
| 77 | + records = read() |
| 78 | + peak_bytes = tracemalloc.get_traced_memory()[1] - baseline_bytes |
| 79 | + return PeakReading(records=records, peak_bytes=peak_bytes) |
| 80 | + |
| 81 | + |
| 82 | +async def measure_peak_async(read: Callable[[], Awaitable[int]]) -> PeakReading: |
| 83 | + gc.collect() |
| 84 | + tracemalloc.reset_peak() |
| 85 | + baseline_bytes = tracemalloc.get_traced_memory()[0] |
| 86 | + records = await read() |
| 87 | + peak_bytes = tracemalloc.get_traced_memory()[1] - baseline_bytes |
| 88 | + return PeakReading(records=records, peak_bytes=peak_bytes) |
| 89 | + |
| 90 | + |
| 91 | +def drain_stream(service: Any, limit: int) -> int: |
| 92 | + """Stream `limit` records, retaining none. |
| 93 | +
|
| 94 | + The stream is consumed to exhaustion, which is what triggers its completeness check. |
| 95 | + """ |
| 96 | + return sum(1 for _ in service.stream(limit=limit)) |
| 97 | + |
| 98 | + |
| 99 | +async def drain_stream_async(service: Any, limit: int) -> int: |
| 100 | + """Stream `limit` records, retaining none. |
| 101 | +
|
| 102 | + The stream is consumed to exhaustion, which is what triggers its completeness check. |
| 103 | + """ |
| 104 | + # WPS519 would have the records collected and counted, which is the buffering this |
| 105 | + # measurement exists to rule out. |
| 106 | + consumed = 0 |
| 107 | + async for _ in service.stream(limit=limit): # noqa: WPS519 |
| 108 | + consumed += 1 |
| 109 | + return consumed |
| 110 | + |
| 111 | + |
| 112 | +def buffer_export(service: Any) -> int: |
| 113 | + """Materialise a whole export, the buffering read the streamed read is compared with.""" |
| 114 | + return len(list(service.stream(limit=VOLUME_RECORDS))) |
| 115 | + |
| 116 | + |
| 117 | +async def buffer_export_async(service: Any) -> int: |
| 118 | + """Materialise a whole export, the buffering read the streamed read is compared with.""" |
| 119 | + return len([record async for record in service.stream(limit=VOLUME_RECORDS)]) |
| 120 | + |
| 121 | + |
| 122 | +def profile_streaming_memory(service: Any) -> StreamingMemoryProfile: |
| 123 | + with tracing_allocations(): |
| 124 | + drain_stream(service, WARMUP_RECORDS) |
| 125 | + sample = measure_peak(lambda: drain_stream(service, SAMPLE_RECORDS)) |
| 126 | + volume = measure_peak(lambda: drain_stream(service, VOLUME_RECORDS)) |
| 127 | + buffered = measure_peak(lambda: buffer_export(service)) |
| 128 | + return StreamingMemoryProfile(sample=sample, volume=volume, buffered=buffered) |
| 129 | + |
| 130 | + |
| 131 | +async def profile_async_streaming_memory(service: Any) -> StreamingMemoryProfile: |
| 132 | + with tracing_allocations(): |
| 133 | + await drain_stream_async(service, WARMUP_RECORDS) |
| 134 | + sample = await measure_peak_async(lambda: drain_stream_async(service, SAMPLE_RECORDS)) |
| 135 | + volume = await measure_peak_async(lambda: drain_stream_async(service, VOLUME_RECORDS)) |
| 136 | + buffered = await measure_peak_async(lambda: buffer_export_async(service)) |
| 137 | + return StreamingMemoryProfile(sample=sample, volume=volume, buffered=buffered) |
0 commit comments