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