Implement a JSON compilation cache

lib/bootsnap.rb

@@ -43,7 +43,8 @@ def self.setup(
     autoload_paths_cache: nil,
     disable_trace: nil,
     compile_cache_iseq: true,
-    compile_cache_yaml: true
+    compile_cache_yaml: true,
+    compile_cache_json: true
   )
     unless autoload_paths_cache.nil?
       warn "[DEPRECATED] Bootsnap's `autoload_paths_cache:` option is deprecated and will be removed. " \
@@ -69,7 +70,8 @@ def self.setup(
     Bootsnap::CompileCache.setup(
       cache_dir: cache_dir + '/bootsnap/compile-cache',
       iseq: compile_cache_iseq,
-      yaml: compile_cache_yaml
+      yaml: compile_cache_yaml,
+      json: compile_cache_json,
     )
   end
 
@@ -113,6 +115,7 @@ def self.default_setup
         load_path_cache:      !ENV['DISABLE_BOOTSNAP_LOAD_PATH_CACHE'],
         compile_cache_iseq:   !ENV['DISABLE_BOOTSNAP_COMPILE_CACHE'] && iseq_cache_supported?,
         compile_cache_yaml:   !ENV['DISABLE_BOOTSNAP_COMPILE_CACHE'],
+        compile_cache_json:   !ENV['DISABLE_BOOTSNAP_COMPILE_CACHE'],
       )
 
       if ENV['BOOTSNAP_LOG']

lib/bootsnap/cli.rb

@@ -21,7 +21,7 @@ def match?(string)
 
     attr_reader :cache_dir, :argv
 
-    attr_accessor :compile_gemfile, :exclude, :verbose, :iseq, :yaml, :jobs
+    attr_accessor :compile_gemfile, :exclude, :verbose, :iseq, :yaml, :json, :jobs
 
     def initialize(argv)
       @argv = argv
@@ -32,37 +32,44 @@ def initialize(argv)
       self.jobs = Etc.nprocessors
       self.iseq = true
       self.yaml = true
+      self.json = true
     end
 
     def precompile_command(*sources)
       require 'bootsnap/compile_cache/iseq'
       require 'bootsnap/compile_cache/yaml'
+      require 'bootsnap/compile_cache/json'
 
       fix_default_encoding do
         Bootsnap::CompileCache::ISeq.cache_dir = self.cache_dir
         Bootsnap::CompileCache::YAML.init!
         Bootsnap::CompileCache::YAML.cache_dir = self.cache_dir
+        Bootsnap::CompileCache::JSON.init!
+        Bootsnap::CompileCache::JSON.cache_dir = self.cache_dir
 
         @work_pool = WorkerPool.create(size: jobs, jobs: {
           ruby: method(:precompile_ruby),
           yaml: method(:precompile_yaml),
+          json: method(:precompile_json),
         })
         @work_pool.spawn
 
         main_sources = sources.map { |d| File.expand_path(d) }
         precompile_ruby_files(main_sources)
         precompile_yaml_files(main_sources)
+        precompile_json_files(main_sources)
 
         if compile_gemfile
           # Some gems embed their tests, they're very unlikely to be loaded, so not worth precompiling.
           gem_exclude = Regexp.union([exclude, '/spec/', '/test/'].compact)
           precompile_ruby_files($LOAD_PATH.map { |d| File.expand_path(d) }, exclude: gem_exclude)
 
