Ruby: Implement Write Barriers (#11793)

Write barrier protected objects are allowed to be promoted to the old generation, which means they only get marked on major GC.

The downside is that the `RB_BJ_WRITE` macro MUST be used to set references, otherwise the referenced object may be garbaged collected.

But the `*Descriptor` classes and `Arena` have very few references and are only set in a few places, so it's relatively easy to implement.

cc @peterzhu2118

Closes #11793

COPYBARA_INTEGRATE_REVIEW=#11793 from casperisfine:descriptor-write-barrier 215e8fa
PiperOrigin-RevId: 511875342

Author: casperisfine

ruby/ext/google/protobuf_c/defs.c

@@ -223,6 +223,8 @@ static void DescriptorPool_register(VALUE module) {
 
 typedef struct {
   const upb_MessageDef* msgdef;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE klass;
   VALUE descriptor_pool;
 } Descriptor;
@@ -238,7 +240,7 @@ static void Descriptor_mark(void* _self) {
 static const rb_data_type_t Descriptor_type = {
     "Google::Protobuf::Descriptor",
     {Descriptor_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static Descriptor* ruby_to_Descriptor(VALUE val) {
@@ -280,7 +282,7 @@ static VALUE Descriptor_initialize(VALUE _self, VALUE cookie,
              "Descriptor objects may not be created from Ruby.");
   }
 
-  self->descriptor_pool = descriptor_pool;
+  RB_OBJ_WRITE(_self, &self->descriptor_pool, descriptor_pool);
   self->msgdef = (const upb_MessageDef*)NUM2ULL(ptr);
 
   return Qnil;
@@ -390,7 +392,7 @@ static VALUE Descriptor_lookup_oneof(VALUE _self, VALUE name) {
 static VALUE Descriptor_msgclass(VALUE _self) {
   Descriptor* self = ruby_to_Descriptor(_self);
   if (self->klass == Qnil) {
-    self->klass = build_class_from_descriptor(_self);
+    RB_OBJ_WRITE(_self, &self->klass, build_class_from_descriptor(_self));
   }
   return self->klass;
 }
@@ -417,6 +419,8 @@ static void Descriptor_register(VALUE module) {
 
 typedef struct {
   const upb_FileDef* filedef;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE descriptor_pool;  // Owns the upb_FileDef.
 } FileDescriptor;
 
@@ -430,7 +434,7 @@ static void FileDescriptor_mark(void* _self) {
 static const rb_data_type_t FileDescriptor_type = {
     "Google::Protobuf::FileDescriptor",
     {FileDescriptor_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static FileDescriptor* ruby_to_FileDescriptor(VALUE val) {
@@ -463,7 +467,7 @@ static VALUE FileDescriptor_initialize(VALUE _self, VALUE cookie,
              "Descriptor objects may not be created from Ruby.");
   }
 
-  self->descriptor_pool = descriptor_pool;
+  RB_OBJ_WRITE(_self, &self->descriptor_pool, descriptor_pool);
   self->filedef = (const upb_FileDef*)NUM2ULL(ptr);
 
   return Qnil;
@@ -519,6 +523,8 @@ static void FileDescriptor_register(VALUE module) {
 
 typedef struct {
   const upb_FieldDef* fielddef;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE descriptor_pool;  // Owns the upb_FieldDef.
 } FieldDescriptor;
 
@@ -532,7 +538,7 @@ static void FieldDescriptor_mark(void* _self) {
 static const rb_data_type_t FieldDescriptor_type = {
     "Google::Protobuf::FieldDescriptor",
     {FieldDescriptor_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static FieldDescriptor* ruby_to_FieldDescriptor(VALUE val) {
@@ -570,7 +576,7 @@ static VALUE FieldDescriptor_initialize(VALUE _self, VALUE cookie,
              "Descriptor objects may not be created from Ruby.");
   }
 
-  self->descriptor_pool = descriptor_pool;
+  RB_OBJ_WRITE(_self, &self->descriptor_pool, descriptor_pool);
   self->fielddef = (const upb_FieldDef*)NUM2ULL(ptr);
 
   return Qnil;
@@ -884,6 +890,8 @@ static void FieldDescriptor_register(VALUE module) {
 
 typedef struct {
   const upb_OneofDef* oneofdef;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE descriptor_pool;  // Owns the upb_OneofDef.
 } OneofDescriptor;
 
@@ -897,7 +905,7 @@ static void OneofDescriptor_mark(void* _self) {
 static const rb_data_type_t OneofDescriptor_type = {
     "Google::Protobuf::OneofDescriptor",
     {OneofDescriptor_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static OneofDescriptor* ruby_to_OneofDescriptor(VALUE val) {
@@ -936,7 +944,7 @@ static VALUE OneofDescriptor_initialize(VALUE _self, VALUE cookie,
              "Descriptor objects may not be created from Ruby.");
   }
 
-  self->descriptor_pool = descriptor_pool;
+  RB_OBJ_WRITE(_self, &self->descriptor_pool, descriptor_pool);
   self->oneofdef = (const upb_OneofDef*)NUM2ULL(ptr);
 
   return Qnil;
@@ -988,6 +996,8 @@ static void OneofDescriptor_register(VALUE module) {
 
 typedef struct {
   const upb_EnumDef* enumdef;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE module;           // begins as nil
   VALUE descriptor_pool;  // Owns the upb_EnumDef.
 } EnumDescriptor;
@@ -1003,7 +1013,7 @@ static void EnumDescriptor_mark(void* _self) {
 static const rb_data_type_t EnumDescriptor_type = {
     "Google::Protobuf::EnumDescriptor",
     {EnumDescriptor_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static EnumDescriptor* ruby_to_EnumDescriptor(VALUE val) {
@@ -1042,7 +1052,7 @@ static VALUE EnumDescriptor_initialize(VALUE _self, VALUE cookie,
              "Descriptor objects may not be created from Ruby.");
   }
 
-  self->descriptor_pool = descriptor_pool;
+  RB_OBJ_WRITE(_self, &self->descriptor_pool, descriptor_pool);
   self->enumdef = (const upb_EnumDef*)NUM2ULL(ptr);
 
   return Qnil;
@@ -1138,7 +1148,7 @@ static VALUE EnumDescriptor_each(VALUE _self) {
 static VALUE EnumDescriptor_enummodule(VALUE _self) {
   EnumDescriptor* self = ruby_to_EnumDescriptor(_self);
   if (self->module == Qnil) {
-    self->module = build_module_from_enumdesc(_self);
+    RB_OBJ_WRITE(_self, &self->module, build_module_from_enumdesc(_self));
   }
   return self->module;
 }

ruby/ext/google/protobuf_c/message.c

@@ -53,6 +53,8 @@ VALUE MessageOrEnum_GetDescriptor(VALUE klass) {
 // -----------------------------------------------------------------------------
 
 typedef struct {
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE arena;
   const upb_Message* msg;  // Can get as mutable when non-frozen.
   const upb_MessageDef*
@@ -65,9 +67,9 @@ static void Message_mark(void* _self) {
 }
 
 static rb_data_type_t Message_type = {
-    "Message",
+    "Google::Protobuf::Message",
     {Message_mark, RUBY_DEFAULT_FREE, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static Message* ruby_to_Message(VALUE msg_rb) {
@@ -105,7 +107,7 @@ upb_Message* Message_GetMutable(VALUE msg_rb, const upb_MessageDef** m) {
 void Message_InitPtr(VALUE self_, upb_Message* msg, VALUE arena) {
   Message* self = ruby_to_Message(self_);
   self->msg = msg;
-  self->arena = arena;
+  RB_OBJ_WRITE(self_, &self->arena, arena);
   ObjectCache_Add(msg, self_);
 }
 

ruby/ext/google/protobuf_c/protobuf.c

@@ -171,6 +171,8 @@ void StringBuilder_PrintMsgval(StringBuilder *b, upb_MessageValue val,
 
 typedef struct {
   upb_Arena *arena;
+  // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+  // macro to update VALUE references, as to trigger write barriers.
   VALUE pinned_objs;
 } Arena;
 
@@ -190,7 +192,7 @@ static VALUE cArena;
 const rb_data_type_t Arena_type = {
     "Google::Protobuf::Internal::Arena",
     {Arena_mark, Arena_free, NULL},
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
 static void* ruby_upb_allocfunc(upb_alloc* alloc, void* ptr, size_t oldsize, size_t size) {
@@ -233,7 +235,7 @@ 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_OBJ_WRITE(_arena, &arena->pinned_objs, rb_ary_new());
   }
   rb_ary_push(arena->pinned_objs, obj);
 }