Compress Tachyon string and frame tables

The file-wide string and frame tables can dominate captures with many
unique frames, but they are currently always written as raw bytes.

Build each table during finalization and compress it independently. Keep
the raw representation when zstd does not make the table smaller. Prefix
the stored data with its decoded size and encoding so the reader can
validate and decode it without changing the rest of the file layout.

Raw tables continue to point directly into the mapped file. Compressed
tables reuse a temporary decode buffer before their entries are parsed.

Author: ivonastojanovic

InternalDocs/profiling_binary_format.md

@@ -333,11 +333,30 @@ in the array). This ordering works well with delta compression: function
 calls typically change frames at the end of the array, while shared frames
 remain at the start.
 
+## Table Header
+
+The string and frame tables both start with a header that describes how
+their entries are stored:
+
+```
++---------------------+-----------+----------------------------------------+
+| decoded_size        | varint    | Size of the decoded table              |
+| encoding            | 1 byte    | 0=raw, 1=zstd                          |
++---------------------+-----------+----------------------------------------+
+```
+
+The stored table follows the header and extends to the start of the next
+section: the frame table for the string table, and the profile statistics
+or footer for the frame table. A raw table is stored as is, so its stored
+size equals `decoded_size`. A zstd table is a single zstd frame whose
+content size is `decoded_size`, at most 1 GiB; larger tables are stored raw.
+
 ## String Table
 
 The string table stores deduplicated UTF-8 strings (filenames and function
-names). It begins at `string_table_offset` and contains entries in order of
-their assignment during writing:
+names). It begins at `string_table_offset` with a table header (see
+[Table Header](#table-header)), followed by the entries in order of their
+assignment during writing:
 
 ```
 +----------------+
@@ -356,7 +375,8 @@ byte.
 ## Frame Table
 
 The frame table stores deduplicated frame entries with full source position
-information and bytecode opcode:
+information and bytecode opcode. It begins at `frame_table_offset` with a
+table header, followed by the entries:
 
 ```
 +----------------------------+
@@ -540,10 +560,10 @@ smallest:
   stored-size varint is smaller than the raw column;
 - **raw** otherwise.
 
-The string table, frame table, and footer remain uncompressed so readers can
-access metadata without decompressing the entire file. A tool that only needs
-to report "this file contains 50,000 samples of 3 threads" can read the header
-and footer without touching the compressed sample data.
+The string and frame tables are each stored as raw or zstd, whichever is
+smaller, and are decoded independently of the samples. The header and footer
+remain uncompressed, so a tool that only needs to report "this file contains
+50,000 samples of 3 threads" can read them without decoding anything else.
 
 Zstd provides an excellent balance of compression ratio and speed. Profiling
 data compresses very well (often 5-10x) due to repetitive patterns: the same
@@ -572,6 +592,9 @@ independent, at the cost of restarting compression history at every block.
 8. Write the footer with final counts
 9. Seek to offset 0 and write the header with actual values
 
+Each table is built in memory first, so its decoded size is known, and then
+written with its header as raw bytes or zstd.
+
 The writer maintains two dictionaries: one mapping strings to indices, one
 mapping (filename_idx, funcname_idx, lineno) tuples to frame indices. These
 enable O(1) lookup during interning. To find `REUSE` candidates, it also keeps
@@ -590,9 +613,10 @@ column, and each column is encoded as described in
 3. Seek to end − 32 and read the footer (byte-swapping counts if needed)
 4. Read measured profile statistics when the optional extension is present
 5. Allocate string array of `string_count` elements
-6. Parse the string table, populating the array
+6. Decode the string table (raw tables are read in place) and parse it,
+   populating the array
 7. Allocate the frame array
-8. Parse the frame table, populating the array
+8. Decode and parse the frame table, populating the array
 9. For each block, decode its columns into a reusable buffer and replay its
    samples in order, resolving indices to strings/frames (byte-swapping the
    stream table if needed)

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -127,18 +127,20 @@ def export(self, filename):
         pass
 
 
+def varint(data, offset):
+    """Decode an unsigned varint, returning (value, next offset)."""
+    value = shift = 0
+    while True:
+        byte = data[offset]
+        offset += 1
+        value |= (byte & 0x7F) << shift
+        shift += 7
+        if byte < 0x80:
+            return value, offset
+
+
 def read_block_columns(filename):
     """Return the (decoded size, encoding) of each column of every block."""
-    def varint(data, offset):
-        value = shift = 0
-        while True:
-            byte = data[offset]
-            offset += 1
-            value |= (byte & 0x7F) << shift
-            shift += 7
-            if byte < 0x80:
-                return value, offset
-
     with open(filename, "rb") as raw:
         data = raw.read()
     sample_end = struct.unpack_from("=Q", data, 40)[0]
@@ -1443,13 +1445,14 @@ def test_replay_rejects_sample_count_mismatch(self):
     def test_replay_rejects_trailing_partial_sample_header(self):
         """Replay rejects partial sample bytes instead of silently stopping."""
         filename = self.create_binary_file([], compression="none")
-        sample_data_end = self.FILE_HEADER_PLACEHOLDER_SIZE + 1
-
-        with open(filename, "r+b") as raw:
-            raw.seek(self.HDR_OFF_STR_TABLE)
-            raw.write(struct.pack("=Q", sample_data_end))
-            raw.seek(self.HDR_OFF_FRAME_TABLE)
-            raw.write(struct.pack("=Q", sample_data_end))
+        data = bytearray(pathlib.Path(filename).read_bytes())
+        data[self.FILE_HEADER_PLACEHOLDER_SIZE:
+             self.FILE_HEADER_PLACEHOLDER_SIZE] = b"\0"
+        for offset in (self.HDR_OFF_STR_TABLE, self.HDR_OFF_FRAME_TABLE):
+            old_offset = struct.unpack_from("=Q", data, offset)[0]
+            struct.pack_into("=Q", data, offset, old_offset + 1)
+        struct.pack_into("=Q", data, len(data) - 24, len(data))
+        pathlib.Path(filename).write_bytes(data)
 
         with BinaryReader(filename) as reader:
             with self.assertRaises(ValueError) as cm:
@@ -1461,26 +1464,47 @@ def test_replay_rejects_trailing_partial_sample_header(self):
     MIN_STRING_ENTRY_SIZE = 1
     MIN_FRAME_ENTRY_SIZE = 7
 
-    def _read_offset(self, filename, hdr_off):
-        with open(filename, "rb") as raw:
-            raw.seek(hdr_off)
-            return struct.unpack("=Q", raw.read(8))[0]
+    def _table_bounds(self, data, hdr_off):
+        start = struct.unpack_from("=Q", data, hdr_off)[0]
+        if hdr_off == self.HDR_OFF_STR_TABLE:
+            end = struct.unpack_from("=Q", data, self.HDR_OFF_FRAME_TABLE)[0]
+        else:
+            end = len(data) - self.FILE_FOOTER_SIZE
+        return start, end
+
+    def _read_table_section(self, filename, hdr_off):
+        """Return the decoded size, encoding and stored bytes of a table."""
+        data = pathlib.Path(filename).read_bytes()
+        start, end = self._table_bounds(data, hdr_off)
+        decoded_size, offset = varint(data, start)
+        return decoded_size, data[offset], data[offset + 1:end]
+
+    def _replace_table_section(self, filename, hdr_off, section):
+        data = bytearray(pathlib.Path(filename).read_bytes())
+        start, end = self._table_bounds(data, hdr_off)
+        data[start:end] = section
+        if hdr_off == self.HDR_OFF_STR_TABLE:
+            frame_offset = start + len(section)
+            struct.pack_into("=Q", data, self.HDR_OFF_FRAME_TABLE,
+                             frame_offset)
+        struct.pack_into("=Q", data, len(data) - 24, len(data))
+        pathlib.Path(filename).write_bytes(data)
 
     def _patch_footer_count(self, filename, ftr_off, value):
         size = os.path.getsize(filename)
         with open(filename, "r+b") as raw:
             raw.seek(size - self.FILE_FOOTER_SIZE + ftr_off)
             raw.write(struct.pack("=I", value))
 
-    def test_open_rejects_string_count_larger_than_file(self):
-        """Open rejects a footer string count larger than the file."""
+    def test_open_rejects_string_count_larger_than_table(self):
+        """Open rejects a footer string count larger than its table."""
         samples = [[make_interpreter(0, [
             make_thread(1, [make_frame("s.py", 10, "s")])
         ])]]
         filename = self.create_binary_file(samples, compression="none")
-        size = os.path.getsize(filename)
-        str_off = self._read_offset(filename, self.HDR_OFF_STR_TABLE)
-        max_strings = (size - str_off) // self.MIN_STRING_ENTRY_SIZE
+        decoded_size, _, _ = self._read_table_section(
+            filename, self.HDR_OFF_STR_TABLE)
+        max_strings = decoded_size // self.MIN_STRING_ENTRY_SIZE
         self._patch_footer_count(filename, self.FTR_OFF_STRINGS, 0xFFFFFFFF)
 
         with self.assertRaises(ValueError) as cm:
@@ -1492,15 +1516,15 @@ def test_open_rejects_string_count_larger_than_file(self):
             f"possible {max_strings}",
         )
 
-    def test_open_rejects_frame_count_larger_than_file(self):
-        """Open rejects a footer frame count larger than the file."""
+    def test_open_rejects_frame_count_larger_than_table(self):
+        """Open rejects a footer frame count larger than its table."""
         samples = [[make_interpreter(0, [
             make_thread(1, [make_frame("f.py", 10, "f")])
         ])]]
         filename = self.create_binary_file(samples, compression="none")
-        size = os.path.getsize(filename)
-        frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
-        max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
+        decoded_size, _, _ = self._read_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE)
+        max_frames = decoded_size // self.MIN_FRAME_ENTRY_SIZE
         self._patch_footer_count(filename, self.FTR_OFF_FRAMES, 0xFFFFFFFF)
 
         # On a 32-bit build this count overflows the allocation size, so the