-          # Gems that include YAML files usually don't put them in `lib/`.
+          # Gems that include JSON or YAML files usually don't put them in `lib/`.
           # So we look at the gem root.
           gem_pattern = %r{^#{Regexp.escape(Bundler.bundle_path.to_s)}/?(?:bundler/)?gems\/[^/]+}
           gem_paths = $LOAD_PATH.map { |p| p[gem_pattern] }.compact.uniq
           precompile_yaml_files(gem_paths, exclude: gem_exclude)
+          precompile_json_files(gem_paths, exclude: gem_exclude)
         end
 
         if exitstatus = @work_pool.shutdown
@@ -137,6 +144,29 @@ def precompile_yaml(*yaml_files)
       end
     end
 
+    def precompile_json_files(load_paths, exclude: self.exclude)
+      return unless json
+
+      load_paths.each do |path|
+        if !exclude || !exclude.match?(path)
+          list_files(path, '**/*.json').each do |json_file|
+            # We ignore hidden files to not match the various .config.json files
+            if !File.basename(json_file).start_with?('.') && (!exclude || !exclude.match?(json_file))
+              @work_pool.push(:json, json_file)
+            end
+          end
+        end
+      end
+    end
+
+    def precompile_json(*json_files)
+      Array(json_files).each do |json_file|
+        if p(CompileCache::JSON.precompile(json_file, cache_dir: cache_dir))
+          STDERR.puts(json_file) if verbose
+        end
+      end
+    end
+
     def precompile_ruby_files(load_paths, exclude: self.exclude)
       return unless iseq
 
@@ -240,6 +270,11 @@ def parser
           Disable YAML precompilation.
         EOS
         opts.on('--no-yaml', help) { self.yaml = false }
+
+        help = <<~EOS
+          Disable JSON precompilation.
+        EOS
+        opts.on('--no-json', help) { self.json = false }
       end
     end
   end

lib/bootsnap/compile_cache.rb

@@ -4,7 +4,7 @@ module CompileCache
     Error           = Class.new(StandardError)
     PermissionError = Class.new(Error)
 
-    def self.setup(cache_dir:, iseq:, yaml:)
+    def self.setup(cache_dir:, iseq:, yaml:, json:)
       if iseq
         if supported?
           require_relative('compile_cache/iseq')
@@ -22,6 +22,15 @@ def self.setup(cache_dir:, iseq:, yaml:)
           warn("[bootsnap/setup] YAML parsing caching is not supported on this implementation of Ruby")
         end
       end
+
+      if json
+        if supported?
+          require_relative('compile_cache/json')
+          Bootsnap::CompileCache::JSON.install!(cache_dir)
+        elsif $VERBOSE
+          warn("[bootsnap/setup] YAML parsing caching is not supported on this implementation of Ruby")
+        end
+      end
     end
 
     def self.permission_error(path)

lib/bootsnap/compile_cache/json.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+require('bootsnap/bootsnap')
+
+module Bootsnap
+  module CompileCache
+    module JSON
+      class << self
+        attr_accessor(:msgpack_factory, :cache_dir, :supported_options)
+
+        def input_to_storage(payload, _)
+          obj = ::JSON.parse(payload)
+          msgpack_factory.dump(obj)
+        end
+
+        def storage_to_output(data, kwargs)
+          if kwargs && kwargs.key?(:symbolize_names)
+            kwargs[:symbolize_keys] = kwargs.delete(:symbolize_names)
+          end
+          msgpack_factory.load(data, kwargs)
+        end
+
+        def input_to_output(data, kwargs)
+          ::JSON.parse(data, **(kwargs || {}))
+        end
+
+        def precompile(path, cache_dir: self.cache_dir)
+          Bootsnap::CompileCache::Native.precompile(
+            cache_dir,
+            path.to_s,
+            self,
+          )
+        end
+
+        def install!(cache_dir)
+          self.cache_dir = cache_dir
+          init!
+          if ::JSON.respond_to?(:load_file)
+            ::JSON.singleton_class.prepend(Patch)
+          end
+        end
+
+        def init!
+          require('json')
+          require('msgpack')
+
+          self.msgpack_factory = MessagePack::Factory.new
+          self.supported_options = [:symbolize_names]
+          if ::JSON.parse('["foo"]', freeze: true).first.frozen?
+            self.supported_options = [:freeze]
+          end
+          self.supported_options.freeze
+        end
+      end
+
+      module Patch
+        def load_file(path, *args)
+          return super if args.size > 1
+          if kwargs = args.first
+            return super unless kwargs.is_a?(Hash)
+            return super unless (kwargs.keys - ::Bootsnap::CompileCache::JSON.supported_options).empty?
+          end
+
+          begin
+            ::Bootsnap::CompileCache::Native.fetch(
+              Bootsnap::CompileCache::JSON.cache_dir,
+              File.realpath(path),
+              ::Bootsnap::CompileCache::JSON,
+              kwargs,
+            )
+          rescue Errno::EACCES
+            ::Bootsnap::CompileCache.permission_error(path)
+          end
+        end
+
+        ruby2_keywords :load_file if respond_to?(:ruby2_keywords, true)
+      end
+    end
+  end
+end

test/compile_cache/json_test.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+require('test_helper')
+
+class CompileCacheJSONTest < Minitest::Test
+  include(TmpdirHelper)
+
+  module FakeJson
+    Fallback = Class.new(StandardError)
+    class << self
+      def load_file(path, symbolize_names: false, freeze: false, fallback: nil)
+        raise Fallback
+      end
+    end
+  end
+
+  def setup
+    super
+    Bootsnap::CompileCache::JSON.init!
+    FakeJson.singleton_class.prepend(Bootsnap::CompileCache::JSON::Patch)
+  end
+
+  def test_json_input_to_output
+    document = ::Bootsnap::CompileCache::JSON.input_to_output(<<~JSON, {})
+      {
+        "foo": 42,
+        "bar": [1]
+      }
+    JSON
+    expected = {
+      'foo' => 42,
+      'bar' => [1],
+    }
+    assert_equal expected, document
+  end
+
+  def test_load_file
+    Help.set_file('a.json', '{"foo": "bar"}', 100)
+    assert_equal({'foo' => 'bar'}, FakeJson.load_file('a.json'))
+  end
+
+  def test_load_file_symbolize_names
+    Help.set_file('a.json', '{"foo": "bar"}', 100)
+    FakeJson.load_file('a.json')
+
+    if ::Bootsnap::CompileCache::JSON.supported_options.include?(:symbolize_names)
+      2.times do
+        assert_equal({foo: 'bar'}, FakeJson.load_file('a.json', symbolize_names: true))
+      end
+    else
+      assert_raises(FakeJson::Fallback) do # would call super
+        FakeJson.load_file('a.json', symbolize_names: true)
+      end
+    end
+  end
+
+  def test_load_file_freeze
+    Help.set_file('a.json', '["foo"]', 100)
+    FakeJson.load_file('a.json')
+
+    if ::Bootsnap::CompileCache::JSON.supported_options.include?(:freeze)
+      2.times do
+        string = FakeJson.load_file('a.json', freeze: true).first
+        assert_equal("foo", string)
+        assert_predicate(string, :frozen?)
+      end
+    else
+      assert_raises(FakeJson::Fallback) do # would call super
+        FakeJson.load_file('a.json', freeze: true)
+      end
+    end
+  end
+
+  def test_load_file_unknown_option
+    Help.set_file('a.json', '["foo"]', 100)
+    FakeJson.load_file('a.json')
+
+    assert_raises(FakeJson::Fallback) do # would call super
+      FakeJson.load_file('a.json', fallback: true)
+    end
+  end
+end

test/setup_test.rb

@@ -20,6 +20,7 @@ def test_default_setup
         load_path_cache: true,
         compile_cache_iseq: Bootsnap.iseq_cache_supported?,
         compile_cache_yaml: true,
+        compile_cache_json: true,
       )
 
       Bootsnap.default_setup
