[JAVA-293]: Cursor is not closed after limit is fully fetched from server.

This turned out to be more of a refactoring than hoped, to clean up code while fixing issue.
Solution:
- positive batchsize tells server the ideal size of a batch returned
- negative batchsize tells server to return 1 batch of maximum that size, or whatever fits in a 4MB batch, and close the cursor
- positive limit does limit the total number of documents to return on cursor. Once limit is reached, cursor should be called automatically by driver.
- negative limit is only supported for legacy reason. It sets the batchSize to that negative value.

Author: agirbal

src/main/com/mongodb/DB.java

@@ -149,7 +149,7 @@ public CommandResult command( DBObject cmd )
     public CommandResult command( DBObject cmd , int options )
         throws MongoException {
 
-        Iterator<DBObject> i = getCollection( "$cmd" ).__find( cmd , new BasicDBObject() , 0 , -1 , options );
+        Iterator<DBObject> i = getCollection("$cmd").__find(cmd, new BasicDBObject(), 0, -1, 0, options);
         if ( i == null || ! i.hasNext() )
             return null;
 
@@ -260,7 +260,7 @@ public Set<String> getCollectionNames()
         if (namespaces == null)
             throw new RuntimeException("this is impossible");
 
-        Iterator<DBObject> i = namespaces.__find(new BasicDBObject(), null, 0, 0, getOptions());
+        Iterator<DBObject> i = namespaces.__find(new BasicDBObject(), null, 0, 0, 0, getOptions());
         if (i == null)
             return new HashSet<String>();
 

src/main/com/mongodb/DBApiLayer.java

@@ -50,7 +50,29 @@ static final boolean willTrace(){
     static final void trace( String s ){
         TRACE_LOGGER.log( TRACE_LEVEL , s );
     }
-    
+
+    static int chooseBatchSize(int batchSize, int limit, int fetched) {
+        int bs = Math.abs(batchSize);
+        int remaining = limit > 0 ? limit - fetched : 0;
+        int res = 0;
+        if (bs == 0 && remaining > 0)
+            res = remaining;
+        else if (bs > 0 && remaining == 0)
+            res = bs;
+        else
+            res = Math.min(bs, remaining);
+
+        if (batchSize < 0) {
+            // force close
+            res = -res;
+        }
+
+        if (res == 1) {
+            // optimization: use negative batchsize to close cursor
+            res = -1;
+        }
+        return res;
+    }
 
     /**
      * @param mongo the Mongo instance
@@ -260,15 +282,15 @@ public WriteResult remove( DBObject o , com.mongodb.WriteConcern concern )
         }
 
         @Override
-        Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int batchSize , int options )
+        Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int batchSize, int limit , int options )
             throws MongoException {
 
             if ( ref == null )
                 ref = new BasicDBObject();
 
             if ( willTrace() ) trace( "find: " + _fullNameSpace + " " + JSON.serialize( ref ) );
 
-            OutMessage query = OutMessage.query( _mongo , options , _fullNameSpace , numToSkip , batchSize , ref , fields );
+            OutMessage query = OutMessage.query( _mongo , options , _fullNameSpace , numToSkip , chooseBatchSize(batchSize, limit, 0) , ref , fields );
 
             Response res = _connector.call( _db , this , query , null , 2 );
 
@@ -282,7 +304,7 @@ Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int
                     throw e;
             }
 
-            return new Result( this , res , batchSize , options );
+            return new Result( this , res , batchSize, limit , options );
         }
 
         @Override
@@ -329,28 +351,34 @@ public void createIndex( final DBObject keys, final DBObject options )
 
     class Result implements Iterator<DBObject> {
 
-        Result( MyCollection coll , Response res , int numToReturn , int options ){
-            init( res );
+        Result( MyCollection coll , Response res , int batchSize, int limit , int options ){
             _collection = coll;
-            _numToReturn = numToReturn;
+            _batchSize = batchSize;
+            _limit = limit;
             _options = options;
             _host = res._host;
+            init( res );
         }
 
         private void init( Response res ){
             _totalBytes += res._len;
             _curResult = res;
             _cur = res.iterator();
             _sizes.add( res.size() );
+            _numFetched += res.size();
 
             if ( ( res._flags & Bytes.RESULTFLAG_CURSORNOTFOUND ) > 0 ){
                 throw new MongoException.CursorNotFound();
             }
+
+            if (res._cursor > 0 && _limit > 0 && _limit - _numFetched <= 0) {
+                // fetched all docs within limit, close cursor server-side
+                killCursor();
+            }
         }
 
         public DBObject next(){
             if ( _cur.hasNext() ) {
-                _numSeen++;
                 return _cur.next();
             }
 
@@ -382,7 +410,7 @@ private void _advance(){
 
             m.writeInt( 0 ); 
             m.writeCString( _collection._fullNameSpace );
-            m.writeInt( _numToReturn - _numSeen ); // num to return
+            m.writeInt( chooseBatchSize(_batchSize, _limit, _numFetched) );
             m.writeLong( _curResult.cursor() );
 
             Response res = _connector.call( DBApiLayer.this , _collection , m , _host );
@@ -394,12 +422,12 @@ public void remove(){
             throw new RuntimeException( "can't remove this way" );
         }
 
-        public int getNumberToReturn(){
-            return _numToReturn;
+        public int getBatchSize(){
+            return _batchSize;
         }
 
-        public void setNumberToReturn(int num){
-            _numToReturn = num;
+        public void setBatchSize(int size){
+            _batchSize = size;
         }
 
         public String toString(){
@@ -439,39 +467,47 @@ List<Integer> getSizes(){
         void close(){
             // not perfectly thread safe here, may need to use an atomicBoolean
             if (_curResult != null) {
-                long curId = _curResult.cursor();
+                killCursor();
                 _curResult = null;
                 _cur = null;
-                
-                if (curId > 0) {
-                    List<Long> l = new ArrayList<Long>();
-                    l.add(curId);
-
-                    try {
-                        killCursors(_host, l);
-                    } catch (Throwable t) {
-                        Bytes.LOGGER.log(Level.WARNING, "can't clean 1 cursor", t);
-                        _deadCursorIds.add(new DeadCursor(curId, _host));
-                    }
-                }
             }
         }
 
+        void killCursor() {
+            if (_curResult == null)
+                return;
+            long curId = _curResult.cursor();
+            if (curId < 0)
+                return;
+
+            List<Long> l = new ArrayList<Long>();
+            l.add(curId);
+
+            try {
+                killCursors(_host, l);
+            } catch (Throwable t) {
+                Bytes.LOGGER.log(Level.WARNING, "can't clean 1 cursor", t);
+                _deadCursorIds.add(new DeadCursor(curId, _host));
+            }
+            _curResult._cursor = 0;
+        }
+
         public ServerAddress getServerAddress() {
             return _host;
         }
 
         Response _curResult;
         Iterator<DBObject> _cur;
-        int _numToReturn;
+        int _batchSize;
+        int _limit;
         final MyCollection _collection;
         final int _options;
         final ServerAddress _host; // host where first went.  all subsequent have to go there
 
         private long _totalBytes = 0;
         private int _numGetMores = 0;
         private List<Integer> _sizes = new ArrayList<Integer>();
-        private int _numSeen = 0;
+        private int _numFetched = 0;
 
     }  // class Result
 

src/main/com/mongodb/DBCollection.java

@@ -199,64 +199,44 @@ public WriteResult remove( DBObject o )
 
 
     /**
-     * Finds an object from the database
-     * @param ref query used to search
-     * @param fields the fields of matching objects to return
-     * @param numToSkip will not return the first <tt>numToSkip</tt> matches
-     * @param batchSize  if positive, represents the size of each batch of objects retrieved. If negative, it limits the total number of objects retrieved.
-     * @param options - see Bytes QUERYOPTION_*
-     * @return the objects, if found
-     * @dochub find
+     * Finds objects
      */
-    abstract Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int batchSize , int options ) throws MongoException ;
+    abstract Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int batchSize , int limit, int options ) throws MongoException ;
 
     /**
-     * Finds an object.
      * Calls {@link DBCollection#find(com.mongodb.DBObject, com.mongodb.DBObject, int, int)} and applies the query options
-     * @param ref query used to search
+     * @param query query used to search
      * @param fields the fields of matching objects to return
-     * @param numToSkip will not return the first <tt>numToSkip</tt> matches
-     * @param batchSize if positive, represents the size of each batch of objects retrieved. If negative, it limits the total number of objects retrieved.
+     * @param numToSkip number of objects to skip
+     * @param batchSize the batch size. This option has a complex behavior, see {@link DBCursor#batchSize(int) }
      * @param options - see Bytes QUERYOPTION_*
-     * @return the objects, if found
+     * @return the cursor
      * @throws MongoException
      * @dochub find
      */
