PageRenderTime 3ms CodeModel.GetById 16ms app.highlight 10ms RepoModel.GetById 1ms app.codeStats 0ms

/lualib/socket.lua

https://github.com/yunnet/skynet
Lua | 269 lines | 236 code | 28 blank | 5 comment | 43 complexity | 1bf6945edcde8f7b80dc1c5d075a3370 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
 14local CLOSED = {} -- fd:true
 15
 16local selfaddr = skynet.self()
 17local sockets = assert(skynet.localname ".socket")
 18
 19local function response(session)
 20	skynet.redirect(selfaddr , 0, "response", session, "")
 21end
 22
 23skynet.register_protocol {
 24	name = "client",
 25	id = 3,	-- PTYPE_CLIENT
 26	pack = buffer.pack,
 27	unpack = buffer.unpack,
 28	dispatch = function (_, _, fd, msg, sz)
 29		local qsz = READREQUEST[fd]
 30		local buf = READBUF[fd]
 31		local bsz
 32		if sz == 0 then
 33			CLOSED[fd] = true
 34		else
 35			buf,bsz = buffer.push(buf, msg, sz)
 36			READBUF[fd] = buf
 37		end
 38		local session = READSESSION[fd]
 39		if qsz == nil or session == nil then
 40			return
 41		end
 42		if sz > 0 then
 43			if type(qsz) == "number" then
 44				if qsz > bsz then
 45					return
 46				end
 47			else
 48				-- request readline
 49				if buffer.readline(buf, qsz, true) == nil then
 50					return
 51				end
 52			end
 53		end
 54
 55		response(session)
 56		READSESSION[fd] = nil
 57	end
 58}
 59
 60skynet.register_protocol {
 61	name = "system",
 62	id = 4, -- PTYPE_SYSTEM
 63	pack = skynet.pack,
 64	unpack = function (...) return ... end,
 65	dispatch = function (session, addr, msg, sz)
 66		fd, t, sz = skynet.unpack(msg,sz)
 67		assert(addr == selfaddr, "PTYPE_SYSTEM message must send by self")
 68		if t > 0 then	-- lock request when t == 0
 69			-- request bytes or readline
 70			local buf = READBUF[fd]
 71			if CLOSED[fd] then
 72				skynet.ret()
 73				return
 74			end
 75			local _,bsz = buffer.push(buf)
 76			if t == 1 then
 77				-- sz is size
 78				if bsz >= sz then
 79					skynet.ret()
 80					return
 81				end
 82			elseif t == 2 then
 83				-- sz is sep
 84				if buffer.readline(buf, sz, true) then -- don't real read
 85					skynet.ret()
 86					return
 87				end
 88			end
 89			READSESSION[fd] = session
 90		end
 91	end
 92}
 93
 94local socket = {}
 95
 96function socket.open(addr, port)
 97	local cmd = "open" .. " " .. (port and (addr..":"..port) or addr)
 98	local r = skynet.call(sockets, "text", cmd)
 99	if r == "" then
100		return nil,  cmd .. " failed"
101	end
102	local fd = tonumber(r)
103	READBUF[fd] = true
104	CLOSED[fd] = nil
105	return fd
106end
107
108function socket.stdin()
109	local r = skynet.call(sockets, "text", "bind 1")
110	if r == "" then
111		error("stdin bind failed")
112	end
113	local fd = tonumber(r)
114	READBUF[fd] = true
115	CLOSED[fd] = nil
116	return fd
117end
118
119function socket.close(fd)
120	socket.lock(fd)
121	skynet.call(sockets, "text", "close", fd)
122	READBUF[fd] = nil
123	READLOCK[fd] = nil
124	CLOSED[fd] = nil
125end
126
127function socket.read(fd, sz)
128	local buf = assert(READBUF[fd])
129	local str, bytes = buffer.pop(buf,sz)
130	if str then
131		return str
132	end
133
134	if CLOSED[fd] then
135		READBUF[fd] = nil
136		CLOSED[fd] = nil
137		str = buffer.pop(buf, bytes)
138		return nil, str
139	end
140
141	READREQUEST[fd] = sz
142	skynet.call(selfaddr, "system",fd,1,sz)	-- singal size 1
143	READREQUEST[fd] = nil
144
145	buf = READBUF[fd]
146
147	str, bytes = buffer.pop(buf,sz)
148	if str then
149		return str
150	end
151
152	if CLOSED[fd] then
153		READBUF[fd] = nil
154		CLOSED[fd] = nil
155		str = buffer.pop(buf, bytes)
156		return nil, str
157	end
158end
159
160function socket.readall(fd)
161	local buf = assert(READBUF[fd])
162	if CLOSED[fd] then
163		CLOSED[fd] = nil
164		READBUF[fd] = nil
165		if buf == nil then
166			return ""
167		end
168		local _, bytes = buffer.push(buf)
169		local ret = buffer.pop(buf, bytes)
170		return ret
171	end
172	READREQUEST[fd] = math.huge
173	skynet.call(selfaddr, "system",fd,3)	-- singal readall
174	READREQUEST[fd] = nil
175	assert(CLOSED[fd])
176	buf = READBUF[fd]
177	READBUF[fd] = nil
178	CLOSED[fd] = nil
179	if buf == nil then
180		return ""
181	end
182	local _, bytes = buffer.push(buf)
183	local ret = buffer.pop(buf, bytes)
184	return ret
185end
186
187function socket.readline(fd, sep)
188	local buf = assert(READBUF[fd])
189	local str = buffer.readline(buf,sep)
190	if str then
191		return str
192	end
193
194	if CLOSED[fd] then
195		READBUF[fd] = nil
196		CLOSED[fd] = nil
197		local _, bytes = buffer.push(buf)
198		str = buffer.pop(buf, bytes)
199		return nil, str
200	end
201
202	READREQUEST[fd] = sep
203	skynet.call(selfaddr, "system",fd,2,sep)	-- singal sep 2
204	READREQUEST[fd] = nil
205
206	buf = READBUF[fd]
207	str = buffer.readline(buf,sep)
208	if str then
209		return str
210	end
211
212	if CLOSED[fd] then
213		READBUF[fd] = nil
214		CLOSED[fd] = nil
215		local _, bytes = buffer.push(buf)
216		str = buffer.pop(buf, bytes)
217		return nil, str
218	end
219end
220
221function socket.write(fd, msg, sz)
222	if CLOSED[fd] or not READBUF[fd] then
223		return
224	end
225	skynet.send(sockets, "client", fd, msg, sz)
226	return true
227end
228
229function socket.invalid(fd)
230	return CLOSED[fd] or not READBUF[fd]
231end
232
233function socket.lock(fd)
234	if CLOSED[fd] or not READBUF[fd] then
235		return
236	end
237	local locked = READTHREAD[fd]
238	if locked then
239		-- lock fd
240		local session = skynet.genid()
241		local q = READLOCK[fd]
242		if q == nil then
243			READLOCK[fd] = { session }
244		else
245			table.insert(q, session)
246		end
247
248		skynet.redirect(selfaddr , 0, "system", session, skynet.pack(fd,0))
249		coroutine.yield("CALL",session)
250	else
251		READTHREAD[fd] = true
252	end
253end
254
255function socket.unlock(fd)
256	READTHREAD[fd] = nil
257	local q = READLOCK[fd]
258	if q then
259		if q[1] then
260			READTHREAD[fd] = true
261			response(q[1])
262			table.remove(q,1)
263		else
264			READLOCK[fd] = nil
265		end
266	end
267end
268
269return socket