Skip to content

Commit 8fa9e09

Browse files
Cache decoded map keys in the C unpacker
Repeated map keys are decoded and interned once per occurrence. Reuse the interned key from a small per-context cache once a context has seen 32 keys, and fill the cache only with keys of complete pairs.
1 parent 8a88e3d commit 8fa9e09

4 files changed

Lines changed: 162 additions & 1 deletion

File tree

‎msgpack/_unpacker.pyx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ cdef extern from "unpack.h":
5454
execute_fn read_map_header
5555

5656
void unpack_init(unpack_context* ctx)
57+
void unpack_init_key_cache(unpack_context* ctx)
5758
# Transfers the result's owned reference to Cython and clears its slot.
5859
object unpack_take_result(unpack_context* ctx)
5960
void unpack_clear(unpack_context* ctx)
@@ -68,6 +69,7 @@ cdef inline init_ctx(unpack_context *ctx,
6869
Py_ssize_t max_array_len, Py_ssize_t max_map_len,
6970
Py_ssize_t max_ext_len):
7071
unpack_init(ctx)
72+
unpack_init_key_cache(ctx)
7173
ctx.user.use_list = use_list
7274
ctx.user.raw = raw
7375
ctx.user.strict_map_key = strict_map_key

‎msgpack/unpack.h‎

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@
1818

1919
#define MSGPACK_EMBED_STACK_SIZE (1024)
2020
#define MSGPACK_MAX_PREALLOC (10000)
21+
/* Direct-mapped cache of decoded, interned ASCII map keys. It is enabled only
22+
* after MSGPACK_KEY_CACHE_THRESHOLD keys, so small messages skip its setup. */
23+
#define MSGPACK_KEY_CACHE_SIZE (256) /* must be a power of two */
24+
#define MSGPACK_KEY_CACHE_MAX_LEN (64)
25+
#define MSGPACK_KEY_CACHE_THRESHOLD (32)
2126
#include "unpack_define.h"
2227

