sequel 3.21.0 → 3.41.0

data/lib/sequel/connection_pool/sharded_single.rb

@@ -1,18 +1,15 @@
 # A ShardedSingleConnectionPool is a single threaded connection pool that
 # works with multiple shards/servers.
 class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
-  # Initializes the instance with the supplied block as the connection_proc.
-  #
   # The single threaded pool takes the following options:
   #
   # * :servers - A hash of servers to use.  Keys should be symbols.  If not
-  #   present, will use a single :default server.  The server name symbol will
-  #   be passed to the connection_proc.
+  #   present, will use a single :default server.
   # * :servers_hash - The base hash to use for the servers.  By default,
   #   Sequel uses Hash.new(:default).  You can use a hash with a default proc
   #   that raises an error if you want to catch all cases where a nonexistent
   #   server is used.
-  def initialize(opts={}, &block)
+  def initialize(db, opts={})
     super
     @conns = {}
     @servers = opts.fetch(:servers_hash, Hash.new(:default))
@@ -26,6 +23,11 @@ class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
   def add_servers(servers)
     servers.each{|s| @servers[s] = s}
   end
+
+  # Yield all of the currently established connections
+  def all_connections
+    @conns.values.each{|c| yield c}
+  end
   
   # The connection for the given server.
   def conn(server=:default)
@@ -36,19 +38,18 @@ class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
   # #hold, the connection is reestablished. Options:
   # * :server - Should be a symbol specifing the server to disconnect from,
   #   or an array of symbols to specify multiple servers.
-  def disconnect(opts={}, &block)
-    block ||= @disconnection_proc
-    (opts[:server] ? Array(opts[:server]) : servers).each{|s| disconnect_server(s, &block)}
+  def disconnect(opts={})
+    (opts[:server] ? Array(opts[:server]) : servers).each{|s| disconnect_server(s)}
   end
   
   # Yields the connection to the supplied block for the given server.
   # This method simulates the ConnectionPool#hold API.
   def hold(server=:default)
     begin
-      server = @servers[server]
+      server = pick_server(server)
       yield(@conns[server] ||= make_new(server))
     rescue Sequel::DatabaseDisconnectError
-      disconnect_server(server, &@disconnection_proc)
+      disconnect_server(server)
       raise
     end
   end
@@ -60,7 +61,7 @@ class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
   def remove_servers(servers)
     raise(Sequel::Error, "cannot remove default server") if servers.include?(:default)
     servers.each do |server|
-      disconnect_server(server, &@disconnection_proc)
+      disconnect_server(server)
       @servers.delete(server)
     end
   end
@@ -75,14 +76,23 @@ class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
     @conns.length
   end
   
+  def pool_type
+    :sharded_single
+  end
+  
   private
   
   # Disconnect from the given server, if connected.
-  def disconnect_server(server, &block)
+  def disconnect_server(server)
     if conn = @conns.delete(server)
-      block.call(conn) if block
+      db.disconnect_connection(conn)
     end
   end
+
+  # If the server given is in the hash, return it, otherwise, return the default server.
+  def pick_server(server)
+    @servers[server]
+  end
   
   CONNECTION_POOL_MAP[[true, true]] = self
 end

data/lib/sequel/connection_pool/sharded_threaded.rb

@@ -8,13 +8,12 @@ Sequel.require 'connection_pool/threaded'
 class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
   # The following additional options are respected:
   # * :servers - A hash of servers to use.  Keys should be symbols.  If not
-  #   present, will use a single :default server.  The server name symbol will
-  #   be passed to the connection_proc.
+  #   present, will use a single :default server.
   # * :servers_hash - The base hash to use for the servers.  By default,
   #   Sequel uses Hash.new(:default).  You can use a hash with a default proc
   #   that raises an error if you want to catch all cases where a nonexistent
   #   server is used.
-  def initialize(opts = {}, &block)
+  def initialize(db, opts = {})
     super
     @available_connections = {}
     @connections_to_remove = []
