RUBY-2132 add address, connection generation & monitoring flag fields to sockets for diagnostics (#1946)

* rename address to human_address

* rework socket closes in encryption i/o

* track connection address, generation and monitoring flag on sockets

* move constructors to the top

* extract a base constructor

* add a socket summary method

* fix unit test

* fix ssl socket tests

Co-authored-by: Oleg Pudeyev <oleg@bsdpower.com>

Author: p-mongo

lib/mongo/address.rb

@@ -192,7 +192,7 @@ def socket(socket_timeout, ssl_options = {}, options = {})
 
         options = {
           connect_timeout: Server::CONNECT_TIMEOUT,
-        }.update(options)
+        }.update(Hash[options.map { |k, v| [k.to_sym, v] }])
 
         # When the driver connects to "localhost", it only attempts IPv4
         # connections. When the driver connects to other hosts, it will

lib/mongo/crypt/encryption_io.rb

@@ -251,38 +251,55 @@ def with_ssl_socket(endpoint)
         host, port = endpoint.split(':')
         port ||= 443 # Default port for AWS KMS API
 
+        # Create TCPSocket and set nodelay option
+        tcp_socket = TCPSocket.open(host, port)
         begin
-          # Create TCPSocket and set nodelay option
-          tcp_socket = TCPSocket.open(host, port)
           tcp_socket.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1)
 
           ssl_socket = OpenSSL::SSL::SSLSocket.new(tcp_socket)
-          ssl_socket.sync_close = true # tcp_socket will be closed when ssl_socket is closed
-          ssl_socket.hostname = "#{host}:#{port}" # perform SNI
+          begin
+            # tcp_socket will be closed when ssl_socket is closed
+            ssl_socket.sync_close = true
+            # perform SNI
+            ssl_socket.hostname = "#{host}:#{port}"
+
+            Timeout.timeout(
+              SOCKET_TIMEOUT,
+              Error::SocketTimeoutError,
+              "KMS socket connection timed out after #{SOCKET_TIMEOUT} seconds",
+            ) do
+              ssl_socket.connect
+            end
 
-          Timeout.timeout(
-            SOCKET_TIMEOUT,
-            Error::SocketTimeoutError,
-            'Socket connection timed out'
-          ) do
-            ssl_socket.connect
+            yield(ssl_socket)
+          ensure
+            begin
+              Timeout.timeout(
+                SOCKET_TIMEOUT,
+                Error::SocketTimeoutError,
+                'KMS SSL socket close timed out'
+              ) do
+                ssl_socket.sysclose
+              end
+            rescue
+            end
           end
-
-          yield(ssl_socket)
-        rescue => e
-          raise Error::KmsError, "Error decrypting data key. #{e.class}: #{e.message}"
         ensure
-          # If there is an error during socket creation, the
-          # ssl_socket object won't exist in this scope and this line will
-          # raise an exception
-          Timeout.timeout(
-            SOCKET_TIMEOUT,
-            Error::SocketTimeoutError,
-            'Socket close timed out'
-          ) do
-            ssl_socket.sysclose rescue nil
+          # Still close tcp socket manually in case ssl socket creation
+          # fails.
+          begin
+            Timeout.timeout(
+              SOCKET_TIMEOUT,
+              Error::SocketTimeoutError,
+              'KMS TCP socket close timed out'
+            ) do
+              tcp_socket.close
+            end
+          rescue
           end
         end
+      rescue => e
+        raise Error::KmsError, "Error decrypting data key: #{e.class}: #{e.message}"
       end
     end
   end

lib/mongo/server/connection.rb

@@ -181,7 +181,8 @@ def connect!
       #   returned socket.
       private def do_connect
         socket = add_server_diagnostics do
-          address.socket(socket_timeout, ssl_options, address.options)
+          address.socket(socket_timeout, ssl_options, address.options.merge(
+            connection_address: address, connection_generation: generation).update(ssl_options))
         end
 
         begin

lib/mongo/server/monitor/connection.rb

@@ -162,7 +162,8 @@ def connect!
           end
 
           @socket = add_server_diagnostics do
-            address.socket(socket_timeout, ssl_options, address.options)
+            address.socket(socket_timeout, ssl_options, address.options.merge(
+              connection_address: address, monitor: true).update(ssl_options))
           end
           true
         end

