fixup! Store Tachyon samples in bounded columnar blocks

The stack reuse cache currently keeps one stack at each hash position,
so a collision immediately evicts the previous stack. Alternating
between two colliding stacks therefore prevents either from being
reused.

Let each cache bucket remember two stacks while keeping the cache the
same overall size. Lookups check both stacks, and inserting a third
stack replaces the older one.

Author: ivonastojanovic

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -1019,6 +1019,38 @@ def test_writer_reports_stack_reuse(self):
         self.assertEqual(stats["total_samples"], 2)
         writer.close()
 
+    def test_writer_reuses_colliding_stacks(self):
+        """Stacks sharing a cache bucket are both stored as STACK_REUSE."""
+        buckets = {}
+        for frame_id in range(2, 4099):
+            bucket = hash(struct.pack("=3I", 0, 1, frame_id)) % 4096
+            if bucket in buckets:
+                first_id = buckets[bucket]
+                second_id = frame_id
+                break
+            buckets[bucket] = frame_id
+
+        frames = [make_frame("a.py", line, "f")
+                  for line in range(second_id + 1)]
+        first = [frames[0], frames[1], frames[first_id]]
+        second = [frames[0], frames[1], frames[second_id]]
+
+        filename = self.temp_filename()
+        writer = _remote_debugging.BinaryWriter(filename, 1000, 0)
+        timestamp = 1000
+        # Assign frame IDs without exceeding MAX_STACK_DEPTH.
+        for start in range(0, len(frames), 256):
+            stack = frames[start:start + 256]
+            writer.write_sample([make_interpreter(0, [make_thread(2, stack)])],
+                                timestamp)
+            timestamp += 1000
+        for stack in (first, second, first, second):
+            writer.write_sample([make_interpreter(0, [make_thread(1, stack)])],
+                                timestamp)
+            timestamp += 1000
+        self.assertEqual(writer.get_stats()["reuse_records"], 2)
+        writer.close()
+
     def test_writer_buffered_bytes_match_flushed_bytes(self):
         """bytes_written before finalize() matches the flushed size."""
         filename = self.temp_filename()

Modules/_remote_debugging/binary_io.h

@@ -171,6 +171,9 @@ static_assert(PROFILE_STATS_SIZE == 56,
 
 /* Writer's stack cache size per block */
 #define STACK_CACHE_SIZE        4096
+#define STACK_CACHE_BUCKET_SIZE 2
+#define STACK_CACHE_BUCKET_COUNT \
+    (STACK_CACHE_SIZE / STACK_CACHE_BUCKET_SIZE)
 
 /* Initial capacities for dynamic arrays - sized to reduce reallocations */
 #define INITIAL_STRING_CAPACITY 4096
@@ -352,7 +355,7 @@ typedef struct {
     ByteBuffer transformed;
     ByteBuffer compressed;
     ByteBuffer output;
-    StackCacheEntry cache[STACK_CACHE_SIZE];
+    StackCacheEntry cache[STACK_CACHE_BUCKET_COUNT][STACK_CACHE_BUCKET_SIZE];
     uint64_t base_timestamp;
     uint32_t sample_count;
     uint32_t stream_count;

Modules/_remote_debugging/binary_io_writer.c

@@ -810,6 +810,26 @@ cached_stack_matches(const StackCacheEntry *cached,
     return 1;
 }
 
+static StackCacheEntry *
+stack_cache_lookup(StackCacheEntry *bucket, const ByteBuffer *frames,
+                   Py_uhash_t hash, const uint32_t *stack, size_t depth)
+{
+    for (size_t i = 0; i < STACK_CACHE_BUCKET_SIZE; i++) {
+        if (cached_stack_matches(&bucket[i], frames, hash, stack, depth)) {
+            return &bucket[i];
+        }
+    }
+    return NULL;
+}
+
+static void
+stack_cache_insert(StackCacheEntry *bucket, StackCacheEntry entry)
+{
+    memmove(&bucket[1], &bucket[0],
+            (STACK_CACHE_BUCKET_SIZE - 1) * sizeof(*bucket));
+    bucket[0] = entry;
+}
+
 static int
 write_stack_repeat(BinaryWriter *writer, ThreadEntry *entry, size_t depth)
 {
@@ -853,7 +873,7 @@ write_stack_reuse(BinaryWriter *writer,
 }
 
 static int
-write_stack_replace_suffix(BinaryWriter *writer, StackCacheEntry *cached,
+write_stack_replace_suffix(BinaryWriter *writer, StackCacheEntry *bucket,
                            Py_uhash_t hash, const uint32_t *stack,
                            size_t depth, size_t shared)
 {
@@ -888,9 +908,9 @@ write_stack_replace_suffix(BinaryWriter *writer, StackCacheEntry *cached,
     }
     /* Complete stacks are cached so later samples can REUSE them */
     if (shared == 0) {
-        if (cached != NULL) {
-            *cached = (StackCacheEntry){
-                hash, (uint32_t)stack_offset, (uint32_t)depth};
+        if (bucket != NULL) {
+            stack_cache_insert(bucket, (StackCacheEntry){
+                hash, (uint32_t)stack_offset, (uint32_t)depth});
         }
         writer->stats.full_records++;
     }
@@ -906,10 +926,11 @@ static int
 choose_stack_encoding(SampleBlock *block, const ThreadEntry *entry,
                       int new_stream, const uint32_t *stack, size_t depth,
                       size_t *shared, Py_uhash_t *hash,
-                      StackCacheEntry **cached)
+                      StackCacheEntry **bucket, StackCacheEntry **cached)
 {
     *shared = common_outer_frames(entry, stack, depth);
     *hash = 0;
+    *bucket = NULL;
     *cached = NULL;
     if (!new_stream && *shared == depth &&
         *shared == entry->prev_stack_depth) {
@@ -919,9 +940,10 @@ choose_stack_encoding(SampleBlock *block, const ThreadEntry *entry,
         return STACK_REPLACE_SUFFIX;
     }
     *hash = stack_hash(stack, depth);
-    *cached = &block->cache[*hash % STACK_CACHE_SIZE];
-    if (cached_stack_matches(*cached, &block->columns[COLUMN_FRAME_IDS],
-                             *hash, stack, depth)) {
+    *bucket = block->cache[*hash % STACK_CACHE_BUCKET_COUNT];
+    *cached = stack_cache_lookup(*bucket, &block->columns[COLUMN_FRAME_IDS],
+                                 *hash, stack, depth);
+    if (*cached != NULL) {
         return STACK_REUSE;
     }
     return STACK_REPLACE_SUFFIX;
@@ -933,10 +955,11 @@ write_stack(BinaryWriter *writer, ThreadEntry *entry, int new_stream,
 {
     size_t shared;
     Py_uhash_t hash;
+    StackCacheEntry *bucket;
     StackCacheEntry *cached;
     int encoding = choose_stack_encoding(writer->block, entry, new_stream,
                                          stack, depth, &shared, &hash,
-                                         &cached);
+                                         &bucket, &cached);
 
     int result;
     switch (encoding) {
@@ -946,7 +969,7 @@ write_stack(BinaryWriter *writer, ThreadEntry *entry, int new_stream,
         result = write_stack_reuse(writer, cached, depth);
         break;
     case STACK_REPLACE_SUFFIX:
-        result = write_stack_replace_suffix(writer, cached, hash, stack,
+        result = write_stack_replace_suffix(writer, bucket, hash, stack,
                                             depth, shared);
         break;
     default: