1 files changed, 63 insertions, 0 deletions
diff --git a/app/lib/connection_pool/shared_connection_pool.rb b/app/lib/connection_pool/shared_connection_pool.rb
new file mode 100644
index 000000000..2865a4108
--- /dev/null
+++ b/app/lib/connection_pool/shared_connection_pool.rb
@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require 'connection_pool'
+require_relative './shared_timed_stack'
+
+class ConnectionPool::SharedConnectionPool < ConnectionPool
+  def initialize(options = {}, &block)
+    super(options, &block)
+
+    @available = ConnectionPool::SharedTimedStack.new(@size, &block)
+  end
+
+  delegate :size, :flush, to: :@available
+
+  def with(preferred_tag, options = {})
+    Thread.handle_interrupt(Exception => :never) do
+      conn = checkout(preferred_tag, options)
+
+      begin
+        Thread.handle_interrupt(Exception => :immediate) do
+          yield conn
+        end
+      ensure
+        checkin(preferred_tag)
+      end
+    end
+  end
+
+  def checkout(preferred_tag, options = {})
+    if ::Thread.current[key(preferred_tag)]
+      ::Thread.current[key_count(preferred_tag)] += 1
+      ::Thread.current[key(preferred_tag)]
+    else
+      ::Thread.current[key_count(preferred_tag)] = 1
+      ::Thread.current[key(preferred_tag)] = @available.pop(preferred_tag, options[:timeout] || @timeout)
+    end
+  end
+
+  def checkin(preferred_tag)
+    if ::Thread.current[key(preferred_tag)]
+      if ::Thread.current[key_count(preferred_tag)] == 1
+        @available.push(::Thread.current[key(preferred_tag)])
+        ::Thread.current[key(preferred_tag)] = nil
+      else
+        ::Thread.current[key_count(preferred_tag)] -= 1
+      end
+    else
+      raise ConnectionPool::Error, 'no connections are checked out'
+    end
+
+    nil
+  end
+
+  private
+
+  def key(tag)
+    :"#{@key}-#{tag}"
+  end
+
+  def key_count(tag)
+    :"#{@key_count}-#{tag}"
+  end
+end