From 7ecee7a4d79718bc46086ee8dec23a00cea39b39 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Tue, 7 Feb 2017 22:09:13 +0800
Subject: [PATCH] Introduce ThreadedConnectionPool for customized pool

This way we could reuse this pool for other migrations

Feedback:

* https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/8987#note_22923350
* https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/8987#note_22923365
---
 ..._remove_inactive_default_email_services.rb | 45 +++++++----------
 .../database/threaded_connection_pool.rb      | 48 +++++++++++++++++++
 2 files changed, 65 insertions(+), 28 deletions(-)
 create mode 100644 lib/gitlab/database/threaded_connection_pool.rb

diff --git a/db/post_migrate/20170206040400_remove_inactive_default_email_services.rb b/db/post_migrate/20170206040400_remove_inactive_default_email_services.rb
index d69f68a13d2..52e7f91bb01 100644
--- a/db/post_migrate/20170206040400_remove_inactive_default_email_services.rb
+++ b/db/post_migrate/20170206040400_remove_inactive_default_email_services.rb
@@ -3,38 +3,27 @@ class RemoveInactiveDefaultEmailServices < ActiveRecord::Migration
 
   DOWNTIME = false
 
+  disable_ddl_transaction!
+
   def up
-    builds_service = spawn <<-SQL.strip_heredoc
-      DELETE FROM services
-        WHERE type = 'BuildsEmailService'
-          AND active IS FALSE
-          AND properties = '{"notify_only_broken_builds":true}';
-    SQL
+    Gitlab::Database::ThreadedConnectionPool.with_pool(2) do |pool|
+      pool.execute_async <<-SQL.strip_heredoc
+        DELETE FROM services
+          WHERE type = 'BuildsEmailService'
+            AND active IS FALSE
+            AND properties = '{"notify_only_broken_builds":true}';
+      SQL
 
-    pipelines_service = spawn <<-SQL.strip_heredoc
-      DELETE FROM services
-        WHERE type = 'PipelinesEmailService'
-          AND active IS FALSE
-          AND properties = '{"notify_only_broken_pipelines":true}';
-    SQL
-
-    [builds_service, pipelines_service].each(&:join)
-  end
-
-  private
-
-  def spawn(query)
-    Thread.new do
-      with_connection do |connection|
-        connection.execute(query)
-      end
+      pool.execute_async <<-SQL.strip_heredoc
+        DELETE FROM services
+          WHERE type = 'PipelinesEmailService'
+            AND active IS FALSE
+            AND properties = '{"notify_only_broken_pipelines":true}';
+      SQL
     end
   end
 
-  def with_connection(&block)
-    pool = ActiveRecord::Base.establish_connection
-    pool.with_connection(&block)
-  ensure
-    pool.disconnect!
+  def down
+    # Nothing can be done to restore the records
   end
 end
diff --git a/lib/gitlab/database/threaded_connection_pool.rb b/lib/gitlab/database/threaded_connection_pool.rb
new file mode 100644
index 00000000000..1316b005741
--- /dev/null
+++ b/lib/gitlab/database/threaded_connection_pool.rb
@@ -0,0 +1,48 @@
+module Gitlab
+  module Database
+    class ThreadedConnectionPool
+      def self.with_pool(pool_size)
+        pool = new(pool_size)
+
+        yield(pool)
+
+      ensure
+        pool.join
+        pool.close
+      end
+
+      def initialize(pool_size)
+        config = ActiveRecord::Base.configurations[Rails.env]
+        @ar_pool = ActiveRecord::Base.establish_connection(
+          config.merge(pool: pool_size))
+        @workers = []
+        @mutex = Mutex.new
+      end
+
+      def execute_async(sql)
+        @mutex.synchronize do
+          @workers << Thread.new do
+            @ar_pool.with_connection do |connection|
+              connection.execute(sql)
+            end
+          end
+        end
+      end
+
+      def join
+        threads = nil
+
+        @mutex.synchronize do
+          threads = @workers.dup
+          @workers.clear
+        end
+
+        threads.each(&:join)
+      end
+
+      def close
+        @ar_pool.disconnect!
+      end
+    end
+  end
+end