mongo 2.17.4 → 2.21.0

data/lib/mongo/socket/ssl.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
@@ -142,26 +142,37 @@ module Mongo
       #
       # @since 2.0.0
       def connect!
-        Timeout.timeout(options[:connect_timeout], Error::SocketTimeoutError, "The socket took over #{options[:connect_timeout]} seconds to connect") do
-          map_exceptions do
-            @tcp_socket.connect(::Socket.pack_sockaddr_in(port, host))
-          end
-          @socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, context)
-          begin
-            @socket.hostname = @host_name
-            @socket.sync_close = true
-            map_exceptions do
-              @socket.connect
+        sockaddr = ::Socket.pack_sockaddr_in(port, host)
+        connect_timeout = options[:connect_timeout]
+        map_exceptions do
+          if connect_timeout && connect_timeout != 0
+            deadline = Utils.monotonic_time + connect_timeout
+            if BSON::Environment.jruby?
+              # We encounter some strange problems with connect_nonblock for
+              # ssl sockets on JRuby. Therefore, we use the old +Timeout.timeout+
+              # solution, even though it is known to be not very reliable.
+              raise Error::SocketTimeoutError, 'connect_timeout expired' if connect_timeout < 0
+
+              Timeout.timeout(connect_timeout, Error::SocketTimeoutError, "The socket took over #{options[:connect_timeout]} seconds to connect") do
+                connect_without_timeout(sockaddr)
+              end
+            else
+              connect_with_timeout(sockaddr, connect_timeout)
             end
+            remaining_timeout = deadline - Utils.monotonic_time
+            verify_certificate!(@socket)
+            verify_ocsp_endpoint!(@socket, remaining_timeout)
+          else
+            connect_without_timeout(sockaddr)
             verify_certificate!(@socket)
             verify_ocsp_endpoint!(@socket)
-          rescue
-            @socket.close
-            @socket = nil
-            raise
           end
-          self
         end
+        self
+      rescue
+        @socket&.close
+        @socket = nil
+        raise
       end
       private :connect!
 
@@ -182,6 +193,87 @@ module Mongo
 
       private
 
+      # Connects the socket without a timeout provided.
+      #
+      # @param [ String ] sockaddr Address to connect to.
+      def connect_without_timeout(sockaddr)
+        @tcp_socket.connect(sockaddr)
+        @socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, context)
+        @socket.hostname = @host_name
+        @socket.sync_close = true
+        @socket.connect
+      end
+
+      # Connects the socket with the connect timeout. The timeout applies to
+      # connecting both ssl socket and the underlying tcp socket.
+      #
+      # @param [ String ] sockaddr Address to connect to.
+      def connect_with_timeout(sockaddr, connect_timeout)
+        if connect_timeout <= 0
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+
+        deadline = Utils.monotonic_time + connect_timeout
+        connect_tcp_socket_with_timeout(sockaddr, deadline, connect_timeout)
+        connnect_ssl_socket_with_timeout(deadline, connect_timeout)
+      end
+
+      def connect_tcp_socket_with_timeout(sockaddr, deadline, connect_timeout)
+        if deadline <= Utils.monotonic_time
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+        begin
+          @tcp_socket.connect_nonblock(sockaddr)
+        rescue IO::WaitWritable
+          with_select_timeout(deadline, connect_timeout) do |select_timeout|
+            IO.select(nil, [@tcp_socket], nil, select_timeout)
+          end
+          retry
+        rescue Errno::EISCONN
+          # Socket is connected, nothing to do.
+        end
+      end
+
+      def connnect_ssl_socket_with_timeout(deadline, connect_timeout)
+        if deadline <= Utils.monotonic_time
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+        @socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, context)
+        @socket.hostname = @host_name
+        @socket.sync_close = true
+
+        # We still have time, connecting ssl socket.
+        begin
+          @socket.connect_nonblock
+        rescue IO::WaitReadable, OpenSSL::SSL::SSLErrorWaitReadable
+          with_select_timeout(deadline, connect_timeout) do |select_timeout|
+            IO.select([@socket], nil, nil, select_timeout)
+          end
+          retry
+        rescue IO::WaitWritable, OpenSSL::SSL::SSLErrorWaitWritable
+          with_select_timeout(deadline, connect_timeout) do |select_timeout|
+            IO.select(nil, [@socket], nil, select_timeout)
+          end
+          retry
+        rescue Errno::EISCONN
+          # Socket is connected, nothing to do
+        end
+      end
+
+      # Raises +Error::SocketTimeoutError+ exception if deadline reached or the
+      # block returns nil. The block should call +IO.select+ with the
+      # +connect_timeout+ value. It returns nil if the +connect_timeout+ expires.
+      def with_select_timeout(deadline, connect_timeout, &block)
+        select_timeout = deadline - Utils.monotonic_time
+        if select_timeout <= 0
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+        rv = block.call(select_timeout)
+        if rv.nil?
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+      end
+
       def verify_certificate?
         # If ssl_verify_certificate is not present, disable only if
         # ssl_verify is explicitly set to false.
