Skip to content

Commit 2e86826

Browse files
hanumantmkhanumantmk
committed
CXX-176 implement nToReturn
Implemented nToReturn with actual limit semantics.
1 parent 87f67a8 commit 2e86826

File tree

8 files changed

+83
-41
lines changed

8 files changed

+83
-41
lines changed

‎src/mongo/client/dbclient.cpp

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@
4040
#include "mongo/util/net/ssl_manager.h"
4141
#include "mongo/util/password_digest.h"
4242

43+
#include <algorithm>
44+
#include <cstdlib>
45+
4346
namespace mongo {
4447

4548
using std::auto_ptr;
@@ -857,11 +860,11 @@ namespace mongo {
857860
/** 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
858861
query() and iterate the cursor.
859862
*/
860-
void DBClientInterface::findN(vector<BSONObj>& out, const string& ns, Query query, int nToReturn, int nToSkip, const BSONObj *fieldsToReturn, int queryOptions) {
861-
out.reserve(nToReturn);
863+
void DBClientInterface::findN(vector<BSONObj>& out, const string& ns, Query query, int nToReturn, int nToSkip, const BSONObj *fieldsToReturn, int queryOptions, int batchSize) {
864+
out.reserve(std::min(batchSize, nToReturn));
862865

863866
auto_ptr<DBClientCursor> c =
864-
this->query(ns, query, nToReturn, nToSkip, fieldsToReturn, queryOptions);
867+
this->query(ns, query, nToReturn, nToSkip, fieldsToReturn, queryOptions, batchSize);
865868

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

@@ -1077,7 +1080,7 @@ namespace mongo {
10771080

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

1080-
c->haveLimit = false;
1083+
c->nToReturn = 0;
10811084

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

11031106
simple->shim.reset(new DBClientCursorShimArray(*simple));
1104-
simple->haveLimit = false;
1107+
simple->nToReturn = 0;
11051108

11061109
return simple;
11071110
}
@@ -1112,7 +1115,7 @@ namespace mongo {
11121115
}
11131116

11141117
auto_ptr<DBClientCursor> DBClientBase::getMore( const string &ns, long long cursorId, int nToReturn, int options ) {
1115-
auto_ptr<DBClientCursor> c( new DBClientCursor( this, ns, cursorId, nToReturn, options ) );
1118+
auto_ptr<DBClientCursor> c( new DBClientCursor( this, ns, cursorId, nToReturn < 0 ? abs(nToReturn) : 0, options, abs(nToReturn)) );
11161119
if ( c->init() )
11171120
return c;
11181121
return auto_ptr< DBClientCursor >( 0 );

‎src/mongo/client/dbclientcursor.cpp

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ namespace mongo {
3535
_client(client),
3636
ns(_ns),
3737
query(_query),
38-
nToReturn(_nToReturn),
39-
haveLimit( _nToReturn > 0 && !(queryOptions & QueryOption_CursorTailable)),
38+
nToReturn( (queryOptions & QueryOption_CursorTailable) ? 0 : _nToReturn ),
4039
nToSkip(_nToSkip),
40+
nReturned(0),
4141
fieldsToReturn(_fieldsToReturn),
4242
opts(queryOptions),
4343
batchSize(bs==1?2:bs),
@@ -48,15 +48,15 @@ namespace mongo {
4848
_finishConsInit();
4949
}
5050

51-
DBClientCursor::DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options ) :
51+
DBClientCursor::DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options, int bs ) :
5252
_client(client),
5353
ns(_ns),
54-
nToReturn( _nToReturn ),
55-
haveLimit( _nToReturn > 0 && !(options & QueryOption_CursorTailable)),
54+
nToReturn( (options & QueryOption_CursorTailable) ? 0 : _nToReturn ),
5655
nToSkip(0),
56+
nReturned(0),
5757
fieldsToReturn(0),
5858
opts( options ),
59-
batchSize(0),
59+
batchSize(bs==1?2:bs),
6060
resultFlags(0),
6161
cursorId(_cursorId),
6262
_ownCursor(true),
@@ -71,14 +71,16 @@ namespace mongo {
7171
}
7272

7373
int DBClientCursor::nextBatchSize() {
74+
if (nToReturn) {
75+
int remaining = nToReturn - nReturned;
7476

75-
if ( nToReturn == 0 )
76-
return batchSize;
77+
if (batchSize && batchSize < remaining)
78+
return batchSize;
7779

78-
if ( batchSize == 0 )
79-
return nToReturn;
80+
return -remaining;
81+
}
8082

81-
return batchSize < nToReturn ? batchSize : nToReturn;
83+
return batchSize;
8284
}
8385

8486
void DBClientCursor::_assembleInit( Message& toSend ) {
@@ -89,7 +91,7 @@ namespace mongo {
8991
BufBuilder b;
9092
b.appendNum( opts );
9193
b.appendStr( ns );
92-
b.appendNum( nToReturn );
94+
b.appendNum( nextBatchSize() );
9395
b.appendNum( cursorId );
9496
toSend.setData( dbGetMore, b.buf(), b.len() );
9597
}
@@ -172,10 +174,6 @@ namespace mongo {
172174
void DBClientCursor::requestMore() {
173175
verify( cursorId && batch.pos == batch.nReturned );
174176

175-
if (haveLimit) {
176-
nToReturn -= batch.nReturned;
177-
verify(nToReturn > 0);
178-
}
179177
BufBuilder b;
180178
b.appendNum(opts);
181179
b.appendStr(ns);
@@ -206,7 +204,7 @@ namespace mongo {
206204
/** with QueryOption_Exhaust, the server just blasts data at us (marked at end with cursorid==0). */
207205
void DBClientCursor::exhaustReceiveMore() {
208206
verify( cursorId && batch.pos == batch.nReturned );
209-
verify( !haveLimit );
207+
verify( !nToReturn );
210208
auto_ptr<Message> response(new Message());
211209
verify( _client );
212210
if (!_client->recv(*response)) {
@@ -254,7 +252,7 @@ namespace mongo {
254252
bool DBClientCursor::rawMore() {
255253
DEV _assertIfNull();
256254

257-
if (haveLimit && batch.pos >= nToReturn)
255+
if (nToReturn && nReturned >= nToReturn)
258256
return false;
259257

260258
if ( batch.pos < batch.nReturned )
@@ -294,6 +292,8 @@ namespace mongo {
294292
BSONObj DBClientCursor::next() {
295293
DEV _assertIfNull();
296294

295+
nReturned++;
296+
297297
if ( !_putBack.empty() ) {
298298
BSONObj ret = _putBack.top();
299299
_putBack.pop();

‎src/mongo/client/dbclientcursor.h

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,8 +136,7 @@ namespace mongo {
136136

137137
DBClientCursor( DBClientBase* client, const std::string &_ns, BSONObj _query, int _nToReturn,
138138
int _nToSkip, const BSONObj *_fieldsToReturn, int queryOptions , int bs );
139-
140-
DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options );
139+
DBClientCursor( DBClientBase* client, const std::string &_ns, long long _cursorId, int _nToReturn, int options, int _batchSize );
141140

142141
virtual ~DBClientCursor();
143142

@@ -204,8 +203,8 @@ namespace mongo {
204203
std::string ns;
205204
BSONObj query;
206205
int nToReturn;
207-
bool haveLimit;
208206
int nToSkip;
207+
long long nReturned;
209208
const BSONObj *fieldsToReturn;
210209
int opts;
211210
int batchSize;

‎src/mongo/client/dbclientinterface.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -546,7 +546,7 @@ namespace mongo {
546546
/** 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
547547
query() and iterate the cursor.
548548
*/
549-
void findN(std::vector<BSONObj>& out, const std::string&ns, Query query, int nToReturn, int nToSkip = 0, const BSONObj *fieldsToReturn = 0, int queryOptions = 0);
549+
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);
550550

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

‎src/mongo/dbtests/mock/mock_dbclient_cursor.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
namespace mongo {
2121
MockDBClientCursor::MockDBClientCursor(mongo::DBClientBase* client,
2222
const mongo::BSONArray& resultSet):
23-
mongo::DBClientCursor(client, "", 0, 0, 0) {
23+
mongo::DBClientCursor(client, "", 0, 0, 0, 0) {
2424
_resultSet = resultSet.copy();
2525
_cursor.reset(new mongo::DBClientMockCursor(BSONArray(_resultSet)));
2626
}

‎src/mongo/dbtests/mock/mock_remote_db_server.cpp

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -161,12 +161,14 @@ namespace mongo {
161161
mongo::BSONArray MockRemoteDBServer::query(
162162
MockRemoteDBServer::InstanceID id,
163163
const string& ns,
164-
mongo::Query query,
164+
mongo::Query, /* don't support */
165165
int nToReturn,
166166
int nToSkip,
167-
const BSONObj* fieldsToReturn,
168-
int queryOptions,
169-
int batchSize) {
167+
const BSONObj*, /* don't support */
168+
int, /* don't support */
169+
int) /* don't support */ {
170+
int seen = 0;
171+
170172
checkIfUp(id);
171173

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

181183
const vector<BSONObj>& coll = _dataMgr[ns];
182184
BSONArrayBuilder result;
183-
for (vector<BSONObj>::const_iterator iter = coll.begin(); iter != coll.end(); ++ iter) {
184-
result.append(iter->copy());
185+
for (vector<BSONObj>::const_iterator iter = coll.begin(); iter != coll.end(); ++iter) {
186+
if (nToSkip) {
187+
nToSkip--;
188+
}
189+
else {
190+
if (nToReturn && (seen == nToReturn)) break;
191+
result.append(iter->copy());
192+
seen++;
193+
}
185194
}
186195

187196
return BSONArray(result.obj());

‎src/mongo/dbtests/mock_dbclient_conn_test.cpp

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,17 @@ namespace mongo_test {
109109

110110
ASSERT(!cursor->more());
111111
}
112+
113+
{
114+
MockDBClientConnection conn(&server);
115+
std::auto_ptr<mongo::DBClientCursor> cursor = conn.query(ns, mongo::Query(), 1);
116+
117+
ASSERT(cursor->more());
118+
BSONObj firstDoc = cursor->next();
119+
ASSERT_EQUALS(1, firstDoc["x"].numberInt());
120+
121+
ASSERT(!cursor->more());
122+
}
112123
}
113124

114125
TEST(MockDBClientConnTest, InsertAndQueryTwice) {

‎src/mongo/unittest/dbclient_test.cpp

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,8 @@ namespace {
252252
for(int i = 0; i < 3; ++i) {
253253
c.insert(TEST_NS, BSON("num" << i));
254254
}
255-
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 2);
255+
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 0, 0, 0, 0, 2);
256+
256257
uint64_t cursor_id = cursor->getCursorId();
257258
cursor->decouple();
258259
cursor.reset();
@@ -662,7 +663,7 @@ namespace {
662663
TEST_F(DBClientTest, LazyCursorCommand) {
663664
c.setRunCommandHook(nop_hook);
664665
c.setPostRunCommandHook(nop_hook_post);
665-
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, -1, 0, 0, 0, 0);
666+
DBClientCursor cursor(&c, TEST_DB + ".$cmd", Query("{dbStats: 1}").obj, 1, 0, 0, 0, 0);
666667
bool is_retry = false;
667668
cursor.initLazy(is_retry);
668669
ASSERT_TRUE(cursor.initLazyFinish(is_retry));
@@ -673,7 +674,7 @@ namespace {
673674
}
674675

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

679680
ASSERT_TRUE(cursor.more());
@@ -682,21 +683,40 @@ namespace {
682683
}
683684

684685
TEST_F(DBClientTest, GetMoreLimit) {
685-
// TODO: really implement limit
686686
c.insert(TEST_NS, BSON("num" << 1));
687687
c.insert(TEST_NS, BSON("num" << 2));
688688
c.insert(TEST_NS, BSON("num" << 3));
689689
c.insert(TEST_NS, BSON("num" << 4));
690690

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

697701
ASSERT_EQUALS(docs.size(), 3U);
698702
}
699703

704+
TEST_F(DBClientTest, NoGetMoreLimit) {
705+
c.insert(TEST_NS, BSON("num" << 1));
706+
c.insert(TEST_NS, BSON("num" << 2));
707+
c.insert(TEST_NS, BSON("num" << 3));
708+
c.insert(TEST_NS, BSON("num" << 4));
709+
710+
// set nToReturn to 2 but batch size to 4
711+
// check to see if a limit of 2 takes despite the larger batch
712+
auto_ptr<DBClientCursor> cursor = c.query(TEST_NS, Query("{}"), 2, 0, 0, 0, 4);
713+
vector<BSONObj> docs;
714+
while(cursor->more())
715+
docs.push_back(cursor->next());
716+
717+
ASSERT_EQUALS(docs.size(), 2U);
718+
}
719+
700720
TEST_F(DBClientTest, PeekError) {
701721
BSONObj result;
702722
c.runCommand("admin", BSON("buildinfo" << true), result);

0 commit comments

Comments
 (0)