From de76c7bc21409a01cffdcfa1cff15244e82a191f Mon Sep 17 00:00:00 2001
From: coderrr <coderrr.contact@gmail.com>
Date: Thu, 15 Jan 2009 05:00:38 +0700
Subject: [PATCH] added cleanup_connection which releases connection after block finishes while protecting against nested calls

---
 .../abstract/connection_pool.rb                    |   15 ++++++
 activerecord/test/cases/pooled_connections_test.rb |   54 ++++++++++++++++++++
 2 files changed, 69 insertions(+), 0 deletions(-)

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index 901b171..0e3a62c 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -116,6 +116,21 @@ module ActiveRecord
         checkin conn
       end
 
+      # Yields to a block and ensures that the connection is released after the
+      # block completes.  Makes sure the connection is only released on the
+      # topmost method call in case they are nested.
+      def cleanup_connection
+        return yield if Thread.current[:__AR__cleanup_connection]
+
+        begin
+          Thread.current[:__AR__cleanup_connection] = true
+          yield
+        ensure
+          release_connection
+          Thread.current[:__AR__cleanup_connection] = false
+        end
+      end
+
       # Returns true if a connection has already been opened.
       def connected?
         !@connections.empty?
diff --git a/activerecord/test/cases/pooled_connections_test.rb b/activerecord/test/cases/pooled_connections_test.rb
index 2649a93..1e24d39 100644
--- a/activerecord/test/cases/pooled_connections_test.rb
+++ b/activerecord/test/cases/pooled_connections_test.rb
@@ -1,4 +1,5 @@
 require "cases/helper"
+require "timeout"
 
 class PooledConnectionsTest < ActiveRecord::TestCase
   def setup
@@ -89,6 +90,59 @@ class PooledConnectionsTest < ActiveRecord::TestCase
   ensure
     ActiveRecord::Base.connection_handler = old_handler
   end
+
+  def test_cleanup_connection
+    ActiveRecord::Base.establish_connection(@connection.merge({:pool => 2, :wait_timeout => 10}))
+    threads = []
+    10.times do
+      threads << Thread.new do
+        ActiveRecord::Base.connection_pool.cleanup_connection do
+          ActiveRecord::Base.connection.execute("select 1")
+        end
+      end
+    end
+
+    assert_nothing_raised do
+      Timeout.timeout(2) do
+        threads.each {|t| t.join }
+      end
+    end
+  end
+
+  def add_record(name)
+    ActiveRecord::Base.connection_pool.cleanup_connection do
+      ActiveRecord::Base.connection.execute("insert into projects (name) values ('#{name}')")
+    end
+  end
+
+  def transaction
+    ActiveRecord::Base.connection_pool.cleanup_connection do
+      ActiveRecord::Base.connection.execute("start transaction")
+      yield
+      ActiveRecord::Base.connection.execute("commit")
+    end
+  end
+
+  def test_cleanup_connection_nesting_safety
+    ActiveRecord::Base.establish_connection(@connection.merge({:pool => 1, :wait_timeout => 0.1}))
+
+    before_count = ActiveRecord::Base.connection.execute("select * from projects").num_rows
+
+    add_record('one')
+
+    transaction do
+      add_record('two')
+      # have another thread try to screw up the transaction
+      Thread.new do
+        ActiveRecord::Base.connection.execute("rollback")
+        ActiveRecord::Base.connection_pool.release_connection
+      end.join rescue nil
+      add_record('three')
+    end
+
+    after_count = ActiveRecord::Base.connection.execute("select * from projects").num_rows
+    assert_equal 3, after_count - before_count
+  end
 end unless %w(FrontBase).include? ActiveRecord::Base.connection.adapter_name
 
 class AllowConcurrencyDeprecatedTest < ActiveRecord::TestCase
-- 
1.5.6.3