@@ -287,7 +379,7 @@ module Mongo
         # for instance, if there is no newline between two certificates
         # this code will extract them both but OpenSSL fails in this situation.
         if cert_text
-          certs = cert_text.scan(/-----BEGIN CERTIFICATE-----(?:.|\n)+?-----END CERTIFICATE-----/)
+          certs = extract_certs(cert_text)
           if certs.length > 1
             context.cert = OpenSSL::X509::Certificate.new(certs.shift)
             context.extra_chain_cert = certs.map do |cert|
@@ -362,7 +454,7 @@ module Mongo
         end
       end
 
-      def verify_ocsp_endpoint!(socket)
+      def verify_ocsp_endpoint!(socket, timeout = nil)
         unless verify_ocsp_endpoint?
           return
         end
@@ -371,7 +463,7 @@ module Mongo
         ca_cert = socket.peer_cert_chain.last
 
         verifier = OcspVerifier.new(@host_name, cert, ca_cert, context.cert_store,
-          **Utils.shallow_symbolize_keys(options))
+          **Utils.shallow_symbolize_keys(options).merge(timeout: timeout))
         verifier.verify_with_cache
       end
 
@@ -390,6 +482,27 @@ module Mongo
           hook.call(@context)
         end
       end
+
+      BEGIN_CERT = "-----BEGIN CERTIFICATE-----"
+      END_CERT = "-----END CERTIFICATE-----"
+
+      # This was originally a scan + regex, but the regex was particularly
+      # inefficient and was flagged as a concern by static analysis.
+      def extract_certs(text)
+        [].tap do |list|
+          pos = 0
+
+          while (begin_idx = text.index(BEGIN_CERT, pos))
+            end_idx = text.index(END_CERT, begin_idx)
+            break unless end_idx
+
+            end_idx += END_CERT.length
+            list.push(text[begin_idx...end_idx])
+
+            pos = end_idx
+          end
+        end
+      end
     end
   end
 end

data/lib/mongo/socket/tcp.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
@@ -46,6 +46,9 @@ module Mongo
       # @since 2.0.0
       # @api private
       def initialize(host, port, timeout, family, options = {})
+        if family.nil?
+          raise ArgumentError, 'family must be specified'
+        end
         super(timeout, options)
         @host, @port = host, port
         @family = family
@@ -76,16 +79,50 @@ module Mongo
       # @return [ TCP ] The connected socket instance.
       #
       # @since 2.0.0
+      # @api private
       def connect!
