Merge pull request rails#47770 from jonathanhefner/message_pack

Add `ActiveSupport::MessagePack`

Author: jonathanhefner

Gemfile

@@ -68,6 +68,7 @@ gem "listen", "~> 3.3", require: false
 gem "libxml-ruby", platforms: :ruby
 gem "connection_pool", require: false
 gem "rexml", require: false
+gem "msgpack", ">= 1.7.0", require: false
 
 # for railties
 gem "bootsnap", ">= 1.4.4", require: false

Gemfile.lock

@@ -320,7 +320,7 @@ GEM
     mixlib-shellout (3.2.7)
       chef-utils
     mono_logger (1.1.1)
-    msgpack (1.6.0)
+    msgpack (1.7.0)
     multi_json (1.15.0)
     multipart-post (2.2.3)
     mustermann (3.0.0)
@@ -580,6 +580,7 @@ DEPENDENCIES
   minitest-bisect
   minitest-ci
   minitest-retry
+  msgpack (>= 1.7.0)
   mysql2 (~> 0.5)
   nokogiri (>= 1.8.1, != 1.11.0)
   pg (~> 1.3)

activerecord/lib/active_record/message_pack.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module MessagePack # :nodoc:
+    FORMAT_VERSION = 1
+
+    class << self
+      def dump(input)
+        encoder = Encoder.new
+        [FORMAT_VERSION, encoder.encode(input), encoder.entries]
+      end
+
+      def load(dumped)
+        format_version, top_level, entries = dumped
+        unless format_version == FORMAT_VERSION
+          raise "Invalid format version: #{format_version.inspect}"
+        end
+        Decoder.new(entries).decode(top_level)
+      end
+    end
+
+    module Extensions
+      extend self
+
+      def install(registry)
+        registry.register_type 119, ActiveModel::Type::Binary::Data,
+          packer: :to_s,
+          unpacker: :new
+
+        registry.register_type 120, ActiveRecord::Base,
+          packer: method(:write_record),
+          unpacker: method(:read_record),
+          recursive: true
+      end
+
+      def write_record(record, packer)
+        packer.write(ActiveRecord::MessagePack.dump(record))
+      end
+
+      def read_record(unpacker)
+        ActiveRecord::MessagePack.load(unpacker.read)
+      end
+    end
+
+    class Encoder
+      attr_reader :entries
+
+      def initialize
+        @entries = []
+        @refs = {}.compare_by_identity
+      end
+
+      def encode(input)
+        if input.is_a?(Array)
+          input.map { |record| encode_record(record) }
+        elsif input
+          encode_record(input)
+        end
+      end
+
+      def encode_record(record)
+        ref = @refs[record]
+
+        if !ref
+          ref = @refs[record] = @entries.size
+          @entries << build_entry(record)
+          add_cached_associations(record, @entries.last)
+        end
+
+        ref
+      end
+
+      def build_entry(record)
+        [
+          ActiveSupport::MessagePack::Extensions.dump_class(record.class),
+          record.attributes_for_database,
+          record.new_record?
+        ]
+      end
+
+      def add_cached_associations(record, entry)
+        record.class.reflections.each_value do |reflection|
+          if record.association_cached?(reflection.name)
+            entry << reflection.name << encode(record.association(reflection.name).target)
+          end
+        end
+      end
+    end
+
+    class Decoder
+      def initialize(entries)
+        @records = entries.map { |entry| build_record(entry) }
+        @records.zip(entries) { |record, entry| resolve_cached_associations(record, entry) }
+      end
+
+      def decode(ref)
+        if ref.is_a?(Array)
+          ref.map { |r| @records[r] }
+        elsif ref
+          @records[ref]
+        end
+      end
+
+      def build_record(entry)
+        class_name, attributes_hash, is_new_record, * = entry
+        klass = ActiveSupport::MessagePack::Extensions.load_class(class_name)
+        attributes = klass.attributes_builder.build_from_database(attributes_hash)
+        klass.allocate.init_with_attributes(attributes, is_new_record)
+      end
+
+      def resolve_cached_associations(record, entry)
+        i = 3 # entry == [class_name, attributes_hash, is_new_record, *associations]
+        while i < entry.length
+          begin
+            record.association(entry[i]).target = decode(entry[i + 1])
+          rescue ActiveRecord::AssociationNotFoundError
+            # The association no longer exists, so just skip it.
+          end
+          i += 2
+        end
+      end
+    end
+  end
+end

activerecord/lib/active_record/railtie.rb

@@ -424,5 +424,14 @@ class Railtie < Rails::Railtie # :nodoc:
         end
       end
     end
+
+    initializer "active_record.message_pack" do
+      ActiveSupport.on_load(:message_pack) do
+        ActiveSupport.on_load(:active_record) do
+          require "active_record/message_pack"
+          ActiveRecord::MessagePack::Extensions.install(ActiveSupport::MessagePack::CacheSerializer)
+        end
+      end
+    end
   end
 end

