Allow extra runtime deps to be defined when using prune_bundler

lib/puma/cli.rb

@@ -161,6 +161,10 @@ def setup_options
             user_config.prune_bundler
           end
 
+          o.on "--extra-runtime-dependencies GEM1,GEM2", "Defines any extra needed gems when using --prune-bundler" do |arg|
+            c.extra_runtime_dependencies arg.split(',')
+          end
+
           o.on "-q", "--quiet", "Do not log requests internally (default true)" do
             user_config.quiet
           end

lib/puma/dsl.rb

@@ -584,6 +584,7 @@ def lowlevel_error_handler(obj=nil, &block)
     # dictates.
     #
     # @note This is incompatible with +preload_app!+.
+    # @note This is only supported for RubyGems 2.2+
     def prune_bundler(answer=true)
       @options[:prune_bundler] = answer
     end
@@ -601,6 +602,21 @@ def raise_exception_on_sigterm(answer=true)
       @options[:raise_exception_on_sigterm] = answer
     end
 
+    # When using prune_bundler, if extra runtime dependencies need to be loaded to
+    # initialize your app, then this setting can be used.
+    #
+    # Before bundler is pruned, the gem names supplied will be looked up in the bundler
+    # context and then loaded again after bundler is pruned.
+    # Only applies if prune_bundler is used.
+    #
+    # @example
+    #   extra_runtime_dependencies ['gem_name_1', 'gem_name_2']
+    # @example
+    #   extra_runtime_dependencies ['puma_worker_killer']
+    def extra_runtime_dependencies(answer = [])
+      @options[:extra_runtime_dependencies] = Array(answer)
+    end
+
     # Additional text to display in process listing.
     #
     # If you do not specify a tag, Puma will infer it. If you do not want Puma

lib/puma/launcher.rb

@@ -259,35 +259,64 @@ def restart!
       end
     end
 
-    def prune_bundler
-      return unless defined?(Bundler)
-      puma = Bundler.rubygems.loaded_specs("puma")
-      dirs = puma.require_paths.map { |x| File.join(puma.full_gem_path, x) }
+    def dependencies_and_files_to_require_after_prune
+      puma = spec_for_gem("puma")
+
+      deps = puma.runtime_dependencies.map do |d|
+        "#{d.name}:#{spec_for_gem(d.name).version}"
+      end
+
+      [deps, require_paths_for_gem(puma) + extra_runtime_deps_directories]
+    end
+
+    def extra_runtime_deps_directories
+      Array(@options[:extra_runtime_dependencies]).map do |d_name|
+        if (spec = spec_for_gem(d_name))
+          require_paths_for_gem(spec)
+        else
+          log "* Could not load extra dependency: #{d_name}"
+          nil
+        end
+      end.flatten.compact
+    end
+
+    def puma_wild_location
+      puma = spec_for_gem("puma")
+      dirs = require_paths_for_gem(puma)
       puma_lib_dir = dirs.detect { |x| File.exist? File.join(x, '../bin/puma-wild') }
+      File.expand_path(File.join(puma_lib_dir, "../bin/puma-wild"))
+    end
 
-      unless puma_lib_dir
+    def prune_bundler
+      return unless defined?(Bundler)
+      require_rubygems_min_version!(Gem::Version.new("2.2"), "prune_bundler")
+      unless puma_wild_location
         log "! Unable to prune Bundler environment, continuing"
         return
       end
 
-      deps = puma.runtime_dependencies.map do |d|
-        spec = Bundler.rubygems.loaded_specs(d.name)
-        "#{d.name}:#{spec.version.to_s}"
-      end
+      deps, dirs = dependencies_and_files_to_require_after_prune
 
       log '* Pruning Bundler environment'
       home = ENV['GEM_HOME']
       Bundler.with_clean_env do
         ENV['GEM_HOME'] = home
         ENV['PUMA_BUNDLER_PRUNED'] = '1'
-        wild = File.expand_path(File.join(puma_lib_dir, "../bin/puma-wild"))
-        args = [Gem.ruby, wild, '-I', dirs.join(':'), deps.join(',')] + @original_argv
+        args = [Gem.ruby, puma_wild_location, '-I', dirs.join(':'), deps.join(',')] + @original_argv
         # Ruby 2.0+ defaults to true which breaks socket activation
         args += [{:close_others => false}]
         Kernel.exec(*args)
       end
     end
 
+    def spec_for_gem(gem_name)
+      Bundler.rubygems.loaded_specs(gem_name)
+    end
+
+    def require_paths_for_gem(gem_spec)
+      gem_spec.full_require_paths
+    end
+
     def log(str)
       @events.log str
     end
