Reduce Aliki’s peak memory usage by streaming search-index generation

lib/rdoc/generator/aliki.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'uri'
+require 'tempfile'
 
 module RDoc
   module Generator
@@ -74,29 +75,7 @@ def write_style_sheet
       # Build a search index array for Aliki's searcher.
 
       def build_search_index
-        setup
-
-        index = []
-
-        @classes.each do |klass|
-          next unless klass.display?
-
-          index << build_class_module_entry(klass)
-
-          klass.constants.each do |const|
-            next unless const.display?
-
-            index << build_constant_entry(const, klass)
-          end
-        end
-
-        @methods.each do |method|
-          next unless method.display?
-
-          index << build_method_entry(method)
-        end
-
-        index
+        enum_for(:each_search_index_entry).to_a
       end
 
       ##
@@ -110,15 +89,38 @@ def build_search_index
       def write_search_index
         debug_msg "Writing Aliki search index"
 
-        index = build_search_index
-
-        FileUtils.mkdir_p 'js' unless @dry_run
-
-        search_index_path = 'js/search_data.js'
-        return if @dry_run
+        if @dry_run
+          each_search_index_entry {}
+          return
+        end
 
-        data = { index: index }
-        File.write search_index_path, "var search_data = #{JSON.generate(data)};"
+        FileUtils.mkdir_p 'js'
+        # Follow symlinks, but replace the target inode rather than updating other hard links.
+        search_index_path = File.realdirpath 'js/search_data.js'
+        stat = begin
+          File.stat search_index_path
+        rescue Errno::ENOENT
+          nil
+        end
+        mode = stat ? stat.mode & 0777 : 0666 & ~File.umask
+
+        Tempfile.create('rdoc-search', File.dirname(search_index_path)) do |file|
+          file.write 'var search_data = {"index":['
+          separator = ''
+          each_search_index_entry do |entry|
+            file.write separator, JSON.generate(entry)
+            separator = ','
+          end
+          file.write ']};'
+          if stat
+            tempfile_stat = file.stat
+            file.chown stat.uid, stat.gid if stat.uid != tempfile_stat.uid || stat.gid != tempfile_stat.gid
+          end
+          # shortcut: extended ACLs are not copied; add platform-specific support if indexes need ACL-based access.
+          file.chmod mode
+          file.close
+          File.rename file.path, search_index_path
+        end
       end
 
       ##
@@ -153,6 +155,28 @@ def resolve_url(rel_prefix, url)
 
     private
 
+      def each_search_index_entry
+        setup
+
+        @classes.each do |klass|
+          next unless klass.display?
+
+          yield build_class_module_entry(klass)
+
+          klass.constants.each do |const|
+            next unless const.display?
+
+            yield build_constant_entry(const, klass)
+          end
+        end
+
+        @methods.each do |method|
+          next unless method.display?
+
+          yield build_method_entry(method)
+        end
+      end
+
       def template_encoding
         ::Encoding::UTF_8
       end

test/rdoc/generator/aliki/search_index_test.rb

@@ -213,4 +213,229 @@ def test_build_search_index_entry_structure
 
     assert_equal %i[name full_name type path].sort, entry.keys.sort
   end
