PageRenderTime 49ms CodeModel.GetById 33ms app.highlight 11ms RepoModel.GetById 2ms app.codeStats 0ms

/lualib/socket.lua

https://github.com/huigenius/skynet
Lua | 189 lines | 161 code | 23 blank | 5 comment | 37 complexity | e22ace262ce78eee4eaa05790047160a MD5 | raw file
  1local buffer = require "socketbuffer"
  2local skynet = require "skynet"
  3local table = table
  4local next = next
  5local assert = assert
  6local coroutine = coroutine
  7local type = type
  8
  9local READBUF = {}	-- fd:buffer
 10local READREQUEST = {}	-- fd:request_size
 11local READSESSION = {}	-- fd:session
 12local READLOCK = {}	-- fd:queue(session)
 13local READTHREAD= {} -- fd:thread
 14
 15local selfaddr = skynet.self()
 16
 17local function response(session)
 18	skynet.redirect(selfaddr , 0, "response", session, "")
 19end
 20
 21skynet.register_protocol {
 22	name = "client",
 23	id = 3,	-- PTYPE_CLIENT
 24	pack = buffer.pack,
 25	unpack = buffer.unpack,
 26	dispatch = function (_, _, fd, msg, sz)
 27		local qsz = READREQUEST[fd]
 28		local buf = READBUF[fd]
 29		local bsz
 30		if sz == 0 or buf == true then
 31			buf,bsz = true, qsz
 32		else
 33			buf,bsz = buffer.push(buf, msg, sz)
 34		end
 35		READBUF[fd] = buf
 36		local session = READSESSION[fd]
 37		if qsz == nil or session == nil then
 38			return
 39		end
 40		if type(qsz) == "number" then
 41			if qsz > bsz then
 42				return
 43			end
 44		else
 45			-- request readline
 46			if buffer.readline(buf, qsz, true) == nil then
 47				return
 48			end
 49		end
 50
 51		response(session)
 52		READSESSION[fd] = nil
 53	end
 54}
 55
 56skynet.register_protocol {
 57	name = "system",
 58	id = 4, -- PTYPE_SYSTEM
 59	pack = skynet.pack,
 60	unpack = function (...) return ... end,
 61	dispatch = function (session, addr, msg, sz)
 62		fd, t, sz = skynet.unpack(msg,sz)
 63		assert(addr == selfaddr, "PTYPE_SYSTEM message must send by self")
 64		if t == 0 then
 65			-- lock fd
 66			if READTHREAD[fd] == nil then
 67				skynet.ret()
 68				return
 69			end
 70
 71			local q = READLOCK[fd]
 72			if q == nil then
 73				READLOCK[fd] = { session }
 74			else
 75				table.insert(q, session)
 76			end
 77		else
 78			-- request bytes or readline
 79			local buf = READBUF[fd]
 80			if buf == true then
 81				skynet.ret()
 82				return
 83			end
 84			local _,bsz = buffer.push(buf)
 85			if t == 1 then
 86				-- sz is size
 87				if bsz >= sz then
 88					skynet.ret()
 89					return
 90				end
 91			else
 92				-- sz is sep
 93				if buffer.readline(buf, sz, true) then -- don't real read
 94					skynet.ret()
 95					return
 96				end
 97			end
 98			READSESSION[fd] = session
 99		end
100	end
101}
102
103local socket = {}
104
105function socket.open(addr, port)
106	local cmd = "open" .. " " .. (port and (addr..":"..port) or addr)
107	local r = skynet.call(".socket", "text", cmd)
108	if r == "" then
109		error(cmd .. " failed")
110	end
111	return tonumber(r)
112end
113
114function socket.stdin()
115	local r = skynet.call(".socket", "text", "bind 1")
116	if r == "" then
117		error("stdin bind failed")
118	end
119	return tonumber(r)
120end
121
122function socket.close(fd)
123	socket.lock(fd)
124	skynet.call(".socket", "text", "close", fd)
125	READBUF[fd] = true
126	READLOCK[fd] = nil
127end
128
129function socket.read(fd, sz)
130	assert(coroutine.running() == READTHREAD[fd], "call socket.lock first")
131
132	local str = buffer.pop(READBUF[fd],sz)
133	if str then
134		return str
135	end
136
137	READREQUEST[fd] = sz
138	skynet.call(selfaddr, "system",fd,1,sz)	-- singal size 1
139	READREQUEST[fd] = nil
140
141	str = buffer.pop(READBUF[fd],sz)
142	return str
143end
144
145function socket.readline(fd, sep)
146	assert(coroutine.running() == READTHREAD[fd], "call socket.lock first")
147
148	local str = buffer.readline(READBUF[fd],sep)
149	if str then
150		return str
151	end
152
153	READREQUEST[fd] = sep
154	skynet.call(selfaddr, "system",fd,2,sep)	-- singal sep 2
155	READREQUEST[fd] = nil
156
157	str = buffer.readline(READBUF[fd],sep)
158	return str
159end
160
161function socket.write(fd, msg, sz)
162	skynet.send(".socket", "client", fd, msg, sz)
163end
164
165function socket.lock(fd)
166	local lt = READTHREAD[fd]
167	local ct = coroutine.running()
168	if lt then
169		assert(lt ~= ct, "already lock")
170		skynet.call(selfaddr, "system",fd,0)
171	end
172	READTHREAD[fd] = ct
173end
174
175function socket.unlock(fd)
176	READTHREAD[fd] = nil
177	local q = READLOCK[fd]
178	if q then
179		if q[1] then
180			response(q[1])
181		end
182		table.remove(q,1)
183		if q[1] == nil then
184			READLOCK[fd] = nil
185		end
186	end
187end
188
189return socket