-        Timeout.timeout(options[:connect_timeout], Error::SocketTimeoutError, "The socket took over #{options[:connect_timeout]} seconds to connect") do
-          socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
-          map_exceptions do
-            socket.connect(::Socket.pack_sockaddr_in(port, host))
+        socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
+        sockaddr = ::Socket.pack_sockaddr_in(port, host)
+        connect_timeout = options[:connect_timeout]
+        map_exceptions do
+          if connect_timeout && connect_timeout != 0
+            connect_with_timeout(sockaddr, connect_timeout)
+          else
+            connect_without_timeout(sockaddr)
+          end
+        end
+        self
+      end
+
+      # @api private
+      def connect_without_timeout(sockaddr)
+        socket.connect(sockaddr)
+      end
+
+      # @api private
+      def connect_with_timeout(sockaddr, connect_timeout)
+        if connect_timeout <= 0
+          raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+        end
+
+        deadline = Utils.monotonic_time + connect_timeout
+        begin
+          socket.connect_nonblock(sockaddr)
+        rescue IO::WaitWritable
+          select_timeout = deadline - Utils.monotonic_time
+          if select_timeout <= 0
+            raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
+          end
+          if IO.select(nil, [socket], nil, select_timeout)
+            retry
+          else
+            socket.close
+            raise Error::SocketTimeoutError, "The socket took over #{connect_timeout} seconds to connect"
           end
-          self
+        rescue Errno::EISCONN
+          # Socket is connected, nothing more to do
         end
       end
-      private :connect!
 
       private
 

data/lib/mongo/socket/unix.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2014-2020 MongoDB Inc.
 #

data/lib/mongo/socket.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
@@ -192,27 +192,24 @@ module Mongo
     #   socket.read(4096)
     #
     # @param [ Integer ] length The number of bytes to read.
-    # @param [ Numeric ] timeout The timeout to use for each chunk read.
+    # @param [ Numeric ] socket_timeout The timeout to use for each chunk read,
+    #   mutually exclusive to +timeout+.
+    # @param [ Numeric ] timeout The total timeout to the whole read operation,
+    #   mutually exclusive to +socket_timeout+.
     #
     # @raise [ Mongo::SocketError ] If not all data is returned.
     #
     # @return [ Object ] The data from the socket.
     #
     # @since 2.0.0
-    def read(length, timeout: nil)
-      map_exceptions do
-        data = read_from_socket(length, timeout: timeout)
-        unless (data.length > 0 || length == 0)
-          raise IOError, "Expected to read > 0 bytes but read 0 bytes"
-        end
-        while data.length < length
-          chunk = read_from_socket(length - data.length, timeout: timeout)
-          unless (chunk.length > 0 || length == 0)
-            raise IOError, "Expected to read > 0 bytes but read 0 bytes"
-          end
-          data << chunk
-        end
-        data
+    def read(length, socket_timeout: nil, timeout: nil)
+      if !socket_timeout.nil? && !timeout.nil?
+        raise ArgumentError, 'Both timeout and socket_timeout cannot be set'
+      end
+      if !socket_timeout.nil? || timeout.nil?
+        read_without_timeout(length, socket_timeout)
+      else
+        read_with_timeout(length, timeout)
       end
     end
 
@@ -233,15 +230,16 @@ module Mongo
     # Writes data to the socket instance.
     #
     # @param [ Array<Object> ] args The data to be written.
+    # @param [ Numeric ] timeout The total timeout to the whole write operation.
     #
     # @return [ Integer ] The length of bytes written to the socket.
     #
     # @raise [ Error::SocketError | Error::SocketTimeoutError ] When there is a network error during the write.
     #
     # @since 2.0.0
-    def write(*args)
+    def write(*args, timeout: nil)
       map_exceptions do
-        do_write(*args)
+        do_write(*args, timeout: timeout)
       end
     end
 
@@ -254,7 +252,7 @@ module Mongo
       true
     end
 
-    # For backwards compatibilty only, do not use.
+    # For backwards compatibility only, do not use.
     #
     # @return [ true ] Always true.
     #
@@ -265,18 +263,76 @@ module Mongo
 
     private
 
