lihao/skynet
1
0
Fork 0
mirror of https://github.com/cloudwu/skynet.git synced 2026-07-23 11:33:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'mongors' into dev

Browse Source
This commit is contained in:
Cloud Wu
2014-07-21 14:28:09 +08:00
parent a8b80cd73e 39ff6eb2a6
commit d9d2a7d2ce
2 changed files with 171 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

210
lualib/mongo.lua
View File

@@ -1,96 +1,125 @@
local bson = require "bson"
local socket = require "socket"
local socketchannel = require "socketchannel"
local socketchannel = require "socketchannel"
local skynet = require "skynet"
local driver = require "mongo.driver"
local md5 = require "md5"
local md5 = require "md5"
local rawget = rawget
local assert = assert
local bson_encode = bson.encode
local bson_encode_order = bson.encode_order
local bson_decode = bson.decode
local bson_encode = bson.encode
local bson_encode_order = bson.encode_order
local bson_decode = bson.decode
local empty_bson = bson_encode {}
local mongo = {}
local mongo = {}
mongo.null = assert(bson.null)
mongo.maxkey = assert(bson.maxkey)
mongo.minkey = assert(bson.minkey)
mongo.type = assert(bson.type)
local mongo_cursor = {}
local cursor_meta = {
__index = mongo_cursor,
local cursor_meta = {
__index = mongo_cursor,
}
local mongo_client = {}
local client_meta = {
__index = function(self, key)
return rawget(mongo_client, key) or self:getDB(key)
local client_meta = {
__index = function(self, key)
return rawget(mongo_client, key) or self:getDB(key)
end,
__tostring = function (self)
local port_string
if self.port then
port_string = ":" .. tostring(self.port)
port_string = ":" .. tostring(self.port)
else
port_string = ""
port_string = ""
end
return "[mongo client : " .. self.host .. port_string .."]"
return "[mongo client : " .. self.host .. port_string .."]"
end,
-- DO NOT need disconnect, because channel will shutdown during gc
-- DO NOT need disconnect, because channel will shutdown during gc
}
local mongo_db = {}
local db_meta = {
__index = function (self, key)
return rawget(mongo_db, key) or self:getCollection(key)
local db_meta = {
__index = function (self, key)
return rawget(mongo_db, key) or self:getCollection(key)
end,
__tostring = function (self)
return "[mongo db : " .. self.name .. "]"
return "[mongo db : " .. self.name .. "]"
end
}
local mongo_collection = {}
local collection_meta = {
__index = function(self, key)
return rawget(mongo_collection, key) or self:getCollection(key)
end ,
local collection_meta = {
__index = function(self, key)
return rawget(mongo_collection, key) or self:getCollection(key)
end ,
__tostring = function (self)
return "[mongo collection : " .. self.full_name .. "]"
return "[mongo collection : " .. self.full_name .. "]"
end
}
local function dispatch_reply(so)
local len_reply = so:read(4)
local reply = so:read(driver.length(len_reply))
local result = { result = {} }
local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result)
result.document = document
local len_reply = so:read(4)
local reply = so:read(driver.length(len_reply))
local result = { result = {} }
local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result)
result.document = document
result.cursor_id = cursor_id
result.startfrom = startfrom
result.data = reply
result.data = reply
return reply_id, succ, result
end
local function mongo_auth(mongoc)
local user = rawget(mongoc, "username")
local pass = rawget(mongoc, "password")
local function __parse_addr(addr)
local host, port = string.match(addr, "([^:]+):(.+)")
return host, tonumber(port)
end
local function mongo_auth(mongoc)
local user = rawget(mongoc, "username")
local pass = rawget(mongoc, "password")
if user == nil or pass == nil then
return
end
return function()
assert(mongoc:auth(user, pass))
if user ~= nil and pass ~= nil then
assert(mongoc:auth(user, pass))
end
local rs_data = mongoc:runCommand("ismaster")
if rs_data.ok == 1 then
if rs_data.hosts then
local backup = {}
for _, v in ipairs(rs_data.hosts) do
local host, port = __parse_addr(v)
table.insert(backup, {host = host, port = port})
end
mongoc.__sock:changebackup(backup)
end
if rs_data.ismaster then
return
else
local host, port = __parse_addr(rs_data.primary)
mongoc.host = host
mongoc.port = port
mongoc.__sock:changehost(host, port)
end
end
end
end
function mongo.client( conf )
local obj = {
host = conf.host,
port = conf.port or 27017,
function mongo.client( conf )
local first = conf
local backup = nil
if conf.rs then
first = conf.rs[1]
backup = conf.rs
end
local obj = {
host = first.host,
port = first.port or 27017,
}
obj.__id = 0
@@ -99,9 +128,10 @@ function mongo.client( conf )
port = obj.port,
response = dispatch_reply,
auth = mongo_auth(obj),
backup = backup,
}
setmetatable(obj, client_meta)
obj.__sock:connect(true) -- try connect only once
obj.__sock:connect(true) -- try connect only once
return obj
end
@@ -109,32 +139,32 @@ function mongo_client:getDB(dbname)
local db = {
connection = self,
name = dbname,
full_name = dbname,
full_name = dbname,
database = false,
__cmd = dbname .. "." .. "$cmd",
__cmd = dbname .. "." .. "$cmd",
}
db.database = db
db.database = db
return setmetatable(db, db_meta)
return setmetatable(db, db_meta)
end
function mongo_client:disconnect()
if self.__sock then
local so = self.__sock
self.__sock = false
self.__sock = false
so:close()
end
end
function mongo_client:genId()
local id = self.__id + 1
self.__id = id
self.__id = id
return id
end
function mongo_client:runCommand(...)
if not self.admin then
self.admin = self:getDB "admin"
self.admin = self:getDB "admin"
end
return self.admin:runCommand(...)
end
@@ -146,14 +176,14 @@ function mongo_client:auth(user,password)
return false
end
local key = md5.sumhexa(string.format("%s%s%s",result.nonce,user,password))
local key = md5.sumhexa(string.format("%s%s%s",result.nonce,user,password))
local result= self:runCommand ("authenticate",1,"user",user,"nonce",result.nonce,"key",key)
return result.ok == 1
return result.ok == 1
end
function mongo_client:logout()
local result = self:runCommand "logout"
return result.ok == 1
return result.ok == 1
end
function mongo_db:runCommand(cmd,cmd_v,...)
@@ -166,18 +196,18 @@ function mongo_db:runCommand(cmd,cmd_v,...)
else
bson_cmd = bson_encode_order(cmd,cmd_v,...)
end
local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id)
local doc = req.document
local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id)
local doc = req.document
return bson_decode(doc)
end
function mongo_db:getCollection(collection)
local col = {
local col = {
connection = self.connection,
name = collection,
full_name = self.full_name .. "." .. collection,
full_name = self.full_name .. "." .. collection,
database = self.database,
}
self[collection] = setmetatable(col, collection_meta)
@@ -188,20 +218,20 @@ mongo_collection.getCollection = mongo_db.getCollection
function mongo_collection:insert(doc)
if doc._id == nil then
doc._id = bson.objectid()
doc._id = bson.objectid()
end
local sock = self.connection.__sock
local pack = driver.insert(0, self.full_name, bson_encode(doc))
-- flags support 1: ContinueOnError
-- flags support 1: ContinueOnError
sock:request(pack)
end
function mongo_collection:batch_insert(docs)
for i=1,#docs do
for i=1,#docs do
if docs[i]._id == nil then
docs[i]._id = bson.objectid()
docs[i]._id = bson.objectid()
end
docs[i] = bson_encode(docs[i])
docs[i] = bson_encode(docs[i])
end
local sock = self.connection.__sock
local pack = driver.insert(0, self.full_name, docs)
@@ -209,7 +239,7 @@ function mongo_collection:batch_insert(docs)
end
function mongo_collection:update(selector,update,upsert,multi)
local flags = (upsert and 1 or 0) + (multi and 2 or 0)
local flags = (upsert and 1 or 0) + (multi and 2 or 0)
local sock = self.connection.__sock
local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update))
sock:request(pack)
@@ -225,24 +255,24 @@ function mongo_collection:findOne(query, selector)
local conn = self.connection
local request_id = conn:genId()
local sock = conn.__sock
local pack = driver.query(request_id, 0, self.full_name, 0, 1, query and bson_encode(query) or empty_bson, selector and bson_encode(selector))
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id)
local doc = req.document
local pack = driver.query(request_id, 0, self.full_name, 0, 1, query and bson_encode(query) or empty_bson, selector and bson_encode(selector))
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id)
local doc = req.document
return bson_decode(doc)
end
function mongo_collection:find(query, selector)
return setmetatable( {
__collection = self,
__query = query and bson_encode(query) or empty_bson,
__query = query and bson_encode(query) or empty_bson,
__selector = selector and bson_encode(selector),
__ptr = nil,
__ptr = nil,
__data = nil,
__cursor = nil,
__document = {},
__flags = 0,
} , cursor_meta)
__flags = 0,
} , cursor_meta)
end
function mongo_cursor:hasNext()
@@ -255,42 +285,42 @@ function mongo_cursor:hasNext()
local sock = conn.__sock
local pack
if self.__data == nil then
pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector)
pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector)
else
if self.__cursor then
pack = driver.more(request_id, self.__collection.full_name,0,self.__cursor)
else
-- no more
self.__document = nil
self.__data = nil
self.__document = nil
self.__data = nil
return false
end
end
local ok, result = pcall(sock.request,sock,pack, request_id)
local doc = result.document
local doc = result.document
local cursor = result.cursor_id
if ok then
if doc then
self.__document = result.result
self.__data = result.data
self.__document = result.result
self.__data = result.data
self.__ptr = 1
self.__cursor = cursor
self.__cursor = cursor
return true
else
self.__document = nil
self.__data = nil
self.__cursor = nil
self.__document = nil
self.__data = nil
self.__cursor = nil
return false
end
else
self.__document = nil
self.__data = nil
self.__cursor = nil
self.__document = nil
self.__data = nil
self.__cursor = nil
if doc then
local err = bson_decode(doc)
local err = bson_decode(doc)
error(err["$err"])
else
error("Reply from mongod error")
@@ -303,11 +333,11 @@ end
function mongo_cursor:next()
if self.__ptr == nil then
error "Call hasNext first"
error "Call hasNext first"
end
local r = bson_decode(self.__document[self.__ptr])
self.__ptr = self.__ptr + 1
if self.__ptr > #self.__document then
local r = bson_decode(self.__document[self.__ptr])
self.__ptr = self.__ptr + 1
if self.__ptr > #self.__document then
self.__ptr = nil
end

