Refactor ConnectionPool to use SizedQueue for connection management and improve cleanup logic

lib/wayback_machine_downloader.rb

@@ -25,69 +25,81 @@ class ConnectionPool
   MAX_RETRIES = 3
 
   def initialize(size)
-    @size = size
+    @pool = SizedQueue.new(size)
-    @pool = Concurrent::Map.new
+    size.times { @pool << build_connection_entry }
-    @creation_times = Concurrent::Map.new
     @cleanup_thread = schedule_cleanup
   end
 
-  def with_connection(&block)
+  def with_connection
-    conn = acquire_connection
+    entry = acquire_connection
     begin
-      yield conn
+      yield entry[:http]
     ensure
-      release_connection(conn)
+      release_connection(entry)
     end
   end
 
   def shutdown
     @cleanup_thread&.exit
-    @pool.each_value { |conn| conn.finish if conn&.started? }
+    drain_pool { |entry| safe_finish(entry[:http]) }
-    @pool.clear
-    @creation_times.clear
   end
 
   private
 
   def acquire_connection
-    thread_id = Thread.current.object_id
+    entry = @pool.pop
-    conn = @pool[thread_id]
+    if stale?(entry)
-
+      safe_finish(entry[:http])
-    if should_create_new?(conn)
+      entry = build_connection_entry
-      conn&.finish if conn&.started?
-      conn = create_connection
-      @pool[thread_id] = conn
-      @creation_times[thread_id] = Time.now
     end
-
+    entry
-    conn
   end
 
-  def release_connection(conn)
+  def release_connection(entry)
-    return unless conn
+    if stale?(entry)
-    if conn.started? && Time.now - @creation_times[Thread.current.object_id] > MAX_AGE
+      safe_finish(entry[:http])
-      conn.finish
+      entry = build_connection_entry
-      @pool.delete(Thread.current.object_id)
+    end
-      @creation_times.delete(Thread.current.object_id)
+    @pool << entry
+  end
+
+  def stale?(entry)
+    http = entry[:http]
+    !http.started? || (Time.now - entry[:created_at] > MAX_AGE)
+  end
+
+  def build_connection_entry
+    { http: create_connection, created_at: Time.now }
+  end
+
+  def safe_finish(http)
+    http.finish if http&.started?
+  rescue StandardError
+    nil
+  end
+
+  def drain_pool
+    loop do
+      entry = begin
+        @pool.pop(true)
+      rescue ThreadError
+        break
+      end
+      yield(entry)
     end
   end
 
-  def should_create_new?(conn)
+  def cleanup_old_connections
-    return true if conn.nil?
+    entry = begin
-    return true unless conn.started?
+      @pool.pop(true)
-    return true if Time.now - @creation_times[Thread.current.object_id] > MAX_AGE
+    rescue ThreadError
-    false
+      return
-  end
+    end
-
+    if stale?(entry)
-  def create_connection
+      safe_finish(entry[:http])
-    http = Net::HTTP.new("web.archive.org", 443)
+      entry = build_connection_entry
-    http.use_ssl = true
+    end
-    http.read_timeout = DEFAULT_TIMEOUT
+    @pool << entry
-    http.open_timeout = DEFAULT_TIMEOUT
-    http.keep_alive_timeout = 30
-    http.max_retries = MAX_RETRIES
-    http.start
-    http
   end
 
   def schedule_cleanup
@@ -99,16 +111,15 @@ class ConnectionPool
     end
   end
 
-  def cleanup_old_connections
+  def create_connection
-    current_time = Time.now
+    http = Net::HTTP.new("web.archive.org", 443)
-    @creation_times.each do |thread_id, creation_time|
+    http.use_ssl = true
-      if current_time - creation_time > MAX_AGE
+    http.read_timeout = DEFAULT_TIMEOUT
-        conn = @pool[thread_id]
+    http.open_timeout = DEFAULT_TIMEOUT
-        conn&.finish if conn&.started?
+    http.keep_alive_timeout = 30
-        @pool.delete(thread_id)
+    http.max_retries = MAX_RETRIES
-        @creation_times.delete(thread_id)
+    http.start
-      end
+    http
-    end
   end
 end