mongodb 交互协议调整为 OP_MSG (#1649)

* mongodb 交互协议调整为 OP_MSG

* Update

Co-authored-by: yy151474 <yy151474@alibaba-inc.com>

lualib-src/lua-mongo.c

@@ -10,21 +10,17 @@
 #include <stdio.h>
 #include <string.h>
 
-#define OP_REPLY 1
-#define OP_MSG	1000
-#define OP_UPDATE 2001
-#define OP_INSERT 2002
-#define OP_QUERY 2004
-#define OP_GET_MORE	2005
-#define OP_DELETE 2006
-#define OP_KILL_CURSORS	2007
+#define OP_COMPRESSED 2012
+#define OP_MSG 2013
+
+typedef enum {
+	MSG_CHECKSUM_PRESENT = 1 << 0,
+	MSG_MORE_TO_COME = 1 << 1,
+	MSG_EXHAUST_ALLOWED = 1 << 16,
+} msg_flags_t;
 
-#define REPLY_CURSORNOTFOUND 1
-#define REPLY_QUERYFAILURE 2
-#define REPLY_AWAITCAPABLE 8	// ignore because mongo 1.6+ always set it
 
 #define DEFAULT_CAP 128
-
 struct connection {
 	int sock;
 	int id;
@@ -106,6 +102,13 @@ write_int32(struct buffer *b, int32_t v) {
 	b->ptr[b->size++] = (uv >> 24)&0xff;
 }
 
+static inline void
+write_int8(struct buffer *b, int8_t v) {
+	uint8_t uv = (uint8_t)v;
+	buffer_reserve(b, 1);
+	b->ptr[b->size++] = uv;
+}
+
 static inline void
 write_bytes(struct buffer *b, const void * buf, int sz) {
 	buffer_reserve(b,sz);
@@ -113,6 +116,7 @@ write_bytes(struct buffer *b, const void * buf, int sz) {
 	b->size += sz;
 }
 
+/*
 static void
 write_string(struct buffer *b, const char *key, size_t sz) {
 	buffer_reserve(b,sz+1);
@@ -120,6 +124,7 @@ write_string(struct buffer *b, const char *key, size_t sz) {
 	b->ptr[b->size+sz] = '\0';
 	b->size+=sz+1;
 }
+*/
 
 static inline int
 reserve_length(struct buffer *b) {
@@ -138,389 +143,69 @@ write_length(struct buffer *b, int32_t v, int off) {
 	b->ptr[off++] = (uv >> 24)&0xff;
 }
 
-// 1 integer id
-// 2 integer flags
-// 3 string collection name
-// 4 integer skip
-// 5 integer return number
-// 6 document query
-// 7 document selector (optional)
-// return string package
-static int
-op_query(lua_State *L) {
-	int id = luaL_checkinteger(L,1);
-	document query = lua_touserdata(L,6);
-	if (query == NULL) {
-		return luaL_error(L, "require query document");
-	}
-	document selector = lua_touserdata(L,7);
-	int flags = luaL_checkinteger(L, 2);
-	size_t nsz = 0;
-	const char *name = luaL_checklstring(L,3,&nsz);
-	int skip = luaL_checkinteger(L, 4);
-	int number = luaL_checkinteger(L, 5);
+struct header_t {
+	//int32_t message_length; 	// total message size, include this
+	int32_t request_id;			// identifier for this message
+	int32_t response_to;		// requestID from the original request(used in responses from the database)
+	int32_t opcode;			// message type
 
-	luaL_Buffer b;
-	luaL_buffinit(L,&b);
-
-	struct buffer buf;
-	buffer_create(&buf);
-		int len = reserve_length(&buf);
-		write_int32(&buf, id);
-		write_int32(&buf, 0);
-		write_int32(&buf, OP_QUERY);
-		write_int32(&buf, flags);
-		write_string(&buf, name, nsz);
-		write_int32(&buf, skip);
-		write_int32(&buf, number);
-
-		int32_t query_len = get_length(query);
-		int total = buf.size + query_len;
-		int32_t selector_len = 0;
-		if (selector) {
-			selector_len = get_length(selector);
-			total += selector_len;
-		}
-
-		write_length(&buf, total, len);
-		luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-
-	luaL_addlstring(&b, (const char *)query, query_len);
-
-	if (selector) {
-		luaL_addlstring(&b, (const char *)selector, selector_len);
-	}
-
-	luaL_pushresult(&b);
-
-	return 1;
-}
+	int32_t flags;
+};
 
 // 1 string data
 // 2 result document table
 // return boolean succ (false -> request id, error document)
 //	number request_id
 //  document first
-//	string cursor_id
-//  integer startfrom
 static int
-op_reply(lua_State *L) {
+unpack_reply(lua_State *L) {
 	size_t data_len = 0;
 	const char * data = luaL_checklstring(L,1,&data_len);
-	struct reply_type {
-//		int32_t length; // total message size, including this
-		int32_t request_id; // identifier for this message
-		int32_t response_id; // requestID from the original request
-							// (used in reponses from db)
-		int32_t opcode; // request type 
-		int32_t flags;
-		int32_t cursor_id[2];
-		int32_t starting;
-		int32_t number;
-	};
-	const struct reply_type* reply = (const struct reply_type*)data;
+	const struct header_t* h = (const struct header_t*)data;
 
-	if (data_len < sizeof(*reply)) {
+	if (data_len < sizeof(*h)) {
 		lua_pushboolean(L, 0);
 		return 1;
 	}
 
-	int id = little_endian(reply->response_id);
-	int flags = little_endian(reply->flags);
-	if (flags & REPLY_QUERYFAILURE) {
-		lua_pushboolean(L,0);
-		lua_pushinteger(L, id);
-		lua_pushlightuserdata(L, (void *)(reply+1));
-		return 3;
+	int opcode = little_endian(h->opcode);
+	if (opcode != OP_MSG) {
+		return luaL_error(L, "Unsupported opcode:%d", opcode);
 	}
 
-	int starting_from = little_endian(reply->starting);
-	int number = little_endian(reply->number);
-	int sz = (int)data_len - sizeof(*reply);
-	const uint8_t * doc = (const uint8_t *)(reply+1);
+	int id = little_endian(h->response_to);
+	int flags = little_endian(h->flags);
 
-	if (lua_istable(L,2)) {
-		int i = 1;
-		while (sz > 4) {
-			lua_pushlightuserdata(L, (void *)doc);
-			lua_rawseti(L, 2, i);
-
-			int32_t doc_len = get_length((document)doc);
-			if (doc_len <= 0) {
-				return luaL_error(L, "Invalid result bson document");
-			}
-
-			doc += doc_len;
-			sz -= doc_len;
-
-			++i;
+	if (flags != 0) {
+		if ((flags & MSG_CHECKSUM_PRESENT) != 0)  {
+			return luaL_error(L, "Unsupported OP_MSG flag checksumPresent");
 		}
-		if (i != number + 1) {
-			lua_pushboolean(L,0);
-			lua_pushinteger(L, id);
-			return 2;
-		}
-		int c = lua_rawlen(L, 2);
-		for (;i<=c;i++) {
-			lua_pushnil(L);
-			lua_rawseti(L, 2, i);
-		}
-	} else {
-		if (sz >= 4) {
-			sz -= get_length((document)doc);
+
+		if ((flags ^ MSG_MORE_TO_COME) != 0) {
+			return luaL_error(L, "Unsupported OP_MSG flag:%d", flags);
 		}
 	}
-	if (sz != 0) {
-		return luaL_error(L, "Invalid result bson document");
+
+	int sz = (int)data_len - sizeof(*h);
+
+	const uint8_t * section = (const uint8_t *)(h+1);
+
+	uint8_t payload_type = *section;
+	const uint8_t * doc = section+1;
+
+	if (payload_type != 0) {
+		return luaL_error(L, "Unsupported OP_MSG payload type: %d", payload_type);
 	}
-	lua_pushboolean(L,1);
+
+	int32_t doc_sz = get_length((document)(doc));
+	if ((sz - 1) != doc_sz) {
+		return luaL_error(L, "Unsupported OP_MSG reply: >1 section");
+	}
+
+	lua_pushboolean(L, 1);
 	lua_pushinteger(L, id);
-	if (number == 0)
-		lua_pushnil(L);
-	else
-		lua_pushlightuserdata(L, (void *)(reply+1));
-	if (reply->cursor_id[0] == 0 && reply->cursor_id[1]==0) {
-		// closed cursor
-		lua_pushnil(L);
-	} else {
-		lua_pushlstring(L, (const char *)(reply->cursor_id), 8);
-	}
-	lua_pushinteger(L, starting_from);
-
-	return 5;
-}
-
-/*
-	1 string cursor_id
-	return string package
- */
-static int
-op_kill(lua_State *L) {
-	size_t cursor_len = 0;
-	const char * cursor_id = luaL_tolstring(L, 1, &cursor_len);
-	if (cursor_len != 8) {
-		return luaL_error(L, "Invalid cursor id");
-	}
-
-	struct buffer buf;
-	buffer_create(&buf);
-
-	int len = reserve_length(&buf);
-	write_int32(&buf, 0);
-	write_int32(&buf, 0);
-	write_int32(&buf, OP_KILL_CURSORS);
-
-	write_int32(&buf, 0);
-	write_int32(&buf, 1);
-	write_bytes(&buf, cursor_id, 8);
-
-	write_length(&buf, buf.size, len);
-
-	lua_pushlstring(L, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-
-	return 1;
-}
-
-/*
-	1 string collection
-	2 integer single remove
-	3 document selector
-
-	return string package
- */
-static int
-op_delete(lua_State *L) {
-	document selector  = lua_touserdata(L,3);
-	if (selector == NULL) {
-		luaL_error(L, "Invalid param");
-	}
-	size_t sz = 0;
-	const char * name = luaL_checklstring(L,1,&sz);
-
-	luaL_Buffer b;
-	luaL_buffinit(L,&b);
-
-	struct buffer buf;
-	buffer_create(&buf);
-		int len = reserve_length(&buf);
-		write_int32(&buf, 0);
-		write_int32(&buf, 0);
-		write_int32(&buf, OP_DELETE);
-		write_int32(&buf, 0);
-		write_string(&buf, name, sz);
-		write_int32(&buf, lua_tointeger(L,2));
-
-		int32_t selector_len = get_length(selector);
-		int total = buf.size + selector_len;
-		write_length(&buf, total, len);
-
-		luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-
-	luaL_addlstring(&b, (const char *)selector, selector_len);
-	luaL_pushresult(&b);
-
-	return 1;
-}
-
-/*
-	1 integer id
-	2 string collection
-	3 integer number
-	4 cursor_id (8 bytes string/ 64bit)
-
-	return string package
- */
-static int
-op_get_more(lua_State *L) {
-	int id = luaL_checkinteger(L, 1);
-	size_t sz = 0;
-	const char * name = luaL_checklstring(L,2,&sz);
-	int number = luaL_checkinteger(L, 3);
-	size_t cursor_len = 0;
-	const char * cursor_id = luaL_tolstring(L, 4, &cursor_len);
-	if (cursor_len != 8) {
-		return luaL_error(L, "Invalid cursor id");
-	}
-
-	struct buffer buf;
-	buffer_create(&buf);
-		int len = reserve_length(&buf);
-		write_int32(&buf, id);
-		write_int32(&buf, 0);
-		write_int32(&buf, OP_GET_MORE);
-		write_int32(&buf, 0);
-		write_string(&buf, name, sz);
-		write_int32(&buf, number);
-		write_bytes(&buf, cursor_id, 8);
-		write_length(&buf, buf.size, len);
-
-		lua_pushlstring(L, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-
-	return 1;
-}
-
-// 1 string collection
-// 2 integer flags
-// 3 document selector
-// 4 document update
-// return string package
-static int
-op_update(lua_State *L) {
-	document selector  = lua_touserdata(L,3);
-	document update = lua_touserdata(L,4);
-	if (selector == NULL || update == NULL) {
-		luaL_error(L, "Invalid param");
-	}
-	size_t sz = 0;
-	const char * name = luaL_checklstring(L,1,&sz);
-
-	luaL_Buffer b;
-	luaL_buffinit(L, &b);
-
-	struct buffer buf;
-	buffer_create(&buf);
-		// make package header, don't raise L error
-		int len = reserve_length(&buf);
-		write_int32(&buf, 0);
-		write_int32(&buf, 0);
-		write_int32(&buf, OP_UPDATE);
-		write_int32(&buf, 0);
-		write_string(&buf, name, sz);
-		write_int32(&buf, lua_tointeger(L,2));
-
-		int32_t selector_len = get_length(selector);
-		int32_t update_len = get_length(update);
-
-		int total = buf.size + selector_len + update_len;
-		write_length(&buf, total, len);
-
-		luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-
-	luaL_addlstring(&b, (const char *)selector, selector_len);
-	luaL_addlstring(&b, (const char *)update, update_len);
-
-	luaL_pushresult(&b);
-
-	return 1;
-}
-
-static int
-document_length(lua_State *L) {
-	if (lua_isuserdata(L, 3)) {
-		document doc = lua_touserdata(L,3);
-		return get_length(doc);
-	}
-	if (lua_istable(L,3)) {
-		int total = 0;
-		int s = lua_rawlen(L,3);
-		int i;
-		for (i=1;i<=s;i++) {
-			lua_rawgeti(L, 3, i);
-			document doc = lua_touserdata(L,-1);
-			if (doc == NULL) {
-				lua_pop(L,1);
-				return luaL_error(L, "Invalid document at %d", i);
-			} else {
-				total += get_length(doc);
-				lua_pop(L,1);
-			}
-		}
-		return total;
-	}
-	return luaL_error(L, "Insert need documents");
-}
-
-// 1 integer flags
-// 2 string collection
-// 3 documents
-// return string package
-static int
-op_insert(lua_State *L) {
-	size_t sz = 0;
-	const char * name = luaL_checklstring(L,2,&sz);
-	int dsz = document_length(L);
-
-	luaL_Buffer b;
-	luaL_buffinit(L, &b);
-
-	struct buffer buf;
-	buffer_create(&buf);
-		// make package header, don't raise L error
-		int len = reserve_length(&buf);
-		write_int32(&buf, 0);
-		write_int32(&buf, 0);
-		write_int32(&buf, OP_INSERT);
-		write_int32(&buf, lua_tointeger(L,1));
-		write_string(&buf, name, sz);
-
-		int total = buf.size + dsz;
-		write_length(&buf, total, len);
-
-		luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
-	buffer_destroy(&buf);
-	
-	if (lua_isuserdata(L,3)) {
-		document doc = lua_touserdata(L,3);
-		luaL_addlstring(&b, (const char *)doc, get_length(doc));
-	} else {
-		int s = lua_rawlen(L, 3);
-		int i;
-		for (i=1;i<=s;i++) {
-			lua_rawgeti(L,3,i);
-			document doc = lua_touserdata(L,-1);
-			lua_pop(L,1);	// must call lua_pop before luaL_addlstring, because addlstring may change stack top
-			luaL_addlstring(&b, (const char *)doc, get_length(doc));
-		}
-	}
-
-	luaL_pushresult(&b);
-
-	return 1;
+	lua_pushlightuserdata(L, (void *)(doc));
+	return 3;
 }
 
 // string 4 bytes length
@@ -535,21 +220,55 @@ reply_length(lua_State *L) {
 	return 1;
 }
 
+// @param 1 request_id int
+// @param 2 flags int
+// @param 3 command bson document
+// @return
+static int
+op_msg(lua_State *L) {
+	int id = luaL_checkinteger(L, 1);
+	int flags = luaL_checkinteger(L, 2);
+	document cmd = lua_touserdata(L, 3);
+
+	if (cmd == NULL) {
+		return luaL_error(L, "opmsg require cmd document");
+	}
+
+	luaL_Buffer b;
+	luaL_buffinit(L, &b);
+
+	struct buffer buf;
+	buffer_create(&buf);
+		int len = reserve_length(&buf);
+		write_int32(&buf, id);
+		write_int32(&buf, 0);
+		write_int32(&buf, OP_MSG);
+		write_int32(&buf, flags);
+		write_int8(&buf, 0);
+
+		int32_t cmd_len = get_length(cmd);
+		int total = buf.size + cmd_len;
+
+		write_length(&buf, total, len);
+		luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
+	buffer_destroy(&buf);
+
+	luaL_addlstring(&b, (const char *)cmd, cmd_len);
+	luaL_pushresult(&b);
+	return 1;
+}
+
+
 LUAMOD_API int
 luaopen_skynet_mongo_driver(lua_State *L) {
 	luaL_checkversion(L);
 	luaL_Reg l[] ={
-		{ "query", op_query },
-		{ "reply", op_reply },
-		{ "kill", op_kill },
-		{ "delete", op_delete },
-		{ "more", op_get_more },
-		{ "update", op_update },
-		{ "insert", op_insert },
+		{ "reply", unpack_reply }, // 接收响应
 		{ "length", reply_length },
+		{ "op_msg", op_msg},
 		{ NULL, NULL },
 	};
 
 	luaL_newlib(L,l);
 	return 1;
-}
+}