@@ -430,5 +459,12 @@ def setup_signals
         log "*** SIGHUP not implemented, signal based logs reopening unavailable!"
       end
     end
+
+    def require_rubygems_min_version!(min_version, feature)
+      return if min_version <= Gem::Version.new(Gem::VERSION)
+
+      raise "#{feature} is not supported on your version of RubyGems. " \
+              "You must have RubyGems #{min_version}+ to use this feature."
+    end
   end
 end

test/config/prune_bundler_with_deps.rb

@@ -0,0 +1,2 @@
+prune_bundler true
+extra_runtime_dependencies ["rdoc"]

test/rackup/hello-last-load-path.ru

@@ -0,0 +1,3 @@
+run lambda { |env|
+  [200, {}, [$LOAD_PATH[-1]]]
+}

test/test_integration.rb

@@ -348,6 +348,15 @@ def test_term_signal_suppress_in_clustered_mode
     @server = nil # prevent `#teardown` from killing already killed server
   end
 
+  def test_load_path_includes_extra_deps
+    skip NO_FORK_MSG unless HAS_FORK
+
+    server("-w 2 -C test/config/prune_bundler_with_deps.rb test/rackup/hello-last-load-path.ru")
+
+    last_load_path = read_body(connect)
+    assert_match(%r{gems/rdoc-[\d.]+/lib$}, last_load_path)
+  end
+
   def test_not_accepts_new_connections_after_term_signal
     skip_on :jruby, :windows
 

test/test_launcher.rb

@@ -0,0 +1,63 @@
+require_relative "helper"
+
+require "puma/configuration"
+require "puma/const"
+require "puma/launcher"
+
+class TestLauncher < Minitest::Test
+  def test_dependencies_and_files_to_require_after_prune_is_correctly_built_for_no_extra_deps
+    skip_on :appveyor, suffix: " - bundler not used in appveyor so prune bundler logic tests unavailable"
+
+    l = Puma::Launcher.new Puma::Configuration.new
+    deps, dirs = l.send(:dependencies_and_files_to_require_after_prune)
+
+    assert_equal(1, deps.length)
+    assert_match(%r{^nio4r:[\d.]+$}, deps.first)
+    assert_equal(2, dirs.length)
+    assert_match(%r{puma/lib$}, dirs[0]) # lib dir
+    assert_match(%r{puma-#{Puma::Const::PUMA_VERSION}$}, dirs[1]) # native extension dir
+  end
+
+  def test_dependencies_and_files_to_require_after_prune_is_correctly_built_with_extra_deps
+    skip_on :appveyor, suffix: " - bundler not used in appveyor so prune bundler logic tests unavailable"
+
+    conf = Puma::Configuration.new do |c|
+      c.extra_runtime_dependencies ['rdoc']
+    end
+    l = Puma::Launcher.new conf
+    deps, dirs = l.send(:dependencies_and_files_to_require_after_prune)
+
+    assert_equal(1, deps.length)
+    assert_match(%r{^nio4r:[\d.]+$}, deps.first)
+    assert_equal(3, dirs.length)
+    assert_match(%r{puma/lib$}, dirs[0]) # lib dir
+    assert_match(%r{puma-#{Puma::Const::PUMA_VERSION}$}, dirs[1]) # native extension dir
+    assert_match(%r{gems/rdoc-[\d.]+/lib$}, dirs[2]) # rdoc dir
+  end
+
+  def test_extra_runtime_deps_directories_is_empty_for_no_config
+    l = Puma::Launcher.new Puma::Configuration.new
+    assert_equal([], l.send(:extra_runtime_deps_directories))
+  end
+
+  def test_extra_runtime_deps_directories_is_correctly_built
+    skip_on :appveyor, suffix: " - bundler not used in appveyor so prune bundler logic tests unavailable"
+    conf = Puma::Configuration.new do |c|
+      c.extra_runtime_dependencies ['rdoc']
+    end
+    l = Puma::Launcher.new conf
+    dep_dirs = l.send(:extra_runtime_deps_directories)
+
+    assert_equal(1, dep_dirs.length)
+    assert_match(%r{gems/rdoc-[\d.]+/lib$}, dep_dirs.first)
+  end
+
+  def test_puma_wild_location_is_an_absolute_path
+    skip_on :appveyor, suffix: " - bundler not used in appveyor so prune bundler logic tests unavailable"
+    l = Puma::Launcher.new Puma::Configuration.new
+    puma_wild_location = l.send(:puma_wild_location)
+    assert_match(%r{bin/puma-wild$}, puma_wild_location)
+    # assert no "/../" in path
+    refute_match(%r{/\.\./}, puma_wild_location)
+  end
+end