pythongh-115491: Keep some fields valid across allocations (free-threading) (python#115573)

This avoids filling the memory occupied by ob_tid, ob_ref_local, and
ob_ref_shared with debug bytes (e.g., 0xDD) in mimalloc in the
free-threaded build.

Author: colesbury

Include/internal/mimalloc/mimalloc/types.h

@@ -312,6 +312,7 @@ typedef struct mi_page_s {
   uint8_t               is_committed : 1;  // `true` if the page virtual memory is committed
   uint8_t               is_zero_init : 1;  // `true` if the page was initially zero initialized
   uint8_t               tag : 4;           // tag from the owning heap
+  uint8_t               debug_offset;      // number of bytes to preserve when filling freed or uninitialized memory
 
   // layout like this to optimize access in `mi_malloc` and `mi_free`
   uint16_t              capacity;          // number of blocks committed, must be the first field, see `segment.c:page_clear`
@@ -553,6 +554,7 @@ struct mi_heap_s {
   mi_heap_t*            next;                                // list of heaps per thread
   bool                  no_reclaim;                          // `true` if this heap should not reclaim abandoned pages
   uint8_t               tag;                                 // custom identifier for this heap
+  uint8_t               debug_offset;                        // number of bytes to preserve when filling freed or uninitialized memory
 };
 
 

Objects/mimalloc/alloc.c

@@ -26,6 +26,15 @@ terms of the MIT license. A copy of the license can be found in the file
 // Allocation
 // ------------------------------------------------------
 
+#if (MI_DEBUG>0)
+static void mi_debug_fill(mi_page_t* page, mi_block_t* block, int c, size_t size) {
+  size_t offset = (size_t)page->debug_offset;
+  if (offset < size) {
+    memset((char*)block + offset, c, size - offset);
+  }
+}
+#endif
+
 // Fast allocation in a page: just pop from the free list.
 // Fall back to generic allocation only if the list is empty.
 extern inline void* _mi_page_malloc(mi_heap_t* heap, mi_page_t* page, size_t size, bool zero) mi_attr_noexcept {
@@ -65,7 +74,7 @@ extern inline void* _mi_page_malloc(mi_heap_t* heap, mi_page_t* page, size_t siz
 
 #if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN
   if (!zero && !mi_page_is_huge(page)) {
-    memset(block, MI_DEBUG_UNINIT, mi_page_usable_block_size(page));
+    mi_debug_fill(page, block, MI_DEBUG_UNINIT, mi_page_usable_block_size(page));
   }
 #elif (MI_SECURE!=0)
   if (!zero) { block->next = 0; } // don't leak internal data
@@ -426,7 +435,7 @@ static mi_decl_noinline void _mi_free_block_mt(mi_page_t* page, mi_block_t* bloc
 
   #if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN        // note: when tracking, cannot use mi_usable_size with multi-threading
   if (segment->kind != MI_SEGMENT_HUGE) {                  // not for huge segments as we just reset the content
-    memset(block, MI_DEBUG_FREED, mi_usable_size(block));
+    mi_debug_fill(page, block, MI_DEBUG_FREED, mi_usable_size(block));
   }
   #endif
 
@@ -480,7 +489,7 @@ static inline void _mi_free_block(mi_page_t* page, bool local, mi_block_t* block
     mi_check_padding(page, block);
     #if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN
     if (!mi_page_is_huge(page)) {   // huge page content may be already decommitted
-      memset(block, MI_DEBUG_FREED, mi_page_block_size(page));
+      mi_debug_fill(page, block, MI_DEBUG_FREED, mi_page_block_size(page));
     }
     #endif
     mi_block_set_next(page, block, page->local_free);
@@ -575,7 +584,7 @@ void mi_free(void* p) mi_attr_noexcept
       mi_check_padding(page, block);
       mi_stat_free(page, block);
       #if (MI_DEBUG>0) && !MI_TRACK_ENABLED  && !MI_TSAN
-      memset(block, MI_DEBUG_FREED, mi_page_block_size(page));
+      mi_debug_fill(page, block, MI_DEBUG_FREED, mi_page_block_size(page));
       #endif
       mi_track_free_size(p, mi_page_usable_size_of(page,block)); // faster then mi_usable_size as we already know the page and that p is unaligned
       mi_block_set_next(page, block, page->local_free);

Objects/mimalloc/init.c

@@ -13,27 +13,7 @@ terms of the MIT license. A copy of the license can be found in the file
 
 
 // Empty page used to initialize the small free pages array
-const mi_page_t _mi_page_empty = {
-  0, false, false, false, 0,
-  0,       // capacity
-  0,       // reserved capacity
-  { 0 },   // flags
-  false,   // is_zero
-  0,       // retire_expire
-  NULL,    // free
-  0,       // used
-  0,       // xblock_size
-  NULL,    // local_free
-  #if (MI_PADDING || MI_ENCODE_FREELIST)
-  { 0, 0 },
-  #endif
-  MI_ATOMIC_VAR_INIT(0), // xthread_free
-  MI_ATOMIC_VAR_INIT(0), // xheap
-  NULL, NULL
-  #if MI_INTPTR_SIZE==8
-  , { 0 }  // padding
-  #endif
-};
+const mi_page_t _mi_page_empty;
 
 #define MI_PAGE_EMPTY() ((mi_page_t*)&_mi_page_empty)
 
@@ -122,6 +102,7 @@ mi_decl_cache_align const mi_heap_t _mi_heap_empty = {
   MI_BIN_FULL, 0,   // page retired min/max
   NULL,             // next
   false,
+  0,
   0
 };
 

Objects/mimalloc/page.c

@@ -661,6 +661,7 @@ static void mi_page_init(mi_heap_t* heap, mi_page_t* page, size_t block_size, mi
   // set fields
   mi_page_set_heap(page, heap);
   page->tag = heap->tag;
+  page->debug_offset = heap->debug_offset;
   page->xblock_size = (block_size < MI_HUGE_BLOCK_SIZE ? (uint32_t)block_size : MI_HUGE_BLOCK_SIZE); // initialize before _mi_segment_page_start
   size_t page_size;
   const void* page_start = _mi_segment_page_start(segment, page, &page_size);

Python/pystate.c

@@ -2845,9 +2845,24 @@ tstate_mimalloc_bind(PyThreadState *tstate)
     // pools to keep Python objects from different interpreters separate.
     tld->segments.abandoned = &tstate->interp->mimalloc.abandoned_pool;
 
+    // Don't fill in the first N bytes up to ob_type in debug builds. We may
+    // access ob_tid and the refcount fields in the dict and list lock-less
+    // accesses, so they must remain valid for a while after deallocation.
+    size_t base_offset = offsetof(PyObject, ob_type);
+    if (_PyMem_DebugEnabled()) {
+        // The debug allocator adds two words at the beginning of each block.
+        base_offset += 2 * sizeof(size_t);
+    }
+    size_t debug_offsets[_Py_MIMALLOC_HEAP_COUNT] = {
+        [_Py_MIMALLOC_HEAP_OBJECT] = base_offset,
+        [_Py_MIMALLOC_HEAP_GC] = base_offset,
+        [_Py_MIMALLOC_HEAP_GC_PRE] = base_offset + 2 * sizeof(PyObject *),
+    };
+
     // Initialize each heap
     for (uint8_t i = 0; i < _Py_MIMALLOC_HEAP_COUNT; i++) {
         _mi_heap_init_ex(&mts->heaps[i], tld, _mi_arena_id_none(), false, i);
+        mts->heaps[i].debug_offset = (uint8_t)debug_offsets[i];
     }
 
     // By default, object allocations use _Py_MIMALLOC_HEAP_OBJECT.