-    def read_from_socket(length, timeout: nil)
+    # Reads the +length+ bytes from the socket, the read operation duration is
+    # limited to +timeout+ second.
+    #
+    # @param [ Integer ] length The number of bytes to read.
+    # @param [ Numeric ] timeout The total timeout to the whole read operation.
+    #
+    # @return [ Object ] The data from the socket.
+    def read_with_timeout(length, timeout)
+      deadline = Utils.monotonic_time + timeout
+      map_exceptions do
+        String.new.tap do |data|
+          while data.length < length
+            socket_timeout = deadline - Utils.monotonic_time
+            if socket_timeout <= 0
+              raise Mongo::Error::TimeoutError
+            end
+            chunk = read_from_socket(length - data.length, socket_timeout: socket_timeout, csot: true)
+            unless chunk.length > 0
+              raise IOError, "Expected to read > 0 bytes but read 0 bytes"
+            end
+            data << chunk
+          end
+        end
+      end
+    end
+
+    # Reads the +length+ bytes from the socket. The read operation may involve
+    # multiple socket reads, each read is limited to +timeout+ second,
+    # if the parameter is provided.
+    #
+    # @param [ Integer ] length The number of bytes to read.
+    # @param [ Numeric ] socket_timeout The timeout to use for each chunk read.
+    #
+    # @return [ Object ] The data from the socket.
+    def read_without_timeout(length, socket_timeout = nil)
+      map_exceptions do
+        String.new.tap do |data|
+          while data.length < length
+            chunk = read_from_socket(length - data.length, socket_timeout: socket_timeout)
+            unless chunk.length > 0
+              raise IOError, "Expected to read > 0 bytes but read 0 bytes"
+            end
+            data << chunk
+          end
+        end
+      end
+    end
+
+
+    # Reads the +length+ bytes from the socket. The read operation may involve
+    # multiple socket reads, each read is limited to +timeout+ second,
+    # if the parameter is provided.
+    #
+    # @param [ Integer ] length The number of bytes to read.
+    # @param [ Numeric ] :socket_timeout The timeout to use for each chunk read.
+    # @param [ true | false ] :csot Whether the CSOT timeout is set for the operation.
+    #
+    # @return [ Object ] The data from the socket.
+    def read_from_socket(length, socket_timeout: nil, csot: false)
       # Just in case
       if length == 0
         return ''.force_encoding('BINARY')
       end
 
-      _timeout = timeout || self.timeout
+      _timeout = socket_timeout || self.timeout
       if _timeout
         if _timeout > 0
           deadline = Utils.monotonic_time + _timeout
         elsif _timeout < 0
-          raise Errno::ETIMEDOUT, "Negative timeout #{_timeout} given to socket"
+          raise_timeout_error!("Negative timeout #{_timeout} given to socket", csot)
         end
       end
 
@@ -331,15 +387,39 @@ module Mongo
         if deadline
           select_timeout = deadline - Utils.monotonic_time
           if select_timeout <= 0
-            raise Errno::ETIMEDOUT, "Took more than #{_timeout} seconds to receive data"
+            raise_timeout_error!("Took more than #{_timeout} seconds to receive data", csot)
           end
         end
+        pipe = options[:pipe]
         if exc.is_a?(IO::WaitReadable)
-          select_args = [[@socket], nil, [@socket], select_timeout]
+          if pipe
+            select_args = [[@socket, pipe], nil, [@socket, pipe], select_timeout]
+          else
+            select_args = [[@socket], nil, [@socket], select_timeout]
+          end
         else
           select_args = [nil, [@socket], [@socket], select_timeout]
         end
+
         rv = Kernel.select(*select_args)
+        if Lint.enabled?
+          if pipe && rv&.include?(pipe)
+            # If the return value of select is the read end of the pipe, and
+            # an IOError is not raised, then that means the socket is still
+            # open. Select is interrupted be closing the write end of the
+            # pipe, which either returns the pipe if the socket is open, or
+            # raises an IOError if it isn't. Select is interrupted after all
+            # of the pending and checked out connections have been interrupted
+            # and closed, and this only happens once the pool is cleared with
+            # interrupt_in_use connections flag. This means that in order for
+            # the socket to still be open when the select is interrupted, and
+            # that socket is being read from, that means after clear was
+            # called, a connection from the previous generation was checked
+            # out of the pool, for reading on its socket. This should be impossible.
+            raise Mongo::LintError, "Select interrupted for live socket. This should be impossible."
+          end
+        end
+
         if BSON::Environment.jruby?
           # Ignore the return value of Kernel.select.
           # On JRuby, select appears to return nil prior to timeout expiration
