gh-161478: Collect large compression outputs in blocks again, except with glibc

Include/internal/pycore_blocks_output_buffer.h

@@ -6,14 +6,17 @@
         stream->next_out:  point to the next output position.
         stream->avail_out: the number of available bytes left in the buffer.
 
-   It maintains a list of bytes object, so there is no overhead of resizing
-   the buffer.
+   The first blocks are written into a PyBytesWriter, which is returned as is
+   when the output fits in them. Further blocks are separate bytes objects,
+   joined once by _BlocksOutputBuffer_Finish(), except where the writer is
+   known to grow in place (glibc). See OUTPUT_BUFFER_HEAD_SIZE.
 
    Usage:
 
    1, Initialize the struct instance like this:
-        _BlocksOutputBuffer buffer = {.list = NULL};
-      Set .list to NULL for _BlocksOutputBuffer_OnError()
+        _BlocksOutputBuffer buffer = {.writer = NULL};
+      This zeroes the other members too. _BlocksOutputBuffer_OnError()
+      relies on .writer and .list being NULL.
 
    2, Initialize the buffer use one of these functions:
         _BlocksOutputBuffer_InitAndGrow()
@@ -28,7 +31,8 @@
    5, Finish the buffer, and return a bytes object:
         _BlocksOutputBuffer_Finish()
 
-   6, Clean up the buffer when an error occurred:
+   6, Clean up the buffer when an error occurred, including when
+      _BlocksOutputBuffer_Finish() returned NULL:
         _BlocksOutputBuffer_OnError()
 */
 
@@ -53,23 +57,58 @@ typedef struct {
     Py_ssize_t max_length;
     // Number of blocks of bytes. Used to calculate next allocation size
     size_t num_blocks;
+    // Blocks after the first OUTPUT_BUFFER_HEAD_SIZE bytes, as a list of
+    // bytes objects. NULL while all the output is in the writer.
+    PyObject *list;
+    // Size of the writer's data once blocks are added to the list
+    Py_ssize_t head_size;
 } _BlocksOutputBuffer;
 
 static const char unable_allocate_msg[] = "Unable to allocate output buffer.";
 
 /* In 32-bit build, the max block size should <= INT32_MAX. */
 #define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)
 
-/* Block size sequence */
-#define KB (1024)
+/* Block size sequence. The first four blocks are the writer's head, see
+   OUTPUT_BUFFER_HEAD_SIZE; the tests in test_zlib, test_gzip, test_bz2,
+   test_lzma and test_zstd probe the boundaries of these sizes. */
+#define OUTPUT_BUFFER_BLOCK_1 (32*1024)
+#define OUTPUT_BUFFER_BLOCK_2 (64*1024)
+#define OUTPUT_BUFFER_BLOCK_3 (256*1024)
+#define OUTPUT_BUFFER_BLOCK_4 (1024*1024)
 #define MB (1024*1024)
 static const Py_ssize_t BUFFER_BLOCK_SIZE[] =