@@ -45,6 +44,23 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
     @allocated[server]
   end
   
+  # Yield all of the available connections, and the ones currently allocated to
+  # this thread.  This will not yield connections currently allocated to other
+  # threads, as it is not safe to operate on them.  This holds the mutex while
+  # it is yielding all of the connections, which means that until
+  # the method's block returns, the pool is locked.
+  def all_connections
+    t = Thread.current
+    sync do
+      @allocated.values.each do |threads|
+        threads.each do |thread, conn|
+          yield conn if t == thread
+        end
+      end
+      @available_connections.values.each{|v| v.each{|c| yield c}}
+    end
+  end
+  
   # An array of connections opened but not currently used, for the given
   # server. Nonexistent servers will return nil. Treat this as read only, do
   # not modify the resulting object.
@@ -70,11 +86,10 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
   # creates new connections to the database. Options:
   # * :server - Should be a symbol specifing the server to disconnect from,
   #   or an array of symbols to specify multiple servers.
-  def disconnect(opts={}, &block)
-    block ||= @disconnection_proc
+  def disconnect(opts={})
     sync do
       (opts[:server] ? Array(opts[:server]) : @servers.keys).each do |s|
-        disconnect_server(s, &block)
+        disconnect_server(s)
       end
     end
   end
@@ -94,7 +109,7 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
   # connection can be acquired, a Sequel::PoolTimeout is 
   # raised.
   def hold(server=:default)
-    sync{server = @servers[server]}
+    server = pick_server(server)
     t = Thread.current
     if conn = owned_connection(t, server)
       return yield(conn)
@@ -128,7 +143,7 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
       raise(Sequel::Error, "cannot remove default server") if servers.include?(:default)
       servers.each do |server|
         if @servers.include?(server)
-          disconnect_server(server, &@disconnection_proc)
+          disconnect_server(server)
           @available_connections.delete(server)
           @allocated.delete(server)
           @servers.delete(server)
@@ -142,6 +157,10 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
     sync{@servers.keys}
   end
 
+  def pool_type
+    :sharded_threaded
+  end
+  
   private
 
   # Assigns a connection to the supplied thread for the given server, if one
@@ -159,16 +178,30 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
   # available, tries to create a new connection. The calling code should already
   # have the mutex before calling this.
   def available(server)
-    available_connections(server).pop || make_new(server)
+    next_available(server) || make_new(server)
+  end
+
+  # Return a connection to the pool of available connections for the server,
+  # returns the connection. The calling code should already have the mutex
+  # before calling this.
+  def checkin_connection(server, conn)
+    case @connection_handling
+    when :queue
+      available_connections(server).unshift(conn)
+    else
+      available_connections(server) << conn
+    end
+
+    conn
   end
 
   # Disconnect from the given server.  Disconnects available connections
   # immediately, and schedules currently allocated connections for disconnection
   # as soon as they are returned to the pool. The calling code should already
   # have the mutex before calling this.
-  def disconnect_server(server, &block)
+  def disconnect_server(server)
     if conns = available_connections(server)
-      conns.each{|conn| block.call(conn)} if block
+      conns.each{|conn| db.disconnect_connection(conn)}
       conns.clear
     end
     @connections_to_remove.concat(allocated(server).values)
@@ -184,12 +217,24 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
     end
     default_make_new(server) if (n || size(server)) < @max_size
   end
+
+  # Return the next available connection in the pool for the given server, or nil
+  # if there is not currently an available connection for the server.
+  # The calling code should already have the mutex before calling this.
+  def next_available(server)
+    available_connections(server).pop
+  end
   
   # Returns the connection owned by the supplied thread for the given server,
   # if any. The calling code should NOT already have the mutex before calling this.
   def owned_connection(thread, server)
     sync{@allocated[server][thread]}
   end
+
+  # If the server given is in the hash, return it, otherwise, return the default server.
+  def pick_server(server)
+    sync{@servers[server]}
+  end
   
   # Releases the connection assigned to the supplied thread and server. If the
   # server or connection given is scheduled for disconnection, remove the
