Skip to content

Commit 6dcdb39

Browse files
byrootbyroot
authored
Merge pull request #349 from Shopify/better-buffer-stack-fix
Better fix for marking buffers held on the stack
2 parents c3d6367 + e4ca627 commit 6dcdb39

3 files changed

Lines changed: 76 additions & 37 deletions

File tree

‎ext/msgpack/buffer_class.c‎

Lines changed: 72 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
#include "buffer.h"
2222
#include "buffer_class.h"
2323

24-
VALUE cMessagePack_Buffer;
24+
VALUE cMessagePack_Buffer = Qnil;
25+
VALUE cMessagePack_HeldBuffer = Qnil;
2526

2627
static ID s_read;
2728
static ID s_readpartial;
@@ -34,6 +35,73 @@ static VALUE sym_read_reference_threshold;
3435
static VALUE sym_write_reference_threshold;
3536
static VALUE sym_io_buffer_size;
3637

38+
typedef struct msgpack_held_buffer_t msgpack_held_buffer_t;
39+
struct msgpack_held_buffer_t {
40+
size_t size;
41+
VALUE mapped_strings[];
42+
};
43+
44+
static void HeldBuffer_mark(void *data)
45+
{
46+
msgpack_held_buffer_t* held_buffer = (msgpack_held_buffer_t*)data;
47+
for (size_t index = 0; index < held_buffer->size; index++) {
48+
rb_gc_mark(held_buffer->mapped_strings[index]);
49+
}
50+
}
51+
52+
static size_t HeldBuffer_memsize(const void *data)
53+
{
54+
const msgpack_held_buffer_t* held_buffer = (msgpack_held_buffer_t*)data;
55+
return sizeof(size_t) + sizeof(VALUE) * held_buffer->size;
56+
}
57+
58+
static const rb_data_type_t held_buffer_data_type = {
59+
.wrap_struct_name = "msgpack:held_buffer",
60+
.function = {
61+
.dmark = HeldBuffer_mark,
62+
.dfree = RUBY_TYPED_DEFAULT_FREE,
63+
.dsize = HeldBuffer_memsize,
64+
},
65+
.flags = RUBY_TYPED_FREE_IMMEDIATELY
66+
};
67+
68+
VALUE MessagePack_Buffer_hold(msgpack_buffer_t* buffer)
69+
{
70+
size_t mapped_strings_count = 0;
71+
msgpack_buffer_chunk_t* c = buffer->head;
72+
while (c != &buffer->tail) {
73+
if (c->mapped_string != NO_MAPPED_STRING) {
74+
mapped_strings_count++;
75+
}
76+
c = c->next;
77+
}
78+
if (c->mapped_string != NO_MAPPED_STRING) {
79+
mapped_strings_count++;
80+
}
81+
82+
if (mapped_strings_count == 0) {
83+
return Qnil;
84+
}
85+
86+
msgpack_held_buffer_t* held_buffer = xmalloc(sizeof(msgpack_held_buffer_t) + mapped_strings_count * sizeof(VALUE));
87+
88+
c = buffer->head;
89+
mapped_strings_count = 0;
90+
while (c != &buffer->tail) {
91+
if (c->mapped_string != NO_MAPPED_STRING) {
92+
held_buffer->mapped_strings[mapped_strings_count] = c->mapped_string;
93+
mapped_strings_count++;
94+
}
95+
c = c->next;
96+
}
97+
if (c->mapped_string != NO_MAPPED_STRING) {
98+
held_buffer->mapped_strings[mapped_strings_count] = c->mapped_string;
99+
mapped_strings_count++;
100+
}
101+
held_buffer->size = mapped_strings_count;
102+
return TypedData_Wrap_Struct(cMessagePack_HeldBuffer, &held_buffer_data_type, held_buffer);
103+
}
104+
37105

38106
#define CHECK_STRING_TYPE(value) \
39107
value = rb_check_string_type(value); \
@@ -520,6 +588,9 @@ void MessagePack_Buffer_module_init(VALUE mMessagePack)
520588

521589
msgpack_buffer_static_init();
522590

591+
cMessagePack_HeldBuffer = rb_define_class_under(mMessagePack, "HeldBuffer", rb_cBasicObject);
592+
rb_undef_alloc_func(cMessagePack_HeldBuffer);
593+
523594
cMessagePack_Buffer = rb_define_class_under(mMessagePack, "Buffer", rb_cObject);
524595

525596
rb_define_alloc_func(cMessagePack_Buffer, Buffer_alloc);

‎ext/msgpack/buffer_class.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ extern VALUE cMessagePack_Buffer;
2525
void MessagePack_Buffer_module_init(VALUE mMessagePack);
2626

2727
VALUE MessagePack_Buffer_wrap(msgpack_buffer_t* b, VALUE owner);
28+
VALUE MessagePack_Buffer_hold(msgpack_buffer_t* b);
2829

2930
void MessagePack_Buffer_set_options(msgpack_buffer_t* b, VALUE io, VALUE options);
3031

‎ext/msgpack/packer.c‎

Lines changed: 3 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
*/
1818

1919
#include "packer.h"
20+
#include "buffer_class.h"
2021

2122
#if !defined(HAVE_RB_PROC_CALL_WITH_BLOCK)
2223
#define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
@@ -114,39 +115,7 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
114115
}
115116

116117
if(ext_flags & MSGPACK_EXT_RECURSIVE) {
117-
// HACK: While we call the proc, the current pk->buffer won't be reachable
118-
// as it will be stored on the stack.
119-
// To ensure all the `mapped_string` reference in that buffer are properly
120-
// marked and pined, we copy them all on the stack.
121-
VALUE* mapped_strings = NULL;
122-
size_t mapped_strings_count = 0;
123-
msgpack_buffer_chunk_t* c = pk->buffer.head;
124-
while (c != &pk->buffer.tail) {
125-
if (c->mapped_string != NO_MAPPED_STRING) {
126-
mapped_strings_count++;
127-
}
128-
c = c->next;
129-
}
130-
if (c->mapped_string != NO_MAPPED_STRING) {
131-
mapped_strings_count++;
132-
}
133-
134-
if (mapped_strings_count > 0) {
135-
mapped_strings = ALLOCA_N(VALUE, mapped_strings_count);
136-
mapped_strings_count = 0;
137-
c = pk->buffer.head;
138-
while (c != &pk->buffer.tail) {
139-
if (c->mapped_string != NO_MAPPED_STRING) {
140-
mapped_strings[mapped_strings_count] = c->mapped_string;
141-
mapped_strings_count++;
142-
}
143-
c = c->next;
144-
}
145-
if (c->mapped_string != NO_MAPPED_STRING) {
146-
mapped_strings[mapped_strings_count] = c->mapped_string;
147-
mapped_strings_count++;
148-
}
149-
}
118+
VALUE held_buffer = MessagePack_Buffer_hold(&pk->buffer);
150119

151120
msgpack_buffer_t parent_buffer = pk->buffer;
152121
msgpack_buffer_init(PACKER_BUFFER_(pk));
@@ -167,9 +136,7 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
167136
msgpack_packer_write_ext(pk, ext_type, payload);
168137
}
169138

170-
if (mapped_strings_count > 0) {
171-
RB_GC_GUARD(mapped_strings[0]);
172-
}
139+
RB_GC_GUARD(held_buffer);
173140
} else {
174141
VALUE payload = rb_proc_call_with_block(proc, 1, &v, Qnil);
175142
StringValue(payload);

0 commit comments

Comments
 (0)