shared_connection_pool.rb 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. # frozen_string_literal: true
  2. require 'connection_pool'
  3. require_relative './shared_timed_stack'
  4. class ConnectionPool::SharedConnectionPool < ConnectionPool
  5. def initialize(options = {}, &block)
  6. super(options, &block)
  7. @available = ConnectionPool::SharedTimedStack.new(@size, &block)
  8. end
  9. delegate :size, :flush, to: :@available
  10. def with(preferred_tag, options = {})
  11. Thread.handle_interrupt(Exception => :never) do
  12. conn = checkout(preferred_tag, options)
  13. begin
  14. Thread.handle_interrupt(Exception => :immediate) do
  15. yield conn
  16. end
  17. ensure
  18. checkin(preferred_tag)
  19. end
  20. end
  21. end
  22. def checkout(preferred_tag, options = {})
  23. if ::Thread.current[key(preferred_tag)]
  24. ::Thread.current[key_count(preferred_tag)] += 1
  25. ::Thread.current[key(preferred_tag)]
  26. else
  27. ::Thread.current[key_count(preferred_tag)] = 1
  28. ::Thread.current[key(preferred_tag)] = @available.pop(preferred_tag, options[:timeout] || @timeout)
  29. end
  30. end
  31. def checkin(preferred_tag)
  32. if ::Thread.current[key(preferred_tag)]
  33. if ::Thread.current[key_count(preferred_tag)] == 1
  34. @available.push(::Thread.current[key(preferred_tag)])
  35. ::Thread.current[key(preferred_tag)] = nil
  36. else
  37. ::Thread.current[key_count(preferred_tag)] -= 1
  38. end
  39. else
  40. raise ConnectionPool::Error, 'no connections are checked out'
  41. end
  42. nil
  43. end
  44. private
  45. def key(tag)
  46. :"#{@key}-#{tag}"
  47. end
  48. def key_count(tag)
  49. :"#{@key_count}-#{tag}"
  50. end
  51. end