Extended example of Lua coroutines

-- ***************************************************************
--
-- Copyright 2013 by Sean Conner.  All Rights Reserved.
--
-- This library is free software; you can redistribute it and/or modify it
-- under the terms of the GNU Lesser General Public License as published by
-- the Free Software Foundation; either version 3 of the License, or (at your
-- option) any later version.
--
-- This library is distributed in the hope that it will be useful, but
-- WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
-- or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
-- License for more details.
--
-- You should have received a copy of the GNU Lesser General Public License
-- along with this library; if not, see <http://www.gnu.org/licenses/>.
--
-- Comments, questions and criticisms can be sent to: [email protected]
--
-- ********************************************************************

local proc    = require "org.conman.process"
local net     = require "org.conman.net"
local pollset = require "org.conman.pollset"
local lpeg    = require "lpeg"

-- **********************************************************************

FILE       = arg[1] or "/home/spc/source/lua/daemon/lua/silly.lua"
setqueue   = pollset()
runqueue   = {}
sleepqueue = {}

-- **********************************************************************

function wsleep(sock,amount)
  sock.awake = os.time() + amount
  table.insert(sleepqueue,sock)
  coroutine.yield('sleep')
end

-- **********************************************************************

local eoln      = lpeg.P"\r"^-1 * lpeg.P"\n"
local lineparse = lpeg.C((lpeg.P(1) - eoln)^0) * eoln * lpeg.C(lpeg.P(1)^0)
local integer   = lpeg.R"09"^1
local frac      = lpeg.P"."  * lpeg.R"09"^1
local exp       = lpeg.S"Ee" * lpeg.S"-+"^-1 * lpeg.R"09"^1
local number    = ((lpeg.P"-"^-1 * integer * frac^-1 * exp^-1) / tonumber)
                * lpeg.locale().space^1 * lpeg.C(lpeg.P(1)^0)

function wread(sock,amount)
  local amount = amount or "*l"

  if type(amount) == 'number' then
    if #sock.readbuf < amount then
      coroutine.yield('read')
      return wread(sock,amount)
    else
      local d      = sock.readbuf:sub(1,amount)
      sock.readbuf = sock.readbuf:sub(amount+1,-1)
      return d
    end

  elseif amount == '*l' then
    local line,rest = lineparse:match(sock.readbuf)
    if line == nil then
      coroutine.yield('read')
      return wread(sock,amount)
    else
      sock.readbuf = rest
      return line
    end

  elseif amount == '*a' then
    if #sock.readbuf == 0 then
      coroutine.yield('read')
      return wread(sock,amount)
    else
      local d = sock.readbuf
      sock.readbuf = ""
      return d
    end
  elseif amount == '*n' then
    local num,rest = number:match(sock.readbuf)
    if num == nil then
      coroutine.yield('read')
      return read(sock,amount)
    else
      sock.readbuf = rest
      return num
    end
  end
end

-- **********************************************************************

function wwrite(sock,data)
  if #sock.writebuf == 0 then
    local bytes = sock.rawsock:send(sock.raddr,data)
    data  = data:sub(bytes + 1,-1)
  end
  sock.writebuf = sock.writebuf .. data
  if #sock.writebuf > 0 then
    setqueue:update(sock.rawsock:fd(),"rw")
  end
end

-- **********************************************************************

function check_signals()
  if proc.sig.caught(proc.sig.INT)
  or proc.sig.caught(proc.sig.QUIT)
  or proc.sig.caught(proc.sig.TERM) then
    os.exit(0)
  elseif proc.sig.caught(proc.sig.USR1) then
    dofile(FILE)
  end
end

-- **********************************************************************

function mainloop()
  local timeout
  local events
  local now
  local socket
  local okay
  local newsleep

  while true do
    check_signals()

    if #runqueue > 0 then
      timeout = 0
    elseif #sleepqueue > 0 then
      timeout = 1
    else
      timeout = -1
    end

    events,err = setqueue:events(timeout)
    if err ~= 0 then
      return mainloop()
    end

    for i = 1 , #events do
      events[i].obj(events[i])
    end

    newsleep = {}
    now      = os.time()

    while #sleepqueue > 0 do
      obj = table.remove(sleepqueue,1)
      if obj.awake > now then
        table.insert(newsleep,obj)
      else
        table.insert(runqueue,obj)
      end
    end

    sleepqueue = newsleep

    socket = table.remove(runqueue,1)
    if socket then
      okay,socket.state = coroutine.resume(socket.f,socket)
      if coroutine.status(socket.f) == 'dead' then
        setqueue:remove(socket.rawsock:fd())
        socket.rawsock:close()
      end
    end
  end
end

-- **********************************************************************

proc.sig.catch(proc.sig.INT)
proc.sig.catch(proc.sig.QUIT)
proc.sig.catch(proc.sig.TERM)
proc.sig.catch(proc.sig.USR1)

laddr = net.address('192.168.1.10','tcp',22222)
sock  = net.socket('ip','tcp')

sock.reuseaddr = true
sock.nonblock  = true
sock:bind(laddr)
sock:listen()

dofile(FILE)

setqueue:insert(sock:fd(),"r",function(event)
  local conn,rem = sock:accept()
  conn.nonblock  = true

  setqueue:insert(conn:fd(),"r",(function(conn,raddr)
    local socket =
    {
      read     = wread,
      write    = wwrite,
      sleep    = wsleep,
      readbuf  = "",
      writebuf = "",
      state    = 'run',
      rawsock  = conn,
      raddr    = raddr,
      f        = coroutine.create(main)
    }

    table.insert(runqueue,socket)

    return function(event)
      if event.read then
        local rem,data = conn:recv()

        if #data == 0 then
          setqueue:remove(conn:fd())
          conn:close()
          if fini then fini(socket) end
       end

        socket.readbuf = socket.readbuf .. data
        if socket.state == 'read' then
          socket.state = 'run'
          table.insert(runqueue,socket)
        end
      end

      if event.write then
        if #socket.writebuf > 0 then
          local bytes = conn:send(raddr,socket.writebuf)
          socket.writebuf = socket.writebuf:sub(bytes+1,-1)
        end

        if #socket.writebuf == 0 then
          setqueue:update(conn:fd(),"r")
        end
      end
    end
  end)(conn,rem))
end)

mainloop()