Skip to content

Commit 1d8f0ec

Browse files
rustyconoverclaude
andcommitted
fix: two cache/rehydrate faults the integration suite caught
Both are mine, both from this release's work, and neither was reachable from the unit suite — 2,100+ tests passed while the HTTP integration suite failed 174 of 304. Recording that, because it says where the risk in this optimization work actually sits. MetaWorker had no `_unwrap_attach_full`. Dropping the decrypted attach from the call state means the rehydrate path re-opens the sealed envelope on whatever the transport hands it as the implementation — and under a MetaWorker that object is not a `Worker`: MetaWorker proxies rather than subclasses. Every HTTP continuation of a catalog-backed stream died on `'MetaWorker' object has no attribute '_unwrap_attach_full'`. It now routes the way `_resolve_function` does, to the sub-worker owning the catalog, falling back to `_workers[0]` — every sub-worker shares one signing key, as `_maybe_worker_for_attach` already assumes when it opens envelopes through `_workers[0]` to read the catalog name. Nothing here drove a MetaWorker over HTTP, which is why unit coverage said nothing. `_partition_fields_cached` was keyed on `pa.Schema`, whose `__eq__` and `__hash__` ignore field metadata — the exact thing the function reads. Its own docstring said "a pure function of the schema's field metadata" above a cache that cannot see it. So an annotated schema and an unannotated one of the same names and types shared an entry and whichever was scanned first answered for both: a function with no partition annotations inherited a neighbour's fields, and the `partition_values` contract check passed where it had to raise. Order-dependent, so it reproduced on one transport and not the other. Fixed with a `_SchemaKey` wrapper that keeps the metadata-blind hash — a hash may collide, it may not lie — and discriminates in `__eq__` via `equals(check_metadata=True)`, measured at 0.042us, so the optimization survives. Regression tests cover both cache-population orders, plus the Arrow equality premise that makes the naive key unsafe in the first place. Integration green on both transports: subprocess 10,334 assertions, HTTP 10,020, 0 failures. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MS3MkkHgE8pA3Aty7BnZKE
1 parent ed1da23 commit 1d8f0ec

4 files changed

Lines changed: 157 additions & 4 deletions

File tree

‎tests/test_call_state_attach.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,3 +303,44 @@ def test_cache_is_bounded(self) -> None:
303303
# Least-recently-used entries evicted; the newest survive.
304304
assert cache.get((bytes([19]), b"id")) == bytes([19])
305305
assert cache.get((bytes([0]), b"id")) is None
306+
307+
308+
class TestMetaWorkerCanOpenTheAttach:
309+
"""A MetaWorker must satisfy the same rehydrate contract a Worker does.
310+
311+
Since the stream state stopped carrying the attach decrypted, the rehydrate
312+
path re-opens the sealed envelope on whatever the transport hands it as the
313+
implementation. Under a MetaWorker that object is *not* a ``Worker`` —
314+
MetaWorker does not subclass it, it proxies — so the method has to exist
315+
there too. It did not, and every HTTP continuation of a catalog-backed
316+
stream died on ``'MetaWorker' object has no attribute
317+
'_unwrap_attach_full'``. Unit coverage missed it because nothing here drove
318+
a MetaWorker over HTTP; the integration suite caught it.
319+
"""
320+
321+
@staticmethod
322+
def _meta_and_worker() -> tuple[Any, Any]:
323+
from vgi._test_fixtures.worker import ExampleWorker
324+
from vgi.meta_worker import MetaWorker
325+
326+
worker = ExampleWorker(quiet=True)
327+
worker._signing_key = _KEY
328+
return MetaWorker([worker]), worker
329+
330+
def test_meta_worker_exposes_the_rehydrate_hook(self) -> None:
331+
"""The attribute the rehydrate path calls must exist on a MetaWorker."""
332+
meta, _ = self._meta_and_worker()
333+
assert hasattr(meta, "_unwrap_attach_full")
334+
335+
def test_meta_worker_unwraps_identically_to_its_sub_worker(self) -> None:
336+
"""Routing through the MetaWorker yields the sub-worker's plaintext."""
337+
meta, worker = self._meta_and_worker()
338+
with _as_principal(_DOMAIN, _PRINCIPAL):
339+
envelope = bytes(worker._seal_attach(_UUID + _SECRET))
340+
assert meta._unwrap_attach_full(envelope) == _UUID + _SECRET
341+
assert meta._unwrap_attach_full(envelope) == worker._unwrap_attach_full(envelope)
342+
343+
def test_no_attach_passes_through(self) -> None:
344+
"""No catalog context means nothing to open."""
345+
meta, _ = self._meta_and_worker()
346+
assert meta._unwrap_attach_full(None) is None