@@ -199,7 +244,13 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
     if @connections_to_remove.include?(conn)
       remove(thread, conn, server)
     else
-      available_connections(server) << allocated(server).delete(thread)
+      conn = allocated(server).delete(thread)
+
+      if @connection_handling == :disconnect
+        db.disconnect_connection(conn)
+      else
+        checkin_connection(server, conn)
+      end
     end
   end
 
@@ -208,7 +259,7 @@ class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
   def remove(thread, conn, server)
     @connections_to_remove.delete(conn)
     allocated(server).delete(thread) if @servers.include?(server)
-    @disconnection_proc.call(conn) if @disconnection_proc
+    db.disconnect_connection(conn)
   end
   
   CONNECTION_POOL_MAP[[false, true]] = self

data/lib/sequel/connection_pool/single.rb

@@ -8,11 +8,15 @@ class Sequel::SingleConnectionPool < Sequel::ConnectionPool
     @conn ? 1 : 0
   end
   
+  # Yield the connection if one has been made.
+  def all_connections
+    yield @conn if @conn
+  end
+  
   # Disconnect the connection from the database.
-  def disconnect(opts=nil, &block)
+  def disconnect(opts=nil)
     return unless @conn
-    block ||= @disconnection_proc
-    block.call(@conn) if block
+    db.disconnect_connection(@conn)
     @conn = nil
   end
   
@@ -26,5 +30,9 @@ class Sequel::SingleConnectionPool < Sequel::ConnectionPool
     end
   end
 
+  def pool_type
+    :single
+  end
+  
   CONNECTION_POOL_MAP[[true, false]] = self
 end

data/lib/sequel/connection_pool/threaded.rb

@@ -13,45 +13,65 @@ class Sequel::ThreadedConnectionPool < Sequel::ConnectionPool
   attr_reader :allocated
 
   # The following additional options are respected:
+  # * :connection_handling - Set how to handle available connections.  By default,
+  #   uses a a stack for performance.  Can be set to :queue to use a queue, which reduces
+  #   the chances of connections becoming stale.
   # * :max_connections - The maximum number of connections the connection pool
   #   will open (default 4)
   # * :pool_sleep_time - The amount of time to sleep before attempting to acquire
   #   a connection again (default 0.001)
   # * :pool_timeout - The amount of seconds to wait to acquire a connection
   #   before raising a PoolTimeoutError (default 5)
-  def initialize(opts = {}, &block)
+  def initialize(db, opts = {})
     super
     @max_size = Integer(opts[:max_connections] || 4)
     raise(Sequel::Error, ':max_connections must be positive') if @max_size < 1
     @mutex = Mutex.new  
+    @connection_handling = opts[:connection_handling]
     @available_connections = []
     @allocated = {}
     @timeout = Integer(opts[:pool_timeout] || 5)
     @sleep_time = Float(opts[:pool_sleep_time] || 0.001)
   end
   
-  # The total number of connections opened for the given server, should
-  # be equal to available_connections.length + allocated.length.
+  # The total number of connections opened, either available or allocated.
+  # This may not be completely accurate as it isn't protected by the mutex.
   def size
     @allocated.length + @available_connections.length
   end
   
-  # Removes all connections currently available on all servers, optionally
+  # Yield all of the available connections, and the one currently allocated to
+  # this thread.  This will not yield connections currently allocated to other
+  # threads, as it is not safe to operate on them.  This holds the mutex while
+  # it is yielding all of the available connections, which means that until
+  # the method's block returns, the pool is locked.
+  def all_connections
+    hold do |c|
+      sync do
+        yield c
+        @available_connections.each{|c| yield c}
+      end
+    end
+  end
+  
+  # Removes all connections currently available, optionally
   # yielding each connection to the given block. This method has the effect of 
   # disconnecting from the database, assuming that no connections are currently
