Mongodb op msg database commands fix (#4004)

* * Fix: MongoDB::OpMsgCursor did not handle zero batch size properly: cursor requests failed.
* Improvement: Add emptyFirstBatch to indicate that the size of the first batch shall be zero for performance to get potential error ASAP from the server.

* Poco::MongoDB: Some database commands do not need collection as an argument. An integer "1" is passed instead.

MongoDB/include/Poco/MongoDB/Database.h

@@ -99,7 +99,7 @@ public:
 		/// Creates OpMsgMessage. (new wire protocol)
 
 	Poco::SharedPtr<Poco::MongoDB::OpMsgMessage> createOpMsgMessage() const;
-		/// Creates OpMsgMessage for database commands. (new wire protocol)
+		/// Creates OpMsgMessage for database commands that do not require collection as an argument. (new wire protocol)
 
 	Poco::SharedPtr<Poco::MongoDB::OpMsgCursor> createOpMsgCursor(const std::string& collectionName) const;
 		/// Creates OpMsgCursor. (new wire protocol)
@@ -215,8 +215,8 @@ Database::createOpMsgMessage(const std::string& collectionName) const
 inline Poco::SharedPtr<Poco::MongoDB::OpMsgMessage>
 Database::createOpMsgMessage() const
 {
-	// Collection name for database commands is ignored and any value will do.
+	// Collection name for database commands is not needed.
-	return createOpMsgMessage("1");
+	return createOpMsgMessage("");
 }
 
 inline Poco::SharedPtr<Poco::MongoDB::OpMsgCursor>

MongoDB/include/Poco/MongoDB/OpMsgCursor.h

@@ -36,11 +36,16 @@ public:
 	virtual ~OpMsgCursor();
 		/// Destroys the OpMsgCursor.
 
+	void setEmptyFirstBatch(bool empty);
+		/// Empty first batch is used to get error response faster with little server processing
+
+	bool emptyFirstBatch() const;
+
 	void setBatchSize(Int32 batchSize);
 		/// Set non-default batch size
 
 	Int32 batchSize() const;
-		/// Current batch size (negative number indicates default batch size)
+		/// Current batch size (zero or negative number indicates default batch size)
 
 	Int64 cursorID() const;
 
@@ -60,8 +65,9 @@ private:
 	OpMsgMessage    _query;
 	OpMsgMessage 	_response;
 
+	bool			_emptyFirstBatch { false };
 	Int32			_batchSize { -1 };
-		/// Batch size used in the cursor. Negative value means that default shall be used.
+		/// Batch size used in the cursor. Zero or negative value means that default shall be used.
 
 	Int64			_cursorID { 0 };
 };

MongoDB/include/Poco/MongoDB/OpMsgMessage.h

@@ -52,6 +52,8 @@ public:
 
 	// Replication and administration 
 	static const std::string CMD_HELLO;
+	static const std::string CMD_REPL_SET_GET_STATUS;
+	static const std::string CMD_REPL_SET_GET_CONFIG;
 
 	static const std::string CMD_CREATE;
 	static const std::string CMD_CREATE_INDEXES;

MongoDB/src/OpMsgCursor.cpp

@@ -71,6 +71,18 @@ OpMsgCursor::~OpMsgCursor()
 }
 
 
+void OpMsgCursor::setEmptyFirstBatch(bool empty)
+{
+	_emptyFirstBatch = empty;
+}
+
+
+bool OpMsgCursor::emptyFirstBatch() const
+{
+	return _emptyFirstBatch;
+}
+
+
 void OpMsgCursor::setBatchSize(Int32 batchSize)
 {
 	_batchSize = batchSize;
@@ -89,16 +101,18 @@ OpMsgMessage& OpMsgCursor::next(Connection& connection)
 	{
 		_response.clear();
 
-		if (_query.commandName() == OpMsgMessage::CMD_FIND)
+		if (_emptyFirstBatch || _batchSize > 0)
 		{
-			if (_batchSize >= 0)
+			Int32 bsize = _emptyFirstBatch ? 0 : _batchSize;
-				_query.body().add("batchSize", _batchSize);
+			if (_query.commandName() == OpMsgMessage::CMD_FIND)
-		}
+			{
-		else if (_query.commandName() == OpMsgMessage::CMD_AGGREGATE)
+				_query.body().add("batchSize", bsize);
-		{
+			}
-			auto cursorDoc = _query.body().addNewDocument("cursor");
+			else if (_query.commandName() == OpMsgMessage::CMD_AGGREGATE)
-			if (_batchSize >= 0)
+			{
-				cursorDoc.add("batchSize", _batchSize);
+				auto& cursorDoc = _query.body().addNewDocument("cursor");
+				cursorDoc.add("batchSize", bsize);
+			}
 		}
 
 		connection.sendRequest(_query, _response);

MongoDB/src/OpMsgMessage.cpp

@@ -38,6 +38,8 @@ const std::string OpMsgMessage::CMD_MAP_REDUCE { "mapReduce" };
 
 // Replication and administration 
 const std::string OpMsgMessage::CMD_HELLO { "hello" };
+const std::string OpMsgMessage::CMD_REPL_SET_GET_STATUS { "replSetGetStatus" };
+const std::string OpMsgMessage::CMD_REPL_SET_GET_CONFIG { "replSetGetConfig" };
 
 const std::string OpMsgMessage::CMD_CREATE { "create" };
 const std::string OpMsgMessage::CMD_CREATE_INDEXES { "createIndexes" };
@@ -100,7 +102,16 @@ void OpMsgMessage::setCommandName(const std::string& command)
 	_body.clear();
 
 	// IMPORTANT: Command name must be first
-	_body.add(_commandName, _collectionName);
+	if (_collectionName.empty())
+	{
+		// Collection is not specified. It is assumed that this particular command does 
+		// not need it.
+		_body.add(_commandName, Int32(1));
+	}
+	else
+	{
+		_body.add(_commandName, _collectionName);
+	}
 	_body.add("$db", _databaseName);
 }
 
@@ -114,7 +125,7 @@ void OpMsgMessage::setCursor(Poco::Int64 cursorID, Poco::Int32 batchSize)
 	_body.add(_commandName, cursorID);
 	_body.add("$db", _databaseName);
 	_body.add("collection", _collectionName);
-	if (batchSize >= 0)
+	if (batchSize > 0)
 	{
 		_body.add("batchSize", batchSize);
 	}

MongoDB/testsuite/src/MongoDBTest.cpp

@@ -556,7 +556,11 @@ CppUnit::Test* MongoDBTest::suite()
 		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdCursor);
 		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdCursorAggregate);
 		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdKillCursor);
