gabriel / musehub public
test_wire_compression.py python
488 lines 19.4 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago
1 """TDD — fetch compression: zstd + path-sorted objects + delta frames.
2
3 Test plan
4 ---------
5 Layer 1 — zstd
6 C1 Server uses zstd encoding (not zlib) for O frames when zstd is available.
7 C2 Client can decode the zstd frame back to original bytes.
8
9 Layer 2 — path-sorted objects
10 C3 O frames for a multi-object fetch arrive sorted by path within each batch.
11 Sorted order is a prerequisite for profitable intra-batch compression and
12 deterministic delta base selection.
13
14 Layer 3 — delta frames
15 C4 Two successive versions of the same file → server sends the second as a
16 delta frame (enc="delta+zstd") with a non-empty base_id.
17 C5 Delta frame content reconstructs to the correct target bytes.
18 C6 Delta is not sent when unprofitable (random content → no shared runs →
19 delta ≥ full object; server falls back to full zstd frame).
20 C7 Full clone with a mix of full + delta frames delivers the right object_id
21 set to the client.
22 C8 Objects from a prior batch serve as delta bases for the next batch
23 (cross-batch delta).
24 """
25 from __future__ import annotations
26
27 import zlib
28 from datetime import datetime, timezone
29
30 import msgpack
31 import pytest
32 from sqlalchemy.ext.asyncio import AsyncSession
33
34 from muse.core.compression import ZSTD_AVAILABLE, apply_delta, compute_delta
35 from muse.core.types import blob_id, fake_id
36 from musehub.db import musehub_models as db
37 from musehub.models.wire import (
38 SFRAME_COMMIT_PACK,
39 SFRAME_END,
40 SFRAME_HEADER,
41 SFRAME_OBJECT,
42 WireFetchRequest,
43 )
44 from tests.factories import create_repo
45
46
47 # ---------------------------------------------------------------------------
48 # Helpers (mirror test_wire_batch_stream.py)
49 # ---------------------------------------------------------------------------
50
51 def _uid(seed: str) -> str:
52 return fake_id(seed)
53
54
55 def _now() -> datetime:
56 return datetime.now(tz=timezone.utc)
57
58
59 async def _make_object(
60 session: AsyncSession,
61 repo_id: str,
62 content: bytes,
63 path: str = "file.dat",
64 *,
65 owner: str = "",
66 slug: str = "",
67 ) -> str:
68 from musehub.services.musehub_wire import get_backend
69 from musehub.storage.backends import repo_root_for
70 from sqlalchemy.dialects.postgresql import insert as pg_insert
71
72 backend = get_backend(owner or None, slug or None)
73 oid = blob_id(content)
74 # Use per-repo root so wire_fetch_stream (which passes repo_root=repo_root_for(...))
75 # reads from the same location we write to.
76 per_repo_root = repo_root_for(owner, slug) if (owner and slug) else None
77 uri = await backend.put(oid, content, repo_root=per_repo_root)
78 await session.execute(
79 pg_insert(db.MusehubObject)
80 .values(
81 object_id=oid,
82 path=path,
83 size_bytes=len(content),
84 disk_path=uri.replace("local://", ""),
85 storage_uri=uri,
86 )
87 .on_conflict_do_nothing(index_elements=["object_id"])
88 )
89 await session.execute(
90 pg_insert(db.MusehubObjectRef)
91 .values(repo_id=repo_id, object_id=oid)
92 .on_conflict_do_nothing()
93 )
94 await session.commit()
95 return oid
96
97
98 async def _make_snapshot(
99 session: AsyncSession,
100 repo_id: str,
101 manifest: dict[str, str],
102 ) -> db.MusehubSnapshot:
103 sid = _uid(str(sorted(manifest.items())))
104 snap = db.MusehubSnapshot(
105 snapshot_id=sid,
106 repo_id=repo_id,
107 directories=[],
108 manifest_blob=msgpack.packb(manifest, use_bin_type=True),
109 entry_count=len(manifest),
110 created_at=_now(),
111 )
112 session.add(snap)
113 await session.commit()
114 return snap
115
116
117 async def _make_commit(
118 session: AsyncSession,
119 repo_id: str,
120 *,
121 parent_ids: list[str] | None = None,
122 snapshot_id: str | None = None,
123 seed: str = "",
124 ) -> db.MusehubCommit:
125 row = db.MusehubCommit(
126 commit_id=_uid(f"commit-{seed}"),
127 repo_id=repo_id,
128 branch="main",
129 parent_ids=parent_ids or [],
130 message=f"commit {seed}",
131 author="gabriel",
132 timestamp=_now(),
133 snapshot_id=snapshot_id,
134 )
135 session.add(row)
136 await session.commit()
137 return row
138
139
140 async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]:
141 unpacker = msgpack.Unpacker(raw=False)
142 async for chunk in gen:
143 unpacker.feed(chunk)
144 return list(unpacker)
145
146
147 def _o_frames(frames: list[dict]) -> list[dict]:
148 return [f for f in frames if f.get("t") == SFRAME_OBJECT]
149
150
151 def _realistic_source(version: str = "v1") -> bytes:
152 """~4 KB of varied Python source — compresses 2:1 with zstd, not to nothing.
153
154 Generates 60 unique functions with distinct hex identifiers so the content
155 is varied enough that zstd does not win with an extreme ratio. A tiny
156 one-character version change creates a delta that is much smaller than the
157 compressed full object — ideal for testing delta profitability.
158 """
159 lines: list[str] = [
160 "# Auto-generated module\n",
161 f"__version__ = '{version}'\n",
162 "import hashlib, struct, os, sys, json, logging\n",
163 "from typing import Dict, List, Optional, Any\n",
164 "\n",
165 "logger = logging.getLogger(__name__)\n",
166 "\n",
167 ]
168 for i in range(60):
169 h = hex(i * 0x9E3779B9 & 0xFFFFFFFF)[2:].zfill(8)
170 lines += [
171 f"def compute_{h}(x: int, ctx: Dict[str, Any]) -> int:\n",
172 f" base = ctx.get('{h}', {i * 7 + 1})\n",
173 f" return (x ^ base) % {0xFFFF - i * 3}\n",
174 "\n",
175 ]
176 return "".join(lines).encode()
177
178
179 # ---------------------------------------------------------------------------
180 # C1 — Server uses zstd encoding on O frames
181 # ---------------------------------------------------------------------------
182
183 @pytest.mark.asyncio
184 @pytest.mark.skipif(not ZSTD_AVAILABLE, reason="zstd not installed")
185 async def test_c1_server_uses_zstd_encoding(db_session: AsyncSession) -> None:
186 repo = await create_repo(db_session, owner="test-comp-c1")
187 repo_id = str(repo.repo_id)
188
189 content = b"hello world " * 100
190 oid = await _make_object(db_session, repo_id, content, path="c1/file.py", owner=repo.owner, slug=repo.slug)
191 snap = await _make_snapshot(db_session, repo_id, {"c1/file.py": oid})
192 commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c1")
193
194 from musehub.services.musehub_wire import wire_fetch_stream
195 frames = await _collect_frames(
196 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
197 )
198 o = _o_frames(frames)
199 assert o, "Expected at least one O frame"
200 assert all(f.get("enc") == "zstd" for f in o), (
201 f"Expected enc=zstd on all O frames, got: {[f.get('enc') for f in o]}"
202 )
203
204
205 # ---------------------------------------------------------------------------
206 # C2 — zstd frame decompresses back to original bytes
207 # ---------------------------------------------------------------------------
208
209 @pytest.mark.asyncio
210 @pytest.mark.skipif(not ZSTD_AVAILABLE, reason="zstd not installed")
211 async def test_c2_zstd_frame_decompresses_correctly(db_session: AsyncSession) -> None:
212 import zstandard
213
214 repo = await create_repo(db_session, owner="test-comp-c2")
215 repo_id = str(repo.repo_id)
216
217 content = b"def foo(): pass\n" * 50
218 oid = await _make_object(db_session, repo_id, content, path="c2/mod.py", owner=repo.owner, slug=repo.slug)
219 snap = await _make_snapshot(db_session, repo_id, {"c2/mod.py": oid})
220 commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c2")
221
222 from musehub.services.musehub_wire import wire_fetch_stream
223 frames = await _collect_frames(
224 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
225 )
226 o = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == oid), None)
227 assert o is not None
228 assert o["enc"] == "zstd"
229
230 reconstructed = zstandard.ZstdDecompressor().decompress(bytes(o["content"]))
231 assert reconstructed == content
232
233
234 # ---------------------------------------------------------------------------
235 # C3 — O frames arrive sorted by path within each batch
236 # ---------------------------------------------------------------------------
237
238 @pytest.mark.asyncio
239 async def test_c3_o_frames_sorted_by_path(db_session: AsyncSession) -> None:
240 repo = await create_repo(db_session, owner="test-comp-c3")
241 repo_id = str(repo.repo_id)
242
243 # Create objects at paths that would be out of alphabetical order if unsorted
244 paths = ["z/last.py", "a/first.py", "m/middle.py", "b/second.py"]
245 manifest: dict[str, str] = {}
246 for p in paths:
247 oid = await _make_object(db_session, repo_id, f"content of {p}".encode(), path=p, owner=repo.owner, slug=repo.slug)
248 manifest[p] = oid
249
250 snap = await _make_snapshot(db_session, repo_id, manifest)
251 commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c3")
252
253 from musehub.services.musehub_wire import wire_fetch_stream
254 frames = await _collect_frames(
255 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
256 )
257 o_paths = [f.get("path", "") for f in frames if f.get("t") == SFRAME_OBJECT]
258 assert o_paths == sorted(o_paths), (
259 f"O frames not sorted by path.\nGot: {o_paths}\nExpected: {sorted(o_paths)}"
260 )
261
262
263 # ---------------------------------------------------------------------------
264 # C4 — Successive versions of same file → delta frame with base_id
265 # ---------------------------------------------------------------------------
266
267 @pytest.mark.asyncio
268 async def test_c4_successive_versions_use_delta_frame(db_session: AsyncSession) -> None:
269 repo = await create_repo(db_session, owner="test-comp-c4")
270 repo_id = str(repo.repo_id)
271
272 # Realistic Python source: large varied content, tiny version change.
273 # Delta must win against zstd full frame for this test to be meaningful.
274 base_content = _realistic_source("v1")
275 new_content = _realistic_source("v2")
276
277 base_oid = await _make_object(db_session, repo_id, base_content, path="src/greet.py", owner=repo.owner, slug=repo.slug)
278 new_oid = await _make_object(db_session, repo_id, new_content, path="src/greet.py", owner=repo.owner, slug=repo.slug)
279
280 snap1 = await _make_snapshot(db_session, repo_id, {"src/greet.py": base_oid})
281 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c4-1")
282
283 snap2 = await _make_snapshot(db_session, repo_id, {"src/greet.py": new_oid})
284 c2 = await _make_commit(
285 db_session, repo_id,
286 parent_ids=[c1.commit_id],
287 snapshot_id=snap2.snapshot_id,
288 seed="c4-2",
289 )
290
291 from musehub.services.musehub_wire import wire_fetch_stream
292 frames = await _collect_frames(
293 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c2.commit_id], have=[]))
294 )
295 o_frames = _o_frames(frames)
296
297 # The first version (base) must be a full frame; the second a delta
298 frame_for_new = next((f for f in o_frames if f.get("id") == new_oid), None)
299 assert frame_for_new is not None, f"No O frame for {new_oid[:16]}"
300 assert "delta" in frame_for_new.get("enc", ""), (
301 f"Expected delta encoding for successive version, got enc={frame_for_new.get('enc')!r}"
302 )
303 assert frame_for_new.get("base_id") == base_oid, (
304 f"base_id should be {base_oid[:16]}, got {str(frame_for_new.get('base_id'))[:16]}"
305 )
306
307
308 # ---------------------------------------------------------------------------
309 # C5 — Delta frame reconstructs to correct target bytes
310 # ---------------------------------------------------------------------------
311
312 @pytest.mark.asyncio
313 async def test_c5_delta_frame_reconstructs_correctly(db_session: AsyncSession) -> None:
314 repo = await create_repo(db_session, owner="test-comp-c5")
315 repo_id = str(repo.repo_id)
316
317 base_content = _realistic_source("a1")
318 new_content = _realistic_source("a2")
319
320 base_oid = await _make_object(db_session, repo_id, base_content, path="lib/foo.py", owner=repo.owner, slug=repo.slug)
321 new_oid = await _make_object(db_session, repo_id, new_content, path="lib/foo.py", owner=repo.owner, slug=repo.slug)
322
323 snap1 = await _make_snapshot(db_session, repo_id, {"lib/foo.py": base_oid})
324 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c5-1")
325 snap2 = await _make_snapshot(db_session, repo_id, {"lib/foo.py": new_oid})
326 c2 = await _make_commit(
327 db_session, repo_id,
328 parent_ids=[c1.commit_id],
329 snapshot_id=snap2.snapshot_id,
330 seed="c5-2",
331 )
332
333 from musehub.services.musehub_wire import wire_fetch_stream
334 frames = await _collect_frames(
335 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c2.commit_id], have=[]))
336 )
337 o_frames = _o_frames(frames)
338
339 frame_base = next((f for f in o_frames if f.get("id") == base_oid), None)
340 frame_new = next((f for f in o_frames if f.get("id") == new_oid), None)
341 assert frame_base is not None and frame_new is not None
342
343 # Decompress base
344 from muse.core.compression import decompress_frame
345 base_enc = frame_base.get("enc", "raw")
346 decoded_base = decompress_frame(bytes(frame_base["content"]), base_enc)
347
348 # Apply delta
349 enc = frame_new.get("enc", "")
350 assert "delta" in enc
351 delta_bytes = bytes(frame_new["content"])
352 # Delta stream is zlib-compressed inside (compute_delta always zlib-wraps)
353 reconstructed = apply_delta(decoded_base, delta_bytes)
354
355 assert reconstructed == new_content
356 assert blob_id(reconstructed) == new_oid
357
358
359 # ---------------------------------------------------------------------------
360 # C6 — Unprofitable delta → falls back to full frame
361 # ---------------------------------------------------------------------------
362
363 @pytest.mark.asyncio
364 async def test_c6_unprofitable_delta_uses_full_frame(db_session: AsyncSession) -> None:
365 import os
366
367 repo = await create_repo(db_session, owner="test-comp-c6")
368 repo_id = str(repo.repo_id)
369
370 # Two completely random (incompressible) objects at the same path
371 rng = __import__("random").Random(42)
372 base_content = bytes(rng.randint(0, 255) for _ in range(4096))
373 new_content = bytes(rng.randint(0, 255) for _ in range(4096))
374
375 base_oid = await _make_object(db_session, repo_id, base_content, path="bin/rand.dat", owner=repo.owner, slug=repo.slug)
376 new_oid = await _make_object(db_session, repo_id, new_content, path="bin/rand.dat", owner=repo.owner, slug=repo.slug)
377
378 snap1 = await _make_snapshot(db_session, repo_id, {"bin/rand.dat": base_oid})
379 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c6-1")
380 snap2 = await _make_snapshot(db_session, repo_id, {"bin/rand.dat": new_oid})
381 c2 = await _make_commit(
382 db_session, repo_id,
383 parent_ids=[c1.commit_id],
384 snapshot_id=snap2.snapshot_id,
385 seed="c6-2",
386 )
387
388 from musehub.services.musehub_wire import wire_fetch_stream
389 frames = await _collect_frames(
390 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c2.commit_id], have=[]))
391 )
392 frame_new = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == new_oid), None)
393 assert frame_new is not None
394 assert "delta" not in frame_new.get("enc", ""), (
395 "Expected full frame for incompressible random content, got delta"
396 )
397
398
399 # ---------------------------------------------------------------------------
400 # C7 — Full clone with delta mix delivers correct object_id set
401 # ---------------------------------------------------------------------------
402
403 @pytest.mark.asyncio
404 async def test_c7_full_clone_correct_object_ids(db_session: AsyncSession) -> None:
405 repo = await create_repo(db_session, owner="test-comp-c7")
406 repo_id = str(repo.repo_id)
407
408 # Three commits: add file, modify it, add another file
409 v1 = b"version one\n" * 50
410 v2 = b"version two\n" * 50
411 other = b"other file\n" * 30
412
413 oid_v1 = await _make_object(db_session, repo_id, v1, path="src/main.py", owner=repo.owner, slug=repo.slug)
414 oid_v2 = await _make_object(db_session, repo_id, v2, path="src/main.py", owner=repo.owner, slug=repo.slug)
415 oid_other = await _make_object(db_session, repo_id, other, path="src/util.py", owner=repo.owner, slug=repo.slug)
416
417 snap1 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v1})
418 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c7-1")
419
420 snap2 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v2})
421 c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], snapshot_id=snap2.snapshot_id, seed="c7-2")
422
423 snap3 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v2, "src/util.py": oid_other})
424 c3 = await _make_commit(db_session, repo_id, parent_ids=[c2.commit_id], snapshot_id=snap3.snapshot_id, seed="c7-3")
425
426 from musehub.services.musehub_wire import wire_fetch_stream
427 frames = await _collect_frames(
428 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c3.commit_id], have=[]))
429 )
430 received_oids = {f["id"] for f in frames if f.get("t") == SFRAME_OBJECT}
431 expected_oids = {oid_v1, oid_v2, oid_other}
432 assert received_oids == expected_oids, (
433 f"Missing: {expected_oids - received_oids}, Extra: {received_oids - expected_oids}"
434 )
435
436
437 # ---------------------------------------------------------------------------
438 # C8 — Cross-batch delta: prior-batch object serves as base for next batch
439 # ---------------------------------------------------------------------------
440
441 @pytest.mark.asyncio
442 async def test_c8_cross_batch_delta(db_session: AsyncSession) -> None:
443 """An object in batch N+1 that shares a path with an object in batch N
444 should be delta-encoded against the batch-N object."""
445 from musehub.services.musehub_wire import _COMMIT_BATCH
446
447 repo = await create_repo(db_session, owner="test-comp-c8")
448 repo_id = str(repo.repo_id)
449
450 base_content = _realistic_source("r1")
451 new_content = _realistic_source("r2")
452
453 base_oid = await _make_object(db_session, repo_id, base_content, path="worker.py", owner=repo.owner, slug=repo.slug)
454 new_oid = await _make_object(db_session, repo_id, new_content, path="worker.py", owner=repo.owner, slug=repo.slug)
455
456 # Put base in batch 1 and new version in batch 2 (requires > _COMMIT_BATCH commits between them)
457 snap_base = await _make_snapshot(db_session, repo_id, {"worker.py": base_oid})
458 c_base = await _make_commit(db_session, repo_id, snapshot_id=snap_base.snapshot_id, seed="c8-base")
459
460 # Pad with _COMMIT_BATCH commits so base ends up in batch 1 (reuse same snap)
461 prev = c_base.commit_id
462 for i in range(_COMMIT_BATCH):
463 c = await _make_commit(
464 db_session, repo_id,
465 parent_ids=[prev],
466 snapshot_id=snap_base.snapshot_id,
467 seed=f"c8-fill-{i}",
468 )
469 prev = c.commit_id
470
471 snap_new = await _make_snapshot(db_session, repo_id, {"worker.py": new_oid})
472 c_new = await _make_commit(
473 db_session, repo_id,
474 parent_ids=[prev],
475 snapshot_id=snap_new.snapshot_id,
476 seed="c8-new",
477 )
478
479 from musehub.services.musehub_wire import wire_fetch_stream
480 frames = await _collect_frames(
481 wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c_new.commit_id], have=[]))
482 )
483 frame_new = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == new_oid), None)
484 assert frame_new is not None
485 assert "delta" in frame_new.get("enc", ""), (
486 f"Expected delta frame for cross-batch successive version, got enc={frame_new.get('enc')!r}"
487 )
488 assert frame_new.get("base_id") == base_oid
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago