diff options
-rw-r--r-- | core/moduleapi.lua | 32 | ||||
-rw-r--r-- | net/http.lua | 17 | ||||
-rw-r--r-- | net/server_event.lua | 76 | ||||
-rw-r--r-- | net/server_select.lua | 55 | ||||
-rw-r--r-- | plugins/mod_pep_plus.lua | 368 | ||||
-rw-r--r-- | plugins/mod_pubsub/mod_pubsub.lua | 2 | ||||
-rw-r--r-- | plugins/mod_pubsub/pubsub.lib.lua | 4 | ||||
-rw-r--r-- | util/indexedbheap.lua | 153 | ||||
-rw-r--r-- | util/timer.lua | 62 |
9 files changed, 692 insertions, 77 deletions
diff --git a/core/moduleapi.lua b/core/moduleapi.lua index 65e00d41..5a24f69c 100644 --- a/core/moduleapi.lua +++ b/core/moduleapi.lua @@ -16,8 +16,10 @@ local timer = require "util.timer"; local t_insert, t_remove, t_concat = table.insert, table.remove, table.concat; local error, setmetatable, type = error, setmetatable, type; -local ipairs, pairs, select, unpack = ipairs, pairs, select, unpack; +local ipairs, pairs, select = ipairs, pairs, select; local tonumber, tostring = tonumber, tostring; +local pack = table.pack or function(...) return {n=select("#",...), ...}; end -- table.pack is only in 5.2 +local unpack = table.unpack or unpack; -- renamed in 5.2 local prosody = prosody; local hosts = prosody.hosts; @@ -347,11 +349,29 @@ function api:send(stanza) return core_post_stanza(hosts[self.host], stanza); end -function api:add_timer(delay, callback) - return timer.add_task(delay, function (t) - if self.loaded == false then return; end - return callback(t); - end); +local timer_methods = { } +local timer_mt = { + __index = timer_methods; +} +function timer_methods:stop( ) + timer.stop(self.id); +end +timer_methods.disarm = timer_methods.stop +function timer_methods:reschedule(delay) + timer.reschedule(self.id, delay) +end + +local function timer_callback(now, id, t) + if t.module_env.loaded == false then return; end + return t.callback(now, unpack(t, 1, t.n)); +end + +function api:add_timer(delay, callback, ...) + local t = pack(...) + t.module_env = self; + t.callback = callback; + t.id = timer.add_task(delay, timer_callback, t); + return setmetatable(t, timer_mt); end local path_sep = package.config:sub(1,1); diff --git a/net/http.lua b/net/http.lua index ab9ec7b6..b87c9396 100644 --- a/net/http.lua +++ b/net/http.lua @@ -6,7 +6,6 @@ -- COPYING file in the source package for more information. -- -local socket = require "socket" local b64 = require "util.encodings".base64.encode; local url = require "socket.url" local httpstream_new = require "net.http.parser".new; @@ -160,21 +159,17 @@ function request(u, ex, callback) end local port_number = port and tonumber(port) or (using_https and 443 or 80); - -- Connect the socket, and wrap it with net.server - local conn = socket.tcp(); - conn:settimeout(10); - local ok, err = conn:connect(host, port_number); - if not ok and err ~= "timeout" then - callback(nil, 0, req); - return nil, err; - end - local sslctx = false; if using_https then sslctx = ex and ex.sslctx or { mode = "client", protocol = "sslv23", options = { "no_sslv2" } }; end - req.handler, req.conn = assert(server.wrapclient(conn, host, port_number, listener, "*a", sslctx)); + local handler, conn = server.addclient(host, port_number, listener, "*a", sslctx) + if not handler then + callback(nil, 0, req); + return nil, conn; + end + req.handler, req.conn = handler, conn req.write = function (...) return req.handler:write(...); end req.callback = function (content, code, request, response) log("debug", "Calling callback, status %s", code or "---"); return select(2, xpcall(function () return callback(content, code, request, response) end, handleerr)); end diff --git a/net/server_event.lua b/net/server_event.lua index 59217a0c..ef0a27d8 100644 --- a/net/server_event.lua +++ b/net/server_event.lua @@ -44,8 +44,9 @@ local setmetatable = use "setmetatable" local t_insert = table.insert local t_concat = table.concat -local ssl = use "ssl" +local has_luasec, ssl = pcall ( require , "ssl" ) local socket = use "socket" or require "socket" +local getaddrinfo = socket.dns.getaddrinfo local log = require ("util.logger").init("socket") @@ -128,7 +129,7 @@ do return self:_destroy(); end - function interface_mt:_start_connection(plainssl) -- should be called from addclient + function interface_mt:_start_connection(plainssl) -- called from wrapclient local callback = function( event ) if EV_TIMEOUT == event then -- timeout during connection self.fatalerror = "connection timeout" @@ -136,7 +137,7 @@ do self:_close() debug( "new connection failed. id:", self.id, "error:", self.fatalerror ) else - if plainssl and ssl then -- start ssl session + if plainssl and has_luasec then -- start ssl session self:starttls(self._sslctx, true) else -- normal connection self:_start_session(true) @@ -506,7 +507,7 @@ do _sslctx = sslctx; -- parameters _usingssl = false; -- client is using ssl; } - if not ssl then interface.starttls = false; end + if not has_luasec then interface.starttls = false; end interface.id = tostring(interface):match("%x+$"); interface.writecallback = function( event ) -- called on write events --vdebug( "new client write event, id/ip/port:", interface, ip, port ) @@ -689,7 +690,7 @@ do interface._connections = interface._connections + 1 -- increase connection count local clientinterface = handleclient( client, client_ip, client_port, interface, pattern, listener, sslctx ) --vdebug( "client id:", clientinterface, "startssl:", startssl ) - if ssl and sslctx then + if has_luasec and sslctx then clientinterface:starttls(sslctx, true) else clientinterface:_start_session( true ) @@ -710,25 +711,17 @@ do end local addserver = ( function( ) - return function( addr, port, listener, pattern, sslcfg, startssl ) -- TODO: check arguments - --vdebug( "creating new tcp server with following parameters:", addr or "nil", port or "nil", sslcfg or "nil", startssl or "nil") + return function( addr, port, listener, pattern, sslctx, startssl ) -- TODO: check arguments + --vdebug( "creating new tcp server with following parameters:", addr or "nil", port or "nil", sslctx or "nil", startssl or "nil") + if sslctx and not has_luasec then + debug "fatal error: luasec not found" + return nil, "luasec not found" + end local server, err = socket.bind( addr, port, cfg.ACCEPT_QUEUE ) -- create server socket if not server then debug( "creating server socket on "..addr.." port "..port.." failed:", err ) return nil, err end - local sslctx - if sslcfg then - if not ssl then - debug "fatal error: luasec not found" - return nil, "luasec not found" - end - sslctx, err = sslcfg - if err then - debug( "error while creating new ssl context for server socket:", err ) - return nil, err - end - end local interface = handleserver( server, addr, port, pattern, listener, sslctx, startssl ) -- new server handler debug( "new server created with id:", tostring(interface)) return interface @@ -744,37 +737,34 @@ do --function handleclient( client, ip, port, server, pattern, listener, _, sslctx ) -- creates an client interface end - function addclient( addr, serverport, listener, pattern, localaddr, localport, sslcfg, startssl ) - local client, err = socket.tcp() -- creating new socket + function addclient( addr, serverport, listener, pattern, sslctx, typ ) + if sslctx and not has_luasec then + debug "need luasec, but not available" + return nil, "luasec not found" + end + if getaddrinfo and not typ then + local addrinfo, err = getaddrinfo(addr) + if not addrinfo then return nil, err end + if addrinfo[1] and addrinfo[1].family == "inet6" then + typ = "tcp6" + end + end + local create = socket[typ or "tcp"] + if type( create ) ~= "function" then + return nil, "invalid socket type" + end + local client, err = create() -- creating new socket if not client then debug( "cannot create socket:", err ) return nil, err end client:settimeout( 0 ) -- set nonblocking - if localaddr then - local res, err = client:bind( localaddr, localport, -1 ) - if not res then - debug( "cannot bind client:", err ) - return nil, err - end - end - local sslctx - if sslcfg then -- handle ssl/new context - if not ssl then - debug "need luasec, but not available" - return nil, "luasec not found" - end - sslctx, err = sslcfg - if err then - debug( "cannot create new ssl context:", err ) - return nil, err - end - end local res, err = client:connect( addr, serverport ) -- connect if res or ( err == "timeout" ) then - local ip, port = client:getsockname( ) - local interface = wrapclient( client, ip, serverport, listener, pattern, sslctx, startssl ) - interface:_start_connection( startssl ) + if client.getsockname then + addr = client:getsockname( ) + end + local interface = wrapclient( client, addr, serverport, listener, pattern, sslctx ) debug( "new connection id:", interface.id ) return interface, err else diff --git a/net/server_select.lua b/net/server_select.lua index c5e0772f..b69b5fc7 100644 --- a/net/server_select.lua +++ b/net/server_select.lua @@ -48,13 +48,14 @@ local coroutine_yield = coroutine.yield --// extern libs //-- -local luasec = use "ssl" +local has_luasec, luasec = pcall ( require , "ssl" ) local luasocket = use "socket" or require "socket" local luasocket_gettime = luasocket.gettime +local getaddrinfo = luasocket.dns.getaddrinfo --// extern lib methods //-- -local ssl_wrap = ( luasec and luasec.wrap ) +local ssl_wrap = ( has_luasec and luasec.wrap ) local socket_bind = luasocket.bind local socket_sleep = luasocket.sleep local socket_select = luasocket.select @@ -585,7 +586,7 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport end ) end - if luasec then + if has_luasec then handler.starttls = function( self, _sslctx) if _sslctx then handler:set_sslctx(_sslctx); @@ -638,7 +639,7 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport _socketlist[ socket ] = handler _readlistlen = addsocket(_readlist, socket, _readlistlen) - if sslctx and luasec then + if sslctx and has_luasec then out_put "server.lua: auto-starting ssl negotiation..." handler.autostart_ssl = true; local ok, err = handler:starttls(sslctx); @@ -713,22 +714,23 @@ end ----------------------------------// PUBLIC //-- addserver = function( addr, port, listeners, pattern, sslctx ) -- this function provides a way for other scripts to reg a server + addr = addr or "*" local err if type( listeners ) ~= "table" then err = "invalid listener table" - end - if type( port ) ~= "number" or not ( port >= 0 and port <= 65535 ) then + elseif type ( addr ) ~= "string" then + err = "invalid address" + elseif type( port ) ~= "number" or not ( port >= 0 and port <= 65535 ) then err = "invalid port" elseif _server[ addr..":"..port ] then err = "listeners on '[" .. addr .. "]:" .. port .. "' already exist" - elseif sslctx and not luasec then + elseif sslctx and not has_luasec then err = "luasec not found" end if err then out_error( "server.lua, [", addr, "]:", port, ": ", err ) return nil, err end - addr = addr or "*" local server, err = socket_bind( addr, port, _tcpbacklog ) if err then out_error( "server.lua, [", addr, "]:", port, ": ", err ) @@ -929,17 +931,44 @@ local wrapclient = function( socket, ip, serverport, listeners, pattern, sslctx return handler, socket end -local addclient = function( address, port, listeners, pattern, sslctx ) - local client, err = luasocket.tcp( ) +local addclient = function( address, port, listeners, pattern, sslctx, typ ) + local err + if type( listeners ) ~= "table" then + err = "invalid listener table" + elseif type ( address ) ~= "string" then + err = "invalid address" + elseif type( port ) ~= "number" or not ( port >= 0 and port <= 65535 ) then + err = "invalid port" + elseif sslctx and not has_luasec then + err = "luasec not found" + end + if getaddrinfo and not typ then + local addrinfo, err = getaddrinfo(address) + if not addrinfo then return nil, err end + if addrinfo[1] and addrinfo[1].family == "inet6" then + typ = "tcp6" + end + end + local create = luasocket[typ or "tcp"] + if type( create ) ~= "function" then + err = "invalid socket type" + end + + if err then + out_error( "server.lua, addclient: ", err ) + return nil, err + end + + local client, err = create( ) if err then return nil, err end client:settimeout( 0 ) - _, err = client:connect( address, port ) - if err then -- try again + local ok, err = client:connect( address, port ) + if ok or err == "timeout" then return wrapclient( client, address, port, listeners, pattern, sslctx ) else - return wrapconnection( nil, listeners, client, address, port, "clientport", pattern, sslctx ) + return nil, err end end diff --git a/plugins/mod_pep_plus.lua b/plugins/mod_pep_plus.lua new file mode 100644 index 00000000..4a74e437 --- /dev/null +++ b/plugins/mod_pep_plus.lua @@ -0,0 +1,368 @@ +local pubsub = require "util.pubsub"; +local jid_bare = require "util.jid".bare; +local jid_split = require "util.jid".split; +local set_new = require "util.set".new; +local st = require "util.stanza"; +local calculate_hash = require "util.caps".calculate_hash; +local is_contact_subscribed = require "core.rostermanager".is_contact_subscribed; + +local xmlns_pubsub = "http://jabber.org/protocol/pubsub"; +local xmlns_pubsub_event = "http://jabber.org/protocol/pubsub#event"; +local xmlns_pubsub_owner = "http://jabber.org/protocol/pubsub#owner"; + +local lib_pubsub = module:require "pubsub"; +local handlers = lib_pubsub.handlers; +local pubsub_error_reply = lib_pubsub.pubsub_error_reply; + +local services = {}; +local recipients = {}; +local hash_map = {}; + +function module.save() + return { services = services }; +end + +function module.restore(data) + services = data.services; +end + +local function subscription_presence(user_bare, recipient) + local recipient_bare = jid_bare(recipient); + if (recipient_bare == user_bare) then return true; end + local username, host = jid_split(user_bare); + return is_contact_subscribed(username, host, recipient_bare); +end + +local function get_broadcaster(name) + local function simple_broadcast(kind, node, jids, item) + if item then + item = st.clone(item); + item.attr.xmlns = nil; -- Clear the pubsub namespace + end + local message = st.message({ from = name, type = "headline" }) + :tag("event", { xmlns = xmlns_pubsub_event }) + :tag(kind, { node = node }) + :add_child(item); + for jid in pairs(jids) do + module:log("debug", "Sending notification to %s from %s: %s", jid, name, tostring(item)); + message.attr.to = jid; + module:send(message); + end + end + return simple_broadcast; +end + +local function get_pep_service(name) + if services[name] then + return services[name]; + end + services[name] = pubsub.new({ + capabilities = { + none = { + create = false; + publish = false; + retract = false; + get_nodes = false; + + subscribe = false; + unsubscribe = false; + get_subscription = false; + get_subscriptions = false; + get_items = false; + + subscribe_other = false; + unsubscribe_other = false; + get_subscription_other = false; + get_subscriptions_other = false; + + be_subscribed = true; + be_unsubscribed = true; + + set_affiliation = false; + }; + subscriber = { + create = false; + publish = false; + retract = false; + get_nodes = true; + + subscribe = true; + unsubscribe = true; + get_subscription = true; + get_subscriptions = true; + get_items = true; + + subscribe_other = false; + unsubscribe_other = false; + get_subscription_other = false; + get_subscriptions_other = false; + + be_subscribed = true; + be_unsubscribed = true; + + set_affiliation = false; + }; + publisher = { + create = false; + publish = true; + retract = true; + get_nodes = true; + + subscribe = true; + unsubscribe = true; + get_subscription = true; + get_subscriptions = true; + get_items = true; + + subscribe_other = false; + unsubscribe_other = false; + get_subscription_other = false; + get_subscriptions_other = false; + + be_subscribed = true; + be_unsubscribed = true; + + set_affiliation = false; + }; + owner = { + create = true; + publish = true; + retract = true; + delete = true; + get_nodes = true; + + subscribe = true; + unsubscribe = true; + get_subscription = true; + get_subscriptions = true; + get_items = true; + + + subscribe_other = true; + unsubscribe_other = true; + get_subscription_other = true; + get_subscriptions_other = true; + + be_subscribed = true; + be_unsubscribed = true; + + set_affiliation = true; + }; + }; + + autocreate_on_publish = true; + autocreate_on_subscribe = true; + + broadcaster = get_broadcaster(name); + get_affiliation = function (jid) + if jid_bare(jid) == name then + return "owner"; + elseif subscription_presence(name, jid) then + return "subscriber"; + end + end; + + normalize_jid = jid_bare; + }); + return services[name]; +end + +function handle_pubsub_iq(event) + local origin, stanza = event.origin, event.stanza; + local pubsub = stanza.tags[1]; + local action = pubsub.tags[1]; + if not action then + return origin.send(st.error_reply(stanza, "cancel", "bad-request")); + end + local service_name = stanza.attr.to or origin.username.."@"..origin.host + local service = get_pep_service(service_name); + local handler = handlers[stanza.attr.type.."_"..action.name]; + if handler then + handler(origin, stanza, action, service); + return true; + end +end + +module:hook("iq/bare/"..xmlns_pubsub..":pubsub", handle_pubsub_iq); +module:hook("iq/bare/"..xmlns_pubsub_owner..":pubsub", handle_pubsub_iq); + +module:add_identity("pubsub", "pep", module:get_option_string("name", "Prosody")); +module:add_feature("http://jabber.org/protocol/pubsub#publish"); + +local function get_caps_hash_from_presence(stanza, current) + local t = stanza.attr.type; + if not t then + local child = stanza:get_child("c", "http://jabber.org/protocol/caps"); + if child then + local attr = child.attr; + if attr.hash then -- new caps + if attr.hash == 'sha-1' and attr.node and attr.ver then + return attr.ver, attr.node.."#"..attr.ver; + end + else -- legacy caps + if attr.node and attr.ver then + return attr.node.."#"..attr.ver.."#"..(attr.ext or ""), attr.node.."#"..attr.ver; + end + end + end + return; -- no or bad caps + elseif t == "unavailable" or t == "error" then + return; + end + return current; -- no caps, could mean caps optimization, so return current +end + +local function resend_last_item(jid, node, service) + local ok, items = service:get_items(node, jid); + if not ok then return; end + for i, id in ipairs(items) do + service.config.broadcaster("items", node, { [jid] = true }, items[id]); + end +end + +local function update_subscriptions(recipient, service_name, nodes) + local service = get_pep_service(service_name); + + recipients[service_name] = recipients[service_name] or {}; + nodes = nodes or set_new(); + local old = recipients[service_name][recipient]; + + if old and type(old) == table then + for node in pairs((old - nodes):items()) do + service:remove_subscription(node, recipient, recipient); + end + end + + for node in nodes:items() do + service:add_subscription(node, recipient, recipient); + resend_last_item(recipient, node, service); + end + recipients[service_name][recipient] = nodes; +end + +module:hook("presence/bare", function(event) + -- inbound presence to bare JID recieved + local origin, stanza = event.origin, event.stanza; + local user = stanza.attr.to or (origin.username..'@'..origin.host); + local t = stanza.attr.type; + local self = not stanza.attr.to; + local service = get_pep_service(user); + + if not t then -- available presence + if self or subscription_presence(user, stanza.attr.from) then + local recipient = stanza.attr.from; + local current = recipients[user] and recipients[user][recipient]; + local hash, query_node = get_caps_hash_from_presence(stanza, current); + if current == hash or (current and current == hash_map[hash]) then return; end + if not hash then + update_subscriptions(recipient, user); + else + recipients[user] = recipients[user] or {}; + if hash_map[hash] then + update_subscriptions(recipient, user, hash_map[hash]); + else + recipients[user][recipient] = hash; + local from_bare = origin.type == "c2s" and origin.username.."@"..origin.host; + if self or origin.type ~= "c2s" or (recipients[from_bare] and recipients[from_bare][origin.full_jid]) ~= hash then + -- COMPAT from ~= stanza.attr.to because OneTeam can't deal with missing from attribute + origin.send( + st.stanza("iq", {from=user, to=stanza.attr.from, id="disco", type="get"}) + :tag("query", {xmlns = "http://jabber.org/protocol/disco#info", node = query_node}) + ); + end + end + end + end + elseif t == "unavailable" then + update_subscriptions(stanza.attr.from, user); + elseif not self and t == "unsubscribe" then + local from = jid_bare(stanza.attr.from); + local subscriptions = recipients[user]; + if subscriptions then + for subscriber in pairs(subscriptions) do + if jid_bare(subscriber) == from then + update_subscriptions(subscriber, user); + end + end + end + end +end, 10); + +module:hook("iq-result/bare/disco", function(event) + local origin, stanza = event.origin, event.stanza; + local disco = stanza:get_child("query", "http://jabber.org/protocol/disco#info"); + if not disco then + return; + end + + -- Process disco response + local self = not stanza.attr.to; + local user = stanza.attr.to or (origin.username..'@'..origin.host); + local contact = stanza.attr.from; + local current = recipients[user] and recipients[user][contact]; + if type(current) ~= "string" then return; end -- check if waiting for recipient's response + local ver = current; + if not string.find(current, "#") then + ver = calculate_hash(disco.tags); -- calculate hash + end + local notify = set_new(); + for _, feature in pairs(disco.tags) do + if feature.name == "feature" and feature.attr.var then + local nfeature = feature.attr.var:match("^(.*)%+notify$"); + if nfeature then notify:add(nfeature); end + end + end + hash_map[ver] = notify; -- update hash map + if self then + for jid, item in pairs(origin.roster) do -- for all interested contacts + if item.subscription == "both" or item.subscription == "from" then + if not recipients[jid] then recipients[jid] = {}; end + update_subscriptions(contact, jid, notify); + end + end + end + update_subscriptions(contact, user, notify); +end); + +module:hook("account-disco-info-node", function(event) + local reply, stanza, origin = event.reply, event.stanza, event.origin; + local service_name = stanza.attr.to or origin.username.."@"..origin.host + local service = get_pep_service(service_name); + local node = event.node; + local ok = service:get_items(node, jid_bare(stanza.attr.from) or true); + if not ok then return; end + event.exists = true; + reply:tag('identity', {category='pubsub', type='leaf'}):up(); +end); + +module:hook("account-disco-info", function(event) + local reply = event.reply; + reply:tag('identity', {category='pubsub', type='pep'}):up(); + reply:tag('feature', {var='http://jabber.org/protocol/pubsub#publish'}):up(); +end); + +module:hook("account-disco-items-node", function(event) + local reply, stanza, origin = event.reply, event.stanza, event.origin; + local node = event.node; + local service_name = stanza.attr.to or origin.username.."@"..origin.host + local service = get_pep_service(service_name); + local ok, ret = service:get_items(node, jid_bare(stanza.attr.from) or true); + if not ok then return; end + event.exists = true; + for _, id in ipairs(ret) do + reply:tag("item", { jid = service_name, name = id }):up(); + end +end); + +module:hook("account-disco-items", function(event) + local reply, stanza, origin = event.reply, event.stanza, event.origin; + + local service_name = reply.attr.from or origin.username.."@"..origin.host + local service = get_pep_service(service_name); + local ok, ret = service:get_nodes(jid_bare(stanza.attr.from)); + if not ok then return; end + + for node, node_obj in pairs(ret) do + reply:tag("item", { jid = service_name, node = node, name = node_obj.config.name }):up(); + end +end); diff --git a/plugins/mod_pubsub/mod_pubsub.lua b/plugins/mod_pubsub/mod_pubsub.lua index c6dbe831..33e729af 100644 --- a/plugins/mod_pubsub/mod_pubsub.lua +++ b/plugins/mod_pubsub/mod_pubsub.lua @@ -100,7 +100,7 @@ module:hook("host-disco-items-node", function (event) return; end - for id, item in pairs(ret) do + for _, id in ipairs(ret) do reply:tag("item", { jid = module.host, name = id }):up(); end event.exists = true; diff --git a/plugins/mod_pubsub/pubsub.lib.lua b/plugins/mod_pubsub/pubsub.lib.lua index 2b015e34..4e9acd68 100644 --- a/plugins/mod_pubsub/pubsub.lib.lua +++ b/plugins/mod_pubsub/pubsub.lib.lua @@ -42,8 +42,8 @@ function handlers.get_items(origin, stanza, items, service) end local data = st.stanza("items", { node = node }); - for _, entry in pairs(results) do - data:add_child(entry); + for _, id in ipairs(results) do + data:add_child(results[id]); end local reply; if data then diff --git a/util/indexedbheap.lua b/util/indexedbheap.lua new file mode 100644 index 00000000..3cb03037 --- /dev/null +++ b/util/indexedbheap.lua @@ -0,0 +1,153 @@ + +local setmetatable = setmetatable; +local math_floor = math.floor; +local t_remove = table.remove; + +local function _heap_insert(self, item, sync, item2, index) + local pos = #self + 1; + while true do + local half_pos = math_floor(pos / 2); + if half_pos == 0 or item > self[half_pos] then break; end + self[pos] = self[half_pos]; + sync[pos] = sync[half_pos]; + index[sync[pos]] = pos; + pos = half_pos; + end + self[pos] = item; + sync[pos] = item2; + index[item2] = pos; +end + +local function _percolate_up(self, k, sync, index) + local tmp = self[k]; + local tmp_sync = sync[k]; + while k ~= 1 do + local parent = math_floor(k/2); + if tmp < self[parent] then break; end + self[k] = self[parent]; + sync[k] = sync[parent]; + index[sync[k]] = k; + k = parent; + end + self[k] = tmp; + sync[k] = tmp_sync; + index[tmp_sync] = k; + return k; +end + +local function _percolate_down(self, k, sync, index) + local tmp = self[k]; + local tmp_sync = sync[k]; + local size = #self; + local child = 2*k; + while 2*k <= size do + if child ~= size and self[child] > self[child + 1] then + child = child + 1; + end + if tmp > self[child] then + self[k] = self[child]; + sync[k] = sync[child]; + index[sync[k]] = k; + else + break; + end + + k = child; + child = 2*k; + end + self[k] = tmp; + sync[k] = tmp_sync; + index[tmp_sync] = k; + return k; +end + +local function _heap_pop(self, sync, index) + local size = #self; + if size == 0 then return nil; end + + local result = self[1]; + local result_sync = sync[1]; + index[result_sync] = nil; + if size == 1 then + self[1] = nil; + sync[1] = nil; + return result, result_sync; + end + self[1] = t_remove(self); + sync[1] = t_remove(sync); + index[sync[1]] = 1; + + _percolate_down(self, 1, sync, index); + + return result, result_sync; +end + +local indexed_heap = {}; + +function indexed_heap:insert(item, priority, id) + if id == nil then + id = self.current_id; + self.current_id = id + 1; + end + self.items[id] = item; + _heap_insert(self.priorities, priority, self.ids, id, self.index); + return id; +end +function indexed_heap:pop() + local priority, id = _heap_pop(self.priorities, self.ids, self.index); + if id then + local item = self.items[id]; + self.items[id] = nil; + return priority, item, id; + end +end +function indexed_heap:peek() + return self.priorities[1]; +end +function indexed_heap:reprioritize(id, priority) + local k = self.index[id]; + if k == nil then return; end + self.priorities[k] = priority; + + k = _percolate_up(self.priorities, k, self.ids, self.index); + k = _percolate_down(self.priorities, k, self.ids, self.index); +end +function indexed_heap:remove_index(k) + local size = #self.priorities; + + local result = self.priorities[k]; + local result_sync = self.ids[k]; + local item = self.items[result_sync]; + if result == nil then return; end + self.index[result_sync] = nil; + self.items[result_sync] = nil; + + self.priorities[k] = self.priorities[size]; + self.ids[k] = self.ids[size]; + self.index[self.ids[k]] = k; + t_remove(self.priorities); + t_remove(self.ids); + + k = _percolate_up(self.priorities, k, self.ids, self.index); + k = _percolate_down(self.priorities, k, self.ids, self.index); + + return result, item, result_sync; +end +function indexed_heap:remove(id) + return self:remove_index(self.index[id]); +end + +local mt = { __index = indexed_heap }; + +local _M = { + create = function() + return setmetatable({ + ids = {}; -- heap of ids, sync'd with priorities + items = {}; -- map id->items + priorities = {}; -- heap of priorities + index = {}; -- map of id->index of id in ids + current_id = 1.5 + }, mt); + end +}; +return _M; diff --git a/util/timer.lua b/util/timer.lua index 0e10e144..23bd6a37 100644 --- a/util/timer.lua +++ b/util/timer.lua @@ -6,6 +6,8 @@ -- COPYING file in the source package for more information. -- +local indexedbheap = require "util.indexedbheap"; +local log = require "util.logger".init("timer"); local server = require "net.server"; local math_min = math.min local math_huge = math.huge @@ -13,6 +15,9 @@ local get_time = require "socket".gettime; local t_insert = table.insert; local pairs = pairs; local type = type; +local debug_traceback = debug.traceback; +local tostring = tostring; +local xpcall = xpcall; local data = {}; local new_data = {}; @@ -78,6 +83,61 @@ else end end -add_task = _add_task; +--add_task = _add_task; + +local h = indexedbheap.create(); +local params = {}; +local next_time = nil; +local _id, _callback, _now, _param; +local function _call() return _callback(_now, _id, _param); end +local function _traceback_handler(err) log("error", "Traceback[timer]: %s", debug_traceback(tostring(err), 2)); end +local function _on_timer(now) + local peek; + while true do + peek = h:peek(); + if peek == nil or peek > now then break; end + local _; + _, _callback, _id = h:pop(); + _now = now; + _param = params[_id]; + params[_id] = nil; + --item(now, id, _param); -- FIXME pcall + local success, err = xpcall(_call, _traceback_handler); + if success and type(err) == "number" then + h:insert(_callback, err + now, _id); -- re-add + params[_id] = _param; + end + end + next_time = peek; + if peek ~= nil then + return peek - now; + end +end +function add_task(delay, callback, param) + local current_time = get_time(); + local event_time = current_time + delay; + + local id = h:insert(callback, event_time); + params[id] = param; + if next_time == nil or event_time < next_time then + next_time = event_time; + _add_task(next_time - current_time, _on_timer); + end + return id; +end +function stop(id) + params[id] = nil; + return h:remove(id); +end +function reschedule(id, delay) + local current_time = get_time(); + local event_time = current_time + delay; + h:reprioritize(id, delay); + if next_time == nil or event_time < next_time then + next_time = event_time; + _add_task(next_time - current_time, _on_timer); + end + return id; +end return _M; |