-  # being used.
+  # being used.  If you want to be able to disconnect connections that are
+  # currently in use, use the ShardedThreadedConnectionPool, which can do that.
+  # This connection pool does not, for performance reasons. To use the sharded pool,
+  # pass the <tt>:servers=>{}</tt> option when connecting to the database.
   # 
   # Once a connection is requested using #hold, the connection pool
   # creates new connections to the database.
-  def disconnect(opts={}, &block)
-    block ||= @disconnection_proc
+  def disconnect(opts={})
     sync do
-      @available_connections.each{|conn| block.call(conn)} if block
+      @available_connections.each{|conn| db.disconnect_connection(conn)}
       @available_connections.clear
     end
   end
   
-  # Chooses the first available connection to the given server, or if none are
+  # Chooses the first available connection, or if none are
   # available, creates a new connection.  Passes the connection to the supplied
   # block:
   # 
@@ -83,18 +103,23 @@ class Sequel::ThreadedConnectionPool < Sequel::ConnectionPool
       end
       yield conn
     rescue Sequel::DatabaseDisconnectError
-      @disconnection_proc.call(conn) if @disconnection_proc && conn
-      @allocated.delete(t)
+      oconn = conn
       conn = nil
+      db.disconnect_connection(oconn) if oconn
+      @allocated.delete(t)
       raise
     ensure
       sync{release(t)} if conn
     end
   end
+
+  def pool_type
+    :threaded
+  end
   
   private
 
-  # Assigns a connection to the supplied thread for the given server, if one
+  # Assigns a connection to the supplied thread, if one
   # is available. The calling code should NOT already have the mutex when
   # calling this.
   def acquire(thread)
@@ -105,11 +130,24 @@ class Sequel::ThreadedConnectionPool < Sequel::ConnectionPool
     end
   end
   
-  # Returns an available connection to the given server. If no connection is
+  # Returns an available connection. If no connection is
   # available, tries to create a new connection. The calling code should already
   # have the mutex before calling this.
   def available
-    @available_connections.pop || make_new(DEFAULT_SERVER)
+    next_available || make_new(DEFAULT_SERVER)
+  end
+
+  # Return a connection to the pool of available connections, returns the connection.
+  # The calling code should already have the mutex before calling this.
+  def checkin_connection(conn)
+    case @connection_handling
+    when :queue
+      @available_connections.unshift(conn)
+    else
+      @available_connections << conn
+    end
+
+    conn
   end
 
   # Alias the default make_new method, so subclasses can call it directly.
@@ -125,19 +163,30 @@ class Sequel::ThreadedConnectionPool < Sequel::ConnectionPool
     end
     super if (n || size) < @max_size
   end
-  
-  # Returns the connection owned by the supplied thread for the given server,
+
+  # Return the next available connection in the pool, or nil if there
+  # is not currently an available connection.  The calling code should already
+  # have the mutex before calling this.
+  def next_available
+    @available_connections.pop
+  end
+
+  # Returns the connection owned by the supplied thread,
   # if any. The calling code should NOT already have the mutex before calling this.
   def owned_connection(thread)
     sync{@allocated[thread]}
   end
   
-  # Releases the connection assigned to the supplied thread and server. If the
-  # server or connection given is scheduled for disconnection, remove the
-  # connection instead of releasing it back to the pool.
+  # Releases the connection assigned to the supplied thread back to the pool.
   # The calling code should already have the mutex before calling this.
   def release(thread)
-    @available_connections << @allocated.delete(thread)
+    conn = @allocated.delete(thread)
+
+    if @connection_handling == :disconnect
+      db.disconnect_connection(conn)
+    else
+      checkin_connection(conn)
+    end
   end
 
   # Yield to the block while inside the mutex. The calling code should NOT

data/lib/sequel/connection_pool.rb

@@ -2,15 +2,13 @@
 # on.  This class is not instantiated directly, but subclasses should at
 # the very least implement the following API:
 #
