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