Cache decoded map keys in the C unpacker

msgpack/_unpacker.pyx

@@ -54,6 +54,7 @@ cdef extern from "unpack.h":
     execute_fn read_map_header
 
     void unpack_init(unpack_context* ctx)
+    void unpack_init_key_cache(unpack_context* ctx)
     # Transfers the result's owned reference to Cython and clears its slot.
     object unpack_take_result(unpack_context* ctx)
     void unpack_clear(unpack_context* ctx)
@@ -68,6 +69,7 @@ cdef inline init_ctx(unpack_context *ctx,
                      Py_ssize_t max_array_len, Py_ssize_t max_map_len,
                      Py_ssize_t max_ext_len):
     unpack_init(ctx)
+    unpack_init_key_cache(ctx)
     ctx.user.use_list = use_list
     ctx.user.raw = raw
     ctx.user.strict_map_key = strict_map_key

msgpack/unpack.h

@@ -18,6 +18,11 @@
 
 #define MSGPACK_EMBED_STACK_SIZE  (1024)
 #define MSGPACK_MAX_PREALLOC  (10000)
+/* Direct-mapped cache of decoded, interned ASCII map keys. It is enabled only
+ * after MSGPACK_KEY_CACHE_THRESHOLD keys, so small messages skip its setup. */
+#define MSGPACK_KEY_CACHE_SIZE  (256)  /* must be a power of two */
+#define MSGPACK_KEY_CACHE_MAX_LEN  (64)
+#define MSGPACK_KEY_CACHE_THRESHOLD  (32)
 #include "unpack_define.h"
 
 typedef struct unpack_user {
@@ -34,8 +39,22 @@ typedef struct unpack_user {
     PyObject *utc;
     const char *unicode_errors;
     Py_ssize_t max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len;
+    unsigned int map_keys_seen;
+    bool key_cache_ready;
+    PyObject *key_cache[MSGPACK_KEY_CACHE_SIZE];  /* valid only if key_cache_ready */
 } unpack_user;
 
+static inline void unpack_key_cache_clear(unpack_user* u)
+{
+    if (u->key_cache_ready) {
+        for (unsigned int i = 0; i < MSGPACK_KEY_CACHE_SIZE; i++) {
+            Py_CLEAR(u->key_cache[i]);
+        }
+        u->key_cache_ready = false;
+    }
+    u->map_keys_seen = 0;
+}
+
 typedef PyObject* msgpack_unpack_object; //TODO: remove typedef and use PyObject* directly
 struct unpack_context;
 typedef struct unpack_context unpack_context;
@@ -203,6 +222,26 @@ static inline int unpack_callback_map(unpack_user* u, unsigned int n, msgpack_un
     return 0;
 }
 
+static inline PyObject** unpack_key_cache_slot(unpack_user* u, const char* p, Py_ssize_t l)
+{
+    uint32_t hash = 2166136261u;  /* FNV-1a */
+    for (Py_ssize_t i = 0; i < l; i++) {
+        hash = (hash ^ (unsigned char)p[i]) * 16777619u;
+    }
+    return &u->key_cache[hash & (MSGPACK_KEY_CACHE_SIZE - 1)];
+}
+
+/* Only ASCII keys are stored, so a cached key's data equals its raw bytes. */
+static inline void unpack_key_cache_store(unpack_user* u, PyObject* key)
+{
+    Py_ssize_t l = PyUnicode_GET_LENGTH(key);
+    if (!u->key_cache_ready || !PyUnicode_IS_ASCII(key) || l > MSGPACK_KEY_CACHE_MAX_LEN)
+        return;
+    PyObject** slot = unpack_key_cache_slot(u, (const char*)PyUnicode_DATA(key), l);
+    Py_INCREF(key);
+    Py_XSETREF(*slot, key);
+}
+
 /* Consumes k and v on both success and failure; borrows the container. */
 static inline int unpack_callback_map_item(unpack_user* u, unsigned int current, msgpack_unpack_object* c, msgpack_unpack_object k, msgpack_unpack_object v)
 {
@@ -212,8 +251,10 @@ static inline int unpack_callback_map_item(unpack_user* u, unsigned int current,
         Py_DECREF(v);
         return -1;
     }
-    if (PyUnicode_CheckExact(k)) {
+    if (PyUnicode_CheckExact(k) && !PyUnicode_CHECK_INTERNED(k)) {
         PyUnicode_InternInPlace(&k);
+        /* Cache only keys of complete pairs, as interned strings may be immortal. */
+        unpack_key_cache_store(u, k);
     }
     if (u->has_pairs_hook) {
         msgpack_unpack_object item = PyTuple_Pack(2, k, v);
@@ -262,6 +303,31 @@ static inline int unpack_callback_raw(unpack_user* u, const char* b, const char*
     return 0;
 }
 
+/* Map keys repeat across objects of the same shape, so reuse the interned
+ * string instead of decoding and interning the same bytes again. */
+static inline int unpack_callback_map_key_raw(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
+{
+    if (u->raw || l > MSGPACK_KEY_CACHE_MAX_LEN || l > u->max_str_len) {
+        return unpack_callback_raw(u, b, p, l, o);
+    }
+    if (!u->key_cache_ready) {
+        if (++u->map_keys_seen < MSGPACK_KEY_CACHE_THRESHOLD) {
+            return unpack_callback_raw(u, b, p, l, o);
+        }
+        memset(u->key_cache, 0, sizeof(u->key_cache));
+        u->key_cache_ready = true;
+    }
+
+    PyObject* cached = *unpack_key_cache_slot(u, p, l);
+    if (cached && PyUnicode_GET_LENGTH(cached) == (Py_ssize_t)l &&
+            memcmp(PyUnicode_DATA(cached), p, l) == 0) {
+        Py_INCREF(cached);
+        *o = cached;
+        return 0;
+    }
+    return unpack_callback_raw(u, b, p, l, o);
+}
+
 static inline int unpack_callback_bin(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
 {
     if (l > u->max_bin_len) {

msgpack/unpack_template.h

@@ -58,6 +58,13 @@ static inline void unpack_init(unpack_context* ctx)
     ctx->execute = NULL;
 }
 
+/* Must run before any unpack_clear() on a context whose memory is uninitialized. */
+static inline void unpack_init_key_cache(unpack_context* ctx)
+{
+    ctx->user.map_keys_seen = 0;
+    ctx->user.key_cache_ready = false;
+}
+
 static inline PyObject* unpack_take_result(unpack_context* ctx)
 {
     return unpack_take_ref(&ctx->result);
@@ -70,6 +77,7 @@ static inline void unpack_clear(unpack_context *ctx)
         Py_CLEAR(ctx->stack[i].obj);
     }
     Py_CLEAR(ctx->result);
+    unpack_key_cache_clear(&ctx->user);
     unpack_init(ctx);
 }
 
@@ -293,6 +301,9 @@ static inline int unpack_execute(bool construct, unpack_context* ctx, const char
                 again_fixed_trail_if_zero(ACS_RAW_VALUE, _msgpack_load32(uint32_t,n), _raw_zero);
             case ACS_RAW_VALUE:
             _raw_zero:
+                if (top > 0 && stack[top-1].ct == CT_MAP_KEY) {
+                    push_variable_value(_map_key_raw, data, n, trail);
+                }
                 push_variable_value(_raw, data, n, trail);
 
             case ACS_EXT_VALUE:

test/test_unpack.py

@@ -1,5 +1,6 @@
 import gc
 import sys
+import uuid
 import weakref
 from io import BytesIO
 
@@ -185,3 +186,84 @@ def test_unpack_raises_extra_data_on_stream_with_trailing_bytes():
         unpack(stream)
     assert exc_info.value.unpacked == 100
     assert exc_info.value.extra == packb(200)
+
+
+# Enough maps with the same keys to exercise repeated-key handling.
+REPEATED_KEY_MAPS = [{"id": i, "state": "open", "ключ": i, "x" * 70: i} for i in range(100)]
+
+
+def test_unpackb_repeated_map_keys_are_equal_and_interned():
+    result = unpackb(packb(REPEATED_KEY_MAPS))
+    assert result == REPEATED_KEY_MAPS
+    for obj in result:
+        for key in obj:
+            assert key is sys.intern(key)
+
+
+def test_unpacker_repeated_map_keys_across_objects():
+    up = Unpacker()
+    up.feed(b"".join(packb(obj) for obj in REPEATED_KEY_MAPS))
+    result = list(up)
+    assert result == REPEATED_KEY_MAPS
+    assert all(key is sys.intern(key) for obj in result for key in obj)
+
+
+def test_unpackb_repeated_map_keys_with_raw():
+    result = unpackb(packb(REPEATED_KEY_MAPS), raw=True)
+    assert [list(obj) for obj in result] == [[k.encode() for k in obj] for obj in REPEATED_KEY_MAPS]
+
+
+def test_unpackb_many_distinct_and_empty_map_keys():
+    keys = [""] + [f"k{i}" for i in range(600)]
+    maps = [{key: j for key in keys} for j in range(3)]
+    result = unpackb(packb(maps))
+    assert result == maps
+    assert all(list(obj) == keys for obj in result)
+
+
+# 100 maps with the key "ab", then one map whose key is `last_key_bytes`.
+def _warm_maps_then_key(last_key_bytes):
+    warm = b"".join(bytes([0x81, 0xA2]) + b"ab" + bytes([i]) for i in range(100))
+    last = bytes([0x81, 0xA0 | len(last_key_bytes)]) + last_key_bytes + b"\x00"
+    return bytes([0xDC, 0x00, 101]) + warm + last
+
+
+def test_unpackb_repeated_map_keys_respect_max_str_len():
+    packed = _warm_maps_then_key(b"abcd")
+    with raises(ValueError, match="exceeds max_str_len"):
+        unpackb(packed, max_str_len=3)
+    assert unpackb(packed, max_str_len=4)[-1] == {"abcd": 0}
+
+
+def test_unpackb_repeated_map_keys_then_invalid_utf8():
+    packed = _warm_maps_then_key(b"a\xff")
+    with raises(UnicodeDecodeError):
+        unpackb(packed)
+    replaced_key = b"a\xff".decode("utf-8", "replace")
+    assert unpackb(packed, unicode_errors="replace")[-1] == {replaced_key: 0}
+
+
+@mark.skipif(
+    Unpacker.__module__ == "msgpack.fallback",
+    reason="the fallback interns a key before reading its value",
+)
+def test_unpackb_does_not_intern_key_of_incomplete_pair():
+    # Interned strings may be immortal, so a key whose value never arrives
+    # must not be interned. The key is unique per call, as tests run in
+    # parallel threads and repeat under pytest-run-parallel.
+    suffix = uuid.uuid4().hex[:28]  # "nc-" + 28 chars fits a fixstr
+    packed = _warm_maps_then_key(f"nc-{suffix}".encode())[:-1]  # drop the last value
+    with raises(ValueError, match="incomplete"):
+        unpackb(packed)
+    fresh = "".join(["nc-", suffix])
+    assert sys.intern(fresh) is fresh
+
+
+def test_unpacker_repeated_map_keys_after_error():
+    up = Unpacker()
+    up.feed(_warm_maps_then_key(b"a\xff"))
+    with raises(UnicodeDecodeError):
+        up.unpack()
+    up.__init__()
+    up.feed(packb(REPEATED_KEY_MAPS))
+    assert up.unpack() == REPEATED_KEY_MAPS