Skip to content

Commit

Permalink
JSON.dump: write directly into the provided IO
Browse files Browse the repository at this point in the history
Ref: ruby#524

Rather than to buffer everything in memory.

Unfortunately Ruby doesn't provide an API to write into
and IO without first allocating a string, which is a bit
wasteful.
  • Loading branch information
byroot committed Nov 5, 2024
1 parent bc2c970 commit 11e6ed7
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 43 deletions.
33 changes: 26 additions & 7 deletions ext/json/ext/fbuffer/fbuffer.h
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ typedef struct FBufferStruct {
unsigned long len;
unsigned long capa;
char *ptr;
VALUE io;
} FBuffer;

#define FBUFFER_STACK_SIZE 512
Expand All @@ -66,7 +67,7 @@ static void fbuffer_append_long(FBuffer *fb, long number);
#endif
static inline void fbuffer_append_char(FBuffer *fb, char newchr);
#ifdef JSON_GENERATOR
static VALUE fbuffer_to_s(FBuffer *fb);
static VALUE fbuffer_finalize(FBuffer *fb);
#endif

static void fbuffer_stack_init(FBuffer *fb, unsigned long initial_length, char *stack_buffer, long stack_buffer_size)
Expand All @@ -86,15 +87,26 @@ static void fbuffer_free(FBuffer *fb)
}
}

#ifndef JSON_GENERATOR
static void fbuffer_clear(FBuffer *fb)
{
fb->len = 0;
}
#endif

static void fbuffer_flush(FBuffer *fb)
{
rb_io_write(fb->io, rb_utf8_str_new(fb->ptr, fb->len));
fbuffer_clear(fb);
}

static void fbuffer_do_inc_capa(FBuffer *fb, unsigned long requested)
{
if (RB_UNLIKELY(fb->io)) {
fbuffer_flush(fb);
if (RB_LIKELY(requested < fb->capa)) {
return;
}
}

unsigned long required;

if (RB_UNLIKELY(!fb->ptr)) {
Expand Down Expand Up @@ -174,11 +186,18 @@ static void fbuffer_append_long(FBuffer *fb, long number)
fbuffer_append(fb, buffer_end - len, len);
}

static VALUE fbuffer_to_s(FBuffer *fb)
static VALUE fbuffer_finalize(FBuffer *fb)
{
VALUE result = rb_utf8_str_new(FBUFFER_PTR(fb), FBUFFER_LEN(fb));
fbuffer_free(fb);
return result;
if (fb->io) {
fbuffer_flush(fb);
fbuffer_free(fb);
rb_io_flush(fb->io);
return fb->io;
} else {
VALUE result = rb_utf8_str_new(FBUFFER_PTR(fb), FBUFFER_LEN(fb));
fbuffer_free(fb);
return result;
}
}
#endif
#endif
49 changes: 23 additions & 26 deletions ext/json/ext/generator/generator.c
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ struct generate_json_data {
};

static VALUE cState_from_state_s(VALUE self, VALUE opts);
static VALUE cState_partial_generate(VALUE self, VALUE obj, generator_func);
static VALUE cState_partial_generate(VALUE self, VALUE obj, generator_func, VALUE io);
static void generate_json(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj);
static void generate_json_object(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj);
static void generate_json_array(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj);
Expand Down Expand Up @@ -453,7 +453,7 @@ static VALUE mHash_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_object);
return cState_partial_generate(Vstate, self, generate_json_object, Qfalse);
}

