From 93b2f3da9bc84fcbd6f61c0b954a64b2dbb55b07 Mon Sep 17 00:00:00 2001
From: Tyler Brock <tyler.brock@gmail.com>
Date: Fri, 2 Mar 2012 19:09:02 -0500
Subject: [PATCH] minor: Added TCPSocket Class w/ pool accessor

Sockets now know what pool they were checked out from
SSLSocket updated as well
---
 lib/mongo.rb                 | 1 +
 lib/mongo/connection.rb      | 2 +-
 lib/mongo/util/pool.rb       | 1 +
 lib/mongo/util/ssl_socket.rb | 3 ++-
 lib/mongo/util/tcp_socket.rb | 6 ++++++
 5 files changed, 11 insertions(+), 2 deletions(-)
 create mode 100644 lib/mongo/util/tcp_socket.rb

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