1 files changed, 56 insertions, 7 deletions

diff --git a/lib/bundler/vendor/connection_pool/lib/connection_pool.rb b/lib/bundler/vendor/connection_pool/lib/connection_pool.rb
index 984c1c3dcb..317088a866 100644
--- a/lib/bundler/vendor/connection_pool/lib/connection_pool.rb
+++ b/lib/bundler/vendor/connection_pool/lib/connection_pool.rb
@@ -1,10 +1,12 @@
-require "timeout"
+require_relative "../../../vendored_timeout"
 require_relative "connection_pool/version"
 
 class Bundler::ConnectionPool
   class Error < ::RuntimeError; end
+
   class PoolShuttingDownError < ::Bundler::ConnectionPool::Error; end
-  class TimeoutError < ::Timeout::Error; end
+
+  class TimeoutError < ::Gem::Timeout::Error; end
 end
 
 # Generic connection pool class for sharing a limited number of objects or network connections
@@ -34,14 +36,57 @@ end
 # Accepts the following options:
 # - :size - number of connections to pool, defaults to 5
 # - :timeout - amount of time to wait for a connection if none currently available, defaults to 5 seconds
+# - :auto_reload_after_fork - automatically drop all connections after fork, defaults to true
 #
 class Bundler::ConnectionPool
-  DEFAULTS = {size: 5, timeout: 5}
+  DEFAULTS = {size: 5, timeout: 5, auto_reload_after_fork: true}
 
   def self.wrap(options, &block)
     Wrapper.new(options, &block)
   end
 
+  if Process.respond_to?(:fork)
+    INSTANCES = ObjectSpace::WeakMap.new
+    private_constant :INSTANCES
+
+    def self.after_fork
+      INSTANCES.values.each do |pool|
+        next unless pool.auto_reload_after_fork
+
+        # We're on after fork, so we know all other threads are dead.
+        # All we need to do is to ensure the main thread doesn't have a
+        # checked out connection
+        pool.checkin(force: true)
+        pool.reload do |connection|
+          # Unfortunately we don't know what method to call to close the connection,
+          # so we try the most common one.
+          connection.close if connection.respond_to?(:close)
+        end
+      end
+      nil
+    end
+
+    if ::Process.respond_to?(:_fork) # MRI 3.1+
+      module ForkTracker
+        def _fork
+          pid = super
+          if pid == 0
+            Bundler::ConnectionPool.after_fork
+          end
+          pid
+        end
+      end
+      Process.singleton_class.prepend(ForkTracker)
+    end
+  else
+    INSTANCES = nil
+    private_constant :INSTANCES
+
+    def self.after_fork
+      # noop
+    end
+  end
+
   def initialize(options = {}, &block)
     raise ArgumentError, "Connection pool requires a block" unless block
 
@@ -49,10 +94,12 @@ class Bundler::ConnectionPool
 
     @size = Integer(options.fetch(:size))
     @timeout = options.fetch(:timeout)
+    @auto_reload_after_fork = options.fetch(:auto_reload_after_fork)
 
     @available = TimedStack.new(@size, &block)
     @key = :"pool-#{@available.object_id}"
     @key_count = :"pool-#{@available.object_id}-count"
+    INSTANCES[self] = self if INSTANCES
   end
 
   def with(options = {})
@@ -67,7 +114,7 @@ class Bundler::ConnectionPool
       end
     end
   end
-  alias then with
+  alias_method :then, :with
 
   def checkout(options = {})
     if ::Thread.current[@key]
@@ -79,16 +126,16 @@ class Bundler::ConnectionPool
     end
   end
 
-  def checkin
+  def checkin(force: false)
     if ::Thread.current[@key]
-      if ::Thread.current[@key_count] == 1
+      if ::Thread.current[@key_count] == 1 || force
         @available.push(::Thread.current[@key])
         ::Thread.current[@key] = nil
         ::Thread.current[@key_count] = nil
       else
         ::Thread.current[@key_count] -= 1
       end
-    else
+    elsif !force
       raise Bundler::ConnectionPool::Error, "no connections are checked out"
     end
 
@@ -115,6 +162,8 @@ class Bundler::ConnectionPool
 
   # Size of this connection pool
   attr_reader :size
+  # Automatically drop all connections after fork
+  attr_reader :auto_reload_after_fork
 
   # Number of pool entries available for checkout at this instant.
   def available