-    public final DBCursor find( DBObject ref , DBObject fields , int numToSkip , int batchSize , int options ) throws MongoException{
-    	return find(ref, fields, numToSkip, batchSize).addOption(options);
+    public final DBCursor find( DBObject query , DBObject fields , int numToSkip , int batchSize , int options ) throws MongoException{
+    	return find(query, fields, numToSkip, batchSize).addOption(options);
     }
 
 
     /**
-     * Finds an object.
-     * @param ref query used to search
+     * Finds objects from the database that match a query.
+     * A DBCursor object is returned, that can be iterated to go through the results.
+     *
+     * @param query query used to search
      * @param fields the fields of matching objects to return
-     * @param numToSkip will not return the first <tt>numToSkip</tt> matches
-     * @param batchSize if positive, represents the size of each batch of objects retrieved. If negative, it limits the total number of objects retrieved.
-     * @return the objects, if found
+     * @param numToSkip number of objects to skip
+     * @param batchSize the batch size. This option has a complex behavior, see {@link DBCursor#batchSize(int) }
+     * @param options - see Bytes QUERYOPTION_*
+     * @return the cursor
+     * @throws MongoException
      * @dochub find
      */
     public final DBCursor find( DBObject ref , DBObject fields , int numToSkip , int batchSize ) {
     	DBCursor cursor = find(ref, fields).skip(numToSkip).batchSize(batchSize);
-    	if ( batchSize < 0 ) 
-    		cursor.limit( Math.abs(batchSize) );
     	return cursor;
     }
 
-    /**
-     * Finds an object.
-     * @param ref query used to search
-     * @param fields the fields of matching objects to return
-     * @param numToSkip will not return the first <tt>numToSkip</tt> matches
-     * @param batchSize if positive, is the # of objects per batch sent back from the db.  all objects that match will be returned.  if batchSize < 0, its a hard limit, and only 1 batch will either batchSize or the # that fit in a batch
-     * @return the objects, if found
-     * @dochub find
-     */
-    Iterator<DBObject> __find( DBObject ref , DBObject fields , int numToSkip , int batchSize ) 
-        throws MongoException {
-        return __find( ref , fields , numToSkip , batchSize , getOptions() );
-    }
-
     // ------
 
     /**
@@ -282,7 +262,7 @@ public final DBObject findOne( Object obj )
      * @dochub find
      */
     public final DBObject findOne( Object obj, DBObject fields ) {
-        Iterator<DBObject> iterator = __find(new BasicDBObject("_id", obj), fields, 0, -1, getOptions() );
+        Iterator<DBObject> iterator = __find(new BasicDBObject("_id", obj), fields, 0, -1, 0, getOptions() );
         return (iterator != null ? iterator.next() : null);
     }
 
@@ -573,7 +553,7 @@ public final DBObject findOne( DBObject o )
      * @dochub find
      */
     public final DBObject findOne( DBObject o, DBObject fields ) {
-        Iterator<DBObject> i = __find( o , fields , 0 , -1 , getOptions() );
+        Iterator<DBObject> i = __find( o , fields , 0 , -1 , 0, getOptions() );
         if ( i == null || ! i.hasNext() )
             return null;
         return i.next();