diff --git a/lib/mongo.rb b/lib/mongo.rb index 3226eb3..3187b49 100644 --- a/lib/mongo.rb +++ b/lib/mongo.rb @@ -61,6 +61,7 @@ require 'mongo/util/pool' require 'mongo/util/pool_manager' require 'mongo/util/server_version' require 'mongo/util/ssl_socket' +require 'mongo/util/tcp_socket' require 'mongo/util/uri_parser' require 'mongo/collection' diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb index 2113070..bedf792 100644 --- a/lib/mongo/connection.rb +++ b/lib/mongo/connection.rb @@ -537,7 +537,7 @@ module Mongo if @ssl @socket_class = Mongo::SSLSocket else - @socket_class = ::TCPSocket + @socket_class = Mongo::TCPSocket end # Authentication objects diff --git a/lib/mongo/util/pool.rb b/lib/mongo/util/pool.rb index 52782a8..92ccaf8 100644 --- a/lib/mongo/util/pool.rb +++ b/lib/mongo/util/pool.rb @@ -167,6 +167,7 @@ module Mongo begin socket = self.connection.socket_class.new(@host, @port) socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1) + socket.pool = self rescue => ex socket.close if socket raise ConnectionFailure, "Failed to connect to host #{@host} and port #{@port}: #{ex}" diff --git a/lib/mongo/util/ssl_socket.rb b/lib/mongo/util/ssl_socket.rb index 2a4bcbe..9861ae1 100644 --- a/lib/mongo/util/ssl_socket.rb +++ b/lib/mongo/util/ssl_socket.rb @@ -7,6 +7,8 @@ module Mongo # mirroring Ruby's TCPSocket, vis., TCPSocket#send and TCPSocket#read. class SSLSocket + attr_accessor :pool + def initialize(host, port) @socket = ::TCPSocket.new(host, port) @ssl = OpenSSL::SSL::SSLSocket.new(@socket) @@ -33,6 +35,5 @@ module Mongo def close @ssl.close end - end end diff --git a/lib/mongo/util/tcp_socket.rb b/lib/mongo/util/tcp_socket.rb new file mode 100644 index 0000000..7d8d97d --- /dev/null +++ b/lib/mongo/util/tcp_socket.rb @@ -0,0 +1,6 @@ +module Mongo + class TCPSocket < ::TCPSocket + attr_accessor :pool + + end +end