Bound the RomFS packer’s parallel reads by file size.

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()

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}