lib/mongo/socket.rb

@@ -45,6 +45,25 @@ class Socket
     # @api private
     WRITE_CHUNK_SIZE = 65536
 
+    # Initializes common socket attributes.
+    #
+    # @param [ Float ] timeout The socket timeout value.
+    # @param [ Hash ] options The options.
+    #
+    # @option options [ Float ] :connect_timeout Connect timeout.
+    # @option options [ Address ] :connection_address Address of the
+    #   connection that created this socket.
+    # @option options [ Integer ] :connection_generation Generation of the
+    #   connection (for non-monitoring connections) that created this socket.
+    # @option options [ true | false ] :monitor Whether this socket was
+    #   created by a monitoring connection.
+    #
+    # @api private
+    def initialize(timeout, options)
+      @timeout = timeout
+      @options = options
+    end
+
     # @return [ Integer ] family The type of host family.
     attr_reader :family
 
@@ -57,6 +76,41 @@ class Socket
     # @return [ Float ] timeout The socket timeout.
     attr_reader :timeout
 
+    # @return [ Address ] Address of the connection that created this socket.
+    #
+    # @api private
+    def connection_address
+      options[:connection_address]
+    end
+
+    # @return [ Integer ] Generation of the connection (for non-monitoring
+    #   connections) that created this socket.
+    #
+    # @api private
+    def connection_generation
+      options[:connection_generation]
+    end
+
+    # @return [ true | false ] Whether this socket was created by a monitoring
+    #   connection.
+    #
+    # @api private
+    def monitor?
+      !!options[:monitor]
+    end
+
+    # @return [ String ] Human-readable summary of the socket for debugging.
+    #
+    # @api private
+    def summary
+      fileno = @socket&.fileno rescue '<no socket>' || '<no socket>'
+      if monitor?
+        "#{connection_address}:m #{fileno}"
+      else
+        "#{connection_address}:c:#{connection_generation} #{fileno}"
+      end
+    end
+
     # Is the socket connection alive?
     #
     # @example Is the socket alive?
@@ -350,15 +404,15 @@ def map_exceptions
       begin
         yield
       rescue Errno::ETIMEDOUT => e
-        raise Error::SocketTimeoutError, "#{e.class}: #{e} (for #{address})"
+        raise Error::SocketTimeoutError, "#{e.class}: #{e} (for #{human_address})"
       rescue IOError, SystemCallError => e
-        raise Error::SocketError, "#{e.class}: #{e} (for #{address})"
+        raise Error::SocketError, "#{e.class}: #{e} (for #{human_address})"
       rescue OpenSSL::SSL::SSLError => e
-        raise Error::SocketError, "#{e.class}: #{e} (for #{address}) (#{SSL_ERROR})"
+        raise Error::SocketError, "#{e.class}: #{e} (for #{human_address}) (#{SSL_ERROR})"
       end
     end
 
-    def address
+    def human_address
       raise NotImplementedError
     end
   end

lib/mongo/socket/ssl.rb

@@ -21,6 +21,42 @@ class Socket
     class SSL < Socket
       include OpenSSL
 
+      # Initializes a new SSL socket.
+      #
+      # @example Create the SSL socket.
+      #   SSL.new('::1', 27017, 30)
+      #
+      # @param [ String ] host The hostname or IP address.
+      # @param [ Integer ] port The port number.
+      # @param [ Float ] timeout The socket timeout value.
+      # @param [ Integer ] family The socket family.
+      # @param [ Hash ] options The options.
+      #
+      # @option options [ Float ] :connect_timeout Connect timeout.
+      # @option options [ Address ] :connection_address Address of the
+      #   connection that created this socket.
+      # @option options [ Integer ] :connection_generation Generation of the
+      #   connection (for non-monitoring connections) that created this socket.
+      # @option options [ true | false ] :monitor Whether this socket was
+      #   created by a monitoring connection.
+      #
+      # @since 2.0.0
+      def initialize(host, port, host_name, timeout, family, options = {})
+        super(timeout, options)
+        @host, @port, @host_name = host, port, host_name
+        @context = create_context(options)
+        @family = family
+        @tcp_socket = ::Socket.new(family, SOCK_STREAM, 0)
+        begin
+          @tcp_socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
+          set_socket_options(@tcp_socket)
+          connect!
+        rescue
+          @tcp_socket.close
+          raise
+        end
+      end
+
       # @return [ SSLContext ] context The ssl context.
       attr_reader :context
 
