gh-117755: Fix mimalloc for huge allocation on s390x

Fix mimalloc allocator for huge memory allocation (around
8,589,934,592 GiB) on s390x.

* Abort allocation early in mimalloc if the number of slices doesn't
  fit into uint32_t, to prevent a integer overflow (cast 64-bit
  size_t to uint32_t).
* Add test_large_alloc() to test_bigaddrspace.
* Reenable test_maxcontext_exact_arith() of test_decimal on s390x.

Author: vstinner

Lib/test/test_bigaddrspace.py

@@ -11,9 +11,25 @@
 from test import support
 from test.support import bigaddrspacetest, MAX_Py_ssize_t
 
-import unittest
+import contextlib
 import operator
+import os
 import sys
+import unittest
+
+
+@contextlib.contextmanager
+def ignore_stderr():
+    fd = 2
+    old_stderr = os.dup(fd)
+    try:
+        # Redirect stderr to /dev/null
+        with open(os.devnull, 'wb') as null:
+            os.dup2(null.fileno(), fd)
+            yield
+    finally:
+        os.dup2(old_stderr, fd)
+        os.close(old_stderr)
 
 
 class BytesTest(unittest.TestCase):
@@ -52,6 +68,48 @@ def test_repeat(self):
         finally:
             x = None
 
+    def test_large_alloc(self):
+        debug_bytes = 0
+        if support.check_impl_detail(cpython=True) and support.Py_DEBUG:
+            try:
+                from _testcapi import SIZEOF_SIZE_T
+            except ImportError:
+                if sys.maxsize > 0xffff_ffff:
+                    SIZEOF_SIZE_T = 8
+                else:
+                    SIZEOF_SIZE_T = 4
+
+            # The debug hook on memory allocator adds 3 size_t per memory block
+            # See PYMEM_DEBUG_EXTRA_BYTES in Objects/obmalloc.c.
+            debug_bytes = SIZEOF_SIZE_T * 3
+
+            try:
+                from _testinternalcapi import pymem_getallocatorsname
+                if not pymem_getallocatorsname().endswith('_debug'):
+                    # PYTHONMALLOC env var is used and disables the debug hook
+                    debug_bytes = 0
+            except (ImportError, RuntimeError):
+                pass
+
+        def allocate(size):
+            length = size - sys.getsizeof(b'') - debug_bytes
+            # allocate 'size' bytes
+            return b'x' * length
+
+        # 64-bit size which cannot be allocated on any reasonable hardware
+        # (in 2024) and must fail immediately with MemoryError.
+        for size in (
+            # gh-114331: test_decimal.test_maxcontext_exact_arith()
+            0x0BAFC246_72035E58,
+            # gh-117755: test_io.test_constructor()
+            0x7FFFFFFF_FFFFFFFF,
+        ):
+            with self.subTest(size=size):
+                with self.assertRaises(MemoryError):
+                    # ignore "mimalloc: error: unable to allocate memory"
+                    with ignore_stderr():
+                        allocate(size)
+
 
 class StrTest(unittest.TestCase):
 

Lib/test/test_decimal.py

@@ -38,8 +38,7 @@
                           check_disallow_instantiation)
 from test.support import (TestFailed,
                           run_with_locale, cpython_only,
-                          darwin_malloc_err_warning, is_emscripten,
-                          skip_on_s390x)
+                          darwin_malloc_err_warning, is_emscripten)
 from test.support.import_helper import import_fresh_module
 from test.support import threading_helper
 from test.support import warnings_helper
@@ -5651,9 +5650,6 @@ def __abs__(self):
     @unittest.skipIf(check_sanitizer(address=True, memory=True),
                      "ASAN/MSAN sanitizer defaults to crashing "
                      "instead of returning NULL for malloc failure.")
-    # gh-114331: The test allocates 784 271 641 GiB and mimalloc does not fail
-    # to allocate it when using mimalloc on s390x.
-    @skip_on_s390x
     def test_maxcontext_exact_arith(self):
 
         # Make sure that exact operations do not raise MemoryError due

Misc/NEWS.d/next/Core and Builtins/2024-04-12-12-28-49.gh-issue-117755.6ct8kU.rst

@@ -0,0 +1,2 @@
+Fix mimalloc allocator for huge memory allocation (around 8,589,934,592 GiB) on
+s390x. Patch by Victor Stinner.

Objects/mimalloc/segment.c

@@ -814,6 +814,9 @@ static mi_segment_t* mi_segment_os_alloc( size_t required, size_t page_alignment
     const size_t extra = align_offset - info_size;
     // recalculate due to potential guard pages
     *psegment_slices = mi_segment_calculate_slices(required + extra, ppre_size, pinfo_slices);
+
+    // mi_page_t.slice_count type is uint32_t
+    if (*psegment_slices > (size_t)UINT32_MAX) return NULL;
   }
 
   const size_t segment_size = (*psegment_slices) * MI_SEGMENT_SLICE_SIZE;
@@ -865,6 +868,9 @@ static mi_segment_t* mi_segment_alloc(size_t required, size_t page_alignment, mi
   size_t pre_size;
   size_t segment_slices = mi_segment_calculate_slices(required, &pre_size, &info_slices);
 
+  // mi_page_t.slice_count type is uint32_t
+  if (segment_slices > (size_t)UINT32_MAX) return NULL;
+
   // Commit eagerly only if not the first N lazy segments (to reduce impact of many threads that allocate just a little)
   const bool eager_delay = (// !_mi_os_has_overcommit() &&             // never delay on overcommit systems
                             _mi_current_thread_count() > 1 &&       // do not delay for the first N threads