2328
typedef struct unpack_user {
@@ -34,8 +39,22 @@ typedef struct unpack_user {
3439
PyObject *utc;
3540
const char *unicode_errors;
3641
Py_ssize_t max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len;
42+
unsigned int map_keys_seen;
43+
bool key_cache_ready;
44+
PyObject *key_cache[MSGPACK_KEY_CACHE_SIZE]; /* valid only if key_cache_ready */
3745
} unpack_user;
3846

47+
static inline void unpack_key_cache_clear(unpack_user* u)
48+
{
49+
if (u->key_cache_ready) {
50+
for (unsigned int i = 0; i < MSGPACK_KEY_CACHE_SIZE; i++) {
51+
Py_CLEAR(u->key_cache[i]);
52+
}
53+
u->key_cache_ready = false;
54+
}
55+
u->map_keys_seen = 0;
56+
}
57+
3958
typedef PyObject* msgpack_unpack_object; //TODO: remove typedef and use PyObject* directly
4059
struct unpack_context;
4160
typedef struct unpack_context unpack_context;
@@ -203,6 +222,26 @@ static inline int unpack_callback_map(unpack_user* u, unsigned int n, msgpack_un
203222
return 0;
204223
}
205224

225+
static inline PyObject** unpack_key_cache_slot(unpack_user* u, const char* p, Py_ssize_t l)
226+
{
227+
uint32_t hash = 2166136261u; /* FNV-1a */
228+
for (Py_ssize_t i = 0; i < l; i++) {
229+
hash = (hash ^ (unsigned char)p[i]) * 16777619u;
230+
}
231+
return &u->key_cache[hash & (MSGPACK_KEY_CACHE_SIZE - 1)];
232+
}
233+
234+
/* Only ASCII keys are stored, so a cached key's data equals its raw bytes. */
235+
static inline void unpack_key_cache_store(unpack_user* u, PyObject* key)
236+
{
237+
Py_ssize_t l = PyUnicode_GET_LENGTH(key);
238+
if (!u->key_cache_ready || !PyUnicode_IS_ASCII(key) || l > MSGPACK_KEY_CACHE_MAX_LEN)
239+
return;
240+
PyObject** slot = unpack_key_cache_slot(u, (const char*)PyUnicode_DATA(key), l);
241+
Py_INCREF(key);
242+
Py_XSETREF(*slot, key);
243+
}
244+
206245
/* Consumes k and v on both success and failure; borrows the container. */
207246
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)
208247
{
@@ -212,8 +251,10 @@ static inline int unpack_callback_map_item(unpack_user* u, unsigned int current,
212251
Py_DECREF(v);
213252
return -1;
214253
}
215-
if (PyUnicode_CheckExact(k)) {
254+
if (PyUnicode_CheckExact(k) && !PyUnicode_CHECK_INTERNED(k)) {
216255
PyUnicode_InternInPlace(&k);
256+
/* Cache only keys of complete pairs, as interned strings may be immortal. */
257+
unpack_key_cache_store(u, k);
217258
}
218259
if (u->has_pairs_hook) {
219260
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*
262303
return 0;
263304
}
264305

306+
/* Map keys repeat across objects of the same shape, so reuse the interned
307+
* string instead of decoding and interning the same bytes again. */
308+
static inline int unpack_callback_map_key_raw(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
309+
{
310+
if (u->raw || l > MSGPACK_KEY_CACHE_MAX_LEN || l > u->max_str_len) {
311+
return unpack_callback_raw(u, b, p, l, o);
312+
}
313+
if (!u->key_cache_ready) {
314+
if (++u->map_keys_seen < MSGPACK_KEY_CACHE_THRESHOLD) {
315+
return unpack_callback_raw(u, b, p, l, o);
316+
}
317+
memset(u->key_cache, 0, sizeof(u->key_cache));
318+
u->key_cache_ready = true;
319+
}
320+
321+
PyObject* cached = *unpack_key_cache_slot(u, p, l);
322+
if (cached && PyUnicode_GET_LENGTH(cached) == (Py_ssize_t)l &&
323+
memcmp(PyUnicode_DATA(cached), p, l) == 0) {
324+
Py_INCREF(cached);
325+
*o = cached;
326+
return 0;
327+
}
328+
return unpack_callback_raw(u, b, p, l, o);
329+
}
330+
265331
static inline int unpack_callback_bin(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
266332
{
267333
if (l > u->max_bin_len) {

‎msgpack/unpack_template.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,13 @@ static inline void unpack_init(unpack_context* ctx)
5858
ctx->execute = NULL;
5959
}
6060

61+
/* Must run before any unpack_clear() on a context whose memory is uninitialized. */
62+
static inline void unpack_init_key_cache(unpack_context* ctx)
63+
{
64+
ctx->user.map_keys_seen = 0;
65+
ctx->user.key_cache_ready = false;
66+
}
67+
6168
static inline PyObject* unpack_take_result(unpack_context* ctx)
6269
{
6370
return unpack_take_ref(&ctx->result);
@@ -70,6 +77,7 @@ static inline void unpack_clear(unpack_context *ctx)
7077
Py_CLEAR(ctx->stack[i].obj);
7178
}
7279
Py_CLEAR(ctx->result);
80+
unpack_key_cache_clear(&ctx->user);
7381
unpack_init(ctx);
7482
}
7583

@@ -293,6 +301,9 @@ static inline int unpack_execute(bool construct, unpack_context* ctx, const char
293301
again_fixed_trail_if_zero(ACS_RAW_VALUE, _msgpack_load32(uint32_t,n), _raw_zero);
294302
case ACS_RAW_VALUE:
295303
_raw_zero:
304+
if (top > 0 && stack[top-1].ct == CT_MAP_KEY) {
305+
push_variable_value(_map_key_raw, data, n, trail);
306+
}
296307
push_variable_value(_raw, data, n, trail);
297308

298309
case ACS_EXT_VALUE:

‎test/test_unpack.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import gc
22
import sys
3+
import uuid
34
import weakref
45
from io import BytesIO
56

@@ -185,3 +186,84 @@ def test_unpack_raises_extra_data_on_stream_with_trailing_bytes():
185186
unpack(stream)
186187
assert exc_info.value.unpacked == 100
187188
assert exc_info.value.extra == packb(200)
189+
190+
191+
# Enough maps with the same keys to exercise repeated-key handling.
192+
REPEATED_KEY_MAPS = [{"id": i, "state": "open", "ключ": i, "x" * 70: i} for i in range(100)]
193+
194+
195+
def test_unpackb_repeated_map_keys_are_equal_and_interned():
196+
result = unpackb(packb(REPEATED_KEY_MAPS))
197+
assert result == REPEATED_KEY_MAPS
198+
for obj in result:
199+
for key in obj:
200+
assert key is sys.intern(key)
201+
202+
203+
def test_unpacker_repeated_map_keys_across_objects():
204+
up = Unpacker()
205+
up.feed(b"".join(packb(obj) for obj in REPEATED_KEY_MAPS))
206+
result = list(up)
207+
assert result == REPEATED_KEY_MAPS
208+
assert all(key is sys.intern(key) for obj in result for key in obj)
209+
210+
211+
def test_unpackb_repeated_map_keys_with_raw():
212+
result = unpackb(packb(REPEATED_KEY_MAPS), raw=True)
213+
assert [list(obj) for obj in result] == [[k.encode() for k in obj] for obj in REPEATED_KEY_MAPS]
214+
215+
216+
def test_unpackb_many_distinct_and_empty_map_keys():
217+
keys = [""] + [f"k{i}" for i in range(600)]
218+
maps = [{key: j for key in keys} for j in range(3)]
219+
result = unpackb(packb(maps))
220+
assert result == maps
221+
assert all(list(obj) == keys for obj in result)
222+
223+
224+
# 100 maps with the key "ab", then one map whose key is `last_key_bytes`.
225+
def _warm_maps_then_key(last_key_bytes):
226+
warm = b"".join(bytes([0x81, 0xA2]) + b"ab" + bytes([i]) for i in range(100))
227+
last = bytes([0x81, 0xA0 | len(last_key_bytes)]) + last_key_bytes + b"\x00"
228+
return bytes([0xDC, 0x00, 101]) + warm + last
229+
230+
231+
def test_unpackb_repeated_map_keys_respect_max_str_len():
232+
packed = _warm_maps_then_key(b"abcd")
233+
with raises(ValueError, match="exceeds max_str_len"):
234+
unpackb(packed, max_str_len=3)
235+
assert unpackb(packed, max_str_len=4)[-1] == {"abcd": 0}
236+
237+
238+
def test_unpackb_repeated_map_keys_then_invalid_utf8():
239+
packed = _warm_maps_then_key(b"a\xff")
240+
with raises(UnicodeDecodeError):
241+
unpackb(packed)
242+
replaced_key = b"a\xff".decode("utf-8", "replace")
243+
assert unpackb(packed, unicode_errors="replace")[-1] == {replaced_key: 0}
244+
245+
246+
@mark.skipif(
247+
Unpacker.__module__ == "msgpack.fallback",
248+
reason="the fallback interns a key before reading its value",
249+
)
250+
def test_unpackb_does_not_intern_key_of_incomplete_pair():
251+
# Interned strings may be immortal, so a key whose value never arrives
252+
# must not be interned. The key is unique per call, as tests run in
253+
# parallel threads and repeat under pytest-run-parallel.
254+
suffix = uuid.uuid4().hex[:28] # "nc-" + 28 chars fits a fixstr
255+
packed = _warm_maps_then_key(f"nc-{suffix}".encode())[:-1] # drop the last value
256+
with raises(ValueError, match="incomplete"):
257+
unpackb(packed)
258+
fresh = "".join(["nc-", suffix])
259+
assert sys.intern(fresh) is fresh
260+
261+
262+
def test_unpacker_repeated_map_keys_after_error():
263+
up = Unpacker()
264+
up.feed(_warm_maps_then_key(b"a\xff"))
265+
with raises(UnicodeDecodeError):
266+
up.unpack()
267+
up.__init__()
268+
up.feed(packb(REPEATED_KEY_MAPS))
269+
assert up.unpack() == REPEATED_KEY_MAPS

0 commit comments

Comments
 (0)