Merge pull request #347 from Shopify/341-store-mapped-strings-on-stack

Packer: store the parent buffer mapped strings on the stack

Author: byroot

.github/workflows/ci.yaml

@@ -52,7 +52,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu]
-        ruby: ['jruby-9.2', 'jruby-9.3', 'jruby-9.4', 'truffleruby']
+        ruby: ['jruby-9.3', 'jruby-9.4', 'truffleruby']
     runs-on: ${{ matrix.os }}-latest
     steps:
       - uses: actions/checkout@v2

ChangeLog

@@ -1,3 +1,4 @@
+* Fix a potential GC bug when packing data using recursive extensions and buffers containing over 512KkiB of data (See #341).
 * Fix a regression where feeding an empty string to an Unpacker would be considered like the end of the buffer.
 
 2023-05-19 1.7.1:

bin/rspec

@@ -0,0 +1,27 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+#
+# This file was generated by Bundler.
+#
+# The application 'rspec' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+ENV["BUNDLE_GEMFILE"] ||= File.expand_path("../Gemfile", __dir__)
+
+bundle_binstub = File.expand_path("bundle", __dir__)
+
+if File.file?(bundle_binstub)
+  if File.read(bundle_binstub, 300).include?("This file was generated by Bundler")
+    load(bundle_binstub)
+  else
+    abort("Your `bin/bundle` was not generated by Bundler, so this binstub cannot run.
+Replace `bin/bundle` by running `bundle binstubs bundler --force`, then run this command again.")
+  end
+end
+
+require "rubygems"
+require "bundler/setup"
+
+load Gem.bin_path("rspec-core", "rspec")

ext/msgpack/packer.c

@@ -114,6 +114,40 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
     }
 
     if(ext_flags & MSGPACK_EXT_RECURSIVE) {
+        // HACK: While we call the proc, the current pk->buffer won't be reachable
+        // as it will be stored on the stack.
+        // To ensure all the `mapped_string` reference in that buffer are properly
+        // marked and pined, we copy them all on the stack.
+        VALUE* mapped_strings = NULL;
+        size_t mapped_strings_count = 0;
+        msgpack_buffer_chunk_t* c = pk->buffer.head;
+        while (c != &pk->buffer.tail) {
+            if (c->mapped_string != NO_MAPPED_STRING) {
+                mapped_strings_count++;
+            }
+            c = c->next;
+        }
+        if (c->mapped_string != NO_MAPPED_STRING) {
+            mapped_strings_count++;
+        }
+
+        if (mapped_strings_count > 0) {
+            mapped_strings = ALLOCA_N(VALUE, mapped_strings_count);
+            mapped_strings_count = 0;
+            c = pk->buffer.head;
+            while (c != &pk->buffer.tail) {
+                if (c->mapped_string != NO_MAPPED_STRING) {
+                    mapped_strings[mapped_strings_count] = c->mapped_string;
+                    mapped_strings_count++;
+                }
+                c = c->next;
+            }
+            if (c->mapped_string != NO_MAPPED_STRING) {
+                mapped_strings[mapped_strings_count] = c->mapped_string;
+                mapped_strings_count++;
+            }
+        }
+
         msgpack_buffer_t parent_buffer = pk->buffer;
         msgpack_buffer_init(PACKER_BUFFER_(pk));
 
@@ -132,6 +166,10 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
             pk->buffer = parent_buffer;
             msgpack_packer_write_ext(pk, ext_type, payload);
         }
+
+        if (mapped_strings_count > 0) {
+            RB_GC_GUARD(mapped_strings[0]);
+        }
     } else {
         VALUE payload = rb_proc_call_with_block(proc, 1, &v, Qnil);
         StringValue(payload);

spec/packer_spec.rb

@@ -589,4 +589,31 @@ def to_msgpack_ext
       GC.stress = stress
     end
   end
+
+  it "doesn't crash when using recursive packers concurrently" do
+    hash_with_indifferent_access = Class.new(Hash)
+    msgpack = MessagePack::Factory.new
+    msgpack.register_type(
+      0x02,
+      hash_with_indifferent_access,
+      packer: ->(value, packer) do
+        packer.write("a".b * 600_0000) # Over MSGPACK_BUFFER_STRING_WRITE_REFERENCE_DEFAULT
+        packer.write(value.to_h)
+        GC.start(full_mark: true, immediate_mark: true, immediate_sweep: true)
+      end,
+      unpacker: ->(unpacker) { hash_with_indifferent_access.new(unpacker.read) },
+      recursive: true
+    )
+
+    packer = msgpack.packer
+
+    top_hash = hash = hash_with_indifferent_access.new
+    10.times do
+      hash["a"] = new_hash = hash_with_indifferent_access.new
+      hash = new_hash
+    end
+    packer.write(top_hash)
+    packer.write(top_hash)
+    packer.full_pack
+  end
 end