-    { 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,
+    { OUTPUT_BUFFER_BLOCK_1, OUTPUT_BUFFER_BLOCK_2, OUTPUT_BUFFER_BLOCK_3,
+      OUTPUT_BUFFER_BLOCK_4, 4*MB, 8*MB, 16*MB, 16*MB,
       32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,
       OUTPUT_BUFFER_MAX_BLOCK_SIZE };
-#undef KB
 #undef MB
 
+/* The writer grows to at most this many bytes: the first four blocks of the
+   sequence, 1.34 MiB. Up to there, returning the writer without a join
+   beats separate blocks; from the 4 MiB block on, the blocks win (measured
+   on macOS and Windows, see gh-161478). Later blocks go to the list, and so
+   does every block after a larger initial size.
+
+   On macOS and Windows, a large block grown with realloc() lands in memory
+   the process has never touched (on Windows, realloc() also copies it), so
+   a writer that keeps growing faults in new pages on every call, which
+   costs more than writing blocks once and joining them at the end, as the
+   code before gh-139877 did. With glibc, realloc() grows a large block in
+   place (mremap()), and there growing the writer is faster than joining
+   blocks. Everything else gets the blocks: platforms that were not
+   measured, and allocators that copy on realloc(): Android's, musl's, and
+   mimalloc, the free-threaded build's object allocator (with
+   PYTHONMALLOC=mimalloc on macOS, the growing writer regressed just like
+   with the system allocator). */
+#if defined(__GLIBC__) && !defined(Py_GIL_DISABLED)
+#  define OUTPUT_BUFFER_HEAD_SIZE PY_SSIZE_T_MAX
+#else
+#  define OUTPUT_BUFFER_HEAD_SIZE \
+     (OUTPUT_BUFFER_BLOCK_1 + OUTPUT_BUFFER_BLOCK_2 \
+      + OUTPUT_BUFFER_BLOCK_3 + OUTPUT_BUFFER_BLOCK_4)
+#endif
+
 /* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole
    allocated size growth step is:
     1   32 KB       +32 KB
@@ -130,6 +169,8 @@ _BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,
     buffer->allocated = block_size;
     buffer->max_length = max_length;
     buffer->num_blocks = 1;
+    buffer->list = NULL;
+    buffer->head_size = 0;
 
     *next_out = PyBytesWriter_GetData(buffer->writer);
     return block_size;
@@ -161,6 +202,8 @@ _BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,
     buffer->allocated = init_size;
     buffer->max_length = -1;
     buffer->num_blocks = 1;
+    buffer->list = NULL;
+    buffer->head_size = 0;
 
     *next_out = PyBytesWriter_GetData(buffer->writer);
     return init_size;
@@ -209,19 +252,44 @@ _BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,
         return -1;
     }
 
-    if (PyBytesWriter_Grow(buffer->writer, block_size)) {
-        PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
-        return -1;
+    if (buffer->list == NULL
+        && buffer->allocated <= OUTPUT_BUFFER_HEAD_SIZE - block_size)
+    {
+        // grow the writer
+        if (PyBytesWriter_Grow(buffer->writer, block_size)) {
+            PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+            return -1;
+        }
+        char *data = PyBytesWriter_GetData(buffer->writer);
+        *next_out = data + buffer->allocated;
+    }
+    else {
+        // add a block to the list
+        if (buffer->list == NULL) {
+            buffer->list = PyList_New(0);
+            if (buffer->list == NULL) {
+                PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+                return -1;
+            }
+            buffer->head_size = buffer->allocated;
+        }
+        PyObject *block = PyBytes_FromStringAndSize(NULL, block_size);
+        if (block == NULL) {
+            PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+            return -1;
+        }
+        if (PyList_Append(buffer->list, block) < 0) {
+            Py_DECREF(block);
+            PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+            return -1;
+        }
+        Py_DECREF(block);  // the list holds it
+        *next_out = PyBytes_AS_STRING(block);
     }
-
-    Py_ssize_t current_size = buffer->allocated;
 
     // set variables
     buffer->allocated += block_size;
     buffer->num_blocks += 1;
-
-    char *data = PyBytesWriter_GetData(buffer->writer);
-    *next_out = data + current_size;
     return block_size;
 }
 
@@ -244,9 +312,50 @@ _BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,
 {
     PyObject *obj;
     assert(buffer->writer != NULL);
-    obj = PyBytesWriter_FinishWithSize(buffer->writer,
-                                       buffer->allocated - avail_out);
+    if (buffer->list == NULL) {
+        obj = PyBytesWriter_FinishWithSize(buffer->writer,
+                                           buffer->allocated - avail_out);
+        buffer->writer = NULL;
+        return obj;
+    }
+
+    Py_ssize_t num = PyList_GET_SIZE(buffer->list);
+    if (num == 1
+        && avail_out == PyBytes_GET_SIZE(PyList_GET_ITEM(buffer->list, 0)))
+    {
+        // the one block was never written to: the writer holds all the data
+        obj = PyBytesWriter_FinishWithSize(buffer->writer, buffer->head_size);
+        buffer->writer = NULL;
+        Py_CLEAR(buffer->list);
+        return obj;
+    }
+
+    // join the writer's data and the blocks into one bytes object
+    Py_ssize_t size = buffer->allocated - avail_out;
+    obj = PyBytes_FromStringAndSize(NULL, size);
+    if (obj == NULL) {
+        PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+        return NULL;  // the caller calls _BlocksOutputBuffer_OnError()
+    }
+    char *dst = PyBytes_AS_STRING(obj);
+    memcpy(dst, PyBytesWriter_GetData(buffer->writer), buffer->head_size);
+    dst += buffer->head_size;
+
+    for (Py_ssize_t i = 0; i < num; i++) {
+        PyObject *block = PyList_GET_ITEM(buffer->list, i);
+        Py_ssize_t block_size = PyBytes_GET_SIZE(block);
+        if (i == num - 1) {
+            // the last block holds avail_out unused bytes
+            block_size -= avail_out;
+        }
+        memcpy(dst, PyBytes_AS_STRING(block), block_size);
+        dst += block_size;
+    }
+    assert(dst == PyBytes_AS_STRING(obj) + size);
+
+    PyBytesWriter_Discard(buffer->writer);
     buffer->writer = NULL;
+    Py_CLEAR(buffer->list);
     return obj;
 }
 
@@ -256,6 +365,7 @@ _BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)
 {
     PyBytesWriter_Discard(buffer->writer);
     buffer->writer = NULL;
+    Py_CLEAR(buffer->list);
 }
 
 #ifdef __cplusplus

Lib/test/test_bz2.py

@@ -1324,6 +1324,50 @@ def test_bzlib_version(self):
         self.assertStartsWith(bz2.bzlib_version, '%d.%d.%d' % v)
 
 
+class OutputBufferTest(unittest.TestCase):
+    # Outputs at and around the boundaries of the internal output buffer's
+    # blocks (Include/internal/pycore_blocks_output_buffer.h): the first four
+    # (32 KiB, 64 KiB, 256 KiB and 1 MiB) are one buffer grown by realloc();
+    # the next one (4 MiB) is a separate block on most platforms.
+    HEAD = (32 + 64 + 256 + 1024) * 1024
+    SIZES = [HEAD - 1, HEAD, HEAD + 1,
+             HEAD + 4 * 1024 * 1024 - 1, HEAD + 4 * 1024 * 1024,
+             HEAD + 4 * 1024 * 1024 + 1]
+    # a pattern every codec compresses quickly: 256 varied bytes per 4 KiB
+    PATTERN = bytes(range(256)) + bytes(4096 - 256)
+
+    @classmethod
+    def data(cls, size):
+        return (cls.PATTERN * (size // len(cls.PATTERN) + 1))[:size]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.compressed = {size: bz2.compress(cls.data(size), 1)
+                          for size in cls.SIZES}
+
+    def test_decompress(self):
+        for size in self.SIZES:
+            with self.subTest(size=size):
+                data = self.data(size)
+                self.assertEqual(bz2.decompress(self.compressed[size]), data)
+
+    def test_decompressor(self):
+        for size in self.SIZES:
+            with self.subTest(size=size):
+                data = self.data(size)
+                compressed = self.compressed[size]
+                self.assertEqual(
+                    BZ2Decompressor().decompress(compressed), data)
+                # max_length past the first blocks and short of the next one
+                max_length = self.HEAD + 100 * 1024
+                d = BZ2Decompressor()
+                first = d.decompress(compressed, max_length)
+                self.assertEqual(len(first), min(size, max_length))
+                rest = b'' if d.eof else d.decompress(b'')
+                self.assertEqual(first + rest, data)
+                self.assertTrue(d.eof)
+
+
 class TestModule(unittest.TestCase):
     def test_deprecated_module_attributes(self):
         support.check_deprecated_module_attributes(

Lib/test/test_gzip.py

@@ -1337,5 +1337,28 @@ def test_decompress_cannot_have_flags_compression(self):
         self.assertEqual(out, b'')
 
 
+class TestOutputBuffer(unittest.TestCase):
+    # Outputs at and around the boundaries of the internal output buffer's
+    # blocks (Include/internal/pycore_blocks_output_buffer.h): the first four
+    # (32 KiB, 64 KiB, 256 KiB and 1 MiB) are one buffer grown by realloc();
+    # the next one (4 MiB) is a separate block on most platforms.
+    HEAD = (32 + 64 + 256 + 1024) * 1024
+    SIZES = [HEAD - 1, HEAD, HEAD + 1,
+             HEAD + 4 * 1024 * 1024 - 1, HEAD + 4 * 1024 * 1024,
+             HEAD + 4 * 1024 * 1024 + 1]
+    # a pattern every codec compresses quickly: 256 varied bytes per 4 KiB
+    PATTERN = bytes(range(256)) + bytes(4096 - 256)
+
+    @classmethod
+    def data(cls, size):
+        return (cls.PATTERN * (size // len(cls.PATTERN) + 1))[:size]
+
+    def test_decompress(self):
+        for size in self.SIZES:
+            with self.subTest(size=size):
+                data = self.data(size)
+                self.assertEqual(gzip.decompress(gzip.compress(data, 1)), data)
+
+
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_lzma.py

@@ -2214,5 +2214,49 @@ def test_filter_properties_roundtrip(self):
 )
 
 
+class OutputBufferTestCase(unittest.TestCase):
+    # Outputs at and around the boundaries of the internal output buffer's
+    # blocks (Include/internal/pycore_blocks_output_buffer.h): the first four
+    # (32 KiB, 64 KiB, 256 KiB and 1 MiB) are one buffer grown by realloc();
+    # the next one (4 MiB) is a separate block on most platforms.
+    HEAD = (32 + 64 + 256 + 1024) * 1024
+    SIZES = [HEAD - 1, HEAD, HEAD + 1,
+             HEAD + 4 * 1024 * 1024 - 1, HEAD + 4 * 1024 * 1024,
+             HEAD + 4 * 1024 * 1024 + 1]
+    # a pattern every codec compresses quickly: 256 varied bytes per 4 KiB
+    PATTERN = bytes(range(256)) + bytes(4096 - 256)
+
+    @classmethod
+    def data(cls, size):
+        return (cls.PATTERN * (size // len(cls.PATTERN) + 1))[:size]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.compressed = {size: lzma.compress(cls.data(size), preset=0)
+                          for size in cls.SIZES}
+
+    def test_decompress(self):
+        for size in self.SIZES:
+            with self.subTest(size=size):
+                data = self.data(size)
+                self.assertEqual(lzma.decompress(self.compressed[size]), data)
+
+    def test_decompressor(self):
+        for size in self.SIZES:
+            with self.subTest(size=size):
+                data = self.data(size)
+                compressed = self.compressed[size]
+                self.assertEqual(
+                    LZMADecompressor().decompress(compressed), data)
+                # max_length past the first blocks and short of the next one
+                max_length = self.HEAD + 100 * 1024
+                d = LZMADecompressor()
+                first = d.decompress(compressed, max_length)
+                self.assertEqual(len(first), min(size, max_length))
+                rest = b'' if d.eof else d.decompress(b'')
+                self.assertEqual(first + rest, data)
+                self.assertTrue(d.eof)
+
+
 if __name__ == "__main__":
     unittest.main()