1 files changed, 91 insertions, 32 deletions

diff --git a/client/dbclientcursor.cpp b/client/dbclientcursor.cpp
index 6c6afc0..5db360e 100644
--- a/client/dbclientcursor.cpp
+++ b/client/dbclientcursor.cpp
@@ -37,8 +37,7 @@ namespace mongo {
         return batchSize < nToReturn ? batchSize : nToReturn;
     }
 
-    bool DBClientCursor::init() {
-        Message toSend;
+    void DBClientCursor::_assembleInit( Message& toSend ) {
         if ( !cursorId ) {
             assembleRequest( ns, query, nextBatchSize() , nToSkip, fieldsToReturn, opts, toSend );
         }
@@ -50,12 +49,18 @@ namespace mongo {
             b.appendNum( cursorId );
             toSend.setData( dbGetMore, b.buf(), b.len() );
         }
-        if ( !_client->call( toSend, *m, false ) ) {
+    }
+
+    bool DBClientCursor::init() {
+        Message toSend;
+        _assembleInit( toSend );
+
+        if ( !_client->call( toSend, *b.m, false ) ) {
             // log msg temp?
             log() << "DBClientCursor::init call() failed" << endl;
             return false;
         }
-        if ( m->empty() ) {
+        if ( b.m->empty() ) {
             // log msg temp?
             log() << "DBClientCursor::init message from call() was empty" << endl;
             return false;
@@ -63,12 +68,41 @@ namespace mongo {
         dataReceived();
         return true;
     }
+    
+    void DBClientCursor::initLazy( bool isRetry ) {
+        verify( 15875 , _client->lazySupported() );
+        Message toSend;
+        _assembleInit( toSend );
+        _client->say( toSend, isRetry );
+    }
+
+    bool DBClientCursor::initLazyFinish( bool& retry ) {
+
+        bool recvd = _client->recv( *b.m );
+
+        // If we get a bad response, return false
+        if ( ! recvd || b.m->empty() ) {
+
+            if( !recvd )
+                log() << "DBClientCursor::init lazy say() failed" << endl;
+            if( b.m->empty() )
+                log() << "DBClientCursor::init message from say() was empty" << endl;
+
+            _client->checkResponse( NULL, -1, &retry, &_lazyHost );
+
+            return false;
+
+        }
+
+        dataReceived( retry, _lazyHost );
+        return ! retry;
+    }
 
     void DBClientCursor::requestMore() {
-        assert( cursorId && pos == nReturned );
+        assert( cursorId && b.pos == b.nReturned );
 
         if (haveLimit) {
-            nToReturn -= nReturned;
+            nToReturn -= b.nReturned;
             assert(nToReturn > 0);
         }
         BufBuilder b;
@@ -83,7 +117,7 @@ namespace mongo {
 
         if ( _client ) {
             _client->call( toSend, *response );
-            m = response;
+            this->b.m = response;
             dataReceived();
         }
         else {
@@ -91,7 +125,7 @@ namespace mongo {
             ScopedDbConnection conn( _scopedHost );
             conn->call( toSend , *response );
             _client = conn.get();
-            m = response;
+            this->b.m = response;
             dataReceived();
             _client = 0;
             conn.done();
@@ -100,19 +134,24 @@ namespace mongo {
 
     /** with QueryOption_Exhaust, the server just blasts data at us (marked at end with cursorid==0). */
     void DBClientCursor::exhaustReceiveMore() {
-        assert( cursorId && pos == nReturned );
+        assert( cursorId && b.pos == b.nReturned );
         assert( !haveLimit );
         auto_ptr<Message> response(new Message());
         assert( _client );
         _client->recv(*response);
-        m = response;
+        b.m = response;
         dataReceived();
     }
 
-    void DBClientCursor::dataReceived() {
-        QueryResult *qr = (QueryResult *) m->singleData();
+    void DBClientCursor::dataReceived( bool& retry, string& host ) {
+
+        QueryResult *qr = (QueryResult *) b.m->singleData();
         resultFlags = qr->resultFlags();
 
+        if ( qr->resultFlags() & ResultFlag_ErrSet ) {
+            wasError = true;
+        }
+
         if ( qr->resultFlags() & ResultFlag_CursorNotFound ) {
             // cursor id no longer valid at the server.
             assert( qr->cursorId == 0 );
@@ -127,11 +166,12 @@ namespace mongo {
             cursorId = qr->cursorId;
         }
 
-        nReturned = qr->nReturned;
-        pos = 0;
-        data = qr->data();
+        b.nReturned = qr->nReturned;
+        b.pos = 0;
+        b.data = qr->data();
+
+        _client->checkResponse( b.data, b.nReturned, &retry, &host ); // watches for "not master"
 
-        _client->checkResponse( data, nReturned );
         /* this assert would fire the way we currently work:
             assert( nReturned || cursorId == 0 );
         */
@@ -144,17 +184,17 @@ namespace mongo {
         if ( !_putBack.empty() )
             return true;
 
-        if (haveLimit && pos >= nToReturn)
+        if (haveLimit && b.pos >= nToReturn)
             return false;
 
-        if ( pos < nReturned )
+        if ( b.pos < b.nReturned )
             return true;
 
         if ( cursorId == 0 )
             return false;
 
         requestMore();
-        return pos < nReturned;
+        return b.pos < b.nReturned;
     }
 
     BSONObj DBClientCursor::next() {
@@ -165,11 +205,11 @@ namespace mongo {
             return ret;
         }
 
-        uassert(13422, "DBClientCursor next() called but more() is false", pos < nReturned);
+        uassert(13422, "DBClientCursor next() called but more() is false", b.pos < b.nReturned);
 
-        pos++;
-        BSONObj o(data);
-        data += o.objsize();
+        b.pos++;
+        BSONObj o(b.data);
+        b.data += o.objsize();
         /* todo would be good to make data null at end of batch for safety */
         return o;
     }
@@ -187,9 +227,9 @@ namespace mongo {
         }
         */
 
-        int p = pos;
-        const char *d = data;
-        while( m && p < nReturned ) {
+        int p = b.pos;
+        const char *d = b.data;
+        while( m && p < b.nReturned ) {
             BSONObj o(d);
             d += o.objsize();
             p++;
@@ -198,6 +238,19 @@ namespace mongo {
         }
     }
     
+    bool DBClientCursor::peekError(BSONObj* error){
+        if( ! wasError ) return false;
+
+        vector<BSONObj> v;
+        peek(v, 1);
+
+        assert( v.size() == 1 );
+        assert( hasErrField( v[0] ) );
+
+        if( error ) *error = v[0].getOwned();
+        return true;
+    }
+
     void DBClientCursor::attach( AScopedConnection * conn ) {
         assert( _scopedHost.size() == 0 );
         assert( conn );
@@ -205,14 +258,20 @@ namespace mongo {
 
         if ( conn->get()->type() == ConnectionString::SET ||
              conn->get()->type() == ConnectionString::SYNC ) {
-            _scopedHost = _client->getServerAddress();
+            if( _lazyHost.size() > 0 )
+                _scopedHost = _lazyHost;
+            else if( _client )
+                _scopedHost = _client->getServerAddress();
+            else
+                massert(14821, "No client or lazy client specified, cannot store multi-host connection.", false);
         }
         else {
             _scopedHost = conn->getHost();
         }
-        
+
         conn->done();
         _client = 0;
+        _lazyHost = "";
     }
 
     DBClientCursor::~DBClientCursor() {
@@ -221,12 +280,12 @@ namespace mongo {
 
         DESTRUCTOR_GUARD (
 
-        if ( cursorId && _ownCursor ) {
-        BufBuilder b;
-        b.appendNum( (int)0 ); // reserved
+        if ( cursorId && _ownCursor && ! inShutdown() ) {
+            BufBuilder b;
+            b.appendNum( (int)0 ); // reserved
             b.appendNum( (int)1 ); // number
             b.appendNum( cursorId );
-
+            
             Message m;
             m.setData( dbKillCursors , b.buf() , b.len() );