@@ -349,11 +429,11 @@ module Mongo
           if deadline
             select_timeout = deadline - Utils.monotonic_time
             if select_timeout <= 0
-              raise Errno::ETIMEDOUT, "Took more than #{_timeout} seconds to receive data"
+              raise_timeout_error!("Took more than #{_timeout} seconds to receive data", csot)
             end
           end
         elsif rv.nil?
-          raise Errno::ETIMEDOUT, "Took more than #{_timeout} seconds to receive data (select call timed out)"
+          raise_timeout_error!("Took more than #{_timeout} seconds to receive data (select call timed out)", csot)
         end
         retry
       end
@@ -378,9 +458,23 @@ module Mongo
     # sholud map exceptions.
     #
     # @param [ Array<Object> ] args The data to be written.
+    # @param [ Numeric ] :timeout The total timeout to the whole write operation.
+    #
+    # @return [ Integer ] The length of bytes written to the socket.
+    def do_write(*args, timeout: nil)
+      if timeout.nil?
+        write_without_timeout(*args)
+      else
+        write_with_timeout(*args, timeout: timeout)
+      end
+    end
+
+    # Writes data to to the socket.
+    #
+    # @param [ Array<Object> ] args The data to be written.
     #
     # @return [ Integer ] The length of bytes written to the socket.
-    def do_write(*args)
+    def write_without_timeout(*args)
       # This method used to forward arguments to @socket.write in a
       # single call like so:
       #
@@ -404,6 +498,57 @@ module Mongo
       end
     end
 
+    # Writes data to to the socket, the write duration is limited to +timeout+.
+    #
+    # @param [ Array<Object> ] args The data to be written.
+    # @param [ Numeric ] :timeout The total timeout to the whole write operation.
+    #
+    # @return [ Integer ] The length of bytes written to the socket.
+    def write_with_timeout(*args, timeout:)
+      raise ArgumentError, 'timeout cannot be nil' if timeout.nil?
+      raise_timeout_error!("Negative timeout #{timeout} given to socket", true) if timeout < 0
+
+      written = 0
+      args.each do |buf|
+        buf = buf.to_s
+        i = 0
+        while i < buf.length
+          chunk = buf[i...(i + WRITE_CHUNK_SIZE)]
+          written += write_chunk(chunk, timeout)
+          i += WRITE_CHUNK_SIZE
+        end
+      end
+      written
+    end
+
+    def write_chunk(chunk, timeout)
+      deadline = Utils.monotonic_time + timeout
+      written = 0
+      begin
+        written += @socket.write_nonblock(chunk[written..-1])
+      rescue IO::WaitWritable, Errno::EINTR
+        select_timeout = deadline - Utils.monotonic_time
+        rv = Kernel.select(nil, [@socket], nil, select_timeout)
+        if BSON::Environment.jruby?
+          # Ignore the return value of Kernel.select.
+          # On JRuby, select appears to return nil prior to timeout expiration
+          # (apparently due to a EAGAIN) which then causes us to fail the read
+          # even though we could have retried it.
+          # Check the deadline ourselves.
+          if deadline
+            select_timeout = deadline - Utils.monotonic_time
+            if select_timeout <= 0
+              raise_timeout_error!("Took more than #{timeout} seconds to receive data", true)
+            end
+          end
+        elsif rv.nil?
+          raise_timeout_error!("Took more than #{timeout} seconds to receive data (select call timed out)", true)
+        end
+        retry
+      end
+      written
+    end
+
     def unix_socket?(sock)
       defined?(UNIXSocket) && sock.is_a?(UNIXSocket)
     end
@@ -414,11 +559,14 @@ module Mongo
 
     DEFAULT_TCP_KEEPIDLE = 120
 
+    DEFAULT_TCP_USER_TIMEOUT = 210
+
     def set_keepalive_opts(sock)
       sock.setsockopt(SOL_SOCKET, SO_KEEPALIVE, true)
       set_option(sock, :TCP_KEEPINTVL, DEFAULT_TCP_KEEPINTVL)
       set_option(sock, :TCP_KEEPCNT, DEFAULT_TCP_KEEPCNT)
       set_option(sock, :TCP_KEEPIDLE, DEFAULT_TCP_KEEPIDLE)