@@ -1516,14 +1540,14 @@ def test_open_rejects_frame_count_larger_than_file(self):
             )
 
     def test_open_accepts_frame_count_at_capacity_boundary(self):
-        """A frame count at the file-size cap opens; one more is rejected."""
+        """A frame count at the table-size cap opens; one more is rejected."""
         samples = [[make_interpreter(0, [
             make_thread(1, [make_frame("f.py", 10, "f")])
         ])]]
         filename = self.create_binary_file(samples, compression="none")
-        size = os.path.getsize(filename)
-        frame_off = self._read_offset(filename, self.HDR_OFF_FRAME_TABLE)
-        max_frames = (size - frame_off) // self.MIN_FRAME_ENTRY_SIZE
+        decoded_size, _, _ = self._read_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE)
+        max_frames = decoded_size // self.MIN_FRAME_ENTRY_SIZE
 
         self._patch_footer_count(filename, self.FTR_OFF_FRAMES, max_frames)
         with BinaryReader(filename):
@@ -1539,6 +1563,120 @@ def test_open_accepts_frame_count_at_capacity_boundary(self):
             f"possible {max_frames}",
         )
 
+    def test_raw_table_sections_record_decoded_size(self):
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("raw.py", 10, "raw")])
+        ])]]
+        filename = self.create_binary_file(samples, compression="none")
+        for hdr_off in (self.HDR_OFF_STR_TABLE, self.HDR_OFF_FRAME_TABLE):
+            with self.subTest(hdr_off=hdr_off):
+                decoded_size, encoding, payload = self._read_table_section(
+                    filename, hdr_off)
+                self.assertEqual(encoding, 0)
+                self.assertEqual(decoded_size, len(payload))
+
+    @unittest.skipUnless(ZSTD_AVAILABLE, "zstd compression not available")
+    def test_repetitive_tables_use_zstd(self):
+        frames = [
+            make_frame(
+                f"package/with/repeated/prefix/module_{index:04}.py",
+                index + 1,
+                f"function_with_repeated_prefix_{index:04}",
+            )
+            for index in range(256)
+        ]
+        samples = [[make_interpreter(0, [make_thread(1, frames)])]]
+        filename = self.create_binary_file(samples, compression="zstd")
+        for hdr_off in (self.HDR_OFF_STR_TABLE, self.HDR_OFF_FRAME_TABLE):
+            with self.subTest(hdr_off=hdr_off):
+                decoded_size, encoding, payload = self._read_table_section(
+                    filename, hdr_off)
+                self.assertEqual(encoding, 1)
+                self.assertLess(len(payload), decoded_size)
+
+    def test_open_rejects_invalid_table_sections(self):
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("table.py", 10, "table")])
+        ])]]
+        cases = [
+            ("Malformed string table header", self.HDR_OFF_STR_TABLE,
+             lambda size, payload: b""),
+            ("Unknown frame table encoding: 255", self.HDR_OFF_FRAME_TABLE,
+             lambda size, payload: self._varint(size) + b"\xff" + payload),
+            ("Raw string table size does not match its header",
+             self.HDR_OFF_STR_TABLE,
+             lambda size, payload: self._varint(size + 1) + b"\0" + payload),
+            ("Frame table has trailing data", self.HDR_OFF_FRAME_TABLE,
+             lambda size, payload:
+                 self._varint(size + 1) + b"\0" + payload + b"\0"),
+            ("zstd frame table in an uncompressed file",
+             self.HDR_OFF_FRAME_TABLE,
+             lambda size, payload: self._varint(size) + b"\1" + payload),
+        ]
+        for message, hdr_off, make_section in cases:
+            with self.subTest(message=message):
+                filename = self.create_binary_file(
+                    samples, compression="none")
+                decoded_size, encoding, payload = self._read_table_section(
+                    filename, hdr_off)
+                self.assertEqual(encoding, 0)
+                self._replace_table_section(
+                    filename, hdr_off, make_section(decoded_size, payload))
+                with self.assertRaisesRegex(ValueError, message):
+                    with BinaryReader(filename):
+                        pass
+
+    @unittest.skipUnless(ZSTD_AVAILABLE, "zstd compression not available")
+    def test_open_rejects_invalid_zstd_table(self):
+        frames = [
+            make_frame(f"repeated/module_{index:04}.py", index + 1,
+                       f"repeated_function_{index:04}")
+            for index in range(256)
+        ]
+        samples = [[make_interpreter(0, [make_thread(1, frames)])]]
+
+        filename = self.create_binary_file(samples, compression="zstd")
+        decoded_size, encoding, payload = self._read_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE)
+        self.assertEqual(encoding, 1)
+        self._replace_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE,
+            self._varint(decoded_size) + bytes([encoding]) + payload[:-1])
+        with self.assertRaisesRegex(ValueError,
+                                   "Corrupted zstd frame header"):
+            with BinaryReader(filename):
+                pass
+
+        filename = self.create_binary_file(samples, compression="zstd")
+        decoded_size, encoding, payload = self._read_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE)
+        self._replace_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE,
+            self._varint(decoded_size + 1) + bytes([encoding]) + payload)
+        with self.assertRaisesRegex(
+                ValueError,
+                "zstd frame size does not match the decoded size"):
+            with BinaryReader(filename):
+                pass
+
+    @unittest.skipUnless(ZSTD_AVAILABLE, "zstd compression not available")
+    def test_open_rejects_oversized_zstd_table(self):
+        # A table claiming 2**40 bytes must be rejected before anything is
+        # allocated for it.
+        samples = [[make_interpreter(0, [
+            make_thread(1, [make_frame("table.py", 10, "table")])
+        ])]]
+        filename = self.create_binary_file(samples, compression="zstd")
+        _, _, payload = self._read_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE)
+        self._replace_table_section(
+            filename, self.HDR_OFF_FRAME_TABLE,
+            self._varint(1 << 40) + b"\1" + payload)
+        with self.assertRaisesRegex(
+                ValueError, "frame table exceeds the maximum decoded size"):
+            with BinaryReader(filename):
+                pass
+
     def test_open_rejects_older_format_version(self):
         filename = self.create_binary_file([], compression="none")
         with BinaryReader(filename) as reader:

Modules/_remote_debugging/binary_io.h

@@ -73,6 +73,9 @@ static_assert(FILE_HEADER_SIZE <= FILE_HEADER_PLACEHOLDER_SIZE,
 #define MAX_BLOCK_SIZE          (1024 * 1024)
 #define MAX_BLOCK_STREAMS       4096
 
+/* Maximum decoded size of a compressed string or frame table (1GB) */
+#define MAX_DECOMPRESS_SIZE     (1ULL << 30)
+
 /* Block columns */
 enum {
     COLUMN_STREAM_TABLE,