-# initialize(Hash, &block) :: The +block+ is used as the connection proc,
-#                             which should accept a single symbol argument.
+# initialize(Database, Hash) :: Initialize using the passed Sequel::Database
+#                               object and options hash.
 # hold(Symbol, &block) :: Yield a connection object (obtained from calling
 #                         the block passed to +initialize+) to the current block. For sharded
 #                         connection pools, the Symbol passed is the shard/server to use.
-# disconnect(Symbol, &block) :: Disconnect the connection object.  If a
-#                               block is given, pass the connection option to it, otherwise use the
-#                               <tt>:disconnection_proc</tt> option in the hash passed to initialize.  For sharded
-#                               connection pools, the Symbol passed is the shard/server to use.
+# disconnect(Symbol) :: Disconnect the connection object.  For sharded
+#                       connection pools, the Symbol passed is the shard/server to use.
 # servers :: An array of shard/server symbols for all shards/servers that this
 #            connection pool recognizes.
 # size :: an integer representing the total number of connections in the pool,
@@ -39,13 +37,13 @@ class Sequel::ConnectionPool
     # option is provided is provided, use that pool class, otherwise
     # use a new instance of an appropriate pool subclass based on the
     # <tt>:single_threaded</tt> and <tt>:servers</tt> options.
-    def get_pool(opts = {}, &block)
+    def get_pool(db, opts = {})
       case v = connection_pool_class(opts)
       when Class
-        v.new(opts, &block)
+        v.new(db, opts)
       when Symbol
         Sequel.ts_require("connection_pool/#{v}")
-        connection_pool_class(opts).new(opts, &block) || raise(Sequel::Error, "No connection pool class found")
+        connection_pool_class(opts).new(db, opts) || raise(Sequel::Error, "No connection pool class found")
       end
     end
     
@@ -53,11 +51,18 @@ class Sequel::ConnectionPool
     
     # Return a connection pool class based on the given options.
     def connection_pool_class(opts)
-      opts[:pool_class] || CONNECTION_POOL_MAP[[!!opts[:single_threaded], !!opts[:servers]]]
+      CONNECTION_POOL_MAP[opts[:pool_class]] || opts[:pool_class] || CONNECTION_POOL_MAP[[!!opts[:single_threaded], !!opts[:servers]]]
     end
   end
   extend ClassMethods
-  
+
+  # The after_connect proc used for this pool.  This is called with each new
+  # connection made, and is usually used to set custom per-connection settings.
+  attr_accessor :after_connect
+
+  # The Sequel::Database object tied to this connection pool.
+  attr_accessor :db
+
   # Instantiates a connection pool with the given options.  The block is called
   # with a single symbol (specifying the server/shard to use) every time a new
   # connection is needed.  The following options are respected for all connection
@@ -65,11 +70,8 @@ class Sequel::ConnectionPool
   # :after_connect :: The proc called after each new connection is made, with the
   #                   connection object, useful for customizations that you want to apply to all
   #                   connections.
-  # :disconnection_proc :: The proc called when removing connections from the pool,
-  #                        which is passed the connection to disconnect.
-  def initialize(opts={}, &block)
-    raise(Sequel::Error, "No connection proc specified") unless @connection_proc = block
-    @disconnection_proc = opts[:disconnection_proc]
+  def initialize(db, opts={})
+    @db = db
     @after_connect = opts[:after_connect]
   end
   
@@ -80,7 +82,7 @@ class Sequel::ConnectionPool
   
   # An array of symbols for all shards/servers, which is a single <tt>:default</tt> by default.
   def servers
-    [:default]
+    [DEFAULT_SERVER]
   end
   
   private
@@ -89,7 +91,7 @@ class Sequel::ConnectionPool
   # and checking for connection errors.
   def make_new(server)
     begin
-      conn = @connection_proc.call(server)
+      conn = @db.connect(server)
       @after_connect.call(conn) if @after_connect
     rescue Exception=>exception
       raise Sequel.convert_exception_class(exception, Sequel::DatabaseConnectionError)