Improve ssl read performance (#981)

https://jira.mongodb.org/browse/RUBY-1364

Author: p-mongo

lib/mongo/socket.rb

@@ -169,11 +169,57 @@ def eof?
     private
 
     def read_from_socket(length)
-      data = String.new
+      # Just in case
+      if length == 0
+        return ''.force_encoding('BINARY')
+      end
+
       deadline = (Time.now + timeout) if timeout
+
+      # We want to have a fixed and reasonably small size buffer for reads
+      # because, for example, OpenSSL reads in 16 kb chunks max.
+      # Having a 16 mb buffer means there will be 1000 reads each allocating
+      # 16 mb of memory and using 16 kb of it.
+      buf_size = read_buffer_size
+      data = nil
+
+      # If we want to read less than the buffer size, just allocate the
+      # memory that is necessary
+      if length < buf_size
+        buf_size = length
+      end
+
+      # The binary encoding is important, otherwise ruby performs encoding
+      # conversions of some sort during the write into the buffer which
+      # kills performance
+      buf = allocate_string(buf_size)
+      retrieved = 0
       begin
-        while (data.length < length)
-          data << @socket.read_nonblock(length - data.length)
+        while retrieved < length
+          retrieve = length - retrieved
+          if retrieve > buf_size
+            retrieve = buf_size
+          end
+          chunk = @socket.read_nonblock(retrieve, buf)
+
+          # If we read the entire wanted length in one operation,
+          # return the data as is which saves one memory allocation and
+          # one copy per read
+          if retrieved == 0 && chunk.length == length
+            return chunk
+          end
+
+          # If we are here, we are reading the wanted length in
+          # multiple operations. Allocate the total buffer here rather
+          # than up front so that the special case above won't be
+          # allocating twice
+          if data.nil?
+            data = allocate_string(length)
+          end
+
+          # ... and we need to copy the chunks at this point
+          data[retrieved, chunk.length] = chunk
+          retrieved += chunk.length
         end
       rescue IO::WaitReadable
         select_timeout = (deadline - Time.now) if deadline
@@ -186,6 +232,20 @@ def read_from_socket(length)
       data
     end
 
+    def allocate_string(capacity)
+      if RUBY_VERSION >= '2.4.0'
+        String.new('', :capacity => capacity, :encoding => 'BINARY')
+      else
+        ('x'*capacity).force_encoding('BINARY')
+      end
+    end
+
+    def read_buffer_size
+      # Buffer size for non-SSL reads
+      # 64kb
+      65536
+    end
+
     def unix_socket?(sock)
       defined?(UNIXSocket) && sock.is_a?(UNIXSocket)
     end

lib/mongo/socket/ssl.rb

@@ -168,6 +168,12 @@ def verify_certificate!(socket)
           end
         end
       end
+
+      def read_buffer_size
+        # Buffer size for SSL reads.
+        # Capped at 16k due to https://linux.die.net/man/3/ssl_read
+        16384
+      end
     end
   end
 end