‎tests/test_protocol_classes.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1248,3 +1248,52 @@ def test_namespace_survives_a_round_trip(self) -> None:
12481248
# The scan itself must accept what came back: _coerce_ns takes the
12491249
# reserved b"_vgi/" prefix only from a genuine enum member.
12501250
assert _coerce_ns(restored.ns) == b"_vgi/streaming_finalize"
1251+
1252+
1253+
class TestPartitionFieldsCacheKeyIsMetadataAware:
1254+
"""Partition-field lookup must not conflate schemas that differ only in metadata.
1255+
1256+
``pa.Schema.__eq__``/``__hash__`` ignore field metadata, so an annotated
1257+
schema is equal to and hashes the same as an unannotated one of the same
1258+
names and types. A cache keyed straight on the schema therefore answers
1259+
whichever was scanned first for both — which made a function with no
1260+
partition annotations inherit an annotated neighbour's fields, so the
1261+
``partition_values`` contract check passed where it had to raise.
1262+
"""
1263+
1264+
@staticmethod
1265+
def _pair() -> tuple[pa.Schema, pa.Schema]:
1266+
from vgi.schema_utils import VGI_PARTITION_COLUMN_KEY
1267+
1268+
plain = pa.schema([("country", pa.string()), ("sales", pa.int64())])
1269+
annotated = pa.schema(
1270+
[
1271+
pa.field("country", pa.string(), metadata={VGI_PARTITION_COLUMN_KEY: b"true"}),
1272+
("sales", pa.int64()),
1273+
]
1274+
)
1275+
return plain, annotated
1276+
1277+
def test_arrow_equality_still_ignores_metadata(self) -> None:
1278+
"""The premise: this is why keying on the schema alone is unsafe."""
1279+
plain, annotated = self._pair()
1280+
assert plain == annotated
1281+
assert hash(plain) == hash(annotated)
1282+
assert not plain.equals(annotated, check_metadata=True)
1283+
1284+
@pytest.mark.parametrize("annotated_first", [True, False])
1285+
def test_lookup_is_correct_in_either_cache_order(self, annotated_first: bool) -> None:
1286+
"""Whichever schema populates the cache first, both answers stay right."""
1287+
from vgi.protocol import _partition_fields_cached, _partition_fields_from_schema
1288+
1289+
plain, annotated = self._pair()
1290+
_partition_fields_cached.cache_clear()
1291+
try:
1292+
first, second = (annotated, plain) if annotated_first else (plain, annotated)
1293+
_partition_fields_from_schema(first)
1294+
_partition_fields_from_schema(second)
1295+
1296+
assert [f.name for f in _partition_fields_from_schema(annotated)] == ["country"]
1297+
assert _partition_fields_from_schema(plain) == []
1298+
finally:
1299+
_partition_fields_cached.cache_clear()

‎vgi/meta_worker.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,37 @@ def _worker_for_unattached(self, function_name: str, schema_name: str | None) ->
247247
raise ValueError(msg)
248248
return candidates[0]
249249

