diff --git a/tests/romfs_packer_test.py b/tests/romfs_packer_test.py new file mode 100644 index 00000000..88ce654d --- /dev/null +++ b/tests/romfs_packer_test.py @@ -0,0 +1,72 @@ +"""RomFS read-ahead and packing with synthetic files. Run: python tests/romfs_packer_test.py""" +import hashlib +import json +from pathlib import Path +import sys +import tempfile +import threading +import unittest +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools')) +import make_game + + +class RomfsPackerTest(unittest.TestCase): + def test_read_ahead_limits_total_bytes_and_file_count(self): + limit = make_game.ROMFS_READ_AHEAD + sizes = [limit - 1, 1, limit + 1, 1, limit] + [0] * 512 + files = [(str(i), i, size) for i, size in enumerate(sizes)] + groups = list(make_game.romfs_read_groups(files)) + self.assertEqual([entry for group in groups for entry in group], files) + for group in groups: + self.assertLessEqual(sum(size for _, _, size in group if size <= limit), limit) + self.assertLessEqual(len(group), 256) + + def test_changed_file_cannot_replace_the_image(self): + source = mock.Mock(romfs=[('changed.bin', 0)]) + source.romfs_size.return_value = 4 + with tempfile.TemporaryDirectory() as tmp: + dst = Path(tmp) / 'romfs' + dst.mkdir() + (dst / 'romfs.bin').write_bytes(b'previous image') + (dst / 'index.json').write_bytes(b'previous index') + for limit, chunks in [(8, [b'123']), (2, [b'123', b'456'])]: + with self.subTest(limit=limit), mock.patch.object(make_game, 'ROMFS_READ_AHEAD', limit): + source.romfs_chunks.return_value = chunks + with self.assertRaises(SystemExit): + make_game.step_romfs(source, dst.parent, 1) + self.assertEqual((dst / 'romfs.bin').read_bytes(), b'previous image') + self.assertEqual((dst / 'index.json').read_bytes(), b'previous index') + + def test_streamed_and_buffered_files_keep_the_same_layout(self): + # Lower the read-ahead limit so these tiny files span several worker groups. + large_file_readers = [] + + class Source: + romfs = [('e', b'123456789'), ('a', b'12345'), ('b', b'678'), ('c', b''), ('d', b'!')] + + def romfs_size(self, ref): + return len(ref) + + def romfs_chunks(self, ref): + if len(ref) > 8: + large_file_readers.append(threading.current_thread()) + yield ref[:3] + yield ref[3:] + + with tempfile.TemporaryDirectory() as tmp, mock.patch.object(make_game, 'ROMFS_READ_AHEAD', 8): + out = Path(tmp) + make_game.step_romfs(Source(), out, 2) + image = (out / 'romfs' / 'romfs.bin').read_bytes() + index = json.loads((out / 'romfs' / 'index.json').read_text()) + self.assertEqual(large_file_readers, [threading.main_thread()]) + self.assertEqual(image, b'12345' + bytes(11) + b'678' + bytes(13) + b'!' + bytes(15) + b'123456789') + self.assertEqual(index['romfs_sha256'], hashlib.sha256(image).hexdigest()) + self.assertEqual(index['romfs_size'], len(image)) + self.assertEqual([(e['path'], e['offset'], e['size']) for e in index['files']], + [('a', 0, 5), ('b', 16, 3), ('c', 32, 0), ('d', 32, 1), ('e', 48, 9)]) + + +if __name__ == '__main__': + unittest.main() diff --git a/tools/make_game.py b/tools/make_game.py index 3319ce7e..01e410e3 100644 --- a/tools/make_game.py +++ b/tools/make_game.py @@ -365,11 +365,28 @@ def exefs_digest(out): # ---------------------------------------------------------------- romfs +ROMFS_READ_AHEAD = 64 << 20 # total file bytes requested from workers in one group + + +def romfs_read_groups(files): + group, buffered = [], 0 + for entry in files: + size = entry[2] + ahead = size if size <= ROMFS_READ_AHEAD else 0 # the writer streams larger files + if buffered + ahead > ROMFS_READ_AHEAD or len(group) == 256: + yield group + group, buffered = [], 0 + group.append(entry) + buffered += ahead + if group: + yield group + + def step_romfs(src, out, jobs): dst = out / 'romfs' dst.mkdir(parents=True, exist_ok=True) - files = sorted(src.romfs, key=lambda e: e[0].encode()) - total = sum(src.romfs_size(ref) for _, ref in files) + files = sorted(((rel, ref, src.romfs_size(ref)) for rel, ref in src.romfs), key=lambda e: e[0].encode()) + total = sum(size for _, _, size in files) free = shutil.disk_usage(dst).free if free < total + (1 << 30): fail(f'{dst} needs {total / 1e9:.1f} GB free for the RomFS image, it has {free / 1e9:.1f} GB') @@ -379,28 +396,28 @@ def step_romfs(src, out, jobs): emit('progress', step='romfs', done=0, total=total, unit='bytes') def load(e): # small files are read whole on a worker; big ones are streamed by the writer - rel, ref = e - return None if src.romfs_size(ref) > (64 << 20) else b''.join(src.romfs_chunks(ref)) + _, ref, size = e + return None if size > ROMFS_READ_AHEAD else b''.join(src.romfs_chunks(ref)) with open(part, 'wb', buffering=16 << 20) as img, concurrent.futures.ThreadPoolExecutor(jobs) as ex: - batch = 256 - for start in range(0, len(files), batch): - group = files[start:start + batch] - for (rel, ref), data in zip(group, ex.map(load, group)): + for batch, group in enumerate(romfs_read_groups(files)): + for (rel, ref, size), data in zip(group, ex.map(load, group)): pad = -pos % ALIGN if pad: img.write(b'\0' * pad) h.update(b'\0' * pad) pos += pad - size = src.romfs_size(ref) - entries.append({'path': rel, 'offset': pos, 'size': size}) + offset = pos for chunk in ((data,) if data is not None else src.romfs_chunks(ref)): img.write(chunk) h.update(chunk) - pos += size + pos += len(chunk) + if pos - offset != size: + fail(f'{rel} changed size while packing the RomFS: expected {size} bytes, read {pos - offset}') + entries.append({'path': rel, 'offset': offset, 'size': size}) done += size emit('progress', step='romfs', done=done, total=total, unit='bytes') - if start // batch % 40 == 0: + if batch % 40 == 0: say(f' RomFS {done / 1e9:5.1f} of {total / 1e9:.1f} GB, {time.time() - t0:.0f} s') part.replace(dst / 'romfs.bin') index = {'status': 'pass', 'romfs_size': pos, 'romfs_sha256': h.hexdigest(), 'files': entries}