55
lualib/socketchannel.lua
View File

@@ -26,6 +26,7 @@ function socket_channel.channel(desc)
local c = {
__host = assert(desc.host),
__port = assert(desc.port),
__backup = desc.backup,
__auth = desc.auth,
__response = desc.response, -- It's for session mode
__request = {}, -- request seq { response func or session } -- It's for order mode
@@ -152,34 +153,66 @@ local function dispatch_function(self)
end
end
local function connect_backup(self)
if self.__backup then
for _, addr in ipairs(self.__backup) do
local host, port
if type(addr) == "table" then
host, port = addr.host, addr.port
else
host = addr
port = self.__port
end
skynet.error("socket: connect to backup host", host, port)
local fd = socket.open(host, port)
if fd then
self.__host = host
self.__port = port
return fd
end
end
end
end
local function connect_once(self)
if self.__closed then
return false
end
assert(not self.__sock and not self.__authcoroutine)
local fd = socket.open(self.__host, self.__port)
if not fd then
return false
fd = connect_backup(self)
if not fd then
return false
end
end
self.__authcoroutine = coroutine.running()
self.__sock = setmetatable( {fd} , channel_socket_meta )
skynet.fork(dispatch_function(self), self)
if self.__auth then
self.__authcoroutine = coroutine.running()
local ok , message = pcall(self.__auth, self)
if not ok then
close_channel_socket(self)
if message ~= socket_error then
self.__authcoroutine = false
skynet.error("socket: auth failed", message)
end
end
self.__authcoroutine = false
if ok and not self.__sock then
-- auth may change host, so connect again
return connect_once(self)
end
return ok
end
self.__authcoroutine = false
return true
end
local function try_connect(self , once)
local t = 100
local t = 0
while not self.__closed do
if connect_once(self) then
if not once then
@@ -292,6 +325,20 @@ function channel:close()
end
end
function channel:changehost(host, port)
self.__host = host
if port then
self.__port = port
end
if not self.__closed then
close_channel_socket(self)
end
end
function channel:changebackup(backup)
self.__backup = backup
end
channel_meta.__gc = channel.close
local function wrapper_socket_function(f)