diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index 3c81808f..a9211e0b 100644 --- a/msgpack/_unpacker.pyx +++ b/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 diff --git a/msgpack/unpack.h b/msgpack/unpack.h index 90622bbc..494cf148 100644 --- a/msgpack/unpack.h +++ b/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) { diff --git a/msgpack/unpack_template.h b/msgpack/unpack_template.h index 9e3e3ce0..47c4a157 100644 --- a/msgpack/unpack_template.h +++ b/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: diff --git a/test/test_unpack.py b/test/test_unpack.py index 03e49ca3..44de7f75 100644 --- a/test/test_unpack.py +++ b/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