Workaround the ActiveRecord / Marshal serialization bug on Ruby 3.2 (#24142)

Co-authored-by: Jean Boussier <jean.boussier@gmail.com>

app/controllers/concerns/cache_concern.rb

@@ -3,6 +3,158 @@
 module CacheConcern
   extend ActiveSupport::Concern
 
+  module ActiveRecordCoder
+    EMPTY_HASH = {}.freeze
+
+    class << self
+      def dump(record)
+        instances = InstanceTracker.new
+        serialized_associations = serialize_associations(record, instances)
+        serialized_records = instances.map { |r| serialize_record(r) }
+        [serialized_associations, *serialized_records]
+      end
+
+      def load(payload)
+        instances = InstanceTracker.new
+        serialized_associations, *serialized_records = payload
+        serialized_records.each { |attrs| instances.push(deserialize_record(*attrs)) }
+        deserialize_associations(serialized_associations, instances)
+      end
+
+      private
+
+      # Records without associations, or which have already been visited before,
+      # are serialized by their id alone.
+      #
+      # Records with associations are serialized as a two-element array including
+      # their id and the record's association cache.
+      #
+      def serialize_associations(record, instances)
+        return unless record
+
+        if (id = instances.lookup(record))
+          payload = id
+        else
+          payload = instances.push(record)
+
+          cached_associations = record.class.reflect_on_all_associations.select do |reflection|
+            record.association_cached?(reflection.name)
+          end
+
+          unless cached_associations.empty?
+            serialized_associations = cached_associations.map do |reflection|
+              association = record.association(reflection.name)
+
+              serialized_target = if reflection.collection?
+                                    association.target.map { |target_record| serialize_associations(target_record, instances) }
+                                  else
+                                    serialize_associations(association.target, instances)
+                                  end
+
+              [reflection.name, serialized_target]
+            end
+
+            payload = [payload, serialized_associations]
+          end
+        end
+
+        payload
+      end
+
+      def deserialize_associations(payload, instances)
+        return unless payload
+
+        id, associations = payload
+        record = instances.fetch(id)
+
+        associations&.each do |name, serialized_target|
+          begin
+            association = record.association(name)
+          rescue ActiveRecord::AssociationNotFoundError
+            raise AssociationMissingError, "undefined association: #{name}"
+          end
+
+          target = if association.reflection.collection?
+                     serialized_target.map! { |serialized_record| deserialize_associations(serialized_record, instances) }
+                   else
+                     deserialize_associations(serialized_target, instances)
+                   end
+
+          association.target = target
+        end
+
+        record
+      end
+
+      def serialize_record(record)
+        arguments = [record.class.name, attributes_for_database(record)]
+        arguments << true if record.new_record?
+        arguments
+      end
+
+      if Rails.gem_version >= Gem::Version.new('7.0')
+        def attributes_for_database(record)
+          attributes = record.attributes_for_database
+          attributes.transform_values! { |attr| attr.is_a?(::ActiveModel::Type::Binary::Data) ? attr.to_s : attr }
+          attributes
+        end
+      else
+        def attributes_for_database(record)
+          attributes = record.instance_variable_get(:@attributes).send(:attributes).transform_values(&:value_for_database)
+          attributes.transform_values! { |attr| attr.is_a?(::ActiveModel::Type::Binary::Data) ? attr.to_s : attr }
+          attributes
+        end
+      end
+
+      def deserialize_record(class_name, attributes_from_database, new_record = false) # rubocop:disable Style/OptionalBooleanParameter
+        begin
+          klass = Object.const_get(class_name)
+        rescue NameError
+          raise ClassMissingError, "undefined class: #{class_name}"
+        end
+
+        # Ideally we'd like to call `klass.instantiate`, however it doesn't allow to pass
+        # wether the record was persisted or not.
+        attributes = klass.attributes_builder.build_from_database(attributes_from_database, EMPTY_HASH)
+        klass.allocate.init_with_attributes(attributes, new_record)
+      end
+    end
+
+    class Error < StandardError
+    end
+
+    class ClassMissingError < Error
+    end
+
+    class AssociationMissingError < Error
+    end
+
+    class InstanceTracker
+      def initialize
+        @instances = []
+        @ids = {}.compare_by_identity
+      end
+
+      def map(&block)
+        @instances.map(&block)
+      end
+
+      def fetch(...)
+        @instances.fetch(...)
+      end
+
+      def push(instance)
+        id = @ids[instance] = @instances.size
+        @instances << instance
+        id
+      end
+
+      def lookup(instance)
+        @ids[instance]
+      end
+    end
+  end
+
   def render_with_cache(**options)
     raise ArgumentError, 'only JSON render calls are supported' unless options.key?(:json) || block_given?
 
@@ -34,8 +186,13 @@ module CacheConcern
     raw = raw.cache_ids.to_a if raw.is_a?(ActiveRecord::Relation)
     return [] if raw.empty?
 
-    cached_keys_with_value = Rails.cache.read_multi(*raw).transform_keys(&:id)
-    uncached_ids           = raw.map(&:id) - cached_keys_with_value.keys
+    cached_keys_with_value = begin
+      Rails.cache.read_multi(*raw, namespace: 'v2').transform_keys(&:id).transform_values { |r| ActiveRecordCoder.load(r) }
+    rescue ActiveRecordCoder::Error
+      {} # The serialization format may have changed, let's pretend it's a cache miss.
+    end
+
+    uncached_ids = raw.map(&:id) - cached_keys_with_value.keys
 
     klass.reload_stale_associations!(cached_keys_with_value.values) if klass.respond_to?(:reload_stale_associations!)
 
@@ -43,7 +200,7 @@ module CacheConcern
       uncached = klass.where(id: uncached_ids).with_includes.index_by(&:id)
 
       uncached.each_value do |item|
-        Rails.cache.write(item, item)
+        Rails.cache.write(item, ActiveRecordCoder.dump(item), namespace: 'v2')
       end
     end
 

config/environments/test.rb

@@ -51,6 +51,12 @@ Rails.application.configure do
 
   config.i18n.default_locale = :en
   config.i18n.fallbacks = true
+
+  config.to_prepare do
+    # Force Status to always be SHAPE_TOO_COMPLEX
+    # Ref: https://github.com/mastodon/mastodon/issues/23644
+    10.times { |i| Status.allocate.instance_variable_set(:"@ivar_#{i}", nil) }
+  end
 end
 
 Paperclip::Attachment.default_options[:path] = "#{Rails.root}/spec/test_files/:class/:id_partition/:style.:extension"