From 136dc79433295aded9ecabb15aae2dc1e228b903 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Mon, 13 Feb 2017 21:45:26 +0800
Subject: Have some simple way to create connection pool

---
 lib/gitlab/database.rb           | 24 ++++++++++++++++++++
 spec/lib/gitlab/database_spec.rb | 48 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index dc2537d36aa..e6612bc3aad 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -69,6 +69,30 @@ module Gitlab
       end
     end
 
+    def self.with_connection_pool(pool_size)
+      pool = create_connection_pool(pool_size)
+
+      yield(pool)
+
+    ensure
+      pool.disconnect!
+    end
+
+    def self.create_connection_pool(pool_size)
+      # See activerecord-4.2.7.1/lib/active_record/connection_adapters/connection_specification.rb
+      env = Rails.env
+      original_config = ActiveRecord::Base.configurations
+      env_config = original_config[env].merge('pool' => pool_size)
+      config = original_config.merge(env => env_config)
+
+      spec =
+        ActiveRecord::
+          ConnectionAdapters::
+          ConnectionSpecification::Resolver.new(config).spec(env.to_sym)
+
+      ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec)
+    end
+
     def self.connection
       ActiveRecord::Base.connection
     end
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index b142b3a2781..c9be832b1a0 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -71,6 +71,54 @@ describe Gitlab::Database, lib: true do
     end
   end
 
+  describe '.with_connection_pool' do
+    it 'creates a new connection pool and disconnect it after used' do
+      closed_pool = nil
+
+      described_class.with_connection_pool(1) do |pool|
+        pool.with_connection do |connection|
+          connection.execute('SELECT 1 AS value')
+        end
+
+        expect(pool).to be_connected
+
+        closed_pool = pool
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+
+    it 'disconnects the pool even an exception was raised' do
+      error = Class.new(RuntimeError)
+      closed_pool = nil
+
+      begin
+        described_class.with_connection_pool(1) do |pool|
+          pool.with_connection do |connection|
+            connection.execute('SELECT 1 AS value')
+          end
+
+          closed_pool = pool
+
+          raise error.new('boom')
+        end
+      rescue error
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+  end
+
+  describe '.create_connection_pool' do
+    it 'creates a new connection pool with specific pool size' do
+      pool = described_class.create_connection_pool(5)
+
+      expect(pool)
+        .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
+      expect(pool.spec.config[:pool]).to eq(5)
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value for PostgreSQL' do
       expect(described_class).to receive(:postgresql?).and_return(true)
-- 
cgit v1.2.1


From 03f1abfcc3e43fce188f94a770f5f7b6af6f36b5 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Tue, 14 Feb 2017 13:48:13 +0800
Subject: Only ensure against yield so that pool should be available

Feedback:
https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/9192#note_23293693
---
 lib/gitlab/database.rb | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index e6612bc3aad..a6e9ea3dead 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -72,10 +72,11 @@ module Gitlab
     def self.with_connection_pool(pool_size)
       pool = create_connection_pool(pool_size)
 
-      yield(pool)
-
-    ensure
-      pool.disconnect!
+      begin
+        yield(pool)
+      ensure
+        pool.disconnect!
+      end
     end
 
     def self.create_connection_pool(pool_size)
-- 
cgit v1.2.1