/*
Expand All @@ -467,7 +467,7 @@ static VALUE mHash_to_json(int argc, VALUE *argv, VALUE self)
static VALUE mArray_to_json(int argc, VALUE *argv, VALUE self) {
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_array);
return cState_partial_generate(Vstate, self, generate_json_array, Qfalse);
}

#ifdef RUBY_INTEGER_UNIFICATION
Expand All @@ -480,7 +480,7 @@ static VALUE mInteger_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_integer);
return cState_partial_generate(Vstate, self, generate_json_integer, Qfalse);
}

#else
Expand All @@ -493,7 +493,7 @@ static VALUE mFixnum_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_fixnum);
return cState_partial_generate(Vstate, self, generate_json_fixnum, Qfalse);
}

/*
Expand All @@ -505,7 +505,7 @@ static VALUE mBignum_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_bignum);
return cState_partial_generate(Vstate, self, generate_json_bignum, Qfalse);
}
#endif

Expand All @@ -518,7 +518,7 @@ static VALUE mFloat_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_float);
return cState_partial_generate(Vstate, self, generate_json_float, Qfalse);
}

/*
Expand All @@ -543,7 +543,7 @@ static VALUE mString_to_json(int argc, VALUE *argv, VALUE self)
{
rb_check_arity(argc, 0, 1);
VALUE Vstate = cState_from_state_s(cState, argc == 1 ? argv[0] : Qnil);
return cState_partial_generate(Vstate, self, generate_json_string);
return cState_partial_generate(Vstate, self, generate_json_string, Qfalse);
}

/*
Expand Down Expand Up @@ -638,7 +638,7 @@ static VALUE mObject_to_json(int argc, VALUE *argv, VALUE self)
rb_scan_args(argc, argv, "01", &state);
Check_Type(string, T_STRING);
state = cState_from_state_s(cState, state);
return cState_partial_generate(state, string, generate_json_string);
return cState_partial_generate(state, string, generate_json_string, Qfalse);
}

static void State_mark(void *ptr)
Expand Down Expand Up @@ -1045,12 +1045,14 @@ static VALUE generate_json_rescue(VALUE d, VALUE exc)
return Qundef;
}

static VALUE cState_partial_generate(VALUE self, VALUE obj, generator_func func)
static VALUE cState_partial_generate(VALUE self, VALUE obj, generator_func func, VALUE io)
{
GET_STATE(self);

char stack_buffer[FBUFFER_STACK_SIZE];
FBuffer buffer = {0};
FBuffer buffer = {
.io = RTEST(io) ? io : Qfalse,
};
fbuffer_stack_init(&buffer, state->buffer_initial_length, stack_buffer, FBUFFER_STACK_SIZE);

struct generate_json_data data = {
Expand All @@ -1062,19 +1064,12 @@ static VALUE cState_partial_generate(VALUE self, VALUE obj, generator_func func)
};
rb_rescue(generate_json_try, (VALUE)&data, generate_json_rescue, (VALUE)&data);

return fbuffer_to_s(&buffer);
return fbuffer_finalize(&buffer);
}

/*
* call-seq: generate(obj)
*
* Generates a valid JSON document from object +obj+ and returns the
* result. If no valid JSON document can be created this method raises a
* GeneratorError exception.
*/
static VALUE cState_generate(VALUE self, VALUE obj)
static VALUE cState_generate(VALUE self, VALUE obj, VALUE io)
{
VALUE result = cState_partial_generate(self, obj, generate_json);
VALUE result = cState_partial_generate(self, obj, generate_json, io);
GET_STATE(self);
(void)state;
return result;
Expand Down Expand Up @@ -1502,14 +1497,16 @@ static VALUE cState_configure(VALUE self, VALUE opts)
return self;
}

static VALUE cState_m_generate(VALUE klass, VALUE obj, VALUE opts)
static VALUE cState_m_generate(VALUE klass, VALUE obj, VALUE opts, VALUE io)
{
JSON_Generator_State state = {0};
state_init(&state);
configure_state(&state, opts);

char stack_buffer[FBUFFER_STACK_SIZE];
FBuffer buffer = {0};
FBuffer buffer = {
.io = RTEST(io) ? io : Qfalse,
};
fbuffer_stack_init(&buffer, state.buffer_initial_length, stack_buffer, FBUFFER_STACK_SIZE);

struct generate_json_data data = {
Expand All @@ -1521,7 +1518,7 @@ static VALUE cState_m_generate(VALUE klass, VALUE obj, VALUE opts)
};
rb_rescue(generate_json_try, (VALUE)&data, generate_json_rescue, (VALUE)&data);

return fbuffer_to_s(&buffer);
return fbuffer_finalize(&buffer);
}

/*
Expand Down Expand Up @@ -1583,9 +1580,9 @@ void Init_generator(void)
rb_define_method(cState, "depth=", cState_depth_set, 1);
rb_define_method(cState, "buffer_initial_length", cState_buffer_initial_length, 0);
rb_define_method(cState, "buffer_initial_length=", cState_buffer_initial_length_set, 1);
rb_define_method(cState, "generate", cState_generate, 1);
rb_define_private_method(cState, "_generate", cState_generate, 2);

rb_define_singleton_method(cState, "generate", cState_m_generate, 2);
rb_define_singleton_method(cState, "generate", cState_m_generate, 3);

VALUE mGeneratorMethods = rb_define_module_under(mGenerator, "GeneratorMethods");

Expand Down
17 changes: 7 additions & 10 deletions lib/json/common.rb
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,7 @@ def generate(obj, opts = nil)
if State === opts
opts.generate(obj)
else
State.generate(obj, opts)
State.generate(obj, opts, nil)
end
end

Expand Down Expand Up @@ -793,18 +793,15 @@ def dump(obj, anIO = nil, limit = nil, kwargs = nil)
opts = opts.merge(:max_nesting => limit) if limit
opts = merge_dump_options(opts, **kwargs) if kwargs

result = begin
generate(obj, opts)
begin
if State === opts
opts.generate(obj, anIO)
else
State.generate(obj, opts, anIO)
end
rescue JSON::NestingError
raise ArgumentError, "exceed depth limit"
end

if anIO.nil?
result
else
anIO.write result
anIO
end
end

# Encodes string using String.encode.
Expand Down
11 changes: 11 additions & 0 deletions lib/json/ext/generator/state.rb
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,17 @@ def configure(opts)

alias_method :merge, :configure

# call-seq:
# generate(obj) -> String
# generate(obj, anIO) -> anIO
#
# Generates a valid JSON document from object +obj+ and returns the
# result. If no valid JSON document can be created this method raises a
# GeneratorError exception.
def generate(obj, io = nil)
_generate(obj, io)
end

# call-seq: to_h
#
# Returns the configuration instance variables as a hash, that can be
Expand Down

0 comments on commit 11e6ed7

Please sign in to comment.