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