6 files changed, 38 insertions, 13 deletions

diff --git a/jstests/slowNightly/sharding_rs2.js b/jstests/slowNightly/sharding_rs2.js
index 7ef0e08d9b0..d71b3d0aebd 100644
--- a/jstests/slowNightly/sharding_rs2.js
+++ b/jstests/slowNightly/sharding_rs2.js
@@ -93,7 +93,7 @@ db.getLastError( 3 , 10000 );
 
 assert.eq( 100 , ts.count() , "B4" )
 assert.eq( 100 , ts.find().itcount() , "B5" )
-//assert.eq( 100 , ts.find().batchSize(5).itcount() , "B6" )
+assert.eq( 100 , ts.find().batchSize(5).itcount() , "B6" )
 
 // --- sharded ----
 
diff --git a/s/cursors.cpp b/s/cursors.cpp
index c28832ca60d..903d588f77a 100644
--- a/s/cursors.cpp
+++ b/s/cursors.cpp
@@ -139,9 +139,9 @@ namespace mongo {
                         << endl;
     }
 
-    ShardedClientCursorPtr CursorCache::get( long long id ) {
+    ShardedClientCursorPtr CursorCache::get( long long id ) const {
         scoped_lock lk( _mutex );
-        MapSharded::iterator i = _cursors.find( id );
+        MapSharded::const_iterator i = _cursors.find( id );
         if ( i == _cursors.end() ) {
             OCCASIONALLY log() << "Sharded CursorCache missing cursor id: " << id << endl;
             return ShardedClientCursorPtr();
@@ -168,6 +168,16 @@ namespace mongo {
         _refs[id] = server;
     }
 
+    string CursorCache::getRef( long long id ) const {
+        assert( id );
+        scoped_lock lk( _mutex );
+        MapNormal::const_iterator i = _refs.find( id );
+        if ( i == _refs.end() )
+            return "";
+        return i->second;
+    }
+
+
     long long CursorCache::genId() {
         while ( true ) {
             long long x = security.getNonce();
@@ -236,7 +246,7 @@ namespace mongo {
         }
     }
 
-    void CursorCache::appendInfo( BSONObjBuilder& result ) {
+    void CursorCache::appendInfo( BSONObjBuilder& result ) const {
         scoped_lock lk( _mutex );
         result.append( "sharded" , (int)_cursors.size() );
         result.appendNumber( "shardedEver" , _shardedTotal );
diff --git a/s/cursors.h b/s/cursors.h
index d1654680057..f11d3fb7b69 100644
--- a/s/cursors.h
+++ b/s/cursors.h
@@ -74,22 +74,25 @@ namespace mongo {
         CursorCache();
         ~CursorCache();
 
-        ShardedClientCursorPtr get( long long id );
+        ShardedClientCursorPtr get( long long id ) const;
         void store( ShardedClientCursorPtr cursor );
         void remove( long long id );
 
         void storeRef( const string& server , long long id );
 
+        /** @return the server for id or "" */
+        string getRef( long long id ) const ;
+        
         void gotKillCursors(Message& m );
 
-        void appendInfo( BSONObjBuilder& result );
+        void appendInfo( BSONObjBuilder& result ) const ;
 
         long long genId();
 
         void doTimeouts();
         void startTimeoutThread();
     private:
-        mongo::mutex _mutex;
+        mutable mongo::mutex _mutex;
 
         MapSharded _cursors;
         MapNormal _refs;
diff --git a/s/request.cpp b/s/request.cpp
index d5aabd1c6d9..4ae28ec21e2 100644
--- a/s/request.cpp
+++ b/s/request.cpp
@@ -154,7 +154,14 @@ namespace mongo {
         assert( _didInit );
         long long cursor =response.header()->getCursor();
         if ( cursor ) {
-            cursorCache.storeRef( fromServer , cursor );
+            if ( fromServer.size() ) {
+                cursorCache.storeRef( fromServer , cursor );
+            }
+            else {
+                // probably a getMore
+                // make sure we have a ref for this
+                assert( cursorCache.getRef( cursor ).size() );
+            }
         }
         _p->reply( _m , response , _id );
     }
diff --git a/s/strategy.cpp b/s/strategy.cpp
index d562507bc65..7c1fb0bff33 100644
--- a/s/strategy.cpp
+++ b/s/strategy.cpp
@@ -49,8 +49,10 @@ namespace mongo {
         ShardConnection dbcon( shard , r.getns() );
         DBClientBase &c = dbcon.conn();
 
+        string actualServer;
+
         Message response;
-        bool ok = c.call( r.m(), response);
+        bool ok = c.call( r.m(), response, true , &actualServer );
         uassert( 10200 , "mongos: error calling db", ok );
 
         {
@@ -61,7 +63,7 @@ namespace mongo {
             }
         }
 
-        r.reply( response , c.getServerAddress() );
+        r.reply( response , actualServer.size() ? actualServer : c.getServerAddress() );
         dbcon.done();
     }
 
diff --git a/s/strategy_single.cpp b/s/strategy_single.cpp
index 514ece3cc43..b3b5502a2c8 100644
--- a/s/strategy_single.cpp
+++ b/s/strategy_single.cpp
@@ -18,6 +18,7 @@
 
 #include "pch.h"
 #include "request.h"
+#include "cursors.h"
 #include "../client/connpool.h"
 #include "../db/commands.h"
 
@@ -83,14 +84,16 @@ namespace mongo {
         virtual void getMore( Request& r ) {
             const char *ns = r.getns();
 
-            log(3) << "single getmore: " << ns << endl;
+            LOG(3) << "single getmore: " << ns << endl;
 
-            ShardConnection conn( r.primaryShard() , ns );
+            long long id = r.d().getInt64( 4 );
+
+            ShardConnection conn( cursorCache.getRef( id ) , ns );
 
             Message response;
             bool ok = conn->callRead( r.m() , response);
             uassert( 10204 , "dbgrid: getmore: error calling db", ok);
-            r.reply( response , conn->getServerAddress() );
+            r.reply( response , "" /*conn->getServerAddress() */ );
 
             conn.done();