activerecord/test/cases/message_pack_test.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/author"
+require "models/binary"
+require "models/comment"
+require "models/post"
+require "active_support/message_pack"
+require "active_record/message_pack"
+
+class ActiveRecordMessagePackTest < ActiveRecord::TestCase
+  test "enshrines type IDs" do
+    expected = {
+      119 => ActiveModel::Type::Binary::Data,
+      120 => ActiveRecord::Base,
+    }
+
+    factory = ::MessagePack::Factory.new
+    ActiveRecord::MessagePack::Extensions.install(factory)
+    actual = factory.registered_types.to_h do |entry|
+      [entry[:type], entry[:class]]
+    end
+
+    assert_equal expected, actual
+  end
+
+  test "roundtrips record and cached associations" do
+    post = Post.create!(title: "A Title", body: "A body.")
+    post.create_author!(name: "An Author")
+    post.comments.create!(body: "A comment.")
+    post.comments.create!(body: "Another comment.", author: post.author)
+    post.comments.load
+
+    assert_no_queries do
+      roundtripped_post = roundtrip(post)
+
+      assert_equal post, roundtripped_post
+      assert_equal post.author, roundtripped_post.author
+      assert_equal post.comments.to_a, roundtripped_post.comments.to_a
+      assert_equal post.comments.map(&:author), roundtripped_post.comments.map(&:author)
+
+      assert_same roundtripped_post, roundtripped_post.comments[0].post
+      assert_same roundtripped_post, roundtripped_post.comments[1].post
+      assert_same roundtripped_post.author, roundtripped_post.comments[1].author
+    end
+  end
+
+  test "roundtrips new_record? status" do
+    post = Post.new(title: "A Title", body: "A body.")
+    post.create_author!(name: "An Author")
+
+    assert_no_queries do
+      roundtripped_post = roundtrip(post)
+
+      assert_equal post.attributes, roundtripped_post.attributes
+      assert_equal post.new_record?, roundtripped_post.new_record?
+      assert_equal post.author, roundtripped_post.author
+      assert_equal post.author.new_record?, roundtripped_post.author.new_record?
+    end
+  end
+
+  test "roundtrips binary attribute" do
+    binary = Binary.new(data: Marshal.dump("data"))
+    assert_equal binary.attributes, roundtrip(binary).attributes
+  end
+
+  test "raises ActiveSupport::MessagePack::MissingClassError if record class no longer exists" do
+    klass = Class.new(Post)
+    def klass.name; "SomeLegacyClass"; end
+    dumped = serializer.dump(klass.new(title: "A Title", body: "A body."))
+
+    assert_raises ActiveSupport::MessagePack::MissingClassError do
+      serializer.load(dumped)
+    end
+  end
+
+  private
+    def serializer
+      @serializer ||= ::MessagePack::Factory.new.tap do |factory|
+        ActiveRecord::MessagePack::Extensions.install(factory)
+        ActiveSupport::MessagePack::Extensions.install(factory)
+        ActiveSupport::MessagePack::Extensions.install_unregistered_type_error(factory)
+      end
+    end
+
+    def roundtrip(input)
+      serializer.load(serializer.dump(input))
+    end
+end

activesupport/lib/active_support/message_pack.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+begin
+  gem "msgpack", ">= 1.7.0"
+  require "msgpack"
+rescue LoadError => error
+  warn "ActiveSupport::MessagePack requires the msgpack gem, version 1.7.0 or later. " \
+    "Please add it to your Gemfile: `gem \"msgpack\", \">= 1.7.0\"`"
+  raise error
+end
+
+require_relative "message_pack/cache_serializer"
+require_relative "message_pack/serializer"
+
+module ActiveSupport
+  module MessagePack
+    extend Serializer
+
+    ##
+    # :singleton-method: dump
+    # :call-seq: dump(object)
+    #
+    # Dumps an object. Raises ActiveSupport::MessagePack::UnserializableObjectError
+    # if the object type is not supported.
+    #
+    #--
+    # Implemented by Serializer#dump.
+
+    ##
+    # :singleton-method: load
+    # :call-seq: load(dumped)
+    #
+    # Loads an object dump created by ::dump.
+    #
+    #--
+    # Implemented by Serializer#load.
+
+    ##
+    # :singleton-method: signature?
+    # :call-seq: signature?(dumped)
+    #
+    # Returns true if the given dump begins with an +ActiveSupport::MessagePack+
+    # signature.
+    #
+    #--
+    # Implemented by Serializer#signature?.
+
+    ActiveSupport.run_load_hooks(:message_pack, self)
+  end
+end

activesupport/lib/active_support/message_pack/cache_serializer.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative "serializer"
+
+module ActiveSupport
+  module MessagePack
+    module CacheSerializer
+      include Serializer
+      extend self
+
+      ZLIB_HEADER = "\x78"
+
+      def dump(entry)
+        super(entry.pack)
+      end
+
+      def dump_compressed(entry, threshold) # :nodoc:
+        dumped = dump(entry)
+        if dumped.bytesize >= threshold
+          compressed = Zlib::Deflate.deflate(dumped)
+          compressed.bytesize < dumped.bytesize ? compressed : dumped
+        else
+          dumped
+        end
+      end
+
+      def load(dumped)
+        dumped = Zlib::Inflate.inflate(dumped) if compressed?(dumped)
+        ActiveSupport::Cache::Entry.unpack(super)
+      rescue ActiveSupport::MessagePack::MissingClassError
+        # Treat missing class as cache miss => return nil
+      end
+
+      private
+        def compressed?(dumped)
+          dumped.start_with?(ZLIB_HEADER)
+        end
+
+        def install_unregistered_type_handler
+          Extensions.install_unregistered_type_fallback(message_pack_factory)
+        end
+    end
+  end
+end