@@ -67,35 +103,6 @@ def connect!
       end
       private :connect!
 
-      # Initializes a new SSL socket.
-      #
-      # @example Create the SSL socket.
-      #   SSL.new('::1', 27017, 30)
-      #
-      # @param [ String ] host The hostname or IP address.
-      # @param [ Integer ] port The port number.
-      # @param [ Float ] timeout The socket timeout value.
-      # @param [ Integer ] family The socket family.
-      # @param [ Hash ] options The options.
-      #
-      # @option options [ Float ] :connect_timeout Connect timeout.
-      #
-      # @since 2.0.0
-      def initialize(host, port, host_name, timeout, family, options = {})
-        @host, @port, @host_name, @timeout, @options = host, port, host_name, timeout, options
-        @context = create_context(options)
-        @family = family
-        @tcp_socket = ::Socket.new(family, SOCK_STREAM, 0)
-        begin
-          @tcp_socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
-          set_socket_options(@tcp_socket)
-          connect!
-        rescue
-          @tcp_socket.close
-          raise
-        end
-      end
-
       # Read a single byte from the socket.
       #
       # @example Read a single byte.
@@ -292,7 +299,7 @@ def read_buffer_size
         16384
       end
 
-      def address
+      def human_address
         "#{host}:#{port} (#{host_name}:#{port}, TLS)"
       end
     end

lib/mongo/socket/tcp.rb

@@ -20,6 +20,41 @@ class Socket
     # @since 2.0.0
     class TCP < Socket
 
+      # Initializes a new TCP socket.
+      #
+      # @example Create the TCP socket.
+      #   TCP.new('::1', 27017, 30, Socket::PF_INET)
+      #   TCP.new('127.0.0.1', 27017, 30, Socket::PF_INET)
+      #
+      # @param [ String ] host The hostname or IP address.
+      # @param [ Integer ] port The port number.
+      # @param [ Float ] timeout The socket timeout value.
+      # @param [ Integer ] family The socket family.
+      # @param [ Hash ] options The options.
+      #
+      # @option options [ Float ] :connect_timeout Connect timeout.
+      # @option options [ Address ] :connection_address Address of the
+      #   connection that created this socket.
+      # @option options [ Integer ] :connection_generation Generation of the
+      #   connection (for non-monitoring connections) that created this socket.
+      # @option options [ true | false ] :monitor Whether this socket was
+      #   created by a monitoring connection.
+      #
+      # @since 2.0.0
+      def initialize(host, port, timeout, family, options = {})
+        super(timeout, options)
+        @host, @port = host, port
+        @family = family
+        @socket = ::Socket.new(family, SOCK_STREAM, 0)
+        begin
+          set_socket_options(@socket)
+          connect!
+        rescue
+          @socket.close
+          raise
+        end
+      end
+
       # @return [ String ] host The host to connect to.
       attr_reader :host
 
@@ -48,37 +83,9 @@ def connect!
       end
       private :connect!
 
-      # Initializes a new TCP socket.
-      #
-      # @example Create the TCP socket.
-      #   TCP.new('::1', 27017, 30, Socket::PF_INET)
-      #   TCP.new('127.0.0.1', 27017, 30, Socket::PF_INET)
-      #
-      # @param [ String ] host The hostname or IP address.
-      # @param [ Integer ] port The port number.
-      # @param [ Float ] timeout The socket timeout value.
-      # @param [ Integer ] family The socket family.
-      # @param [ Hash ] options The options.
-      #
-      # @option options [ Float ] :connect_timeout Connect timeout.
-      #
-      # @since 2.0.0
-      def initialize(host, port, timeout, family, options = {})
-        @host, @port, @timeout, @options = host, port, timeout, options
-        @family = family
-        @socket = ::Socket.new(family, SOCK_STREAM, 0)
-        begin
-          set_socket_options(@socket)
-          connect!
-        rescue
-          @socket.close
-          raise
-        end
-      end
-
       private
 
-      def address
+      def human_address
         "#{host}:#{port} (no TLS)"
       end
     end