Addressed code-review comments.

2 files changed, 26 insertions, 12 deletions

diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index 3a441be3..b9ecd456 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -349,8 +349,6 @@ static void *oneofsubmsg_handler(void *closure,
   MessageHeader* msg = closure;
   const oneof_handlerdata_t *oneofdata = hd;
   uint32_t oldcase = DEREF(msg, oneofdata->case_ofs, uint32_t);
-  DEREF(msg, oneofdata->case_ofs, uint32_t) =
-      oneofdata->oneof_case_num;
 
   VALUE subdesc =
       get_def_obj((void*)oneofdata->md);
@@ -361,6 +359,11 @@ static void *oneofsubmsg_handler(void *closure,
     DEREF(msg, oneofdata->ofs, VALUE) =
         rb_class_new_instance(0, NULL, subklass);
   }
+  // Set the oneof case *after* allocating the new class instance -- see comment
+  // in layout_set() as to why. There are subtle interactions with possible GC
+  // points and oneof field type transitions.
+  DEREF(msg, oneofdata->case_ofs, uint32_t) =
+      oneofdata->oneof_case_num;
 
   VALUE submsg_rb = DEREF(msg, oneofdata->ofs, VALUE);
   MessageHeader* submsg;
@@ -965,11 +968,11 @@ static void putmsg(VALUE msg_rb, const Descriptor* desc,
     uint32_t offset =
         desc->layout->fields[upb_fielddef_index(f)].offset +
         sizeof(MessageHeader);
-    uint32_t oneof_case_offset =
-        desc->layout->fields[upb_fielddef_index(f)].case_offset +
-        sizeof(MessageHeader);
 
     if (upb_fielddef_containingoneof(f)) {
+      uint32_t oneof_case_offset =
+          desc->layout->fields[upb_fielddef_index(f)].case_offset +
+          sizeof(MessageHeader);
       // For a oneof, check that this field is actually present -- skip all the
       // below if not.
       if (DEREF(msg, oneof_case_offset, uint32_t) !=
diff --git a/ruby/ext/google/protobuf_c/storage.c b/ruby/ext/google/protobuf_c/storage.c
index fac06e0a..662101f8 100644
--- a/ruby/ext/google/protobuf_c/storage.c
+++ b/ruby/ext/google/protobuf_c/storage.c
@@ -579,15 +579,26 @@ void layout_set(MessageLayout* layout,
       *oneof_case = 0;
       memset(memory, 0, NATIVE_SLOT_MAX_SIZE);
     } else {
-      // Set the oneof case *first* in case a GC is triggered during
-      // native_slot_set(): layout_mark() depends on oneof_case to know whether
-      // the slot may be a Ruby VALUE and so we need that lifetime to start
-      // before we could possibly stick a VALUE in it.
-      *oneof_case = upb_fielddef_number(field);
-      // We just overwrite the value directly if we changed oneof cases:
-      // native_slot_set() does not depend on the old value in memory.
+      // The transition between field types for a single oneof (union) slot is
+      // somewhat complex because we need to ensure that a GC triggered at any
+      // point by a call into the Ruby VM sees a valid state for this field and
+      // does not either go off into the weeds (following what it thinks is a
+      // VALUE but is actually a different field type) or miss an object (seeing
+      // what it thinks is a primitive field but is actually a VALUE for the new
+      // field type).
+      //
+      // native_slot_set() has two parts: (i) conversion of some sort, and (ii)
+      // setting the in-memory content to the new value. It guarantees that all
+      // calls to the Ruby VM are completed before the memory slot is altered.
+      //
+      // In order for the transition to be safe, the oneof case slot must be in
+      // sync with the value slot whenever the Ruby VM has been called. Because
+      // we are guaranteed that no Ruby VM calls occur after native_slot_set()
+      // alters the memory slot and before it returns, we set the oneof case
+      // immediately after native_slot_set() returns.
       native_slot_set(upb_fielddef_type(field), field_type_class(field),
                       memory, val);
+      *oneof_case = upb_fielddef_number(field);
     }
   } else if (is_map_field(field)) {
     check_map_field_type(val, field);