mirror of
https://github.com/cloudwu/skynet.git
synced 2026-07-23 19:43:09 +00:00
new socket lib
This commit is contained in:
云风
@@ -1,273 +1,240 @@
|
||||
local buffer = require "socketbuffer"
|
||||
local driver = require "socketdriver"
|
||||
local skynet = require "skynet"
|
||||
local table = table
|
||||
local next = next
|
||||
local assert = assert
|
||||
local coroutine = coroutine
|
||||
local type = type
|
||||
|
||||
local READBUF = {} -- fd:buffer
|
||||
local READREQUEST = {} -- fd:request_size
|
||||
local READSESSION = {} -- fd:session
|
||||
local READLOCK = {} -- fd:queue(session)
|
||||
local READTHREAD= {} -- fd:thread
|
||||
local CLOSED = {} -- fd:true
|
||||
local socket = {} -- api
|
||||
local buffer_pool = {} -- store all message buffer object
|
||||
local socket_pool = setmetatable( -- store all socket object
|
||||
{},
|
||||
{ __gc = function(p)
|
||||
for id,v in pairs(p) do
|
||||
driver.close(id)
|
||||
-- don't need clear v.buffer, because buffer pool will be free at the end
|
||||
p[id] = nil
|
||||
end
|
||||
end
|
||||
}
|
||||
)
|
||||
|
||||
local selfaddr = skynet.self()
|
||||
local sockets = assert(skynet.localname ".socket")
|
||||
local socket_message = {}
|
||||
|
||||
local function response(session)
|
||||
skynet.redirect(selfaddr , 0, "response", session, "")
|
||||
local function wakeup(s)
|
||||
local co = s.co
|
||||
if co then
|
||||
s.co = nil
|
||||
skynet.wakeup(co)
|
||||
end
|
||||
end
|
||||
|
||||
local function suspend(s)
|
||||
assert(not s.co)
|
||||
s.co = coroutine.running()
|
||||
skynet.wait()
|
||||
end
|
||||
|
||||
-- read skynet_socket.h for these macro
|
||||
-- SKYNET_SOCKET_TYPE_DATA = 1
|
||||
socket_message[1] = function(id, size, data)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
print("socket: drop package from " .. id)
|
||||
driver.drop(data)
|
||||
return
|
||||
end
|
||||
|
||||
local sz = driver.push(s.buffer, buffer_pool, data, size)
|
||||
local rr = s.read_required
|
||||
local rrt = type(rr)
|
||||
if rrt == "number" then
|
||||
-- read size
|
||||
if sz >= rr then
|
||||
s.read_required = nil
|
||||
wakeup(s)
|
||||
end
|
||||
elseif rrt == "string" then
|
||||
-- read line
|
||||
if driver.readline(s.buffer,nil,rr) then
|
||||
s.read_required = nil
|
||||
wakeup(s)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- SKYNET_SOCKET_TYPE_CONNECT = 2
|
||||
socket_message[2] = function(id, _ , addr)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
return
|
||||
end
|
||||
-- log remote addr
|
||||
s.connected = true
|
||||
wakeup(s)
|
||||
end
|
||||
|
||||
-- SKYNET_SOCKET_TYPE_CLOSE = 3
|
||||
socket_message[3] = function(id)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
return
|
||||
end
|
||||
s.connected = false
|
||||
wakeup(s)
|
||||
end
|
||||
|
||||
-- SKYNET_SOCKET_TYPE_ACCEPT = 4
|
||||
socket_message[4] = function(id, newid, addr)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
driver.close(newid)
|
||||
return
|
||||
end
|
||||
skynet.fork(s.callback, newid, addr)
|
||||
end
|
||||
|
||||
-- SKYNET_SOCKET_TYPE_ERROR = 5
|
||||
socket_message[5] = function(id)
|
||||
print("error on ", id)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
return
|
||||
end
|
||||
s.connected = false
|
||||
wakeup(s)
|
||||
end
|
||||
|
||||
skynet.register_protocol {
|
||||
name = "client",
|
||||
id = 3, -- PTYPE_CLIENT
|
||||
pack = buffer.pack,
|
||||
unpack = buffer.unpack,
|
||||
dispatch = function (_, _, fd, msg, sz)
|
||||
local qsz = READREQUEST[fd]
|
||||
local buf = READBUF[fd]
|
||||
local bsz
|
||||
if sz == 0 then
|
||||
CLOSED[fd] = true
|
||||
else
|
||||
buf,bsz = buffer.push(buf, msg, sz)
|
||||
READBUF[fd] = buf
|
||||
end
|
||||
local session = READSESSION[fd]
|
||||
if qsz == nil or session == nil then
|
||||
return
|
||||
end
|
||||
if sz > 0 then
|
||||
if type(qsz) == "number" then
|
||||
if qsz > bsz then
|
||||
return
|
||||
end
|
||||
else
|
||||
-- request readline
|
||||
if buffer.readline(buf, qsz, true) == nil then
|
||||
return
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
response(session)
|
||||
READSESSION[fd] = nil
|
||||
name = "socket",
|
||||
id = 6, -- PTYPE_SOCKET
|
||||
unpack = driver.unpack,
|
||||
dispatch = function (_, _, t, n1, n2, data)
|
||||
socket_message[t](n1,n2,data)
|
||||
end
|
||||
}
|
||||
|
||||
skynet.register_protocol {
|
||||
name = "system",
|
||||
id = 4, -- PTYPE_SYSTEM
|
||||
pack = skynet.pack,
|
||||
unpack = function (...) return ... end,
|
||||
dispatch = function (session, addr, msg, sz)
|
||||
fd, t, sz = skynet.unpack(msg,sz)
|
||||
assert(addr == selfaddr, "PTYPE_SYSTEM message must send by self")
|
||||
if t > 0 then -- lock request when t == 0
|
||||
-- request bytes or readline
|
||||
local buf = READBUF[fd]
|
||||
if CLOSED[fd] then
|
||||
skynet.ret()
|
||||
return
|
||||
end
|
||||
local _,bsz = buffer.push(buf)
|
||||
if t == 1 then
|
||||
-- sz is size
|
||||
if bsz >= sz then
|
||||
skynet.ret()
|
||||
return
|
||||
end
|
||||
elseif t == 2 then
|
||||
-- sz is sep
|
||||
if buffer.readline(buf, sz, true) then -- don't real read
|
||||
skynet.ret()
|
||||
return
|
||||
end
|
||||
end
|
||||
READSESSION[fd] = session
|
||||
end
|
||||
local function connect(id)
|
||||
local s = {
|
||||
id = id,
|
||||
buffer = driver.buffer(),
|
||||
connected = false,
|
||||
read_require = false,
|
||||
co = false,
|
||||
}
|
||||
socket_pool[id] = s
|
||||
suspend(s)
|
||||
if s.connected then
|
||||
return id
|
||||
end
|
||||
}
|
||||
|
||||
local socket = {}
|
||||
end
|
||||
|
||||
function socket.open(addr, port)
|
||||
local cmd = "open" .. " " .. (port and (addr..":"..port) or addr)
|
||||
local r = skynet.call(sockets, "text", cmd)
|
||||
if r == "" then
|
||||
return nil, cmd .. " failed"
|
||||
end
|
||||
local fd = tonumber(r)
|
||||
READBUF[fd] = true
|
||||
CLOSED[fd] = nil
|
||||
return fd
|
||||
end
|
||||
|
||||
function socket.bind(sock)
|
||||
local r = skynet.call(sockets, "text", "bind " .. tonumber(sock))
|
||||
if r == "" then
|
||||
error("stdin bind failed")
|
||||
end
|
||||
local fd = tonumber(r)
|
||||
READBUF[fd] = true
|
||||
CLOSED[fd] = nil
|
||||
return fd
|
||||
local id = driver.connect(addr,port)
|
||||
return connect(id)
|
||||
end
|
||||
|
||||
function socket.stdin()
|
||||
return socket.bind(1)
|
||||
local id = driver.bind(1)
|
||||
return connect(id)
|
||||
end
|
||||
|
||||
function socket.accept(id)
|
||||
driver.accept(id)
|
||||
return connect(id)
|
||||
end
|
||||
|
||||
function socket.close(fd)
|
||||
socket.lock(fd)
|
||||
skynet.call(sockets, "text", "close", fd)
|
||||
READBUF[fd] = nil
|
||||
READLOCK[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
-- socket.lock(fd)
|
||||
local s = socket_pool[id]
|
||||
if s == nil then
|
||||
return
|
||||
end
|
||||
if s.connected then
|
||||
driver.close(s.id)
|
||||
suspend(s)
|
||||
end
|
||||
if s.buffer then
|
||||
driver.clear(s.buffer)
|
||||
end
|
||||
socket_pool[id] = nil
|
||||
end
|
||||
|
||||
function socket.read(fd, sz)
|
||||
local buf = assert(READBUF[fd])
|
||||
local str, bytes = buffer.pop(buf,sz)
|
||||
if str then
|
||||
return str
|
||||
end
|
||||
|
||||
if CLOSED[fd] then
|
||||
READBUF[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
str = buffer.pop(buf, bytes)
|
||||
return nil, str
|
||||
end
|
||||
|
||||
READREQUEST[fd] = sz
|
||||
skynet.call(selfaddr, "system",fd,1,sz) -- singal size 1
|
||||
READREQUEST[fd] = nil
|
||||
|
||||
buf = READBUF[fd]
|
||||
|
||||
str, bytes = buffer.pop(buf,sz)
|
||||
if str then
|
||||
return str
|
||||
end
|
||||
|
||||
if CLOSED[fd] then
|
||||
READBUF[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
str = buffer.pop(buf, bytes)
|
||||
return nil, str
|
||||
end
|
||||
end
|
||||
|
||||
function socket.readall(fd)
|
||||
local buf = assert(READBUF[fd])
|
||||
if CLOSED[fd] then
|
||||
CLOSED[fd] = nil
|
||||
READBUF[fd] = nil
|
||||
if buf == nil then
|
||||
return ""
|
||||
end
|
||||
local _, bytes = buffer.push(buf)
|
||||
local ret = buffer.pop(buf, bytes)
|
||||
function socket.read(id, sz)
|
||||
local s = socket_pool[id]
|
||||
assert(s)
|
||||
local ret = driver.pop(s.buffer, buffer_pool, sz)
|
||||
if ret then
|
||||
return ret
|
||||
end
|
||||
READREQUEST[fd] = math.huge
|
||||
skynet.call(selfaddr, "system",fd,3) -- singal readall
|
||||
READREQUEST[fd] = nil
|
||||
assert(CLOSED[fd])
|
||||
buf = READBUF[fd]
|
||||
READBUF[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
if buf == nil then
|
||||
return ""
|
||||
end
|
||||
local _, bytes = buffer.push(buf)
|
||||
local ret = buffer.pop(buf, bytes)
|
||||
return ret
|
||||
end
|
||||
|
||||
function socket.readline(fd, sep)
|
||||
local buf = assert(READBUF[fd])
|
||||
local str = buffer.readline(buf,sep)
|
||||
if str then
|
||||
return str
|
||||
if not s.connected then
|
||||
return false, driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
|
||||
if CLOSED[fd] then
|
||||
READBUF[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
local _, bytes = buffer.push(buf)
|
||||
str = buffer.pop(buf, bytes)
|
||||
return nil, str
|
||||
end
|
||||
|
||||
READREQUEST[fd] = sep
|
||||
skynet.call(selfaddr, "system",fd,2,sep) -- singal sep 2
|
||||
READREQUEST[fd] = nil
|
||||
|
||||
buf = READBUF[fd]
|
||||
str = buffer.readline(buf,sep)
|
||||
if str then
|
||||
return str
|
||||
end
|
||||
|
||||
if CLOSED[fd] then
|
||||
READBUF[fd] = nil
|
||||
CLOSED[fd] = nil
|
||||
local _, bytes = buffer.push(buf)
|
||||
str = buffer.pop(buf, bytes)
|
||||
return nil, str
|
||||
end
|
||||
end
|
||||
|
||||
function socket.write(fd, msg, sz)
|
||||
if CLOSED[fd] or not READBUF[fd] then
|
||||
return
|
||||
end
|
||||
skynet.send(sockets, "client", fd, msg, sz)
|
||||
return true
|
||||
end
|
||||
|
||||
function socket.invalid(fd)
|
||||
return CLOSED[fd] or not READBUF[fd]
|
||||
end
|
||||
|
||||
function socket.lock(fd)
|
||||
if CLOSED[fd] or not READBUF[fd] then
|
||||
return
|
||||
end
|
||||
local locked = READTHREAD[fd]
|
||||
if locked then
|
||||
-- lock fd
|
||||
local session = skynet.genid()
|
||||
local q = READLOCK[fd]
|
||||
if q == nil then
|
||||
READLOCK[fd] = { session }
|
||||
else
|
||||
table.insert(q, session)
|
||||
end
|
||||
|
||||
skynet.redirect(selfaddr , 0, "system", session, skynet.pack(fd,0))
|
||||
coroutine.yield("CALL",session)
|
||||
assert(not s.read_required)
|
||||
s.read_required = sz
|
||||
suspend(s)
|
||||
ret = driver.pop(s.buffer, buffer_pool, sz)
|
||||
if ret then
|
||||
return ret
|
||||
else
|
||||
READTHREAD[fd] = true
|
||||
return false, driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
end
|
||||
|
||||
function socket.readall(id)
|
||||
local s = socket_pool[id]
|
||||
assert(s)
|
||||
if not s.connected then
|
||||
return driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
assert(not s.read_required)
|
||||
s.read_required = true
|
||||
suspend(s)
|
||||
assert(s.connected == false)
|
||||
return driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
|
||||
function socket.readline(id, sep)
|
||||
sep = sep or "\n"
|
||||
local s = socket_pool[id]
|
||||
assert(s)
|
||||
local ret = driver.readline(s.buffer, buffer_pool, sep)
|
||||
if ret then
|
||||
return ret
|
||||
end
|
||||
if not s.connected then
|
||||
return false, driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
assert(not s.read_required)
|
||||
s.read_required = sep
|
||||
suspend(s)
|
||||
if s.connected then
|
||||
return driver.readline(s.buffer, buffer_pool, sep)
|
||||
else
|
||||
return false, driver.readall(s.buffer, buffer_pool)
|
||||
end
|
||||
end
|
||||
|
||||
socket.write = assert(driver.send)
|
||||
|
||||
function socket.invalid(id)
|
||||
return socket_pool[id] == nil
|
||||
end
|
||||
|
||||
function socket.listen(host,port,func)
|
||||
local id = driver.listen(host,port)
|
||||
local s = {
|
||||
id = id,
|
||||
connected = true,
|
||||
callback = func
|
||||
}
|
||||
socket_pool[id] = s
|
||||
return id
|
||||
end
|
||||
|
||||
function socket.lock(id)
|
||||
end
|
||||
|
||||
function socket.unlock(fd)
|
||||
READTHREAD[fd] = nil
|
||||
local q = READLOCK[fd]
|
||||
if q then
|
||||
if q[1] then
|
||||
READTHREAD[fd] = true
|
||||
response(q[1])
|
||||
table.remove(q,1)
|
||||
else
|
||||
READLOCK[fd] = nil
|
||||
end
|
||||
end
|
||||
|
||||
end
|
||||
|
||||
return socket
|
||||
|
||||
Reference in New Issue
Block a user