Ruby로 Connection Pool 구현하기

구경회·2023년 1월 15일
1

커넥션 풀?

아마 대부분의 개발자에게 익숙한 단어일 것이다. 생성하는 비용이 있는 것을 대상으로 여러 Object를 만들어두고 재활용하는 방식이다. 대표적으로는 쓰레드풀, DB의 커넥션 풀등이 있으며 Http Client도 커넥션 풀을 활용하는 경우가 잦다.

개인적으로 이걸 구현하며 공부해보면 더 재미있을 거 같아 글을 쓴다. 언어는 요새 가장 편한 루비로 하자.

나이브한 구현

단순한 테스트 먼저 작성하자. RSpec을 더 좋아하지만, 간단한 테스트니 Minitest를 사용해보겠다.

require_relative 'helper'

class TestPool < Minitest::Test
  class Foo
    def initialize
      puts 'Foo#initialize'
    end
  end
  
  def test_create
    pool = Pool.new(3) { Foo.new }
    
    assert_equal 3, pool.size
  end
end

아주 단순한 Pool의 시그니처와 Pool#size 메서드에 대한 테스트를 작성했다. 다음과 같이 이제 구현하자.

class Pool
  attr_reader :size

  def initialize(size, &block)
    @size = size
  end
end

좀 더 나아가서, 일반적인 Pool의 개념과 같으려면 다음과 같이 사용할 수 있어야 할 거 같다.

class Foo
  def initialize
    puts 'Foo#initialize'
  end

  def bar
    42
  end
end

def test_with
  pool = Pool.new(3) { Foo.new }

  assert_equal 3, pool.size
  assert_equal(pool.with { |conn| conn.bar }, 42)
end
class Pool
  attr_reader :size

  def initialize(size, &block)
    @size = size
    @pools = Array.new(size, &block)
  end

  def with
    yield @pools.pop
  end
end

위와 같이 단순하게 구현해보자. TC가 통과한다. 하지만 문제점이 보이나? 우리는 커넥션 풀에서 커넥션을 꺼내온 후에 반환하지 않는다. 아마 에러가 뜰 것이다. 에러를 띄우는 TC를 만들어보자.

def test_with
  pool = Pool.new(3) { Foo.new }

  assert_equal 3, pool.size
  assert_equal(pool.with { |conn| conn.bar }, 42)
  assert_equal(pool.with { |conn| conn.bar }, 42)
  assert_equal(pool.with { |conn| conn.bar }, 42)
  assert_equal(pool.with { |conn| conn.bar }, 42)
end
NoMethodError: undefined method `bar' for nil:NilClass

예상한대로 에러가 발생한다.

우선 클래스의 사용성을 다듬기 위해, @pools에 가능한 게 하나도 없을 때 커스텀 에러를 발생시키도록 하자.

class Pool
  PoolExhaustedError = Class.new(StandardError)

  # ...

  def with
    raise PoolExhaustedError if @pools.zero?

    yield @pools.pop
  end
end

좀 낫다. 다음과 같은 TC로 검증하자.

def test_pool_exhausted
  pool = Pool.new(0) { Foo.new }

  assert_raises(Pool::PoolExhaustedError) do
    pool.with(&:bar)
  end
end

이제 커넥션 풀을 쓴 후에 반환할 차례이다. 이 부분은 단순하다.

class Pool
  # ...

  def with
    raise PoolExhaustedError if @pools.zero?

    connection = @pools.pop
    yield connection
  ensure
    @pools.push(connection)
  end
end

이제 문제의 TC가 잘 동작하는 걸 볼 수 있다.

Thread-safety

하지만 위의 구현에도 문제는 있다.

위와 같이 귀여운 강아지 사진을 보면 알 수 있는데, 다음과 같은 코드를 생각해보자.

class Foo
  def delayed_bar
    sleep 0.5
    42
  end
end
  
def test_all_occupied
  Thread.abort_on_exception = true

  pool = Pool.new(2) { Foo.new }

  threads = []

  5.times do
    threads << Thread.new { pool.with(&:delayed_bar) }
  end

  threads.each(&:join)
end
Traceback (most recent call last):
	1: from /Users/username/workspace/connection_pool/test/test_pool.rb:48:in `block (2 levels) in test_all_occupied'
/Users/username/workspace/connection_pool/lib/pool.rb:19:in `with': undefined method `delayed_bar' for nil:NilClass (NoMethodError)

이 때는 또 다시 에러가 나게 된다. 어떻게 해결해야할까? 바로 Mutex의 도입이다. 이제 우리 커넥션 풀에 뮤텍스를 도입하고, 자원을 점유하려 할 때 싱크를 해주자.

class Pool
  PoolExhaustedError = Class.new(StandardError)

  attr_reader :size

  def initialize(size, &block)
    # ...
    @mutex = Mutex.new
  end

  def with
    @mutex.synchronize do
      # ...
    end
  end
  
  # ...
end

이제 문제의 해결이다.

Timeout 구현하기

레일즈를 쓰다보면 다음과 같은 에러를 아마 수도 없이 마주쳤을 것이다.

(ActiveRecord::ConnectionTimeoutError) "could not obtain a database connection within 5 seconds (waited 5.000144774 seconds).

우리도 비슷하게 커넥션 풀에 타임아웃을 구현해보자.

class Pool
  PoolExhaustedError = Class.new(StandardError)

  attr_reader :size

  def initialize(size, timeout: 0.5, &block)
    @size = size
    @available_size = size
    @pools = Array.new(size, &block)
    @mutex = Mutex.new
    @conditional_var = ConditionVariable.new
    @timeout = timeout
  end

  def with
    @mutex.synchronize do
      @conditional_var.wait(@mutex, @timeout) if exhausted?
      raise PoolExhaustedError if exhausted?

      connection = @pools.pop
      yield connection
    ensure
      @pools.push(connection)
      @conditional_var.signal
    end
  end

  def available_pool_size
    @pools.size
  end

  private

  def exhausted?
    available_pool_size.zero?
  end
end

위와 같이 ConditionVariable을 사용하면 된다.

profile
즐기는 거야

0개의 댓글