Add the ability to discard the current connection (#204)

Author: womblep

README.md

@@ -129,6 +129,25 @@ Thread.new do
 end
 ```
 
+## Discarding Connections
+
+You can discard connections in the ConnectionPool instance to remove connections that are broken and can't be restarted. 
+
+NOTE: the connection is not closed. It will just be removed from the pool so it won't be selected again.
+
+It can only be done inside the block passed to `with` or `with_timeout`.
+
+```ruby
+   pool.with do |conn|
+     begin
+       conn.execute("SELECT 1")
+     rescue SomeConnectionError
+       pool.discard_current_connection  # remove the connection from the pool
+       raise
+     end
+   end
+```
+
 ## Current State
 
 There are several methods that return information about a pool.

lib/connection_pool.rb

@@ -99,6 +99,7 @@ def initialize(options = {}, &block)
     @available = TimedStack.new(@size, &block)
     @key = :"pool-#{@available.object_id}"
     @key_count = :"pool-#{@available.object_id}-count"
+    @discard_key = :"pool-#{@available.object_id}-discard"
     INSTANCES[self] = self if @auto_reload_after_fork && INSTANCES
   end
 
@@ -116,6 +117,32 @@ def with(options = {})
   end
   alias_method :then, :with
 
+  ##
+  # Marks the current thread's checked-out connection for discard.
+  #
+  # When a connection is marked for discard, it will not be returned to the pool
+  # when checked in. Instead, the connection will be discarded.
+  # This is useful when a connection has become invalid or corrupted
+  # and should not be reused.
+  #
+  # Note: This only affects the connection currently checked out by the calling thread.
+  # The connection will be discarded when +checkin+ is called.
+  #
+  # @return [void]
+  #
+  # @example
+  #   pool.with do |conn|
+  #     begin
+  #       conn.execute("SELECT 1")
+  #     rescue SomeConnectionError
+  #       pool.discard_current_connection  # Mark connection as bad
+  #       raise
+  #     end
+  #   end
+  def discard_current_connection
+    ::Thread.current[@discard_key] = true
+  end
+
   def checkout(options = {})
     if ::Thread.current[@key]
       ::Thread.current[@key_count] += 1
@@ -129,7 +156,12 @@ def checkout(options = {})
   def checkin(force: false)
     if ::Thread.current[@key]
       if ::Thread.current[@key_count] == 1 || force
-        @available.push(::Thread.current[@key])
+        if ::Thread.current[@discard_key]
+          @available.decrement_created
+          ::Thread.current[@discard_key] = nil
+        else
+          @available.push(::Thread.current[@key])
+        end
         ::Thread.current[@key] = nil
         ::Thread.current[@key_count] = nil
       else

lib/connection_pool/timed_stack.rb

@@ -134,6 +134,12 @@ def idle
     @que.length
   end
 
+  ##
+  # Reduce the created count
+  def decrement_created
+    @created -= 1 unless @created == 0
+  end
+
   private
 
   def current_time

test/test_connection_pool.rb

@@ -186,6 +186,7 @@ def test_explicit_return
       def mock.disconnect!
         raise "should not disconnect upon explicit return"
       end
+
       mock
     }
 
@@ -246,6 +247,24 @@ def test_checkin
     assert_same conn, Thread.new { pool.checkout }.value
   end
 
+  def test_discard
+    pool = ConnectionPool.new(timeout: 0, size: 1) { NetworkConnection.new }
+    pool.checkout
+
+    Thread.new {
+      assert_raises Timeout::Error do
+        pool.checkout
+      end
+    }.join
+
+    pool.discard_current_connection
+    pool.checkin
+
+    assert_equal 1, pool.size
+    assert_equal 0, pool.idle
+    assert_equal 1, pool.available
+  end
+
   def test_returns_value
     pool = ConnectionPool.new(timeout: 0, size: 1) { Object.new }
     assert_equal 1, pool.with { |o| 1 }
@@ -410,6 +429,33 @@ def test_nested_checkout
     assert_equal ["inner", "outer", "other"], recorder.calls
   end
 
+  def test_nested_discard
+    recorder = Recorder.new
+    pool = ConnectionPool.new(size: 1) { {recorder: recorder} }
+    pool.with do |r_outer|
+      @other = Thread.new { |t|
+        pool.with do |r_other|
+          r_other[:recorder].do_work("other")
+        end
+      }
+
+      pool.with do |r_inner|
+        @inner = r_inner
+        r_inner[:recorder].do_work("inner")
+        pool.discard_current_connection
+      end
+
+      Thread.pass
+
+      r_outer[:recorder].do_work("outer")
+    end
+
+    @other.join
+
+    assert_equal ["inner", "outer", "other"], recorder.calls
+    refute_same @inner, pool.checkout
+  end
+
   def test_shutdown_is_executed_for_all_connections
     recorders = []