Skip to content

Commit

Permalink
CXX-176 implement nToReturn
Browse files Browse the repository at this point in the history 
Implemented nToReturn with actual limit semantics.
  • Loading branch information
@hanumantmk
hanumantmk committed May 22, 2014
1 parent 87f67a8 commit 2e86826
Show file tree
Hide file tree
Showing 8 changed files with 83 additions and 41 deletions.
15 changes: 9 additions & 6 deletions src/mongo/client/dbclient.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@
#include "mongo/util/net/ssl_manager.h"
#include "mongo/util/password_digest.h"

#include <algorithm>
#include <cstdlib>

namespace mongo {

using std::auto_ptr;
Expand Down Expand Up @@ -857,11 +860,11 @@ namespace mongo {
/** query N objects from the database into an array. makes sense mostly when you want a small number of results. if a huge number, use
query() and iterate the cursor.
*/
void DBClientInterface::findN(vector<BSONObj>& out, const string& ns, Query query, int nToReturn, int nToSkip, const BSONObj *fieldsToReturn, int queryOptions) {
out.reserve(nToReturn);
void DBClientInterface::findN(vector<BSONObj>& out, const string& ns, Query query, int nToReturn, int nToSkip, const BSONObj *fieldsToReturn, int queryOptions, int batchSize) {
out.reserve(std::min(batchSize, nToReturn));

auto_ptr<DBClientCursor> c =
this->query(ns, query, nToReturn, nToSkip, fieldsToReturn, queryOptions);
this->query(ns, query, nToReturn, nToSkip, fieldsToReturn, queryOptions, batchSize);

uassert( 10276 , str::stream() << "DBClientBase::findN: transport error: " << getServerAddress() << " ns: " << ns << " query: " << query.toString(), c.get() );

Expand Down Expand Up @@ -1077,7 +1080,7 @@ namespace mongo {

c->shim.reset((cursor_shim = new DBClientCursorShimCursorID(*c)));

c->haveLimit = false;
c->nToReturn = 0;

if (c->rawMore()) {
BSONObj res = cursor_shim->get_cursor();
Expand All @@ -1101,7 +1104,7 @@ namespace mongo {
"cursor"), 1, 0, NULL, queryOptions, 0);

simple->shim.reset(new DBClientCursorShimArray(*simple));
simple->haveLimit = false;
simple->nToReturn = 0;

return simple;
}
Expand All @@ -1112,7 +1115,7 @@ namespace mongo {
}

auto_ptr<DBClientCursor> DBClientBase::getMore( const string &ns, long long cursorId, int nToReturn, int options ) {
auto_ptr<DBClientCursor> c( new DBClientCursor( this, ns, cursorId, nToReturn, options ) );
auto_ptr<DBClientCursor> c( new DBClientCursor( this, ns, cursorId, nToReturn < 0 ? abs(nToReturn) : 0, options, abs(nToReturn)) );
if ( c->init() )
return c;
return auto_ptr< DBClientCursor >( 0 );
Expand Down
36 changes: 18 additions & 18 deletions src/mongo/client/dbclientcursor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,9 @@ namespace mongo {
_client(client),
ns(_ns),
query(_query),
nToReturn(_nToReturn),
haveLimit( _nToReturn > 0 && !(queryOptions & QueryOption_CursorTailable)),
nToReturn( (queryOptions & QueryOption_CursorTailable) ? 0 : _nToReturn ),
nToSkip(_nToSkip),
nReturned(0),
fieldsToReturn(_fieldsToReturn),
opts(queryOptions),
batchSize(bs==1?2:bs),
Expand All @@ -48,15 +48,15 @@ namespace mongo {
_finishConsInit();
}

DBClientCursor::DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options ) :
DBClientCursor::DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options, int bs ) :
_client(client),
ns(_ns),
nToReturn( _nToReturn ),
haveLimit( _nToReturn > 0 && !(options & QueryOption_CursorTailable)),
nToReturn( (options & QueryOption_CursorTailable) ? 0 : _nToReturn ),
nToSkip(0),
nReturned(0),
fieldsToReturn(0),
opts( options ),
batchSize(0),
batchSize(bs==1?2:bs),
resultFlags(0),
cursorId(_cursorId),
_ownCursor(true),
Expand All @@ -71,14 +71,16 @@ namespace mongo {
}

int DBClientCursor::nextBatchSize() {
if (nToReturn) {
int remaining = nToReturn - nReturned;

if ( nToReturn == 0 )
return batchSize;
if (batchSize && batchSize < remaining)
return batchSize;

if ( batchSize == 0 )
return nToReturn;
return -remaining;
}

return batchSize < nToReturn ? batchSize : nToReturn;
return batchSize;
}

void DBClientCursor::_assembleInit( Message& toSend ) {
Expand All @@ -89,7 +91,7 @@ namespace mongo {
BufBuilder b;
b.appendNum( opts );
b.appendStr( ns );
b.appendNum( nToReturn );
b.appendNum( nextBatchSize() );
b.appendNum( cursorId );
toSend.setData( dbGetMore, b.buf(), b.len() );
}
Expand Down Expand Up @@ -172,10 +174,6 @@ namespace mongo {
void DBClientCursor::requestMore() {
verify( cursorId && batch.pos == batch.nReturned );

if (haveLimit) {
nToReturn -= batch.nReturned;
verify(nToReturn > 0);
}
BufBuilder b;
b.appendNum(opts);
b.appendStr(ns);
Expand Down Expand Up @@ -206,7 +204,7 @@ namespace mongo {
/** with QueryOption_Exhaust, the server just blasts data at us (marked at end with cursorid==0). */
void DBClientCursor::exhaustReceiveMore() {
verify( cursorId && batch.pos == batch.nReturned );
verify( !haveLimit );
verify( !nToReturn );
auto_ptr<Message> response(new Message());
verify( _client );
if (!_client->recv(*response)) {
Expand Down Expand Up @@ -254,7 +252,7 @@ namespace mongo {
bool DBClientCursor::rawMore() {
DEV _assertIfNull();

if (haveLimit && batch.pos >= nToReturn)
if (nToReturn && nReturned >= nToReturn)
return false;

if ( batch.pos < batch.nReturned )
Expand Down Expand Up @@ -294,6 +292,8 @@ namespace mongo {
BSONObj DBClientCursor::next() {
DEV _assertIfNull();

nReturned++;

if ( !_putBack.empty() ) {
BSONObj ret = _putBack.top();
_putBack.pop();
Expand Down
5 changes: 2 additions & 3 deletions src/mongo/client/dbclientcursor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,8 +136,7 @@ namespace mongo {

DBClientCursor( DBClientBase* client, const std::string &_ns, BSONObj _query, int _nToReturn,
int _nToSkip, const BSONObj *_fieldsToReturn, int queryOptions , int bs );

DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options );
DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options, int _batchSize );

virtual ~DBClientCursor();

Expand Down Expand Up @@ -204,8 +203,8 @@ namespace mongo {
std::string ns;
BSONObj query;
int nToReturn;
bool haveLimit;
int nToSkip;
long long nReturned;
const BSONObj *fieldsToReturn;
int opts;
int batchSize;
Expand Down
2 changes: 1 addition & 1 deletion src/mongo/client/dbclientinterface.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -546,7 +546,7 @@ namespace mongo {
/** query N objects from the database into an array. makes sense mostly when you want a small number of results. if a huge number, use
query() and iterate the cursor.
*/
void findN(std::vector<BSONObj>& out, const std::string&ns, Query query, int nToReturn, int nToSkip = 0, const BSONObj *fieldsToReturn = 0, int queryOptions = 0);
void findN(std::vector<BSONObj>& out, const std::string&ns, Query query, int nToReturn, int nToSkip = 0, const BSONObj *fieldsToReturn = 0, int queryOptions = 0, int batchSize = 0);

virtual std::string getServerAddress() const = 0;

Expand Down
2 changes: 1 addition & 1 deletion src/mongo/dbtests/mock/mock_dbclient_cursor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
namespace mongo {
MockDBClientCursor::MockDBClientCursor(mongo::DBClientBase* client,
const mongo::BSONArray& resultSet):
mongo::DBClientCursor(client, "", 0, 0, 0) {
mongo::DBClientCursor(client, "", 0, 0, 0, 0) {
_resultSet = resultSet.copy();
_cursor.reset(new mongo::DBClientMockCursor(BSONArray(_resultSet)));
}
Expand Down
21 changes: 15 additions & 6 deletions src/mongo/dbtests/mock/mock_remote_db_server.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,12 +161,14 @@ namespace mongo {
mongo::BSONArray MockRemoteDBServer::query(
MockRemoteDBServer::InstanceID id,
const string& ns,
mongo::Query query,
mongo::Query, /* don't support */
int nToReturn,
int nToSkip,
const BSONObj* fieldsToReturn,
int queryOptions,
int batchSize) {
const BSONObj*, /* don't support */
int, /* don't support */
int) /* don't support */ {
int seen = 0;

checkIfUp(id);

if (_delayMilliSec > 0) {
Expand All @@ -180,8 +182,15 @@ namespace mongo {

const vector<BSONObj>& coll = _dataMgr[ns];
BSONArrayBuilder result;
for (vector<BSONObj>::const_iterator iter = coll.begin(); iter != coll.end(); ++ iter) {
result.append(iter->copy());
for (vector<BSONObj>::const_iterator iter = coll.begin(); iter != coll.end(); ++iter) {
if (nToSkip) {
nToSkip--;
}
else {
if (nToReturn && (seen == nToReturn)) break;
result.append(iter->copy());
seen++;
}
}

return BSONArray(result.obj());
Expand Down
11 changes: 11 additions & 0 deletions src/mongo/dbtests/mock_dbclient_conn_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,17 @@ namespace mongo_test {

ASSERT(!cursor->more());
}

{
MockDBClientConnection conn(&server);
std::auto_ptr<mongo::DBClientCursor> cursor = conn.query(ns, mongo::Query(), 1);

ASSERT(cursor->more());
BSONObj firstDoc = cursor->next();
ASSERT_EQUALS(1, firstDoc["x"].numberInt());

ASSERT(!cursor->more());
}
}

TEST(MockDBClientConnTest, InsertAndQueryTwice) {
Expand Down
32 changes: 26 additions & 6 deletions src/mongo/unittest/dbclient_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,8 @@ namespace {
for(int i = 0; i < 3; ++i) {
c.insert(TEST_NS, BSON("num" << i));
}
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 2);
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 0, 0, 0, 0, 2);

uint64_t cursor_id = cursor->getCursorId();
cursor->decouple();
cursor.reset();
Expand Down Expand Up @@ -662,7 +663,7 @@ namespace {
TEST_F(DBClientTest, LazyCursorCommand) {
c.setRunCommandHook(nop_hook);
c.setPostRunCommandHook(nop_hook_post);
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, -1, 0, 0, 0, 0);
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, 1, 0, 0, 0, 0);
bool is_retry = false;
cursor.initLazy(is_retry);
ASSERT_TRUE(cursor.initLazyFinish(is_retry));
Expand All @@ -673,7 +674,7 @@ namespace {
}

TEST_F(DBClientTest, InitCommand) {
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, -1, 0, 0, 0, 0);
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, 1, 0, 0, 0, 0);
ASSERT_TRUE(cursor.initCommand());

ASSERT_TRUE(cursor.more());
Expand All @@ -682,21 +683,40 @@ namespace {
}

TEST_F(DBClientTest, GetMoreLimit) {
// TODO: really implement limit
c.insert(TEST_NS, BSON("num" << 1));
c.insert(TEST_NS, BSON("num" << 2));
c.insert(TEST_NS, BSON("num" << 3));
c.insert(TEST_NS, BSON("num" << 4));

// set nToReturn to 3 but batch size to 2
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 3, 0, 0, 0, 2);
// set nToReturn to 3 but batch size to 1
// This verifies:
// * we can manage with multiple batches
// * we can correctly upgrade batchSize 1 to 2 to avoid automatic
// cursor closing when nReturn = 1 (wire protocol edge case)
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 3, 0, 0, 0, 1);
vector<BSONObj> docs;
while(cursor->more())
docs.push_back(cursor->next());

ASSERT_EQUALS(docs.size(), 3U);
}

TEST_F(DBClientTest, NoGetMoreLimit) {
c.insert(TEST_NS, BSON("num" << 1));
c.insert(TEST_NS, BSON("num" << 2));
c.insert(TEST_NS, BSON("num" << 3));
c.insert(TEST_NS, BSON("num" << 4));

// set nToReturn to 2 but batch size to 4
// check to see if a limit of 2 takes despite the larger batch
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 2, 0, 0, 0, 4);
vector<BSONObj> docs;
while(cursor->more())
docs.push_back(cursor->next());

ASSERT_EQUALS(docs.size(), 2U);
}

TEST_F(DBClientTest, PeekError) {
BSONObj result;
c.runCommand("admin", BSON("buildinfo" << true), result);
Expand Down

0 comments on commit 2e86826

Please sign in to comment.