RUBY-189 use result of ismaster's maxBsonObjectSize

ext/cbson/cbson.c

@@ -88,6 +88,8 @@ static VALUE InvalidDocument;
 static VALUE DigestMD5;
 static VALUE RB_HASH;
 
+static int max_bson_size;
+
 #if HAVE_RUBY_ENCODING_H
 #include "ruby/encoding.h"
 #define STR_NEW(p,n)                                                    \
@@ -582,9 +584,9 @@ static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys, VALUE move_
     length = buffer_get_position(buffer) - start_position;
 
     // make sure that length doesn't exceed 4MB
-    if (length > 4 * 1024 * 1024) {
+    if (length > max_bson_size) {
       buffer_free(buffer);
-      rb_raise(InvalidDocument, "Document too large: BSON documents are limited to 4MB.");
+      rb_raise(InvalidDocument, "Document too large: BSON documents are limited to %d bytes.", max_bson_size);
       return;
     }
     SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&length, 4);
@@ -902,11 +904,18 @@ static VALUE objectid_generate(VALUE self)
     return oid;
 }
 
+static void method_update_max_bson_size(VALUE self, VALUE connection) {
+    max_bson_size = FIX2INT(rb_funcall(connection, rb_intern("max_bson_size"), 0));
+}
+
+static VALUE method_max_bson_size(VALUE self) {
+    return INT2FIX(max_bson_size);
+}
 
 void Init_cbson() {
     VALUE bson, CBson, Digest, ext_version, digest;
     static char hostname[MAX_HOSTNAME_LENGTH];
-    
+
     element_assignment_method = rb_intern("[]=");
     unpack_method = rb_intern("unpack");
     utc_method = rb_intern("utc");
@@ -939,6 +948,8 @@ void Init_cbson() {
     rb_define_const(CBson, "VERSION", ext_version);
     rb_define_module_function(CBson, "serialize", method_serialize, 3);
     rb_define_module_function(CBson, "deserialize", method_deserialize, 1);
+    rb_define_module_function(CBson, "max_bson_size", method_max_bson_size, 0);
+    rb_define_module_function(CBson, "update_max_bson_size", method_update_max_bson_size, 1);
 
     rb_require("digest/md5");
     Digest = rb_const_get(rb_cObject, rb_intern("Digest"));
@@ -953,4 +964,6 @@ void Init_cbson() {
         rb_str_new2(hostname));
     memcpy(hostname_digest, RSTRING_PTR(digest), 16);
     hostname_digest[16] = '\0';
+
+    max_bson_size = 4 * 1024 * 1024;
 }

ext/java/src/org/jbson/RubyBSONEncoder.java

@@ -61,6 +61,8 @@ public class RubyBSONEncoder extends BSONEncoder {
     private boolean _check_keys;
     private boolean _move_id;
 
+    private static final int DEFAULT_MAX_BSON_SIZE = 4 * 1024 * 1024;
+    private static int _max_bson_size = DEFAULT_MAX_BSON_SIZE;
     private static final int BIT_SIZE = 64;
     private static final long MAX = (1L << (BIT_SIZE - 1)) - 1;
     private static final BigInteger LONG_MAX = BigInteger.valueOf(MAX);
@@ -84,6 +86,17 @@ public class RubyBSONEncoder extends BSONEncoder {
       }
     }
 
+    public static RubyFixnum max_bson_size(RubyObject obj) {
+        Ruby _run = obj.getRuntime();
+        return _run.newFixnum(_max_bson_size);
+    }
+
+    public static void update_max_bson_size(RubyObject obj, RubyObject conn) {
+        Ruby _run = obj.getRuntime();
+        _max_bson_size = ((Long)JavaEmbedUtils.invokeMethod( _run, conn, "max_bson_size",
+          new Object[] {}, Object.class)).intValue();
+    }
+
     public RubyString encode( Object arg ) {
         RubyHash o = (RubyHash)arg;
         BasicOutputBuffer buf = new BasicOutputBuffer();
@@ -95,12 +108,12 @@ public class RubyBSONEncoder extends BSONEncoder {
         return b;
     }
 
-    public void set( OutputBuffer out ){
+    public void set( OutputBuffer out ) {
         if ( _buf != null ) {
             done();
             throw new IllegalStateException( "in the middle of something" );
         }
-        
+
         _buf = out;
     }
  
@@ -202,10 +215,10 @@ public class RubyBSONEncoder extends BSONEncoder {
           }
 
         // Make sure we're within the 4MB limit
-        if ( _buf.size() > 4 * 1024 * 1024 ) {
+        if ( _buf.size() > _max_bson_size ) {
             _rbRaise( (RubyClass)_rbclsInvalidDocument,
-              "Document is too large (" + _buf.size() + "). BSON documents are limited to 4MB (" +
-               4 * 1024 * 1024 + ").");
+              "Document is too large (" + _buf.size() + "). BSON documents are limited to " +
+              _max_bson_size + " bytes." );
         }
 
         _buf.write( EOO );

lib/bson/bson_c.rb

@@ -28,5 +28,12 @@ module BSON
       CBson.deserialize(ByteBuffer.new(buf).to_s)
     end
 
+    def self.max_bson_size
+      CBson.max_bson_size
+    end
+
+    def self.update_max_bson_size(connection)
+      CBson.update_max_bson_size(connection)
+    end
   end
 end

lib/bson/bson_java.rb

@@ -17,5 +17,12 @@ module BSON
       callback.get
     end
 
+    def self.max_bson_size
+      Java::OrgJbson::RubyBSONEncoder.max_bson_size(self)
+    end
+
+    def self.update_max_bson_size(connection)
+      Java::OrgJbson::RubyBSONEncoder.update_max_bson_size(self, connection)
+    end
   end
 end

lib/bson/bson_ruby.rb

@@ -20,6 +20,10 @@ module BSON
   # A BSON seralizer/deserializer in pure Ruby.
   class BSON_RUBY
 
+    DEFAULT_MAX_BSON_SIZE = 4 * 1024 * 1024
+
+    @@max_bson_size = DEFAULT_MAX_BSON_SIZE
+
     MINKEY = -1
     EOO = 0
     NUMBER = 1
@@ -51,13 +55,13 @@ module BSON
       NULL_BYTE       = "\0".force_encoding('binary').freeze
       UTF8_ENCODING   = Encoding.find('utf-8')
       BINARY_ENCODING = Encoding.find('binary')
-      
+
       def self.to_utf8_binary(str)
         str.encode(UTF8_ENCODING).force_encoding(BINARY_ENCODING)
       end
     else
       NULL_BYTE = "\0"
-      
+
       def self.to_utf8_binary(str)
         begin
           str.unpack("U*")
@@ -68,6 +72,14 @@ module BSON
       end
     end
 
+    def self.update_max_bson_size(connection)
+      @@max_bson_size = connection.max_bson_size
+    end
+
+    def self.max_bson_size
+      @@max_bson_size
+    end
+
     def self.serialize_cstr(buf, val)
       buf.put_binary(to_utf8_binary(val.to_s))
       buf.put_binary(NULL_BYTE)
@@ -120,8 +132,8 @@ module BSON
       end
 
       serialize_eoo_element(@buf)
-      if @buf.size > 4 * 1024 * 1024
-        raise InvalidDocument, "Document is too large (#{@buf.size}). BSON documents are limited to 4MB (#{4 * 1024 * 1024})."
+      if @buf.size > @@max_bson_size
+        raise InvalidDocument, "Document is too large (#{@buf.size}). BSON documents are limited to #{@@max_bson_size} bytes."
       end
       @buf.put_int(@buf.size, 0)
       @buf

lib/mongo.rb

@@ -11,6 +11,8 @@ module Mongo
   DESCENDING = -1
   GEO2D      = '2d'
 
+  DEFAULT_MAX_BSON_SIZE = 4 * 1024 * 1024
+
   module Constants
     OP_REPLY        = 1
     OP_MSG          = 1000

lib/mongo/connection.rb

@@ -440,7 +440,9 @@ module Mongo
         set_primary(@host_to_try)
       end
 
-      if !connected?
+      if connected?
+        BSON::BSON_CODER.update_max_bson_size(self)
+      else
         raise ConnectionFailure, "Failed to connect to a master node at #{@host_to_try[0]}:#{@host_to_try[1]}"
       end
     end
@@ -470,6 +472,15 @@ module Mongo
       @primary_pool = nil
     end
 
+    # Returns the maximum BSON object size as returned by the core server.
+    # Use the 4MB default when the server doesn't report this.
+    #
+    # @return [Integer]
+    def max_bson_size
+      config = self['admin'].command({:ismaster => 1})
+      config['maxBsonObjectSize'] || Mongo::DEFAULT_MAX_BSON_SIZE
+    end
+
     # Checkout a socket for reading (i.e., a secondary node).
     # Note: this is overridden in ReplSetConnection.
     def checkout_reader

lib/mongo/repl_set_connection.rb

@@ -111,7 +111,9 @@ module Mongo
 
       pick_secondary_for_read if @read_secondary
 
-      if !connected?
+      if connected?
+        BSON::BSON_CODER.update_max_bson_size(self)
+      else
         if @secondary_pools.empty?
           raise ConnectionFailure, "Failed to connect any given host:port"
         else

test/bson/bson_test.rb

@@ -67,13 +67,17 @@ class BSONTest < Test::Unit::TestCase
     assert_doc_pass(doc)
   end
 
-  def test_document_length
-    doc = {'name' => 'a' * 5 * 1024 * 1024}
+  def test_limit_max_bson_size
+    doc = {'name' => 'a' * BSON_CODER.max_bson_size}
     assert_raise InvalidDocument do
       assert @encoder.serialize(doc)
     end
   end
 
+  def test_max_bson_size
+    assert BSON_CODER.max_bson_size >= BSON::DEFAULT_MAX_BSON_SIZE
+  end
+
   def test_round_trip
     doc = {'doc' => 123}
     @encoder.deserialize(@encoder.serialize(doc))

test/connection_test.rb

@@ -165,6 +165,41 @@ class TestConnection < Test::Unit::TestCase
     assert unlocked, "mongod failed to unlock"
   end
 
+  def test_max_bson_size_value
+    conn = standard_connection
+    if conn.server_version > "1.6"
+      assert_equal conn['admin'].command({:ismaster => 1})['maxBsonObjectSize'], conn.max_bson_size
+    end
+
+    conn.connect
+    assert_equal BSON::BSON_CODER.max_bson_size, conn.max_bson_size
+    doc = {'n' => 'a' * (BSON_CODER.max_bson_size - 11)}
+    assert_raise InvalidDocument do
+      assert BSON::BSON_CODER.serialize(doc)
+    end
+
+    limit = 7 * 1024 * 1024
+    conn.stubs(:max_bson_size).returns(limit)
+    conn.connect
+    assert_equal limit, conn.max_bson_size
+    assert_equal limit, BSON::BSON_CODER.max_bson_size
+    doc = {'n' => 'a' * ((limit) - 11)}
+    assert_raise_error InvalidDocument, "limited to #{limit}" do
+      assert BSON::BSON_CODER.serialize(doc)
+    end
+  end
+
+  def test_max_bson_size_with_old_mongod
+    conn = standard_connection(:connect => false)
+
+    admin_db = Object.new
+    admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1}).twice
+    conn.expects(:[]).with('admin').returns(admin_db).twice
+
+    conn.connect
+    assert_equal Mongo::DEFAULT_MAX_BSON_SIZE, BSON::BSON_CODER.max_bson_size
+  end
+
   context "Saved authentications" do
     setup do
       @conn = standard_connection