250+
def _unwrap_attach_full(self, envelope: bytes | None) -> bytes | None:
251+
"""Open an attach envelope on the sub-worker that owns its catalog.
252+
253+
A stream's state no longer carries the attach decrypted, so the
254+
rehydrate path re-opens the sealed envelope on whatever object the
255+
transport hands it as the implementation — which under a MetaWorker is
256+
*this* object, not a `[`Worker`][]`. Without this the HTTP continuation
257+
of any catalog-backed stream dies on ``'MetaWorker' object has no
258+
attribute '_unwrap_attach_full'``.
259+
260+
Routes the same way :meth:`_resolve_function` does, so the plaintext is
261+
produced by the sub-worker that owns the catalog. Every sub-worker in a
262+
process shares one signing key (``_maybe_worker_for_attach`` itself
263+
opens the envelope through ``_workers[0]`` to read the catalog name),
264+
so an attach naming no catalog this process knows can still be opened —
265+
fall back to the first worker rather than failing, and let the sub-worker
266+
raise if the seal genuinely does not verify.
267+
268+
Args:
269+
envelope: The sealed ``attach_opaque_data``, or ``None``.
270+
271+
Returns:
272+
The full framework plaintext (``uuid || catalog_bytes``), or
273+
``None`` when there is no attach.
274+
275+
"""
276+
if envelope is None:
277+
return None
278+
worker = self._maybe_worker_for_attach(envelope) or self._workers[0]
279+
return worker._unwrap_attach_full(envelope)
280+
250281
def _resolve_function(self, request: BindRequest) -> Any:
251282
"""Dispatch function-class resolution to the worker that hosts it.
252283

‎vgi/protocol.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -954,12 +954,44 @@ def _partition_fields_from_schema(bind_schema: pa.Schema) -> list[pa.Field[Any]]
954954
once at wrapper construction, so per-emit validation only does an
955955
O(P) walk where P is the partition column count.
956956
"""
957-
return list(_partition_fields_cached(bind_schema))
957+
return list(_partition_fields_cached(_SchemaKey(bind_schema)))
958+
959+
960+
class _SchemaKey:
961+
"""Cache key for a schema whose *field metadata* is significant.
962+
963+
``pa.Schema.__eq__`` and ``__hash__`` ignore field metadata, so a schema
964+
annotated with partition columns is equal to — and hashes the same as — an
965+
unannotated schema of the same names and types. Keying a cache directly on
966+
a schema therefore conflates the two, and whichever was scanned first wins
967+
for both. That is not hypothetical: it made a function with no annotations
968+
inherit an annotated neighbour's partition fields, so the
969+
``partition_values`` contract check silently passed where it had to raise,
970+
and it reproduced only on whichever transport happened to populate the
971+
cache in the losing order.
972+
973+
The metadata-blind hash is kept deliberately — a hash may collide, it just
974+
may not lie — and equality does the discriminating via
975+
``equals(check_metadata=True)``, measured at 0.042us.
976+
"""
977+
978+
__slots__ = ("schema",)
979+
980+
def __init__(self, schema: pa.Schema) -> None:
981+
self.schema = schema
982+
983+
def __hash__(self) -> int:
984+
return hash(self.schema)
985+
986+
def __eq__(self, other: object) -> bool:
987+
if not isinstance(other, _SchemaKey):
988+
return NotImplemented
989+
return bool(self.schema.equals(other.schema, check_metadata=True))
958990

959991

960992
@lru_cache(maxsize=512)
961-
def _partition_fields_cached(bind_schema: pa.Schema) -> tuple[pa.Field[Any], ...]:
962-
"""Memoized partition-field scan, keyed on the schema.
993+
def _partition_fields_cached(key: _SchemaKey) -> tuple[pa.Field[Any], ...]:
994+
"""Memoized partition-field scan, keyed on the schema *including* metadata.
963995
964996
Which columns are partition columns is a pure function of the schema's
965997
field metadata, and the schema is fixed for a stream -- but this ran on
@@ -975,7 +1007,7 @@ def _partition_fields_cached(bind_schema: pa.Schema) -> tuple[pa.Field[Any], ...
9751007
from vgi.schema_utils import VGI_PARTITION_COLUMN_KEY
9761008

9771009
return tuple(
978-
f for f in bind_schema if f.metadata is not None and f.metadata.get(VGI_PARTITION_COLUMN_KEY) == b"true"
1010+
f for f in key.schema if f.metadata is not None and f.metadata.get(VGI_PARTITION_COLUMN_KEY) == b"true"
9791011
)
9801012

9811013

0 commit comments

Comments
 (0)