+# Licensed under the Moderne Source Available License (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +#
+# https://docs.moderne.io/licensing/moderne-source-available-license +#
+# 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. + +"""Ref-id bookkeeping for the send side of the RPC protocol.""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Tuple + + +class ReferenceMap: + """Ref ids assigned to objects sent to a peer, keyed by object identity. + + The peer builds the mirror of this map as it receives, so an id may only be + cited once the object behind it has arrived. Its lifetime is therefore the + peer connection's, and ``rollback_to`` undoes the ids of an exchange the peer + never received. Mirrors ``RewriteRpc.localRefs`` and the JavaScript and Go + ``ReferenceMap``. + """ + + def __init__(self) -> None: + # The object is kept beside its ref so a recycled id() cannot alias a + # freed object onto another object's ref. + self._by_id: Dict[int, Tuple[Any, int]] = {} + self._next: int = 0 + + def get(self, obj: Any) -> Optional[int]: + """The ref already assigned to ``obj``, or None when it has not been sent.""" + entry = self._by_id.get(id(obj)) + return entry[1] if entry is not None and entry[0] is obj else None + + def create(self, obj: Any) -> int: + self._next += 1 + self._by_id[id(obj)] = (obj, self._next) + return self._next + + def snapshot(self) -> int: + """The high-water ref id, for a later ``rollback_to``.""" + return self._next + + def rollback_to(self, snapshot: int) -> None: + """Drop every ref assigned since ``snapshot``.""" + for key in [k for k, (_, ref) in self._by_id.items() if ref > snapshot]: + del self._by_id[key] + self._next = snapshot + + def clear(self) -> None: + self._by_id.clear() + self._next = 0 + + def __len__(self) -> int: + return len(self._by_id) diff --git a/rewrite-python/rewrite/src/rewrite/rpc/send_queue.py b/rewrite-python/rewrite/src/rewrite/rpc/send_queue.py index 76cd1b7e26..92ddf635a2 100644 --- a/rewrite-python/rewrite/src/rewrite/rpc/send_queue.py +++ b/rewrite-python/rewrite/src/rewrite/rpc/send_queue.py @@ -13,6 +13,8 @@ from typing import Any, Dict, List, Optional, Callable, TypeVar from uuid import UUID +from rewrite.rpc.reference import ReferenceMap + class RpcObjectState(str, Enum): NO_CHANGE = "NO_CHANGE" @@ -30,10 +32,11 @@ class RpcObjectState(str, Enum): class RpcSendQueue: """Queue for generating RpcObjectData array from Python LST using visitor pattern.""" - def __init__(self, source_file_type: Optional[str] = None): + def __init__(self, source_file_type: Optional[str] = None, + refs: Optional[ReferenceMap] = None): self.q: List[Dict[str, Any]] = [] - self.refs: Dict[int, tuple] = {} # id(obj) -> (obj, ref_number) — verified with `is` - self.next_ref: int = 0 + # A caller-supplied map spans the peer connection; the default spans this queue. + self.refs: ReferenceMap = refs if refs is not None else ReferenceMap() self.source_file_type = source_file_type self._before: Any = None @@ -220,21 +223,19 @@ def _add_as_ref(self, obj: Any, on_change: Optional[Callable[[], None]] = None) self.put({'state': RpcObjectState.DELETE}) return - obj_id = id(obj) - entry = self.refs.get(obj_id) - if entry is not None and entry[0] is obj: + ref = self.refs.get(obj) + if ref is not None: # Already sent — emit ref number only, no onChange - self.put({'state': RpcObjectState.ADD, 'ref': entry[1]}) + self.put({'state': RpcObjectState.ADD, 'ref': ref}) return # First time — assign ref number and serialize fully - self.next_ref += 1 - self.refs[obj_id] = (obj, self.next_ref) + ref = self.refs.create(obj) value_type = self._get_value_type(obj) codec = self._get_rpc_codec(obj) value = None if on_change is not None or codec is not None else self._get_primitive_value(obj) - self.put({'state': RpcObjectState.ADD, 'valueType': value_type, 'value': value, 'ref': self.next_ref}) + self.put({'state': RpcObjectState.ADD, 'valueType': value_type, 'value': value, 'ref': ref}) self._do_change(obj, None, on_change, codec) def _do_change(self, after: Any, before: Any, diff --git a/rewrite-python/rewrite/src/rewrite/rpc/server.py b/rewrite-python/rewrite/src/rewrite/rpc/server.py index 50ae0a4a45..9fa266ebfe 100644 --- a/rewrite-python/rewrite/src/rewrite/rpc/server.py +++ b/rewrite-python/rewrite/src/rewrite/rpc/server.py @@ -45,6 +45,7 @@ resource = None from rewrite.discovery import RecipeAttribution, RecipeName, _normalize_package_name +from rewrite.rpc.reference import ReferenceMap # Deeply nested LST nodes (e.g., 256 implicitly concatenated strings) can # overflow the default recursion limit (1000) during RPC serialization. @@ -67,9 +68,13 @@ remote_objects: Dict[str, Any] = {} # Remote refs - maps reference IDs to objects for cyclic graph handling remote_refs: Dict[int, Any] = {} -# Per-source-file remote_refs high-water, captured before a file is first visited so -# handle_evict can roll back exactly the refs that file introduced. Keyed by tree id. +# Refs sent to Java, so a type sent for one source file is cited rather than resent +# by the next (mirrors RewriteRpc.localRefs). +local_refs = ReferenceMap() +# Per-source-file ref high-water on each side, captured before a file is first visited +# so handle_evict can roll back exactly the refs that file introduced. Keyed by tree id. _ref_checkpoints: Dict[str, int] = {} +_local_ref_checkpoints: Dict[str, int] = {} # Per-call metrics CSV (--metrics-csv), same schema as Go: cache-size ramp vs per-file-Evict sawtooth. _metrics_file = None @@ -1008,8 +1013,15 @@ def handle_get_object(params: dict) -> List[dict]: # Get the "before" state - what we previously sent to Java before = remote_objects.get(obj_id) - q = RpcSendQueue(source_file_type) - result = q.generate(obj, before) + saved_refs = local_refs.snapshot() + q = RpcSendQueue(source_file_type, local_refs) + try: + result = q.generate(obj, before) + except BaseException: + # Java receives nothing of this exchange, so refs assigned during it name + # objects it never got; citing them later would fail on its side. + local_refs.rollback_to(saved_refs) + raise # Update remote_objects to track that Java now has this version remote_objects[obj_id] = obj @@ -1111,6 +1123,8 @@ def handle_reset(params: dict) -> bool: _recipe_accumulators.clear() _recipe_phases.clear() _ref_checkpoints.clear() + _local_ref_checkpoints.clear() + local_refs.clear() logger.info("Reset: cleared all cached state") return True @@ -1120,6 +1134,9 @@ def handle_evict(params: dict) -> bool: """Handle an Evict RPC notification - drop one source file's tree and roll back the refs it introduced, bounding memory to roughly one source file at a time. Recipe, accumulator, and execution-context state (keyed separately) is left intact. + + Both directions roll back together, matching RewriteRpc.evict: Java drops this file's + refs from both of its maps, so any kept here would name ids it no longer holds. """ obj_id = params.get('id') if obj_id is None: @@ -1130,6 +1147,9 @@ def handle_evict(params: dict) -> bool: if checkpoint is not None: for ref_id in [k for k in remote_refs if k > checkpoint]: del remote_refs[ref_id] + local_checkpoint = _local_ref_checkpoints.pop(obj_id, None) + if local_checkpoint is not None: + local_refs.rollback_to(local_checkpoint) return True @@ -2062,8 +2082,10 @@ def handle_visit(params: dict) -> dict: ctx = _context_for(p_id) - # Snapshot the remote_refs high-water for this file before fetching its tree (first visit wins). + # Snapshot both directions' ref high-water for this file before fetching its tree + # (first visit wins). _ref_checkpoints.setdefault(tree_id, max(remote_refs.keys(), default=-1)) + _local_ref_checkpoints.setdefault(tree_id, local_refs.snapshot()) # Always fetch the tree from Java to ensure we have the latest version. # Java may have modified the tree (e.g., via a Java-side recipe) since our last sync. @@ -2120,8 +2142,9 @@ def handle_batch_visit(params: dict) -> dict: ctx = _context_for(p_id) - # Snapshot the remote_refs high-water for this file before fetching its tree. + # Snapshot both directions' ref high-water for this file before fetching its tree. _ref_checkpoints.setdefault(tree_id, max(remote_refs.keys(), default=-1)) + _local_ref_checkpoints.setdefault(tree_id, local_refs.snapshot()) # Fetch tree once from Java tree = get_object_from_java(tree_id, source_file_type) @@ -2342,8 +2365,7 @@ def handle_generate(params: dict) -> dict: # pair, so one child's ref numbering never has to mean anything to another child. Relaying a # child's stream directly to a sibling is unsafe. _hub_tree: Dict[str, Any] = {} # obj_id -> the facade's authoritative tree -_hub_send_refs: Dict[str, Dict] = {} # bundle -> send ref map (facade -> child) -_hub_send_next: Dict[str, int] = {} # bundle -> next send ref number +_hub_send_refs: Dict[str, ReferenceMap] = {} # bundle -> send ref map (facade -> child) _hub_served: Dict[tuple, Any] = {} # (bundle, obj_id) -> what that child was last served _hub_send_checkpoint: Dict[tuple, int] = {} # (bundle, obj_id) -> send ref counter before this file @@ -2373,14 +2395,11 @@ def _hub_serve_child(bundle: str, obj_id: str, source_file_type: Optional[str]) if tree is None: return [{'state': 'DELETE'}, {'state': 'END_OF_OBJECT'}] - q = RpcSendQueue(source_file_type) - q.refs = _hub_send_refs.setdefault(bundle, {}) - q.next_ref = _hub_send_next.get(bundle, 0) + refs = _hub_send_refs.setdefault(bundle, ReferenceMap()) # Remember where this child's ref numbering stood before this file, so Evict can roll it back # in lockstep with the child's own rollback (see _hub_release). - _hub_send_checkpoint.setdefault((bundle, obj_id), q.next_ref) - data = q.generate(tree, _hub_served.get((bundle, obj_id))) - _hub_send_next[bundle] = q.next_ref + _hub_send_checkpoint.setdefault((bundle, obj_id), refs.snapshot()) + data = RpcSendQueue(source_file_type, refs).generate(tree, _hub_served.get((bundle, obj_id))) _hub_served[(bundle, obj_id)] = tree return data @@ -2417,13 +2436,10 @@ def _hub_release(obj_id: str) -> None: for key in [k for k in _hub_served if k[1] == obj_id]: del _hub_served[key] for key in [k for k in _hub_send_checkpoint if k[1] == obj_id]: - bundle = key[0] checkpoint = _hub_send_checkpoint.pop(key) - refs = _hub_send_refs.get(bundle) + refs = _hub_send_refs.get(key[0]) if refs is not None: - for ref_key in [k for k, (_, num) in refs.items() if num > checkpoint]: - del refs[ref_key] - _hub_send_next[bundle] = checkpoint + refs.rollback_to(checkpoint) def _hub_is_builtin_visitor(visitor_name: Optional[str]) -> bool: @@ -2868,8 +2884,8 @@ def _rss_bytes(): def _record_metric(method: str, duration_ms: float, error: str) -> None: - """Append one row of timing + cache residency. refs counts remote_refs only: Python's send-side - refs live on a per-call RpcSendQueue, the only cross-call ref cache handle_evict rolls back.""" + """Append one row of timing + cache residency. refs counts both connection-scoped ref + tables, the ones handle_evict rolls back: what Java sent us and what we sent Java.""" if _metrics_writer is None: return used, peak = _rss_bytes() @@ -2879,7 +2895,8 @@ def _record_metric(method: str, duration_ms: float, error: str) -> None: try: _metrics_writer.writerow([ datetime.now(timezone.utc).isoformat(), method, f"{duration_ms:.0f}", error, - used, peak, len(local_objects), len(remote_objects), len(remote_refs)]) + used, peak, len(local_objects), len(remote_objects), + len(remote_refs) + len(local_refs)]) _metrics_file.flush() except OSError as e: logger.warning(f"metrics-csv: write failed: {e}") diff --git a/rewrite-python/rewrite/tests/python/test_type_interning.py b/rewrite-python/rewrite/tests/python/test_type_interning.py new file mode 100644 index 0000000000..25104e9319 --- /dev/null +++ b/rewrite-python/rewrite/tests/python/test_type_interning.py @@ -0,0 +1,179 @@ +# Copyright 2025 the original author or authors. +#
+# Licensed under the Moderne Source Available License (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +#
+# https://docs.moderne.io/licensing/moderne-source-available-license +#
+# 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. +"""How many distinct JavaType objects a whole-project parse materializes.""" +import dataclasses +import shutil +from enum import Enum + +import pytest + +from rewrite.java import JavaType +from rewrite.rpc import server + +requires_ty_types_cli = pytest.mark.skipif( + shutil.which('ty-types') is None, + reason="ty-types CLI is not installed (ensure ty-types binary is on PATH)", +) + +# Every concrete JavaType shape except Primitive, whose members are enum +# singletons and so are shared no matter how the caches are scoped. +_JAVA_TYPES = ( + JavaType.FullyQualified, + JavaType.Method, + JavaType.Variable, + JavaType.GenericTypeVariable, + JavaType.Union, + JavaType.Intersection, +) + +_LEAVES = (str, bytes, int, float, bool, complex, Enum, type) + +# A first-party module every generated file imports, so the corpus has a type +# whose definition and references are spread across files. +_COMMON = '''\ +from typing import List + + +class Shared: + name: str + values: List[int] + + def total(self) -> int: + return sum(self.values) +''' + +_SOURCE = '''\ +import json +import os.path +from typing import Dict, List + +from common import Shared + + +class Record{i}: + name: str + values: List[int] + + def to_json(self) -> str: + return json.dumps({{"name": self.name, "values": self.values}}) + + +def load{i}(path: str) -> Dict[str, object]: + with open(os.path.join(path, "data.json")) as fh: + return json.load(fh) + + +def totals{i}(shared: List[Shared]) -> List[int]: + return [s.total() for s in shared] +''' + + +def _fields_of(obj): + """The referenced values of an arbitrary object, whatever it uses to store them.""" + if dataclasses.is_dataclass(obj): + return [getattr(obj, f.name, None) for f in dataclasses.fields(obj)] + instance_dict = getattr(obj, '__dict__', None) + if instance_dict is not None: + return list(instance_dict.values()) + return [getattr(obj, slot, None) + for klass in type(obj).__mro__ + for slot in getattr(klass, '__slots__', ()) or ()] + + +def _java_types(root): + """Every JavaType object reachable from ``root``, keyed by identity.""" + found = {} + seen = set() + stack = [root] + while stack: + obj = stack.pop() + if obj is None or isinstance(obj, _LEAVES): + continue + if id(obj) in seen: + continue + seen.add(id(obj)) + if isinstance(obj, _JAVA_TYPES): + found[id(obj)] = obj + if isinstance(obj, (list, tuple, set, frozenset)): + stack.extend(obj) + elif isinstance(obj, dict): + stack.extend(obj.values()) + else: + stack.extend(_fields_of(obj)) + return found + + +def _parse_project(tmp_path, module_count): + """Parse a generated project, returning each source path's JavaType set.""" + (tmp_path / "common.py").write_text(_COMMON) + for i in range(module_count): + (tmp_path / f"mod{i}.py").write_text(_SOURCE.format(i=i)) + + results = server.handle_parse_project( + {'projectPath': str(tmp_path), 'relativeTo': str(tmp_path)}) + assert len(results) == module_count + 1 + return {result['sourcePath']: _java_types(server.local_objects[result['id']]) + for result in results} + + +@requires_ty_types_cli +def test_java_types_are_shared_across_the_files_of_one_parse_project(tmp_path): + per_file = _parse_project(tmp_path, 10) + + largest_file = max(len(types) for types in per_file.values()) + assert largest_file > 100, "type attribution did not run; the count is meaningless" + + corpus = {} + for types in per_file.values(): + corpus.update(types) + + # One object per type puts the corpus within reach of the richest single + # file; one object per type per file puts it near the sum. + assert len(corpus) < largest_file * 2 + + +@requires_ty_types_cli +def test_same_named_typed_dicts_in_two_modules_keep_their_own_fields(tmp_path): + (tmp_path / "a_mod.py").write_text( + 'from typing import TypedDict\n\n\n' + 'class Movie(TypedDict):\n name: str\n\n\n' + 'def pick() -> Movie:\n return Movie(name="x")\n') + (tmp_path / "b_mod.py").write_text( + 'from typing import TypedDict\n\n\n' + 'class Movie(TypedDict):\n year: int\n director: str\n\n\n' + 'def pick() -> Movie:\n return Movie(year=1, director="y")\n') + + results = server.handle_parse_project( + {'projectPath': str(tmp_path), 'relativeTo': str(tmp_path)}) + fields = {} + for result in results: + for java_type in _java_types(server.local_objects[result['id']]).values(): + if getattr(java_type, 'fully_qualified_name', None) == 'Movie': + fields[result['sourcePath']] = sorted( + m.name for m in (getattr(java_type, '_members', None) or [])) + + assert fields == {'a_mod.py': ['name'], 'b_mod.py': ['director', 'year']} + + +@requires_ty_types_cli +def test_a_class_defined_in_one_file_and_used_in_another_is_one_object(tmp_path): + per_file = _parse_project(tmp_path, 2) + + def shared_class(source_path): + return {oid for oid, t in per_file[source_path].items() + if getattr(t, 'fully_qualified_name', None) == 'common.Shared'} + + in_definition = shared_class('common.py') + assert in_definition, "common.Shared was not attributed in its own module" + assert in_definition == shared_class('mod0.py') == shared_class('mod1.py') diff --git a/rewrite-python/rewrite/tests/rpc/test_facade.py b/rewrite-python/rewrite/tests/rpc/test_facade.py index 699398ea07..309f5102cc 100644 --- a/rewrite-python/rewrite/tests/rpc/test_facade.py +++ b/rewrite-python/rewrite/tests/rpc/test_facade.py @@ -297,20 +297,23 @@ def generate(self, params): ... def test_hub_release_rolls_each_childs_ref_table_back_in_lockstep(): import rewrite.rpc.server as server + from rewrite.rpc.reference import ReferenceMap server._hub_tree["T"] = object() server._hub_served[("A", "T")] = object() # child A had 2 refs before this file, then the file introduced refs 3 and 4 - server._hub_send_refs["A"] = {10: ("before-1", 1), 11: ("before-2", 2), - 12: ("from-file", 3), 13: ("from-file", 4)} - server._hub_send_next["A"] = 4 + refs = server._hub_send_refs["A"] = ReferenceMap() + survivors = [object(), object()] + for obj in survivors + [object(), object()]: + refs.create(obj) server._hub_send_checkpoint[("A", "T")] = 2 server._hub_release("T") # only the refs this file introduced are dropped; the pre-file ones survive - assert sorted(n for _, n in server._hub_send_refs["A"].values()) == [1, 2] - assert server._hub_send_next["A"] == 2 # counter rewound so the next file re-ADDs + assert [refs.get(obj) for obj in survivors] == [1, 2] + assert len(refs) == 2 + assert refs.snapshot() == 2 # counter rewound so the next file re-ADDs assert "T" not in server._hub_tree assert ("A", "T") not in server._hub_served assert ("A", "T") not in server._hub_send_checkpoint @@ -404,7 +407,7 @@ def _print_python(cu) -> str: def _isolated_hub(monkeypatch, server): """Fresh hub state so this test neither sees nor leaves module-global tables.""" - for name in ("_hub_tree", "_hub_served", "_hub_send_refs", "_hub_send_next", + for name in ("_hub_tree", "_hub_served", "_hub_send_refs", "_hub_send_checkpoint", "local_objects"): monkeypatch.setattr(server, name, {}) @@ -496,7 +499,7 @@ def visit_compilation_unit(self, cu, p): server._hub_acquire(tree_id, sft) server._hub_serve_child(bundle, tree_id, sft) # child now holds this file and its refs - assert server._hub_send_next[bundle] > 0 + assert server._hub_send_refs[bundle].snapshot() > 0 results = server._hub_local_visit([{"visitor": "Delete"}, {"visitor": "Later"}], {"treeId": tree_id, "sourceFileType": sft}) @@ -510,8 +513,8 @@ def visit_compilation_unit(self, cu, p): assert tree_id not in server._hub_tree assert (bundle, tree_id) not in server._hub_served assert (bundle, tree_id) not in server._hub_send_checkpoint - assert server._hub_send_refs[bundle] == {} - assert server._hub_send_next[bundle] == 0 + assert len(server._hub_send_refs[bundle]) == 0 + assert server._hub_send_refs[bundle].snapshot() == 0 # So a child asking for it again is told the file is gone rather than served a stale tree. assert server._hub_serve_child(bundle, tree_id, sft) == [ diff --git a/rewrite-python/rewrite/tests/rpc/test_reference_map.py b/rewrite-python/rewrite/tests/rpc/test_reference_map.py new file mode 100644 index 0000000000..69b885f804 --- /dev/null +++ b/rewrite-python/rewrite/tests/rpc/test_reference_map.py @@ -0,0 +1,77 @@ +# Copyright 2025 the original author or authors. +#
+# Licensed under the Moderne Source Available License (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +#
+# https://docs.moderne.io/licensing/moderne-source-available-license +#
+# 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. +"""How ref ids are assigned, reused, and rolled back.""" +from rewrite.rpc.reference import ReferenceMap + + +def test_ids_start_at_one_and_ascend(): + refs = ReferenceMap() + assert refs.create(object()) == 1 + assert refs.create(object()) == 2 + + +def test_an_object_keeps_the_id_it_was_given(): + refs = ReferenceMap() + obj = object() + assert refs.create(obj) == refs.get(obj) + + +def test_equal_but_distinct_objects_get_their_own_ids(): + refs = ReferenceMap() + first, second = ['a'], ['a'] + assert refs.create(first) != refs.create(second) + + +def test_an_unsent_object_has_no_id(): + assert ReferenceMap().get(object()) is None + + +def test_an_id_recycled_by_the_allocator_is_not_mistaken_for_a_sent_object(): + refs = ReferenceMap() + recycled = id(refs.create([1, 2, 3]) and [4, 5, 6]) + # Whatever now occupies that address, nothing was sent under it. + assert all(refs.get(obj) is None for obj in ([4, 5, 6], [1, 2, 3])) + assert isinstance(recycled, int) + + +def test_rollback_drops_only_what_came_after_the_snapshot(): + refs = ReferenceMap() + kept = object() + refs.create(kept) + snapshot = refs.snapshot() + dropped = object() + refs.create(dropped) + + refs.rollback_to(snapshot) + assert refs.get(kept) == 1 + assert refs.get(dropped) is None + + +def test_ids_resume_from_the_snapshot_after_a_rollback(): + refs = ReferenceMap() + refs.create(object()) + snapshot = refs.snapshot() + refs.create(object()) + refs.rollback_to(snapshot) + + assert refs.create(object()) == 2 + + +def test_clear_restarts_the_sequence(): + refs = ReferenceMap() + refs.create(object()) + refs.clear() + + assert len(refs) == 0 + assert refs.create(object()) == 1 diff --git a/rewrite-python/rewrite/tests/rpc/test_send_ref_reuse.py b/rewrite-python/rewrite/tests/rpc/test_send_ref_reuse.py new file mode 100644 index 0000000000..d2aea3255c --- /dev/null +++ b/rewrite-python/rewrite/tests/rpc/test_send_ref_reuse.py @@ -0,0 +1,129 @@ +# Copyright 2025 the original author or authors. +#
+# Licensed under the Moderne Source Available License (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +#
+# https://docs.moderne.io/licensing/moderne-source-available-license +#
+# 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. +"""Whether refs sent for one source file are reusable by the next.""" +import shutil + +import pytest + +from rewrite.rpc import server + +requires_ty_types_cli = pytest.mark.skipif( + shutil.which('ty-types') is None, + reason="ty-types CLI is not installed (ensure ty-types binary is on PATH)", +) + +_SOURCE = '''\ +import json +from typing import Dict + + +def load{i}(text: str) -> Dict[str, object]: + return json.loads(text) + + +def dump{i}(value: Dict[str, object]) -> str: + return json.dumps(value) +''' + +_CU_TYPE = 'org.openrewrite.python.tree.Py$CompilationUnit' + + +def defined(response): + """Ref ids this response assigns, i.e. sends the full object for.""" + return {d['ref'] for d in response if 'ref' in d and 'valueType' in d} + + +def used(response): + """Ref ids this response cites without resending the object.""" + return {d['ref'] for d in response if 'ref' in d and 'valueType' not in d} + + +@pytest.fixture(autouse=True) +def clean_connection_state(): + server.handle_reset({}) + yield + server.handle_reset({}) + + +def parse_two_files(tmp_path): + for i in range(2): + (tmp_path / f"mod{i}.py").write_text(_SOURCE.format(i=i)) + results = server.handle_parse_project( + {'projectPath': str(tmp_path), 'relativeTo': str(tmp_path)}) + assert len(results) == 2 + return sorted(r['id'] for r in results) + + +def get_object(obj_id): + return server.handle_get_object({'id': obj_id, 'sourceFileType': _CU_TYPE}) + + +@requires_ty_types_cli +def test_the_second_file_cites_refs_the_first_file_sent(tmp_path): + first_id, second_id = parse_two_files(tmp_path) + + first = get_object(first_id) + second = get_object(second_id) + + assert defined(first), "the first file sent no refs; the test proves nothing" + # Ref ids are one ascending sequence across the connection, so the second file's + # ids continue past the first's rather than restarting and renaming them. + assert not (defined(first) & defined(second)), "a ref id was assigned twice" + + inherited = used(second) - defined(second) + assert inherited, "the second file resent everything" + assert inherited <= defined(first) + assert len(defined(second)) < len(defined(first)) / 2 + + +@requires_ty_types_cli +def test_evicting_a_file_releases_the_refs_it_introduced(tmp_path): + first_id, second_id = parse_two_files(tmp_path) + + # The sequence Java drives per file: checkpoint on first visit, pull the tree, evict. + server._local_ref_checkpoints.setdefault(first_id, server.local_refs.snapshot()) + get_object(first_id) + server.handle_evict({'id': first_id}) + + second = get_object(second_id) + assert used(second) <= defined(second), "cited a ref Java rolled back" + assert len(server.local_refs) == len(defined(second)) + + +@requires_ty_types_cli +def test_a_failed_transfer_releases_the_refs_it_assigned(tmp_path, monkeypatch): + from rewrite.rpc import send_queue + + first_id, second_id = parse_two_files(tmp_path) + get_object(first_id) + before_failure = server.local_refs.snapshot() + + def fail_midway(self, after, before=None): + self.refs.create(object()) + raise RuntimeError("transfer failed") + + monkeypatch.setattr(send_queue.RpcSendQueue, 'generate', fail_midway) + + assert get_object(second_id) == [{'state': 'END_OF_OBJECT'}] + assert server.local_refs.snapshot() == before_failure + + +@requires_ty_types_cli +def test_reset_clears_the_sent_refs(tmp_path): + first_id, _ = parse_two_files(tmp_path) + get_object(first_id) + assert len(server.local_refs) > 0 + + server.handle_reset({}) + assert len(server.local_refs) == 0 diff --git a/rewrite-python/rewrite/tests/rpc/test_server.py b/rewrite-python/rewrite/tests/rpc/test_server.py index 8e3062a36e..88574e2e3e 100644 --- a/rewrite-python/rewrite/tests/rpc/test_server.py +++ b/rewrite-python/rewrite/tests/rpc/test_server.py @@ -531,27 +531,27 @@ def test_hub_release_rewinds_send_refs_in_lockstep_with_the_child(): it would emit a GET_REF for a ref the child no longer holds; if it rewound while the child did not, it would reuse a number still bound to the old object and serve a wrong tree silently.""" import rewrite.rpc.server as server + from rewrite.rpc.reference import ReferenceMap bundle, first, second = "pkg", "file-1", "file-2" - server._hub_send_refs[bundle] = {} - server._hub_send_next[bundle] = 0 + refs = server._hub_send_refs[bundle] = ReferenceMap() # Serving the first file advances this child's numbering and records where it started. - server._hub_send_checkpoint.setdefault((bundle, first), server._hub_send_next[bundle]) - server._hub_send_refs[bundle].update({"obj-a": (object(), 1), "obj-b": (object(), 2)}) - server._hub_send_next[bundle] = 2 + server._hub_send_checkpoint.setdefault((bundle, first), refs.snapshot()) + refs.create(object()) + refs.create(object()) server._hub_served[(bundle, first)] = object() server._hub_tree[first] = object() server._hub_release(first) # Everything that file introduced is gone, and the counter is back where it began. - assert server._hub_send_next[bundle] == 0 - assert server._hub_send_refs[bundle] == {} + assert refs.snapshot() == 0 + assert len(refs) == 0 assert (bundle, first) not in server._hub_served assert (bundle, first) not in server._hub_send_checkpoint assert first not in server._hub_tree # So the next file reuses the same ref numbers rather than continuing past them. - server._hub_send_checkpoint.setdefault((bundle, second), server._hub_send_next[bundle]) + server._hub_send_checkpoint.setdefault((bundle, second), refs.snapshot()) assert server._hub_send_checkpoint[(bundle, second)] == 0