@@ -34,6 +35,7 @@ def test_default_setup_with_ENV_not_dev
         load_path_cache: true,
         compile_cache_iseq: Bootsnap.iseq_cache_supported?,
         compile_cache_yaml: true,
+        compile_cache_json: true,
       )
 
       Bootsnap.default_setup
@@ -48,6 +50,7 @@ def test_default_setup_with_DISABLE_BOOTSNAP_LOAD_PATH_CACHE
         load_path_cache: false,
         compile_cache_iseq: Bootsnap.iseq_cache_supported?,
         compile_cache_yaml: true,
+        compile_cache_json: true,
       )
 
       Bootsnap.default_setup
@@ -62,6 +65,7 @@ def test_default_setup_with_DISABLE_BOOTSNAP_COMPILE_CACHE
         load_path_cache: true,
         compile_cache_iseq: false,
         compile_cache_yaml: false,
+        compile_cache_json: false,
       )
 
       Bootsnap.default_setup
@@ -83,6 +87,7 @@ def test_default_setup_with_BOOTSNAP_LOG
         load_path_cache: true,
         compile_cache_iseq: Bootsnap.iseq_cache_supported?,
         compile_cache_yaml: true,
+        compile_cache_json: true,
       )
       Bootsnap.expects(:logger=).with($stderr.method(:puts))
 

test/test_helper.rb

@@ -10,6 +10,7 @@
 require('bundler/setup')
 require('bootsnap')
 require('bootsnap/compile_cache/yaml')
+require('bootsnap/compile_cache/json')
 
 require('tmpdir')
 require('fileutils')
@@ -18,7 +19,7 @@
 require('mocha/minitest')
 
 cache_dir = File.expand_path('../../tmp/bootsnap/compile-cache', __FILE__)
-Bootsnap::CompileCache.setup(cache_dir: cache_dir, iseq: true, yaml: false)
+Bootsnap::CompileCache.setup(cache_dir: cache_dir, iseq: true, yaml: false, json: false)
 
 if GC.respond_to?(:verify_compaction_references)
   # This method was added in Ruby 3.0.0. Calling it this way asks the GC to
@@ -99,6 +100,7 @@ def setup
     @prev = Bootsnap::CompileCache::ISeq.cache_dir
     Bootsnap::CompileCache::ISeq.cache_dir = @tmp_dir
     Bootsnap::CompileCache::YAML.cache_dir = @tmp_dir
+    Bootsnap::CompileCache::JSON.cache_dir = @tmp_dir
   end
 
   def teardown
@@ -107,5 +109,6 @@ def teardown
     FileUtils.remove_entry(@tmp_dir)
     Bootsnap::CompileCache::ISeq.cache_dir = @prev
     Bootsnap::CompileCache::YAML.cache_dir = @prev
+    Bootsnap::CompileCache::JSON.cache_dir = @prev
   end
 end