Ruby <2.7now uses WeakMap too, which prevents memory leaks. Ruby <2.7 does not allow non-finalizable objects to be WeakMap keys: https://bugs.ruby-lang.org/issues/16035 We work around this by using a secondary map for Ruby <2.7 which maps the non-finalizable integer to a distinct object. For now we accept that the entries in the secondary map wil never be collected. If this becomes a problem we can perform a GC pass every so often that looks at the contents of the object cache to decide what can be deleted from the secondary map.
diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c index 737cd28..c27f30a 100644 --- a/ruby/ext/google/protobuf_c/protobuf.c +++ b/ruby/ext/google/protobuf_c/protobuf.c
@@ -167,30 +167,55 @@ // Arena // ----------------------------------------------------------------------------- -void Arena_free(void* data) { upb_arena_free(data); } +typedef struct { + upb_arena *arena; + VALUE pinned_objs; +} Arena; + +static void Arena_mark(void *data) { + Arena *arena = data; + rb_gc_mark(arena->pinned_objs); +} + +static void Arena_free(void *data) { + Arena *arena = data; + upb_arena_free(arena->arena); +} static VALUE cArena; const rb_data_type_t Arena_type = { "Google::Protobuf::Internal::Arena", - { NULL, Arena_free, NULL }, + { Arena_mark, Arena_free, NULL }, + .flags = RUBY_TYPED_FREE_IMMEDIATELY, }; static VALUE Arena_alloc(VALUE klass) { - upb_arena *arena = upb_arena_new(); + Arena *arena = ALLOC(Arena); + arena->arena = upb_arena_new(); + arena->pinned_objs = Qnil; return TypedData_Wrap_Struct(klass, &Arena_type, arena); } upb_arena *Arena_get(VALUE _arena) { - upb_arena *arena; - TypedData_Get_Struct(_arena, upb_arena, &Arena_type, arena); - return arena; + Arena *arena; + TypedData_Get_Struct(_arena, Arena, &Arena_type, arena); + return arena->arena; } VALUE Arena_new() { return Arena_alloc(cArena); } +void Arena_Pin(VALUE _arena, VALUE obj) { + Arena *arena; + TypedData_Get_Struct(_arena, Arena, &Arena_type, arena); + if (arena->pinned_objs == Qnil) { + arena->pinned_objs = rb_ary_new(); + } + rb_ary_push(arena->pinned_objs, obj); +} + void Arena_register(VALUE module) { VALUE internal = rb_define_module_under(module, "Internal"); VALUE klass = rb_define_class_under(internal, "Arena", rb_cObject); @@ -209,15 +234,37 @@ // different wrapper objects for the same C object, which saves memory and // preserves object identity. // -// We use Hash and/or WeakMap for the cache. WeakMap is faster overall -// (probably due to removal being integrated with GC) but doesn't work for Ruby -// <2.7 (see note below). We need Hash for Ruby <2.7 and for cases where we -// need to GC-root the object (notably when the object has been frozen). +// We use WeakMap for the cache. For Ruby <2.7 we also need a secondary Hash +// to store WeakMap keys because Ruby <2.7 WeakMap doesn't allow non-finalizable +// keys. #if RUBY_API_VERSION_CODE >= 20700 -#define USE_WEAK_MAP 1 +#define USE_SECONDARY_MAP 0 #else -#define USE_WEAK_MAP 0 +#define USE_SECONDARY_MAP 1 +#endif + +#if USE_SECONDARY_MAP + +// Maps Numeric -> Object. The object is then used as a key into the WeakMap. +// This is needed for Ruby <2.7 where a number cannot be a key to WeakMap. +// The object is used only for its identity; it does not contain any data. +VALUE secondary_map = Qnil; + +static void SecondaryMap_Init() { + rb_gc_register_address(&secondary_map); + secondary_map = rb_hash_new(); +} + +static VALUE SecondaryMap_Get(VALUE key) { + VALUE ret = rb_hash_lookup(secondary_map, key); + if (ret == Qnil) { + ret = rb_eval_string("Object.new"); + rb_hash_aset(secondary_map, key, ret); + } + return ret; +} + #endif static VALUE ObjectCache_GetKey(const void* key) { @@ -225,106 +272,41 @@ memcpy(&buf, &key, sizeof(key)); intptr_t key_int = (intptr_t)key; PBRUBY_ASSERT((key_int & 3) == 0); - return LL2NUM(key_int >> 2); -} - -// Strong object cache, uses regular Hash and GC-roots objects. -// - For Ruby <2.7, used for all objects. -// - For Ruby >=2.7, used only for frozen objects, so we preserve the "frozen" -// bit (since this information is not preserved at the upb level). - -VALUE strong_obj_cache = Qnil; - -static void StrongObjectCache_Init() { - rb_gc_register_address(&strong_obj_cache); - strong_obj_cache = rb_hash_new(); -} - -static void StrongObjectCache_Remove(void* key) { - VALUE key_rb = ObjectCache_GetKey(key); - PBRUBY_ASSERT(rb_hash_lookup(strong_obj_cache, key_rb) != Qnil); - rb_hash_delete(strong_obj_cache, key_rb); -} - -static VALUE StrongObjectCache_Get(const void* key) { - VALUE key_rb = ObjectCache_GetKey(key); - return rb_hash_lookup(strong_obj_cache, key_rb); -} - -static void StrongObjectCache_Add(const void* key, VALUE val, - upb_arena* arena) { - PBRUBY_ASSERT(StrongObjectCache_Get(key) == Qnil); - VALUE key_rb = ObjectCache_GetKey(key); - rb_hash_aset(strong_obj_cache, key_rb, val); - upb_arena_addcleanup(arena, (void*)key, StrongObjectCache_Remove); -} - -// Weak object cache. This speeds up the test suite significantly, so we -// presume it speeds up real code also. However we can only use it in Ruby -// >=2.7 due to: -// https://bugs.ruby-lang.org/issues/16035 - -#if USE_WEAK_MAP - -VALUE weak_obj_cache = Qnil; - -static void WeakObjectCache_Init() { - rb_gc_register_address(&weak_obj_cache); - VALUE klass = rb_eval_string("ObjectSpace::WeakMap"); - weak_obj_cache = rb_class_new_instance(0, NULL, klass); -} - -static VALUE WeakObjectCache_Get(const void* key) { - VALUE key_rb = ObjectCache_GetKey(key); - VALUE ret = rb_funcall(weak_obj_cache, rb_intern("[]"), 1, key_rb); + VALUE ret = LL2NUM(key_int >> 2); +#if USE_SECONDARY_MAP + ret = SecondaryMap_Get(ret); +#endif return ret; } -static void WeakObjectCache_Add(const void* key, VALUE val) { - PBRUBY_ASSERT(WeakObjectCache_Get(key) == Qnil); - VALUE key_rb = ObjectCache_GetKey(key); - rb_funcall(weak_obj_cache, rb_intern("[]="), 2, key_rb, val); - PBRUBY_ASSERT(WeakObjectCache_Get(key) == val); -} - -#endif - // Public ObjectCache API. +VALUE weak_obj_cache = Qnil; +ID item_get; +ID item_set; + static void ObjectCache_Init() { - StrongObjectCache_Init(); -#if USE_WEAK_MAP - WeakObjectCache_Init(); + rb_gc_register_address(&weak_obj_cache); + VALUE klass = rb_eval_string("ObjectSpace::WeakMap"); + weak_obj_cache = rb_class_new_instance(0, NULL, klass); + item_get = rb_intern("[]"); + item_set = rb_intern("[]="); +#if USE_SECONDARY_MAP + SecondaryMap_Init(); #endif } -void ObjectCache_Add(const void* key, VALUE val, upb_arena *arena) { -#if USE_WEAK_MAP - (void)arena; - WeakObjectCache_Add(key, val); -#else - StrongObjectCache_Add(key, val, arena); -#endif +void ObjectCache_Add(const void* key, VALUE val) { + PBRUBY_ASSERT(ObjectCache_Get(key) == Qnil); + VALUE key_rb = ObjectCache_GetKey(key); + rb_funcall(weak_obj_cache, item_set, 2, key_rb, val); + PBRUBY_ASSERT(ObjectCache_Get(key) == val); } // Returns the cached object for this key, if any. Otherwise returns Qnil. VALUE ObjectCache_Get(const void* key) { -#if USE_WEAK_MAP - return WeakObjectCache_Get(key); -#else - return StrongObjectCache_Get(key); -#endif -} - -void ObjectCache_Pin(const void* key, VALUE val, upb_arena *arena) { -#if USE_WEAK_MAP - PBRUBY_ASSERT(WeakObjectCache_Get(key) == val); - // This will GC-root the object, but we'll still use the weak map for - // actual lookup. - StrongObjectCache_Add(key, val, arena); -#else - // Value is already pinned, nothing to do. -#endif + VALUE key_rb = ObjectCache_GetKey(key); + return rb_funcall(weak_obj_cache, item_get, 1, key_rb); } /*