Merge commit from fork

* Limit unpacker container preallocation to 1000 items

* Increase unpacker preallocation limit to 10000 items

Author: methane

msgpack/fallback.py

@@ -518,7 +518,7 @@ def _unpack(self, execute=EX_CONSTRUCT):
                     # TODO check whether we need to call `list_hook`
                     self._unpack(EX_SKIP)
                 return
-            ret = newlist_hint(n)
+            ret = newlist_hint(min(n, 10000))
             for i in range(n):
                 ret.append(self._unpack(EX_CONSTRUCT))
             if self._list_hook is not None:

msgpack/unpack.h

@@ -17,6 +17,7 @@
  */
 
 #define MSGPACK_EMBED_STACK_SIZE  (1024)
+#define MSGPACK_MAX_PREALLOC  (10000)
 #include "unpack_define.h"
 
 typedef struct unpack_user {
@@ -131,28 +132,48 @@ static inline int unpack_callback_array(unpack_user* u, unsigned int n, msgpack_
         PyErr_Format(PyExc_ValueError, "%u exceeds max_array_len(%zd)", n, u->max_array_len);
         return -1;
     }
-    PyObject *p = u->use_list ? PyList_New(n) : PyTuple_New(n);
+    /* Large tuples are built as lists and converted when complete. */
+    PyObject *p = n > MSGPACK_MAX_PREALLOC ? PyList_New(MSGPACK_MAX_PREALLOC) :
+        (u->use_list ? PyList_New(n) : PyTuple_New(n));
 
     if (!p)
         return -1;
     *o = p;
     return 0;
 }
 
+/* Consumes o on both success and failure; borrows the list. */
+static inline int unpack_list_item(PyObject *list, unsigned int current, PyObject *o)
+{
+    if (current < MSGPACK_MAX_PREALLOC) {
+        PyList_SET_ITEM(list, current, o);
+        return 0;
+    }
+    int ret = PyList_Append(list, o);
+    Py_DECREF(o);
+    return ret;
+}
+
 /* Consumes o on both success and failure; borrows the container. */
 static inline int unpack_callback_array_item(unpack_user* u, unsigned int current, msgpack_unpack_object* c, msgpack_unpack_object o)
 {
-    if (u->use_list)
-        PyList_SET_ITEM(*c, current, o);
-    else
-        PyTuple_SET_ITEM(*c, current, o);
+    if (PyList_CheckExact(*c))
+        return unpack_list_item(*c, current, o);
+    PyTuple_SET_ITEM(*c, current, o);
     return 0;
 }
 
 /* Replaces the owned container reference with the hook result, or NULL on
  * failure. The original reference is consumed in either case. */
 static inline int unpack_callback_array_end(unpack_user* u, msgpack_unpack_object* c)
 {
+    if (!u->use_list && PyList_CheckExact(*c)) {
+        PyObject *tuple = PyList_AsTuple(*c);
+        Py_DECREF(*c);
+        *c = tuple;
+        if (!tuple)
+            return -1;
+    }
     if (u->list_hook) {
         PyObject *new_c = PyObject_CallFunctionObjArgs(u->list_hook, *c, NULL);
         Py_DECREF(*c);
@@ -171,7 +192,7 @@ static inline int unpack_callback_map(unpack_user* u, unsigned int n, msgpack_un
     }
     PyObject *p;
     if (u->has_pairs_hook) {
-        p = PyList_New(n); // Or use tuple?
+        p = PyList_New(n > MSGPACK_MAX_PREALLOC ? MSGPACK_MAX_PREALLOC : n);
     }
     else {
         p = PyDict_New();
@@ -200,8 +221,7 @@ static inline int unpack_callback_map_item(unpack_user* u, unsigned int current,
         Py_DECREF(v);
         if (!item)
             return -1;
-        PyList_SET_ITEM(*c, current, item);
-        return 0;
+        return unpack_list_item(*c, current, item);
     }
     int ret = PyDict_SetItem(*c, k, v);
     Py_DECREF(k);

test/test_limits.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python
 import sysconfig
+import tracemalloc
 
 import pytest
 
 from msgpack import (
     ExtType,
+    OutOfData,
     Packer,
     PackOverflowError,
     PackValueError,
@@ -118,6 +120,62 @@ def test_max_map_len():
         unpacker.unpack()
 
 
+@pytest.mark.parametrize("size", [9999, 10000, 10001, 20001])
+@pytest.mark.parametrize("use_list", [False, True])
+def test_array_preallocation_boundary(size, use_list):
+    value = list(range(size))
+    expected = value if use_list else tuple(value)
+    packed = packb(value)
+    assert unpackb(packed, use_list=use_list, list_hook=lambda obj: obj) == expected
+
+    unpacker = Unpacker(use_list=use_list)
+    # Suspend after crossing the preallocation boundary.
+    unpacker.feed(packed[:-1])
+    with pytest.raises(OutOfData):
+        unpacker.unpack()
+    unpacker.feed(packed[-1:])
+    assert unpacker.unpack() == expected
+
+
+@pytest.mark.parametrize("size", [9999, 10000, 10001, 20001])
+def test_pairs_preallocation_boundary(size):
+    value = {str(i): i for i in range(size)}
+    expected = list(value.items())
+    packed = packb(value)
+    assert unpackb(packed, object_pairs_hook=list) == expected
+
+    unpacker = Unpacker(object_pairs_hook=list)
+    unpacker.feed(packed[:-1])
+    with pytest.raises(OutOfData):
+        unpacker.unpack()
+    unpacker.feed(packed[-1:])
+    assert unpacker.unpack() == expected
+
+
+@pytest.mark.parametrize("kind", ["list", "tuple", "pairs"])
+def test_nested_headers_preallocation(kind):
+    packer = Packer()
+    header = (
+        packer.pack_map_header(100_000) if kind == "pairs" else packer.pack_array_header(100_000)
+    )
+    unpacker = Unpacker(
+        use_list=kind != "tuple", object_pairs_hook=list if kind == "pairs" else None
+    )
+    # Map headers nest in key position, so no completed pairs are needed.
+    unpacker.feed(header * 16)
+    already_tracing = tracemalloc.is_tracing()
+    tracemalloc.start()
+    try:
+        before = tracemalloc.get_traced_memory()[0]
+        with pytest.raises(OutOfData):
+            unpacker.unpack()
+        # Full preallocation would retain over 12 MB in the C unpacker.
+        assert tracemalloc.get_traced_memory()[0] - before < 4 * 1024 * 1024
+    finally:
+        if not already_tracing:
+            tracemalloc.stop()
+
+
 def test_max_ext_len():
     d = ExtType(42, b"abc")
     packed = packb(d)