+
+  def test_write_search_index_matches_array_serialization
+    klass = @top_level.add_class RDoc::NormalClass, 'Outer'
+    klass.add_comment "Quotes \" and \\, café 日本語, <tags> & links.\n", @top_level
+    klass.add_class RDoc::NormalClass, 'Inner'
+    klass.add_constant RDoc::Constant.new('VALUE', '1', 'A constant')
+    klass.add_method RDoc::AnyMethod.new('[]')
+    method = RDoc::AnyMethod.new('build')
+    method.singleton = true
+    klass.add_method method
+    @store.complete :private
+    expected = "var search_data = #{JSON.generate(index: @g.build_search_index)};"
+
+    @g.write_search_index
+
+    assert_equal expected, File.read('js/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('js')
+  end
+
+  def test_write_search_index_empty
+    @g.write_search_index
+
+    assert_equal 'var search_data = {"index":[]};', File.read('js/search_data.js')
+  end
+
+  def test_write_search_index_serializes_before_building_next_entry
+    @top_level.add_class RDoc::NormalClass, 'First'
+    @top_level.add_class RDoc::NormalClass, 'Second'
+    @store.complete :private
+    built = []
+    serialized = []
+    @g.define_singleton_method(:build_class_module_entry) do |klass|
+      raise 'Previous entries have not been serialized' unless built == serialized
+
+      built << klass.name
+      entry = Object.new
+      entry.define_singleton_method(:to_json) do |*|
+        serialized << klass.name
+        JSON.generate(name: klass.name)
+      end
+      entry
+    end
+
+    @g.write_search_index
+
+    assert_equal ['First', 'Second'], serialized
+    assert_equal built, serialized
+  end
+
+  def test_write_search_index_preserves_existing_file_on_serialization_error
+    @top_level.add_class RDoc::NormalClass, 'First'
+    @top_level.add_class RDoc::NormalClass, 'Second'
+    @store.complete :private
+    @g.define_singleton_method(:build_class_module_entry) do |klass|
+      entry = super(klass)
+      entry[:snippet] = "\xFF" if klass.name == 'Second'
+      entry
+    end
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+
+    assert_raise JSON::GeneratorError do
+      @g.write_search_index
+    end
+
+    assert_equal 'previous index', File.read('js/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('js')
+  end
+
+  def test_write_search_index_replaces_existing_file
+    omit 'open files cannot always be replaced on Windows' if /mswin|mingw/ =~ RUBY_PLATFORM
+
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+
+    File.open 'js/search_data.js' do |previous|
+      @g.write_search_index
+
+      assert_equal 'previous index', previous.read
+    end
+
+    assert_equal 'var search_data = {"index":[]};', File.read('js/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('js')
+  end
+
+  def test_write_search_index_preserves_existing_file_on_publication_error
+    original = File.method :rename
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+    File.singleton_class.remove_method :rename
+    File.define_singleton_method(:rename) { |*| raise Errno::ENOSPC }
+
+    assert_raise Errno::ENOSPC do
+      @g.write_search_index
+    end
+
+    assert_equal 'previous index', File.read('js/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('js')
+  ensure
+    File.singleton_class.remove_method :rename
+    File.define_singleton_method :rename, original
+  end
+
+  def test_write_search_index_preserves_symlink
+    FileUtils.mkdir_p %w[js target]
+    File.write 'target/search_data.js', 'previous index'
+    begin
+      File.symlink '../target/search_data.js', 'js/search_data.js'
+    rescue NotImplementedError, Errno::EACCES, Errno::EPERM
+      omit 'symlinks are not supported'
+    end
+
+    @g.write_search_index
+
+    assert File.symlink?('js/search_data.js')
+    assert_equal 'var search_data = {"index":[]};', File.read('target/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('target')
+  end
+
+  def test_write_search_index_follows_dangling_symlink
+    FileUtils.mkdir_p %w[js target]
+    begin
+      File.symlink '../target/search_data.js', 'js/search_data.js'
+    rescue NotImplementedError, Errno::EACCES, Errno::EPERM
+      omit 'symlinks are not supported'
+    end
+
+    @g.write_search_index
+
+    assert File.symlink?('js/search_data.js')
+    assert_equal 'var search_data = {"index":[]};', File.read('target/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('target')
+  end
+
+  def test_write_search_index_preserves_permissions
+    omit 'assumes UNIX permission model' if /mswin|mingw/ =~ RUBY_PLATFORM
+
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+    File.chmod 0640, 'js/search_data.js'
+
+    @g.write_search_index
+
+    assert_equal 0640, File.stat('js/search_data.js').mode & 0777
+  end
+
+  def test_write_search_index_preserves_ownership
+    omit 'assumes UNIX ownership model' if /mswin|mingw/ =~ RUBY_PLATFORM
+
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+    previous_stat = File.stat 'js/search_data.js'
+    group = Process.groups.find { |gid| gid != previous_stat.gid }
+    omit 'requires a secondary group' unless group
+    File.chown nil, group, 'js/search_data.js'
+    File.chmod 0640, 'js/search_data.js'
+
+    @g.write_search_index
+
+    stat = File.stat 'js/search_data.js'
+    assert_equal previous_stat.uid, stat.uid
+    assert_equal group, stat.gid
+    assert_equal 0640, stat.mode & 0777
+  end
+
+  def test_write_search_index_preserves_existing_file_on_ownership_error
+    original = Tempfile.method :create
+    FileUtils.mkdir_p 'js'
+    File.write 'js/search_data.js', 'previous index'
+    different_owner = File.stat 'js/search_data.js'
+    uid = different_owner.uid + 1
+    different_owner.define_singleton_method(:uid) { uid }
+    Tempfile.singleton_class.remove_method :create
+    Tempfile.define_singleton_method(:create) do |*args, &block|
+      original.call(*args) do |file|
+        file.define_singleton_method(:stat) { different_owner }
+        file.define_singleton_method(:chown) { |*| raise Errno::EPERM }
+        block.call(file)
+      end
+    end
+
+    assert_raise Errno::EPERM do
+      @g.write_search_index
+    end
+
+    assert_equal 'previous index', File.read('js/search_data.js')
+    assert_equal ['search_data.js'], Dir.children('js')
+  ensure
+    Tempfile.singleton_class.remove_method :create
+    Tempfile.define_singleton_method :create, original
+  end
+
+  def test_write_search_index_uses_default_permissions
+    omit 'assumes UNIX permission model' if /mswin|mingw/ =~ RUBY_PLATFORM
+
+    old_umask = File.umask 0027
+
+    @g.write_search_index
+
+    assert_equal 0640, File.stat('js/search_data.js').mode & 0777
+  ensure
+    File.umask old_umask if old_umask
+  end
+
+  def test_write_search_index_dry_run
+    @g.dry_run = true
+    @g.define_singleton_method(:build_search_index) { raise 'must not build an array' }
+
+    @g.write_search_index
+
+    assert_equal [], Dir.children(@tmpdir)
+  end
+
+  def test_write_search_index_uses_shared_entry_iterator
+    entries = [{ name: 'Custom entry' }]
+    @g.define_singleton_method(:each_search_index_entry) do |&block|
+      entries.each(&block)
+    end
+    assert_equal entries, @g.build_search_index
+    @g.define_singleton_method(:build_search_index) { raise 'must not build an array' }
+
+    @g.write_search_index
+
+    assert_equal "var search_data = #{JSON.generate(index: entries)};", File.read('js/search_data.js')
+  end
 end