+		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdCursorEmptyFirstBatch);
+		
 		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdUUID);
+
+		CppUnit_addTest(pSuite, MongoDBTest, testOpCmdDropDatabase);		
 	}
 
 	return pSuite;

MongoDB/testsuite/src/MongoDBTest.h

@@ -57,11 +57,13 @@ public:
 	void testOpCmdFind();
 	void testOpCmdCursor();
 	void testOpCmdCursorAggregate();
+	void testOpCmdCursorEmptyFirstBatch();
 	void testOpCmdKillCursor();
 	void testOpCmdCount();
 	void testOpCmdDelete();
 	void testOpCmdUnaknowledgedInsert();
 	void testOpCmdConnectionPool();
+	void testOpCmdDropDatabase();
 
 	static CppUnit::Test* suite();
 

MongoDB/testsuite/src/MongoDBTestOpMsg.cpp

@@ -327,7 +327,11 @@ void MongoDBTest::testOpCmdCursorAggregate()
 	auto cresponse = cursor->next(*_mongo);
 	while(true)
 	{
-		n += static_cast<int>(cresponse.documents().size());
+		int batchDocSize = cresponse.documents().size();
+		if (cursor->cursorID() != 0)
+			assertEquals (1000, batchDocSize);
+
+		n += batchDocSize;
 		if ( cursor->cursorID() == 0 )
 			break;
 		cresponse = cursor->next(*_mongo);
@@ -400,6 +404,53 @@ void MongoDBTest::testOpCmdCount()
 }
 
 
+void MongoDBTest::testOpCmdCursorEmptyFirstBatch()
+{
+	Database db("team");
+
+	Poco::SharedPtr<OpMsgMessage> request = db.createOpMsgMessage("numbers");
+	OpMsgMessage response;
+
+	request->setCommandName(OpMsgMessage::CMD_DROP);
+	_mongo->sendRequest(*request, response);
+
+	request->setCommandName(OpMsgMessage::CMD_INSERT);
+	for(int i = 0; i < 10000; ++i)
+	{
+		Document::Ptr doc = new Document();
+		doc->add("number", i);
+		request->documents().push_back(doc);
+	}
+	_mongo->sendRequest(*request, response);
+	assertTrue(response.responseOk());
+
+	Poco::SharedPtr<OpMsgCursor> cursor = db.createOpMsgCursor("numbers");
+	cursor->query().setCommandName(OpMsgMessage::CMD_AGGREGATE);
+	cursor->setEmptyFirstBatch(true);
+	cursor->setBatchSize(0); // Will be ignored, default is used
+
+	// Empty pipeline: get all documents
+	cursor->query().body().addNewArray("pipeline");
+
+	auto cresponse = cursor->next(*_mongo);
+	assertEquals (0, cresponse.documents().size()); // First batch is empty
+
+	int n = 0;
+	while(true)
+	{
+		n += static_cast<int>(cresponse.documents().size());
+		if ( cursor->cursorID() == 0 )
+			break;
+		cresponse = cursor->next(*_mongo);
+	}
+	assertEquals (10000, n);
+
+	request->setCommandName(OpMsgMessage::CMD_DROP);
+	_mongo->sendRequest(*request, response);
+	assertTrue(response.responseOk());
+}
+
+
 void MongoDBTest::testOpCmdDelete()
 {
 	Database db("team");
@@ -442,3 +493,21 @@ void MongoDBTest::testOpCmdConnectionPool()
 	assertEquals (1, doc.getInteger("n"));
 }
 
+
+void MongoDBTest::testOpCmdDropDatabase()
+{
+	Database db("team");
+	Poco::SharedPtr<OpMsgMessage> request = db.createOpMsgMessage();
+	request->setCommandName(OpMsgMessage::CMD_DROP_DATABASE);
+
+	OpMsgMessage response;
+	_mongo->sendRequest(*request, response);
+
+	std::cout << request->body().toString(2) << std::endl;
+	std::cout << response.body().toString(2) << std::endl;
+
+	assertTrue(response.responseOk());
+}
+
+
+