+      set_option(sock, :TCP_USER_TIMEOUT, DEFAULT_TCP_USER_TIMEOUT)
     rescue
       # JRuby 9.2.13.0 and lower do not define TCP_KEEPINTVL etc. constants.
       # JRuby 9.2.14.0 defines the constants but does not allow to get or
@@ -455,5 +603,13 @@ module Mongo
     def human_address
       raise NotImplementedError
     end
+
+    def raise_timeout_error!(message = nil, csot = false)
+      if csot
+        raise Mongo::Error::TimeoutError
+      else
+        raise Errno::ETIMEDOUT, message
+      end
+    end
   end
 end

data/lib/mongo/srv/monitor.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
@@ -70,8 +70,6 @@ module Mongo
       end
 
       def scan!
-        old_hosts = last_result.address_strs
-
         begin
           last_result = Timeout.timeout(timeout) do
             @resolver.get_records(@srv_uri.query_hostname)

data/lib/mongo/srv/resolver.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2017-2020 MongoDB Inc.
 #
@@ -29,6 +29,16 @@ module Mongo
       #   before querying SRV records.
       RECORD_PREFIX = '_mongodb._tcp.'.freeze
 
+      # Generates the record prefix with a custom SRV service name if it is
+      # provided.
+      #
+      # @option srv_service_name [ String | nil ] The SRV service name to use
+      #   in the record prefix.
+      # @return [ String ] The generated record prefix.
+      def record_prefix(srv_service_name=nil)
+        return srv_service_name ? "_#{srv_service_name}._tcp." : RECORD_PREFIX
+      end
+
       # Creates a new Resolver.
       #
       # @option opts [ Float ] :timeout The timeout, in seconds, to use for
@@ -51,13 +61,18 @@ module Mongo
         options[:timeout] || Monitor::DEFAULT_TIMEOUT
       end
 
-      # Obtains all of the SRV records for a given hostname.
+      # Obtains all of the SRV records for a given hostname. If a srv_max_hosts
+      # is specified and it is greater than 0, return maximum srv_max_hosts records.
       #
       # In the event that a record with a mismatched domain is found or no
       # records are found, if the :raise_on_invalid option is true,
       # an exception will be raised, otherwise a warning will be logged.
       #
       # @param [ String ] hostname The hostname whose records should be obtained.
+      # @param [ String | nil ] srv_service_name The SRV service name for the DNS query.
+      #   If nil, 'mongodb' is used.
+      # @param [ Integer | nil ] srv_max_hosts The maximum number of records to return.
+      #   If this value is nil, return all of the records.
       #
       # @raise [ Mongo::Error::MismatchedDomain ] If the :raise_in_invalid
       #   Resolver option is true and a record with a domain name that does
@@ -66,8 +81,8 @@ module Mongo
       #   option is true and no records are found.
       #
       # @return [ Mongo::Srv::Result ] SRV lookup result.
-      def get_records(hostname)
-        query_name = RECORD_PREFIX + hostname
+      def get_records(hostname, srv_service_name=nil, srv_max_hosts=nil)
+        query_name = record_prefix(srv_service_name) + hostname
         resources = @resolver.getresources(query_name, Resolv::DNS::Resource::IN::SRV)
 
         # Collect all of the records into a Result object, raising an error
@@ -97,6 +112,12 @@ module Mongo
           end
         end
 
+        # if srv_max_hosts is in [1, #addresses)
+        if (1...result.address_strs.length).include? srv_max_hosts
+          sampled_records = resources.shuffle.first(srv_max_hosts)
+          result = Srv::Result.new(hostname)
+          sampled_records.each { |record| result.add_record(record) }
+        end
         result
       end
 

data/lib/mongo/srv/result.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-# encoding: utf-8
+# rubocop:todo all
 
 # Copyright (C) 2017-2020 MongoDB Inc.
 #